feat: add per-user tenant scoping
This commit is contained in:
@@ -11,6 +11,7 @@ import (
|
||||
"git.warky.dev/wdevs/amcs/internal/config"
|
||||
"git.warky.dev/wdevs/amcs/internal/observability"
|
||||
"git.warky.dev/wdevs/amcs/internal/requestip"
|
||||
"git.warky.dev/wdevs/amcs/internal/tenancy"
|
||||
)
|
||||
|
||||
type contextKey string
|
||||
@@ -50,6 +51,10 @@ func Middleware(cfg config.AuthConfig, keyring *Keyring, oauthRegistry *OAuthReg
|
||||
)
|
||||
}
|
||||
}
|
||||
withTenant := func(ctx context.Context, keyID string) context.Context {
|
||||
ctx = context.WithValue(ctx, keyIDContextKey, keyID)
|
||||
return tenancy.WithTenantKey(ctx, keyID)
|
||||
}
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
remoteAddr := requestip.FromRequest(r)
|
||||
@@ -63,7 +68,7 @@ func Middleware(cfg config.AuthConfig, keyring *Keyring, oauthRegistry *OAuthReg
|
||||
return
|
||||
}
|
||||
recordAccess(r, keyID)
|
||||
next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID)))
|
||||
next.ServeHTTP(w, r.WithContext(withTenant(r.Context(), keyID)))
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -73,14 +78,14 @@ func Middleware(cfg config.AuthConfig, keyring *Keyring, oauthRegistry *OAuthReg
|
||||
if tokenStore != nil {
|
||||
if keyID, ok := tokenStore.Lookup(bearer); ok {
|
||||
recordAccess(r, keyID)
|
||||
next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID)))
|
||||
next.ServeHTTP(w, r.WithContext(withTenant(r.Context(), keyID)))
|
||||
return
|
||||
}
|
||||
}
|
||||
if keyring != nil {
|
||||
if keyID, ok := keyring.Lookup(bearer); ok {
|
||||
recordAccess(r, keyID)
|
||||
next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID)))
|
||||
next.ServeHTTP(w, r.WithContext(withTenant(r.Context(), keyID)))
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -103,7 +108,7 @@ func Middleware(cfg config.AuthConfig, keyring *Keyring, oauthRegistry *OAuthReg
|
||||
return
|
||||
}
|
||||
recordAccess(r, keyID)
|
||||
next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID)))
|
||||
next.ServeHTTP(w, r.WithContext(withTenant(r.Context(), keyID)))
|
||||
return
|
||||
}
|
||||
|
||||
@@ -117,7 +122,7 @@ func Middleware(cfg config.AuthConfig, keyring *Keyring, oauthRegistry *OAuthReg
|
||||
return
|
||||
}
|
||||
recordAccess(r, keyID)
|
||||
next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID)))
|
||||
next.ServeHTTP(w, r.WithContext(withTenant(r.Context(), keyID)))
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"git.warky.dev/wdevs/amcs/internal/config"
|
||||
"git.warky.dev/wdevs/amcs/internal/tenancy"
|
||||
)
|
||||
|
||||
func TestMiddlewareAddsTenantKeyToContext(t *testing.T) {
|
||||
keyring, err := NewKeyring([]config.APIKey{{ID: "user-a", Value: "secret-a"}})
|
||||
if err != nil {
|
||||
t.Fatalf("NewKeyring error = %v", err)
|
||||
}
|
||||
|
||||
var gotKeyID, gotTenant string
|
||||
handler := Middleware(config.AuthConfig{}, keyring, nil, nil, nil, slog.Default())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var ok bool
|
||||
gotKeyID, ok = KeyIDFromContext(r.Context())
|
||||
if !ok {
|
||||
t.Fatal("KeyIDFromContext ok = false")
|
||||
}
|
||||
gotTenant, ok = tenancy.KeyFromContext(r.Context())
|
||||
if !ok {
|
||||
t.Fatal("tenancy.KeyFromContext ok = false")
|
||||
}
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}))
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/mcp", nil)
|
||||
req.Header.Set("x-brain-key", "secret-a")
|
||||
rr := httptest.NewRecorder()
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
if rr.Code != http.StatusNoContent {
|
||||
t.Fatalf("status = %d, want %d", rr.Code, http.StatusNoContent)
|
||||
}
|
||||
if gotKeyID != "user-a" || gotTenant != "user-a" {
|
||||
t.Fatalf("keyID=%q tenant=%q, want user-a/user-a", gotKeyID, gotTenant)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user