Made security review, applied recommendations.

sts-agent/.dockerignore

@@ -0,0 +1,5 @@
+*
+!go.mod
+!go.sum
+!main.go
+!main_test.go

sts-agent/Dockerfile

@@ -1,13 +1,14 @@
-FROM golang:1.27-trixie AS build
+FROM golang:1.27.3-trixie AS build
 
+ENV GOTOOLCHAIN=local CGO_ENABLED=0
 WORKDIR /src
 COPY go.mod go.sum ./
 RUN go mod download
 COPY main.go main_test.go ./
-RUN CGO_ENABLED=0 go test ./... \
-    && CGO_ENABLED=0 go build -trimpath -ldflags="-s -w" -o /imds-emulator .
+RUN go test ./... \
+    && go build -trimpath -ldflags="-s -w" -o /imds-emulator .
 
-FROM gcr.io/distroless/static-debian13
+FROM gcr.io/distroless/static-debian13:latest
 
 COPY --from=build /imds-emulator /imds-emulator
 

sts-agent/compose.yaml

@@ -2,6 +2,8 @@ name: sts
 services:
   imds-emulator:
     build: .
+    restart: unless-stopped
+    stop_grace_period: 30s
     volumes:
       - /etc/sts-agent:/etc/sts-agent:ro
     environment:
@@ -9,11 +11,20 @@ services:
       GOOGLE_API_CERTIFICATE_CONFIG: /etc/sts-agent/certificate_config.json
       PROJECT_ID: ${PROJECT_ID}
       PROJECT_NUMBER: ${PROJECT_NUMBER}
+    healthcheck:
+      test: ["CMD", "/imds-emulator", "-check"]
+      interval: 10s
+      timeout: 3s
+      start_period: 10s
+      start_interval: 1s
+      retries: 3
 
   agent:
     image: gcr.io/cloud-ingest/tsop-agent:latest
+    restart: unless-stopped
     depends_on:
-      - imds-emulator
+      imds-emulator:
+        condition: service_healthy
     ulimits:
       memlock: 64000000
     volumes:

sts-agent/go.mod

@@ -1,6 +1,6 @@
 module imds-emulator
 
