From 78d6705f84dc4a5a7da2c629a4e942eb1dc4d86e Mon Sep 17 00:00:00 2001
From: ArtoLord <32140169+ArtoLord@users.noreply.github.com>
Date: Tue, 12 Dec 2023 14:31:24 +0300
Subject: [PATCH] Validate image (#1157)

* Validate image

* Fix
---
 .../java/ai/lzy/env/base/DockerEnvironment.java    | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

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
index f3d6dc9e7..d12578e5f 100644
--- 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);
+            }
+        }
+    }
 }