diff --git a/analysis/plugin/configurator/script/script.py b/analysis/plugin/configurator/script/script.py index 273c2d24e4804c3a4acaac8083e43ace5e642f06..e71e85c12a507e8378d32a75b441c1af164ad7e3 100755 --- a/analysis/plugin/configurator/script/script.py +++ b/analysis/plugin/configurator/script/script.py @@ -31,11 +31,15 @@ class Script(Configurator): _module = "SCRIPT" _submod = "SCRIPT" cmd_delimiter = "|" + scripts_path = "/usr/libexec/atuned/scripts" def __init__(self, user=None): Configurator.__init__(self, user) def _set(self, key, value): + script_dir = os.path.dirname(key) + if self.scripts_path != os.path.realpath(script_dir): + raise SetConfigError("key:{} is invalid".format(key)) name = os.path.basename(key) script = "{}/set.sh".format(key) if not os.path.exists(script): @@ -48,6 +52,9 @@ class Script(Configurator): return 0 def _get(self, key, value): + script_dir = os.path.dirname(key) + if self.scripts_path != os.path.realpath(script_dir): + raise GetConfigError("key:{} is invalid".format(key)) name = os.path.basename(key) script = "{}/get.sh".format(key) if not os.path.exists(script): diff --git a/modules/client/profile/profile_analysis.go b/modules/client/profile/profile_analysis.go index b049c50e6327b87bb1f40c4b68052c7be6acb845..f55ccb443c517d001abd9cbd837767468153eacb 100644 --- a/modules/client/profile/profile_analysis.go +++ b/modules/client/profile/profile_analysis.go @@ -72,6 +72,9 @@ func profileAnalysis(ctx *cli.Context) error { } if ctx.NArg() == 1 { appname = ctx.Args().Get(0) + if !utils.IsInputStringValid(appname) { + return fmt.Errorf("input:%s is invalid", appname) + } } c, err := client.NewClientFromContext(ctx) @@ -81,6 +84,10 @@ func profileAnalysis(ctx *cli.Context) error { defer c.Close() modelFile := ctx.String("model") + if modelFile != "" && !utils.IsInputStringValid(modelFile) { + return fmt.Errorf("input:%s is invalid", modelFile) + } + svc := PB.NewProfileMgrClient(c.Connection()) stream, _ := svc.Analysis(CTX.Background(), &PB.AnalysisMessage{Name: appname, Model: modelFile}) diff --git a/modules/client/profile/profile_define.go b/modules/client/profile/profile_define.go index 8c7788ba4a57d5a9653d493436b32fb2e003e767..d304730603af3c352d0d02653d2969758c8bb1a8 100644 --- a/modules/client/profile/profile_define.go +++ b/modules/client/profile/profile_define.go @@ -64,6 +64,10 @@ func profileDefineCheck(ctx *cli.Context) error { } file := ctx.Args().Get(2) + if !utils.IsInputStringValid(file) { + return fmt.Errorf("input:%s is invalid", file) + } + exist, err := utils.PathExist(file) if err != nil { return err @@ -85,7 +89,14 @@ func profileDefined(ctx *cli.Context) error { return err } workloadType := ctx.Args().Get(0) + if !utils.IsInputStringValid(workloadType) { + return fmt.Errorf("input:%s is invalid", workloadType) + } + profileName := ctx.Args().Get(1) + if !utils.IsInputStringValid(profileName) { + return fmt.Errorf("input:%s is invalid", profileName) + } data, err := ioutil.ReadFile(ctx.Args().Get(2)) if err != nil { diff --git a/modules/client/profile/profile_train.go b/modules/client/profile/profile_train.go index 338b62c3e28bafa7d7ce8a2309b4cdb4b2982bae..49d490af1a0227ab346136cad46a0c8e519713b5 100644 --- a/modules/client/profile/profile_train.go +++ b/modules/client/profile/profile_train.go @@ -77,12 +77,18 @@ func checkTrainCtx(ctx *cli.Context) error { _ = cli.ShowCommandHelp(ctx, "train") return fmt.Errorf("error: data_path must be specified") } + if !utils.IsInputStringValid(dataPath) { + return fmt.Errorf("input:%s is invalid", dataPath) + } outputPath := ctx.String("output_file") if outputPath == "" { _ = cli.ShowCommandHelp(ctx, "train") return fmt.Errorf("error: output_file must be specified") } + if !utils.IsInputStringValid(outputPath) { + return fmt.Errorf("input:%s is invalid", outputPath) + } return nil } diff --git a/modules/client/profile/profile_update.go b/modules/client/profile/profile_update.go index e9a78e3d88d90fa5457c8503f8ade5b5546bf766..abd10c4df9e31a7674b6288b8e1c67799a1b1e6e 100644 --- a/modules/client/profile/profile_update.go +++ b/modules/client/profile/profile_update.go @@ -61,6 +61,10 @@ func profileUpdateCheck(ctx *cli.Context) error { } file := ctx.Args().Get(2) + if !utils.IsInputStringValid(file) { + return fmt.Errorf("input:%s is invalid", file) + } + exist, err := utils.PathExist(file) if err != nil { return err @@ -77,7 +81,14 @@ func profileUpdate(ctx *cli.Context) error { return err } workloadType := ctx.Args().Get(0) + if !utils.IsInputStringValid(workloadType) { + return fmt.Errorf("input:%s is invalid", workloadType) + } + profileName := ctx.Args().Get(1) + if !utils.IsInputStringValid(profileName) { + return fmt.Errorf("input:%s is invalid", profileName) + } data, err := ioutil.ReadFile(ctx.Args().Get(2)) if err != nil { diff --git a/modules/client/profile/profile_upgrade.go b/modules/client/profile/profile_upgrade.go index 21cf6a794a2ab5238f5a6dc777ed7daeb18608ad..ab6dd24c1adb71af5f3a8f90eee51b4861d23999 100644 --- a/modules/client/profile/profile_upgrade.go +++ b/modules/client/profile/profile_upgrade.go @@ -60,6 +60,10 @@ func profileUpgrade(ctx *cli.Context) error { } dbPath := ctx.Args().Get(0) + if !utils.IsInputStringValid(dbPath) { + return fmt.Errorf("input:%s is invalid", dbPath) + } + exist, err := utils.PathExist(dbPath) if err != nil { return err diff --git a/modules/server/profile/profile.go b/modules/server/profile/profile.go index b8fbf4df79f6bd6bb3b79a2e277add1cd9b933d5..41860f895cf45e17b28fe5bdd85abaff6a32ab38 100644 --- a/modules/server/profile/profile.go +++ b/modules/server/profile/profile.go @@ -677,6 +677,14 @@ func (s *ProfileServer) ProfileRollback(profileInfo *PB.ProfileInfo, stream PB.P Collection method call collection script to collect system data. */ func (s *ProfileServer) Collection(message *PB.CollectFlag, stream PB.ProfileMgr_CollectionServer) error { + isLocalAddr, err := SVC.CheckRpcIsLocalAddr(stream.Context()) + if err != nil { + return err + } + if !isLocalAddr { + return fmt.Errorf("the collection command can not be remotely operated") + } + if valid := utils.IsInputStringValid(message.GetWorkload()); !valid { return fmt.Errorf("input:%s is invalid", message.GetWorkload()) } @@ -698,7 +706,7 @@ func (s *ProfileServer) Collection(message *PB.CollectFlag, stream PB.ProfileMgr } classApps := &sqlstore.GetClassApp{Class: message.GetType()} - err := sqlstore.GetClassApps(classApps) + err = sqlstore.GetClassApps(classApps) if err != nil { return err } @@ -785,6 +793,14 @@ func (s *ProfileServer) Collection(message *PB.CollectFlag, stream PB.ProfileMgr Training method train the collected data to generate the model */ func (s *ProfileServer) Training(message *PB.TrainMessage, stream PB.ProfileMgr_TrainingServer) error { + isLocalAddr, err := SVC.CheckRpcIsLocalAddr(stream.Context()) + if err != nil { + return err + } + if !isLocalAddr { + return fmt.Errorf("the train command can not be remotely operated") + } + DataPath := message.GetDataPath() OutputPath := message.GetOutputPath() diff --git a/tests/configurator/test_script.py b/tests/configurator/test_script.py index 17ce1715d15116b33bc321d5a4cbc9ad0491d93c..2dea133c664a25606c89a37c12554460c0268c8c 100644 --- a/tests/configurator/test_script.py +++ b/tests/configurator/test_script.py @@ -20,7 +20,7 @@ from analysis.plugin.configurator.script.script import Script class TestScript: """ test script""" user = "UT" - path = "scripts/hugepage" + path = "/usr/libexec/atuned/scripts/hugepage" def test_get_script_with_hugepage(self): """test get script result with hugepage"""