Made security review, applied recommendations.

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)
 	}