-go 1.27.1
+go 1.27.3
 
 require (
 	cloud.google.com/go/auth v0.24.1

sts-agent/main.go

@@ -6,6 +6,8 @@ import (
 	"cmp"
 	"context"
 	"encoding/json"
+	"flag"
+	"fmt"
 	"io"
 	"log"
 	"net"
@@ -31,8 +33,11 @@ const (
 	accounts       = "instance/service-accounts/"
 	defaultAccount = accounts + "default/"
 
-	// shutdownGrace is shorter than the 10 seconds Docker waits after SIGTERM.
-	shutdownGrace = 5 * time.Second
+	// shutdownGrace is shorter than the stop grace period of the container.
+	shutdownGrace = 20 * time.Second
+
+	// maxProviders bounds the token providers kept for distinct scope sets.
+	maxProviders = 16
 )
 
 // tokenFunc returns an access token for the scopes.
@@ -96,11 +101,7 @@ func (s *metadataServer) handle(w http.ResponseWriter, r *http.Request) int {
 
 	switch {
 	case path == defaultAccount+"token":
-		scopes := []string{defaultScope}
-		if value := query.Get("scopes"); value != "" {
-			scopes = strings.Split(value, ",")
-		}
-		token, err := s.token(r.Context(), scopes)
+		token, err := s.token(r.Context(), parseScopes(query.Get("scopes")))
 		if err != nil {
 			log.Printf("token error: %v", err)
 			return text(w, http.StatusInternalServerError, "Internal Server Error")
@@ -138,11 +139,22 @@ func (s *metadataServer) handle(w http.ResponseWriter, r *http.Request) int {
 	return text(w, http.StatusNotFound, "Not Found")
 }
 
+// parseScopes returns the comma separated scopes of the query, or the default scope.
+func parseScopes(value string) []string {
+	scopes := slices.DeleteFunc(strings.Split(value, ","), func(scope string) bool { return scope == "" })
+	if len(scopes) == 0 {
+		return []string{defaultScope}
+	}
+	return scopes
+}
+
 // text writes a plain text response and returns its status.
 func text(w http.ResponseWriter, status int, body string) int {
 	w.Header().Set("Content-Type", "application/text")
 	w.WriteHeader(status)
-	io.WriteString(w, body)
+	if _, err := io.WriteString(w, body); err != nil {
+		log.Printf("write error: %v", err)
+	}
 	return status
 }
 
@@ -150,10 +162,31 @@ func text(w http.ResponseWriter, status int, body string) int {
 func jsonBody(w http.ResponseWriter, value any) int {
 	w.Header().Set("Content-Type", "application/json")
 	w.WriteHeader(http.StatusOK)
-	json.NewEncoder(w).Encode(value)
+	if err := json.NewEncoder(w).Encode(value); err != nil {
+		log.Printf("write error: %v", err)
+	}
 	return http.StatusOK
 }
 
+// ping checks that the emulator at the address answers the metadata ping.
+func ping(address string) error {
+	request, err := http.NewRequest(http.MethodGet, "http://"+address+"/", nil)
+	if err != nil {
+		return err
+	}
+	request.Header.Set("Metadata-Flavor", "Google")
+	client := &http.Client{Timeout: 2 * time.Second}
+	response, err := client.Do(request)
+	if err != nil {
+		return err
+	}
+	defer response.Body.Close()
+	if response.StatusCode != http.StatusOK || response.Header.Get("Metadata-Flavor") != "Google" {
+		return fmt.Errorf("unexpected response: %s", response.Status)
+	}
+	return nil
+}
+
 var impersonationURL = regexp.MustCompile(`/serviceAccounts/([^/:]+):generateAccessToken$`)
 
 // serviceAccountEmail returns the service account of a credentials file, or an empty string.
@@ -193,6 +226,10 @@ func (c *providerCache) get(scopes []string) (auth.TokenProvider, error) {
 	if c.providers == nil {
 		c.providers = map[string]auth.TokenProvider{}
 	}
+	// Drop all providers once the limit is reached, so the cache never grows unbounded.
+	if len(c.providers) >= maxProviders {
+		clear(c.providers)
+	}
 	c.providers[key] = provider
 	return provider, nil
 }
@@ -230,6 +267,18 @@ func serve(ctx context.Context, server *http.Server, listener net.Listener, grac
 }
 
 func main() {
+	check := flag.Bool("check", false, "check that the emulator answers and exit")
+	flag.Parse()
+
+	// IMDS_HOST empty binds all interfaces; the check then uses the loopback address.
+	host, port := os.Getenv("IMDS_HOST"), cmp.Or(os.Getenv("IMDS_PORT"), "8080")
+	if *check {
+		if err := ping(net.JoinHostPort(cmp.Or(host, "127.0.0.1"), port)); err != nil {
+			log.Fatalf("Check failed: %v", err)
+		}
+		return
+	}
+
 	// The first signal stops the server gracefully, the second one exits at once.
 	signals := make(chan os.Signal, 2)
 	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
@@ -279,7 +328,7 @@ func main() {
 		IdleTimeout:       60 * time.Second,
 	}
 
-	listener, err := net.Listen("tcp", ":"+cmp.Or(os.Getenv("IMDS_PORT"), "8080"))
+	listener, err := net.Listen("tcp", net.JoinHostPort(host, port))
 	if err != nil {
 		log.Fatalf("Cannot listen: %v", err)
 	}

sts-agent/main_test.go

@@ -205,10 +205,80 @@ func TestToken(t *testing.T) {
 			if !slices.Equal(scopes, test.wantScopes) {
 				t.Errorf("scopes = %v, want %v", scopes, test.wantScopes)
 			}
+			if got := recorder.Header().Get("Content-Type"); got != "application/json" {
+				t.Errorf("Content-Type = %q, want application/json", got)
+			}
 		})
 	}
 }
 
+func TestParseScopes(t *testing.T) {
+	tests := []struct {
+		value string
+		want  []string
+	}{
+		{"", []string{defaultScope}},
+		{",,", []string{defaultScope}},
+		{"a", []string{"a"}},
+		{"a,,b,", []string{"a", "b"}},
+	}
+	for _, test := range tests {
+		if got := parseScopes(test.value); !slices.Equal(got, test.want) {
+			t.Errorf("parseScopes(%q) = %v, want %v", test.value, got, test.want)
+		}
+	}
+}
+
+func TestRequestLogOmitsQuery(t *testing.T) {
+	var logged bytes.Buffer
+	log.SetOutput(&logged)
+	defer log.SetOutput(io.Discard)
+
+	get(newTestServer(nil), root+"instance/service-accounts/default/token?scopes=secret-scope")
+	if !strings.Contains(logged.String(), "GET "+root+"instance/service-accounts/default/token 200") {
+		t.Errorf("the request was not logged:\n%s", logged.String())
+	}
+	if strings.Contains(logged.String(), "secret-scope") {
+		t.Errorf("the log contains the query:\n%s", logged.String())
+	}
+}
+
+func TestHeadHasNoBody(t *testing.T) {
+	server := httptest.NewServer(newTestServer(nil))
+	defer server.Close()
+
+	request, _ := http.NewRequest(http.MethodHead, server.URL+root+"project/project-id", nil)
+	request.Header.Set("Metadata-Flavor", "Google")
+	response, err := http.DefaultClient.Do(request)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer response.Body.Close()
+	body, _ := io.ReadAll(response.Body)
+	if response.StatusCode != 200 || len(body) != 0 {
+		t.Errorf("got %d with body %q, want 200 without body", response.StatusCode, body)
+	}
+}
+
+func TestPing(t *testing.T) {
+	emulator := httptest.NewServer(newTestServer(nil))
+	defer emulator.Close()
+	other := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+	defer other.Close()
+	closed := httptest.NewServer(nil)
+	closed.Close()
+
+	if err := ping(strings.TrimPrefix(emulator.URL, "http://")); err != nil {
+		t.Errorf("ping(emulator) = %v, want nil", err)
+	}
+	if err := ping(strings.TrimPrefix(other.URL, "http://")); err == nil {
+		t.Error("ping(other server) = nil, want an error")
+	}
+	if err := ping(strings.TrimPrefix(closed.URL, "http://")); err == nil {
+		t.Error("ping(closed) = nil, want an error")
+	}
+}
+
 func TestTokenExpiry(t *testing.T) {
 	tests := []struct {
 		name     string
@@ -315,6 +385,20 @@ func TestProviderCache(t *testing.T) {
 	}
 }
 
+func TestProviderCacheBound(t *testing.T) {
+	cache := &providerCache{detect: func([]string) (auth.TokenProvider, error) {
+		return tokenProviderFunc(func(context.Context) (*auth.Token, error) { return &auth.Token{Value: "t"}, nil }), nil
+	}}
+	for i := range 10 * maxProviders {
+		if _, err := cache.token(context.Background(), []string{strings.Repeat("s", i+1)}); err != nil {
+			t.Fatal(err)
+		}
+		if len(cache.providers) > maxProviders {
+			t.Fatalf("%d providers cached after %d scope sets, want at most %d", len(cache.providers), i+1, maxProviders)
+		}
+	}
+}
+
 func TestProviderCacheConcurrent(t *testing.T) {
 	var detected atomic.Int32
 	cache := &providerCache{detect: func([]string) (auth.TokenProvider, error) {
@@ -497,6 +581,98 @@ func writeServiceAccountKey(t *testing.T) string {
 	return path
 }
 
+// safeBuffer collects the output of a process while it runs.
+type safeBuffer struct {
+	mu  sync.Mutex
+	buf bytes.Buffer
+}
+
+func (b *safeBuffer) Write(p []byte) (int, error) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return b.buf.Write(p)
+}
+
+func (b *safeBuffer) String() string {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return b.buf.String()
+}
+
+// emulator is the test binary running as the emulator process.
+type emulator struct {
+	address string
+	process *os.Process
+	exited  <-chan error
+	output  *safeBuffer
+}
+
+// freePort returns a free TCP port of the loopback interface.
+func freePort(t *testing.T) string {
+	t.Helper()
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer listener.Close()
+	_, port, _ := net.SplitHostPort(listener.Addr().String())
+	return port
+}
+
+// runEmulator starts the test binary as the emulator with the arguments and
+// the environment on the loopback interface, and returns it once it answers.
+func runEmulator(t *testing.T, env []string, args ...string) *emulator {
+	t.Helper()
+	port := freePort(t)
+	e := &emulator{address: "127.0.0.1:" + port, output: &safeBuffer{}}
+
+	command := exec.Command(os.Args[0], args...)
+	command.Env = append(os.Environ(), mainEnv+"=1", "GOOGLE_APPLICATION_CREDENTIALS="+writeServiceAccountKey(t), "IMDS_HOST=127.0.0.1", "IMDS_PORT="+port)
+	command.Env = append(command.Env, env...)
+	command.Stderr = e.output
+	if err := command.Start(); err != nil {
+		t.Fatal(err)
+	}
+	exited := make(chan error, 1)
+	go func() { exited <- command.Wait() }()
+	e.process, e.exited = command.Process, exited
+	t.Cleanup(func() { command.Process.Kill() })
+
+	for deadline := time.Now().Add(10 * time.Second); time.Now().Before(deadline); {
+		select {
+		case err := <-exited:
+			t.Fatalf("the emulator exited before it was ready: %v\n%s", err, e.output.String())
+		default:
+		}
+		if ping(e.address) == nil {
+			return e
+		}
+		time.Sleep(50 * time.Millisecond)
+	}
+	t.Fatalf("the emulator did not answer\n%s", e.output.String())
+	return nil
+}
+
+// wait returns the exit error of the emulator, failing the test when it does not stop.
+func (e *emulator) wait(t *testing.T) error {
+	t.Helper()
+	select {
+	case err := <-e.exited:
+		return err
+	case <-time.After(10 * time.Second):
+		t.Fatalf("the emulator did not stop\n%s", e.output.String())
+		return nil
+	}
+}
+
+// runCheck runs the test binary in check mode against the port and returns its exit error.
+func runCheck(t *testing.T, port string) error {
+	t.Helper()
+	command := exec.Command(os.Args[0], "-check")
+	command.Env = append(os.Environ(), mainEnv+"=1", "IMDS_PORT="+port)
+	return command.Run()
+}
+
 // TestSignals runs the emulator as a process and stops it with signals.
 func TestSignals(t *testing.T) {
 	if runtime.GOOS == "windows" {
@@ -504,59 +680,96 @@ func TestSignals(t *testing.T) {
 	}
 	for _, signal := range []syscall.Signal{syscall.SIGTERM, syscall.SIGINT} {
 		t.Run(signal.String(), func(t *testing.T) {
-			listener, err := net.Listen("tcp", "127.0.0.1:0")
-			if err != nil {
+			e := runEmulator(t, nil)
+			if err := e.process.Signal(signal); err != nil {
 				t.Fatal(err)
 			}
-			address := listener.Addr().String()
-			_, port, _ := net.SplitHostPort(address)
-			listener.Close()
-
-			var output bytes.Buffer
-			command := exec.Command(os.Args[0])
-			command.Env = append(os.Environ(), mainEnv+"=1", "GOOGLE_APPLICATION_CREDENTIALS="+writeServiceAccountKey(t), "IMDS_PORT="+port)
-			command.Stderr = &output
-			if err := command.Start(); err != nil {
-				t.Fatal(err)
+			if err := e.wait(t); err != nil {
+				t.Errorf("the emulator exited with %v, want status 0\n%s", err, e.output.String())
 			}
-			exited := make(chan error, 1)
-			go func() { exited <- command.Wait() }()
-			t.Cleanup(func() { command.Process.Kill() })
-
-			// Wait until the emulator answers the ping.
-			ready := false
-			for deadline := time.Now().Add(10 * time.Second); time.Now().Before(deadline) && !ready; {
-				select {
-				case err := <-exited:
-					t.Fatalf("the emulator exited before it was ready: %v\n%s", err, output.String())
-				default:
-				}
-				request, _ := http.NewRequest(http.MethodGet, "http://"+address+"/", nil)
-				request.Header.Set("Metadata-Flavor", "Google")
-				if response, err := http.DefaultClient.Do(request); err == nil {
-					response.Body.Close()
-					ready = response.StatusCode == http.StatusOK
-				}
-				time.Sleep(50 * time.Millisecond)
-			}
-			if !ready {
-				t.Fatalf("the emulator did not answer\n%s", output.String())
-			}
-
-			if err := command.Process.Signal(signal); err != nil {
-				t.Fatal(err)
-			}
-			select {
-			case err := <-exited:
-				if err != nil {
-					t.Errorf("the emulator exited with %v, want status 0\n%s", err, output.String())
-				}
-			case <-time.After(10 * time.Second):
-				t.Fatalf("the emulator did not stop\n%s", output.String())
-			}
-			if !strings.Contains(output.String(), "shutting down") || !strings.Contains(output.String(), "Stopped") {
-				t.Errorf("unexpected log:\n%s", output.String())
+			if !strings.Contains(e.output.String(), "shutting down") || !strings.Contains(e.output.String(), "Stopped") {
+				t.Errorf("unexpected log:\n%s", e.output.String())
 			}
 		})
 	}
+
+	t.Run("second signal", func(t *testing.T) {
+		e := runEmulator(t, nil)
+
+		// A partial request keeps a connection active, so the shutdown waits for it.
+		connection, err := net.Dial("tcp", e.address)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer connection.Close()
+		if _, err := io.WriteString(connection, "GET / HTTP/1.1\r\nHost: imds\r\n"); err != nil {
+			t.Fatal(err)
+		}
+
+		if err := e.process.Signal(syscall.SIGTERM); err != nil {
+			t.Fatal(err)
+		}
+		select {
+		case err := <-e.exited:
+			t.Fatalf("the emulator exited with %v while a connection was active", err)
+		case <-time.After(500 * time.Millisecond):
+		}
+		if err := e.process.Signal(syscall.SIGTERM); err != nil {
+			t.Fatal(err)
+		}
+		var exitError *exec.ExitError
+		if err := e.wait(t); !errors.As(err, &exitError) || exitError.ExitCode() != 1 {
+			t.Errorf("the emulator exited with %v, want status 1\n%s", err, e.output.String())
+		}
+		if !strings.Contains(e.output.String(), "again, exiting") {
+			t.Errorf("unexpected log:\n%s", e.output.String())
+		}
+	})
+}
+
+// TestProcess checks the environment handling and the check mode of the emulator process.
+func TestProcess(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("POSIX signals")
+	}
+	email := func(t *testing.T, address string) string {
+		t.Helper()
+		request, _ := http.NewRequest(http.MethodGet, "http://"+address+root+"instance/service-accounts/default/email", nil)
+		request.Header.Set("Metadata-Flavor", "Google")
+		response, err := http.DefaultClient.Do(request)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer response.Body.Close()
+		body, _ := io.ReadAll(response.Body)
+		return string(body)
+	}
+
+	t.Run("email from credentials", func(t *testing.T) {
+		e := runEmulator(t, nil)
+		if got := email(t, e.address); got != testEmail {
+			t.Errorf("email = %q, want %q", got, testEmail)
+		}
+		if !strings.Contains(e.output.String(), "Serving tokens of "+testEmail+" for project dom-lab on "+e.address) {
+			t.Errorf("unexpected log:\n%s", e.output.String())
+		}
+	})
+
+	t.Run("email override", func(t *testing.T) {
+		e := runEmulator(t, []string{"SERVICE_ACCOUNT_EMAIL=other@dom-lab.iam.gserviceaccount.com"})
+		if got := email(t, e.address); got != "other@dom-lab.iam.gserviceaccount.com" {
+			t.Errorf("email = %q, want the override", got)
+		}
+	})
+
+	t.Run("check", func(t *testing.T) {
+		e := runEmulator(t, nil)
+		_, port, _ := net.SplitHostPort(e.address)
+		if err := runCheck(t, port); err != nil {
+			t.Errorf("check against the emulator = %v, want nil", err)
+		}
+		if err := runCheck(t, freePort(t)); err == nil {
+			t.Error("check against a closed port = nil, want an error")
+		}
+	})
 }