Add --discovery-base-url flag to policy server

Enables serving discovery URLs through a CDN by allowing the policy server
to return absolute URLs instead of relative ones. When set, both the Link
header and redirect Location use the configured base URL.

Author: brianm

.tasks/closed/307ktjx8.md

@@ -0,0 +1,26 @@
+---
+yatl_version: 1
+title: Add --discovery-base-url flag to policy server CLI
+id: 307ktjx8
+created: 2025-12-28T00:14:59.357075Z
+updated: 2025-12-28T00:20:29.291136Z
+author: Brian McCallister
+priority: high
+tags:
+- feature
+---
+
+---
+# Log: 2025-12-28T00:14:59Z Brian McCallister
+
+Created task.
+
+---
+# Log: 2025-12-28T00:17:17Z Brian McCallister
+
+Started working.
+
+---
+# Log: 2025-12-28T00:20:29Z Brian McCallister
+
+Closed: Implemented --discovery-base-url flag for CDN support

cmd/epithet/policy.go

@@ -40,6 +40,9 @@ type PolicyServerCLI struct {
 
 	// Default expiration
 	DefaultExpiration string `help:"Default certificate expiration (e.g., 5m)" name:"default-expiration"`
+
+	// Discovery base URL for CDN support
+	DiscoveryBaseURL string `help:"Base URL for discovery endpoints (e.g., https://cdn.example.com)" name:"discovery-base-url"`
 }
 
 func (c *PolicyServerCLI) Run(logger *slog.Logger, tlsCfg tlsconfig.Config, unifiedConfig cue.Value) error {
@@ -82,10 +85,11 @@ func (c *PolicyServerCLI) Run(logger *slog.Logger, tlsCfg tlsconfig.Config, unif
 
 	// Create policy server handler
 	handler := policyserver.NewHandler(policyserver.Config{
-		CAPublicKey:   sshcert.RawPublicKey(caPubkey),
-		Validator:     validator,
-		Evaluator:     eval,
-		DiscoveryHash: cfg.DiscoveryHash(),
+		CAPublicKey:      sshcert.RawPublicKey(caPubkey),
+		Validator:        validator,
+		Evaluator:        eval,
+		DiscoveryHash:    cfg.DiscoveryHash(),
+		DiscoveryBaseURL: c.DiscoveryBaseURL,
 	})
 
 	// Create discovery handler
@@ -110,7 +114,7 @@ func (c *PolicyServerCLI) Run(logger *slog.Logger, tlsCfg tlsconfig.Config, unif
 
 	r.Post("/", handler)
 	// Redirect endpoint: /d/current -> /d/{hash} (cached 5 min)
-	r.Get("/d/current", policyserver.NewDiscoveryRedirectHandler(cfg.DiscoveryHash()))
+	r.Get("/d/current", policyserver.NewDiscoveryRedirectHandler(cfg.DiscoveryHash(), c.DiscoveryBaseURL))
 	// Content-addressed endpoint: /d/{hash} (immutable)
 	r.Get("/d/*", discoveryHandler)
 

pkg/policyserver/discovery.go

@@ -43,10 +43,15 @@ func NewDiscoveryHandler(config DiscoveryConfig) http.HandlerFunc {
 // The redirect response is cached for 5 minutes to allow policy changes to propagate.
 // Clients should request /d/current and follow the redirect to /d/{hash}.
 // Uses 302 Found (temporary) rather than 301 (permanent) since the redirect target may change.
-func NewDiscoveryRedirectHandler(hash string) http.HandlerFunc {
+// If baseURL is set, redirects to an absolute URL on that base; otherwise uses relative URLs.
+func NewDiscoveryRedirectHandler(hash string, baseURL string) http.HandlerFunc {
+	location := "/d/" + hash
+	if baseURL != "" {
+		location = strings.TrimSuffix(baseURL, "/") + "/d/" + hash
+	}
 	return func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Cache-Control", "max-age=300")
-		w.Header().Set("Location", "/d/"+hash)
+		w.Header().Set("Location", location)
 		w.WriteHeader(http.StatusFound)
 	}
 }

pkg/policyserver/discovery_test.go

@@ -175,7 +175,7 @@ func TestDiscoveryHandler_EmptyPatterns(t *testing.T) {
 }
 
 func TestDiscoveryRedirectHandler(t *testing.T) {
-	handler := policyserver.NewDiscoveryRedirectHandler("abc123def456")
+	handler := policyserver.NewDiscoveryRedirectHandler("abc123def456", "")
 
 	req := httptest.NewRequest(http.MethodGet, "/d/current", nil)
 	w := httptest.NewRecorder()
@@ -199,3 +199,38 @@ func TestDiscoveryRedirectHandler(t *testing.T) {
 		t.Errorf("expected Cache-Control 'max-age=300', got %q", cacheControl)
 	}
 }
+
+func TestDiscoveryRedirectHandler_WithBaseURL(t *testing.T) {
+	handler := policyserver.NewDiscoveryRedirectHandler("abc123def456", "https://cdn.example.com")
+
+	req := httptest.NewRequest(http.MethodGet, "/d/current", nil)
+	w := httptest.NewRecorder()
+
+	handler(w, req)
+
+	// Check status code is 302 Found (temporary redirect)
+	if w.Code != http.StatusFound {
+		t.Errorf("expected status 302, got %d", w.Code)
+	}
+
+	// Check Location header points to absolute URL on base
+	location := w.Header().Get("Location")
+	if location != "https://cdn.example.com/d/abc123def456" {
+		t.Errorf("expected Location 'https://cdn.example.com/d/abc123def456', got %q", location)
+	}
+}
+
+func TestDiscoveryRedirectHandler_WithBaseURLTrailingSlash(t *testing.T) {
+	handler := policyserver.NewDiscoveryRedirectHandler("abc123def456", "https://cdn.example.com/")
+
+	req := httptest.NewRequest(http.MethodGet, "/d/current", nil)
+	w := httptest.NewRecorder()
+
+	handler(w, req)
+
+	// Check Location header correctly handles trailing slash
+	location := w.Header().Get("Location")
+	if location != "https://cdn.example.com/d/abc123def456" {
+		t.Errorf("expected Location 'https://cdn.example.com/d/abc123def456', got %q", location)
+	}
+}

pkg/policyserver/policyserver.go

@@ -114,6 +114,11 @@ type Config struct {
 	// If empty, no Link header is set.
 	// The path is hardcoded to "/d/" + hash.
 	DiscoveryHash string
+
+	// DiscoveryBaseURL is the base URL for discovery endpoints.
+	// If set, discovery URLs will be absolute URLs on this base (e.g., "https://cdn.example.com").
+	// If empty, discovery URLs will be relative (e.g., "/d/current").
+	DiscoveryBaseURL string
 }
 
 // handler holds the config and implements the HTTP handler methods
@@ -217,9 +222,14 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 // setDiscoveryHeader sets the Link header for discovery if configured.
 // Always points to /d/current which redirects to the content-addressed URL.
+// If DiscoveryBaseURL is set, uses absolute URLs; otherwise uses relative URLs.
 func (h *handler) setDiscoveryHeader(w http.ResponseWriter) {
 	if h.config.DiscoveryHash != "" {
-		w.Header().Set("Link", "</d/current>; rel=\"discovery\"")
+		url := "/d/current"
+		if h.config.DiscoveryBaseURL != "" {
+			url = strings.TrimSuffix(h.config.DiscoveryBaseURL, "/") + "/d/current"
+		}
+		w.Header().Set("Link", "<"+url+">; rel=\"discovery\"")
 	}
 }
 

pkg/policyserver/policyserver_test.go

@@ -399,3 +399,98 @@ func TestHandler_DiscoveryLinkHeader_NotSetWhenEmpty(t *testing.T) {
 		t.Errorf("expected no Link header, got %q", link)
 	}
 }
+
+func TestHandler_DiscoveryLinkHeader_WithBaseURL(t *testing.T) {
+	evaluator := &mockEvaluator{
+		response: &policyserver.Response{
+			CertParams: ca.CertParams{
+				Identity:   "test@example.com",
+				Names:      []string{"testuser"},
+				Expiration: 5 * time.Minute,
+			},
+			Policy: policy.Policy{
+				HostUsers: map[string][]string{
+					"*": {"testuser"},
+				},
+			},
+		},
+	}
+
+	handler := policyserver.NewHandler(policyserver.Config{
+		Validator:        &mockValidator{},
+		Evaluator:        evaluator,
+		DiscoveryHash:    "abc123def456",
+		DiscoveryBaseURL: "https://cdn.example.com",
+	})
+
+	req := policyserver.Request{
+		Token: encodeToken("test-token"),
+		Connection: policy.Connection{
+			RemoteHost: "server.example.com",
+			RemoteUser: "testuser",
+			Port:       22,
+		},
+	}
+	body, _ := json.Marshal(req)
+
+	httpReq := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(body))
+	w := httptest.NewRecorder()
+
+	handler(w, httpReq)
+
+	if w.Code != http.StatusOK {
+		t.Errorf("expected status 200, got %d: %s", w.Code, w.Body.String())
+	}
+
+	link := w.Header().Get("Link")
+	// Link header uses absolute URL when DiscoveryBaseURL is set
+	expected := "<https://cdn.example.com/d/current>; rel=\"discovery\""
+	if link != expected {
+		t.Errorf("expected Link header %q, got %q", expected, link)
+	}
+}
+
+func TestHandler_DiscoveryLinkHeader_WithBaseURLTrailingSlash(t *testing.T) {
+	evaluator := &mockEvaluator{
+		response: &policyserver.Response{
+			CertParams: ca.CertParams{
+				Identity:   "test@example.com",
+				Names:      []string{"testuser"},
+				Expiration: 5 * time.Minute,
+			},
+		},
+	}
+
+	handler := policyserver.NewHandler(policyserver.Config{
+		Validator:        &mockValidator{},
+		Evaluator:        evaluator,
+		DiscoveryHash:    "abc123def456",
+		DiscoveryBaseURL: "https://cdn.example.com/", // trailing slash
+	})
+
+	req := policyserver.Request{
+		Token: encodeToken("test-token"),
+		Connection: policy.Connection{
+			RemoteHost: "server.example.com",
+			RemoteUser: "testuser",
+			Port:       22,
+		},
+	}
+	body, _ := json.Marshal(req)
+
+	httpReq := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(body))
+	w := httptest.NewRecorder()
+
+	handler(w, httpReq)
+
+	if w.Code != http.StatusOK {
+		t.Errorf("expected status 200, got %d: %s", w.Code, w.Body.String())
+	}
+
+	link := w.Header().Get("Link")
+	// Trailing slash should be stripped to avoid double slashes
+	expected := "<https://cdn.example.com/d/current>; rel=\"discovery\""
+	if link != expected {
+		t.Errorf("expected Link header %q, got %q", expected, link)
+	}
+}