RestServer - add image regex to limit the valid image (#107)

remove force acr image and add image regex to limit the valid image

Author: yukirora

src/rest-server/deploy/rest-server.yaml.template

@@ -174,6 +174,8 @@ spec:
           value: {{ cluster_cfg['rest-server']['hived-computing-device-envs'] }}
         - name: ALERT_MANAGER_URL
           value: "{{ cluster_cfg['alert-manager']['url'] }}"
+        - name: IMAGE_REGEX
+          value: "{{ cluster_cfg['rest-server']['image-regex'] | default("") }}"
         ports:
         - name: rest-server
           containerPort: 8080

src/rest-server/src/config/launcher.js

@@ -84,6 +84,7 @@ const k8sLauncherConfigSchema = Joi.object()
     jobRestrictionGitScriptName: Joi.string().required(),
     clstoreHostPath: Joi.string().empty(''),
     clstoreJobPath: Joi.string().empty(''),
+    imageRegex: Joi.string().empty(''),
   })
   .required();
 
@@ -156,6 +157,7 @@ if (launcherType === 'k8s') {
     jobRestrictionGitScriptName: process.env.JOB_RESTRICTION_GIT_SCRIPT_NAME || 'unset',
     clstoreHostPath: process.env.CLUSTER_LOCAL_STORAGE_HOST_PATH,
     clstoreJobPath: process.env.CLUSTER_LOCAL_STORAGE_JOB_PATH,
+    imageRegex: process.env.IMAGE_REGEX || '',
   };
 
   const { error, value } = k8sLauncherConfigSchema.validate(launcherConfig);

src/rest-server/src/middlewares/v2/protocol.js

@@ -24,6 +24,8 @@ const hived = require('@pai/middlewares/v2/hived');
 const { enabledHived } = require('@pai/config/launcher');
 const protocolSchema = require('@pai/config/v2/protocol');
 const asyncHandler = require('@pai/middlewares/v2/asyncHandler');
+const logger = require('@pai/config/logger');
+const launcherConfig = require('@pai/config/launcher'); 
 
 const mustacheWriter = new mustache.Writer();
 
@@ -53,6 +55,19 @@ const render = (template, dict, tags = ['<%', '%>']) => {
   return result.trim();
 };
 
+const getImageName = (prerequisites) => { 
+  if (
+    typeof prerequisites !== 'object' ||
+    typeof prerequisites.dockerimage !== 'object'
+  ) {
+    return null;
+  }
+
+  const dockerImages = Object.values(prerequisites.dockerimage);
+  const img = dockerImages.find(p => p.type === 'dockerimage' && p.uri);
+  return img?.uri ?? null;
+};
+
 const protocolValidate = (protocolYAML) => {
   const protocolObj = yaml.load(protocolYAML);
   if (!protocolSchema.validate(protocolObj)) {
@@ -87,6 +102,42 @@ const protocolValidate = (protocolYAML) => {
         prerequisiteSet.add(item.name);
       }
     }
+    const imageRegexPattern = launcherConfig.imageRegex;
+    let imageRegex;
+
+    try {
+      if (imageRegexPattern && imageRegexPattern.length > 0) {
+        imageRegex = new RegExp(imageRegexPattern);
+      }
+    } catch (error) {
+      logger.info(`Invalid imageRegex pattern: ${imageRegexPattern}. Error: ${error.message}`);
+      throw createError(
+          'Internal Server Error',
+          'InvalidImageRegexError',
+          `The provided imageRegex pattern "${imageRegexPattern}" is invalid.`
+      );
+    }
+
+    if (imageRegex) {
+      const imageName = getImageName(protocolObj.prerequisites);
+      // Check if the imageName matches the imageRegex
+      if (!imageName) {
+        throw createError(
+            'Bad Request',
+            'NoDockerImageError',
+            'No valid docker image found in prerequisites.'
+        );
+      }
+      // Check if the imageName matches the imageRegex
+      const match = imageRegex.test(imageName);
+      if (!match) {
+        throw createError(
+            'Bad Request',
+            'InvalidImageError',
+            `The image ${imageName} is not allowed.`
+        );
+      }
+    }
   }
   protocolObj.prerequisites = prerequisites;
   // convert deployments list to dict

src/rest-server/src/middlewares/v2/quota.js

@@ -75,19 +75,6 @@ const getReqeustedGpuCount = (protocol) => {
   return gpuCount;
 };
 
-function getImageName(prerequisites) {
-  if (
-    typeof prerequisites !== 'object' ||
-    typeof prerequisites.dockerimage !== 'object'
-  ) {
-    return null;
-  }
-
-  const dockerImages = Object.values(prerequisites.dockerimage);
-  const img = dockerImages.find(p => p.type === 'dockerimage' && p.uri);
-  return img?.uri ?? null;
-}
-
 const getJobPriority = (protocol) => {
   return protocol.extras?.hivedScheduler?.jobPriorityClass || null;
 };
@@ -104,21 +91,6 @@ const check = async (req, res, next) => {
     const userPrioritySet = userInfo.extension?.jobPriority ?? null;
     const userPriorityExpiration = userInfo.extension?.jobExpiration ?? null;
 
-    const imageName = getImageName(jobProtocol.prerequisites);
-    const imageRepo = imageName.split('/')[0].trim().toLowerCase();
-    const isAzureCR = /\.azurecr\.io$/i.test(imageRepo);
-    // Reject if it’s not ACR
-    if (!isAzureCR) {
-      return next(
-        createError(
-          'Forbidden',
-          'InvalidImageError',
-          `The image ${imageRepo || imageName} is not allowed. ` +
-          `Please use an image from Azure Container Registry (ACR).`
-        )
-      );
-    }
-
     if (jobPriority === 'prod') {
       if (userPrioritySet !== 1) {
         logger.debug(

src/rest-server/src/utils/error.d.ts

@@ -33,8 +33,11 @@ declare type Code =
     'ForbiddenUserError' |
     'ForbiddenKeyError' |
     'IncorrectPasswordError' |
+    'InvalidImageError' |
+    'InvalidImageRegexError' |
     'InvalidParametersError' |
     'NoApiError' |
+    'NoDockerImageError' |
     'NoJobError' |
     'NoJobConfigError' |
     'NoJobSshInfoError' |