Apply API key middleware to /v2/check and /v2/authorize

errors.go

@@ -22,6 +22,7 @@ import (
 
 const (
 	ErrorCacheNotReady = "cache_not_ready"
+	ErrorUnauthorized  = "unauthorized"
 )
 
 // CacheNotReady type
@@ -39,3 +40,19 @@ func NewCacheNotReady() *CacheNotReady {
 		),
 	}
 }
+
+// Unauthorized type
+type Unauthorized struct {
+	*service.GenericError
+}
+
+func NewUnauthorized() *Unauthorized {
+	return &Unauthorized{
+		GenericError: service.NewGenericError(
+			"Unauthorized",
+			ErrorUnauthorized,
+			http.StatusUnauthorized,
+			"Unauthorized",
+		),
+	}
+}

middleware.go

@@ -15,11 +15,14 @@
 package edge
 
 import (
+	"crypto/subtle"
 	"fmt"
 	"log"
 	"net"
 	"net/http"
 	"time"
+
+	"github.com/warrant-dev/warrant/pkg/service"
 )
 
 type loggingResponseWriter struct {
@@ -36,6 +39,24 @@ func (rw *loggingResponseWriter) WriteHeader(code int) {
 	rw.ResponseWriter.WriteHeader(code)
 }
 
+// authMiddleware enforces that inbound requests present the configured API key
+// via an "Authorization: ApiKey <key>" header, matching the WorkOS cloud API.
+// Requests without a valid credential are rejected with 401 before any handler
+// runs. If no API key is configured, the agent fails closed and rejects all
+// requests rather than serving the authorization graph anonymously.
+func authMiddleware(apiKey string, next http.Handler) http.Handler {
+	expected := fmt.Sprintf("ApiKey %s", apiKey)
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		provided := r.Header.Get("Authorization")
+		if apiKey == "" || subtle.ConstantTimeCompare([]byte(provided), []byte(expected)) != 1 {
+			service.SendErrorResponse(w, NewUnauthorized())
+			return
+		}
+
+		next.ServeHTTP(w, r)
+	})
+}
+
 func loggingMiddleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		startTime := time.Now()

server.go

@@ -147,12 +147,15 @@ func (server *Server) check(w http.ResponseWriter, r *http.Request) {
 	})
 }
 
-func (server *Server) Run() error {
+func (server *Server) Handler() http.Handler {
 	mux := http.NewServeMux()
 	mux.Handle("/health", loggingMiddleware(http.HandlerFunc(server.health)))
-	mux.Handle(fmt.Sprintf("/%s/authorize", ApiVersion), loggingMiddleware(http.HandlerFunc(server.check)))
-	mux.Handle(fmt.Sprintf("/%s/check", ApiVersion), loggingMiddleware(http.HandlerFunc(server.check)))
+	mux.Handle(fmt.Sprintf("/%s/authorize", ApiVersion), loggingMiddleware(authMiddleware(server.config.ApiKey, http.HandlerFunc(server.check))))
+	mux.Handle(fmt.Sprintf("/%s/check", ApiVersion), loggingMiddleware(authMiddleware(server.config.ApiKey, http.HandlerFunc(server.check))))
+	return mux
+}
 
+func (server *Server) Run() error {
 	log.Printf("Edge agent ready to serve authz requests on port %d", server.config.Port)
-	return http.ListenAndServe(fmt.Sprintf(":%d", server.config.Port), mux)
+	return http.ListenAndServe(fmt.Sprintf(":%d", server.config.Port), server.Handler())
 }

server_test.go

@@ -0,0 +1,124 @@
+// Copyright 2024 WorkOS, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package edge
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+)
+
+const testApiKey = "sk_test_edge_agent_key"
+
+// warrant present in the seeded cache
+const seededWarrant = "role:admin#member@user:alice@example.com"
+
+func newTestServer(t *testing.T) *httptest.Server {
+	t.Helper()
+	repo := NewMemoryRepository()
+	if err := repo.Set(seededWarrant, 1); err != nil {
+		t.Fatalf("failed to seed repository: %v", err)
+	}
+	server, err := NewServer(ServerConfig{
+		ApiKey:     testApiKey,
+		Port:       3000,
+		Repository: repo,
+	})
+	if err != nil {
+		t.Fatalf("failed to create server: %v", err)
+	}
+	return httptest.NewServer(server.Handler())
+}
+
+func postCheck(t *testing.T, baseURL, authHeader, body string) int {
+	t.Helper()
+	req, err := http.NewRequest(http.MethodPost, baseURL+"/v2/check", strings.NewReader(body))
+	if err != nil {
+		t.Fatalf("failed to build request: %v", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+	if authHeader != "" {
+		req.Header.Set("Authorization", authHeader)
+	}
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		t.Fatalf("request failed: %v", err)
+	}
+	defer resp.Body.Close()
+	return resp.StatusCode
+}
+
+// A request for the seeded warrant that an attacker-observable oracle would
+// answer "Authorized" (HTTP 200) if it reached the handler.
+const checkBody = `{"warrants":[{"objectType":"role","objectId":"admin","relation":"member","subject":{"objectType":"user","objectId":"alice@example.com"}}]}`
+
+func TestCheckRejectsUnauthenticatedRequest(t *testing.T) {
+	ts := newTestServer(t)
+	defer ts.Close()
+
+	if status := postCheck(t, ts.URL, "", checkBody); status != http.StatusUnauthorized {
+		t.Fatalf("unauthenticated /v2/check: expected 401, got %d", status)
+	}
+}
+
+func TestCheckRejectsWrongApiKey(t *testing.T) {
+	ts := newTestServer(t)
+	defer ts.Close()
+
+	if status := postCheck(t, ts.URL, "ApiKey wrong", checkBody); status != http.StatusUnauthorized {
+		t.Fatalf("wrong-key /v2/check: expected 401, got %d", status)
+	}
+}
+
+func TestCheckAllowsValidApiKey(t *testing.T) {
+	ts := newTestServer(t)
+	defer ts.Close()
+
+	status := postCheck(t, ts.URL, "ApiKey "+testApiKey, checkBody)
+	if status != http.StatusOK {
+		t.Fatalf("authenticated /v2/check: expected 200, got %d", status)
+	}
+}
+
+func TestAuthorizeRejectsUnauthenticatedRequest(t *testing.T) {
+	ts := newTestServer(t)
+	defer ts.Close()
+
+	req, _ := http.NewRequest(http.MethodPost, ts.URL+"/v2/authorize", strings.NewReader(checkBody))
+	req.Header.Set("Content-Type", "application/json")
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		t.Fatalf("request failed: %v", err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusUnauthorized {
+		t.Fatalf("unauthenticated /v2/authorize: expected 401, got %d", resp.StatusCode)
+	}
+}
+
+func TestHealthRemainsUnauthenticated(t *testing.T) {
+	ts := newTestServer(t)
+	defer ts.Close()
+
+	resp, err := http.Get(ts.URL + "/health")
+	if err != nil {
+		t.Fatalf("health request failed: %v", err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("/health: expected 200, got %d", resp.StatusCode)
+	}
+}