feat: support multiple jwks endpoints

Author: adrianmxb

cmd/serve.go

@@ -4,11 +4,12 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	_ "net/http/pprof"
 	"os"
 	"path/filepath"
-	_ "net/http/pprof"
 	"runtime"
 	"slices"
+	"strings"
 	"sync"
 	"time"
 
@@ -24,7 +25,8 @@ import (
 	"go.uber.org/zap"
 )
 
-var jwksUrl, userId string
+var jwksUrls []string
+var userId string
 var ignoreVersionCheck bool
 var port int
 var shutdownWait int
@@ -55,7 +57,8 @@ to interact and monitor the Scroll Application`,
 		}
 
 		logger.Log().Info("Starting Scroll Daemon")
-		authorizer, err := services.NewAuthorizer(jwksUrl, userId)
+		jwksURLs := buildJWKSURLs(jwksUrls)
+		authorizer, err := services.NewAuthorizer(jwksURLs, userId)
 		if err != nil {
 			return err
 		}
@@ -154,7 +157,7 @@ to interact and monitor the Scroll Application`,
 		signalHandler := signals.NewSignalHandler(ctx, queueManager, processManager, nil, shutdownWait)
 		daemonHander := handler.NewDaemonHandler(signalHandler)
 
-		s := web.NewServer(jwksUrl, scrollHandler, scrollLogHandler, scrollMetricHandler, annotationHandler, processHandler, queueHandler, websocketHandler, portHandler, healthHandler, coldstarterHandler, daemonHander, authorizer, uiDevHandler, cwd, scrollService.GetDir())
+		s := web.NewServer(jwksURLs, scrollHandler, scrollLogHandler, scrollMetricHandler, annotationHandler, processHandler, queueHandler, websocketHandler, portHandler, healthHandler, coldstarterHandler, daemonHander, authorizer, uiDevHandler, cwd, scrollService.GetDir())
 
 		a := s.Initialize()
 
