Validate image (#1157)

* Validate image * Fix
lambdazy · Dec 12, 2023 · 78d6705 · 78d6705
1 parent 9897fee
commit 78d6705
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java b/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java
@@ -69,6 +69,7 @@ public void install(LogStream outStream, LogStream errStream) throws Environment
         String sourceImage = config.image();
         try {
             prepareImage(sourceImage, outStream);
+            validateImage(sourceImage);
         } catch (InterruptedException e) {
             LOG.error("Image pulling was interrupted");
             errStream.log("Image pulling was interrupted");
@@ -291,4 +292,17 @@ private void prepareImage(String image, LogStream out) throws Exception {
         LOG.info(msg);
         out.log(msg);
     }
+
+    private void validateImage(String image) throws Exception {
+        var inspectResp = client.inspectImageCmd(image).exec();
+        var config = inspectResp.getConfig();
+
+        if (config != null) {
+            var user = config.getUser();
+            if (user != null && !user.isEmpty() && !user.equals("root")) {
+                LOG.error("Got custom user {} in image {}", user, image);
+                throw new Exception("Cannot use docker image " + image + "with custom user " + user);
+            }
+        }
+    }
 }