mirror of
https://github.com/jcreek/CosmicClash.git
synced 2026-09-11 00:14:00 +00:00
feat: add authenticated supervisor drain
This commit is contained in:
@@ -31,6 +31,8 @@ type Config struct {
|
||||
SDKBaseURL string
|
||||
ReadyURL string
|
||||
Transport string
|
||||
DrainURL string
|
||||
DrainToken string
|
||||
ReadyTimeout time.Duration
|
||||
PollInterval time.Duration
|
||||
HTTPClient *http.Client
|
||||
@@ -115,6 +117,29 @@ func (s *Supervisor) Wait() error {
|
||||
return s.cmd.Wait()
|
||||
}
|
||||
|
||||
// Drain asks the allocated Godot process to stop accepting new work. The
|
||||
// token is sent only over the configured localhost control endpoint and is
|
||||
// never placed in command arguments or logs.
|
||||
func (s *Supervisor) Drain(ctx context.Context) error {
|
||||
if s.config.DrainURL == "" || s.config.DrainToken == "" {
|
||||
return fmt.Errorf("authenticated drain endpoint is required")
|
||||
}
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodPost, s.config.DrainURL, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
request.Header.Set("Authorization", "Bearer "+s.config.DrainToken)
|
||||
response, err := s.client.Do(request)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer response.Body.Close()
|
||||
if response.StatusCode/100 != 2 {
|
||||
return fmt.Errorf("drain endpoint returned %s", response.Status)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Supervisor) assignedEndpoint(ctx context.Context) (int, string, error) {
|
||||
var server GameServer
|
||||
if err := s.sdkGet(ctx, "/gameserver", &server); err != nil {
|
||||
|
||||
@@ -111,3 +111,34 @@ func TestDirectModeDoesNotRequireAgonesReadiness(t *testing.T) {
|
||||
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")
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user