From 9b516eb7d762e5cbcecb0f016b4d001f3b2d19a9 Mon Sep 17 00:00:00 2001 From: Krrish Date: Sat, 9 May 2026 01:04:23 +0530 Subject: [PATCH] test: make PicoD execute tests portable Signed-off-by: Krrish --- pkg/picod/execute_test.go | 32 +++--- pkg/picod/picod_test.go | 23 +++-- pkg/picod/test_helpers_test.go | 171 +++++++++++++++++++++++++++++++++ 3 files changed, 202 insertions(+), 24 deletions(-) create mode 100644 pkg/picod/test_helpers_test.go diff --git a/pkg/picod/execute_test.go b/pkg/picod/execute_test.go index d6cc4bdc..2c7c8408 100644 --- a/pkg/picod/execute_test.go +++ b/pkg/picod/execute_test.go @@ -41,6 +41,8 @@ func init() { } func setupExecuteTestServer(t *testing.T) (*Server, string) { + t.Setenv(testHelperProcessEnv, "1") + // Generate RSA key pair privateKey, err := rsa.GenerateKey(rand.Reader, 2048) require.NoError(t, err) @@ -169,7 +171,7 @@ func TestExecuteHandler_TimeoutFormats(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { req := ExecuteRequest{ - Command: []string{"echo", "test"}, + Command: testCommand("echo", "test"), Timeout: tt.timeout, } body, _ := json.Marshal(req) @@ -234,7 +236,7 @@ func TestExecuteHandler_WorkingDirectory(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { req := ExecuteRequest{ - Command: []string{"pwd"}, + Command: testCommand("pwd"), WorkingDir: tt.workingDir, } body, _ := json.Marshal(req) @@ -264,11 +266,11 @@ func TestExecuteHandler_WorkingDirectory_SymlinkEscape(t *testing.T) { // Plant a symlink inside the workspace that points outside it. outsideDir := t.TempDir() symlinkPath := filepath.Join(tmpDir, "escape") - require.NoError(t, os.Symlink(outsideDir, symlinkPath)) + requireSymlink(t, outsideDir, symlinkPath) // Attempt to execute in a subdirectory of the symlink — mkdirSafe should reject this. req := ExecuteRequest{ - Command: []string{"pwd"}, + Command: testCommand("pwd"), WorkingDir: "escape/newdir", } body, _ := json.Marshal(req) @@ -291,7 +293,7 @@ func TestExecuteHandler_DefaultsToWorkspace(t *testing.T) { // No WorkingDir set — command should run inside the workspace directory. req := ExecuteRequest{ - Command: []string{"pwd"}, + Command: testCommand("pwd"), } body, _ := json.Marshal(req) @@ -331,17 +333,17 @@ func TestExecuteHandler_ExitCodes(t *testing.T) { }{ { name: "success (exit 0)", - command: []string{"true"}, + command: testCommand("exit", "0"), expectedCode: 0, }, { name: "failure (exit 1)", - command: []string{"false"}, + command: testCommand("exit", "1"), expectedCode: 1, }, { name: "custom exit code 42", - command: []string{"sh", "-c", "exit 42"}, + command: testCommand("exit", "42"), expectedCode: 42, }, } @@ -374,7 +376,7 @@ func TestExecuteHandler_TimeoutHandling(t *testing.T) { defer os.Unsetenv(PublicKeyEnvVar) req := ExecuteRequest{ - Command: []string{"sleep", "1"}, + Command: testCommand("sleep", "1s"), Timeout: "100ms", } body, _ := json.Marshal(req) @@ -401,7 +403,7 @@ func TestExecuteHandler_EnvironmentVariables(t *testing.T) { defer os.Unsetenv(PublicKeyEnvVar) req := ExecuteRequest{ - Command: []string{"sh", "-c", "echo $TEST_VAR"}, + Command: testCommand("env", "TEST_VAR"), Env: map[string]string{ "TEST_VAR": "test-value", }, @@ -429,7 +431,7 @@ func TestExecuteHandler_ResponseStructure(t *testing.T) { defer os.Unsetenv(PublicKeyEnvVar) req := ExecuteRequest{ - Command: []string{"echo", "hello", "world"}, + Command: testCommand("echo", "hello", "world"), } body, _ := json.Marshal(req) @@ -465,7 +467,7 @@ func TestExecuteHandler_StderrCapture(t *testing.T) { defer os.Unsetenv(PublicKeyEnvVar) req := ExecuteRequest{ - Command: []string{"sh", "-c", "echo 'error message' >&2"}, + Command: testCommand("stderr-exit", "error message", "0"), } body, _ := json.Marshal(req) @@ -490,7 +492,7 @@ func TestExecuteHandler_CommandWithArguments(t *testing.T) { defer os.Unsetenv(PublicKeyEnvVar) req := ExecuteRequest{ - Command: []string{"sh", "-c", "echo arg1 arg2 arg3"}, + Command: testCommand("echo", "arg1", "arg2", "arg3"), } body, _ := json.Marshal(req) @@ -515,7 +517,7 @@ func TestExecuteHandler_EmptyEnvVars(t *testing.T) { defer os.Unsetenv(PublicKeyEnvVar) req := ExecuteRequest{ - Command: []string{"echo", "test"}, + Command: testCommand("echo", "test"), Env: map[string]string{}, } body, _ := json.Marshal(req) @@ -536,7 +538,7 @@ func TestExecuteHandler_MultipleEnvVars(t *testing.T) { defer os.Unsetenv(PublicKeyEnvVar) req := ExecuteRequest{ - Command: []string{"sh", "-c", "echo $VAR1 $VAR2 $VAR3"}, + Command: testCommand("echo-env", "VAR1", "VAR2", "VAR3"), Env: map[string]string{ "VAR1": "value1", "VAR2": "value2", diff --git a/pkg/picod/picod_test.go b/pkg/picod/picod_test.go index ea116d1b..793660b2 100644 --- a/pkg/picod/picod_test.go +++ b/pkg/picod/picod_test.go @@ -64,6 +64,8 @@ func createToken(t *testing.T, key *rsa.PrivateKey, claims jwt.MapClaims) string // setupTestServer creates a test server with public key loaded from env func setupTestServer(t *testing.T, pubPEM string) (*Server, *httptest.Server, string) { + t.Setenv(testHelperProcessEnv, "1") + tmpDir, err := os.MkdirTemp("", "picod_test") require.NoError(t, err) @@ -113,7 +115,9 @@ func TestPicoD_EndToEnd(t *testing.T) { t.Run("Unauthenticated Access", func(t *testing.T) { // Execute without auth header - execReq := ExecuteRequest{Command: []string{"echo", "hello"}} + execReq := ExecuteRequest{ + Command: testCommand("echo", "hello"), + } body, _ := json.Marshal(execReq) req, _ := http.NewRequest("POST", ts.URL+"/api/execute", bytes.NewBuffer(body)) resp, err := client.Do(req) @@ -152,29 +156,29 @@ func TestPicoD_EndToEnd(t *testing.T) { } // 1. Basic Execution - resp := doExec([]string{"echo", "hello"}, nil, "") + resp := doExec(testCommand("echo", "hello"), nil, "") assert.Equal(t, "hello\n", resp.Stdout) assert.Equal(t, 0, resp.ExitCode) assert.False(t, resp.StartTime.IsZero()) assert.False(t, resp.EndTime.IsZero()) // 2. Environment Variables - resp = doExec([]string{"sh", "-c", "echo $TEST_VAR"}, map[string]string{"TEST_VAR": "picod_env"}, "") + resp = doExec(testCommand("env", "TEST_VAR"), map[string]string{"TEST_VAR": "picod_env"}, "") assert.Equal(t, "picod_env\n", resp.Stdout) // 3. Stderr and Exit Code - resp = doExec([]string{"sh", "-c", "echo error_msg >&2; exit 1"}, nil, "") + resp = doExec(testCommand("stderr-exit", "error_msg", "1"), nil, "") assert.Equal(t, "error_msg\n", resp.Stderr) assert.Equal(t, 1, resp.ExitCode) // 4. Timeout - resp = doExec([]string{"sleep", "2"}, nil, "0.5s") + resp = doExec(testCommand("sleep", "2s"), nil, "0.5s") assert.Equal(t, 124, resp.ExitCode) assert.Contains(t, resp.Stderr, "Command timed out") // 5. Working Directory Escape (Should Fail) escapeReq := ExecuteRequest{ - Command: []string{"ls"}, + Command: testCommand("echo", "escape"), WorkingDir: "../", } escapeBody, _ := json.Marshal(escapeReq) @@ -311,7 +315,9 @@ func TestPicoD_EndToEnd(t *testing.T) { } token := createToken(t, wrongPriv, claims) - reqBody := ExecuteRequest{Command: []string{"echo", "malicious"}} + reqBody := ExecuteRequest{ + Command: testCommand("echo", "malicious"), + } realBody, _ := json.Marshal(reqBody) req, _ := http.NewRequest("POST", ts.URL+"/api/execute", bytes.NewBuffer(realBody)) @@ -376,8 +382,7 @@ func TestPicoD_SetWorkspace(t *testing.T) { // Create a symlink linkDir := filepath.Join(tmpDir, "link") - err = os.Symlink(realDir, linkDir) - require.NoError(t, err) + requireSymlink(t, realDir, linkDir) server := &Server{} diff --git a/pkg/picod/test_helpers_test.go b/pkg/picod/test_helpers_test.go new file mode 100644 index 00000000..8c6c8191 --- /dev/null +++ b/pkg/picod/test_helpers_test.go @@ -0,0 +1,171 @@ +/* +Copyright The Volcano Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package picod + +import ( + "errors" + "fmt" + "os" + "runtime" + "strconv" + "strings" + "syscall" + "testing" + "time" +) + +const testHelperProcessEnv = "GO_WANT_PICOD_HELPER_PROCESS" + +func init() { + if os.Getenv(testHelperProcessEnv) == "1" { + picodHelperProcessMain() + os.Exit(0) + } +} + +func testCommand(args ...string) []string { + exe, err := os.Executable() + if err != nil { + exe = os.Args[0] + } + command := []string{exe, "--"} + return append(command, args...) +} + +func requireSymlink(t *testing.T, oldname, newname string) { + t.Helper() + if err := os.Symlink(oldname, newname); err != nil { + if isWindowsSymlinkPrivilegeError(err) { + t.Skipf("creating symlinks requires privileges on Windows: %v", err) + } + t.Fatalf("create symlink: %v", err) + } +} + +func isWindowsSymlinkPrivilegeError(err error) bool { + if runtime.GOOS != "windows" { + return false + } + var linkErr *os.LinkError + if errors.As(err, &linkErr) { + var errno syscall.Errno + if errors.As(linkErr.Err, &errno) { + return errno == syscall.Errno(1314) + } + } + return false +} + +func picodHelperProcessArgs() []string { + args := os.Args + for len(args) > 0 && args[0] != "--" { + args = args[1:] + } + if len(args) == 0 { + fmt.Fprintln(os.Stderr, "missing helper separator") + os.Exit(2) + } + args = args[1:] + if len(args) == 0 { + fmt.Fprintln(os.Stderr, "missing helper command") + os.Exit(2) + } + return args +} + +func picodHelperProcessMain() { + args := picodHelperProcessArgs() + switch args[0] { + case "echo": + fmt.Println(strings.Join(args[1:], " ")) + case "env": + helperPrintEnv(args) + case "echo-env": + helperEchoEnv(args) + case "exit": + os.Exit(helperExitCode(args)) + case "pwd": + helperPrintWorkingDirectory() + case "sleep": + helperSleep(args) + case "stderr-exit": + helperStderrExit(args) + default: + fmt.Fprintf(os.Stderr, "unknown helper command: %s\n", args[0]) + os.Exit(2) + } +} + +func helperPrintEnv(args []string) { + if len(args) != 2 { + fmt.Fprintln(os.Stderr, "env helper requires variable name") + os.Exit(2) + } + fmt.Println(os.Getenv(args[1])) +} + +func helperEchoEnv(args []string) { + values := make([]string, 0, len(args)-1) + for _, key := range args[1:] { + values = append(values, os.Getenv(key)) + } + fmt.Println(strings.Join(values, " ")) +} + +func helperExitCode(args []string) int { + if len(args) != 2 { + fmt.Fprintln(os.Stderr, "exit helper requires code") + os.Exit(2) + } + code, err := strconv.Atoi(args[1]) + if err != nil { + fmt.Fprintf(os.Stderr, "invalid exit code: %v\n", err) + os.Exit(2) + } + return code +} + +func helperPrintWorkingDirectory() { + wd, err := os.Getwd() + if err != nil { + fmt.Fprintf(os.Stderr, "getwd: %v\n", err) + os.Exit(2) + } + fmt.Println(wd) +} + +func helperSleep(args []string) { + if len(args) != 2 { + fmt.Fprintln(os.Stderr, "sleep helper requires duration") + os.Exit(2) + } + duration, err := time.ParseDuration(args[1]) + if err != nil { + fmt.Fprintf(os.Stderr, "invalid duration: %v\n", err) + os.Exit(2) + } + time.Sleep(duration) +} + +func helperStderrExit(args []string) { + if len(args) != 3 { + fmt.Fprintln(os.Stderr, "stderr-exit helper requires message and code") + os.Exit(2) + } + fmt.Fprintln(os.Stderr, args[1]) + os.Exit(helperExitCode([]string{"exit", args[2]})) +}