From 6e16b720517728939c92a4a7f3490a48860a2a5f Mon Sep 17 00:00:00 2001
From: Brian Ginsburg <7957636+bgins@users.noreply.github.com>
Date: Thu, 23 Jan 2025 09:49:22 -0800
Subject: [PATCH] feat: Add validation token key ID config (#495)

---
 pkg/http/types.go     | 1 +
 pkg/options/server.go | 9 +++++++++
 pkg/solver/server.go  | 2 +-
 stack                 | 1 +
 4 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/pkg/http/types.go b/pkg/http/types.go
index cbcb884d..c50d143f 100644
--- a/pkg/http/types.go
+++ b/pkg/http/types.go
@@ -11,6 +11,7 @@ type ServerOptions struct {
 type AccessControlOptions struct {
 	ValidationTokenSecret     string
 	ValidationTokenExpiration int
+	ValidationTokenKid        string
 }
 
 type ValidationToken struct {
diff --git a/pkg/options/server.go b/pkg/options/server.go
index 39f74926..29b1776d 100644
--- a/pkg/options/server.go
+++ b/pkg/options/server.go
@@ -21,6 +21,7 @@ func GetDefaultAccessControlOptions() http.AccessControlOptions {
 	return http.AccessControlOptions{
 		ValidationTokenSecret:     GetDefaultServeOptionString("SERVER_VALIDATION_TOKEN_SECRET", ""),
 		ValidationTokenExpiration: GetDefaultServeOptionInt("SERVER_VALIDATION_TOKEN_EXPIRATION", 604800), // one week
+		ValidationTokenKid:        GetDefaultServeOptionString("SERVER_VALIDATION_TOKEN_KID", ""),
 	}
 }
 
@@ -54,6 +55,11 @@ func AddServerCliFlags(cmd *cobra.Command, serverOptions *http.ServerOptions) {
 		serverOptions.AccessControl.ValidationTokenExpiration,
 		`Validation service JWT expiration in seconds (SERVER_VALIDATION_TOKEN_EXPIRATION).`,
 	)
+	cmd.PersistentFlags().StringVar(
+		&serverOptions.AccessControl.ValidationTokenKid, "server-validation-token-kid",
+		serverOptions.AccessControl.ValidationTokenKid,
+		`Key ID header for validation service JWTs (SERVER_VALIDATION_TOKEN_KID).`,
+	)
 	cmd.PersistentFlags().IntVar(
 		&serverOptions.RateLimiter.RequestLimit, "server-rate-request-limit", serverOptions.RateLimiter.RequestLimit,
 		`The max requests over the rate window length (SERVER_RATE_REQUEST_LIMIT).`,
@@ -71,5 +77,8 @@ func CheckServerOptions(options http.ServerOptions) error {
 	if options.AccessControl.ValidationTokenSecret == "" {
 		return fmt.Errorf("SERVER_VALIDATION_TOKEN_SECRET is required")
 	}
+	if options.AccessControl.ValidationTokenKid == "" {
+		return fmt.Errorf("SERVER_VALIDATION_TOKEN_KID is required")
+	}
 	return nil
 }
diff --git a/pkg/solver/server.go b/pkg/solver/server.go
index 0d7a53e5..37887542 100644
--- a/pkg/solver/server.go
+++ b/pkg/solver/server.go
@@ -647,7 +647,7 @@ func (solverServer *solverServer) getValidationToken(res corehttp.ResponseWriter
 	})
 
 	// Add the key ID to the token header
-	token.Header["kid"] = "key-1"
+	token.Header["kid"] = solverServer.options.AccessControl.ValidationTokenKid
 
 	// Sign the token
 	secret := []byte(solverServer.options.AccessControl.ValidationTokenSecret)
diff --git a/stack b/stack
index 51db9439..26e276f8 100755
--- a/stack
+++ b/stack
@@ -208,6 +208,7 @@ function solver() {
   export STORE_CONN_STR=postgres://postgres:postgres@localhost:5432/solver-db?sslmode=disable
   export STORE_GORM_LOG_LEVEL=silent
   export SERVER_VALIDATION_TOKEN_SECRET=912dd001a6613632c066ca10a19254430db2986a84612882a18f838a6360880e
+  export SERVER_VALIDATION_TOKEN_KID=key-dev
   go run . solver --network dev "$@"
 }