From e12df2b1649f99fa5be097b0a33a42957add4a78 Mon Sep 17 00:00:00 2001
From: SparkSnail <shinaiyang@pku.edu.cn>
Date: Fri, 21 Jun 2019 10:09:48 +0700
Subject: [PATCH] Validate file name in codeDir for PAI platform (#1168)

---
 src/nni_manager/common/utils.ts               | 36 ++++++++++++++++++-
 .../training_service/common/util.ts           | 15 ++++++--
 2 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/src/nni_manager/common/utils.ts b/src/nni_manager/common/utils.ts
index cb5770026f..3c11c0c835 100644
--- a/src/nni_manager/common/utils.ts
+++ b/src/nni_manager/common/utils.ts
@@ -374,6 +374,40 @@ function countFilesRecursively(directory: string, timeoutMilliSeconds?: number):
     });
 }
 
+function validateFileName(fileName: string): boolean {
+    let pattern: string = '^[a-z0-9A-Z\.-_]+$';
+    const validateResult = fileName.match(pattern);
+    if(validateResult) {
+        return true;
+    }
+    return false;
+}
+
+async function validateFileNameRecursively(directory: string): Promise<boolean> {
+    if(!fs.existsSync(directory)) {
+        throw Error(`Direcotory ${directory} doesn't exist`);
+    }
+
+    const fileNameArray: string[] = fs.readdirSync(directory);
+    let result = true;
+    for(var name of fileNameArray){
+        const fullFilePath: string = path.join(directory, name);
+        try {
+            // validate file names and directory names
+            result = validateFileName(name);
+            if (fs.lstatSync(fullFilePath).isDirectory()) {
+                result = result && await validateFileNameRecursively(fullFilePath);
+            }
+            if(!result) {
+                return Promise.reject(new Error(`file name in ${fullFilePath} is not valid!`));
+            }
+        } catch(error) {
+            return Promise.reject(error);
+        }
+    }
+    return Promise.resolve(result);   
+}
+
 /**
  * get the version of current package
  */
@@ -474,6 +508,6 @@ function unixPathJoin(...paths: any[]): string {
     return dir;
 }
 
-export {countFilesRecursively, getRemoteTmpDir, generateParamFileName, getMsgDispatcherCommand, getCheckpointDir,
+export {countFilesRecursively, validateFileNameRecursively, getRemoteTmpDir, generateParamFileName, getMsgDispatcherCommand, getCheckpointDir,
     getLogDir, getExperimentRootDir, getJobCancelStatus, getDefaultDatabaseDir, getIPV4Address, unixPathJoin,
     mkDirP, delay, prepareUnitTest, parseArg, cleanupUnitTest, uniqueString, randomSelect, getLogLevel, getVersion, getCmdPy, getTunerProc, isAlive, killPid, getNewLine };
diff --git a/src/nni_manager/training_service/common/util.ts b/src/nni_manager/training_service/common/util.ts
index 71c255e848..aab276bd30 100644
--- a/src/nni_manager/training_service/common/util.ts
+++ b/src/nni_manager/training_service/common/util.ts
@@ -25,7 +25,7 @@ import * as fs from 'fs';
 import * as os from 'os';
 import * as path from 'path';
 import { String } from 'typescript-string-operations';
-import { countFilesRecursively, getNewLine } from '../../common/utils';
+import { countFilesRecursively, getNewLine, validateFileNameRecursively } from '../../common/utils';
 import { file } from '../../node_modules/@types/tmp';
 import { GPU_INFO_COLLECTOR_FORMAT_LINUX, GPU_INFO_COLLECTOR_FORMAT_WINDOWS } from './gpuData';
 
@@ -38,22 +38,33 @@ import { GPU_INFO_COLLECTOR_FORMAT_LINUX, GPU_INFO_COLLECTOR_FORMAT_WINDOWS } fr
 // tslint:disable: no-redundant-jsdoc
 export async function validateCodeDir(codeDir: string) : Promise<number> {
     let fileCount: number | undefined;
-
+    let fileNameValid: boolean = true;
     try {
         fileCount = await countFilesRecursively(codeDir);
     } catch (error) {
         throw new Error(`Call count file error: ${error}`);
     }
+    try {
+        fileNameValid = await validateFileNameRecursively(codeDir);
+    } catch(error) {
+        throw new Error(`Validate file name error: ${error}`);
+    }
 
     if (fileCount !== undefined && fileCount > 1000) {
         const errMessage: string = `Too many files(${fileCount} found}) in ${codeDir},`
                                     + ` please check if it's a valid code dir`;
         throw new Error(errMessage);
     }
+    
+    if(!fileNameValid) {
+        const errMessage: string = `File name in ${codeDir} is not valid, please check file names, only support digit number、alphabet and (.-_) in file name.`; 
+        throw new Error(errMessage);    
+    }
 
     return fileCount;
 }
 
+
 /**
  * crete a new directory
  * @param directory