@@ -312,16 +315,35 @@ to interact and monitor the Scroll Application`,
 	},
 }
 
+func buildJWKSURLs(values []string) []string {
+	urls := make([]string, 0, len(values))
+	seen := make(map[string]struct{}, len(values))
+
+	for _, url := range values {
+		url = strings.TrimSpace(url)
+		if url == "" {
+			continue
+		}
+		if _, ok := seen[url]; ok {
+			continue
+		}
+		seen[url] = struct{}{}
+		urls = append(urls, url)
+	}
+
+	return urls
+}
+
 func init() {
 	ServeCommand.Flags().StringVarP(&pprofBind, "pprof", "", "", "Enable pprof on the given bind. This is useful for debugging purposes. E.g. --pprof=localhost:6060 or --pprof=:6060")
 
 	ServeCommand.Flags().IntVarP(&port, "port", "p", 8081, "Port")
 
 	ServeCommand.Flags().IntVarP(&shutdownWait, "shutdown-wait", "", 10, "Wait interval how long the process is allowed to shutdown. First normal shutdown, then forced shutdown")
 
-	ServeCommand.Flags().StringVarP(&jwksUrl, "jwks-server", "", "", "JWKS Server to authenticate requests against")
+	ServeCommand.Flags().StringSliceVarP(&jwksUrls, "jwks-server", "", nil, "JWKS servers to authenticate requests against. Can be comma-separated or set multiple times")
 
-	ServeCommand.Flags().StringVarP(&userId, "user-id", "u", "", "Allowed user ID, if JWKS is not set. It checks claims.sub of the JWT token")
+	ServeCommand.Flags().StringVarP(&userId, "user-id", "u", "", "Allowed user ID. When JWKS authentication is enabled, checks claims.sub of the JWT token")
 
 	ServeCommand.Flags().BoolVarP(&idleScroll, "idle", "", false, "Don't start the queue manager, just use coldstarter")
 

cmd/serve_test.go

@@ -0,0 +1,34 @@
+package cmd
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestBuildJWKSURLsNormalizesAndDeduplicates(t *testing.T) {
+	got := buildJWKSURLs([]string{
+		" https://old.example/jwks ",
+		"",
+		"https://api.druid.gg/auth/v2/jwks",
+		"https://next.example/jwks",
+		" https://old.example/jwks ",
+		"https://api.druid.gg/auth/v2/jwks",
+	})
+
+	want := []string{
+		"https://old.example/jwks",
+		"https://api.druid.gg/auth/v2/jwks",
+		"https://next.example/jwks",
+	}
+
+	if !reflect.DeepEqual(got, want) {
+		t.Fatalf("buildJWKSURLs() = %#v, want %#v", got, want)
+	}
+}
+
+func TestBuildJWKSURLsEmptyWhenUnset(t *testing.T) {
+	got := buildJWKSURLs([]string{"", "   "})
+	if len(got) != 0 {
+		t.Fatalf("buildJWKSURLs() = %#v, want empty slice", got)
+	}
+}

cmd/server/web/server.go

@@ -45,7 +45,7 @@ type Server struct {
 }
 
 func NewServer(
-	jwlsUrl string,
+	jwksURLs []string,
 	scrollHandler ports.ScrollHandlerInterface,
 	scrollLogHandler ports.ScrollLogHandlerInterface,
 	scrollMetricHandler ports.ScrollMetricHandlerInterface,
@@ -88,9 +88,9 @@ func NewServer(
 		watchHandler:                  watchHandler,
 	}
 
-	if jwlsUrl != "" {
+	if len(jwksURLs) > 0 {
 		server.jwtMiddleware = jwtware.New(jwtware.Config{
-			KeySetURLs: []string{jwlsUrl},
+			KeySetURLs: jwksURLs,
 		})
 	}
 

internal/core/services/authorizer_service.go

@@ -3,6 +3,7 @@ package services
 import (
 	"encoding/json"
 	"errors"
+	"fmt"
 	"strings"
 	"time"
 
@@ -16,13 +17,14 @@ import (
 )
 
 type AuthorizerService struct {
-	jwksUrl string
-	jwks    *keyfunc.JWKS
-	userId  string
-	tokens  map[string]time.Time
+	jwksUrls []string
+	jwks     []*keyfunc.JWKS
+	userId   string
+	tokens   map[string]time.Time
 }
 
-func NewAuthorizer(jwksURL string, userId string) (ports.AuthorizerServiceInterface, error) {
+func NewAuthorizer(jwksURLs []string, userId string) (ports.AuthorizerServiceInterface, error) {
+	jwksURLs = normalizeJWKSURLs(jwksURLs)
 
 	// Create the keyfunc options. Refresh the JWKS every hour and log errors.
 	var refreshInterval = time.Hour
@@ -33,18 +35,22 @@ func NewAuthorizer(jwksURL string, userId string) (ports.AuthorizerServiceInterf
 		},
 	}
 
-	if jwksURL != "" {
-		// Create the JWKS from the resource at the given URL.
-		var jwks, err = keyfunc.Get(jwksURL, options)
-		if err != nil {
-			return nil, err
+	if len(jwksURLs) > 0 {
+		jwksList := make([]*keyfunc.JWKS, 0, len(jwksURLs))
+		for _, jwksURL := range jwksURLs {
+			// Create the JWKS from the resource at the given URL.
+			jwks, err := keyfunc.Get(jwksURL, options)
+			if err != nil {
+				return nil, fmt.Errorf("failed to get JWKS from %q: %w", jwksURL, err)
+			}
+			jwksList = append(jwksList, jwks)
 		}
 
 		return &AuthorizerService{
-			jwks:    jwks,
-			jwksUrl: jwksURL,
-			userId:  userId,
-			tokens:  make(map[string]time.Time),
+			jwks:     jwksList,
+			jwksUrls: jwksURLs,
+			userId:   userId,
+			tokens:   make(map[string]time.Time),
 		}, nil
 
 	} else {
@@ -54,9 +60,26 @@ func NewAuthorizer(jwksURL string, userId string) (ports.AuthorizerServiceInterf
 	}
 }
 
+func normalizeJWKSURLs(jwksURLs []string) []string {
+	normalized := make([]string, 0, len(jwksURLs))
+	seen := make(map[string]struct{}, len(jwksURLs))
+	for _, jwksURL := range jwksURLs {
+		jwksURL = strings.TrimSpace(jwksURL)
+		if jwksURL == "" {
+			continue
+		}
+		if _, ok := seen[jwksURL]; ok {
+			continue
+		}
+		seen[jwksURL] = struct{}{}
+		normalized = append(normalized, jwksURL)
+	}
+	return normalized
+}
+
 func (auth *AuthorizerService) CheckHeader(c *fiber.Ctx) (*time.Time, error) {
 
-	if auth.jwksUrl == "" {
+	if len(auth.jwksUrls) == 0 {
 		return nil, nil
 	}
 
@@ -71,8 +94,7 @@ func (auth *AuthorizerService) CheckHeader(c *fiber.Ctx) (*time.Time, error) {
 	jwtToken := splitToken[1]
 
 	// Parse the JWT.
-	token, err := jwt.Parse(jwtToken, auth.jwks.Keyfunc)
-
+	token, err := auth.parseJWT(jwtToken)
 	if err != nil {
 		return nil, errors.New("Failed to parse the JWT.\nError: " + err.Error())
 	}
@@ -105,6 +127,22 @@ func (auth *AuthorizerService) CheckHeader(c *fiber.Ctx) (*time.Time, error) {
 	return &tm, nil
 }
 
+func (auth *AuthorizerService) parseJWT(jwtToken string) (*jwt.Token, error) {
+	var parseErrors []string
+	for index, jwks := range auth.jwks {
+		token, err := jwt.Parse(jwtToken, jwks.Keyfunc)
+		if err == nil && token.Valid {
+			return token, nil
+		}
+		if err != nil {
+			parseErrors = append(parseErrors, fmt.Sprintf("JWKS %d: %s", index+1, err.Error()))
+			continue
+		}
+		parseErrors = append(parseErrors, fmt.Sprintf("JWKS %d: token is not valid", index+1))
+	}
+	return nil, errors.New(strings.Join(parseErrors, "; "))
+}
+
 func (auth *AuthorizerService) CheckQuery(token string) (*time.Time, error) {
 	if validUntil, ok := auth.tokens[token]; ok {
 		defer delete(auth.tokens, token)