Yicong-Huang commented on code in PR #58962: URL: https://github.com/apache/spark/pull/58962#discussion_r4086705231
########## sql/connect/common/src/main/scala/org/apache/spark/sql/connect/common/UdfSerialization.scala: ########## @@ -0,0 +1,108 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.connect.common + +import java.io.{ByteArrayInputStream, InputStream, ObjectInputStream, ObjectStreamClass} + +/** + * Java deserialization for Scala UDF payloads that tolerates the `serialVersionUID` drift of + * `org.apache.spark.sql.types` classes between Spark versions. + * + * A Scala UDF payload embeds `org.apache.spark.sql.types` classes (the UDF's input/output schema). + * Those classes carry no explicit `@SerialVersionUID`, so the JVM auto-computes it from the whole + * class shape, which folds in compiler-synthesized members that are irrelevant to serialization -- + * most notably the `$anonfun$` public static methods Scala emits for lambdas. A source change that + * only reshapes a lambda (e.g. rewriting a helper to use `existsRecursively { ... }`) changes the + * auto-computed `serialVersionUID` without changing any serialized field, which makes a plain + * [[ObjectInputStream]] reject a payload produced by a different Spark version with an + * `InvalidClassException`, even though the payload is field-compatible. + * + * The property that actually governs compatibility is the serialized field layout. When a + * `sql.types` class arrives with a mismatched SUID but an identical serialized field layout, this + * reader rebinds the stream descriptor to the local class; any field-layout difference is left + * untouched so the standard SUID check still fails fast rather than misreading the stream. + * + * Only `org.apache.spark.sql.types` descriptors are treated tolerantly; every other class keeps + * the standard `serialVersionUID` compatibility check. + */ +private[spark] object UdfSerialization { + + private val suidTolerantPackagePrefix = "org.apache.spark.sql.types." + + /** Deserialize `bytes` resolving classes with `loader`, tolerating `sql.types` SUID drift. */ + def deserialize[T](bytes: Array[Byte], loader: ClassLoader): T = { + val ois = new SuidTolerantObjectInputStream(new ByteArrayInputStream(bytes), loader) + try ois.readObject().asInstanceOf[T] + finally ois.close() + } + + /** Deserialize from `in` with default class resolution, tolerating `sql.types` SUID drift. */ + def deserialize[T](in: InputStream): T = { + new SuidTolerantObjectInputStream(in, null).readObject().asInstanceOf[T] + } + + /** + * The complete serialized field layout of a descriptor: the set of persistent field name + JVM + * type signature. This is what governs whether [[ObjectInputStream]] can consume the producer's + * class-data block through the local descriptor. Every persistent slot is included, in + * particular the Scala lazy-val init `bitmap$*` slots: rebinding to a local descriptor whose + * slot shape differs would misalign the stream, so a bitmap difference must also block rebinding. + */ + private[connect] def fieldSignature(desc: ObjectStreamClass): Set[String] = { + // getTypeString is null for primitives, where the single-char type code is the signature. + desc.getFields + .map(f => s"${f.getName}:${Option(f.getTypeString).getOrElse(f.getTypeCode.toString)}") + .toSet + } + + private class SuidTolerantObjectInputStream(in: InputStream, loader: ClassLoader) + extends ObjectInputStream(in) { + + override def resolveClass(desc: ObjectStreamClass): Class[_] = { + if (loader != null) { + // scalastyle:off classforname + Class.forName(desc.getName, false, loader) + // scalastyle:on classforname + } else { + super.resolveClass(desc) + } + } + + override def readClassDescriptor(): ObjectStreamClass = { + val streamDesc = super.readClassDescriptor() + if (!streamDesc.getName.startsWith(suidTolerantPackagePrefix)) { + return streamDesc + } + val localClass = + try { + resolveClass(streamDesc) + } catch { + case _: ClassNotFoundException => return streamDesc + } + val localDesc = ObjectStreamClass.lookup(localClass) + if (localDesc == null || + localDesc.getSerialVersionUID == streamDesc.getSerialVersionUID || + fieldSignature(streamDesc) != fieldSignature(localDesc)) { Review Comment: Agreed, a local-only predicate cannot prove the producer side. Switched to an exact audited table of `(className, streamSuid, localSuid)` transitions. Rebinding now requires the triple to be audited, and additionally that the local class passes `isRebindSafe` (now also excluding `Externalizable`) and that the complete persistent field layout matches. Any unlisted pair, including an unchanged SUID, keeps `streamDesc`. The table is audited against the published `spark-sql-api_2.13` jars of 4.0.0 to 4.0.4, 4.1.0 to 4.1.3 and 4.2.0 against this build: 24 classes whose computed SUID changed while both sides use default serialization, declare no `serialVersionUID`, and have identical `ObjectStreamClass.getFields`. A test checks that every entry's `localSuid` matches this build, so any later change to these classes forces a re-audit. `StructField`, `StructType`, `CharType` and `VarcharType` are not listed because their persistent layout differs from 4.x. Added the asymmetric fixtures: a producer with a custom `writeObject` and a producer with an explicit `serialVersionUID`, each read by a default-serialized consumer, are both rejected. The audited cases are also rejected when the layout differs, the consumer declares a SUID or custom `readObject`, or the class is outside `sql.types`. ########## sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/common/UdfSerializationSuite.scala: ########## @@ -0,0 +1,107 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.connect.common + +import java.io.{ByteArrayInputStream, ByteArrayOutputStream, InvalidClassException} +import java.io.{ObjectInputStream, ObjectOutputStream} + +import org.apache.spark.sql.connect.test.ConnectFunSuite +import org.apache.spark.sql.types.{SuidCompatV1, SuidCompatV2, SuidCustomV1, SuidExplicitV1} +import org.apache.spark.sql.types.{SuidLayoutV1, SuidLayoutV2} + +/** + * Tests for [[UdfSerialization]]'s tolerance of `serialVersionUID` drift for + * `org.apache.spark.sql.types` classes. + * + * Each fixture pair has same-length class names, so serializing the `*V1` instance and patching + * the class name in the byte stream to `*V2` yields exactly what a `*V2` reader would see from a + * `*V1` producer whose `serialVersionUID` differs -- a deterministic stand-in for a cross-version + * payload without needing two builds. + */ +class UdfSerializationSuite extends ConnectFunSuite { + + private def serialize(o: AnyRef): Array[Byte] = { + val bos = new ByteArrayOutputStream() + val oos = new ObjectOutputStream(bos) + oos.writeObject(o) + oos.close() + bos.toByteArray + } + + /** Replace the ASCII `from` class name with the same-length `to` in a serialized stream. */ + private def patchClassName(bytes: Array[Byte], from: String, to: String): Array[Byte] = { + assert(from.length == to.length, "class names must be the same length for an in-place patch") + val fromBytes = from.getBytes("US-ASCII") + val toBytes = to.getBytes("US-ASCII") + val idx = bytes.indexOfSlice(fromBytes) + assert(idx >= 0, s"$from not found in the serialized stream") + val patched = bytes.clone() + System.arraycopy(toBytes, 0, patched, idx, toBytes.length) + patched + } + + private val loader = getClass.getClassLoader + + test("tolerates serialVersionUID drift for a sql.types class (bytes overload)") { Review Comment: Added `UdfPacket.apply` and `SparkConnectPlanner.unpackScalaUDF` cases. Each builds a `UdfPacket` that references a production-audited object, rewrites that object's SUID in the stream to the audited release value, and checks that a plain `ObjectInputStream` rejects the stream while the entry point accepts it, so reverting either delegation fails. `unpackScalaUDF` is now `private[connect]` for this. The suite moved to the server module so it can reach the planner. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
