[Serving] Implement SageMaker Secure Mode & support for multiple data sources (#2042)

Author: ethnzhng

gradle/libs.versions.toml

@@ -17,6 +17,7 @@ httpcomponents = "4.5.14"
 
 testng = "7.10.2"
 junit = "4.13.2"
+mockitocore = "5.12.0"
 
 [libraries]
 huggingface-tokenizers = { module = "ai.djl.huggingface:tokenizers" }
@@ -41,3 +42,4 @@ prometheus-exposition-formats = { module = "io.prometheus:prometheus-metrics-exp
 
 testng = { module = "org.testng:testng", version.ref = "testng" }
 junit = { module = "junit:junit", version.ref = "junit" }
+mockito-core = { module = "org.mockito:mockito-core", version.ref = "mockitocore" }

plugins/secure-mode/README.md

@@ -0,0 +1,3 @@
+# DJL Serving - Secure Mode Plugin
+
+This plugin implements SageMaker Secure Mode for the model server, by performing checks for potentially unsafe files and options. It is configured by the SageMaker platform.

plugins/secure-mode/build.gradle.kts

@@ -0,0 +1,23 @@
+plugins {
+    ai.djl.javaProject
+}
+
+dependencies {
+    implementation(project(":serving"))
+    implementation(project(":wlm"))
+
+    testImplementation(libs.testng) {
+        exclude(group = "junit", module = "junit")
+    }
+    testImplementation(libs.mockito.core)
+
+}
+
+tasks {
+    register<Copy>("copyJar") {
+        from(jar) // here it automatically reads jar file produced from jar task
+        into("../../serving/plugins")
+    }
+
+    jar { finalizedBy("copyJar") }
+}

plugins/secure-mode/gradlew

@@ -0,0 +1 @@
+../../gradlew

plugins/secure-mode/src/main/java/ai/djl/serving/plugins/securemode/IllegalConfigurationException.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.serving.plugins.securemode;
+
+/** Thrown when Secure Mode encounters a violation during security checks. */
+public class IllegalConfigurationException extends RuntimeException {
+
+    static final long serialVersionUID = 1L;
+
+    /**
+     * Constructs a {@link IllegalConfigurationException} with the specified detail message.
+     *
+     * @param message The detail message (which is saved for later retrieval by the {@link
+     *     #getMessage()} method)
+     */
+    public IllegalConfigurationException(String message) {
+        super(message);
+    }
+
+    /**
+     * Constructs a {@link IllegalConfigurationException} with the specified detail message and
+     * cause.
+     *
+     * <p>Note that the detail message associated with {@code cause} is <i>not</i> automatically
+     * incorporated into this exception's detail message.
+     *
+     * @param message The detail message (which is saved for later retrieval by the {@link
+     *     #getMessage()} method)
+     * @param cause The cause (which is saved for later retrieval by the {@link #getCause()}
+     *     method). (A null value is permitted, and indicates that the cause is nonexistent or
+     *     unknown.)
+     */
+    public IllegalConfigurationException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}

plugins/secure-mode/src/main/java/ai/djl/serving/plugins/securemode/SecureModeModelServerListener.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.serving.plugins.securemode;
+
+import ai.djl.Device;
+import ai.djl.ModelException;
+import ai.djl.serving.wlm.ModelInfo;
+import ai.djl.serving.wlm.util.ModelServerListenerAdapter;
+
+import java.io.IOException;
+import java.net.URISyntaxException;
+
+class SecureModeModelServerListener extends ModelServerListenerAdapter {
+
+    @Override
+    public void onModelLoading(ModelInfo<?, ?> model, Device device) {
+        super.onModelLoading(model, device);
+
+        if (SecureModeUtils.isSecureMode()) {
+            try {
+                SecureModeUtils.validateSecurity();
+                SecureModeUtils.reconcileSources(model.getModelUrl());
+            } catch (ModelException e) {
+                throw new IllegalConfigurationException("Secure Mode check failed", e);
+            } catch (IOException | URISyntaxException e) {
+                throw new IllegalConfigurationException(
+                        "Error while running Secure Mode checks", e);
+            }
+        }
+    }
+}

plugins/secure-mode/src/main/java/ai/djl/serving/plugins/securemode/SecureModePlugin.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.serving.plugins.securemode;
+
+import ai.djl.serving.plugins.RequestHandler;
+import ai.djl.serving.wlm.util.EventManager;
+
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.handler.codec.http.FullHttpRequest;
+import io.netty.handler.codec.http.QueryStringDecoder;
+
+/** A plugin for Secure Mode. */
+public class SecureModePlugin implements RequestHandler<Void> {
+
+    /** Constructs a new {@code SecureModePlugin} instance. */
+    public SecureModePlugin() {
+        if (SecureModeUtils.isSecureMode()) {
+            EventManager.getInstance().addListener(new SecureModeModelServerListener());
+        }
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public boolean acceptInboundMessage(Object msg) {
+        return false;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Void handleRequest(
+            ChannelHandlerContext ctx,
+            FullHttpRequest req,
+            QueryStringDecoder decoder,
+            String[] segments) {
+        return null;
+    }
+}