agents/sec-rev: final tweaks to agent code after review

Signed-off-by: Jacob Howard <[email protected]>

Author: xenoscopic

agents/security-reviewer/Dockerfile

@@ -54,7 +54,7 @@ RUN useradd --create-home --shell /bin/bash agent \
     && install -d -o agent -g agent /workspace/input \
     && install -d -o agent -g agent /workspace/output
 
-# Copy the reviewer assets and report template.
+# Copy the reviewer entrypoint and assets.
 COPY --from=gobuilder /security-reviewer/build/reviewer /opt/security-reviewer/reviewer
 COPY --chown=agent:agent prompt-template.md /opt/security-reviewer/prompt-template.md
 COPY --chown=agent:agent report-template.md /opt/security-reviewer/report-template.md
@@ -67,17 +67,22 @@ WORKDIR /workspace
 ENTRYPOINT ["/opt/security-reviewer/reviewer"]
 
 
-# Lightweight image hosting the API proxy.
+# Use a debian image to host the proxy.
 FROM debian:trixie-slim AS proxy-image
 
+# Install dependencies, including those needed for healthchecks.
 ENV DEBIAN_FRONTEND=noninteractive
 RUN apt-get update \
     && apt-get install -y --no-install-recommends \
         ca-certificates \
         wget \
     && rm -rf /var/lib/apt/lists/*
 
+# Copy the proxy entrypoint.
 COPY --from=gobuilder /security-reviewer/build/proxy /opt/security-reviewer/proxy
 
+# Expose the proxy's default port.
 EXPOSE 4000
+
+# Set the proxy entrypoint.
 ENTRYPOINT ["/opt/security-reviewer/proxy"]

agents/security-reviewer/proxy/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"crypto/subtle"
 	"errors"
 	"fmt"
 	"log"
@@ -10,27 +11,44 @@ import (
 	"net/http/httputil"
 	"net/url"
 	"os"
+	"os/signal"
 	"strings"
+	"syscall"
 	"time"
 )
 
 const (
-	defaultListenAddr       = ":4000"
-	defaultOpenAIBaseURL    = "https://api.openai.com/v1"
+	// defaultListenAddr is the fallback bind address when none is provided via PROXY_LISTEN_ADDR.
+	defaultListenAddr = ":4000"
+	// defaultOpenAIBaseURL is the upstream OpenAI API base path used when none is provided.
+	defaultOpenAIBaseURL = "https://api.openai.com/v1"
+	// defaultAnthropicBaseURL is the upstream Anthropic API base path used when none is provided.
+	// NOTE: Anthropic clients are expected to include /v1 in their requests, so
+	// it is not idiomatic to include it in the base URL.
 	defaultAnthropicBaseURL = "https://api.anthropic.com"
-	openAIInboundPrefix     = "/openai/"
-	anthropicInboundPrefix  = "/anthropic/"
-	healthPath              = "/health/liveliness"
-	headerAuthorization     = "Authorization"
-	headerAnthropicAPIKey   = "X-Api-Key"
+	// openAIInboundPrefix is the path prefix used to route requests to OpenAI.
+	openAIInboundPrefix = "/openai/"
+	// anthropicInboundPrefix is the path prefix used to route requests to Anthropic.
+	anthropicInboundPrefix = "/anthropic/"
+	// healthPath is the HTTP endpoint used for container health checks.
+	healthPath = "/health/liveliness"
+	// headerAuthorization is the inbound HTTP header that carries bearer tokens.
+	headerAuthorization = "Authorization"
+	// headerAnthropicAPIKey is the Anthropic-specific header carrying API keys.
+	headerAnthropicAPIKey = "X-Api-Key"
 )
 
 // providerProxy defines how to forward requests to a specific upstream API.
 type providerProxy struct {
-	Prefix      string
-	Target      *url.URL
-	HeaderName  string
+	// Prefix is the inbound path prefix handled by the provider.
+	Prefix string
+	// Target is the upstream endpoint used to service requests for the provider.
+	Target *url.URL
+	// HeaderName is the outbound header carrying the provider-specific credential.
+	HeaderName string
+	// HeaderValue is the credential value set on outbound requests.
 	HeaderValue string
+	// DisplayName is the human-readable name of the provider used in logs.
 	DisplayName string
 }
 
@@ -41,6 +59,9 @@ func main() {
 		log.Fatalf("proxy configuration error: %v", err)
 	}
 
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
+	defer stop()
+
 	mux := http.NewServeMux()
 	mux.HandleFunc(healthPath, handleHealth)
 
@@ -63,6 +84,15 @@ func main() {
 	log.Printf("proxy listening on %s (OpenAI -> %s, Anthropic -> %s)",
 		cfg.listenAddr, cfg.openAIProxy.Target.String(), cfg.anthropicProxy.Target.String())
 
+	go func() {
+		<-ctx.Done()
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		if err := server.Shutdown(shutdownCtx); err != nil {
+			log.Printf("proxy shutdown error: %v", err)
+		}
+	}()
+
 	if err = server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
 		log.Fatalf("proxy server error: %v", err)
 	}
@@ -289,6 +319,8 @@ func firstNonEmpty(values ...string) string {
 	return ""
 }
 
+// validateClientToken ensures inbound requests present the proxy bearer secret using
+// a constant-time comparison to avoid leaking timing information.
 func validateClientToken(headerValue, expectedToken string) bool {
 	if expectedToken == "" {
 		return false
@@ -300,5 +332,6 @@ func validateClientToken(headerValue, expectedToken string) bool {
 	if !strings.EqualFold(parts[0], "bearer") {
 		return false
 	}
-	return strings.TrimSpace(parts[1]) == expectedToken
+	provided := strings.TrimSpace(parts[1])
+	return subtle.ConstantTimeCompare([]byte(provided), []byte(expectedToken)) == 1
 }