diff --git a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/language/AnalyzeTextLROSuite.scala b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/language/AnalyzeTextLROSuite.scala index 3d1846a535f..b276251c884 100644 --- a/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/language/AnalyzeTextLROSuite.scala +++ b/cognitive/src/test/scala/com/microsoft/azure/synapse/ml/services/language/AnalyzeTextLROSuite.scala @@ -364,7 +364,7 @@ class SentimentAnalysisLROSuite extends TransformerFuzzing[AnalyzeTextLongRunnin override def testObjects(): Seq[TestObject[AnalyzeTextLongRunningOperations]] = Seq(new TestObject[AnalyzeTextLongRunningOperations](model, df)) - override def reader: MLReadable[_] = AnalyzeText + override def reader: MLReadable[_] = AnalyzeTextLongRunningOperations } @@ -602,7 +602,7 @@ class EntityRecognitionLROSuite extends TransformerFuzzing[AnalyzeTextLongRunnin override def testObjects(): Seq[TestObject[AnalyzeTextLongRunningOperations]] = Seq(new TestObject[AnalyzeTextLongRunningOperations](model, df)) - override def reader: MLReadable[_] = AnalyzeText + override def reader: MLReadable[_] = AnalyzeTextLongRunningOperations } class CustomEntityRecognitionSuite extends TransformerFuzzing[AnalyzeTextLongRunningOperations] @@ -648,7 +648,7 @@ class CustomEntityRecognitionSuite extends TransformerFuzzing[AnalyzeTextLongRun override def testObjects(): Seq[TestObject[AnalyzeTextLongRunningOperations]] = Seq(new TestObject[AnalyzeTextLongRunningOperations](model, df)) - override def reader: MLReadable[_] = AnalyzeText + override def reader: MLReadable[_] = AnalyzeTextLongRunningOperations } @@ -697,8 +697,7 @@ class MultiLableClassificationSuite extends TransformerFuzzing[AnalyzeTextLongRu override def testObjects(): Seq[TestObject[AnalyzeTextLongRunningOperations]] = Seq(new TestObject[AnalyzeTextLongRunningOperations](model, df)) - override def reader: MLReadable[_] = AnalyzeText + override def reader: MLReadable[_] = AnalyzeTextLongRunningOperations } - diff --git a/core/src/main/python/synapse/ml/core/schema/Utils.py b/core/src/main/python/synapse/ml/core/schema/Utils.py index aa945e3b059..9d4d1605404 100644 --- a/core/src/main/python/synapse/ml/core/schema/Utils.py +++ b/core/src/main/python/synapse/ml/core/schema/Utils.py @@ -10,6 +10,7 @@ from pyspark.ml.wrapper import JavaParams from pyspark.ml.common import inherit_doc, _java2py from pyspark import SparkContext +from pyspark.sql import SparkSession from synapse.ml.core.serialize._safe_import import secure_import_class @@ -58,6 +59,11 @@ def read(cls): @inherit_doc class ComplexParamsMixin(MLReadable): + @classmethod + def read(cls): + """Returns a reader bound to the active Spark session.""" + return JavaMMLReader(cls) + def _transfer_params_from_java(self): """ Transforms the embedded com.microsoft.azure.synapse.ml.core.serialize.params from the companion Java object. @@ -131,6 +137,7 @@ class JavaMMLReader(JavaMLReader): def __init__(self, clazz): super(JavaMMLReader, self).__init__(clazz) + self.session(SparkSession.builder.getOrCreate()) @classmethod def _java_loader_class(cls, clazz): diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/core/serialize/ComplexParam.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/core/serialize/ComplexParam.scala index 58168b9e9e3..f4d92348562 100644 --- a/core/src/main/scala/com/microsoft/azure/synapse/ml/core/serialize/ComplexParam.scala +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/core/serialize/ComplexParam.scala @@ -3,6 +3,7 @@ package com.microsoft.azure.synapse.ml.core.serialize +import com.microsoft.azure.synapse.ml.core.utils.DeserializationClassFilter import com.microsoft.azure.synapse.ml.param.WrappableParam import org.apache.hadoop.fs.Path import org.apache.spark.ml.Serializer @@ -16,12 +17,36 @@ abstract class ComplexParam[T: TypeTag](parent: Params, name: String, doc: Strin def ttag: TypeTag[T] = typeTag[T] + /** Class policy for legacy Java object streams. No policy means loading is disabled unless the + * Spark session explicitly opts into trusted legacy deserialization. + */ + protected def deserializationClassFilter: Option[DeserializationClassFilter] = None + + protected def supportsUntrustedDeserialization: Boolean = true + + def isSafeForUntrustedDeserialization: Boolean = { + supportsUntrustedDeserialization && + (!Serializer.usesObjectSerializer(ttag.tpe) || deserializationClassFilter.isDefined) + } + def save(obj: T, sparkSession: SparkSession, path: Path, overwrite: Boolean): Unit = { - Serializer.typeToSerializer[T](ttag.tpe, sparkSession).write(obj, path, overwrite) + Serializer.typeToSerializer[T](ttag.tpe, sparkSession, deserializationClassFilter) + .write(obj, path, overwrite) } def load(sparkSession: SparkSession, path: Path): T = { - Serializer.typeToSerializer[T](ttag.tpe, sparkSession).read(path) + if ( + !isSafeForUntrustedDeserialization && + !Serializer.trustedLoadEnabled(sparkSession) + ) { + throw new SecurityException( + s"Complex parameter $name requires a trusted artifact. Set " + + s"${Serializer.LegacyObjectDeserializationConfig}=true and load through " + + "read.session(sparkSession).load(path), or wrap a native Pipeline load in " + + "Serializer.withTrustedArtifactLoad(sparkSession), only when loading trusted data." + ) + } + Serializer.typeToSerializer[T](ttag.tpe, sparkSession, deserializationClassFilter).read(path) } override def jsonEncode(value: T): String = { diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/core/utils/Jep290ObjectInputFilter.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/core/utils/Jep290ObjectInputFilter.scala new file mode 100644 index 00000000000..d1610156bdf --- /dev/null +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/core/utils/Jep290ObjectInputFilter.scala @@ -0,0 +1,288 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.core.utils + +import java.io.{InvalidClassException, ObjectInputStream} +import java.lang.reflect.{InvocationHandler, InvocationTargetException, Method, Proxy} + +import scala.util.{Failure, Success, Try} + +private[utils] sealed trait DeserializationFilterStatus { + private[utils] def runtimeName: String +} + +private[utils] object DeserializationFilterStatus { + case object Allowed extends DeserializationFilterStatus { + override private[utils] val runtimeName: String = "ALLOWED" + } + + case object Rejected extends DeserializationFilterStatus { + override private[utils] val runtimeName: String = "REJECTED" + } + + case object Undecided extends DeserializationFilterStatus { + override private[utils] val runtimeName: String = "UNDECIDED" + } +} + +private[utils] final case class DeserializationFilterInfo( + serialClass: Class[_], + arrayLength: Long, + depth: Long, + references: Long, + streamBytes: Long) + +private[utils] trait DeserializationInputFilter { + def checkInput(info: DeserializationFilterInfo): DeserializationFilterStatus +} + +/** Installs a per-stream JEP 290 filter without linking against a post-Java 8 API. + * + * Java 9 moved ObjectInputFilter from sun.misc to java.io. SynapseML artifacts + * still compile on Java 8, so the runtime-specific interface is implemented by + * a dynamic proxy while the security policy remains statically typed. + */ +private[utils] object Jep290ObjectInputFilter { + + import DeserializationFilterStatus._ + + private val ModernFilterClass = "java.io.ObjectInputFilter" + private val LegacyFilterClass = "sun.misc.ObjectInputFilter" + private val LegacyConfigClass = "sun.misc.ObjectInputFilter$Config" + + private lazy val ApiResult: Try[FilterApi] = { + Try(FilterApi.modern()).recoverWith { + case _: ClassNotFoundException => Try(FilterApi.legacy()) + } + } + + def install(stream: ObjectInputStream, resourceFilter: DeserializationInputFilter): Unit = { + val api = ApiResult match { + case Success(value) => value + case Failure(error) => throw installationFailure(error) + } + + try { + val inheritedFilter = Option(api.getFilter(stream)) + val proxy = Proxy.newProxyInstance( + classOf[SafeObjectInputStream].getClassLoader, + Array(api.filterClass), + new FilterInvocationHandler(api, inheritedFilter, resourceFilter) + ) + api.setFilter(stream, proxy) + } catch { + case error: ReflectiveOperationException => throw installationFailure(error) + case error: SecurityException => throw installationFailure(error) + case error: IllegalArgumentException => throw installationFailure(error) + } + } + + private[utils] def composeStatuses( + inheritedStatus: DeserializationFilterStatus, + resourceStatus: DeserializationFilterStatus + ): DeserializationFilterStatus = { + if (inheritedStatus == Rejected || resourceStatus == Rejected) { + Rejected + } else if (inheritedStatus == Allowed && resourceStatus == Allowed) { + Allowed + } else { + Undecided + } + } + + private def installationFailure(cause: Throwable): InvalidClassException = { + val error = new InvalidClassException( + "Safe Java deserialization requires JEP 290 ObjectInputFilter support" + ) + error.initCause(cause) + error + } + + private def invokeReflectively( + method: Method, + receiver: AnyRef, + arguments: AnyRef* + ): AnyRef = { + try { + method.invoke(receiver, arguments: _*) + } catch { + case error: InvocationTargetException => + throw Option(error.getCause).getOrElse(error) + } + } + + private final class FilterInvocationHandler( + api: FilterApi, + inheritedFilter: Option[AnyRef], + resourceFilter: DeserializationInputFilter + ) extends InvocationHandler { + + override def invoke( + proxy: Any, + method: Method, + arguments: Array[AnyRef] + ): AnyRef = { + method.getName match { + case "checkInput" => checkInput(arguments) + case _ => invokeObjectMethod(proxy, method, arguments) + } + } + + private def checkInput(arguments: Array[AnyRef]): AnyRef = { + if (arguments == null || arguments.length != 1) { + throw new IllegalArgumentException("ObjectInputFilter.checkInput requires one argument") + } + val runtimeInfo = arguments(0) + val inheritedStatus = inheritedFilter.map(api.checkInput(_, runtimeInfo)) + val resourceStatus = resourceFilter.checkInput(api.toFilterInfo(runtimeInfo)) + val combinedStatus = inheritedStatus match { + case Some(status) => composeStatuses(status, resourceStatus) + case None => resourceStatus + } + api.toRuntimeStatus(combinedStatus) + } + + private def invokeObjectMethod( + proxy: Any, + method: Method, + arguments: Array[AnyRef] + ): AnyRef = { + method.getName match { + case "equals" => Boolean.box(isSameProxy(proxy, arguments)) + case "hashCode" => Int.box(System.identityHashCode(proxy)) + case "toString" => "SynapseML JEP 290 deserialization filter" + case name => + throw new UnsupportedOperationException(s"Unsupported ObjectInputFilter method: $name") + } + } + + private def isSameProxy(proxy: Any, arguments: Array[AnyRef]): Boolean = { + arguments != null && + arguments.length == 1 && + (proxy.asInstanceOf[AnyRef] eq arguments(0)) + } + } + + private final class FilterApi( + val filterClass: Class[_], + filterInfoClass: Class[_], + statusClass: Class[_], + getFilterMethod: Method, + setFilterMethod: Method, + staticAccessReceiver: Option[AnyRef]) { + + private val checkInputMethod = filterClass.getMethod("checkInput", filterInfoClass) + private val serialClassMethod = filterInfoClass.getMethod("serialClass") + private val arrayLengthMethod = filterInfoClass.getMethod("arrayLength") + private val depthMethod = filterInfoClass.getMethod("depth") + private val referencesMethod = filterInfoClass.getMethod("references") + private val streamBytesMethod = filterInfoClass.getMethod("streamBytes") + + private val allowedStatus = runtimeStatus(Allowed) + private val rejectedStatus = runtimeStatus(Rejected) + private val undecidedStatus = runtimeStatus(Undecided) + + def getFilter(stream: ObjectInputStream): AnyRef = { + staticAccessReceiver match { + case Some(receiver) => invokeReflectively(getFilterMethod, receiver, stream) + case None => invokeReflectively(getFilterMethod, stream) + } + } + + def setFilter(stream: ObjectInputStream, filter: AnyRef): Unit = { + staticAccessReceiver match { + case Some(receiver) => invokeReflectively(setFilterMethod, receiver, stream, filter) + case None => invokeReflectively(setFilterMethod, stream, filter) + } + } + + def checkInput(filter: AnyRef, runtimeInfo: AnyRef): DeserializationFilterStatus = { + fromRuntimeStatus(invokeReflectively(checkInputMethod, filter, runtimeInfo)) + } + + def toFilterInfo(runtimeInfo: AnyRef): DeserializationFilterInfo = { + DeserializationFilterInfo( + serialClass = invokeReflectively(serialClassMethod, runtimeInfo).asInstanceOf[Class[_]], + arrayLength = longValue(arrayLengthMethod, runtimeInfo), + depth = longValue(depthMethod, runtimeInfo), + references = longValue(referencesMethod, runtimeInfo), + streamBytes = longValue(streamBytesMethod, runtimeInfo) + ) + } + + def toRuntimeStatus(status: DeserializationFilterStatus): AnyRef = { + status match { + case Allowed => allowedStatus + case Rejected => rejectedStatus + case Undecided => undecidedStatus + } + } + + private def fromRuntimeStatus(status: AnyRef): DeserializationFilterStatus = { + if (allowedStatus eq status) { + Allowed + } else if (rejectedStatus eq status) { + Rejected + } else if (undecidedStatus eq status) { + Undecided + } else { + throw new IllegalStateException("ObjectInputFilter returned an unknown status") + } + } + + private def runtimeStatus(status: DeserializationFilterStatus): AnyRef = { + statusClass.getField(status.runtimeName).get(statusClass) + } + + private def longValue(method: Method, runtimeInfo: AnyRef): Long = { + invokeReflectively(method, runtimeInfo).asInstanceOf[java.lang.Long].longValue() + } + } + + private object FilterApi { + + def modern(): FilterApi = { + val filterClass = Class.forName(ModernFilterClass) + create( + filterClass, + classOf[ObjectInputStream].getMethod("getObjectInputFilter"), + classOf[ObjectInputStream].getMethod("setObjectInputFilter", filterClass), + staticAccessReceiver = None + ) + } + + def legacy(): FilterApi = { + val filterClass = Class.forName(LegacyFilterClass) + val configClass = Class.forName(LegacyConfigClass) + create( + filterClass, + configClass.getMethod("getObjectInputFilter", classOf[ObjectInputStream]), + configClass.getMethod( + "setObjectInputFilter", + classOf[ObjectInputStream], + filterClass + ), + staticAccessReceiver = Some(configClass) + ) + } + + private def create( + filterClass: Class[_], + getFilterMethod: Method, + setFilterMethod: Method, + staticAccessReceiver: Option[AnyRef] + ): FilterApi = { + val filterInfoClass = Class.forName(s"${filterClass.getName}$$FilterInfo") + val statusClass = Class.forName(s"${filterClass.getName}$$Status") + new FilterApi( + filterClass, + filterInfoClass, + statusClass, + getFilterMethod, + setFilterMethod, + staticAccessReceiver + ) + } + } +} diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/core/utils/SafeObjectInputStream.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/core/utils/SafeObjectInputStream.scala index 2a6b1ef46ec..dab123350e3 100644 --- a/core/src/main/scala/com/microsoft/azure/synapse/ml/core/utils/SafeObjectInputStream.scala +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/core/utils/SafeObjectInputStream.scala @@ -3,21 +3,62 @@ package com.microsoft.azure.synapse.ml.core.utils -import java.io.{InputStream, InvalidClassException, ObjectStreamClass} +import java.io.{ + InputStream, + InvalidClassException, + InvalidObjectException, + ObjectStreamClass, + StreamCorruptedException +} + +final case class DeserializationClassFilter( + allowedPrefixes: Set[String] = Set.empty, + allowedClasses: Set[String] = Set.empty) { + + require(!allowedPrefixes.contains(""), "Deserialization class prefixes cannot contain an empty prefix") + + private[utils] def allows(className: String): Boolean = { + allowedClasses.contains(className) || allowedPrefixes.exists(className.startsWith) + } +} -/** An ObjectInputStream that restricts deserialization to an allowlist of class name prefixes. +/** An ObjectInputStream that restricts deserialization to an allowlist of class names and prefixes. * * This mitigates Java deserialization attacks (CWE-502) by rejecting any class - * whose fully-qualified name does not start with one of the allowed prefixes. + * whose fully-qualified name is not explicitly allowed. * It also inherits the context-classloader resolution from [[ContextObjectInputStream]]. * - * @param input the underlying input stream - * @param allowedPrefixes set of class name prefixes that are permitted for deserialization + * @param input the underlying input stream + * @param classFilter exact class names and class-name prefixes permitted for deserialization */ -class SafeObjectInputStream( +class SafeObjectInputStream private[utils] ( input: InputStream, - allowedPrefixes: Set[String] -) extends ContextObjectInputStream(input) { + classFilter: DeserializationClassFilter, + resourceLimits: DeserializationResourceLimits +) extends ContextObjectInputStream( + new BoundedDeserializationInputStream(input, resourceLimits.maxStreamBytes) +) { + + Jep290ObjectInputFilter.install( + this, + SafeObjectInputStream.newResourceFilter(resourceLimits) + ) + enableResolveObject(true) + + private val alwaysRejectedClasses = Set( + "java.lang.invoke.SerializedLambda" + ) + + private var resolvedObjects = 0L + private var resolvedStringBytes = 0L + + def this(input: InputStream, classFilter: DeserializationClassFilter) = { + this(input, classFilter, SafeObjectInputStream.defaultResourceLimits) + } + + def this(input: InputStream, allowedPrefixes: Set[String]) = { + this(input, DeserializationClassFilter(allowedPrefixes = allowedPrefixes)) + } /** Extracts the component type name from a JVM array descriptor. * Primitive arrays (e.g. `[I`, `[D`) return None since they are always safe. @@ -34,7 +75,7 @@ class SafeObjectInputStream( } private def isAllowed(className: String): Boolean = { - allowedPrefixes.exists(prefix => className.startsWith(prefix)) + !alwaysRejectedClasses.contains(className) && classFilter.allows(className) } protected override def resolveClass(desc: ObjectStreamClass): Class[_] = { @@ -58,6 +99,34 @@ class SafeObjectInputStream( super.resolveClass(desc) } + protected override def resolveObject(value: AnyRef): AnyRef = { + if (resolvedObjects >= resourceLimits.maxResolvedObjects) { + throw new InvalidObjectException( + s"Deserialized object count exceeds ${resourceLimits.maxResolvedObjects}" + ) + } + resolvedObjects += 1 + value match { + case stringValue: String => + if (!isAllowed(classOf[String].getName)) { + throw new InvalidClassException( + classOf[String].getName, + "Deserialization of this class is not allowed. " + + "Only classes allowed by the configured class policy may be deserialized." + ) + } + val stringBytes = stringValue.length.toLong * Character.BYTES + if (stringBytes > resourceLimits.maxStringBytes - resolvedStringBytes) { + throw new InvalidObjectException( + s"Deserialized String data exceeds ${resourceLimits.maxStringBytes} bytes" + ) + } + resolvedStringBytes += stringBytes + case _ => + } + super.resolveObject(value) + } + /** Rejects dynamic proxy deserialization unless every interface is allowlisted. * * Dynamic proxies are a known deserialization attack vector (e.g. via @@ -80,16 +149,199 @@ class SafeObjectInputStream( object SafeObjectInputStream { - /** Default allowlist suitable for deserializing SynapseML nn package objects - * (BallTree, ConditionalBallTree, and their object graphs). - */ - val DefaultNNAllowedPrefixes: Set[String] = Set( - "com.microsoft.azure.synapse.ml.nn.", - "breeze.", - "scala.", + private val MaxDeserializationDepth = 100L // scalastyle:ignore magic.number + private val MaxDeserializationReferences = 1000000L // scalastyle:ignore magic.number + private val AbsoluteMaxAllocationBytes = 1L << 30 // scalastyle:ignore magic.number + private val AbsoluteMaxStreamBytes = 1L << 29 // scalastyle:ignore magic.number + private val MinimumResourceLimit = 1L << 20 // scalastyle:ignore magic.number + + private def allocationLimit: Long = { + Math.max( + MinimumResourceLimit, + Math.min(AbsoluteMaxAllocationBytes, Runtime.getRuntime.maxMemory() / 4) + ) + } + + private def streamLimit: Long = { + Math.max( + MinimumResourceLimit, + Math.min(AbsoluteMaxStreamBytes, Runtime.getRuntime.maxMemory() / 4) + ) + } + + private def referenceLimit: Long = { + Math.max( + 10000L, // scalastyle:ignore magic.number + Math.min(MaxDeserializationReferences, Runtime.getRuntime.maxMemory() / 64) // scalastyle:ignore magic.number + ) + } + + private[utils] def defaultResourceLimits: DeserializationResourceLimits = { + val references = referenceLimit + DeserializationResourceLimits( + maxDepth = MaxDeserializationDepth, + maxReferences = references, + maxStreamBytes = streamLimit, + maxArrayBytes = allocationLimit, + maxResolvedObjects = references, + maxStringBytes = allocationLimit + ) + } + + private[utils] def newResourceFilter( + resourceLimits: DeserializationResourceLimits = defaultResourceLimits + ): DeserializationInputFilter = { + new DeserializationResourceFilter(resourceLimits) + } + + val CommonDataAllowedPrefixes: Set[String] = Set( "java.lang.", + "java.math.", "java.util.", - "java.io.", - "java.math." + "scala." + ) + + /** Retained for binary compatibility. Generic BallTree values make a safe package-prefix + * allowlist impossible, so the default policy intentionally permits no object classes. + */ + val DefaultNNAllowedPrefixes: Set[String] = Set.empty + + /** Retained for binary compatibility. Legacy BallTree streams require explicit trusted loading. */ + val DefaultNNFilter: DeserializationClassFilter = DeserializationClassFilter() +} + +private[utils] final case class DeserializationResourceLimits( + maxDepth: Long, + maxReferences: Long, + maxStreamBytes: Long, + maxArrayBytes: Long, + maxResolvedObjects: Long, + maxStringBytes: Long) { + require( + Seq( + maxDepth, + maxReferences, + maxStreamBytes, + maxArrayBytes, + maxResolvedObjects, + maxStringBytes + ).forall(_ > 0), + "Deserialization resource limits must be positive" ) } + +private[utils] final class DeserializationResourceFilter( + limits: DeserializationResourceLimits +) extends DeserializationInputFilter { + + private var declaredArrayBytes = 0L + + private def arrayElementSize(arrayClass: Class[_]): Long = { + val component = arrayClass.getComponentType + if (!component.isPrimitive) { + 8L // scalastyle:ignore magic.number + } else if (component == java.lang.Long.TYPE || component == java.lang.Double.TYPE) { + 8L // scalastyle:ignore magic.number + } else if (component == java.lang.Integer.TYPE || component == java.lang.Float.TYPE) { + 4L // scalastyle:ignore magic.number + } else if (component == java.lang.Character.TYPE || component == java.lang.Short.TYPE) { + 2L // scalastyle:ignore magic.number + } else { + 1L + } + } + + override def checkInput(info: DeserializationFilterInfo): DeserializationFilterStatus = { + try { + val arrayBytes = Option(info.serialClass).flatMap { serializedClass => + if (serializedClass.isArray && info.arrayLength >= 0) { + Some(Math.multiplyExact(info.arrayLength, arrayElementSize(serializedClass))) + } else { + None + } + } + val updatedArrayBytes = Math.addExact(declaredArrayBytes, arrayBytes.getOrElse(0L)) + if (info.depth > limits.maxDepth || + info.references > limits.maxReferences || + info.streamBytes > limits.maxStreamBytes || + updatedArrayBytes > limits.maxArrayBytes) { + DeserializationFilterStatus.Rejected + } else { + declaredArrayBytes = updatedArrayBytes + DeserializationFilterStatus.Undecided + } + } catch { + case _: ArithmeticException => DeserializationFilterStatus.Rejected + } + } +} + +private[utils] final class BoundedDeserializationInputStream( + input: InputStream, + maxBytes: Long +) extends InputStream { + + require(maxBytes > 0, "Deserialization byte limit must be positive") + + private var bytesRead = 0L + + private def limitExceeded(): Nothing = { + throw new StreamCorruptedException( + s"Serialized input exceeds the $maxBytes byte deserialization limit" + ) + } + + override def read(): Int = { + val value = input.read() + if (value >= 0) { + if (bytesRead >= maxBytes) { + limitExceeded() + } + bytesRead += 1 + } + value + } + + override def read(buffer: Array[Byte], offset: Int, length: Int): Int = { + if (length == 0) { + 0 + } else if (bytesRead >= maxBytes) { + val value = input.read() + if (value < 0) { + -1 + } else { + limitExceeded() + } + } else { + val allowedLength = Math.min(length.toLong, maxBytes - bytesRead).toInt + val count = input.read(buffer, offset, allowedLength) + if (count > 0) { + bytesRead += count + } + count + } + } + + override def skip(length: Long): Long = { + if (length <= 0) { + 0 + } else { + val buffer = new Array[Byte](Math.min(length, 8192L).toInt) // scalastyle:ignore magic.number + var remaining = length + var skipped = 0L + var count = read(buffer, 0, Math.min(remaining, buffer.length).toInt) + while (remaining > 0 && count >= 0) { // scalastyle:ignore while + remaining -= count + skipped += count + count = read(buffer, 0, Math.min(remaining, buffer.length).toInt) + } + skipped + } + } + + override def available(): Int = { + Math.min(input.available().toLong, Math.max(0L, maxBytes - bytesRead)).toInt + } + + override def close(): Unit = input.close() +} diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/nn/BallTree.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/nn/BallTree.scala index ea6ac9f2621..c6e0684aa20 100644 --- a/core/src/main/scala/com/microsoft/azure/synapse/ml/nn/BallTree.scala +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/nn/BallTree.scala @@ -6,7 +6,11 @@ package com.microsoft.azure.synapse.ml.nn import breeze.linalg.{DenseVector, norm, _} import com.microsoft.azure.synapse.ml.core.env.StreamUtilities.using -import com.microsoft.azure.synapse.ml.core.utils.SafeObjectInputStream +import com.microsoft.azure.synapse.ml.core.utils.{ + ContextObjectInputStream, + DeserializationClassFilter, + SafeObjectInputStream +} import java.io.{FileInputStream, FileOutputStream, ObjectOutputStream, Serializable} import scala.collection.JavaConverters._ @@ -172,14 +176,32 @@ object ConditionalBallTree { } def load[L, V](filename: String): ConditionalBallTree[L, V] = { - load(filename, SafeObjectInputStream.DefaultNNAllowedPrefixes) + throw new SecurityException( + "Java deserialization is disabled for ConditionalBallTree because its generic object graph " + + "cannot be constrained safely. Call loadUnsafe only for a trusted legacy artifact." + ) + } + + /** Loads a legacy Java-serialized tree. The caller must already trust the artifact. */ + def loadUnsafe[L, V](filename: String): ConditionalBallTree[L, V] = { + using(new FileInputStream(filename)) { fileIn => + using(new ContextObjectInputStream(fileIn)) { in => + in.readObject().asInstanceOf[ConditionalBallTree[L, V]] + } + }.get.get } def load[L, V](filename: String, allowedPrefixes: Set[String] ): ConditionalBallTree[L, V] = { + load(filename, DeserializationClassFilter(allowedPrefixes = allowedPrefixes)) + } + + def load[L, V]( + filename: String, + classFilter: DeserializationClassFilter): ConditionalBallTree[L, V] = { using(new FileInputStream(filename)) { fileIn => - using(new SafeObjectInputStream(fileIn, allowedPrefixes)) { in => + using(new SafeObjectInputStream(fileIn, classFilter)) { in => in.readObject().asInstanceOf[ConditionalBallTree[L, V]] } }.get.get diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/param/BallTreeParam.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/param/BallTreeParam.scala index 5c4af1c21a3..acca4b776ac 100644 --- a/core/src/main/scala/com/microsoft/azure/synapse/ml/param/BallTreeParam.scala +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/param/BallTreeParam.scala @@ -15,6 +15,8 @@ class BallTreeParam(parent: Params, name: String, doc: String, isValid: BallTree def this(parent: Params, name: String, doc: String) = this(parent, name, doc, (_: BallTree[_]) => true) + override protected def supportsUntrustedDeserialization: Boolean = false + } class ConditionalBallTreeParam(parent: Params, @@ -26,4 +28,6 @@ class ConditionalBallTreeParam(parent: Params, def this(parent: Params, name: String, doc: String) = this(parent, name, doc, (_: ConditionalBallTree[_, _]) => true) + override protected def supportsUntrustedDeserialization: Boolean = false + } diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/param/ByteArrayParam.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/param/ByteArrayParam.scala index c4e42183012..d4eb710ca12 100644 --- a/core/src/main/scala/com/microsoft/azure/synapse/ml/param/ByteArrayParam.scala +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/param/ByteArrayParam.scala @@ -4,6 +4,7 @@ package com.microsoft.azure.synapse.ml.param import com.microsoft.azure.synapse.ml.core.serialize.ComplexParam +import com.microsoft.azure.synapse.ml.core.utils.DeserializationClassFilter import org.apache.spark.ml.param.Params /** Param for ByteArray. Needed as spark has explicit params for many different @@ -15,4 +16,7 @@ class ByteArrayParam(parent: Params, name: String, doc: String, isValid: Array[B def this(parent: Params, name: String, doc: String) = this(parent, name, doc, (_: Array[Byte]) => true) + override protected def deserializationClassFilter: Option[DeserializationClassFilter] = + Some(DeserializationClassFilter()) + } diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/param/DataFrameParam.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/param/DataFrameParam.scala index 5df0771b33a..6a3a6b19a11 100644 --- a/core/src/main/scala/com/microsoft/azure/synapse/ml/param/DataFrameParam.scala +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/param/DataFrameParam.scala @@ -120,6 +120,8 @@ class DataFrameParam(parent: Params, name: String, doc: String, isValid: DataFra with DataFrameEquality with ExternalWrappableParam[DataFrame] { + override protected def supportsUntrustedDeserialization: Boolean = false + def this(parent: Params, name: String, doc: String) = this(parent, name, doc, (_: DataFrame) => true) diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/param/DataTypeParam.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/param/DataTypeParam.scala index f0150042f8f..6cade04666d 100644 --- a/core/src/main/scala/com/microsoft/azure/synapse/ml/param/DataTypeParam.scala +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/param/DataTypeParam.scala @@ -5,7 +5,7 @@ package com.microsoft.azure.synapse.ml.param import com.microsoft.azure.synapse.ml.core.serialize.ComplexParam import org.apache.spark.ml.param.Params -import org.apache.spark.sql.types.{DataType, StructType} +import org.apache.spark.sql.types.DataType /** Param for DataType */ class DataTypeParam(parent: Params, name: String, doc: String, isValid: DataType => Boolean) diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/param/EstimatorArrayParam.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/param/EstimatorArrayParam.scala index c0a5cd70182..fb3dfe69b7d 100644 --- a/core/src/main/scala/com/microsoft/azure/synapse/ml/param/EstimatorArrayParam.scala +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/param/EstimatorArrayParam.scala @@ -4,8 +4,11 @@ package com.microsoft.azure.synapse.ml.param import com.microsoft.azure.synapse.ml.core.serialize.ComplexParam +import org.apache.hadoop.fs.Path import org.apache.spark.ml.Estimator +import org.apache.spark.ml.{PipelineArraySerializer, PipelineStage} import org.apache.spark.ml.param.{ParamPair, Params} +import org.apache.spark.sql.SparkSession import scala.collection.JavaConverters._ @@ -19,4 +22,23 @@ class EstimatorArrayParam(parent: Params, name: String, doc: String, isValid: Ar /** Creates a param pair with the given value (for Java). */ def w(value: java.util.List[Estimator[_]]): ParamPair[Array[Estimator[_]]] = w(value.asScala.toArray) + override def save( + obj: Array[Estimator[_]], + sparkSession: SparkSession, + path: Path, + overwrite: Boolean): Unit = { + new PipelineArraySerializer(sparkSession) + .write(obj.map(stage => stage: PipelineStage), path, overwrite) + } + + override def load(sparkSession: SparkSession, path: Path): Array[Estimator[_]] = { + new PipelineArraySerializer(sparkSession).read(path).map { + case estimator: Estimator[_] => estimator + case stage => + throw new IllegalArgumentException( + s"Expected an Estimator in $path but found ${stage.getClass.getName}" + ) + } + } + } diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/param/EstimatorParam.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/param/EstimatorParam.scala index 1d308a13026..aeeea582fab 100644 --- a/core/src/main/scala/com/microsoft/azure/synapse/ml/param/EstimatorParam.scala +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/param/EstimatorParam.scala @@ -19,9 +19,15 @@ trait PipelineStageWrappable[T <: PipelineStage] override def pyLoadLine(modelNum: Int): String = { s""" - |from pyspark.ml import Pipeline - |${name}Model = Pipeline.load(join(test_data_dir, "model-$modelNum.model", "complexParams", "$name")) - |${name}Model = ${name}Model.getStages()[0] + |from pyspark.ml.wrapper import JavaParams + |_jvm = spark.sparkContext._jvm + |${name}Model = JavaParams._from_java( + | _jvm.org.apache.spark.ml.PipelineSerializer(spark._jsparkSession).read( + | _jvm.org.apache.hadoop.fs.Path( + | join(test_data_dir, "model-$modelNum.model", "complexParams", "$name") + | ) + | ) + |) |""".stripMargin } diff --git a/core/src/main/scala/com/microsoft/azure/synapse/ml/param/TransformerArrayParam.scala b/core/src/main/scala/com/microsoft/azure/synapse/ml/param/TransformerArrayParam.scala index b865a7fa116..a9c3a6c7da7 100644 --- a/core/src/main/scala/com/microsoft/azure/synapse/ml/param/TransformerArrayParam.scala +++ b/core/src/main/scala/com/microsoft/azure/synapse/ml/param/TransformerArrayParam.scala @@ -4,8 +4,10 @@ package com.microsoft.azure.synapse.ml.param import com.microsoft.azure.synapse.ml.core.serialize.ComplexParam -import org.apache.spark.ml.Transformer +import org.apache.hadoop.fs.Path +import org.apache.spark.ml.{PipelineArraySerializer, PipelineStage, Transformer} import org.apache.spark.ml.param.{ParamPair, Params} +import org.apache.spark.sql.SparkSession import scala.collection.JavaConverters._ @@ -19,4 +21,23 @@ class TransformerArrayParam(parent: Params, name: String, doc: String, isValid: /** Creates a param pair with the given value (for Java). */ def w(value: java.util.List[Transformer]): ParamPair[Array[Transformer]] = w(value.asScala.toArray) + override def save( + obj: Array[Transformer], + sparkSession: SparkSession, + path: Path, + overwrite: Boolean): Unit = { + new PipelineArraySerializer(sparkSession) + .write(obj.map(stage => stage: PipelineStage), path, overwrite) + } + + override def load(sparkSession: SparkSession, path: Path): Array[Transformer] = { + new PipelineArraySerializer(sparkSession).read(path).map { + case transformer: Transformer => transformer + case stage => + throw new IllegalArgumentException( + s"Expected a Transformer in $path but found ${stage.getClass.getName}" + ) + } + } + } diff --git a/core/src/main/scala/org/apache/spark/ml/ArtifactPathResolver.scala b/core/src/main/scala/org/apache/spark/ml/ArtifactPathResolver.scala new file mode 100644 index 00000000000..4afb8b88aee --- /dev/null +++ b/core/src/main/scala/org/apache/spark/ml/ArtifactPathResolver.scala @@ -0,0 +1,436 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package org.apache.spark.ml + +import com.microsoft.azure.synapse.ml.core.env.StreamUtilities.using + +import scala.annotation.tailrec +import scala.collection.JavaConverters._ +import scala.io.Source +import java.io.{ByteArrayOutputStream, IOException, InputStream} +import java.nio.charset.StandardCharsets +import java.nio.file.{Files, Paths} + +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.{ + FileContext, + FileStatus, + FileSystem, + Path, + UnsupportedFileSystemException +} +import org.apache.hadoop.io.compress.CompressionCodecFactory +import org.apache.spark.ml.util.DefaultParamsReader +import org.apache.spark.sql.SparkSession + +private[ml] object ArtifactPathResolver { + + private val MaxMetadataEntries = 256 // scalastyle:ignore magic.number + private val MaxMetadataPartFiles = 256 // scalastyle:ignore magic.number + private val MaxMetadataRecordBytes = 1048576 // scalastyle:ignore magic.number + private val MaxMetadataInputBytes = 2L * MaxMetadataRecordBytes + private val DistributedFileSystemClass = "org.apache.hadoop.hdfs.DistributedFileSystem" + private val S3AFileSystemClass = "org.apache.hadoop.fs.s3a.S3AFileSystem" + private val AzureBlobFileSystemClass = "org.apache.hadoop.fs.azurebfs.AzureBlobFileSystem" + private val AbfsListIteratorConfig = "fs.azure.enable.abfslistiterator" + + private final class MetadataInputBudget( + modelContext: Option[ModelLoadContext]) { + private var reservedBytes = 0L + private var consumedBytes = 0L + + def reserve(path: Path, bytes: Long): Unit = { + require( + bytes <= MaxMetadataInputBytes - reservedBytes, + s"Metadata physical input at $path exceeds the $MaxMetadataInputBytes byte directory limit" + ) + modelContext.foreach(_.reserveMetadataInput(path, bytes)) + reservedBytes += bytes + } + + def remaining: Long = { + Math.min( + MaxMetadataInputBytes - consumedBytes, + modelContext.map(_.remainingMetadataInput).getOrElse(Long.MaxValue) + ) + } + + def consume(bytes: Long): Unit = { + consumedBytes += bytes + modelContext.foreach(_.consumeMetadataInput(bytes)) + } + } + + private final class BoundedMetadataInputStream( + input: InputStream, + budget: MetadataInputBudget) extends InputStream { + + private def limitExceeded(): Nothing = { + throw new IOException( + "Metadata physical input exceeds its aggregate byte limit" + ) + } + + override def read(): Int = { + val value = input.read() + if (value >= 0) { + if (budget.remaining <= 0) { + limitExceeded() + } + budget.consume(1) + } + value + } + + override def read(buffer: Array[Byte], offset: Int, length: Int): Int = { + if (length == 0) { + 0 + } else if (budget.remaining <= 0) { + val value = input.read() + if (value < 0) -1 else limitExceeded() + } else { + val allowedLength = Math.min(length.toLong, budget.remaining).toInt + val count = input.read(buffer, offset, allowedLength) + if (count > 0) { + budget.consume(count) + } + count + } + } + + override def skip(length: Long): Long = { + if (length <= 0) { + 0 + } else { + val buffer = new Array[Byte](Math.min(length, 8192L).toInt) // scalastyle:ignore magic.number + var remaining = length + var skipped = 0L + var count = read(buffer, 0, Math.min(remaining, buffer.length).toInt) + while (remaining > 0 && count >= 0) { // scalastyle:ignore while + remaining -= count + skipped += count + count = read(buffer, 0, Math.min(remaining, buffer.length).toInt) + } + skipped + } + } + + override def available(): Int = { + Math.min( + input.available().toLong, + Math.max(0L, budget.remaining) + ).toInt + } + + override def close(): Unit = input.close() + } + + private def normalizedQualifiedUri( + hadoopConf: Configuration, + path: Path): java.net.URI = { + val fs = path.getFileSystem(hadoopConf) + path.makeQualified(fs.getUri, fs.getWorkingDirectory).toUri.normalize() + } + + private def sameFileSystem(left: java.net.URI, right: java.net.URI): Boolean = { + Option(left.getScheme) == Option(right.getScheme) && + Option(left.getAuthority) == Option(right.getAuthority) + } + + private def isWithin( + root: java.net.URI, + candidate: java.net.URI, + allowRoot: Boolean): Boolean = { + val rootPath = root.getPath.stripSuffix("/") + sameFileSystem(root, candidate) && + ((allowRoot && candidate.getPath == rootPath) || candidate.getPath.startsWith(rootPath + "/")) + } + + private def resolveExistingPath( + spark: SparkSession, + path: java.net.URI): java.net.URI = { + val hadoopConf = Serializer.sessionHadoopConf(spark) + if (Option(path.getScheme).exists(_.equalsIgnoreCase("file"))) { + Paths.get(path).toRealPath().toUri.normalize() + } else { + try { + FileContext.getFileContext(path, hadoopConf) + .resolvePath(new Path(path)) + .toUri + .normalize() + } catch { + case error: UnsupportedFileSystemException => + if (Serializer.legacyObjectDeserializationEnabled(spark)) { + val hadoopPath = new Path(path) + val fs = hadoopPath.getFileSystem(hadoopConf) + fs.getFileStatus(hadoopPath) + hadoopPath.makeQualified(fs.getUri, fs.getWorkingDirectory).toUri.normalize() + } else { + val securityError = new SecurityException( + s"Filesystem scheme ${path.getScheme} cannot resolve links safely. Set " + + s"${Serializer.LegacyObjectDeserializationConfig}=true only for trusted artifacts." + ) + securityError.initCause(error) + throw securityError + } + } + } + } + + def resolvePathInside( + spark: SparkSession, + root: Path, + candidate: Path, + description: String, + allowRoot: Boolean = false): Path = { + val hadoopConf = Serializer.sessionHadoopConf(spark) + val qualifiedRoot = normalizedQualifiedUri(hadoopConf, root) + val qualifiedCandidate = normalizedQualifiedUri(hadoopConf, candidate) + require( + isWithin(qualifiedRoot, qualifiedCandidate, allowRoot), + s"$description $candidate resolves outside $root" + ) + val resolvedRoot = resolveExistingPath(spark, qualifiedRoot) + val resolvedCandidate = resolveExistingPath(spark, qualifiedCandidate) + require( + isWithin(resolvedRoot, resolvedCandidate, allowRoot), + s"$description $candidate resolves outside $root through a filesystem link" + ) + new Path(resolvedCandidate) + } + + def loadMetadata( + spark: SparkSession, + path: Path, + expectedClassName: String = ""): DefaultParamsReader.Metadata = { + val metadataPath = resolvePathInside( + spark, + path, + new Path(path, "metadata"), + "Model metadata directory" + ) + val hadoopConf = Serializer.sessionHadoopConf(spark) + val fs = metadataPath.getFileSystem(hadoopConf) + val codecFactory = new CompressionCodecFactory(hadoopConf) + val record = readMetadataRecord( + spark, + metadataPath, + fs, + codecFactory + ) + require(record.isDefined, s"Model metadata directory $metadataPath has no metadata record") + val metadata = DefaultParamsReader.parseMetadata(record.get) + require( + expectedClassName.isEmpty || metadata.className == expectedClassName, + s"Expected model class $expectedClassName but metadata declared ${metadata.className}" + ) + metadata + } + + private def readMetadataRecord( + spark: SparkSession, + metadataPath: Path, + fs: FileSystem, + codecFactory: CompressionCodecFactory): Option[String] = { + var record = Option.empty[String] + var entryCount = 0 + var partCount = 0 + val inputBudget = new MetadataInputBudget(ModelLoadContext.current) + val resolvedPartPaths = scala.collection.mutable.Set.empty[String] + foreachMetadataEntry(spark, metadataPath, fs) { status => + require( + entryCount < MaxMetadataEntries, + s"Model metadata directory $metadataPath has too many entries" + ) + entryCount += 1 + if (status.getPath.getName.startsWith("part-")) { + require( + partCount < MaxMetadataPartFiles, + s"Model metadata directory $metadataPath has too many part files" + ) + partCount += 1 + val resolved = resolvePathInside( + spark, + metadataPath, + status.getPath, + "Model metadata file" + ) + val resolvedPathKey = resolved.toUri.normalize().toString + require( + !resolvedPartPaths.contains(resolvedPathKey), + s"Model metadata part $resolved was encountered more than once" + ) + resolvedPartPaths += resolvedPathKey + val resolvedStatus = fs.getFileStatus(resolved) + require(resolvedStatus.isFile, s"Model metadata part $resolved is not a file") + inputBudget.reserve(resolved, resolvedStatus.getLen) + val partRecords = readMetadataPart(fs, codecFactory, resolved, inputBudget) + require( + partRecords.length <= 1 && (record.isEmpty || partRecords.isEmpty), + s"Expected one metadata record at $metadataPath but found more than one" + ) + record = partRecords.headOption.orElse(record) + } + } + record + } + + private def foreachMetadataEntry( + spark: SparkSession, + metadataPath: Path, + fs: FileSystem)(consume: FileStatus => Unit): Unit = { + if (Option(metadataPath.toUri.getScheme).exists(_.equalsIgnoreCase("file"))) { + val entries = Files.newDirectoryStream(Paths.get(metadataPath.toUri)) + try { + entries.iterator().asScala.foreach { entry => + consume(fs.getFileStatus(new Path(entry.toUri))) + } + } finally { + entries.close() + } + } else { + val configuration = Serializer.sessionHadoopConf(spark) + require( + guaranteesIncrementalMetadataListing(fs, configuration) || + Serializer.legacyObjectDeserializationEnabled(spark), + s"Filesystem ${fs.getClass.getName} is not configured for bounded metadata enumeration. " + + s"Set ${Serializer.LegacyObjectDeserializationConfig}=true only for trusted artifacts." + ) + val entries = fs.listStatusIterator(metadataPath) + while (entries.hasNext) { // scalastyle:ignore while + consume(entries.next()) + } + } + } + + private[ml] def guaranteesIncrementalMetadataListing( + fs: FileSystem, + configuration: Configuration): Boolean = { + guaranteesIncrementalMetadataListing( + filesystemClassNames(fs.getClass), + fs.getUri, + configuration + ) + } + + private[ml] def guaranteesIncrementalMetadataListing( + filesystemClasses: Set[String], + filesystemUri: java.net.URI, + configuration: Configuration): Boolean = { + filesystemClasses.contains(DistributedFileSystemClass) || + filesystemClasses.contains(S3AFileSystemClass) || + (filesystemClasses.contains(AzureBlobFileSystemClass) && + abfsIncrementalListingEnabled(filesystemUri, configuration)) + } + + private def abfsIncrementalListingEnabled( + filesystemUri: java.net.URI, + configuration: Configuration): Boolean = { + val globalValue = configuration.getBoolean(AbfsListIteratorConfig, true) + Option(filesystemUri.getAuthority) + .map(_.split("@").last) + .map(account => configuration.getBoolean(s"$AbfsListIteratorConfig.$account", globalValue)) + .getOrElse(globalValue) + } + + private def filesystemClassNames(filesystemClass: Class[_]): Set[String] = { + Option(filesystemClass.getSuperclass) + .map(filesystemClassNames) + .getOrElse(Set.empty) + filesystemClass.getName + } + + private def readMetadataPart( + fs: FileSystem, + codecFactory: CompressionCodecFactory, + path: Path, + inputBudget: MetadataInputBudget): Array[String] = { + val bytes = using(fs.open(path)) { rawInput => + val boundedInput = new BoundedMetadataInputStream(rawInput, inputBudget) + Option(codecFactory.getCodec(path)) match { + case Some(codec) => + using(codec.createInputStream(boundedInput)) { + readUpTo(_, MaxMetadataRecordBytes + 1) + }.get + case None => readUpTo(boundedInput, MaxMetadataRecordBytes + 1) + } + }.get + require( + bytes.length <= MaxMetadataRecordBytes, + s"Model metadata record at $path exceeds $MaxMetadataRecordBytes bytes" + ) + ModelLoadContext.current.foreach(_.consumeDecodedMetadata(path, bytes.length)) + using(Source.fromString(new String(bytes, StandardCharsets.UTF_8))) { + source => + val lines = source.getLines() + if (!lines.hasNext) { + Array.empty[String] + } else { + val first = lines.next() + require(!lines.hasNext, s"Model metadata part $path contains more than one record") + Array(first) + } + }.get + } + + private[ml] def validateMetadataForWrite(path: Path, bytes: Long): Unit = { + require( + bytes <= MaxMetadataRecordBytes, + s"Model metadata record at $path exceeds $MaxMetadataRecordBytes bytes" + ) + ModelLoadContext.consumeWrittenMetadata(path, bytes) + } + + private[ml] def readUpTo(input: InputStream, maxBytes: Int): Array[Byte] = { + require(maxBytes >= 0, "Maximum byte count must be non-negative") + val output = new ByteArrayOutputStream(Math.min(maxBytes, 8192)) // scalastyle:ignore magic.number + val buffer = new Array[Byte](Math.min(maxBytes, 8192)) // scalastyle:ignore magic.number + + @tailrec + def copyRemaining(remaining: Int): Unit = { + if (remaining > 0) { + val count = input.read(buffer, 0, Math.min(buffer.length, remaining)) + if (count > 0) { + output.write(buffer, 0, count) + copyRemaining(remaining - count) + } + } + } + + copyRemaining(maxBytes) + output.toByteArray + } + + def resolveTreeInside( + spark: SparkSession, + path: Path, + description: String): Path = { + val resolvedRoot = resolvePathInside(spark, path, path, description, allowRoot = true) + val fs = resolvedRoot.getFileSystem(Serializer.sessionHadoopConf(spark)) + val visitedDirectories = scala.collection.mutable.Set[String]() + + def validateDirectory(directory: Path): Unit = { + val directoryKey = directory.toUri.normalize().toString + require( + visitedDirectories.add(directoryKey), + s"$description contains a filesystem-link cycle at $directory" + ) + fs.listStatus(directory).foreach { status => + val resolvedChild = resolvePathInside( + spark, + resolvedRoot, + status.getPath, + s"$description entry" + ) + if (fs.getFileStatus(resolvedChild).isDirectory) { + validateDirectory(resolvedChild) + } + } + } + + if (fs.getFileStatus(resolvedRoot).isDirectory) { + validateDirectory(resolvedRoot) + } + resolvedRoot + } +} diff --git a/core/src/main/scala/org/apache/spark/ml/ComplexParamsSerializer.scala b/core/src/main/scala/org/apache/spark/ml/ComplexParamsSerializer.scala index 72a8c82ce94..d82970b6381 100644 --- a/core/src/main/scala/org/apache/spark/ml/ComplexParamsSerializer.scala +++ b/core/src/main/scala/org/apache/spark/ml/ComplexParamsSerializer.scala @@ -3,6 +3,7 @@ package org.apache.spark.ml +import com.microsoft.azure.synapse.ml.core.env.StreamUtilities.using import com.microsoft.azure.synapse.ml.core.serialize.ComplexParam import org.apache.hadoop.fs.Path import org.apache.spark.ml.param.{ParamPair, Params} @@ -14,6 +15,9 @@ import org.json4s.JsonDSL._ import org.json4s.jackson.JsonMethods._ import org.json4s.{JObject, _} +import java.io.IOException +import java.nio.charset.StandardCharsets + trait ComplexParamsWritable extends MLWritable { self: Params => @@ -22,7 +26,13 @@ trait ComplexParamsWritable extends MLWritable { trait ComplexParamsReadable[T] extends MLReadable[T] { - override def read: MLReader[T] = new ComplexParamsReader[T] + override def read: MLReader[T] = { + new ComplexParamsReader[T](getClass.getName.stripSuffix("$")) + } + + override def load(path: String): T = { + read.session(SparkSession.builder().getOrCreate()).load(path) + } } /** Default [[MLWriter]] implementation for transformers and estimators that contain basic @@ -34,15 +44,44 @@ trait ComplexParamsReadable[T] extends MLReadable[T] { private[ml] class ComplexParamsWriter(instance: Params) extends MLWriter { override protected def saveImpl(path: String): Unit = { - val complexParamLocs = ComplexParamsWriter.getComplexParamLocations(instance, path) - val complexParamJson = ComplexParamsWriter.getComplexMetadata(complexParamLocs) - ComplexParamsWriter.saveMetadata(instance, path, sparkSession, complexParamJson) - ComplexParamsWriter.saveComplexParams(path, complexParamLocs, shouldOverwrite) + if ( + !ModelLoadContext.writeContextActive && + ComplexParamsWriter.isNativePipelineStagePath(instance, path) + ) { + throw new SecurityException( + s"Native Spark Pipeline persistence does not provide a model-wide metadata budget for " + + s"ComplexParams stage ${instance.uid}. Use SynapseML PipelineSerializer, or wrap native " + + "Pipeline compatibility in Serializer.withTrustedArtifactLoad for trusted artifacts." + ) + } + ModelLoadContext.withWriteContext { + val complexParamLocs = ComplexParamsWriter.getComplexParamLocations(instance, path) + val complexParamJson = ComplexParamsWriter.getComplexMetadata(complexParamLocs) + ComplexParamsWriter.saveMetadata(instance, path, sparkSession, complexParamJson) + ComplexParamsWriter.saveComplexParams(path, complexParamLocs, shouldOverwrite, sparkSession) + } } } private[ml] object ComplexParamsWriter { + private val MetadataLineSeparatorBytes = + "\n".getBytes(StandardCharsets.UTF_8).length.toLong + private val MetadataPartFileName = "part-00000" + private val MetadataSuccessFileName = "_SUCCESS" + + private def isNativePipelineStagePath(instance: Params, path: String): Boolean = { + val stagePath = new Path(path) + val stageSuffix = s"_${instance.uid}" + val stageName = stagePath.getName + val stageIndex = + if (stageName.endsWith(stageSuffix)) stageName.dropRight(stageSuffix.length) + else "" + Option(stagePath.getParent).exists(_.getName == "stages") && + stageIndex.nonEmpty && + stageIndex.forall(_.isDigit) + } + def getComplexMetadata(complexParamLocs: Map[ParamPair[_], Path]): Option[JObject] = { if (complexParamLocs.nonEmpty) { Some(JObject("complexParamLocs" -> complexParamLocs.map { @@ -55,7 +94,8 @@ private[ml] object ComplexParamsWriter { def getComplexParamLocations(instance: Params, path: String): Map[ParamPair[_], Path] = { val complexParams = instance.extractParamMap().toSeq.filter { - case ParamPair(_: ComplexParam[_], _) => true + case ParamPair(p: ComplexParam[_], _) => + instance.isSet(p) || p.isSafeForUntrustedDeserialization case _ => false } complexParams.map { @@ -63,9 +103,11 @@ private[ml] object ComplexParamsWriter { }.toMap } - def saveComplexParams(basePath: String, complexParamLocs: Map[ParamPair[_], Path], - shouldOverwrite: Boolean): Unit = { - val spark = SparkSession.builder().getOrCreate() + def saveComplexParams( + basePath: String, + complexParamLocs: Map[ParamPair[_], Path], + shouldOverwrite: Boolean, + spark: SparkSession): Unit = { complexParamLocs.foreach { case (ParamPair(p: ComplexParam[Any], v), loc) => p.save(v, spark, new Path(basePath, loc), shouldOverwrite) @@ -92,9 +134,24 @@ private[ml] object ComplexParamsWriter { spark: SparkSession, extraMetadata: Option[JObject] = None, paramMap: Option[JValue] = None): Unit = { - val metadataPath = new Path(path, "metadata").toString - val metadataJson = getMetadataToSave(instance, spark, extraMetadata, paramMap) - spark.createDataFrame(Seq(Tuple1(metadataJson))).toDF("value").write.text(metadataPath) + ModelLoadContext.withWriteContext { + val metadataPath = new Path(path, "metadata") + val metadataJson = getMetadataToSave(instance, spark, extraMetadata, paramMap) + ArtifactPathResolver.validateMetadataForWrite( + metadataPath, + metadataJson.getBytes(StandardCharsets.UTF_8).length.toLong + + MetadataLineSeparatorBytes + ) + val fs = metadataPath.getFileSystem(Serializer.sessionHadoopConf(spark)) + if (fs.exists(metadataPath) || !fs.mkdirs(metadataPath)) { + throw new IOException(s"Metadata path $metadataPath already exists or could not be created") + } + using(fs.create(new Path(metadataPath, MetadataPartFileName), false)) { output => + output.write(metadataJson.getBytes(StandardCharsets.UTF_8)) + output.write('\n') + }.get + using(fs.create(new Path(metadataPath, MetadataSuccessFileName), false))(_ => ()).get + } } /** Helper for [[saveMetadata()]] which extracts the JSON to save. @@ -141,46 +198,120 @@ private[ml] object ComplexParamsWriter { * data (e.g., models with coefficients). * * @tparam T ML instance type - * TODO: Consider adding check for correct class name. */ -private[ml] class ComplexParamsReader[T] extends MLReader[T] { +private[ml] class ComplexParamsReader[T]( + expectedClassName: String, + expectedClass: Option[Class[_]] = None, + preloadedMetadata: Option[Metadata] = None) extends MLReader[T] { + + private var assignedSession = Option.empty[SparkSession] + + override def session(session: SparkSession): this.type = { + assignedSession = Some(session) + super.session(session) + } override def load(path: String): T = { - // Mirrors Spark 4.0's DefaultParamsReader.loadMetadata(path, spark, expectedClassName), - // which reads the metadata through the session rather than SparkContext.textFile. Spark 3.5 - // only offers the SparkContext overload, so inline the same three lines here: otherwise - // loading a model still needs a SparkContext and remains unusable on Databricks Unity - // Catalog shared access mode and Spark Connect, which is the whole point of the writer - // change above. - val metadataPath = new Path(path, "metadata").toString - val metadataStr = sparkSession.read.text(metadataPath).first().getString(0) - val metadata = DefaultParamsReader.parseMetadata(metadataStr) - val cls = Utils.classForName(metadata.className) - val instance = - cls.getConstructor(classOf[String]).newInstance(metadata.uid).asInstanceOf[Params] - metadata.getAndSetParams(instance) - ComplexParamsReader.getAndSetComplexParams(instance, metadata, path) - instance.asInstanceOf[T] + if ( + assignedSession.isEmpty && + ModelLoadContext.current.isEmpty && + Serializer.currentTrustedArtifactLoad.forall(trusted => !trusted) + ) { + throw new SecurityException( + s"Unscoped native Spark Pipeline loading is not bounded for ComplexParams at $path. " + + "Use SynapseML PipelineSerializer for untrusted artifacts, or " + + "Serializer.withTrustedArtifactLoad only for trusted compatibility." + ) + } + val loadSession = assignedSession.getOrElse(sparkSession) + val configuredTrust = assignedSession.map(Serializer.configuredLegacyDeserialization) + val trustedArtifact = Serializer.currentTrustedArtifactLoad match { + case Some(scopedTrust) => scopedTrust && configuredTrust.forall(identity) + case None => configuredTrust.getOrElse(false) + } + Serializer.withTrustedArtifactLoad(trustedArtifact) { + Serializer.withActiveSession(loadSession) { + val requestedPath = new Path(path) + val resolvedPath = ArtifactPathResolver.resolvePathInside( + loadSession, + requestedPath, + requestedPath, + "ComplexParams model root", + allowRoot = true + ) + ModelLoadContext.withContext(resolvedPath) { context => + val metadata = preloadedMetadata.getOrElse { + ArtifactPathResolver.loadMetadata( + loadSession, + resolvedPath, + expectedClassName + ) + } + require( + expectedClassName.isEmpty || metadata.className == expectedClassName, + s"Expected model class $expectedClassName but metadata declared ${metadata.className}" + ) + val cls = expectedClass.getOrElse(Utils.classForName(metadata.className)) + val instance = + cls.getConstructor(classOf[String]).newInstance(metadata.uid).asInstanceOf[Params] + metadata.getAndSetParams(instance) + ComplexParamsReader.getAndSetComplexParams( + instance, + metadata, + resolvedPath.toString, + loadSession, + context + ) + instance.asInstanceOf[T] + } + } + } } } private[ml] object ComplexParamsReader { - def getAndSetComplexParams(instance: Params, metadata: Metadata, basePath: String): Unit = { - val spark = SparkSession.builder().getOrCreate() + def getAndSetComplexParams( + instance: Params, + metadata: Metadata, + basePath: String, + spark: SparkSession, + context: ModelLoadContext): Unit = { implicit val format: DefaultFormats.type = DefaultFormats - val complexParamLocs = (metadata.metadata \ "complexParamLocs") match { + val serializedParamLocs = (metadata.metadata \ "complexParamLocs") match { case JNothing => - Map[String, Path]() + Map[String, String]() case j => - j.extract[Map[String, String]].map { case (name, path) => (name, new Path(basePath, path)) } + j.extract[Map[String, String]] + } + val complexParamNames = instance.params.collect { case p: ComplexParam[_] => p.name }.toSet + val unexpectedParams = serializedParamLocs.keySet -- complexParamNames + require( + unexpectedParams.isEmpty, + s"Metadata contains unknown complex parameters: ${unexpectedParams.toSeq.sorted.mkString(", ")}" + ) + val complexParamLocs = serializedParamLocs.map { case (name, serializedPath) => + val expectedPath = new Path("complexParams", name) + require( + serializedPath == expectedPath.toString, + s"Complex parameter $name must use relative path $expectedPath, not $serializedPath" + ) + name -> ArtifactPathResolver.resolvePathInside( + spark, + new Path(basePath), + new Path(basePath, expectedPath), + s"Complex parameter $name" + ) } instance.params.foreach { case p: ComplexParam[_] => complexParamLocs.get(p.name) match { case Some(loc) => - instance.set(p, p.load(spark,loc)) + context.enterPath(loc, s"Complex parameter ${p.name}") + context.withNestedArtifact(loc) { + instance.set(p, p.load(spark, loc)) + } case None => } case _ => diff --git a/core/src/main/scala/org/apache/spark/ml/DataTypeSerializer.scala b/core/src/main/scala/org/apache/spark/ml/DataTypeSerializer.scala new file mode 100644 index 00000000000..03e43b9dc04 --- /dev/null +++ b/core/src/main/scala/org/apache/spark/ml/DataTypeSerializer.scala @@ -0,0 +1,125 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package org.apache.spark.ml + +import com.microsoft.azure.synapse.ml.core.env.StreamUtilities.using +import org.apache.hadoop.fs.Path +import org.apache.spark.ml.linalg.SQLDataTypes +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.types.DataType +import org.json4s.jackson.JsonMethods.parse +import org.json4s.{JArray, JObject, JString, JValue} + +import java.io.ByteArrayInputStream +import java.nio.charset.StandardCharsets + +private[ml] class DataTypeSerializer(spark: SparkSession) extends Serializer[DataType] { + + private val maxJsonBytes = 16 * 1024 * 1024 // scalastyle:ignore magic.number + private val javaStreamMagicFirst = 0xac.toByte // scalastyle:ignore magic.number + private val javaStreamMagicSecond = 0xed.toByte // scalastyle:ignore magic.number + private val allowedUdtClasses = Set( + SQLDataTypes.VectorType.getClass.getName, + SQLDataTypes.MatrixType.getClass.getName + ) + + override def write( + dataType: DataType, + path: Path, + overwrite: Boolean): Unit = { + val bytes = dataType.json.getBytes(StandardCharsets.UTF_8) + require(bytes.length <= maxJsonBytes, s"DataType JSON exceeds $maxJsonBytes bytes") + using(path.getFileSystem(Serializer.sessionHadoopConf(spark)).create(path, overwrite)) { + _.write(bytes) + }.get + } + + override def read(path: Path): DataType = { + val fs = path.getFileSystem(Serializer.sessionHadoopConf(spark)) + val status = fs.getFileStatus(path) + require(status.isFile, s"DataType parameter $path is not a file") + require(status.getLen <= maxJsonBytes, s"DataType parameter $path exceeds $maxJsonBytes bytes") + val bytes = using(fs.open(path)) { + ArtifactPathResolver.readUpTo(_, maxJsonBytes + 1) + }.get + require(bytes.length <= maxJsonBytes, s"DataType parameter $path exceeds $maxJsonBytes bytes") + + if (isJavaObjectStream(bytes)) { + readLegacy(path, bytes) + } else { + readJson(path, bytes) + } + } + + private def isJavaObjectStream(bytes: Array[Byte]): Boolean = { + bytes.length >= 2 && + bytes(0) == javaStreamMagicFirst && + bytes(1) == javaStreamMagicSecond + } + + private def readLegacy(path: Path, bytes: Array[Byte]): DataType = { + if (!Serializer.legacyObjectDeserializationEnabled(spark)) { + throw new SecurityException( + s"Legacy Java-serialized DataType parameter $path requires a trusted artifact. Set " + + s"${Serializer.LegacyObjectDeserializationConfig}=true only when loading trusted data." + ) + } + Serializer.readUnsafe[DataType](new ByteArrayInputStream(bytes)) + } + + private def readJson(path: Path, bytes: Array[Byte]): DataType = { + val json = new String(bytes, StandardCharsets.UTF_8) + val customClasses = customUdtClasses(parse(json)).filterNot(allowedUdtClasses) + if (customClasses.nonEmpty && !Serializer.legacyObjectDeserializationEnabled(spark)) { + throw new SecurityException( + s"DataType parameter $path declares custom UDT classes " + + s"${customClasses.distinct.sorted.mkString(", ")}. Set " + + s"${Serializer.LegacyObjectDeserializationConfig}=true only when loading trusted data." + ) + } + DataType.fromJson(json) + } + + private def customUdtClasses(value: JValue): Seq[String] = { + value match { + case JObject(fields) => customUdtClasses(fields.toMap) + case _ => Seq.empty + } + } + + private def customUdtClasses(fields: Map[String, JValue]): Seq[String] = { + fields.get("type") match { + case Some(JString("array")) => + fields.get("elementType").toSeq.flatMap(customUdtClasses) + case Some(JString("map")) => + Seq("keyType", "valueType").flatMap(fields.get).flatMap(customUdtClasses) + case Some(JString("struct")) => + fields.get("fields").toSeq.flatMap(structUdtClasses) + case Some(JString("udt")) => + declaredUdtClass(fields) ++ fields.get("sqlType").toSeq.flatMap(customUdtClasses) + case _ => Seq.empty + } + } + + private def structUdtClasses(value: JValue): Seq[String] = { + value match { + case JArray(structFields) => structFields.flatMap(structFieldUdtClasses) + case _ => Seq.empty + } + } + + private def structFieldUdtClasses(value: JValue): Seq[String] = { + value match { + case JObject(fields) => fields.toMap.get("type").toSeq.flatMap(customUdtClasses) + case _ => Seq.empty + } + } + + private def declaredUdtClass(fields: Map[String, JValue]): Seq[String] = { + fields.get("class") match { + case Some(JString(className)) => Seq(className) + case _ => Seq("") + } + } +} diff --git a/core/src/main/scala/org/apache/spark/ml/ModelLoadContext.scala b/core/src/main/scala/org/apache/spark/ml/ModelLoadContext.scala new file mode 100644 index 00000000000..be55af57a3d --- /dev/null +++ b/core/src/main/scala/org/apache/spark/ml/ModelLoadContext.scala @@ -0,0 +1,209 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package org.apache.spark.ml + +import org.apache.hadoop.fs.Path +import org.apache.spark.sql.SparkSession + +/** Same-thread compatibility scope for language bindings that cannot pass a Scala closure. + * + * Call [[close()]] in a `finally` block on the same gateway thread that created this scope. + */ +final class TrustedArtifactLoadScope private[ml]( + ownerThread: Thread, + closeAction: () => Unit) extends AutoCloseable { + + private var closed = false + + override def close(): Unit = synchronized { + if (!closed) { + if (Thread.currentThread() ne ownerThread) { + throw new IllegalStateException( + "Trusted artifact scope must be closed on the thread that created it" + ) + } + closeAction() + closed = true + } + } +} + +private[ml] object TrustedArtifactLoadScope { + + def begin( + spark: SparkSession, + legacyDeserializationEnabled: Boolean, + trustedArtifactLoad: ThreadLocal[java.lang.Boolean]): TrustedArtifactLoadScope = { + if (!legacyDeserializationEnabled) { + throw new SecurityException( + s"Set ${Serializer.LegacyObjectDeserializationConfig}=true on the supplied " + + "SparkSession only when loading a trusted legacy artifact." + ) + } + val ownerThread = Thread.currentThread() + val previousTrust = Option(trustedArtifactLoad.get()) + val previousSession = SparkSession.getActiveSession + val writeContextCreated = ModelLoadContext.beginWriteContext() + + def restore(): Unit = { + try { + previousSession.fold(SparkSession.clearActiveSession())(SparkSession.setActiveSession) + } finally { + try { + previousTrust.fold(trustedArtifactLoad.remove())(trustedArtifactLoad.set) + } finally { + ModelLoadContext.endWriteContext(writeContextCreated) + } + } + } + + var initialized = false + try { + trustedArtifactLoad.set(true) + SparkSession.setActiveSession(spark) + val scope = new TrustedArtifactLoadScope(ownerThread, () => restore()) + initialized = true + scope + } finally { + if (!initialized) { + restore() + } + } + } +} + +private[ml] final class ModelLoadContext private(root: Path) { + private val visitedPaths = scala.collection.mutable.Set(root.toUri.normalize().toString) + private val decodedMetadataBudget = new ModelLoadContext.DecodedMetadataBudget + private var nodeCount = 1 + private var nestingDepth = 0 + private var reservedMetadataBytes = 0L + private var consumedMetadataBytes = 0L + + def enterPath(path: Path, description: String): Unit = { + val pathKey = path.toUri.normalize().toString + require(!visitedPaths.contains(pathKey), s"$description $path was encountered more than once") + require( + nodeCount < ModelLoadContext.MaxNodes, + s"Model persistence contains more than ${ModelLoadContext.MaxNodes} nodes" + ) + visitedPaths += pathKey + nodeCount += 1 + } + + def ensurePath(path: Path): Unit = { + val pathKey = path.toUri.normalize().toString + if (!visitedPaths.contains(pathKey)) { + enterPath(path, "Nested model root") + } + } + + def withNestedArtifact[T](path: Path)(action: => T): T = { + require( + nestingDepth < ModelLoadContext.MaxDepth, + s"Model persistence nesting exceeds ${ModelLoadContext.MaxDepth} levels at $path" + ) + nestingDepth += 1 + try { + action + } finally { + nestingDepth -= 1 + } + } + + def reserveMetadataInput(path: Path, bytes: Long): Unit = { + require( + bytes <= ModelLoadContext.MaxMetadataInputBytes - reservedMetadataBytes, + s"Model metadata physical input exceeds " + + s"${ModelLoadContext.MaxMetadataInputBytes} bytes at $path" + ) + reservedMetadataBytes += bytes + } + + def remainingMetadataInput: Long = + ModelLoadContext.MaxMetadataInputBytes - consumedMetadataBytes + + def consumeMetadataInput(bytes: Long): Unit = { + consumedMetadataBytes += bytes + } + + def consumeDecodedMetadata(path: Path, bytes: Long): Unit = { + decodedMetadataBudget.consume(path, bytes) + } +} + +private[ml] object ModelLoadContext { + private val MaxDepth = 100 // scalastyle:ignore magic.number + private val MaxNodes = 10000 // scalastyle:ignore magic.number + private val MaxMetadataInputBytes = 64L << 20 // scalastyle:ignore magic.number + private val MaxDecodedMetadataBytes = Math.max( + 1L << 20, // scalastyle:ignore magic.number + Math.min( + 64L << 20, // scalastyle:ignore magic.number + Runtime.getRuntime.maxMemory() / 16 // scalastyle:ignore magic.number + ) + ) + private val Active = new ThreadLocal[ModelLoadContext] + private val ActiveWriteBudget = new ThreadLocal[DecodedMetadataBudget] + + private final class DecodedMetadataBudget { + private var consumedBytes = 0L + + def consume(path: Path, bytes: Long): Unit = { + require(bytes >= 0, s"Decoded metadata byte count at $path must be non-negative") + require( + bytes <= MaxDecodedMetadataBytes - consumedBytes, + s"Model decoded metadata exceeds $MaxDecodedMetadataBytes bytes at $path" + ) + consumedBytes += bytes + } + } + + def withContext[T](root: Path)(action: ModelLoadContext => T): T = { + Option(Active.get()) match { + case Some(context) => + context.ensurePath(root) + action(context) + case None => + val context = new ModelLoadContext(root) + Active.set(context) + try { + action(context) + } finally { + Active.remove() + } + } + } + + def current: Option[ModelLoadContext] = Option(Active.get()) + + def withWriteContext[T](action: => T): T = { + val created = beginWriteContext() + try { + action + } finally { + endWriteContext(created) + } + } + + private[ml] def beginWriteContext(): Boolean = { + val created = ActiveWriteBudget.get() == null + if (created) { + ActiveWriteBudget.set(new DecodedMetadataBudget) + } + created + } + + private[ml] def endWriteContext(created: Boolean): Unit = { + if (created) { + ActiveWriteBudget.remove() + } + } + + def consumeWrittenMetadata(path: Path, bytes: Long): Unit = { + Option(ActiveWriteBudget.get()).foreach(_.consume(path, bytes)) + } + + def writeContextActive: Boolean = ActiveWriteBudget.get() != null +} diff --git a/core/src/main/scala/org/apache/spark/ml/RuntimeTypedPipelineArraySerializer.scala b/core/src/main/scala/org/apache/spark/ml/RuntimeTypedPipelineArraySerializer.scala new file mode 100644 index 00000000000..b1f6c5f32db --- /dev/null +++ b/core/src/main/scala/org/apache/spark/ml/RuntimeTypedPipelineArraySerializer.scala @@ -0,0 +1,47 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package org.apache.spark.ml + +import org.apache.hadoop.fs.Path +import org.apache.spark.sql.SparkSession + +private[ml] class RuntimeTypedPipelineArraySerializer[T]( + spark: SparkSession, + componentClass: Class[_ <: PipelineStage]) extends Serializer[T] { + + private val delegate = new PipelineArraySerializer(spark) + + private def className(value: Any): String = { + if (value == null) "null" else value.getClass.getName + } + + override def write(value: T, path: Path, overwrite: Boolean): Unit = { + require(value != null, s"Expected an array of ${componentClass.getName} but found null") + val stages = value.asInstanceOf[Array[_]].zipWithIndex.map { + case (stage: PipelineStage, _) if componentClass.isInstance(stage) => stage + case (stage: PipelineStage, index) => + throw new IllegalArgumentException( + s"Expected ${componentClass.getName} at array index $index but found ${className(stage)}" + ) + case (other, index) => + throw new IllegalArgumentException( + s"Expected a PipelineStage at array index $index but found ${className(other)}" + ) + } + delegate.write(stages, path, overwrite) + } + + override def read(path: Path): T = { + val stages = delegate.read(path) + val result = java.lang.reflect.Array.newInstance(componentClass, stages.length) + stages.zipWithIndex.foreach { case (stage, index) => + require( + componentClass.isInstance(stage), + s"Expected ${componentClass.getName} at array index $index but found ${className(stage)}" + ) + java.lang.reflect.Array.set(result, index, stage) + } + result.asInstanceOf[T] + } +} diff --git a/core/src/main/scala/org/apache/spark/ml/Serializer.scala b/core/src/main/scala/org/apache/spark/ml/Serializer.scala index c9296087ffa..90595793d12 100644 --- a/core/src/main/scala/org/apache/spark/ml/Serializer.scala +++ b/core/src/main/scala/org/apache/spark/ml/Serializer.scala @@ -4,15 +4,29 @@ package org.apache.spark.ml import com.microsoft.azure.synapse.ml.core.env.StreamUtilities._ -import com.microsoft.azure.synapse.ml.core.utils.ContextObjectInputStream +import com.microsoft.azure.synapse.ml.core.utils.{ + ContextObjectInputStream, + DeserializationClassFilter, + SafeObjectInputStream +} import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.Path -import org.apache.spark.ml.util.MLWritable +import org.apache.spark.ml.param.Params +import org.apache.spark.ml.util.{ + DefaultParamsReader, + DefaultParamsWriter, + MLReader, + MLWritable +} import org.apache.spark.sql._ +import org.apache.spark.sql.types.DataType +import org.apache.spark.util.Utils +import org.json4s.{DefaultFormats, JArray, JObject, JString} -import java.io.{InputStream, ObjectOutputStream, OutputStream} +import java.io.{IOException, InputStream, InvalidClassException, ObjectOutputStream, OutputStream} import scala.language.existentials import scala.reflect.runtime.universe._ +import scala.util.control.NonFatal abstract class Serializer[O] { def write(obj: O, path: Path, overwrite: Boolean): Unit @@ -21,6 +35,12 @@ abstract class Serializer[O] { object Serializer { + /** Explicit compatibility switch for trusted artifacts whose persistence cannot be constrained, + * including Java object graphs, custom readers, UDF closures, and lazy DataFrame parameters. + */ + val LegacyObjectDeserializationConfig: String = + "spark.synapseml.legacy.allowUnsafeJavaDeserialization" + val ContextClassLoader: ClassLoader = Thread.currentThread().getContextClassLoader val Mirror: Mirror = runtimeMirror(Serializer.ContextClassLoader) @@ -38,11 +58,52 @@ object Serializer { }) } + private def isPipelineStageArray(tpe: Type): Boolean = { + val normalizedType = tpe.dealias + normalizedType.typeSymbol == typeOf[Array[_]].typeSymbol && + normalizedType.typeArgs.headOption.exists(_ <:< typeOf[PipelineStage]) + } + + private def pipelineStageArrayComponentClass(tpe: Type): Class[_ <: PipelineStage] = { + val componentType = tpe.dealias.typeArgs.headOption.getOrElse { + throw new IllegalArgumentException(s"Expected a pipeline-stage array type but found $tpe") + } + val componentClass = Mirror.runtimeClass(componentType.erasure) + require( + classOf[PipelineStage].isAssignableFrom(componentClass), + s"Expected a PipelineStage array component but found ${componentClass.getName}" + ) + componentClass.asSubclass(classOf[PipelineStage]) + } + + def usesObjectSerializer(tpe: Type): Boolean = { + !(tpe <:< typeOf[PipelineStage]) && + !isPipelineStageArray(tpe) && + !(tpe <:< typeOf[Dataset[_]]) && + !(tpe <:< typeOf[DataType]) + } + def typeToSerializer[T](tpe: Type, sparkSession: SparkSession): Serializer[T] = { - (if (tpe <:< typeOf[PipelineStage]) new PipelineSerializer() - else if (tpe <:< typeOf[Array[PipelineStage]]) new PipelineArraySerializer() + typeToSerializer(tpe, sparkSession, None) + } + + def typeToSerializer[T]( + tpe: Type, + sparkSession: SparkSession, + classFilter: Option[DeserializationClassFilter]): Serializer[T] = { + (if (tpe <:< typeOf[PipelineStage]) new PipelineSerializer(sparkSession) + else if (isPipelineStageArray(tpe)) new RuntimeTypedPipelineArraySerializer[T]( + sparkSession, + pipelineStageArrayComponentClass(tpe) + ) else if (tpe <:< typeOf[Dataset[_]]) new DFSerializer(sparkSession) - else new ObjectSerializer(sparkSession)(typeToTypeTag(tpe))) + else if (tpe <:< typeOf[DataType]) new DataTypeSerializer(sparkSession) + else classFilter match { + case Some(filter) => + new FilteredObjectSerializer(sparkSession, filter)(typeToTypeTag(tpe)) + case None => + new ObjectSerializer(sparkSession)(typeToTypeTag(tpe)) + }) .asInstanceOf[Serializer[T]] } @@ -58,7 +119,107 @@ object Serializer { }.get } + private val PrimitiveByteArrayFilter = DeserializationClassFilter() + private val TrustedArtifactLoad = new ThreadLocal[java.lang.Boolean] + + private def defaultClassFilter(tpe: Type): Option[DeserializationClassFilter] = { + if (tpe =:= typeOf[Array[Byte]]) Some(PrimitiveByteArrayFilter) else None + } + + private def disabledDeserializationException(tpe: Type, trustedLoadInstructions: String): SecurityException = { + new SecurityException( + s"Java deserialization is disabled for $tpe because its object graph is not constrained. " + + s"$trustedLoadInstructions This may execute arbitrary code." + ) + } + + private[ml] def configuredLegacyDeserialization(spark: SparkSession): Boolean = { + spark.conf.getOption(LegacyObjectDeserializationConfig).exists(_.equalsIgnoreCase("true")) + } + + def trustedLoadEnabled(spark: SparkSession): Boolean = { + Option(TrustedArtifactLoad.get()) + .map(_.booleanValue()) + .getOrElse(configuredLegacyDeserialization(spark)) + } + + private[ml] def legacyObjectDeserializationEnabled(spark: SparkSession): Boolean = + trustedLoadEnabled(spark) + + private[ml] def currentTrustedArtifactLoad: Option[Boolean] = + Option(TrustedArtifactLoad.get()).map(_.booleanValue()) + + private[ml] def withTrustedArtifactLoad[T](enabled: Boolean)(action: => T): T = { + val previous = Option(TrustedArtifactLoad.get()) + TrustedArtifactLoad.set(enabled) + try { + action + } finally { + previous.fold(TrustedArtifactLoad.remove())(TrustedArtifactLoad.set) + } + } + + /** Opens a same-thread trusted scope for language bindings that cannot pass a Scala closure. + * The supplied session must explicitly enable [[LegacyObjectDeserializationConfig]]. + */ + def beginTrustedArtifactLoad(spark: SparkSession): TrustedArtifactLoadScope = { + TrustedArtifactLoadScope.begin( + spark, configuredLegacyDeserialization(spark), TrustedArtifactLoad) + } + /** Runs a legacy load with trust bound to the supplied session. + * + * This is required for native Spark Pipeline readers because Spark does not propagate their + * assigned session to nested stage readers. + */ + def withTrustedArtifactLoad[T](spark: SparkSession)(action: => T): T = { + val scope = beginTrustedArtifactLoad(spark) + try { + action + } finally { + scope.close() + } + } + + private[ml] def withActiveSession[T](spark: SparkSession)(action: => T): T = { + val previousSession = SparkSession.getActiveSession + SparkSession.setActiveSession(spark) + try { + action + } finally { + previousSession.fold(SparkSession.clearActiveSession())(SparkSession.setActiveSession) + } + } + + def read[A]( + is: InputStream, + classFilter: DeserializationClassFilter)(implicit ttag: TypeTag[A]): A = { + val safeInput = try new SafeObjectInputStream(is, classFilter) catch { + case NonFatal(error) => + try is.close() catch { case NonFatal(closeError) => error.addSuppressed(closeError) } + throw error + } + using(safeInput) { input => + input.readObject.asInstanceOf[A] + }.get + } + def read[A](is: InputStream)(implicit ttag: TypeTag[A]): A = { + defaultClassFilter(ttag.tpe) match { + case Some(filter) => read(is, filter) + case None => + try { + throw disabledDeserializationException( + ttag.tpe, + "Call Serializer.readUnsafe only if the input stream is a trusted legacy artifact." + ) + } finally { + is.close() + } + } + } + + /** Reads an object without a class filter. The caller must already trust the artifact. */ + def readUnsafe[A](is: InputStream)(implicit ttag: TypeTag[A]): A = { using(new ContextObjectInputStream(is)) { in => in.readObject.asInstanceOf[A] }.get @@ -93,11 +254,56 @@ object Serializer { * @return The loaded object. */ def readFromHDFS[O](spark: SparkSession, path: Path)(implicit ttag: TypeTag[O]): O = { + defaultClassFilter(ttag.tpe) match { + case Some(filter) => readFromHDFS(spark, path, filter) + case None if legacyObjectDeserializationEnabled(spark) => + readFromHDFSUnsafe(spark, path) + case None => + throw disabledDeserializationException( + ttag.tpe, + s"Set $LegacyObjectDeserializationConfig=true on the SparkSession only when loading " + + "a trusted legacy model." + ) + } + } + + def readFromHDFS[O]( + spark: SparkSession, + path: Path, + classFilter: DeserializationClassFilter)(implicit ttag: TypeTag[O]): O = { + try { + using(path.getFileSystem(sessionHadoopConf(spark)).open(path)) { in => + read[O](in, classFilter)(ttag) + }.get + } catch { + case error: InvalidClassException if legacyObjectDeserializationEnabled(spark) => + readFromHDFSUnsafe(spark, path) + case error: InvalidClassException => + val securityError = disabledDeserializationException( + ttag.tpe, + s"The object graph contains ${error.classname}, which is outside its approved class policy. " + + s"Set $LegacyObjectDeserializationConfig=true only when loading a trusted legacy model." + ) + securityError.initCause(error) + throw securityError + } + } + + /** Reads an object from Hadoop storage without a class filter. + * The caller must already trust the artifact. + */ + def readFromHDFSUnsafe[O]( + spark: SparkSession, + path: Path)(implicit ttag: TypeTag[O]): O = { using(path.getFileSystem(sessionHadoopConf(spark)).open(path)) { in => - read[O](in)(ttag) + readUnsafe[O](in)(ttag) }.get } + def isDirectory(spark: SparkSession, path: Path): Boolean = { + path.getFileSystem(sessionHadoopConf(spark)).getFileStatus(path).isDirectory + } + def makeQualifiedPath(spark: SparkSession, path: String): Path = { makeQualifiedPath(sessionHadoopConf(spark), path) } @@ -118,7 +324,19 @@ class ObjectSerializer[O](spark: SparkSession)(implicit ttag: TypeTag[O]) extend def read(path: Path): O = Serializer.readFromHDFS(spark, path) } +private[ml] class FilteredObjectSerializer[O]( + spark: SparkSession, + classFilter: DeserializationClassFilter)(implicit ttag: TypeTag[O]) extends Serializer[O] { + + def write(obj: O, path: Path, overwrite: Boolean): Unit = + Serializer.writeToHDFS(spark, obj, path, overwrite) + + def read(path: Path): O = Serializer.readFromHDFS(spark, path, classFilter) +} + class DFSerializer(spark: SparkSession) extends Serializer[DataFrame] { + private val globPathsOption = "__globPaths__" + def write(df: DataFrame, outputPath: Path, overwrite: Boolean): Unit = { val saveMode = if (overwrite) SaveMode.Overwrite @@ -128,28 +346,455 @@ class DFSerializer(spark: SparkSession) extends Serializer[DataFrame] { } def read(path: Path): DataFrame = { - spark.read.format("parquet").load(path.toString) + val resolvedPath = ArtifactPathResolver.resolveTreeInside(spark, path, "DataFrame parameter") + spark.read + .option(globPathsOption, "false") + .format("parquet") + .load(resolvedPath.toString) + } +} + +private[ml] object PipelineSerializer { + + private implicit val JsonFormats: DefaultFormats.type = DefaultFormats + + private val TrustedStageClassPrefixes = Set( + "com.microsoft.azure.synapse.ml.", + "org.apache.spark.ml." + ) + + private val PipelinePersistence = "pipeline" + private val DefaultParamsPersistence = "defaultParams" + private val ComplexParamsPersistence = "complexParams" + private val NativePersistence = "native" + private val LegacyPersistence = "legacy" + private val PersistenceKindsKey = "stagePersistenceKinds" + private val GlobCharacters = Set('*', '?', '[', ']', '{', '}', '\\') + + private def metadata( + spark: SparkSession, + path: Path, + expectedClassName: String = ""): DefaultParamsReader.Metadata = { + ArtifactPathResolver.loadMetadata(spark, path, expectedClassName) + } + + private def normalizedQualifiedPath(spark: SparkSession, path: Path): java.net.URI = { + val conf = Serializer.sessionHadoopConf(spark) + val fs = path.getFileSystem(conf) + path.makeQualified(fs.getUri, fs.getWorkingDirectory).toUri.normalize() + } + + private def sameFileSystem(left: java.net.URI, right: java.net.URI): Boolean = { + Option(left.getScheme) == Option(right.getScheme) && + Option(left.getAuthority) == Option(right.getAuthority) + } + + private def stagePath( + spark: SparkSession, + path: Path, + uid: String, + index: Int, + stageCount: Int): Path = { + val stagesPath = new Path(path, "stages") + val normalizedStagesUri = normalizedQualifiedPath(spark, stagesPath) + val stagesPrefix = normalizedStagesUri.getPath.stripSuffix("/") + "/" + val resolvedStagePath = new Path(stagesPath, stageName(uid, index, stageCount)) + val normalizedStageUri = normalizedQualifiedPath(spark, resolvedStagePath) + require( + sameFileSystem(normalizedStagesUri, normalizedStageUri) && + normalizedStageUri.getPath.startsWith(stagesPrefix), + s"Pipeline stage UID $uid resolves outside its stages directory at $stagesPath" + ) + resolvedStagePath + } + + private def stageName(uid: String, index: Int, stageCount: Int): String = { + val indexString = ("%0" + stageCount.toString.length + "d").format(index) + val name = s"${indexString}_$uid" + val relativePath = new Path(name) + require( + !relativePath.isAbsolute && + relativePath.toUri.getScheme == null && + relativePath.getName == name && + !name.exists(GlobCharacters), + s"Pipeline stage UID $uid resolves outside its stages directory" + ) + name + } + + private def readStagePath( + spark: SparkSession, + path: Path, + uid: String, + index: Int, + stageCount: Int): Path = { + val stagesPath = ArtifactPathResolver.resolvePathInside( + spark, + path, + new Path(path, "stages"), + "Pipeline stages directory" + ) + val candidate = new Path(stagesPath, stageName(uid, index, stageCount)) + ArtifactPathResolver.resolvePathInside(spark, stagesPath, candidate, s"Pipeline stage $uid") + } + + private def prepareOutputPath( + spark: SparkSession, + path: Path, + overwrite: Boolean): Unit = { + val conf = Serializer.sessionHadoopConf(spark) + val fs = path.getFileSystem(conf) + val qualifiedPath = path.makeQualified(fs.getUri, fs.getWorkingDirectory) + if (fs.exists(qualifiedPath)) { + if (overwrite) { + fs.delete(qualifiedPath, true) + } else { + throw new IOException(s"Path $path already exists and overwrite is disabled") + } + } + } + + private def canWriteStage(stage: PipelineStage): Boolean = { + stage match { + case pipeline: Pipeline if pipeline.getClass == classOf[Pipeline] => + pipeline.getStages.forall(canWriteStage) + case model: PipelineModel if model.getClass == classOf[PipelineModel] => + model.stages.forall(canWriteStage) + case _: MLWritable => true + case _ => false + } + } + + def canWriteStages(stages: Array[PipelineStage]): Boolean = stages.forall(canWriteStage) + + private def writeDefaultParams( + spark: SparkSession, + stage: PipelineStage, + path: Path): Unit = { + ComplexParamsWriter.saveMetadata(stage, path.toString, spark) + } + + private def writeWritableStage( + spark: SparkSession, + stage: PipelineStage, + writable: MLWritable, + path: Path): Unit = { + Serializer.withActiveSession(spark) { + val writer = writable.write + if (writer.getClass == classOf[DefaultParamsWriter]) { + writeDefaultParams(spark, stage, path) + } else { + writer.session(spark).save(path.toString) + } + } + } + + private def writeStage(spark: SparkSession, stage: PipelineStage, path: Path): Unit = { + stage match { + case pipeline: Pipeline if pipeline.getClass == classOf[Pipeline] => + writePipeline(spark, pipeline, pipeline.getStages, path, overwrite = false) + case model: PipelineModel if model.getClass == classOf[PipelineModel] => + writePipeline( + spark, + model, + model.stages.map(stage => stage: PipelineStage), + path, + overwrite = false + ) + case writable: MLWritable => + writeWritableStage(spark, stage, writable, path) + case _ => + throw new UnsupportedOperationException( + s"Cannot safely persist non-writable pipeline stage ${stage.uid} of type ${stage.getClass.getName}" + ) + } + } + + private def writePipeline( + spark: SparkSession, + instance: Params, + stages: Array[PipelineStage], + path: Path, + overwrite: Boolean): Unit = { + prepareOutputPath(spark, path, overwrite) + val stageUids = JArray(stages.map(stage => JString(stage.uid)).toList) + val persistenceKinds = JArray(stages.map(stage => JString(persistenceKind(stage))).toList) + ComplexParamsWriter.saveMetadata( + instance, + path.toString, + spark, + paramMap = Some(JObject( + "stageUids" -> stageUids, + PersistenceKindsKey -> persistenceKinds + )) + ) + stages.zipWithIndex.foreach { case (stage, index) => + writeStage(spark, stage, stagePath(spark, path, stage.uid, index, stages.length)) + } + } + + def writeWrappedStages( + spark: SparkSession, + stages: Array[PipelineStage], + path: Path, + overwrite: Boolean): Unit = { + ModelLoadContext.withWriteContext { + val wrapper = new Pipeline().setStages(stages) + writePipeline(spark, wrapper, stages, path, overwrite) + } + } + + private def readDefaultParams( + spark: SparkSession, + stageClass: Class[_], + stageMetadata: DefaultParamsReader.Metadata): PipelineStage = { + Serializer.withActiveSession(spark) { + val instance = stageClass.getConstructor(classOf[String]) + .newInstance(stageMetadata.uid) + .asInstanceOf[Params] + stageMetadata.getAndSetParams(instance) + instance.asInstanceOf[PipelineStage] + } + } + + private def persistenceKind( + stage: PipelineStage): String = { + if (stage.getClass == classOf[Pipeline] || stage.getClass == classOf[PipelineModel]) { + PipelinePersistence + } else if (StageReaderInspector.usesDefaultParamsReader(stage.getClass)) { + DefaultParamsPersistence + } else if (StageReaderInspector.usesComplexParamsReader(stage.getClass)) { + ComplexParamsPersistence + } else { + NativePersistence + } + } + + private def codePersistenceKind(stageClass: Class[_]): String = { + if (StageReaderInspector.usesDefaultParamsReader(stageClass)) { + DefaultParamsPersistence + } else if (StageReaderInspector.usesComplexParamsReader(stageClass)) { + ComplexParamsPersistence + } else { + NativePersistence + } + } + + private def readComplexParams( + spark: SparkSession, + path: Path, + className: String, + stageClass: Class[_], + stageMetadata: DefaultParamsReader.Metadata): PipelineStage = { + Serializer.withActiveSession(spark) { + new ComplexParamsReader[PipelineStage]( + className, + Some(stageClass), + Some(stageMetadata) + ) + .session(spark) + .load(path.toString) + } + } + + private def loadWithReader( + spark: SparkSession, + path: Path, + reader: MLReader[_]): PipelineStage = { + Serializer.withActiveSession(spark) { + reader.session(spark).load(path.toString).asInstanceOf[PipelineStage] + } + } + + private def readCustomStage( + spark: SparkSession, + path: Path, + stageMetadata: DefaultParamsReader.Metadata, + persistence: String): PipelineStage = { + val className = stageMetadata.className + val trustedClassName = TrustedStageClassPrefixes.exists(className.startsWith) + if (!trustedClassName && !Serializer.legacyObjectDeserializationEnabled(spark)) { + throw new SecurityException( + s"Stage class $className at $path is outside the approved package policy. " + + s"Set ${Serializer.LegacyObjectDeserializationConfig}=true only when loading " + + "a trusted legacy model." + ) + } + val stageClass = Utils.classForName[PipelineStage](className, initialize = false) + require( + classOf[PipelineStage].isAssignableFrom(stageClass), + s"Expected a PipelineStage at $path but metadata declared $className" + ) + val expectedPersistence = codePersistenceKind(stageClass) + require( + persistence == LegacyPersistence || persistence == expectedPersistence, + s"Stage $className at $path declares persistence kind $persistence, " + + s"but its installed reader requires $expectedPersistence" + ) + persistence match { + case DefaultParamsPersistence => readDefaultParams(spark, stageClass, stageMetadata) + case ComplexParamsPersistence => + readComplexParams(spark, path, className, stageClass, stageMetadata) + case NativePersistence | LegacyPersistence => + if (!Serializer.legacyObjectDeserializationEnabled(spark)) { + throw new SecurityException( + s"Session-backed persistence is unavailable for stage $className at $path. " + + s"Set ${Serializer.LegacyObjectDeserializationConfig}=true only when loading " + + "a trusted legacy model." + ) + } + val reader = Serializer.withActiveSession(spark) { + stageClass.getMethod("read") + .invoke(None.orNull) + .asInstanceOf[MLReader[_]] + } + loadWithReader(spark, path, reader) + case _ => + throw new SecurityException( + s"Stage $className at $path declares unsupported persistence kind $persistence" + ) + } + } + + private def readPipelineModel( + spark: SparkSession, + path: Path, + stageMetadata: DefaultParamsReader.Metadata, + context: ModelLoadContext): PipelineModel = { + val stages = readPipelineStages(spark, path, stageMetadata, context).map { + case transformer: org.apache.spark.ml.Transformer => transformer + case stage => + throw new IllegalArgumentException( + s"PipelineModel at $path contains non-transformer stage ${stage.uid}" + ) + } + new PipelineModel(stageMetadata.uid, stages) + } + + private def readStage( + spark: SparkSession, + path: Path, + persistence: String, + context: ModelLoadContext): PipelineStage = { + context.enterPath(path, "Pipeline stage path") + val stageMetadata = metadata(spark, path) + stageMetadata.className match { + case className if className == classOf[Pipeline].getName => + require( + persistence == PipelinePersistence || persistence == LegacyPersistence, + s"Pipeline at $path declares incompatible persistence kind $persistence" + ) + context.withNestedArtifact(path) { + val stages = readPipelineStages(spark, path, stageMetadata, context) + new Pipeline(stageMetadata.uid).setStages(stages) + } + case className if className == classOf[PipelineModel].getName => + require( + persistence == PipelinePersistence || persistence == LegacyPersistence, + s"PipelineModel at $path declares incompatible persistence kind $persistence" + ) + context.withNestedArtifact(path) { + readPipelineModel(spark, path, stageMetadata, context) + } + case _ => + require( + persistence != PipelinePersistence, + s"Non-pipeline stage ${stageMetadata.className} at $path declares pipeline persistence" + ) + readCustomStage(spark, path, stageMetadata, persistence) + } + } + + private def readPipelineStages( + spark: SparkSession, + path: Path, + pipelineMetadata: DefaultParamsReader.Metadata, + context: ModelLoadContext): Array[PipelineStage] = { + val stageUids = (pipelineMetadata.params \ "stageUids").extract[Seq[String]] + val persistenceKinds = (pipelineMetadata.params \ PersistenceKindsKey) match { + case org.json4s.JNothing => Seq.fill(stageUids.length)(LegacyPersistence) + case value => value.extract[Seq[String]] + } + require( + persistenceKinds.length == stageUids.length, + s"Pipeline at $path declares ${persistenceKinds.length} persistence kinds " + + s"for ${stageUids.length} stages" + ) + stageUids.zip(persistenceKinds).zipWithIndex.map { + case ((uid, persistence), index) => + readStage( + spark, + readStagePath(spark, path, uid, index, stageUids.length), + persistence, + context + ) + }.toArray + } + + def readWrappedStages(spark: SparkSession, path: Path): Array[PipelineStage] = { + val resolvedPath = ArtifactPathResolver.resolvePathInside( + spark, + path, + path, + "Pipeline root", + allowRoot = true + ) + ModelLoadContext.withContext(resolvedPath) { context => + val wrapperMetadata = metadata(spark, resolvedPath, classOf[Pipeline].getName) + readPipelineStages(spark, resolvedPath, wrapperMetadata, context) + } + } + + def prepareUnsafeOutput( + spark: SparkSession, + path: Path, + overwrite: Boolean): Unit = { + prepareOutputPath(spark, path, overwrite) } } -class PipelineSerializer extends Serializer[PipelineStage] { +class PipelineSerializer(spark: SparkSession) extends Serializer[PipelineStage] { + + def this() = this(SparkSession.builder().getOrCreate()) + def write(stage: PipelineStage, outputPath: Path, overwrite: Boolean): Unit = { - val pipe = new Pipeline().setStages(Array(stage)) - Serializer.writeMLWritable(pipe, outputPath, overwrite) + if (PipelineSerializer.canWriteStages(Array(stage))) { + PipelineSerializer.writeWrappedStages(spark, Array(stage), outputPath, overwrite) + } else { + PipelineSerializer.prepareUnsafeOutput(spark, outputPath, overwrite) + Serializer.writeToHDFS(spark, stage, outputPath, overwrite = false) + } } def read(path: Path): PipelineStage = { - Pipeline.load(path.toString).getStages(0) + if (Serializer.isDirectory(spark, path)) { + val stages = PipelineSerializer.readWrappedStages(spark, path) + require(stages.length == 1, s"Expected one pipeline stage at $path but found ${stages.length}") + stages.head + } else { + Serializer.readFromHDFS[PipelineStage](spark, path) + } } } -class PipelineArraySerializer extends Serializer[Array[PipelineStage]] { +class PipelineArraySerializer(spark: SparkSession) extends Serializer[Array[PipelineStage]] { + + def this() = this(SparkSession.builder().getOrCreate()) + def write(stages: Array[PipelineStage], outputPath: Path, overwrite: Boolean): Unit = { - val pipe = new Pipeline().setStages(stages) - Serializer.writeMLWritable(pipe, outputPath, overwrite) + if (PipelineSerializer.canWriteStages(stages)) { + PipelineSerializer.writeWrappedStages(spark, stages, outputPath, overwrite) + } else { + PipelineSerializer.prepareUnsafeOutput(spark, outputPath, overwrite) + Serializer.writeToHDFS(spark, stages, outputPath, overwrite = false) + } } def read(path: Path): Array[PipelineStage] = { - Pipeline.load(path.toString).getStages + if (Serializer.isDirectory(spark, path)) { + PipelineSerializer.readWrappedStages(spark, path) + } else { + Serializer.readFromHDFS[Array[PipelineStage]](spark, path) + } } } diff --git a/core/src/main/scala/org/apache/spark/ml/StageReaderInspector.scala b/core/src/main/scala/org/apache/spark/ml/StageReaderInspector.scala new file mode 100644 index 00000000000..c02621080c8 --- /dev/null +++ b/core/src/main/scala/org/apache/spark/ml/StageReaderInspector.scala @@ -0,0 +1,225 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package org.apache.spark.ml + +import com.microsoft.azure.synapse.ml.core.env.StreamUtilities.using + +import java.io.{BufferedInputStream, ByteArrayInputStream, DataInputStream, IOException} +import java.nio.charset.StandardCharsets + +private[ml] object StageReaderInspector { + + // The JVM class-file format is defined by numeric tags and length-driven loops. + //scalastyle:off magic.number + //scalastyle:off while + + private sealed trait ConstantPoolEntry + private case class Utf8Entry(value: String) extends ConstantPoolEntry + private case class ClassEntry(nameIndex: Int) extends ConstantPoolEntry + private case class NameAndTypeEntry(nameIndex: Int, descriptorIndex: Int) extends ConstantPoolEntry + private case class MethodEntry(classIndex: Int, nameAndTypeIndex: Int) extends ConstantPoolEntry + private case object OtherEntry extends ConstantPoolEntry + + private val DefaultParamsReadableName = "org/apache/spark/ml/util/DefaultParamsReadable" + private val ComplexParamsReadableName = "org/apache/spark/ml/ComplexParamsReadable" + private val ReaderDescriptor = "()Lorg/apache/spark/ml/util/MLReader;" + private val ReaderReturnDescriptor = ")Lorg/apache/spark/ml/util/MLReader;" + + def usesDefaultParamsReader(stageClass: Class[_]): Boolean = + delegatesToReaderTrait(stageClass, DefaultParamsReadableName) + + def usesComplexParamsReader(stageClass: Class[_]): Boolean = + delegatesToReaderTrait(stageClass, ComplexParamsReadableName) + + private def delegatesToReaderTrait( + stageClass: Class[_], + readerTraitName: String): Boolean = { + val className = stageClass.getName + val resourceName = className.replace('.', '/') + "$.class" + Option(stageClass.getResourceAsStream("/" + resourceName)) match { + case Some(stream) => + using(new DataInputStream(new BufferedInputStream(stream))) { input => + inspectClass(input, resourceName.stripSuffix(".class"), readerTraitName) + }.get + case None => false + } + } + + private def inspectClass( + input: DataInputStream, + expectedClassName: String, + readerTraitName: String): Boolean = { + require(input.readInt() == 0xcafebabe, "Invalid JVM class file") + input.readUnsignedShort() + input.readUnsignedShort() + val constantPool = readConstantPool(input) + input.readUnsignedShort() + val className = constantPool(input.readUnsignedShort()) match { + case ClassEntry(nameIndex) => utf8(constantPool, nameIndex) + case _ => throw new IOException("JVM class file has no this_class entry") + } + require(className == expectedClassName, s"Expected class $expectedClassName but inspected $className") + input.readUnsignedShort() + skipInterfaces(input) + skipMembers(input, constantPool) + inspectMethods(input, constantPool, readerTraitName) + } + + //scalastyle:off cyclomatic.complexity + private def readConstantPool(input: DataInputStream): Array[ConstantPoolEntry] = { + val entries = Array.fill[ConstantPoolEntry](input.readUnsignedShort())(OtherEntry) + var index = 1 + while (index < entries.length) { + input.readUnsignedByte() match { + case 1 => + val bytes = new Array[Byte](input.readUnsignedShort()) + input.readFully(bytes) + entries(index) = Utf8Entry(new String(bytes, StandardCharsets.UTF_8)) + case 7 => + entries(index) = ClassEntry(input.readUnsignedShort()) + case 9 | 10 | 11 => + entries(index) = MethodEntry(input.readUnsignedShort(), input.readUnsignedShort()) + case 12 => + entries(index) = NameAndTypeEntry(input.readUnsignedShort(), input.readUnsignedShort()) + case 3 | 4 => + skipFully(input, 4) + case 5 | 6 => + skipFully(input, 8) + index += 1 + case 8 | 16 | 19 | 20 => + skipFully(input, 2) + case 15 => + skipFully(input, 3) + case 17 | 18 => + skipFully(input, 4) + case tag => + throw new IOException(s"Unsupported class-file constant-pool tag $tag") + } + index += 1 + } + entries + } + + private def skipInterfaces(input: DataInputStream): Unit = { + skipFully(input, input.readUnsignedShort().toLong * 2) + } + + private def skipMembers( + input: DataInputStream, + constantPool: Array[ConstantPoolEntry]): Unit = { + (0 until input.readUnsignedShort()).foreach { _ => + input.readUnsignedShort() + input.readUnsignedShort() + input.readUnsignedShort() + skipAttributes(input, constantPool) + } + } + + private def inspectMethods( + input: DataInputStream, + constantPool: Array[ConstantPoolEntry], + readerTraitName: String): Boolean = { + var readerDelegatesToTrait = false + (0 until input.readUnsignedShort()).foreach { _ => + input.readUnsignedShort() + val methodName = utf8(constantPool, input.readUnsignedShort()) + val descriptor = utf8(constantPool, input.readUnsignedShort()) + (0 until input.readUnsignedShort()).foreach { _ => + val attributeName = utf8(constantPool, input.readUnsignedShort()) + val attributeLength = readUnsignedInt(input) + val attributeBytes = readBytes(input, attributeLength) + if (methodName == "read" && descriptor == ReaderDescriptor && attributeName == "Code") { + readerDelegatesToTrait ||= isTraitForwarder( + attributeBytes, + constantPool, + readerTraitName + ) + } + } + } + readerDelegatesToTrait + } + + private def isTraitForwarder( + codeAttribute: Array[Byte], + constantPool: Array[ConstantPoolEntry], + readerTraitName: String): Boolean = { + using(new DataInputStream(new ByteArrayInputStream(codeAttribute))) { input => + input.readUnsignedShort() + input.readUnsignedShort() + val code = readBytes(input, readUnsignedInt(input)) + code.length == 5 && + (code(0) & 0xff) == 0x2a && + (code(1) & 0xff) == 0xb8 && + (code(4) & 0xff) == 0xb0 && + methodMatches( + constantPool, + ((code(2) & 0xff) << 8) | (code(3) & 0xff), + readerTraitName + ) + }.get + } + + private def methodMatches( + constantPool: Array[ConstantPoolEntry], + methodIndex: Int, + ownerName: String): Boolean = { + constantPool(methodIndex) match { + case MethodEntry(classIndex, nameAndTypeIndex) => + val className = constantPool(classIndex) match { + case ClassEntry(nameIndex) => utf8(constantPool, nameIndex) + case _ => "" + } + val (methodName, descriptor) = constantPool(nameAndTypeIndex) match { + case NameAndTypeEntry(nameIndex, descriptorIndex) => + utf8(constantPool, nameIndex) -> utf8(constantPool, descriptorIndex) + case _ => "" -> "" + } + className == ownerName && methodName == "read$" && + descriptor.endsWith(ReaderReturnDescriptor) + case _ => false + } + } + + private def skipAttributes( + input: DataInputStream, + constantPool: Array[ConstantPoolEntry]): Unit = { + (0 until input.readUnsignedShort()).foreach { _ => + utf8(constantPool, input.readUnsignedShort()) + skipFully(input, readUnsignedInt(input)) + } + } + + private def utf8( + constantPool: Array[ConstantPoolEntry], + index: Int): String = { + constantPool(index) match { + case Utf8Entry(value) => value + case _ => throw new IOException(s"Constant-pool entry $index is not UTF-8") + } + } + + private def readUnsignedInt(input: DataInputStream): Long = + Integer.toUnsignedLong(input.readInt()) + + private def readBytes(input: DataInputStream, length: Long): Array[Byte] = { + require(length <= Int.MaxValue, s"Class-file attribute is too large: $length bytes") + val bytes = new Array[Byte](length.toInt) + input.readFully(bytes) + bytes + } + + private def skipFully(input: DataInputStream, length: Long): Unit = { + var remaining = length + while (remaining > 0) { + val skipped = input.skip(remaining) + if (skipped == 0) { + require(input.read() >= 0, "Unexpected end of JVM class file") + remaining -= 1 + } else { + remaining -= skipped + } + } + } +} diff --git a/core/src/test/python/synapsemltest/recommendation/test_ranking.py b/core/src/test/python/synapsemltest/recommendation/test_ranking.py index 3d58e8354a3..eab750e59a0 100644 --- a/core/src/test/python/synapsemltest/recommendation/test_ranking.py +++ b/core/src/test/python/synapsemltest/recommendation/test_ranking.py @@ -184,7 +184,16 @@ def test_sar_string_model_save_load(self): with tempfile.TemporaryDirectory() as directory: path = directory + "/sar-model" model.write().overwrite().save(path) - loaded = SARModel.load(path) + config = "spark.synapseml.legacy.allowUnsafeJavaDeserialization" + previous = spark.conf.get(config, None) + spark.conf.set(config, "true") + try: + loaded = SARModel.read().session(spark).load(path) + finally: + if previous is None: + spark.conf.unset(config) + else: + spark.conf.set(config, previous) self.assertEqual( loaded.recommendForAllUsers(2).orderBy("user").collect(), model.recommendForAllUsers(2).orderBy("user").collect(), diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/causal/VerifySyntheticDiffInDiffEstimator.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/causal/VerifySyntheticDiffInDiffEstimator.scala index da13314f171..19cd8f0f2ba 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/causal/VerifySyntheticDiffInDiffEstimator.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/causal/VerifySyntheticDiffInDiffEstimator.scala @@ -129,7 +129,7 @@ class VerifySyntheticDiffInDiffEstimator new TestObject(estimator, df, df) ) - override def reader: MLReadable[_] = SyntheticControlEstimator + override def reader: MLReadable[_] = SyntheticDiffInDiffEstimator override def modelReader: MLReadable[_] = DiffInDiffModel } diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/core/serialize/ValidateComplexParamSerializer.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/core/serialize/ValidateComplexParamSerializer.scala index d349e0b20b9..b6fa5585580 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/core/serialize/ValidateComplexParamSerializer.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/core/serialize/ValidateComplexParamSerializer.scala @@ -5,16 +5,74 @@ package com.microsoft.azure.synapse.ml.core.serialize import com.microsoft.azure.synapse.ml.core.env.StreamUtilities.using import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import com.microsoft.azure.synapse.ml.core.utils.DeserializationClassFilter import com.microsoft.azure.synapse.ml.param.ByteArrayParam +import com.microsoft.azure.synapse.ml.stages.{Timer, TimerModel} import org.apache.commons.io.FileUtils import org.apache.hadoop.fs.Path +import org.apache.spark.ml.classification.LogisticRegression +import org.apache.spark.ml.feature.{HashingTF, SQLTransformer} import org.apache.spark.ml.param.{Param, ParamMap, Params} import org.apache.spark.ml.util._ -import org.apache.spark.ml.{ComplexParamsReadable, ComplexParamsWritable, ObjectSerializer, Serializer, Transformer} +import org.apache.spark.ml.{ + ComplexParamsReadable, + ComplexParamsWritable, + ObjectSerializer, + Pipeline, + PipelineArraySerializer, + PipelineStage, + Serializer, + Transformer +} import org.apache.spark.sql.types.StructType import org.apache.spark.sql.{DataFrame, Dataset} -import java.io.File +import java.io.{ + ByteArrayInputStream, + ByteArrayOutputStream, + File, + FileOutputStream, + InvalidClassException, + ObjectInputStream, + StreamCorruptedException +} +import java.nio.file.Files +import java.util.concurrent.atomic.AtomicBoolean +import java.util.zip.GZIPOutputStream + +object DeserializationTripwire { + val Triggered = new AtomicBoolean(false) +} + +@SerialVersionUID(1L) +class DeserializationTripwire extends Serializable { + private def readObject(input: ObjectInputStream): Unit = { + DeserializationTripwire.Triggered.set(true) + input.defaultReadObject() + } +} + +object TestableReaderFactoryTripwireStage extends MLReadable[TestableReaderFactoryTripwireStage] { + val Triggered = new AtomicBoolean(false) + + override def read: MLReader[TestableReaderFactoryTripwireStage] = { + Triggered.set(true) + new MLReader[TestableReaderFactoryTripwireStage] { + override def load(path: String): TestableReaderFactoryTripwireStage = + new TestableReaderFactoryTripwireStage("loaded") + } + } +} + +class TestableReaderFactoryTripwireStage(override val uid: String) extends TestEstimatorBase(uid) + +class UnsafePayloadParam(parent: Params, name: String, doc: String) + extends ComplexParam[DeserializationTripwire]( + parent, + name, + doc, + (_: DeserializationTripwire) => true + ) class TestEstimatorBase(val uid: String) extends Transformer { def this() = this(Identifiable.randomUID("TestEstimatorBase")) @@ -43,6 +101,14 @@ trait HasStringParam extends Params { def setStringParam(value: String): this.type = set(stringParam, value) } +trait HasUnsafePayloadParam extends Params { + val unsafePayload = new UnsafePayloadParam(this, "unsafePayload", "test-only unsafe payload") + + def getUnsafePayload: DeserializationTripwire = $(unsafePayload) + + def setUnsafePayload(value: DeserializationTripwire): this.type = set(unsafePayload, value) +} + class ComplexParamTest(override val uid: String) extends TestEstimatorBase(uid) with HasByteArrayParam with ComplexParamsWritable { def this() = this(Identifiable.randomUID("ComplexParamTest")) @@ -64,10 +130,73 @@ class MixedParamTest(override val uid: String) extends TestEstimatorBase(uid) object MixedParamTest extends ComplexParamsReadable[MixedParamTest] +class TestableUnsafeComplexParamTest(override val uid: String) extends TestEstimatorBase(uid) + with HasUnsafePayloadParam with ComplexParamsWritable { + def this() = this(Identifiable.randomUID("TestableUnsafeComplexParamTest")) +} + +object TestableUnsafeComplexParamTest extends ComplexParamsReadable[TestableUnsafeComplexParamTest] + +class TestableDefaultUnsafeComplexParamTest(override val uid: String) extends TestEstimatorBase(uid) + with HasUnsafePayloadParam with ComplexParamsWritable { + def this() = this(Identifiable.randomUID("TestableDefaultUnsafeComplexParamTest")) + + setDefault(unsafePayload -> new DeserializationTripwire) +} + +object TestableDefaultUnsafeComplexParamTest + extends ComplexParamsReadable[TestableDefaultUnsafeComplexParamTest] + class ValidateComplexParamSerializer extends TestBase { val saveFile = new File(tmpDir.toFile, "m1.model").toString val saveFile2 = new File(tmpDir.toFile, "m2.model").toString + private class CloseTrackingInputStream(bytes: Array[Byte]) extends ByteArrayInputStream(bytes) { + var closed: Boolean = false + + override def close(): Unit = { + closed = true + super.close() + } + } + + private def rewriteMetadata(path: String)(update: String => String): Unit = { + val metadataDir = new File(path, "metadata") + val metadataJson = spark.read.text(metadataDir.toString).first().getString(0) + val updatedMetadata = update(metadataJson) + assert(updatedMetadata !== metadataJson) + FileUtils.deleteDirectory(metadataDir) + spark.createDataFrame(Seq(Tuple1(updatedMetadata))).toDF("value").write.text(metadataDir.toString) + } + + private def restoreConfig(key: String, previous: Option[String]): Unit = { + previous match { + case Some(value) => spark.conf.set(key, value) + case None => spark.conf.unset(key) + } + } + + private def compressMetadata(path: String): Unit = { + val metadataDir = new File(path, "metadata") + val partFile = metadataDir.listFiles() + .find(file => file.getName.startsWith("part-") && !file.getName.endsWith(".crc")) + .get + val compressedFile = new File(metadataDir, partFile.getName + ".gz") + val bytes = Files.readAllBytes(partFile.toPath) + using(new GZIPOutputStream(new FileOutputStream(compressedFile))) { + _.write(bytes) + }.get + Files.delete(partFile.toPath) + } + + private def gzipBytes(bytes: Array[Byte]): Array[Byte] = { + val output = new ByteArrayOutputStream() + using(new GZIPOutputStream(output)) { + _.write(bytes) + }.get + output.toByteArray + } + test("Complex Param serialization should work on all complex, all normal, or mixed") { spark @@ -126,6 +255,125 @@ class ValidateComplexParamSerializer extends TestBase { assert(mpt1.getStringParam === mpt2.getStringParam) } + test("Complex Param serialization should read compressed metadata") { + spark + val original = new MixedParamTest("compressed") + .setByteArray(Array[Byte](1, 2, 3)) + .setStringParam("value") + original.write.overwrite().save(saveFile) + compressMetadata(saveFile) + + val loaded = MixedParamTest.load(saveFile) + assert(loaded.getByteArray === original.getByteArray) + assert(loaded.getStringParam === original.getStringParam) + } + + test("Complex Param serialization bounds compressed metadata input bytes") { + spark + new MixedParamTest("physical-metadata-limit") + .setByteArray(Array[Byte](1)) + .setStringParam("value") + .write.overwrite().save(saveFile) + val metadataDir = new File(saveFile, "metadata") + val partFile = metadataDir.listFiles() + .find(file => file.getName.startsWith("part-") && !file.getName.endsWith(".crc")) + .get + val metadataBytes = Files.readAllBytes(partFile.toPath) + val compressedFile = new File(metadataDir, partFile.getName + ".gz") + val emptyMember = gzipBytes(Array.empty[Byte]) + val maxPhysicalBytes = 2 * 1048576 // scalastyle:ignore magic.number + val emptyMemberCount = maxPhysicalBytes / emptyMember.length + 1 + using(new FileOutputStream(compressedFile)) { output => + (0 until emptyMemberCount).foreach { _ => + output.write(emptyMember) + } + output.write(gzipBytes(metadataBytes)) + }.get + Files.delete(partFile.toPath) + + val error = intercept[IllegalArgumentException] { + MixedParamTest.load(saveFile) + } + assert(error.getMessage.contains("physical input")) + } + + test("Complex Param serialization rejects canonical metadata part aliases") { + spark + new MixedParamTest("metadata-alias") + .setByteArray(Array[Byte](1)) + .setStringParam("value") + .write.overwrite().save(saveFile) + val metadataDir = new File(saveFile, "metadata") + val partFile = metadataDir.listFiles() + .find(file => file.getName.startsWith("part-") && !file.getName.endsWith(".crc")) + .get + Files.createSymbolicLink( + new File(metadataDir, "part-alias").toPath, + partFile.toPath + ) + + val error = intercept[IllegalArgumentException] { + MixedParamTest.load(saveFile) + } + assert(error.getMessage.contains("encountered more than once")) + } + + test("Complex Param serialization bounds aggregate metadata part bytes") { + spark + new MixedParamTest("metadata-aggregate") + .setByteArray(Array[Byte](1)) + .setStringParam("value") + .write.overwrite().save(saveFile) + val metadataDir = new File(saveFile, "metadata") + val target = new File(metadataDir, "empty-members.gz") + val emptyMember = gzipBytes(Array.empty[Byte]) + val targetBytes = 1100000 // scalastyle:ignore magic.number + using(new FileOutputStream(target)) { output => + (0 until targetBytes / emptyMember.length + 1).foreach { _ => + output.write(emptyMember) + } + }.get + Files.createLink(new File(metadataDir, "part-empty-a.gz").toPath, target.toPath) + Files.createLink(new File(metadataDir, "part-empty-b.gz").toPath, target.toPath) + + val error = intercept[IllegalArgumentException] { + MixedParamTest.load(saveFile) + } + assert(error.getMessage.contains("directory limit")) + } + + test("Complex Param serialization rejects oversized metadata") { + spark + new MixedParamTest("oversized") + .setByteArray(Array[Byte](1)) + .setStringParam("value") + .write.overwrite().save(saveFile) + val oversizedSuffix = "x" * 1048576 // scalastyle:ignore magic.number + rewriteMetadata(saveFile)(_ + oversizedSuffix) + + val error = intercept[IllegalArgumentException] { + MixedParamTest.load(saveFile) + } + assert(error.getMessage.contains("exceeds")) + } + + test("Complex Param serialization bounds local metadata directory enumeration") { + spark + new MixedParamTest("metadata-entries") + .setByteArray(Array[Byte](1)) + .setStringParam("value") + .write.overwrite().save(saveFile) + val metadataPath = new File(saveFile, "metadata").toPath + (0 to 256).foreach { index => // scalastyle:ignore magic.number + Files.createFile(metadataPath.resolve(f"extra-$index%03d")) + } + + val error = intercept[IllegalArgumentException] { + MixedParamTest.load(saveFile) + } + assert(error.getMessage.contains("too many entries")) + } + test("Objects written the way earlier versions wrote them still load through the session path") { spark val obj = "round-trip payload".toCharArray.map(_.toByte) @@ -143,6 +391,388 @@ class ValidateComplexParamSerializer extends TestBase { assert(Serializer.readFromHDFS[Array[Byte]](spark, legacyPath) === obj) } + test("Serializer rejects unconstrained Java objects before deserialization callbacks run") { + val output = new ByteArrayOutputStream() + Serializer.write(new DeserializationTripwire, output) + DeserializationTripwire.Triggered.set(false) + + assertThrows[SecurityException] { + Serializer.read[DeserializationTripwire](new ByteArrayInputStream(output.toByteArray)) + } + assert(!DeserializationTripwire.Triggered.get()) + } + + test("Serializer closes the source stream when safe stream construction fails") { + val input = new CloseTrackingInputStream(Array[Byte](0, 1, 2, 3)) + val filter = DeserializationClassFilter(allowedClasses = Set(classOf[String].getName)) + + assertThrows[StreamCorruptedException] { + Serializer.read[String](input, filter) + } + assert(input.closed) + } + + test("Per-parameter filters reject crafted payloads before deserialization callbacks run") { + spark + new MixedParamTest("filtered").setByteArray(Array[Byte](1, 2, 3)).setStringParam("safe") + .write.overwrite().save(saveFile) + val payloadPath = new Path(new File(saveFile, "complexParams/byteArray").toString) + Serializer.writeToHDFS(spark, new DeserializationTripwire, payloadPath, overwrite = true) + DeserializationTripwire.Triggered.set(false) + + val error = intercept[SecurityException] { + MixedParamTest.load(saveFile) + } + assert(error.getCause.isInstanceOf[InvalidClassException]) + assert(!DeserializationTripwire.Triggered.get()) + } + + test("Unsafe ComplexParams require an explicit trusted legacy opt-in") { + spark + new TestableUnsafeComplexParamTest("unsafe").setUnsafePayload(new DeserializationTripwire) + .write.overwrite().save(saveFile) + val config = Serializer.LegacyObjectDeserializationConfig + val previous = spark.conf.getOption(config) + spark.conf.unset(config) + DeserializationTripwire.Triggered.set(false) + + try { + val error = intercept[SecurityException] { + TestableUnsafeComplexParamTest.load(saveFile) + } + assert(error.getMessage.contains(config)) + assert(!DeserializationTripwire.Triggered.get()) + + spark.conf.set(config, "true") + val loaded = TestableUnsafeComplexParamTest.read.session(spark).load(saveFile) + assert(Option(loaded.getUnsafePayload).nonEmpty) + assert(DeserializationTripwire.Triggered.get()) + } finally { + restoreConfig(config, previous) + } + } + + test("Unsafe default ComplexParams are reconstructed instead of serialized") { + spark + new TestableDefaultUnsafeComplexParamTest("default-unsafe").write.overwrite().save(saveFile) + + assert(!new File(saveFile, "complexParams/unsafePayload").exists()) + val loaded = TestableDefaultUnsafeComplexParamTest.load(saveFile) + assert(Option(loaded.getUnsafePayload).nonEmpty) + } + + test("ComplexParamsReader rejects a metadata class pivot before instantiation") { + spark + new MixedParamTest("class-check").setByteArray(Array[Byte](1)).setStringParam("safe") + .write.overwrite().save(saveFile) + rewriteMetadata(saveFile) { + _.replace(classOf[MixedParamTest].getName, classOf[ComplexParamTest].getName) + } + + val error = intercept[IllegalArgumentException] { + MixedParamTest.load(saveFile) + } + assert(error.getMessage.contains("Expected model class")) + } + + test("ComplexParamsReader rejects redirected complex parameter paths") { + spark + val invalidPaths = Seq("../payload", "/absolute/payload", "file:/external/payload") + invalidPaths.foreach { invalidPath => + new MixedParamTest("path-check").setByteArray(Array[Byte](1)).setStringParam("safe") + .write.overwrite().save(saveFile) + rewriteMetadata(saveFile) { + _.replace("complexParams/byteArray", invalidPath) + } + + val error = intercept[IllegalArgumentException] { + MixedParamTest.load(saveFile) + } + assert(error.getMessage.contains("must use relative path")) + } + } + + test("ComplexParamsReader rejects duplicate canonical parameter aliases") { + spark + val modelPath = new File(tmpDir.toFile, "aliased-complex-params").toString + val canonical = new SQLTransformer("canonicalTimerStage") + .setStatement("SELECT * FROM __THIS__") + val alias = new SQLTransformer("aliasedTimerStage") + .setStatement("SELECT * FROM __THIS__") + val timer = new TimerModel("aliasedTimer") + .setTransformer(alias) + timer.set(timer.stage, canonical) + timer.write.overwrite().save(modelPath) + val canonicalPath = new File(modelPath, "complexParams/stage") + val aliasPath = new File(modelPath, "complexParams/transformer") + FileUtils.deleteDirectory(aliasPath) + Files.createSymbolicLink(aliasPath.toPath, canonicalPath.toPath) + + val error = intercept[IllegalArgumentException] { + TimerModel.load(modelPath) + } + assert(error.getMessage.contains("encountered more than once")) + } + + test("ComplexParamsReader shares nesting limits with stage serializers") { + spark + val modelPath = new File(tmpDir.toFile, "deeply-nested-timers").toString + var nestedStage: PipelineStage = new LogisticRegression("timerDepthLeaf") + (0 to 100).foreach { index => // scalastyle:ignore magic.number + nestedStage = new Timer(s"t$index").setStage(nestedStage) + } + nestedStage.asInstanceOf[Timer].write.overwrite().save(modelPath) + + val error = intercept[IllegalArgumentException] { + Timer.load(modelPath) + } + assert(error.getMessage.contains("Model persistence nesting exceeds")) + } + + test("Pipeline array loading rejects nested stage paths outside the model") { + spark + val pipelinePath = new File(tmpDir.toFile, "pipeline-array").toString + val stage = new LogisticRegression("safeStage") + val nestedPipeline = new Pipeline("nestedPipeline").setStages(Array(stage)) + val serializer = new PipelineArraySerializer(spark) + serializer.write(Array[PipelineStage](nestedPipeline), new Path(pipelinePath), overwrite = true) + val nestedPath = new File(pipelinePath, "stages/0_nestedPipeline").toString + rewriteMetadata(nestedPath) { + _.replace(stage.uid, s"${stage.uid}/../../../../outside") + } + + val error = intercept[IllegalArgumentException] { + serializer.read(new Path(pipelinePath)) + } + assert(error.getMessage.contains("resolves outside its stages directory")) + } + + test("Pipeline array persistence rejects glob characters in stage UIDs") { + spark + val serializer = new PipelineArraySerializer(spark) + Seq("glob*", "glob?", "glob[0]", "glob{a,b}", "glob\\name").zipWithIndex.foreach { + case (uid, index) => + val pipelinePath = new File(tmpDir.toFile, s"glob-stage-$index").toString + val error = intercept[IllegalArgumentException] { + serializer.write( + Array[PipelineStage](new LogisticRegression(uid)), + new Path(pipelinePath), + overwrite = true + ) + } + assert(error.getMessage.contains("resolves outside its stages directory")) + } + } + + test("Pipeline array loading rejects linked stage paths outside the model") { + spark + val pipelinePath = new File(tmpDir.toFile, "linked-pipeline-array").toString + val outsidePath = new File(tmpDir.toFile, "outside-pipeline-array").toString + val stage = new LogisticRegression("linkedStage") + val outsideStage = new LogisticRegression("outsideStage") + val serializer = new PipelineArraySerializer(spark) + serializer.write(Array[PipelineStage](stage), new Path(pipelinePath), overwrite = true) + serializer.write(Array[PipelineStage](outsideStage), new Path(outsidePath), overwrite = true) + val linkedStageDir = new File(pipelinePath, s"stages/0_${stage.uid}") + val outsideStageDir = new File(outsidePath, s"stages/0_${outsideStage.uid}") + FileUtils.deleteDirectory(linkedStageDir) + Files.createSymbolicLink(linkedStageDir.toPath, outsideStageDir.toPath) + + val error = intercept[IllegalArgumentException] { + serializer.read(new Path(pipelinePath)) + } + assert(error.getMessage.contains("through a filesystem link")) + } + + test("Pipeline array loading rejects self-referential linked pipelines") { + spark + val pipelinePath = new File(tmpDir.toFile, "self-referential-pipeline").toString + val nestedPipeline = new Pipeline("nestedCycle") + .setStages(Array(new LogisticRegression("cycleLeaf"))) + val serializer = new PipelineArraySerializer(spark) + serializer.write( + Array[PipelineStage](nestedPipeline), + new Path(pipelinePath), + overwrite = true + ) + val nestedPath = new File(pipelinePath, s"stages/0_${nestedPipeline.uid}") + FileUtils.deleteDirectory(nestedPath) + Files.createSymbolicLink(nestedPath.toPath, new File(pipelinePath).toPath) + + val error = intercept[IllegalArgumentException] { + serializer.read(new Path(pipelinePath)) + } + assert(error.getMessage.contains("through a filesystem link")) + } + + test("Pipeline array loading rejects duplicate canonical stage aliases") { + spark + val pipelinePath = new File(tmpDir.toFile, "aliased-pipeline-stage").toString + val canonical = new Pipeline("canonicalNested") + .setStages(Array(new LogisticRegression("canonicalLeaf"))) + val alias = new Pipeline("aliasedNested") + .setStages(Array(new LogisticRegression("aliasedLeaf"))) + val serializer = new PipelineArraySerializer(spark) + serializer.write( + Array[PipelineStage](canonical, alias), + new Path(pipelinePath), + overwrite = true + ) + val canonicalPath = new File(pipelinePath, s"stages/0_${canonical.uid}") + val aliasPath = new File(pipelinePath, s"stages/1_${alias.uid}") + FileUtils.deleteDirectory(aliasPath) + Files.createSymbolicLink(aliasPath.toPath, canonicalPath.toPath) + + val error = intercept[IllegalArgumentException] { + serializer.read(new Path(pipelinePath)) + } + assert(error.getMessage.contains("encountered more than once")) + } + + test("Pipeline array loading limits recursive pipeline depth") { + spark + val pipelinePath = new File(tmpDir.toFile, "deeply-nested-pipeline").toString + var nestedStage: PipelineStage = new LogisticRegression("depthLeaf") + (0 to 100).foreach { index => // scalastyle:ignore magic.number + nestedStage = new Pipeline(s"nestedDepth$index").setStages(Array(nestedStage)) + } + val serializer = new PipelineArraySerializer(spark) + serializer.write( + Array(nestedStage), + new Path(pipelinePath), + overwrite = true + ) + + val error = intercept[IllegalArgumentException] { + serializer.read(new Path(pipelinePath)) + } + assert(error.getMessage.contains("Model persistence nesting exceeds")) + } + + test("Pipeline array loading accepts a linked model root") { + spark + val pipelinePath = new File(tmpDir.toFile, "real-pipeline-array") + val linkedPath = new File(tmpDir.toFile, "linked-pipeline-root") + val stage = new LogisticRegression("linkedRootStage") + val serializer = new PipelineArraySerializer(spark) + serializer.write( + Array[PipelineStage](stage), + new Path(pipelinePath.toString), + overwrite = true + ) + Files.createSymbolicLink(linkedPath.toPath, pipelinePath.toPath) + + val loaded = serializer.read(new Path(linkedPath.toString)) + assert(loaded.length === 1) + assert(loaded.head.uid === stage.uid) + } + + test("Pipeline array loading reads compressed metadata") { + spark + val pipelinePath = new File(tmpDir.toFile, "compressed-pipeline-array").toString + val stage = new LogisticRegression("compressedStage") + val serializer = new PipelineArraySerializer(spark) + serializer.write(Array[PipelineStage](stage), new Path(pipelinePath), overwrite = true) + compressMetadata(pipelinePath) + compressMetadata(new File(pipelinePath, s"stages/0_${stage.uid}").toString) + + val loaded = serializer.read(new Path(pipelinePath)) + assert(loaded.length === 1) + assert(loaded.head.uid === stage.uid) + } + + test("Pipeline array loading preserves version-aware native readers behind trust") { + spark + val pipelinePath = new File(tmpDir.toFile, "legacy-hashing-tf").toString + val stage = new HashingTF("legacyHashingTF") + val serializer = new PipelineArraySerializer(spark) + serializer.write(Array[PipelineStage](stage), new Path(pipelinePath), overwrite = true) + rewriteMetadata(new File(pipelinePath, s"stages/0_${stage.uid}").toString) { + _.replace( + s""""sparkVersion":"${spark.version}"""", + """"sparkVersion":"2.4.8"""" + ) + } + val config = Serializer.LegacyObjectDeserializationConfig + val previous = spark.conf.getOption(config) + spark.conf.unset(config) + + try { + assertThrows[SecurityException] { + serializer.read(new Path(pipelinePath)) + } + rewriteMetadata(pipelinePath) { + _.replace( + """"stagePersistenceKinds":["native"]""", + """"stagePersistenceKinds":["defaultParams"]""" + ) + } + val defaultError = intercept[IllegalArgumentException] { + serializer.read(new Path(pipelinePath)) + } + assert(defaultError.getMessage.contains("installed reader requires native")) + + rewriteMetadata(pipelinePath) { + _.replace( + """"stagePersistenceKinds":["defaultParams"]""", + """"stagePersistenceKinds":["complexParams"]""" + ) + } + val complexError = intercept[IllegalArgumentException] { + serializer.read(new Path(pipelinePath)) + } + assert(complexError.getMessage.contains("installed reader requires native")) + + rewriteMetadata(pipelinePath) { + _.replace( + """"stagePersistenceKinds":["complexParams"]""", + """"stagePersistenceKinds":["native"]""" + ) + } + spark.conf.set(config, "true") + val loaded = serializer.read(new Path(pipelinePath)).head.asInstanceOf[HashingTF] + assert(loaded.hashFuncVersion === 1) + assert(loaded.hashFuncVersion !== new HashingTF().hashFuncVersion) + } finally { + restoreConfig(config, previous) + } + } + + test("Pipeline array loading gates custom reader factories before invoking them") { + spark + val pipelinePath = new File(tmpDir.toFile, "reader-tripwire").toString + val stage = new LogisticRegression("safeStage") + val serializer = new PipelineArraySerializer(spark) + serializer.write(Array[PipelineStage](stage), new Path(pipelinePath), overwrite = true) + rewriteMetadata(pipelinePath) { + _.replace( + """"stagePersistenceKinds":["defaultParams"]""", + """"stagePersistenceKinds":["native"]""" + ) + } + rewriteMetadata(new File(pipelinePath, s"stages/0_${stage.uid}").toString) { + _.replace(classOf[LogisticRegression].getName, classOf[TestableReaderFactoryTripwireStage].getName) + } + val config = Serializer.LegacyObjectDeserializationConfig + val previous = spark.conf.getOption(config) + spark.conf.unset(config) + TestableReaderFactoryTripwireStage.Triggered.set(false) + + try { + val error = intercept[SecurityException] { + serializer.read(new Path(pipelinePath)) + } + assert(error.getMessage.contains(config)) + assert(!TestableReaderFactoryTripwireStage.Triggered.get()) + + spark.conf.set(config, "true") + assert(serializer.read(new Path(pipelinePath)).head.isInstanceOf[TestableReaderFactoryTripwireStage]) + assert(TestableReaderFactoryTripwireStage.Triggered.get()) + } finally { + restoreConfig(config, previous) + } + } + override def afterAll(): Unit = { new File(saveFile).delete() new File(saveFile2).delete() diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/core/serialize/VerifyArtifactPathResolver.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/core/serialize/VerifyArtifactPathResolver.scala new file mode 100644 index 00000000000..3f247c6b100 --- /dev/null +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/core/serialize/VerifyArtifactPathResolver.scala @@ -0,0 +1,6 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.core.serialize + +class VerifyArtifactPathResolver extends org.apache.spark.ml.ArtifactPathResolverTestBase diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/core/serialize/VerifyMetadataBudgets.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/core/serialize/VerifyMetadataBudgets.scala new file mode 100644 index 00000000000..9509498e439 --- /dev/null +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/core/serialize/VerifyMetadataBudgets.scala @@ -0,0 +1,88 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.core.serialize + +import com.microsoft.azure.synapse.ml.core.env.StreamUtilities.using +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import org.apache.commons.io.FileUtils +import org.apache.hadoop.fs.Path +import org.apache.spark.ml.feature.SQLTransformer +import org.apache.spark.ml.{PipelineArraySerializer, PipelineStage} + +import java.io.{File, FileOutputStream} +import java.nio.charset.StandardCharsets +import java.nio.file.Files +import java.util.zip.GZIPOutputStream + +class VerifyMetadataBudgets extends TestBase { + + private def rewriteMetadata(path: String)(update: String => String): Unit = { + val metadataDir = new File(path, "metadata") + val metadataJson = spark.read.text(metadataDir.toString).first().getString(0) + val updatedMetadata = update(metadataJson) + assert(updatedMetadata !== metadataJson) + FileUtils.deleteDirectory(metadataDir) + spark.createDataFrame(Seq(Tuple1(updatedMetadata))).toDF("value").write.text(metadataDir.toString) + } + + test("Pipeline array loading bounds aggregate decoded compressed metadata") { + val pipelinePath = new File(tmpDir.toFile, "decoded-metadata-pipeline").toString + val originalStage = new SQLTransformer("decodedMetadataStage") + .setStatement("SELECT * FROM __THIS__") + val serializer = new PipelineArraySerializer(spark) + serializer.write(Array[PipelineStage](originalStage), new Path(pipelinePath), overwrite = true) + + val originalStagePath = new File(pipelinePath, s"stages/0_${originalStage.uid}") + val originalMetadataDir = new File(originalStagePath, "metadata") + val originalPart = originalMetadataDir.listFiles() + .find(file => file.getName.startsWith("part-") && !file.getName.endsWith(".crc")) + .get + val originalMetadata = new String(Files.readAllBytes(originalPart.toPath), StandardCharsets.UTF_8) + val statement = "x" * 900000 // scalastyle:ignore magic.number + val decodedLimit = Math.max( + 1L << 20, // scalastyle:ignore magic.number + Math.min( + 64L << 20, // scalastyle:ignore magic.number + Runtime.getRuntime.maxMemory() / 16 // scalastyle:ignore magic.number + ) + ) + val stageCount = (decodedLimit / statement.length + 2).toInt + val stageUids = (0 until stageCount).map(index => s"decodedMetadataStage$index") + val stageUidJson = stageUids.map(uid => "\"" + uid + "\"").mkString(",") + val persistenceJson = Seq.fill(stageCount)("\"defaultParams\"").mkString(",") + rewriteMetadata(pipelinePath) { metadata => + metadata + .replace( + s""""stageUids":["${originalStage.uid}"]""", + "\"stageUids\":[" + stageUidJson + "]" + ) + .replace( + """"stagePersistenceKinds":["defaultParams"]""", + "\"stagePersistenceKinds\":[" + persistenceJson + "]" + ) + } + + FileUtils.deleteDirectory(originalStagePath) + val indexWidth = stageCount.toString.length + stageUids.zipWithIndex.foreach { case (uid, index) => + val stagePath = new File( + pipelinePath, + s"stages/${("%0" + indexWidth + "d").format(index)}_$uid" + ) + val metadataDir = new File(stagePath, "metadata") + assert(metadataDir.mkdirs()) + val metadata = originalMetadata + .replace(originalStage.uid, uid) + .replace("SELECT * FROM __THIS__", statement) + using(new GZIPOutputStream(new FileOutputStream(new File(metadataDir, "part-00000.gz")))) { + _.write(metadata.getBytes(StandardCharsets.UTF_8)) + }.get + } + + val error = intercept[IllegalArgumentException] { + serializer.read(new Path(pipelinePath)) + } + assert(error.getMessage.contains("decoded metadata")) + } +} diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/core/serialize/VerifyModelLoadEnvironment.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/core/serialize/VerifyModelLoadEnvironment.scala new file mode 100644 index 00000000000..d54bf5ded54 --- /dev/null +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/core/serialize/VerifyModelLoadEnvironment.scala @@ -0,0 +1,260 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.core.serialize + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import com.microsoft.azure.synapse.ml.stages.TimerModel +import org.apache.hadoop.fs.{FileSystem, Path, RawLocalFileSystem} +import org.apache.spark.ml.util.DefaultParamsWritable +import org.apache.spark.ml.{ + ComplexParamsReadable, + ComplexParamsWritable, + Pipeline, + PipelineArraySerializer, + PipelineModel, + PipelineStage, + Serializer +} +import org.apache.spark.sql.SparkSession + +import java.io.File +import java.net.URI +import java.util.concurrent.atomic.AtomicReference + +object TestableDefaultSessionStage + extends org.apache.spark.ml.util.DefaultParamsReadable[TestableDefaultSessionStage] { + val ActiveSession = new AtomicReference[Option[SparkSession]](None) +} + +class TestableDefaultSessionStage(override val uid: String) extends TestEstimatorBase(uid) + with DefaultParamsWritable { + TestableDefaultSessionStage.ActiveSession.set(SparkSession.getActiveSession) +} + +object TestableDirectSessionComplexStage + extends ComplexParamsReadable[TestableDirectSessionComplexStage] { + val ActiveSession = new AtomicReference[Option[SparkSession]](None) +} + +class TestableDirectSessionComplexStage(override val uid: String) extends TestEstimatorBase(uid) + with ComplexParamsWritable { + TestableDirectSessionComplexStage.ActiveSession.set(SparkSession.getActiveSession) +} + +class FileSystemOnlyLocalFileSystem extends RawLocalFileSystem { + override def getScheme: String = "modeltest" + + override def getUri: URI = URI.create("modeltest:///") +} + +class VerifyModelLoadEnvironment extends TestBase { + + private def restoreConfig( + session: SparkSession, + key: String, + previous: Option[String]): Unit = { + previous match { + case Some(value) => session.conf.set(key, value) + case None => session.conf.unset(key) + } + } + + test("Trusted loading supports FileSystem-only Hadoop providers") { + spark + val modelPath = new File(tmpDir.toFile, "filesystem-only-model").toString + val original = new MixedParamTest("filesystemOnly") + .setByteArray(Array[Byte](1, 2, 3)) + .setStringParam("value") + original.write.overwrite().save(modelPath) + val hadoopConf = spark.sparkContext.hadoopConfiguration + val implementationKey = "fs.modeltest.impl" + val cacheKey = "fs.modeltest.impl.disable.cache" + val previousImplementation = Option(hadoopConf.get(implementationKey)) + val previousCache = Option(hadoopConf.get(cacheKey)) + hadoopConf.setClass( + implementationKey, + classOf[FileSystemOnlyLocalFileSystem], + classOf[FileSystem] + ) + hadoopConf.setBoolean(cacheKey, true) + val providerPath = s"modeltest://${new File(modelPath).toURI.getPath}" + val trustConfig = Serializer.LegacyObjectDeserializationConfig + val previousTrust = spark.conf.getOption(trustConfig) + spark.conf.unset(trustConfig) + + try { + val securityError = intercept[SecurityException] { + MixedParamTest.load(providerPath) + } + assert(securityError.getMessage.contains("cannot resolve links safely")) + + spark.conf.set(trustConfig, "true") + val loaded = MixedParamTest.read.session(spark).load(providerPath) + assert(loaded.getByteArray === original.getByteArray) + assert(loaded.getStringParam === original.getStringParam) + } finally { + restoreConfig(spark, trustConfig, previousTrust) + previousImplementation.fold(hadoopConf.unset(implementationKey)) { + hadoopConf.set(implementationKey, _) + } + previousCache.fold(hadoopConf.unset(cacheKey)) { + hadoopConf.set(cacheKey, _) + } + } + } + + test("Default stage construction uses the supplied Spark session") { + spark + val pipelinePath = new File(tmpDir.toFile, "default-session-stage").toString + val stage = new TestableDefaultSessionStage("defaultSessionStage") + new PipelineArraySerializer(spark) + .write(Array[PipelineStage](stage), new Path(pipelinePath), overwrite = true) + val isolatedSession = spark.newSession() + val previousSession = SparkSession.getActiveSession + SparkSession.setActiveSession(spark) + TestableDefaultSessionStage.ActiveSession.set(None) + + try { + val loaded = new PipelineArraySerializer(isolatedSession).read(new Path(pipelinePath)) + assert(loaded.head.uid === stage.uid) + assert(TestableDefaultSessionStage.ActiveSession.get().contains(isolatedSession)) + } finally { + previousSession.fold(SparkSession.clearActiveSession())(SparkSession.setActiveSession) + } + } + + test("Direct ComplexParams construction uses the supplied Spark session") { + spark + val modelPath = new File(tmpDir.toFile, "direct-session-complex").toString + new TestableDirectSessionComplexStage("directSessionComplex").write.overwrite().save(modelPath) + val isolatedSession = spark.newSession() + val previousSession = SparkSession.getActiveSession + SparkSession.setActiveSession(spark) + TestableDirectSessionComplexStage.ActiveSession.set(None) + + try { + val loaded = TestableDirectSessionComplexStage.read.session(isolatedSession).load(modelPath) + assert(loaded.uid === "directSessionComplex") + assert(TestableDirectSessionComplexStage.ActiveSession.get().contains(isolatedSession)) + } finally { + previousSession.fold(SparkSession.clearActiveSession())(SparkSession.setActiveSession) + } + } + + test("Direct ComplexParams load assigns the default Spark session") { + spark + val modelPath = new File(tmpDir.toFile, "direct-load-session-complex").toString + new TestableDirectSessionComplexStage("directLoadSessionComplex").write.overwrite().save(modelPath) + val previousSession = SparkSession.getActiveSession + SparkSession.setActiveSession(spark) + TestableDirectSessionComplexStage.ActiveSession.set(None) + + try { + val loaded = TestableDirectSessionComplexStage.load(modelPath) + assert(loaded.uid === "directLoadSessionComplex") + assert(TestableDirectSessionComplexStage.ActiveSession.get().contains(spark)) + } finally { + previousSession.fold(SparkSession.clearActiveSession())(SparkSession.setActiveSession) + } + } + + test("Native Pipeline compatibility requires an explicit trusted scope") { + spark + val modelPath = new File(tmpDir.toFile, "native-pipeline-write").toString + val pipeline = new Pipeline("nativePipelineWrite") + .setStages(Array(new TestableDirectSessionComplexStage("nativePipelineComplex"))) + + val error = intercept[SecurityException] { + pipeline.write.overwrite().save(modelPath) + } + assert(error.getMessage.contains("model-wide metadata budget")) + + val trustedPath = new File(tmpDir.toFile, "trusted-native-pipeline-write").toString + val config = Serializer.LegacyObjectDeserializationConfig + val previousTrust = spark.conf.getOption(config) + try { + spark.conf.unset(config) + val disabledScopeError = intercept[SecurityException] { + Serializer.beginTrustedArtifactLoad(spark) + } + assert(disabledScopeError.getMessage.contains(s"Set $config=true")) + + spark.conf.set(config, "true") + val scope = Serializer.beginTrustedArtifactLoad(spark) + try { + pipeline.write.overwrite().save(trustedPath) + val loaded = Pipeline.load(trustedPath) + assert(loaded.getStages.map(_.uid) === Array("nativePipelineComplex")) + + val closeError = new AtomicReference[Throwable]() + val closeThread = new Thread(new Runnable { + override def run(): Unit = { + try { + scope.close() + } catch { + case error: Throwable => closeError.set(error) + } + } + }) + closeThread.start() + closeThread.join() + assert(closeError.get().isInstanceOf[IllegalStateException]) + } finally { + scope.close() + } + scope.close() + assert(new File(trustedPath, "metadata").isDirectory) + + val unscopedError = intercept[SecurityException] { + Pipeline.load(trustedPath) + } + assert(unscopedError.getMessage.contains("Unscoped native Spark Pipeline loading")) + } finally { + restoreConfig(spark, config, previousTrust) + } + } + + test("Native PipelineModel loading cannot inherit unsafe trust from an ambient session") { + spark + val modelPath = new File(tmpDir.toFile, "native-pipeline-session").toString + val unsafeStage = new TestableUnsafeComplexParamTest("nativePipelineUnsafe") + .setUnsafePayload(new DeserializationTripwire) + val stage = new TimerModel("nativePipelineTimer").setTransformer(unsafeStage) + val isolatedSession = spark.newSession() + val config = Serializer.LegacyObjectDeserializationConfig + val previousDefault = spark.conf.getOption(config) + val previousIsolated = isolatedSession.conf.getOption(config) + val previousActive = SparkSession.getActiveSession + isolatedSession.conf.set(config, "true") + Serializer.withTrustedArtifactLoad(isolatedSession) { + new Pipeline("nativePipeline").setStages(Array(stage)) + .fit(isolatedSession.range(1).toDF()) + .write.overwrite().save(modelPath) + } + spark.conf.set(config, "true") + isolatedSession.conf.unset(config) + SparkSession.setActiveSession(spark) + DeserializationTripwire.Triggered.set(false) + + try { + val error = intercept[SecurityException] { + PipelineModel.read.session(isolatedSession).load(modelPath) + } + assert(error.getMessage.contains("withTrustedArtifactLoad")) + assert(!DeserializationTripwire.Triggered.get()) + + spark.conf.unset(config) + isolatedSession.conf.set(config, "true") + val loaded = Serializer.withTrustedArtifactLoad(isolatedSession) { + PipelineModel.read.session(isolatedSession).load(modelPath) + } + assert(loaded.stages.length === 1) + assert(DeserializationTripwire.Triggered.get()) + } finally { + restoreConfig(spark, config, previousDefault) + restoreConfig(isolatedSession, config, previousIsolated) + previousActive.fold(SparkSession.clearActiveSession())(SparkSession.setActiveSession) + } + } +} diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/core/test/base/TestBase.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/core/test/base/TestBase.scala index a90f1ea81eb..1ce01db4ff8 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/core/test/base/TestBase.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/core/test/base/TestBase.scala @@ -198,6 +198,20 @@ abstract class TestBase extends AnyFunSuite with BeforeAndAfterEachTestData with // Utilities + protected def withTrustedLegacyModelPersistence[T](action: => T): T = { + val config = Serializer.LegacyObjectDeserializationConfig + val previous = spark.conf.getOption(config) + spark.conf.set(config, "true") + try { + Serializer.withTrustedArtifactLoad(spark)(action) + } finally { + previous match { + case Some(value) => spark.conf.set(config, value) + case None => spark.conf.unset(config) + } + } + } + def tryWithRetries[T](times: Array[Int] = Array(0, 100, 500, 1000, 3000, 5000))(block: () => T): T = { for ((t, i) <- times.zipWithIndex) { try { diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/core/test/fuzzing/Fuzzing.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/core/test/fuzzing/Fuzzing.scala index 86573324d7b..5579618d494 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/core/test/fuzzing/Fuzzing.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/core/test/fuzzing/Fuzzing.scala @@ -75,7 +75,9 @@ trait PyTestFuzzing[S <: PipelineStage] extends TestBase with DataFrameEquality def savePyModel(conf: CodegenConfig, model: S, name: String): Unit = { model match { case writable: MLWritable => - writable.write.overwrite().save(new File(pyTestDataDir(conf), s"$name.model").toString) + def save(): Unit = writable.write.overwrite() + .save(new File(pyTestDataDir(conf), s"$name.model").toString) + withTrustedLegacyModelPersistence(save()) case _ => throw new IllegalArgumentException(s"${model.getClass.getName} is not writable") } @@ -187,6 +189,43 @@ trait PyTestFuzzing[S <: PipelineStage] extends TestBase with DataFrameEquality val stageName = getClassName(stage) val importPath = stage.getClass.getName.split(".".toCharArray).dropRight(1) val importPathString = importPath.mkString(".").replaceAllLiterally("com.microsoft.azure.synapse.ml", "synapse.ml") + // TestGen creates the expected models locally; trust is scoped to each generated test. + val trustedFixtureLifecycle = + s""" def setUp(self): + | self._legacy_deserialization_config = "${Serializer.LegacyObjectDeserializationConfig}" + | self._previous_legacy_deserialization = spark.conf.get( + | self._legacy_deserialization_config, None + | ) + | spark.conf.set(self._legacy_deserialization_config, "true") + | try: + | self._trusted_artifact_scope = ( + | sc._jvm.org.apache.spark.ml.Serializer.beginTrustedArtifactLoad( + | spark._jsparkSession + | ) + | ) + | except Exception: + | if self._previous_legacy_deserialization is None: + | spark.conf.unset(self._legacy_deserialization_config) + | else: + | spark.conf.set( + | self._legacy_deserialization_config, + | self._previous_legacy_deserialization, + | ) + | raise + | + | def tearDown(self): + | try: + | self._trusted_artifact_scope.close() + | finally: + | if self._previous_legacy_deserialization is None: + | spark.conf.unset(self._legacy_deserialization_config) + | else: + | spark.conf.set( + | self._legacy_deserialization_config, + | self._previous_legacy_deserialization, + | ) + | + |""".stripMargin val testClass = s"""import unittest |from pyspark.sql import SQLContext @@ -204,6 +243,7 @@ trait PyTestFuzzing[S <: PipelineStage] extends TestBase with DataFrameEquality | | |class $testClassName(unittest.TestCase): + |$trustedFixtureLifecycle | def assert_correspondence(self, model, name, num): | model.write().overwrite().save(join(test_data_dir, name)) | sc._jvm.com.microsoft.azure.synapse.ml.core.utils.ModelEquality.assertEqual( @@ -230,6 +270,34 @@ trait PyTestFuzzing[S <: PipelineStage] extends TestBase with DataFrameEquality } +private[synapse] object RTestFuzzing { + + private def pipelineStageLoadLine(param: PipelineStageWrappable[_], modelNum: Int): String = { + val name = param.name + s""" + |${name}Model <- invoke( + | invoke_new(sc, "org.apache.spark.ml.PipelineSerializer", spark_session(sc)), + | "read", + | invoke_new( + | sc, + | "org.apache.hadoop.fs.Path", + | file.path(test_data_dir, "model-$modelNum.model", "complexParams", "$name"))) + |${name}Model <- sparklyr::ml_call_constructor(${name}Model) + """.stripMargin + } + + def loadLine(param: Param[_], modelNum: Int): Option[String] = { + param match { + case ep: PipelineStageWrappable[_] => + Some(pipelineStageLoadLine(ep, modelNum)) + case ep: ExternalRWrappableParam[_] => + Some(ep.rLoadLine(modelNum)) + case _ => None + } + } + +} + trait RTestFuzzing[S <: PipelineStage] extends TestBase with DataFrameEquality with TestFuzzingUtil { def rTestObjects(): Seq[TestObject[S]] @@ -244,7 +312,9 @@ trait RTestFuzzing[S <: PipelineStage] extends TestBase with DataFrameEquality w def saveRModel(conf: CodegenConfig, model: S, name: String): Unit = { model match { case writable: MLWritable => - writable.write.overwrite().save(new File(rTestDataDir(conf), s"$name.model").toString) + def save(): Unit = writable.write.overwrite() + .save(new File(rTestDataDir(conf), s"$name.model").toString) + withTrustedLegacyModelPersistence(save()) case _ => throw new IllegalArgumentException(s"${model.getClass.getName} is not writable") } @@ -272,11 +342,7 @@ trait RTestFuzzing[S <: PipelineStage] extends TestBase with DataFrameEquality w def instantiateModel(paramMap: Seq[ParamPair[_]]): String = { val externalLoadingLines = paramMap.flatMap { pp => - pp.param match { - case ep: ExternalRWrappableParam[_] => - Some(ep.rLoadLine(num)) - case _ => None - } + RTestFuzzing.loadLine(pp.param, num) }.mkString("\n") val modelArg = stage match { @@ -348,9 +414,34 @@ trait RTestFuzzing[S <: PipelineStage] extends TestBase with DataFrameEquality w } case _ => "" } + // TestGen creates the expected models locally; trust is scoped to each generated test. + val trustedFixtureSetup = + s""" + |legacy_deserialization_config <- "${Serializer.LegacyObjectDeserializationConfig}" + |runtime_config <- invoke(spark_session(sc), "conf") + |legacy_deserialization_unset <- "__synapseml_unset__" + |previous_legacy_deserialization <- invoke( + | runtime_config, + | "get", + | legacy_deserialization_config, + | legacy_deserialization_unset) + |invoke(runtime_config, "set", legacy_deserialization_config, "true") + |on.exit({ + | if (identical(previous_legacy_deserialization, legacy_deserialization_unset)) { + | invoke(runtime_config, "unset", legacy_deserialization_config) + | } else { + | invoke( + | runtime_config, + | "set", + | legacy_deserialization_config, + | previous_legacy_deserialization) + | } + |}, add = TRUE) + |""".stripMargin s""" |test_that("${stageName}_constructor_$num", { + | ${indent(trustedFixtureSetup, 1)} | ${indent(rTestInstantiateModel(stage, num), 1)} | | ${indent(s"""assert_correspondence_$stageName(model, "r-constructor-model-$num.model", $num)""", 1)} @@ -453,7 +544,8 @@ trait ExperimentFuzzing[S <: PipelineStage] extends TestBase with DataFrameEqual } -trait SerializationFuzzing[S <: PipelineStage with MLWritable] extends TestBase with DataFrameEquality { +trait SerializationFuzzing[S <: PipelineStage with MLWritable] + extends TestBase with DataFrameEquality with TestFuzzingUtil { def serializationTestObjects(): Seq[TestObject[S]] def reader: MLReadable[_] @@ -486,9 +578,11 @@ trait SerializationFuzzing[S <: PipelineStage with MLWritable] extends TestBase reader: MLReadable[_], fitDF: DataFrame, transDF: DataFrame): Unit = { try { - stage.write.overwrite().save(path) - assert(new File(path).exists()) - val loadedStage = reader.load(path) + val loadedStage = withTrustedLegacyModelPersistence { + stage.write.overwrite().save(path) + assert(new File(path).exists()) + reader.read.session(spark).load(path) + } (stage, loadedStage) match { case (e1: Estimator[_], e2: Estimator[_]) => val df1 = e1.fit(fitDF).transform(transDF) @@ -531,7 +625,13 @@ trait SerializationFuzzing[S <: PipelineStage with MLWritable] extends TestBase testSerializationHelper(savePath + "/pipe", pipe, Pipeline, req.fitDF, req.transDF) } val fitPipe = pipe.fit(req.fitDF) - testSerializationHelper(savePath + "/fitPipe", fitPipe, PipelineModel, req.transDF, req.transDF) + testSerializationHelper( + savePath + "/fitPipe", + fitPipe, + PipelineModel, + req.transDF, + req.transDF + ) } } finally { if (new File(savePath).exists) FileUtils.forceDelete(new File(savePath)) @@ -540,15 +640,19 @@ trait SerializationFuzzing[S <: PipelineStage with MLWritable] extends TestBase val retrySerializationFuzzing = false - test("Serialization Fuzzing") { - if (!ignoreSerializationFuzzing) { - if (retrySerializationFuzzing) { - tryWithRetries() { () => - testSerialization() - } - } else { + private def runSerializationFuzzing(): Unit = { + if (retrySerializationFuzzing) { + tryWithRetries() { () => testSerialization() } + } else { + testSerialization() + } + } + + test("Serialization Fuzzing") { + if (!ignoreSerializationFuzzing) { + runSerializationFuzzing() } } diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/core/utils/VerifySafeObjectInputStream.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/core/utils/VerifySafeObjectInputStream.scala new file mode 100644 index 00000000000..fa8e2eece04 --- /dev/null +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/core/utils/VerifySafeObjectInputStream.scala @@ -0,0 +1,154 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.core.utils + +import com.microsoft.azure.synapse.ml.core.env.StreamUtilities.using +import com.microsoft.azure.synapse.ml.core.test.base.TestBase + +import java.io.{ + ByteArrayInputStream, + ByteArrayOutputStream, + InvalidClassException, + InvalidObjectException, + ObjectOutputStream, + StreamCorruptedException +} + +class VerifySafeObjectInputStream extends TestBase { + + private def serialize(value: AnyRef): Array[Byte] = { + val output = new ByteArrayOutputStream() + using(new ObjectOutputStream(output)) { + _.writeObject(value) + }.get + output.toByteArray + } + + private def limits( + maxStreamBytes: Long = 1024, + maxArrayBytes: Long = 4096, + maxResolvedObjects: Long = 100, + maxStringBytes: Long = 1024): DeserializationResourceLimits = { + DeserializationResourceLimits( + maxDepth = 100, + maxReferences = 100, + maxStreamBytes = maxStreamBytes, + maxArrayBytes = maxArrayBytes, + maxResolvedObjects = maxResolvedObjects, + maxStringBytes = maxStringBytes + ) + } + + private def filterInfo( + depthValue: Long = 0, + referencesValue: Long = 0, + streamBytesValue: Long = 0, + arrayLengthValue: Long = -1, + serialClassValue: Class[_] = classOf[String]): DeserializationFilterInfo = { + DeserializationFilterInfo( + serialClass = serialClassValue, + arrayLength = arrayLengthValue, + depth = depthValue, + references = referencesValue, + streamBytes = streamBytesValue + ) + } + + test("Safe object input filter rejects excessive graph resources") { + val rejected = DeserializationFilterStatus.Rejected + val resourceFilter = SafeObjectInputStream.newResourceFilter() + assert(resourceFilter.checkInput( + filterInfo(depthValue = Long.MaxValue) + ) === rejected) + assert(resourceFilter.checkInput( + filterInfo(referencesValue = Long.MaxValue) + ) === rejected) + assert(resourceFilter.checkInput( + filterInfo(streamBytesValue = Long.MaxValue) + ) === rejected) + assert(resourceFilter.checkInput( + filterInfo(arrayLengthValue = Int.MaxValue, serialClassValue = classOf[Array[Byte]]) + ) === rejected) + } + + test("Safe object input stream bounds serialized String input bytes") { + val serialized = serialize("x" * 1024) + val result = using(new SafeObjectInputStream( + new ByteArrayInputStream(serialized), + DeserializationClassFilter(allowedPrefixes = Set("java.lang.")), + limits(maxStreamBytes = 128) + )) { + _.readObject() + } + assert(result.isFailure) + assert(result.failed.get.isInstanceOf[StreamCorruptedException]) + } + + test("Safe object input stream counts Strings omitted by JEP 290 callbacks") { + val strings = Array.tabulate(20)(index => new String(s"value-$index")) + val result = using(new SafeObjectInputStream( + new ByteArrayInputStream(serialize(strings)), + DeserializationClassFilter(allowedPrefixes = Set("java.lang.")), + limits(maxStreamBytes = 4096, maxResolvedObjects = 10, maxStringBytes = 4096) + )) { + _.readObject() + } + assert(result.isFailure) + assert(result.failed.get.isInstanceOf[InvalidObjectException]) + } + + test("Safe object input stream bounds aggregate declared array allocations") { + val arrays = Array(Array.fill[Byte](80)(1), Array.fill[Byte](80)(2)) + val result = using(new SafeObjectInputStream( + new ByteArrayInputStream(serialize(arrays)), + DeserializationClassFilter(), + limits(maxStreamBytes = 4096, maxArrayBytes = 100) + )) { + _.readObject() + } + assert(result.isFailure) + assert(result.failed.get.isInstanceOf[InvalidClassException]) + } + + test("Safe object input filter rejects declared array byte multiplication overflow") { + val resourceFilter = new DeserializationResourceFilter( + limits(maxArrayBytes = Long.MaxValue) + ) + assert(resourceFilter.checkInput( + filterInfo(arrayLengthValue = Long.MaxValue, serialClassValue = classOf[Array[AnyRef]]) + ) === DeserializationFilterStatus.Rejected) + } + + test("Safe object input filter rejects aggregate declared array byte addition overflow") { + val resourceFilter = new DeserializationResourceFilter( + limits(maxArrayBytes = Long.MaxValue) + ) + assert(resourceFilter.checkInput( + filterInfo(arrayLengthValue = Long.MaxValue, serialClassValue = classOf[Array[Byte]]) + ) === DeserializationFilterStatus.Undecided) + assert(resourceFilter.checkInput( + filterInfo(arrayLengthValue = 1, serialClassValue = classOf[Array[Byte]]) + ) === DeserializationFilterStatus.Rejected) + } + + test("Safe object input stream applies class policy to special-case Strings") { + val result = using(new SafeObjectInputStream( + new ByteArrayInputStream(serialize("payload")), + DeserializationClassFilter(), + limits() + )) { + _.readObject() + } + assert(result.isFailure) + assert(result.failed.get.isInstanceOf[InvalidClassException]) + } + + test("Safe object input stream composes inherited JVM filter rejection") { + val composed = Jep290ObjectInputFilter.composeStatuses( + DeserializationFilterStatus.Rejected, + DeserializationFilterStatus.Undecided + ) + assert(composed === DeserializationFilterStatus.Rejected) + } +} diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/featurize/VerifyFeaturize.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/featurize/VerifyFeaturize.scala index 8da57e42b11..3f395809f82 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/featurize/VerifyFeaturize.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/featurize/VerifyFeaturize.scala @@ -275,8 +275,10 @@ class VerifyFeaturize extends TestBase with EstimatorFuzzing[Featurize] { AttributeGroup.fromStructField(resultWithConstantText.schema(featuresColumn)).numAttributes) val modelDir = new File(tmpDir.toFile, "issue1667-featurize-model") - modelWithConstantText.write.overwrite().save(modelDir.toString) - val loadedModel = PipelineModel.load(modelDir.toString) + val loadedModel = withTrustedLegacyModelPersistence { + modelWithConstantText.write.overwrite().save(modelDir.toString) + PipelineModel.load(modelDir.toString) + } val loadedResult = loadedModel.transform(scoringDataset).select(featuresColumn) assert(verifyResult(resultWithConstantText, loadedResult)) } @@ -568,8 +570,10 @@ class VerifyFeaturize extends TestBase with EstimatorFuzzing[Featurize] { .setIndices(Array(0)))) val model = pipeline.fit(missingValueDataset) val path = new File(tmpDir.toFile, "featurize-pipeline-model-keep").toString - model.write.overwrite().save(path) - val loadedModel = PipelineModel.load(path) + val loadedModel = withTrustedLegacyModelPersistence { + model.write.overwrite().save(path) + PipelineModel.load(path) + } val result = loadedModel.transform(missingValueDataset) assert(result.count() == 4) val byLabel = firstSlotByLabel(result) diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/nn/VerifySchemas.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/nn/VerifySchemas.scala index 550b0e701be..0ed5fd4d290 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/nn/VerifySchemas.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/nn/VerifySchemas.scala @@ -7,9 +7,37 @@ import breeze.linalg.DenseVector import com.microsoft.azure.synapse.ml.core.test.base.TestBase import com.microsoft.azure.synapse.ml.core.env.StreamUtilities.using -import com.microsoft.azure.synapse.ml.core.utils.SafeObjectInputStream +import com.microsoft.azure.synapse.ml.core.utils.{ + DeserializationClassFilter, + SafeObjectInputStream +} +import org.apache.spark.ml.Serializer +import org.apache.spark.ml.linalg.Vectors +import org.apache.spark.sql.Row -import java.io.{ByteArrayInputStream, ByteArrayOutputStream, ObjectOutputStream} +import java.io.{ + ByteArrayInputStream, + ByteArrayOutputStream, + File, + FileOutputStream, + ObjectInputStream, + ObjectOutputStream +} +import java.nio.ByteBuffer +import java.sql.Date +import java.util.concurrent.atomic.AtomicBoolean + +private object ConditionalBallTreeLoadTripwire { + val Triggered = new AtomicBoolean(false) +} + +@SerialVersionUID(1L) +private class ConditionalBallTreeLoadTripwire extends Serializable { + private def readObject(input: ObjectInputStream): Unit = { + ConditionalBallTreeLoadTripwire.Triggered.set(true) + input.defaultReadObject() + } +} class VerifySchemas extends TestBase { @@ -60,7 +88,8 @@ class VerifySchemas extends TestBase { } val bais = new ByteArrayInputStream(baos.toByteArray) - val deserialized = using(new SafeObjectInputStream(bais, SafeObjectInputStream.DefaultNNAllowedPrefixes)) { ois => + val filter = DeserializationClassFilter(allowedClasses = Set(classOf[BestMatch].getName)) + val deserialized = using(new SafeObjectInputStream(bais, filter)) { ois => ois.readObject().asInstanceOf[BestMatch] }.get @@ -99,6 +128,48 @@ class VerifySchemas extends TestBase { assert(deserialized.sameElements(data)) } + test("SafeObjectInputStream rejects oversized declared arrays before allocation") { + val data = Array[Byte](1) + val baos = new ByteArrayOutputStream() + using(new ObjectOutputStream(baos)) { oos => + oos.writeObject(data) + } + val serialized = baos.toByteArray + val arrayLengthOffset = serialized.length - data.length - Integer.BYTES + ByteBuffer.wrap(serialized, arrayLengthOffset, Integer.BYTES).putInt(Int.MaxValue) + + val result = using(new SafeObjectInputStream( + new ByteArrayInputStream(serialized), + Set.empty[String] + )) { ois => + ois.readObject() + } + assert(result.isFailure) + assert(result.failed.get.isInstanceOf[java.io.InvalidClassException]) + } + + test("SafeObjectInputStream rejects excessively deep object graphs") { + var nested: AnyRef = "leaf" + (1 to 110).foreach { _ => // scalastyle:ignore magic.number + val parent = new java.util.ArrayList[AnyRef]() + parent.add(nested) + nested = parent + } + val baos = new ByteArrayOutputStream() + using(new ObjectOutputStream(baos)) { oos => + oos.writeObject(nested) + } + + val result = using(new SafeObjectInputStream( + new ByteArrayInputStream(baos.toByteArray), + SafeObjectInputStream.CommonDataAllowedPrefixes + )) { ois => + ois.readObject() + } + assert(result.isFailure) + assert(result.failed.get.isInstanceOf[java.io.InvalidClassException]) + } + test("SafeObjectInputStream allows object arrays with permitted component type") { val data = Array("hello", "world") val baos = new ByteArrayOutputStream() @@ -148,4 +219,102 @@ class VerifySchemas extends TestBase { assert(result.failed.get.isInstanceOf[java.io.InvalidClassException]) } + test("SafeObjectInputStream rejects SerializedLambda class-resolution callbacks") { + val function: Int => Int = value => value + 1 + val baos = new ByteArrayOutputStream() + using(new ObjectOutputStream(baos)) { oos => + oos.writeObject(function) + } + + val bais = new ByteArrayInputStream(baos.toByteArray) + val result = using(new SafeObjectInputStream(bais, Set("java.lang.", "scala."))) { ois => + ois.readObject() + } + assert(result.isFailure) + assert(result.failed.get.isInstanceOf[java.io.InvalidClassException]) + assert(result.failed.get.getMessage.contains("SerializedLambda")) + } + + test("Default NN filter fails closed for legacy BallTree object graphs") { + val keys = IndexedSeq( + DenseVector(1.0, 0.0), + DenseVector(0.0, 1.0), + DenseVector(1.0, 1.0) + ) + val values = IndexedSeq[Any]( + Date.valueOf("2026-08-24"), + Vectors.dense(1.0, 2.0), + Row("value", 1) + ) + val tree = BallTree(keys, values) + val baos = new ByteArrayOutputStream() + using(new ObjectOutputStream(baos)) { oos => + oos.writeObject(tree) + } + + val result = using(new SafeObjectInputStream( + new ByteArrayInputStream(baos.toByteArray), + SafeObjectInputStream.DefaultNNFilter + )) { ois => + ois.readObject() + } + + assert(result.isFailure) + assert(result.failed.get.isInstanceOf[java.io.InvalidClassException]) + } + + test("ConditionalBallTree default load rejects before deserialization callbacks") { + val path = new File(tmpDir.toFile, "conditional-ball-tree-tripwire.bin") + using(new ObjectOutputStream(new FileOutputStream(path))) { output => + output.writeObject(new ConditionalBallTreeLoadTripwire) + }.get + ConditionalBallTreeLoadTripwire.Triggered.set(false) + + val error = intercept[SecurityException] { + ConditionalBallTree.load[String, Row](path.toString) + } + assert(error.getMessage.contains("loadUnsafe")) + assert(!ConditionalBallTreeLoadTripwire.Triggered.get()) + } + + test("ConditionalBallTree trusted load supports Spark Row values") { + val tree = ConditionalBallTree( + IndexedSeq(DenseVector(1.0, 0.0), DenseVector(0.0, 1.0)), + IndexedSeq(Row("first", 1), Row("second", 2)), + IndexedSeq("a", "b"), + leafSize = 1 + ) + val path = new File(tmpDir.toFile, "conditional-ball-tree.bin") + tree.save(path.toString) + + val loaded = ConditionalBallTree.loadUnsafe[String, Row](path.toString) + assert(loaded.values === tree.values) + } + + test("KNN BallTree persistence requires an explicitly session-scoped trusted load") { + spark + val tree = BallTree( + IndexedSeq(DenseVector(1.0, 0.0), DenseVector(0.0, 1.0)), + IndexedSeq(1, 2), + leafSize = 1 + ) + val path = new File(tmpDir.toFile, "trusted-knn-model").toString + new KNNModel("trustedKnn").setBallTree(tree).write.overwrite().save(path) + val config = Serializer.LegacyObjectDeserializationConfig + val previous = spark.conf.getOption(config) + spark.conf.unset(config) + + try { + assertThrows[SecurityException] { + KNNModel.read.session(spark).load(path) + } + + spark.conf.set(config, "true") + val loaded = KNNModel.read.session(spark).load(path) + assert(loaded.getBallTree.values === tree.values) + } finally { + previous.fold(spark.conf.unset(config))(spark.conf.set(config, _)) + } + } + } diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyArrayParamMapParam.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyArrayParamMapParam.scala new file mode 100644 index 00000000000..846e2e0d14f --- /dev/null +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyArrayParamMapParam.scala @@ -0,0 +1,46 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.param + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import org.apache.hadoop.fs.Path +import org.apache.spark.ml.Serializer +import org.apache.spark.ml.param.{IntParam, ParamMap, Params} + +import java.io.File + +class VerifyArrayParamMapParam extends TestBase { + + private class TestParamsHolder extends Params { + override val uid: String = "param-map-holder" + val number = new IntParam(this, "number", "A numeric test param") + val maps = new ArrayParamMapParam(this, "maps", "Array of parameter maps") + + override def copy(extra: ParamMap): Params = this + } + + test("ArrayParamMapParam persistence requires explicit trust for serialized parameter maps") { + val holder = new TestParamsHolder + val path = new Path(new File(tmpDir.toFile, "param-maps").toString) + val expected = Array(ParamMap(holder.number -> 7)) + holder.maps.save(expected, spark, path, overwrite = true) + val config = Serializer.LegacyObjectDeserializationConfig + val previous = spark.conf.getOption(config) + spark.conf.unset(config) + + try { + assertThrows[SecurityException] { + holder.maps.load(spark, path) + } + spark.conf.set(config, "true") + val loaded = holder.maps.load(spark, path) + assert(loaded.head.get(holder.number).contains(7)) + } finally { + previous match { + case Some(value) => spark.conf.set(config, value) + case None => spark.conf.unset(config) + } + } + } +} diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyDataFrameParam.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyDataFrameParam.scala index 485be61190f..1696db55a5d 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyDataFrameParam.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyDataFrameParam.scala @@ -4,11 +4,16 @@ package com.microsoft.azure.synapse.ml.param import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import org.apache.hadoop.fs.Path +import org.apache.spark.ml.Serializer import org.apache.spark.ml.linalg.{DenseVector, Vectors} import org.apache.spark.ml.param.{ParamMap, Params} import org.apache.spark.sql.{DataFrame, Row} import org.apache.spark.sql.types._ +import java.io.File +import java.nio.file.Files + class VerifyDataFrameParam extends TestBase { import spark.implicits._ @@ -156,6 +161,46 @@ class VerifyDataFrameParam extends TestBase { holder.set(holder.nonEmptyDf, df) } + test("DataFrameParam rejects linked Parquet files outside its artifact") { + val holder = new TestParamsHolder + val parameterPath = new File(tmpDir.toFile, "linked-dataframe-param[1]") + val externalPath = new File(tmpDir.toFile, "external-dataframe") + holder.dfParam.save( + Seq(1, 2).toDF("value"), + spark, + new Path(parameterPath.toString), + overwrite = true + ) + Seq(99).toDF("value").write.parquet(externalPath.toString) + val config = Serializer.LegacyObjectDeserializationConfig + val previous = spark.conf.getOption(config) + spark.conf.unset(config) + + try { + assertThrows[SecurityException] { + holder.dfParam.load(spark, new Path(parameterPath.toString)) + } + + spark.conf.set(config, "true") + assert(holder.dfParam.load(spark, new Path(parameterPath.toString)).count() === 2) + val parameterPart = parameterPath.listFiles() + .find(_.getName.endsWith(".parquet")) + .get + val externalPart = externalPath.listFiles() + .find(_.getName.endsWith(".parquet")) + .get + Files.delete(parameterPart.toPath) + Files.createSymbolicLink(parameterPart.toPath, externalPart.toPath) + + val error = intercept[IllegalArgumentException] { + holder.dfParam.load(spark, new Path(parameterPath.toString)).count() + } + assert(error.getMessage.contains("resolves outside")) + } finally { + previous.fold(spark.conf.unset(config))(spark.conf.set(config, _)) + } + } + test("DataFrameParam sortInDataframeEquality is true") { val holder = new TestParamsHolder assert(holder.dfParam.sortInDataframeEquality) diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyDataTypeParam.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyDataTypeParam.scala index afeb1a4d793..94dc7ccb431 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyDataTypeParam.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyDataTypeParam.scala @@ -4,9 +4,34 @@ package com.microsoft.azure.synapse.ml.param import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import org.apache.hadoop.fs.Path +import org.apache.spark.ml.Serializer +import org.apache.spark.ml.linalg.SQLDataTypes import org.apache.spark.ml.param.{ParamMap, Params} import org.apache.spark.sql.types._ +import java.io.{File, ObjectInputStream} +import java.util.concurrent.atomic.AtomicBoolean + +private object DataTypeLoadTripwire { + val Triggered = new AtomicBoolean(false) +} + +@SerialVersionUID(1L) +private class DataTypeLoadTripwire extends Serializable { + private def readObject(input: ObjectInputStream): Unit = { + DataTypeLoadTripwire.Triggered.set(true) + input.defaultReadObject() + } +} + +class TestStringUDT extends UserDefinedType[String] { + override def sqlType: DataType = StringType + override def serialize(obj: String): Any = obj + override def deserialize(datum: Any): String = datum.asInstanceOf[String] + override def userClass: Class[String] = classOf[String] +} + class VerifyDataTypeParam extends TestBase { private class TestParamsHolder extends Params { @@ -123,4 +148,109 @@ class VerifyDataTypeParam extends TestBase { holder.clear(holder.dataTypeParam) assert(!holder.isSet(holder.dataTypeParam)) } + + test("DataTypeParam safely persists Spark ML vector types") { + val holder = new TestParamsHolder + val path = new Path(new File(tmpDir.toFile, "vector-data-type").toString) + + holder.dataTypeParam.save(SQLDataTypes.VectorType, spark, path, overwrite = true) + + assert(holder.dataTypeParam.load(spark, path) === SQLDataTypes.VectorType) + } + + test("DataTypeParam safely persists nested Spark SQL schemas as JSON") { + val holder = new TestParamsHolder + val path = new Path(new File(tmpDir.toFile, "nested-data-type").toString) + val expected = StructType(Seq( + StructField("items", ArrayType(MapType(StringType, DecimalType(12, 4)))) + )) + + holder.dataTypeParam.save(expected, spark, path, overwrite = true) + + assert(holder.dataTypeParam.load(spark, path) === expected) + } + + test("DataTypeParam ignores UDT-shaped values in StructField metadata") { + val holder = new TestParamsHolder + val path = new Path(new File(tmpDir.toFile, "metadata-data-type").toString) + val metadata = new MetadataBuilder() + .putString("type", "udt") + .putString("class", "com.example.NotAType") + .build() + val expected = StructType(Seq(StructField("value", StringType, nullable = true, metadata))) + + holder.dataTypeParam.save(expected, spark, path, overwrite = true) + + assert(holder.dataTypeParam.load(spark, path) === expected) + } + + test("DataTypeParam requires explicit trust for custom UDT JSON") { + val holder = new TestParamsHolder + val path = new Path(new File(tmpDir.toFile, "custom-data-type").toString) + val config = Serializer.LegacyObjectDeserializationConfig + val previous = spark.conf.getOption(config) + holder.dataTypeParam.save(new TestStringUDT, spark, path, overwrite = true) + spark.conf.unset(config) + + try { + val error = intercept[SecurityException] { + holder.dataTypeParam.load(spark, path) + } + assert(error.getMessage.contains(classOf[TestStringUDT].getName)) + assert(error.getMessage.contains(config)) + + spark.conf.set(config, "true") + assert(holder.dataTypeParam.load(spark, path).isInstanceOf[TestStringUDT]) + } finally { + previous match { + case Some(value) => spark.conf.set(config, value) + case None => spark.conf.unset(config) + } + } + } + + test("DataTypeParam rejects legacy Java streams before deserialization callbacks") { + val holder = new TestParamsHolder + val path = new Path(new File(tmpDir.toFile, "data-type-tripwire").toString) + Serializer.writeToHDFS[AnyRef]( + spark, + new DataTypeLoadTripwire, + path, + overwrite = true + ) + val config = Serializer.LegacyObjectDeserializationConfig + val previous = spark.conf.getOption(config) + spark.conf.unset(config) + DataTypeLoadTripwire.Triggered.set(false) + + try { + val error = intercept[SecurityException] { + holder.dataTypeParam.load(spark, path) + } + assert(error.getMessage.contains(config)) + assert(!DataTypeLoadTripwire.Triggered.get()) + } finally { + previous.fold(spark.conf.unset(config))(spark.conf.set(config, _)) + } + } + + test("DataTypeParam loads legacy Java DataTypes only with explicit trust") { + val holder = new TestParamsHolder + val path = new Path(new File(tmpDir.toFile, "legacy-data-type").toString) + val expected = StructType(Seq(StructField("value", StringType))) + Serializer.writeToHDFS[DataType](spark, expected, path, overwrite = true) + val config = Serializer.LegacyObjectDeserializationConfig + val previous = spark.conf.getOption(config) + spark.conf.unset(config) + + try { + assertThrows[SecurityException] { + holder.dataTypeParam.load(spark, path) + } + spark.conf.set(config, "true") + assert(holder.dataTypeParam.load(spark, path) === expected) + } finally { + previous.fold(spark.conf.unset(config))(spark.conf.set(config, _)) + } + } } diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyEstimatorArrayParam.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyEstimatorArrayParam.scala index 5c0ed4a6089..59877c21700 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyEstimatorArrayParam.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyEstimatorArrayParam.scala @@ -3,19 +3,26 @@ package com.microsoft.azure.synapse.ml.param +import com.microsoft.azure.synapse.ml.core.serialize.TestEstimatorBase import com.microsoft.azure.synapse.ml.core.test.base.TestBase -import org.apache.spark.ml.Estimator +import com.microsoft.azure.synapse.ml.train.ComputePerInstanceStatistics +import org.apache.hadoop.fs.Path +import org.apache.spark.ml.{Estimator, Serializer, Transformer} import org.apache.spark.ml.classification.{LogisticRegression, DecisionTreeClassifier} -import org.apache.spark.ml.feature.StringIndexer +import org.apache.spark.ml.feature.{SQLTransformer, StringIndexer} +import org.apache.spark.ml.linalg.Vectors import org.apache.spark.ml.param.{ParamMap, Params} +import java.io.File import java.util.{ArrayList => JArrayList} +import scala.reflect.runtime.universe.typeOf class VerifyEstimatorArrayParam extends TestBase { private class TestParamsHolder extends Params { override val uid: String = "test-holder" val estimatorsParam = new EstimatorArrayParam(this, "estimators", "An array of estimators") + val transformersParam = new TransformerArrayParam(this, "transformers", "An array of transformers") override def copy(extra: ParamMap): Params = this } @@ -90,4 +97,106 @@ class VerifyEstimatorArrayParam extends TestBase { val holder = new TestParamsHolder assert(holder.get(holder.estimatorsParam).isEmpty) } + + test("Stage array params persist writable stages without Java deserialization") { + val holder = new TestParamsHolder + val estimatorPath = new Path(new File(tmpDir.toFile, "estimators").toString) + val transformerPath = new Path(new File(tmpDir.toFile, "transformers").toString) + val estimators = Array[Estimator[_]](new LogisticRegression(), new DecisionTreeClassifier()) + val transformers = Array[Transformer]( + new SQLTransformer().setStatement("SELECT * FROM __THIS__"), + new ComputePerInstanceStatistics() + ) + + holder.estimatorsParam.save(estimators, spark, estimatorPath, overwrite = true) + holder.transformersParam.save(transformers, spark, transformerPath, overwrite = true) + + assert(holder.estimatorsParam.load(spark, estimatorPath).map(_.getClass) + .sameElements(estimators.map(_.getClass))) + assert(holder.transformersParam.load(spark, transformerPath).map(_.getClass) + .sameElements(transformers.map(_.getClass))) + } + + test("Generic stage-array serializers preserve the requested runtime array type") { + val estimatorPath = new Path(new File(tmpDir.toFile, "generic-estimators").toString) + val transformerPath = new Path(new File(tmpDir.toFile, "generic-transformers").toString) + val estimators = Array[Estimator[_]](new LogisticRegression()) + val transformers = Array[Transformer]( + new SQLTransformer().setStatement("SELECT * FROM __THIS__") + ) + val estimatorSerializer = Serializer.typeToSerializer[Array[Estimator[_]]]( + typeOf[Array[Estimator[_]]], + spark + ) + val transformerSerializer = Serializer.typeToSerializer[Array[Transformer]]( + typeOf[Array[Transformer]], + spark + ) + + estimatorSerializer.write(estimators, estimatorPath, overwrite = true) + transformerSerializer.write(transformers, transformerPath, overwrite = true) + val loadedEstimators = estimatorSerializer.read(estimatorPath) + val loadedTransformers = transformerSerializer.read(transformerPath) + + assert(loadedEstimators.getClass.getComponentType === classOf[Estimator[_]]) + assert(loadedTransformers.getClass.getComponentType === classOf[Transformer]) + assert(loadedEstimators.map(_.getClass).sameElements(estimators.map(_.getClass))) + assert(loadedTransformers.map(_.getClass).sameElements(transformers.map(_.getClass))) + } + + test("Non-writable stage arrays require explicit trust") { + val holder = new TestParamsHolder + val path = new Path(new File(tmpDir.toFile, "non-writable-transformers").toString) + val config = Serializer.LegacyObjectDeserializationConfig + val previous = spark.conf.getOption(config) + holder.transformersParam.save( + Array[Transformer](new TestEstimatorBase("non-writable")), + spark, + path, + overwrite = true + ) + spark.conf.unset(config) + + try { + assertThrows[SecurityException] { + holder.transformersParam.load(spark, path) + } + spark.conf.set(config, "true") + assert(holder.transformersParam.load(spark, path).head.uid === "non-writable") + } finally { + previous.fold(spark.conf.unset(config))(spark.conf.set(config, _)) + } + } + + test("Data-bearing stage arrays require explicit trust") { + import spark.implicits._ + + val holder = new TestParamsHolder + val path = new Path(new File(tmpDir.toFile, "fitted-transformers").toString) + val training = Seq( + (0.0, Vectors.dense(0.0, 1.0)), + (1.0, Vectors.dense(1.0, 0.0)), + (0.0, Vectors.dense(0.1, 0.9)), + (1.0, Vectors.dense(0.9, 0.1)) + ).toDF("label", "features") + val model = new LogisticRegression().setMaxIter(1).fit(training) + val config = Serializer.LegacyObjectDeserializationConfig + val previous = spark.conf.getOption(config) + + holder.transformersParam.save(Array[Transformer](model), spark, path, overwrite = true) + spark.conf.unset(config) + + try { + assertThrows[SecurityException] { + holder.transformersParam.load(spark, path) + } + spark.conf.set(config, "true") + val loaded = holder.transformersParam.load(spark, path).head + assert(loaded.uid === model.uid) + assert(loaded.transform(training).count() === training.count()) + } finally { + previous.fold(spark.conf.unset(config))(spark.conf.set(config, _)) + } + } + } diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyEvaluatorParam.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyEvaluatorParam.scala index f7ef0222800..d5b7da46cdf 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyEvaluatorParam.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyEvaluatorParam.scala @@ -4,11 +4,15 @@ package com.microsoft.azure.synapse.ml.param import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import org.apache.hadoop.fs.Path +import org.apache.spark.ml.Serializer import org.apache.spark.ml.evaluation.{ BinaryClassificationEvaluator, MulticlassClassificationEvaluator, RegressionEvaluator } import org.apache.spark.ml.param.{ParamMap, Params} +import java.io.File + class VerifyEvaluatorParam extends TestBase { private class TestParamsHolder extends Params { @@ -92,4 +96,27 @@ class VerifyEvaluatorParam extends TestBase { val holder = new TestParamsHolder assert(holder.get(holder.evaluatorParam).isEmpty) } + + test("EvaluatorParam persistence requires explicit trust for serialized evaluators") { + val holder = new TestParamsHolder + val path = new Path(new File(tmpDir.toFile, "evaluator").toString) + val evaluator = new BinaryClassificationEvaluator() + holder.evaluatorParam.save(evaluator, spark, path, overwrite = true) + val config = Serializer.LegacyObjectDeserializationConfig + val previous = spark.conf.getOption(config) + spark.conf.unset(config) + + try { + assertThrows[SecurityException] { + holder.evaluatorParam.load(spark, path) + } + spark.conf.set(config, "true") + assert(holder.evaluatorParam.load(spark, path).getClass === evaluator.getClass) + } finally { + previous match { + case Some(value) => spark.conf.set(config, value) + case None => spark.conf.unset(config) + } + } + } } diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyModelParam.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyModelParam.scala index 6723898d278..25ae4b7f793 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyModelParam.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyModelParam.scala @@ -63,7 +63,8 @@ class VerifyModelParam extends TestBase { test("ModelParam pyLoadLine generates Python code") { val holder = new TestParamsHolder val pyCode = holder.modelParam.pyLoadLine(1) - assert(pyCode.contains("Pipeline.load")) + assert(pyCode.contains("PipelineSerializer")) + assert(pyCode.contains("JavaParams._from_java")) assert(pyCode.contains("model-1.model")) assert(pyCode.contains("complexParams")) } diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyPipelineStageParams.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyPipelineStageParams.scala index 2047cbde558..c690e0911bb 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyPipelineStageParams.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/param/VerifyPipelineStageParams.scala @@ -4,6 +4,7 @@ package com.microsoft.azure.synapse.ml.param import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import com.microsoft.azure.synapse.ml.core.test.fuzzing.RTestFuzzing import org.apache.spark.ml.{Transformer, Estimator, Model, PipelineStage} import org.apache.spark.ml.feature.{Tokenizer, HashingTF, StringIndexer, StringIndexerModel} import org.apache.spark.ml.param.{ParamMap, Params} @@ -125,10 +126,10 @@ class VerifyPipelineStageParams extends TestBase { test("PipelineStageWrappable pyLoadLine generates Python code") { val holder = new TestParamsHolder val pyCode = holder.transformerParam.pyLoadLine(1) - assert(pyCode.contains("Pipeline.load")) + assert(pyCode.contains("PipelineSerializer")) + assert(pyCode.contains("JavaParams._from_java")) assert(pyCode.contains("model-1.model")) assert(pyCode.contains("complexParams")) - assert(pyCode.contains("getStages()")) } test("PipelineStageWrappable rValue returns model reference") { @@ -138,6 +139,17 @@ class VerifyPipelineStageParams extends TestBase { assert(rVal === "transformerModel") } + test("R test generation uses bounded pipeline-stage loading") { + val holder = new TestParamsHolder + val rCode = RTestFuzzing.loadLine(holder.transformerParam, 4).get + assert(rCode.contains("PipelineSerializer")) + assert(rCode.contains("spark_session(sc)")) + assert(rCode.contains("ml_call_constructor")) + assert(rCode.contains("model-4.model")) + assert(rCode.contains("complexParams")) + assert(!rCode.contains("ml_load")) + } + test("PipelineStageWrappable assertEquality succeeds for same transformer") { val holder = new TestParamsHolder val t1 = new Tokenizer().setInputCol("a").setOutputCol("b") diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/recommendation/RecommendationIndexerSpec.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/recommendation/RecommendationIndexerSpec.scala index d9d66d7641b..a5fbe251694 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/recommendation/RecommendationIndexerSpec.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/recommendation/RecommendationIndexerSpec.scala @@ -14,7 +14,7 @@ class RecommendationIndexerSpec extends RankingTestBase with EstimatorFuzzing[Re override def reader: MLReadable[_] = RecommendationIndexer - override def modelReader: MLReadable[_] = RankingAdapterModel + override def modelReader: MLReadable[_] = RecommendationIndexerModel test("ALS") { @@ -45,5 +45,5 @@ class RecommendationIndexerModelSpec extends RankingTestBase with TransformerFuz List(new TestObject(recommendationIndexer.fit(df), df)) } - override def reader: MLReadable[_] = RankingAdapterModel + override def reader: MLReadable[_] = RecommendationIndexerModel } diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/recommendation/SARIdentifierSpec.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/recommendation/SARIdentifierSpec.scala index 8b6c28ca94f..e2505f697c8 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/recommendation/SARIdentifierSpec.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/recommendation/SARIdentifierSpec.scala @@ -358,8 +358,10 @@ class SARIdentifierSpec extends TestBase { val path = root.resolve("model").toString try { - model.write.overwrite().save(path) - val loaded = SARModel.load(path) + val loaded = withTrustedLegacyModelPersistence { + model.write.overwrite().save(path) + SARModel.load(path) + } assert(loaded.isSet(loaded.userIdsFitInt)) assert(loaded.isSet(loaded.itemIdsFitInt)) diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/stages/UDFTransformerSuite.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/stages/UDFTransformerSuite.scala index 2e83d65138d..3d449a3ff1b 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/stages/UDFTransformerSuite.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/stages/UDFTransformerSuite.scala @@ -103,6 +103,95 @@ class UDFTransformerSuite extends TestBase with TransformerFuzzing[UDFTransforme } } + test("Persisted UDFs require explicit trust before Java deserialization") { + import org.apache.commons.io.FileUtils + import org.apache.spark.ml.{Pipeline, PipelineModel, Serializer} + + import java.io.File + + spark + val path = new File(tmpDir.toFile, "udf-transformer.model") + val pipelinePath = new File(tmpDir.toFile, "udf-pipeline.model") + val transformer = new UDFTransformer().setUDF(stringToIntegerUDF) + .setInputCol("words").setOutputCol(outCol) + val config = Serializer.LegacyObjectDeserializationConfig + val previous = spark.conf.getOption(config) + transformer.write.overwrite().save(path.toString) + spark.conf.set(config, "true") + Serializer.withTrustedArtifactLoad(spark) { + new Pipeline().setStages(Array(transformer)).fit(baseDF) + .write.overwrite().save(pipelinePath.toString) + } + spark.conf.unset(config) + + try { + val error = intercept[SecurityException] { + UDFTransformer.load(path.toString) + } + assert(error.getMessage.contains(config)) + assertThrows[SecurityException] { + PipelineModel.load(pipelinePath.toString) + } + + spark.conf.set(config, "true") + val loaded = UDFTransformer.read.session(spark).load(path.toString) + assertDFEq(transformer.transform(baseDF), loaded.transform(baseDF)) + val loadedPipeline = Serializer.withTrustedArtifactLoad(spark) { + PipelineModel.read.session(spark).load(pipelinePath.toString) + } + assert(loadedPipeline.stages.length === 1) + } finally { + previous match { + case Some(value) => spark.conf.set(config, value) + case None => spark.conf.unset(config) + } + FileUtils.deleteDirectory(path) + FileUtils.deleteDirectory(pipelinePath) + } + } + + test("PipelineSerializer applies the supplied session to nested stages") { + import org.apache.commons.io.FileUtils + import org.apache.hadoop.fs.Path + import org.apache.spark.ml.{PipelineSerializer, Serializer} + import org.apache.spark.sql.SparkSession + + import java.io.File + + spark + val path = new File(tmpDir.toFile, "udf-pipeline-stage.model") + val transformer = new UDFTransformer().setUDF(stringToIntegerUDF) + .setInputCol("words").setOutputCol(outCol) + new PipelineSerializer(spark).write(transformer, new Path(path.toString), overwrite = true) + + val isolated = spark.newSession() + val config = Serializer.LegacyObjectDeserializationConfig + val previousDefault = spark.conf.getOption(config) + val previousIsolated = isolated.conf.getOption(config) + + def restore(session: SparkSession, value: Option[String]): Unit = { + value.fold(session.conf.unset(config))(session.conf.set(config, _)) + } + + try { + spark.conf.set(config, "true") + isolated.conf.unset(config) + val error = intercept[SecurityException] { + new PipelineSerializer(isolated).read(new Path(path.toString)) + } + assert(error.getMessage.contains(config)) + + spark.conf.unset(config) + isolated.conf.set(config, "true") + val loaded = new PipelineSerializer(isolated).read(new Path(path.toString)) + assert(loaded.isInstanceOf[UDFTransformer]) + } finally { + restore(spark, previousDefault) + restore(isolated, previousIsolated) + FileUtils.deleteDirectory(path) + } + } + def testObjects(): Seq[TestObject[UDFTransformer]] = { List(new TestObject( new UDFTransformer().setUDF(stringToIntegerUDF) diff --git a/core/src/test/scala/com/microsoft/azure/synapse/ml/train/VerifyTrainRegressor.scala b/core/src/test/scala/com/microsoft/azure/synapse/ml/train/VerifyTrainRegressor.scala index 1679b0b446f..5584566dd3f 100644 --- a/core/src/test/scala/com/microsoft/azure/synapse/ml/train/VerifyTrainRegressor.scala +++ b/core/src/test/scala/com/microsoft/azure/synapse/ml/train/VerifyTrainRegressor.scala @@ -129,8 +129,10 @@ class VerifyTrainRegressor extends EstimatorFuzzing[TrainRegressor] { assert(predictions.forall(value => !value.isNaN && !value.isInfinity)) val modelFile = new File(tmpDir.toFile, "rf") - model.write.overwrite().save(modelFile.toString) - val loadedModel = TrainedRegressorModel.load(modelFile.toString) + val loadedModel = withTrustedLegacyModelPersistence { + model.write.overwrite().save(modelFile.toString) + TrainedRegressorModel.load(modelFile.toString) + } val loadedPredictions = loadedModel.transform(dataset).select("prediction").collect().map(_.getDouble(0)) assert(predictions.toSeq === loadedPredictions.toSeq) } @@ -168,14 +170,15 @@ class VerifyTrainRegressor extends EstimatorFuzzing[TrainRegressor] { val model = linearRegressor.fit(dataset) val modelFile = new File(tmpDir.toFile, "testModel") - model.write.overwrite().save(modelFile.toString) - // write a second time with overwrite flag, verify still works - model.write.overwrite().save(modelFile.toString) - // assert directory exists - assert(modelFile.exists()) - - // load the model - val loadedModel = TrainedRegressorModel.load(modelFile.toString) + val loadedModel = withTrustedLegacyModelPersistence { + model.write.overwrite().save(modelFile.toString) + // write a second time with overwrite flag, verify still works + model.write.overwrite().save(modelFile.toString) + // assert directory exists + assert(modelFile.exists()) + + TrainedRegressorModel.load(modelFile.toString) + } // verify model data loaded assert(loadedModel.getLabelCol == model.getLabelCol) diff --git a/core/src/test/scala/org/apache/spark/ml/VerifyArtifactPathResolver.scala b/core/src/test/scala/org/apache/spark/ml/VerifyArtifactPathResolver.scala new file mode 100644 index 00000000000..66435d5bb1c --- /dev/null +++ b/core/src/test/scala/org/apache/spark/ml/VerifyArtifactPathResolver.scala @@ -0,0 +1,138 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package org.apache.spark.ml + +import com.microsoft.azure.synapse.ml.core.env.StreamUtilities.using +import com.microsoft.azure.synapse.ml.core.serialize.StandardParamTest +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.Path + +import java.io.{ByteArrayInputStream, ByteArrayOutputStream} +import java.io.File +import java.net.URI +import java.nio.charset.StandardCharsets + +abstract class ArtifactPathResolverTestBase extends TestBase { + + test("ABFS metadata enumeration requires its incremental iterator configuration") { + val filesystemClasses = Set("org.apache.hadoop.fs.azurebfs.AzureBlobFileSystem") + val filesystemUri = URI.create("abfss://container@account.dfs.core.windows.net") + val configuration = new Configuration(false) + val configKey = "fs.azure.enable.abfslistiterator" + + assert(ArtifactPathResolver.guaranteesIncrementalMetadataListing( + filesystemClasses, + filesystemUri, + configuration + )) + + configuration.setBoolean(configKey, false) + assert(!ArtifactPathResolver.guaranteesIncrementalMetadataListing( + filesystemClasses, + filesystemUri, + configuration + )) + + configuration.setBoolean(s"$configKey.account.dfs.core.windows.net", true) + assert(ArtifactPathResolver.guaranteesIncrementalMetadataListing( + filesystemClasses, + filesystemUri, + configuration + )) + } + + test("S3A metadata enumeration recognizes its incremental iterator") { + assert(ArtifactPathResolver.guaranteesIncrementalMetadataListing( + Set("org.apache.hadoop.fs.s3a.S3AFileSystem"), + URI.create("s3a://model-bucket"), + new Configuration(false) + )) + } + + test("Metadata input reads only through the requested byte limit") { + val input = new ByteArrayInputStream(Array.tabulate[Byte](32)(_.toByte)) + + val bytes = ArtifactPathResolver.readUpTo(input, 17) + + assert(bytes.sameElements(Array.tabulate[Byte](17)(_.toByte))) + assert(input.read() === 17) + } + + test("Metadata reader has no post-Java-8 InputStream method linkage") { + val input = Option(ArtifactPathResolver.getClass + .getResourceAsStream("ArtifactPathResolver$.class")) + .getOrElse(fail("ArtifactPathResolver class resource was unavailable")) + val output = new ByteArrayOutputStream() + val buffer = new Array[Byte](8192) // scalastyle:ignore magic.number + using(input) { stream => + var count = stream.read(buffer) + while (count >= 0) { // scalastyle:ignore while + if (count > 0) { + output.write(buffer, 0, count) + } + count = stream.read(buffer) + } + }.get + + val classFile = new String(output.toByteArray, StandardCharsets.ISO_8859_1) + assert(!classFile.contains("readNBytes")) + } + + test("ComplexParamsReader uses preloaded metadata without reopening it") { + spark + val modelPath = new File(tmpDir.toFile, "preloaded-metadata").toString + val original = new StandardParamTest("preloadedMetadata").setStringParam("value") + original.write.overwrite().save(modelPath) + val rootPath = new Path(modelPath) + val className = classOf[StandardParamTest].getName + val metadata = ArtifactPathResolver.loadMetadata(spark, rootPath, className) + val metadataPath = new Path(rootPath, "metadata") + assert(metadataPath.getFileSystem(spark.sparkContext.hadoopConfiguration) + .delete(metadataPath, true)) + + val loaded = new ComplexParamsReader[StandardParamTest]( + className, + Some(classOf[StandardParamTest]), + Some(metadata) + ).session(spark).load(modelPath) + + assert(loaded.uid === original.uid) + assert(loaded.getStringParam === original.getStringParam) + } + + test("ComplexParams metadata uses the Spark ML part file layout") { + spark + val modelPath = new File(tmpDir.toFile, "spark-ml-metadata-layout").toString + new StandardParamTest("sparkMlMetadataLayout") + .setStringParam("value") + .write.overwrite().save(modelPath) + + val metadataPath = new Path(modelPath, "metadata") + val fs = metadataPath.getFileSystem(Serializer.sessionHadoopConf(spark)) + assert(fs.getFileStatus(new Path(metadataPath, "part-00000")).isFile) + assert(fs.getFileStatus(new Path(metadataPath, "_SUCCESS")).isFile) + } + + test("ComplexParams metadata write limit includes Spark text framing") { + spark + val maxMetadataRecordBytes = 1048576 + val model = new StandardParamTest("metadataFraming").setStringParam("") + val baseBytes = ComplexParamsWriter.getMetadataToSave(model, spark) + .getBytes(StandardCharsets.UTF_8).length + model.setStringParam("a" * (maxMetadataRecordBytes - baseBytes)) + val exactLimitBytes = ComplexParamsWriter.getMetadataToSave(model, spark) + .getBytes(StandardCharsets.UTF_8).length + assert(exactLimitBytes === maxMetadataRecordBytes) + + val error = intercept[IllegalArgumentException] { + ComplexParamsWriter.saveMetadata( + model, + new File(tmpDir.toFile, "metadata-framing").toString, + spark + ) + } + assert(error.getMessage.contains("exceeds")) + } +} diff --git a/lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/params/LightGBMBoosterParam.scala b/lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/params/LightGBMBoosterParam.scala index e2d0dfeb6a2..be637fb16e5 100644 --- a/lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/params/LightGBMBoosterParam.scala +++ b/lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/params/LightGBMBoosterParam.scala @@ -4,6 +4,7 @@ package com.microsoft.azure.synapse.ml.lightgbm.params import com.microsoft.azure.synapse.ml.core.serialize.ComplexParam +import com.microsoft.azure.synapse.ml.core.utils.DeserializationClassFilter import com.microsoft.azure.synapse.ml.lightgbm.booster.LightGBMBooster import com.microsoft.azure.synapse.ml.param.WrappableParam import org.apache.spark.ml.param.Params @@ -19,5 +20,17 @@ class LightGBMBoosterParam(parent: Params, name: String, doc: String, def this(parent: Params, name: String, doc: String) = this(parent, name, doc, { _ => true }) + override protected def deserializationClassFilter: Option[DeserializationClassFilter] = { + Some(DeserializationClassFilter( + allowedClasses = Set( + classOf[LightGBMBooster].getName, + "java.lang.String", + "scala.None$", + "scala.Option", + "scala.Some", + "scala.runtime.ModuleSerializationProxy" + ) + )) + } } diff --git a/lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split6/VerifyLightGBMBoosterParam.scala b/lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split6/VerifyLightGBMBoosterParam.scala new file mode 100644 index 00000000000..c8b8f8d8723 --- /dev/null +++ b/lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split6/VerifyLightGBMBoosterParam.scala @@ -0,0 +1,33 @@ +// Copyright (C) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. See LICENSE in project root for information. + +package com.microsoft.azure.synapse.ml.lightgbm.split6 + +import com.microsoft.azure.synapse.ml.core.test.base.TestBase +import com.microsoft.azure.synapse.ml.lightgbm.booster.LightGBMBooster +import com.microsoft.azure.synapse.ml.lightgbm.params.LightGBMBoosterParam +import org.apache.hadoop.fs.Path +import org.apache.spark.ml.param.{ParamMap, Params} + +import java.io.File + +class VerifyLightGBMBoosterParam extends TestBase { + + private class TestParamsHolder extends Params { + override val uid: String = "lightgbm-booster-holder" + val booster = new LightGBMBoosterParam(this, "booster", "A LightGBM booster param") + + override def copy(extra: ParamMap): Params = this + } + + test("LightGBMBoosterParam loads its constrained legacy object graph") { + val holder = new TestParamsHolder + val path = new Path(new File(tmpDir.toFile, "booster").toString) + val expected = new LightGBMBooster("model-data") + + holder.booster.save(expected, spark, path, overwrite = true) + val loaded = holder.booster.load(spark, path) + + assert(loaded.getNativeModel() === expected.getNativeModel()) + } +}