feat(ai.triton.server): Add Model Encryption support for Triton Serve…

…r Service (#3986) * feat(EncryptionUtils): add class and methods skeleton Signed-off-by: Mattia Dal Ben <[email protected]> * test(EncryptionUtils): add test skeleton Signed-off-by: Mattia Dal Ben <[email protected]> * test(EncryptionUtils): add basic tests for createDecryptionFolder method Signed-off-by: Mattia Dal Ben <[email protected]> * test(EncryptionUtils): refactor to Gherkin style Signed-off-by: Mattia Dal Ben <[email protected]> * test(EncryptionUtils): removed comments, renamed method Signed-off-by: Mattia Dal Ben <[email protected]> * test(EncryptionUtils): rename method Signed-off-by: Mattia Dal Ben <[email protected]> * feature(EncryptionUtils): add createDecryptionFolder method implementation Signed-off-by: Mattia Dal Ben <[email protected]> * test(EncryptionUtils): improve test with workdir setup/teardown Signed-off-by: Mattia Dal Ben <[email protected]> * test(EncryptionUtils): renamed method cleanup -> teardown Signed-off-by: Mattia Dal Ben <[email protected]> * test(resources): add encrypted files for testing Signed-off-by: Mattia Dal Ben <[email protected]> * test(EncryptionUtils): add basic tests for decryptModel method Signed-off-by: Mattia Dal Ben <[email protected]> * test(EncryptionUtils): simplify teardown method * test(EncryptionUtils): move test WORKDIR in java.io.tmpdir Signed-off-by: Mattia Dal Ben <[email protected]> * test(EcryptionUtils): add Gherkin style comments Signed-off-by: Mattia Dal Ben <[email protected]> * test(EncryptionUtils): add test for ASCII-armored file Signed-off-by: Mattia Dal Ben <[email protected]> * test(EncryptionUtils): stricter checks on exception throw Signed-off-by: Mattia Dal Ben <[email protected]> * docs: add copyright notice Signed-off-by: Mattia Dal Ben <[email protected]> * test: fix expectation param order Signed-off-by: Mattia Dal Ben <[email protected]> * build: add org.bouncycastle.openpgp as local dependency Signed-off-by: Mattia Dal Ben <[email protected]> * feat(EncryptionUtils): add first implementation of decryptModel method Signed-off-by: Mattia Dal Ben <[email protected]> * feat(EncryptionUtils): fix BouncyCastle provider addition Signed-off-by: Mattia Dal Ben <[email protected]> * feat(EncryptionUtils): add preconditions check Signed-off-by: Mattia Dal Ben <[email protected]> * refactor: improve var names Signed-off-by: Mattia Dal Ben <[email protected]> * refactor: remove unnecessary toString Signed-off-by: Mattia Dal Ben <[email protected]> * feat(EncryptionUtils): move stream instantiation in decryptFile method Signed-off-by: Mattia Dal Ben <[email protected]> * refactor: move stream instantiation close to use Signed-off-by: Mattia Dal Ben <[email protected]> * refactor: improve var names Signed-off-by: Mattia Dal Ben <[email protected]> * feat(EncryptionUtils): rethrow PGPException on decryption failure Signed-off-by: Mattia Dal Ben <[email protected]> * test(resources): rename resourse files Signed-off-by: Mattia Dal Ben <[email protected]> * test(EncryptionUtils): add zip file test Signed-off-by: Mattia Dal Ben <[email protected]> * test: refactor to Gerkhin style Signed-off-by: Mattia Dal Ben <[email protected]> * test: remove unnecessary vars Signed-off-by: Mattia Dal Ben <[email protected]> * test: remove unnecessary "this." Signed-off-by: Mattia Dal Ben <[email protected]> * test: renamed method Signed-off-by: Mattia Dal Ben <[email protected]> * test: renamed method givenAFileExistsAtPath Signed-off-by: Mattia Dal Ben <[email protected]> * test: add basic test for deleteModel method Signed-off-by: Mattia Dal Ben <[email protected]> * test: add test for missing directory for deleteModel Signed-off-by: Mattia Dal Ben <[email protected]> * test: add check for parent folder in deleteModel test Signed-off-by: Mattia Dal Ben <[email protected]> * feat(EncryptionUtils): add deleteModel method implemention Signed-off-by: Mattia Dal Ben <[email protected]> * test: refactor to Gerkhin style Signed-off-by: Mattia Dal Ben <[email protected]> * refactor: renamed method to thenAFolderExistsAtPath Signed-off-by: Mattia Dal Ben <[email protected]> * test: remove forgotten comment Signed-off-by: Mattia Dal Ben <[email protected]> * test: add .zip test file Signed-off-by: Mattia Dal Ben <[email protected]> * test(EncryptionUtils): add unzipModel tests Signed-off-by: Mattia Dal Ben <[email protected]> * test: add unzipModel method test for wrong file format Signed-off-by: Mattia Dal Ben <[email protected]> * test: add decryptModel test for wrong file format Signed-off-by: Mattia Dal Ben <[email protected]> * feat(EncryptionUtils): add check for wrong file format Signed-off-by: Mattia Dal Ben <[email protected]> * test: improved test naming Signed-off-by: Mattia Dal Ben <[email protected]> * test(resources): fixed wrong test .zip file Signed-off-by: Mattia Dal Ben <[email protected]> * feat(EncryptionUtils): initial implementation of unzipModel method Signed-off-by: Mattia Dal Ben <[email protected]> * build: add dependency on org.eclipse.kura.core.cloud for unZip class Signed-off-by: Mattia Dal Ben <[email protected]> * feat(EncryptionUtils): add ZIP magic number check Signed-off-by: Mattia Dal Ben <[email protected]> * test: refactor to Gerkhin style Signed-off-by: Mattia Dal Ben <[email protected]> * feat(EncryptionUtils): remove unnecessary todo Signed-off-by: Mattia Dal Ben <[email protected]> * test: improve error handling in given/then methods Signed-off-by: Mattia Dal Ben <[email protected]> * feat(EncryptionUtils): add getEncryptedModelPath method skeleton Signed-off-by: Mattia Dal Ben <[email protected]> * test: add getEncryptedModelPath tests Signed-off-by: Mattia Dal Ben <[email protected]> * test: add forgotten no exception check Signed-off-by: Mattia Dal Ben <[email protected]> * test: fixed typo in WORKDIR path Signed-off-by: Mattia Dal Ben <[email protected]> * test: fix wrong "when" method Signed-off-by: Mattia Dal Ben <[email protected]> * feat(EncryptionUtils): add getEncryptedModelPath method implementation Signed-off-by: Mattia Dal Ben <[email protected]> * test: updated test with new behaviour for cleanModelRepository Signed-off-by: Mattia Dal Ben <[email protected]> * feat(EncryptionUtils): updated cleanRepository behaviour Signed-off-by: Mattia Dal Ben <[email protected]> * build: add bouncycastle lib in build properties Signed-off-by: Mattia Dal Ben <[email protected]> * build: add necessary dependencies Signed-off-by: Mattia Dal Ben <[email protected]> * refactor(EncryptionUtils): do not throw from cleanRepository method Signed-off-by: Mattia Dal Ben <[email protected]> * feat(triton): add password field Signed-off-by: Mattia Dal Ben <[email protected]> * feat(triton): add decrypted model folder path Signed-off-by: Mattia Dal Ben <[email protected]> * feat(triton): add encrypted model support Signed-off-by: Mattia Dal Ben <[email protected]> * test: temporary path fix Signed-off-by: Mattia Dal Ben <[email protected]> * fix(build): use range for kura.crypto version dependency Signed-off-by: Mattia Dal Ben <[email protected]> * refactor: use try-with-resource Signed-off-by: Mattia Dal Ben <[email protected]> * fix: remove debug print Signed-off-by: Mattia Dal Ben <[email protected]> * feat: improve decrypted model directory creation Signed-off-by: Mattia Dal Ben <[email protected]> * feat(EncryptionUtils): use createTempDirectory for decryption dir Signed-off-by: Mattia Dal Ben <[email protected]> * feat(EncryptionUtils): dynamic decrypted model folder Signed-off-by: Mattia Dal Ben <[email protected]> * refactor: create startLocalInstance method Signed-off-by: Mattia Dal Ben <[email protected]> * refactor: add isModelEncryptionActive method * refactor: rename method to modelsAreEcnrypted Signed-off-by: Mattia Dal Ben <[email protected]> * test: temporary path fix Signed-off-by: Mattia Dal Ben <[email protected]> * test: fix test work dir path Signed-off-by: Mattia Dal Ben <[email protected]> * test: add dedicated var for temp dir path Signed-off-by: Mattia Dal Ben <[email protected]> * feat(ai.triton): handle decrypt folder deletion Signed-off-by: Mattia Dal Ben <[email protected]> * test: add decryptModelShouldThrowWithDirectory test Signed-off-by: Mattia Dal Ben <[email protected]> * test: fix wrong path in unzipModel test Signed-off-by: Mattia Dal Ben <[email protected]> * build: use bouncy castle provider from target platform Signed-off-by: Mattia Dal Ben <[email protected]> * revert: use bouncy castle provider from target platform This reverts commit a2b746e. * fix: add check for non-existent folder in getEncryptedModelPath Signed-off-by: Mattia Dal Ben <[email protected]> * build: use bouncy castle provider from target platform Signed-off-by: Mattia Dal Ben <[email protected]> * feat(EncryptionUtils): log exception stacktrace on failure Signed-off-by: Mattia Dal Ben <[email protected]>
eclipse-kura · May 24, 2022 · e0aa32f · e0aa32f
1 parent 7ac4cf3
commit e0aa32f
Show file tree

Hide file tree

Showing 15 changed files with 869 additions and 16 deletions.
diff --git a/kura/org.eclipse.kura.ai.triton.server/META-INF/MANIFEST.MF b/kura/org.eclipse.kura.ai.triton.server/META-INF/MANIFEST.MF
@@ -7,6 +7,7 @@ Bundle-Vendor: Eclipse Kura
 Bundle-RequiredExecutionEnvironment: JavaSE-1.8
 Service-Component: OSGI-INF/*.xml
 Bundle-ClassPath: .,
+ lib/bcpg-jdk15on-1.68.jar,
  lib/com.google.protobuf_3.19.3.jar,
  lib/grpc-netty-shaded-1.33.1.jar,
  lib/grpc-protobuf-1.33.1.jar,
@@ -22,10 +23,22 @@ Import-Package: com.google.common.base;version="25.0.0",
  com.google.common.util.concurrent;version="25.0.0",
  com.google.gson;version="2.7.0",
  javax.annotation;version="1.2.0",
+ org.bouncycastle.jcajce;version="1.68.0",
+ org.bouncycastle.jcajce.util;version="1.68.0",
+ org.bouncycastle.jcajce.interfaces;version="1.68.0",
+ org.bouncycastle.jcajce.io;version="1.68.0",
+ org.bouncycastle.jcajce.spec;version="1.68.0",
+ org.bouncycastle.jce.provider;version="1.68.0",
+ org.bouncycastle.util;version="1.68.0",
+ org.bouncycastle.util.encoders;version="1.68.0",
+ org.bouncycastle.util.io;version="1.68.0",
  org.apache.commons.io;version="2.4.0",
+ org.apache.commons.io.filefilter;version="2.11.0",
  org.eclipse.kura;version="[1.0,2.0)",
  org.eclipse.kura.ai.inference;version="[1.0,1.1)",
+ org.eclipse.kura.cloud.app.command;version="[1.0,2.0)",
  org.eclipse.kura.configuration;version="[1.1,2.0)",
+ org.eclipse.kura.crypto;version="[1.0,2.0)",
  org.eclipse.kura.core.linux.executor;version="[1.0,2.0)",
  org.eclipse.kura.core.util;version="[1.0,2.0)",
  org.eclipse.kura.executor;version="[1.0,2.0)",

diff --git a/...triton.server/OSGI-INF/metatype/org.eclipse.kura.ai.triton.server.TritonServerService.xml b/...triton.server/OSGI-INF/metatype/org.eclipse.kura.ai.triton.server.TritonServerService.xml
@@ -64,6 +64,15 @@
             default="" 
             description="Only for local instance, specify the path on the filesystem where the models are stored.">
        </AD>
+
+       <AD id="local.model.repository.password"  
+            name="Local model decryption password"
+            type="Password"
+            cardinality="0" 
+            required="false"
+            default="" 
+            description="Only for local instance, specify the password to be used for decrypting models stored in the model repository. If none is specified, models are supposed to be plaintext.">
+       </AD>
 
         <AD id="local.backends.path"  
             name="Local backends path"

diff --git a/....kura.ai.triton.server/OSGI-INF/org.eclipse.kura.ai.triton.server.TritonServerService.xml b/....kura.ai.triton.server/OSGI-INF/org.eclipse.kura.ai.triton.server.TritonServerService.xml
@@ -21,4 +21,5 @@
    </service>
    <property name="service.pid" type="String" value="org.eclipse.kura.ai.triton.server.TritonServerService"/>
    <reference bind="setCommandExecutorService" cardinality="1..1" interface="org.eclipse.kura.executor.PrivilegedExecutorService" name="PrivilegedExecutorService" policy="static" />
+   <reference bind="setCryptoService" cardinality="1..1" interface="org.eclipse.kura.crypto.CryptoService" name="CryptoService" policy="static"/>
 </scr:component>
diff --git a/kura/org.eclipse.kura.ai.triton.server/build.properties b/kura/org.eclipse.kura.ai.triton.server/build.properties
@@ -15,6 +15,7 @@ bin.includes = .,\
                OSGI-INF/,\
                about.html,\
                about_files/,\
+               lib/bcpg-jdk15on-1.68.jar/,\
                lib/com.google.protobuf_3.19.3.jar,\
                lib/grpc-netty-shaded-1.33.1.jar,\
                lib/grpc-protobuf-1.33.1.jar,\

diff --git a/kura/org.eclipse.kura.ai.triton.server/lib/bcpg-jdk15on-1.68.jar b/kura/org.eclipse.kura.ai.triton.server/lib/bcpg-jdk15on-1.68.jar
diff --git a/...n.server/src/main/java/org/eclipse/kura/ai/triton/server/TritonServerEncryptionUtils.java b/...n.server/src/main/java/org/eclipse/kura/ai/triton/server/TritonServerEncryptionUtils.java
@@ -0,0 +1,202 @@
+/*******************************************************************************
+ * Copyright (c) 2022 Eurotech and/or its affiliates and others
+ *
+ * This program and the accompanying materials are made
+ * available under the terms of the Eclipse Public License 2.0
+ * which is available at https://www.eclipse.org/legal/epl-2.0/
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ * Contributors:
+ *  Eurotech
+ ******************************************************************************/
+
+package org.eclipse.kura.ai.triton.server;
+
+import java.io.BufferedInputStream;
+import java.io.File;
+import java.io.FileFilter;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.nio.file.attribute.PosixFilePermission;
+import java.nio.file.attribute.PosixFilePermissions;
+import java.security.Security;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+
+import org.apache.commons.io.FileUtils;
+import org.apache.commons.io.filefilter.WildcardFileFilter;
+import org.bouncycastle.jce.provider.BouncyCastleProvider;
+import org.bouncycastle.openpgp.PGPCompressedData;
+import org.bouncycastle.openpgp.PGPEncryptedDataList;
+import org.bouncycastle.openpgp.PGPException;
+import org.bouncycastle.openpgp.PGPLiteralData;
+import org.bouncycastle.openpgp.PGPPBEEncryptedData;
+import org.bouncycastle.openpgp.PGPUtil;
+import org.bouncycastle.openpgp.jcajce.JcaPGPObjectFactory;
+import org.bouncycastle.openpgp.operator.jcajce.JcaPGPDigestCalculatorProviderBuilder;
+import org.bouncycastle.openpgp.operator.jcajce.JcePBEDataDecryptorFactoryBuilder;
+import org.bouncycastle.util.io.Streams;
+import org.eclipse.kura.KuraIOException;
+import org.eclipse.kura.cloud.app.command.UnZip;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class TritonServerEncryptionUtils {
+
+    private static final Logger logger = LoggerFactory.getLogger(TritonServerEncryptionUtils.class);
+
+    private TritonServerEncryptionUtils() {
+    }
+
+    protected static String getEncryptedModelPath(String modelName, String folderPath) throws KuraIOException {
+        if (!Files.isDirectory(Paths.get(folderPath))) {
+            throw new KuraIOException("Model repository folder " + folderPath + " does not exist/is not a folder");
+        }
+
+        File dir = new File(folderPath);
+        FileFilter fileFilter = new WildcardFileFilter(modelName + ".*");
+        File[] files = dir.listFiles(fileFilter);
+
+        if (files.length != 1) {
+            throw new KuraIOException("No good match found in folder path");
+        }
+
+        return files[0].toString();
+    }
+
+    protected static String createDecryptionFolder(String prefix) throws IOException {
+        Set<PosixFilePermission> permissions = new HashSet<>(Arrays.asList(PosixFilePermission.OWNER_READ,
+                PosixFilePermission.OWNER_WRITE, PosixFilePermission.OWNER_EXECUTE));
+        Path tempFolderPath = Files.createTempDirectory(prefix, PosixFilePermissions.asFileAttribute(permissions));
+
+        logger.debug("Created temporary directory at path {}", tempFolderPath);
+
+        return tempFolderPath.toString();
+    }
+
+    protected static void decryptModel(String password, String inputFilePath, String outputFilePath)
+            throws IOException, KuraIOException {
+        if (Security.getProvider("BC") == null) {
+            Security.addProvider(new BouncyCastleProvider());
+        }
+
+        if (!Files.isRegularFile(Paths.get(inputFilePath))) {
+            throw new IOException("Input file " + inputFilePath + " does not exist/is not a file");
+        }
+
+        if (Files.exists(Paths.get(outputFilePath))) {
+            throw new IOException("Output file " + outputFilePath + " already exists");
+        }
+
+        try {
+            decryptFile(password, inputFilePath, outputFilePath);
+        } catch (PGPException e) {
+            throw new KuraIOException(e, "File decryption failed.");
+        }
+    }
+
+    protected static void unzipModel(String inputFilePath, String outputFolder) throws IOException {
+        if (!Files.isRegularFile(Paths.get(inputFilePath))) {
+            throw new IOException("Input file " + inputFilePath + " does not exist/is not a file");
+        }
+
+        if (!isZipCompressed(inputFilePath)) {
+            throw new IOException("ZIP magic number check failed. Wrong file format");
+        }
+
+        if (!Files.isDirectory(Paths.get(outputFolder))) {
+            throw new IOException("Output folder " + outputFolder + " does not exist/is not a folder");
+        }
+
+        UnZip.unZipFile(inputFilePath, outputFolder);
+    }
+
+    private static boolean isZipCompressed(String filePath) throws IOException {
+        byte b1 = 0;
+        byte b2 = 0;
+
+        try (InputStream is = new FileInputStream(filePath)) {
+            b1 = (byte) is.read();
+            b2 = (byte) is.read();
+        } catch (IOException e) {
+            throw new IOException(e);
+        }
+
+        return b1 == 0x50 && b2 == 0x4B;
+    }
+
+    protected static void cleanRepository(String modelRootPath) {
+        if (!Files.isDirectory(Paths.get(modelRootPath))) {
+            logger.warn("Model root folder {} does not exist", modelRootPath);
+            return;
+        }
+
+        try {
+            FileUtils.cleanDirectory(new File(modelRootPath));
+        } catch (IOException e) {
+            logger.warn("Cannot clean directory at path {}", modelRootPath, e);
+        }
+    }
+
+    private static void decryptFile(String password, String inputFilePath, String outputFilePath)
+            throws IOException, PGPException {
+        InputStream inStream = new BufferedInputStream(new FileInputStream(inputFilePath));
+
+        inStream = PGPUtil.getDecoderStream(inStream);
+
+        JcaPGPObjectFactory pgpFactory = new JcaPGPObjectFactory(inStream);
+        Object currentObj = pgpFactory.nextObject();
+
+        PGPEncryptedDataList enc;
+        if (currentObj instanceof PGPEncryptedDataList) {
+            enc = (PGPEncryptedDataList) currentObj;
+        } else {
+            enc = (PGPEncryptedDataList) pgpFactory.nextObject();
+        }
+
+        if (enc == null) {
+            inStream.close();
+            throw new IOException("PGP PBE File format read failed");
+        }
+
+        PGPPBEEncryptedData pbe = (PGPPBEEncryptedData) enc.get(0);
+        InputStream clear = pbe.getDataStream(new JcePBEDataDecryptorFactoryBuilder(
+                new JcaPGPDigestCalculatorProviderBuilder().setProvider("BC").build()).setProvider("BC")
+                        .build(password.toCharArray()));
+
+        pgpFactory = new JcaPGPObjectFactory(clear);
+        currentObj = pgpFactory.nextObject();
+
+        if (currentObj instanceof PGPCompressedData) {
+            PGPCompressedData compressedData = (PGPCompressedData) currentObj;
+            pgpFactory = new JcaPGPObjectFactory(compressedData.getDataStream());
+            currentObj = pgpFactory.nextObject();
+        }
+
+        PGPLiteralData literalData = (PGPLiteralData) currentObj;
+        InputStream decryptedStream = literalData.getInputStream();
+
+        OutputStream outStream = new FileOutputStream(outputFilePath);
+        Streams.pipeAll(decryptedStream, outStream);
+        outStream.close();
+        inStream.close();
+
+        if (pbe.isIntegrityProtected()) {
+            if (!pbe.verify()) {
+                logger.error("File failed integrity check");
+            } else {
+                logger.info("File integrity check passed");
+            }
+        } else {
+            logger.info("No file integrity check");
+        }
+    }
+}
diff --git a/...iton.server/src/main/java/org/eclipse/kura/ai/triton/server/TritonServerLocalManager.java b/...iton.server/src/main/java/org/eclipse/kura/ai/triton/server/TritonServerLocalManager.java
@@ -36,6 +36,7 @@ public class TritonServerLocalManager {
     private static final int MONITOR_PERIOD = 30;
     private static final String[] TRITONSERVER = new String[] { "tritonserver" };
 
+    private final String decryptionFolderPath;
     private final CommandExecutorService commandExecutorService;
     private final TritonServerServiceOptions options;
     private Command serverCommand;
@@ -44,9 +45,10 @@ public class TritonServerLocalManager {
     private ScheduledFuture<?> scheduledFuture;
 
     protected TritonServerLocalManager(TritonServerServiceOptions options,
-            CommandExecutorService commandExecutorService) {
+            CommandExecutorService commandExecutorService, String decryptionFolderPath) {
         this.options = options;
         this.commandExecutorService = commandExecutorService;
+        this.decryptionFolderPath = decryptionFolderPath;
         this.scheduledExecutorService = Executors.newSingleThreadScheduledExecutor();
     }
 
@@ -137,7 +139,11 @@ protected static void sleepFor(long timeout) {
     private Command createServerCommand() {
         List<String> commandString = new ArrayList<>();
         commandString.add("tritonserver");
-        commandString.add("--model-repository=" + this.options.getModelRepositoryPath());
+        if (this.options.modelsAreEncrypted()) {
+            commandString.add("--model-repository=" + this.decryptionFolderPath);
+        } else {
+            commandString.add("--model-repository=" + this.options.getModelRepositoryPath());
+        }
         commandString.add("--backend-directory=" + this.options.getBackendsPath());
         if (!this.options.getBackendsConfigs().isEmpty()) {
             this.options.getBackendsConfigs().forEach(config -> commandString.add("--backend-config=" + config));