package supervisor import ( "context" "io" "net/http" "net/http/httptest" "os" "path/filepath" "strings" "sync" "testing" "time" ) func TestWithAllocatedConfigOverridesAuthoritativeChildFlags(t *testing.T) { command := []string{ "game-server", "--", "--allocated-mode", "--match-id=stale-match", "--server-id=stale-server", "--server-image-digest=sha256:stale", "--assignment-expiry-unix=1", "--region=EU", "--custom-flag=preserved", } expiry := time.Unix(1_900_000_000, 0).UTC() got := withAllocatedConfig(command, "match-live", "server-live", "sha256:live", expiry) want := []string{ "game-server", "--", "--allocated-mode", "--match-id=match-live", "--server-id=server-live", "--server-image-digest=sha256:live", "--assignment-expiry-unix=1900000000", "--region=EU", "--custom-flag=preserved", } if strings.Join(got, "\x00") != strings.Join(want, "\x00") { t.Fatalf("allocated command = %#v, want %#v", got, want) } if strings.Join(command, "\x00") == strings.Join(got, "\x00") { t.Fatal("withAllocatedConfig mutated the caller's command slice") } } func TestAllocatedStartInjectsDynamicEndpointAndCallsReadyAfterProbe(t *testing.T) { ready := false readyCalled := false server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/gameserver": _, _ = w.Write([]byte(`{"status":{"address":"203.0.113.9","ports":[{"name":"game","port":31001}]}}`)) case "/ready-probe": if ready { w.WriteHeader(http.StatusOK) } else { w.WriteHeader(http.StatusServiceUnavailable) } case "/ready": readyCalled = true w.WriteHeader(http.StatusOK) default: w.WriteHeader(http.StatusNotFound) } })) defer server.Close() ready = true path := filepath.Join(t.TempDir(), "env.txt") argsPath := filepath.Join(t.TempDir(), "args.txt") command := []string{"/bin/sh", "-c", "env > " + path + "; printf '%s' \"$@\" > " + argsPath, "shell"} s, err := New(Config{Command: command, SDKBaseURL: server.URL, ReadyURL: server.URL + "/ready-probe", Transport: "steam_sdr", ReadyTimeout: time.Second, PollInterval: time.Millisecond}) if err != nil { t.Fatal(err) } if err := s.Start(context.Background()); err != nil { t.Fatal(err) } if err := s.Wait(); err != nil { t.Fatal(err) } contents, err := os.ReadFile(path) if err != nil { t.Fatal(err) } if !strings.Contains(string(contents), "SDR_LISTEN_PORT=31001") || !strings.Contains(string(contents), "SDR_IP=203.0.113.9:31001") { t.Fatalf("dynamic endpoint not injected: %s", contents) } args, err := os.ReadFile(argsPath) if err != nil { t.Fatal(err) } if !strings.Contains(string(args), "--port=31001") { t.Fatalf("dynamic port argument not injected: %s", args) } if !readyCalled { t.Fatal("Agones Ready was called before process-ready probe") } } func TestAllocatedStartMaterializesWorkloadAuthenticatedRosterBeforeChild(t *testing.T) { rosterPath := filepath.Join(t.TempDir(), "join-roster.json") sdk := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/gameserver": _, _ = w.Write([]byte(`{"object_meta":{"annotations":{"cosmic-clash.io/match-id":"match-1","cosmic-clash.io/workload-token":"workload-token"}},"status":{"address":"127.0.0.1","ports":[{"name":"game","port":31001}]}}`)) case "/ready-probe", "/ready": w.WriteHeader(http.StatusOK) default: w.WriteHeader(http.StatusNotFound) } })) defer sdk.Close() controlPlane := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path == "/v1/servers/server-1/roster" { if r.Method != http.MethodGet || r.Header.Get("Authorization") != "Bearer workload-token" { w.WriteHeader(http.StatusUnauthorized) return } _, _ = w.Write([]byte(`[{"authorisation":{"player_id":"player-1","expires_at":"2030-01-01T00:00:00Z"},"signature":"sig"}]`)) return } if r.URL.Path == "/v1/servers/server-1/register" { w.WriteHeader(http.StatusNoContent) return } w.WriteHeader(http.StatusNotFound) })) defer controlPlane.Close() command := []string{"/bin/sh", "-c", "test -s '" + rosterPath + "'"} s, err := New(Config{ Command: command, SDKBaseURL: sdk.URL, ReadyURL: sdk.URL + "/ready-probe", ControlPlaneURL: controlPlane.URL, ServerID: "server-1", ProtocolVersion: 1, ImageDigest: "sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", RosterPath: rosterPath, ReadyTimeout: time.Second, PollInterval: time.Millisecond, }) if err != nil { t.Fatal(err) } if err := s.Start(context.Background()); err != nil { t.Fatal(err) } if err := s.Wait(); err != nil { t.Fatal(err) } contents, err := os.ReadFile(rosterPath) if err != nil || !strings.Contains(string(contents), "player-1") { t.Fatalf("materialized roster=%q err=%v", contents, err) } } func TestControlPlaneRegistrationRejectsIncompleteConfig(t *testing.T) { base := Config{Command: []string{"/bin/true"}, ControlPlaneURL: "https://control-plane.invalid"} if _, err := New(base); err == nil { t.Fatal("registration enabled with no server/protocol/digest was accepted") } complete := base complete.ServerID, complete.MatchID, complete.ProtocolVersion, complete.ImageDigest = "server-1", "match-1", 1, "sha256:aa" if _, err := New(complete); err != nil { t.Fatalf("fully configured registration rejected: %v", err) } // WorkloadTokenPath is deliberately not required at construction time -- // it can instead be resolved at Start time from the GameServer's own // cosmic-clash.io/workload-token annotation (see workloadToken() and // TestControlPlaneRegistrationFallsBackToGameServerAnnotationForWorkloadToken). withTokenPath := complete withTokenPath.WorkloadTokenPath = "/tmp/token" if _, err := New(withTokenPath); err != nil { t.Fatalf("configured token path rejected: %v", err) } } func TestControlPlaneRegistrationReportsProcessReadyThenAssignmentReady(t *testing.T) { var mu sync.Mutex var gotAuth, gotIdempotency string var bodies []string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch { case r.URL.Path == "/gameserver": _, _ = w.Write([]byte(`{"status":{"address":"203.0.113.9","ports":[{"name":"game","port":31001}]}}`)) case r.URL.Path == "/ready-probe": w.WriteHeader(http.StatusOK) case r.URL.Path == "/ready": w.WriteHeader(http.StatusOK) case r.URL.Path == "/v1/servers/server-1/register": mu.Lock() gotAuth = r.Header.Get("Authorization") gotIdempotency = r.Header.Get("Idempotency-Key") body, _ := io.ReadAll(r.Body) bodies = append(bodies, string(body)) mu.Unlock() w.WriteHeader(http.StatusNoContent) default: w.WriteHeader(http.StatusNotFound) } })) defer server.Close() tokenPath := filepath.Join(t.TempDir(), "token") if err := os.WriteFile(tokenPath, []byte(" workload-jwt-abc123 \n"), 0o600); err != nil { t.Fatal(err) } s, err := New(Config{ Command: []string{"/bin/sh", "-c", "true"}, SDKBaseURL: server.URL, ReadyURL: server.URL + "/ready-probe", ReadyTimeout: time.Second, PollInterval: time.Millisecond, ControlPlaneURL: server.URL, WorkloadTokenPath: tokenPath, ServerID: "server-1", MatchID: "match-1", ProtocolVersion: 1, ImageDigest: "sha256:aa", AssignmentReadyAttempts: 3, AssignmentReadyBackoff: time.Millisecond, }) if err != nil { t.Fatal(err) } if err := s.Start(context.Background()); err != nil { t.Fatal(err) } _ = s.Wait() if gotAuth != "Bearer workload-jwt-abc123" { t.Fatalf("Authorization header = %q, want the trimmed token file contents", gotAuth) } if len(gotIdempotency) < 16 { t.Fatalf("Idempotency-Key = %q, too short", gotIdempotency) } mu.Lock() defer mu.Unlock() if len(bodies) != 2 { t.Fatalf("expected exactly 2 register calls (process-ready, assignment-ready), got %d: %v", len(bodies), bodies) } if !strings.Contains(bodies[0], `"match_id":"match-1"`) || !strings.Contains(bodies[0], `"assignment_ready":false`) || !strings.Contains(bodies[0], `"image_digest":"sha256:aa"`) { t.Fatalf("process-ready register body = %s", bodies[0]) } if !strings.Contains(bodies[1], `"assignment_ready":true`) { t.Fatalf("assignment-ready register body = %s", bodies[1]) } } func TestAssignmentReadyRegistrationRetriesUntilTheControlPlaneCatchesUp(t *testing.T) { var mu sync.Mutex assignmentReadyAttempts := 0 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch { case r.URL.Path == "/gameserver": _, _ = w.Write([]byte(`{"status":{"address":"203.0.113.9","ports":[{"name":"game","port":31001}]}}`)) case r.URL.Path == "/ready-probe", r.URL.Path == "/ready": w.WriteHeader(http.StatusOK) case r.URL.Path == "/v1/servers/server-1/register": body, _ := io.ReadAll(r.Body) if !strings.Contains(string(body), `"assignment_ready":true`) { w.WriteHeader(http.StatusNoContent) return } mu.Lock() assignmentReadyAttempts++ attempt := assignmentReadyAttempts mu.Unlock() if attempt < 3 { // Simulates the durable `assignments` rows not having // propagated yet -- the API's own real gate for this. w.WriteHeader(http.StatusConflict) return } w.WriteHeader(http.StatusNoContent) default: w.WriteHeader(http.StatusNotFound) } })) defer server.Close() tokenPath := filepath.Join(t.TempDir(), "token") if err := os.WriteFile(tokenPath, []byte("workload-jwt"), 0o600); err != nil { t.Fatal(err) } s, err := New(Config{ Command: []string{"/bin/sh", "-c", "true"}, SDKBaseURL: server.URL, ReadyURL: server.URL + "/ready-probe", ReadyTimeout: time.Second, PollInterval: time.Millisecond, ControlPlaneURL: server.URL, WorkloadTokenPath: tokenPath, ServerID: "server-1", MatchID: "match-1", ProtocolVersion: 1, ImageDigest: "sha256:aa", AssignmentReadyAttempts: 5, AssignmentReadyBackoff: time.Millisecond, }) if err != nil { t.Fatal(err) } // Start must still succeed -- a slow-to-propagate assignment-ready must // never be treated as a Start() failure (which would kill the child). if err := s.Start(context.Background()); err != nil { t.Fatalf("Start failed despite assignment-ready eventually succeeding: %v", err) } if err := s.Wait(); err != nil { t.Fatalf("child was killed despite Start succeeding: %v", err) } mu.Lock() defer mu.Unlock() if assignmentReadyAttempts != 3 { t.Fatalf("assignment-ready attempts = %d, want exactly 3 (2 conflicts then success)", assignmentReadyAttempts) } } func TestControlPlaneRegistrationFallsBackToGameServerAnnotationForMatchID(t *testing.T) { var gotBody string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/gameserver": _, _ = w.Write([]byte(`{"object_meta":{"annotations":{"cosmic-clash.io/match-id":"match-from-annotation","cosmic-clash.io/allocation-id":"allocation-xyz"}},"status":{"address":"203.0.113.9","ports":[{"name":"game","port":31001}]}}`)) case "/ready-probe", "/ready": w.WriteHeader(http.StatusOK) case "/v1/servers/server-1/register": body, _ := io.ReadAll(r.Body) gotBody = string(body) w.WriteHeader(http.StatusNoContent) default: w.WriteHeader(http.StatusNotFound) } })) defer server.Close() tokenPath := filepath.Join(t.TempDir(), "token") if err := os.WriteFile(tokenPath, []byte("workload-jwt"), 0o600); err != nil { t.Fatal(err) } // Deliberately no MatchID in config -- only the GameServer's own // annotation supplies it, proving the fallback path itself, not just // that an explicitly configured value gets sent. s, err := New(Config{ Command: []string{"/bin/sh", "-c", "true"}, SDKBaseURL: server.URL, ReadyURL: server.URL + "/ready-probe", ReadyTimeout: time.Second, PollInterval: time.Millisecond, ControlPlaneURL: server.URL, WorkloadTokenPath: tokenPath, ServerID: "server-1", ProtocolVersion: 1, ImageDigest: "sha256:aa", }) if err != nil { t.Fatal(err) } if err := s.Start(context.Background()); err != nil { t.Fatal(err) } _ = s.Wait() if !strings.Contains(gotBody, `"match_id":"match-from-annotation"`) { t.Fatalf("register body did not use the GameServer annotation's match ID: %s", gotBody) } } func TestControlPlaneRegistrationWithoutMatchIDOrAnnotationFailsClosed(t *testing.T) { registerCalled := false server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/gameserver": _, _ = w.Write([]byte(`{"status":{"address":"203.0.113.9","ports":[{"name":"game","port":31001}]}}`)) case "/ready-probe", "/ready": w.WriteHeader(http.StatusOK) case "/v1/servers/server-1/register": registerCalled = true w.WriteHeader(http.StatusNoContent) default: w.WriteHeader(http.StatusNotFound) } })) defer server.Close() tokenPath := filepath.Join(t.TempDir(), "token") if err := os.WriteFile(tokenPath, []byte("workload-jwt"), 0o600); err != nil { t.Fatal(err) } s, err := New(Config{ Command: []string{"/bin/sh", "-c", "sleep 30"}, SDKBaseURL: server.URL, ReadyURL: server.URL + "/ready-probe", ReadyTimeout: time.Second, PollInterval: time.Millisecond, ControlPlaneURL: server.URL, WorkloadTokenPath: tokenPath, ServerID: "server-1", ProtocolVersion: 1, ImageDigest: "sha256:aa", }) if err != nil { t.Fatal(err) } if err := s.Start(context.Background()); err == nil { t.Fatal("Start succeeded with no match ID available from either config or annotations") } if registerCalled { t.Fatal("register was called despite having no match ID to send") } } // TestControlPlaneRegistrationFallsBackToGameServerAnnotationForWorkloadToken // proves the primary intended delivery channel for the control-plane's // self-issued signed token (server/workload/signed_token.go): with no // --workload-token-path configured at all, a token arriving only via the // cosmic-clash.io/workload-token annotation Agones applies to this // GameServer (server/agones.Client.Allocate) is what gets sent as the // Authorization bearer. func TestControlPlaneRegistrationFallsBackToGameServerAnnotationForWorkloadToken(t *testing.T) { var gotAuth string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/gameserver": _, _ = w.Write([]byte(`{"object_meta":{"annotations":{"cosmic-clash.io/match-id":"match-1","cosmic-clash.io/workload-token":"signed-token-from-annotation"}},"status":{"address":"203.0.113.9","ports":[{"name":"game","port":31001}]}}`)) case "/ready-probe", "/ready": w.WriteHeader(http.StatusOK) case "/v1/servers/server-1/register": gotAuth = r.Header.Get("Authorization") w.WriteHeader(http.StatusNoContent) default: w.WriteHeader(http.StatusNotFound) } })) defer server.Close() // Deliberately no WorkloadTokenPath -- only the GameServer annotation // supplies a token, proving the fallback path itself. s, err := New(Config{ Command: []string{"/bin/sh", "-c", "true"}, SDKBaseURL: server.URL, ReadyURL: server.URL + "/ready-probe", ReadyTimeout: time.Second, PollInterval: time.Millisecond, ControlPlaneURL: server.URL, ServerID: "server-1", ProtocolVersion: 1, ImageDigest: "sha256:aa", }) if err != nil { t.Fatal(err) } if err := s.Start(context.Background()); err != nil { t.Fatal(err) } _ = s.Wait() if gotAuth != "Bearer signed-token-from-annotation" { t.Fatalf("Authorization header = %q, want the annotation-sourced token", gotAuth) } } // TestControlPlaneRegistrationWithoutWorkloadTokenPathOrAnnotationFailsClosed // is the workload-token counterpart to the match-ID fails-closed test above: // with neither a configured token path nor an annotation present, Start must // fail rather than register unauthenticated or with an empty token. func TestControlPlaneRegistrationWithoutWorkloadTokenPathOrAnnotationFailsClosed(t *testing.T) { registerCalled := false server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/gameserver": _, _ = w.Write([]byte(`{"object_meta":{"annotations":{"cosmic-clash.io/match-id":"match-1"}},"status":{"address":"203.0.113.9","ports":[{"name":"game","port":31001}]}}`)) case "/ready-probe", "/ready": w.WriteHeader(http.StatusOK) case "/v1/servers/server-1/register": registerCalled = true w.WriteHeader(http.StatusNoContent) default: w.WriteHeader(http.StatusNotFound) } })) defer server.Close() s, err := New(Config{ Command: []string{"/bin/sh", "-c", "sleep 30"}, SDKBaseURL: server.URL, ReadyURL: server.URL + "/ready-probe", ReadyTimeout: time.Second, PollInterval: time.Millisecond, ControlPlaneURL: server.URL, ServerID: "server-1", ProtocolVersion: 1, ImageDigest: "sha256:aa", }) if err != nil { t.Fatal(err) } if err := s.Start(context.Background()); err == nil { t.Fatal("Start succeeded with no workload token available from either config or annotations") } if registerCalled { t.Fatal("register was called despite having no workload token to send") } } func TestControlPlaneRegistrationFailureKillsChildRatherThanRunningUnregistered(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/gameserver": _, _ = w.Write([]byte(`{"status":{"address":"203.0.113.9","ports":[{"name":"game","port":31001}]}}`)) case "/ready-probe": w.WriteHeader(http.StatusOK) case "/ready": w.WriteHeader(http.StatusOK) case "/v1/servers/server-1/register": w.WriteHeader(http.StatusInternalServerError) default: w.WriteHeader(http.StatusNotFound) } })) defer server.Close() tokenPath := filepath.Join(t.TempDir(), "token") if err := os.WriteFile(tokenPath, []byte("workload-jwt"), 0o600); err != nil { t.Fatal(err) } // A long-running child: if Start's failure path did not actually kill it, // Wait would block for the full sleep instead of returning promptly with // a "signal: killed" style exit. s, err := New(Config{ Command: []string{"/bin/sh", "-c", "sleep 30"}, SDKBaseURL: server.URL, ReadyURL: server.URL + "/ready-probe", ReadyTimeout: time.Second, PollInterval: time.Millisecond, ControlPlaneURL: server.URL, WorkloadTokenPath: tokenPath, ServerID: "server-1", MatchID: "match-1", ProtocolVersion: 1, ImageDigest: "sha256:aa", }) if err != nil { t.Fatal(err) } if err := s.Start(context.Background()); err == nil { t.Fatal("Start succeeded despite the control-plane rejecting registration") } done := make(chan error, 1) go func() { done <- s.Wait() }() select { case err := <-done: if err == nil { t.Fatal("child was not actually killed after a failed registration") } case <-time.After(5 * time.Second): t.Fatal("child was still running 5s after a failed registration should have killed it") } } func TestAllocatedENetDoesNotReceiveSDRVariables(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path == "/gameserver" { _, _ = w.Write([]byte(`{"status":{"address":"127.0.0.1","ports":[{"name":"game","port":31002}]}}`)) return } if r.URL.Path == "/ready-probe" { w.WriteHeader(http.StatusOK) return } w.WriteHeader(http.StatusOK) })) defer server.Close() path := filepath.Join(t.TempDir(), "env.txt") s, err := New(Config{Command: []string{"/bin/sh", "-c", "env > " + path}, SDKBaseURL: server.URL, ReadyURL: server.URL + "/ready-probe", Transport: "enet", ReadyTimeout: time.Second}) if err != nil { t.Fatal(err) } if err := s.Start(context.Background()); err != nil { t.Fatal(err) } if err := s.Wait(); err != nil { t.Fatal(err) } contents, err := os.ReadFile(path) if err != nil { t.Fatal(err) } if strings.Contains(string(contents), "SDR_LISTEN_PORT=") || strings.Contains(string(contents), "SDR_IP=") { t.Fatalf("ENet received SDR variables: %s", contents) } } func TestDirectModeDoesNotRequireAgonesReadiness(t *testing.T) { s, err := New(Config{Command: []string{"/bin/sh", "-c", "exit 0"}}) if err != nil { t.Fatal(err) } if err := s.Start(context.Background()); err != nil { t.Fatal(err) } if err := s.Wait(); err != nil { t.Fatal(err) } } func TestDrainRequiresAndUsesAuthenticatedLocalEndpoint(t *testing.T) { seenToken := "" server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/drain" { w.WriteHeader(http.StatusNotFound) return } seenToken = r.Header.Get("Authorization") if seenToken != "Bearer secret-token" { w.WriteHeader(http.StatusUnauthorized) return } w.WriteHeader(http.StatusAccepted) })) defer server.Close() s, err := New(Config{Command: []string{"/bin/sh", "-c", "exit 0"}, DrainURL: server.URL + "/drain", DrainToken: "secret-token"}) if err != nil { t.Fatal(err) } if err := s.Drain(context.Background()); err != nil { t.Fatal(err) } if seenToken != "Bearer secret-token" { t.Fatalf("unexpected drain token: %q", seenToken) } missing, _ := New(Config{Command: []string{"/bin/sh", "-c", "exit 0"}}) if err := missing.Drain(context.Background()); err == nil { t.Fatal("unauthenticated drain was allowed") } } func TestAssignedEndpointRejectsMalformedAddressAndPort(t *testing.T) { for _, response := range []string{ `{"status":{"address":"203.0.113.9","ports":[{"name":"game","port":65536}]}}`, `{"status":{"address":" ","ports":[{"name":"game","port":31001}]}}`, `{"status":{"address":"203.0.113.9 bad","ports":[{"name":"game","port":31001}]}}`, } { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path == "/gameserver" { _, _ = w.Write([]byte(response)) return } w.WriteHeader(http.StatusOK) })) s, err := New(Config{Command: []string{"/bin/sh", "-c", "exit 0"}, SDKBaseURL: server.URL, ReadyURL: server.URL + "/probe", ReadyTimeout: time.Second}) if err != nil { server.Close() t.Fatal(err) } if err := s.Start(context.Background()); err == nil { t.Errorf("malformed endpoint was accepted: %s", response) } server.Close() } } func TestSupervisorRejectsRemoteOrPartialDrainConfiguration(t *testing.T) { for _, config := range []Config{ {Command: []string{"/bin/sh", "-c", "exit 0"}, DrainURL: "https://example.com/drain", DrainToken: "token-1234567890123456"}, {Command: []string{"/bin/sh", "-c", "exit 0"}, DrainURL: "http://127.0.0.1/drain"}, {Command: []string{"/bin/sh", "-c", "exit 0"}, DrainToken: "token-1234567890123456"}, {Command: []string{"/bin/sh", "-c", "exit 0"}, DrainURL: "http://127.0.0.1/drain?token=leaked", DrainToken: "token-1234567890123456"}, } { if _, err := New(config); err == nil { t.Fatalf("unsafe drain configuration accepted: %+v", config) } } } func TestRunDrainsBeforeChildExit(t *testing.T) { marker := filepath.Join(t.TempDir(), "drained") server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/drain" { w.WriteHeader(http.StatusNotFound) return } if r.Header.Get("Authorization") != "Bearer run-secret" { w.WriteHeader(http.StatusUnauthorized) return } if err := os.WriteFile(marker, []byte("drained"), 0600); err != nil { http.Error(w, err.Error(), http.StatusInternalServerError) return } w.WriteHeader(http.StatusAccepted) })) defer server.Close() command := []string{"/bin/sh", "-c", "while [ ! -f '" + marker + "' ]; do sleep 0.01; done"} s, err := New(Config{Command: command, DrainURL: server.URL + "/drain", DrainToken: "run-secret"}) if err != nil { t.Fatal(err) } ctx, cancel := context.WithCancel(context.Background()) time.AfterFunc(30*time.Millisecond, cancel) if err := s.Run(ctx, time.Second); err != nil { t.Fatalf("graceful run: %v", err) } if _, err := os.Stat(marker); err != nil { t.Fatalf("drain endpoint was not called: %v", err) } } func TestRunAcknowledgesControlledShutdownWithWorkloadCredential(t *testing.T) { marker := filepath.Join(t.TempDir(), "drained") tokenPath := filepath.Join(t.TempDir(), "token") if err := os.WriteFile(tokenPath, []byte("workload-secret"), 0600); err != nil { t.Fatal(err) } var shutdownCalls int server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/drain": if r.Header.Get("Authorization") != "Bearer run-secret" { w.WriteHeader(http.StatusUnauthorized) return } if err := os.WriteFile(marker, []byte("drained"), 0600); err != nil { http.Error(w, err.Error(), http.StatusInternalServerError) return } w.WriteHeader(http.StatusAccepted) case "/v1/servers/server-1/shutdown": if r.Header.Get("Authorization") != "Bearer workload-secret" || r.Header.Get("Idempotency-Key") == "" { w.WriteHeader(http.StatusUnauthorized) return } shutdownCalls++ w.WriteHeader(http.StatusNoContent) default: w.WriteHeader(http.StatusNotFound) } })) defer server.Close() s, err := New(Config{ Command: []string{"/bin/sh", "-c", "while [ ! -f '" + marker + "' ]; do sleep 0.01; done"}, DrainURL: server.URL + "/drain", DrainToken: "run-secret", ControlPlaneURL: server.URL, WorkloadTokenPath: tokenPath, ServerID: "server-1", MatchID: "match-1", ProtocolVersion: 1, ImageDigest: "sha256:aa", }) if err != nil { t.Fatal(err) } ctx, cancel := context.WithCancel(context.Background()) time.AfterFunc(30*time.Millisecond, cancel) if err := s.Run(ctx, time.Second); err != nil { t.Fatalf("graceful run: %v", err) } if shutdownCalls != 1 { t.Fatalf("shutdown calls = %d, want 1", shutdownCalls) } } func TestRunDoesNotAcknowledgeWhenLocalDrainFails(t *testing.T) { var shutdownCalls int server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path == "/v1/servers/server-1/shutdown" { shutdownCalls++ } w.WriteHeader(http.StatusInternalServerError) })) defer server.Close() s, err := New(Config{ Command: []string{"/bin/sh", "-c", "sleep 5"}, DrainURL: server.URL + "/drain", DrainToken: "run-secret", ControlPlaneURL: server.URL, WorkloadTokenPath: filepath.Join(t.TempDir(), "missing-token"), ServerID: "server-1", MatchID: "match-1", ProtocolVersion: 1, ImageDigest: "sha256:aa", }) if err != nil { t.Fatal(err) } ctx, cancel := context.WithCancel(context.Background()) time.AfterFunc(30*time.Millisecond, cancel) err = s.Run(ctx, 50*time.Millisecond) if err == nil || shutdownCalls != 0 { t.Fatalf("failed drain result=%v shutdown calls=%d", err, shutdownCalls) } } func TestRunForceKillsUnresponsiveChildAtDeadline(t *testing.T) { s, err := New(Config{Command: []string{"/bin/sh", "-c", "trap '' TERM; sleep 5"}}) if err != nil { t.Fatal(err) } ctx, cancel := context.WithCancel(context.Background()) time.AfterFunc(30*time.Millisecond, cancel) started := time.Now() err = s.Run(ctx, 50*time.Millisecond) if err == nil || !strings.Contains(err.Error(), "force-killed") { t.Fatalf("unresponsive child result = %v", err) } if elapsed := time.Since(started); elapsed > time.Second { t.Fatalf("force-kill exceeded bounded deadline: %s", elapsed) } }