From 58ab3ed0b336b823713c649e629586c764470e02 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Thu, 20 Aug 2026 00:42:50 +0100 Subject: [PATCH] feat: auth store with constant-time key identification and group allow rules --- auth.go | 217 ++++++++++++++++++++++++++++++++++++++ auth_test.go | 282 ++++++++++++++++++++++++++++++++++++++++++++++++++ llmgateway.go | 11 +- 3 files changed, 507 insertions(+), 3 deletions(-) create mode 100644 auth.go create mode 100644 auth_test.go diff --git a/auth.go b/auth.go new file mode 100644 index 0000000..94e9231 --- /dev/null +++ b/auth.go @@ -0,0 +1,217 @@ +package traefikllmgateway + +import ( + "crypto/sha256" + "crypto/subtle" + "fmt" + "net/http" + "path" + "strings" + "sync" +) + +// user is an authenticated API-key holder resolved by authStore.identify. +type user struct { + limits *LimitsConfig + name, groupName string +} + +// group describes a group's access rules and default limits. Empty +// providers, models, mcpServers, or agents means all are allowed. +type group struct { + limits *LimitsConfig + name string + providers, models, mcpServers, agents []string +} + +// allowsProvider reports whether name matches one of the group's provider +// glob patterns. An empty pattern list allows every provider. +func (grp *group) allowsProvider(name string) bool { + return matchesGlob(grp.providers, name) +} + +// allowsModel reports whether id matches one of the group's model glob +// patterns. An empty pattern list allows every model. A pattern is checked +// against id as given and, when id has a "provider/model" form, against the +// part after the "/" — so a pattern like "gpt-5*" matches both "gpt-5-mini" +// and "openai/gpt-5-mini". +func (grp *group) allowsModel(id string) bool { + if matchesGlob(grp.models, id) { + return true + } + if _, bare, found := strings.Cut(id, "/"); found { + return matchesGlob(grp.models, bare) + } + return false +} + +// allowsMCP reports whether name matches one of the group's MCP-server glob +// patterns. An empty pattern list allows every server. +func (grp *group) allowsMCP(name string) bool { + return matchesGlob(grp.mcpServers, name) +} + +// allowsAgent reports whether name matches one of the group's agent glob +// patterns. An empty pattern list allows every agent. +func (grp *group) allowsAgent(name string) bool { + return matchesGlob(grp.agents, name) +} + +// matchesGlob reports whether candidate matches any shell-style pattern in +// patterns (path.Match semantics). An empty pattern list matches anything. +// A malformed pattern (path.ErrBadPattern) is treated as a non-match rather +// than a panic or a silent allow. +func matchesGlob(patterns []string, candidate string) bool { + if len(patterns) == 0 { + return true + } + for _, p := range patterns { + if ok, err := path.Match(p, candidate); ok && err == nil { + return true + } + } + return false +} + +// authEntry is the internal record stored per API-key digest. +type authEntry struct { + user *user + group *group + digest [32]byte +} + +// authStore identifies requests by API key and resolves them to a user and +// their group. Raw API keys are never retained after construction — only +// their SHA-256 digests. +type authStore struct { + groups map[string]*group + inline map[[32]byte]*authEntry // built once by newAuthStore; never mutated afterward + byDigest map[[32]byte]*authEntry // inline entries plus the current file-sourced set; guarded by mu + mu sync.RWMutex +} + +// newAuthStore builds an authStore from cfg's groups and inline users. It +// errors if a user references an unknown group, if a user's resolved API +// key is empty, or if two users resolve to the same API key. +func newAuthStore(cfg *Config) (*authStore, error) { + a := &authStore{ + groups: make(map[string]*group, len(cfg.Groups)), + inline: make(map[[32]byte]*authEntry), + byDigest: make(map[[32]byte]*authEntry), + } + for name, gc := range cfg.Groups { + a.groups[name] = &group{ + limits: gc.Limits, + name: name, + providers: gc.Providers, + models: gc.Models, + mcpServers: gc.MCPServers, + agents: gc.Agents, + } + } + + if cfg.Users != nil { + for _, uc := range cfg.Users.Inline { + entry, err := a.buildEntry(uc) + if err != nil { + return nil, err + } + if _, exists := a.inline[entry.digest]; exists { + return nil, fmt.Errorf("llmgateway: duplicate API key for user %q", uc.Name) + } + a.inline[entry.digest] = entry + a.byDigest[entry.digest] = entry + } + } + return a, nil +} + +// buildEntry resolves uc's group and API key into an authEntry. The API key +// goes through resolveSecret before digesting, so env:/file: references +// work the same as for provider keys. +func (a *authStore) buildEntry(uc *UserConfig) (*authEntry, error) { + grp, ok := a.groups[uc.Group] + if !ok { + return nil, fmt.Errorf("llmgateway: user %q references unknown group %q", uc.Name, uc.Group) + } + key, err := resolveSecret(uc.APIKey) + if err != nil { + return nil, fmt.Errorf("llmgateway: user %q: %w", uc.Name, err) + } + if key == "" { + return nil, fmt.Errorf("llmgateway: user %q has an empty API key", uc.Name) + } + digest := sha256.Sum256([]byte(key)) + return &authEntry{ + digest: digest, + user: &user{limits: uc.Limits, name: uc.Name, groupName: uc.Group}, + group: grp, + }, nil +} + +// replaceFileUsers rebuilds the file-sourced portion of the store from us, +// leaving the inline users untouched. It is used by the config-reload path +// (Task 4) to pick up changes to Users.File without restarting the plugin. +// On error the store is left exactly as it was before the call — the new +// set is validated and built in full before it replaces the old one. +func (a *authStore) replaceFileUsers(us []*UserConfig) error { + next := make(map[[32]byte]*authEntry, len(a.inline)+len(us)) + for digest, entry := range a.inline { + next[digest] = entry + } + for _, uc := range us { + entry, err := a.buildEntry(uc) + if err != nil { + return err + } + if _, exists := next[entry.digest]; exists { + return fmt.Errorf("llmgateway: duplicate API key for user %q", uc.Name) + } + next[entry.digest] = entry + } + + a.mu.Lock() + a.byDigest = next + a.mu.Unlock() + return nil +} + +// identify resolves r's presented API key to a user and their group. It +// reads "Authorization: Bearer " (case-insensitive "bearer" prefix) or +// "x-api-key: "; Bearer wins when both are present. The presented key +// is looked up by SHA-256 digest, then verified with a constant-time +// compare — the key itself is never retained. +func (a *authStore) identify(r *http.Request) (*user, *group, bool) { + key, ok := presentedKey(r) + if !ok { + return nil, nil, false + } + digest := sha256.Sum256([]byte(key)) + + a.mu.RLock() + entry, ok := a.byDigest[digest] + a.mu.RUnlock() + if !ok { + return nil, nil, false + } + if subtle.ConstantTimeCompare(digest[:], entry.digest[:]) != 1 { + return nil, nil, false + } + return entry.user, entry.group, true +} + +// presentedKey extracts the API key from r's Authorization Bearer header or +// its x-api-key header. Bearer takes priority when both are set. +func presentedKey(r *http.Request) (string, bool) { + if authz := r.Header.Get("Authorization"); authz != "" { + if scheme, rest, found := strings.Cut(authz, " "); found && strings.EqualFold(scheme, "bearer") { + if key := strings.TrimSpace(rest); key != "" { + return key, true + } + } + } + if key := strings.TrimSpace(r.Header.Get("x-api-key")); key != "" { + return key, true + } + return "", false +} diff --git a/auth_test.go b/auth_test.go new file mode 100644 index 0000000..6f29a7d --- /dev/null +++ b/auth_test.go @@ -0,0 +1,282 @@ +package traefikllmgateway + +import ( + "net/http/httptest" + "testing" +) + +// TestIdentify is the brief's Step-1 table: bearer auth resolves the right +// user and group, and the group's model glob matches both the bare id and +// the provider-prefixed id. +func TestIdentify(t *testing.T) { + cfg := &Config{ + Providers: map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}, + Groups: map[string]*GroupConfig{"eng": {Models: []string{"gpt-5*"}}}, + Users: &UsersConfig{Inline: []*UserConfig{{Name: "a", Group: "eng", APIKey: "sk-secret"}}}, + } + a, err := newAuthStore(cfg) + if err != nil { + t.Fatal(err) + } + r := httptest.NewRequest("POST", "/v1/chat/completions", nil) + r.Header.Set("Authorization", "Bearer sk-secret") + u, grp, ok := a.identify(r) + if !ok || u.name != "a" || grp.name != "eng" { + t.Fatalf("identify failed: %v %v %v", u, grp, ok) + } + if !grp.allowsModel("gpt-5-mini") { + t.Error("glob should match") + } + if !grp.allowsModel("openai/gpt-5-mini") { + t.Error("prefixed form should match") + } + if grp.allowsModel("claude-4") { + t.Error("should not match") + } +} + +func testAuthCfg() *Config { + return &Config{ + Providers: map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}, + Groups: map[string]*GroupConfig{"eng": {}}, + Users: &UsersConfig{Inline: []*UserConfig{{Name: "a", Group: "eng", APIKey: "sk-secret"}}}, + } +} + +func TestIdentify_XAPIKeyHeader_OK(t *testing.T) { + a, err := newAuthStore(testAuthCfg()) + if err != nil { + t.Fatal(err) + } + r := httptest.NewRequest("POST", "/v1/chat/completions", nil) + r.Header.Set("x-api-key", "sk-secret") + u, grp, ok := a.identify(r) + if !ok || u.name != "a" || grp.name != "eng" { + t.Fatalf("identify via x-api-key failed: %v %v %v", u, grp, ok) + } +} + +func TestIdentify_BearerCaseInsensitivePrefix(t *testing.T) { + a, err := newAuthStore(testAuthCfg()) + if err != nil { + t.Fatal(err) + } + for _, prefix := range []string{"Bearer", "bearer", "BEARER", "BeArEr"} { + r := httptest.NewRequest("POST", "/v1/chat/completions", nil) + r.Header.Set("Authorization", prefix+" sk-secret") + if _, _, ok := a.identify(r); !ok { + t.Errorf("prefix %q: want identify to succeed", prefix) + } + } +} + +func TestIdentify_BearerWinsOverXAPIKey(t *testing.T) { + cfg := &Config{ + Providers: map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}, + Groups: map[string]*GroupConfig{"eng": {}, "ops": {}}, + Users: &UsersConfig{Inline: []*UserConfig{ + {Name: "bearer-user", Group: "eng", APIKey: "sk-bearer"}, + {Name: "header-user", Group: "ops", APIKey: "sk-header"}, + }}, + } + a, err := newAuthStore(cfg) + if err != nil { + t.Fatal(err) + } + r := httptest.NewRequest("POST", "/v1/chat/completions", nil) + r.Header.Set("Authorization", "Bearer sk-bearer") + r.Header.Set("x-api-key", "sk-header") + u, _, ok := a.identify(r) + if !ok || u.name != "bearer-user" { + t.Fatalf("want bearer to win, got %v ok=%v", u, ok) + } +} + +func TestIdentify_WrongKeyFails(t *testing.T) { + a, err := newAuthStore(testAuthCfg()) + if err != nil { + t.Fatal(err) + } + r := httptest.NewRequest("POST", "/v1/chat/completions", nil) + r.Header.Set("Authorization", "Bearer sk-wrong") + if _, _, ok := a.identify(r); ok { + t.Fatal("want identify to fail for wrong key") + } +} + +func TestIdentify_MissingHeaderFails(t *testing.T) { + a, err := newAuthStore(testAuthCfg()) + if err != nil { + t.Fatal(err) + } + r := httptest.NewRequest("POST", "/v1/chat/completions", nil) + if _, _, ok := a.identify(r); ok { + t.Fatal("want identify to fail with no auth header") + } +} + +func TestNewAuthStore_UnknownGroup_ReturnsError(t *testing.T) { + cfg := &Config{ + Providers: map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}, + Groups: map[string]*GroupConfig{"eng": {}}, + Users: &UsersConfig{Inline: []*UserConfig{{Name: "a", Group: "nonexistent", APIKey: "sk-secret"}}}, + } + if _, err := newAuthStore(cfg); err == nil { + t.Fatal("want error for user referencing unknown group") + } +} + +func TestNewAuthStore_DuplicateAPIKey_ReturnsError(t *testing.T) { + cfg := &Config{ + Providers: map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}, + Groups: map[string]*GroupConfig{"eng": {}}, + Users: &UsersConfig{Inline: []*UserConfig{ + {Name: "a", Group: "eng", APIKey: "sk-same"}, + {Name: "b", Group: "eng", APIKey: "sk-same"}, + }}, + } + if _, err := newAuthStore(cfg); err == nil { + t.Fatal("want error for duplicate API key across users") + } +} + +func TestNewAuthStore_EmptyAPIKeyAfterResolution_ReturnsError(t *testing.T) { + cfg := &Config{ + Providers: map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}, + Groups: map[string]*GroupConfig{"eng": {}}, + Users: &UsersConfig{Inline: []*UserConfig{{Name: "a", Group: "eng", APIKey: ""}}}, + } + if _, err := newAuthStore(cfg); err == nil { + t.Fatal("want error for user with empty API key after resolution") + } +} + +func TestNewAuthStore_NoUsers_OK(t *testing.T) { + cfg := &Config{ + Providers: map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}, + } + a, err := newAuthStore(cfg) + if err != nil { + t.Fatalf("newAuthStore: %v", err) + } + r := httptest.NewRequest("POST", "/v1/chat/completions", nil) + r.Header.Set("x-api-key", "anything") + if _, _, ok := a.identify(r); ok { + t.Fatal("want identify to fail when no users are configured") + } +} + +func TestGroup_AllowsProvider_EmptyListAllowsAll(t *testing.T) { + grp := &group{name: "eng"} + if !grp.allowsProvider("openai") { + t.Fatal("empty providers list should allow all") + } +} + +func TestGroup_AllowsProvider_Glob(t *testing.T) { + grp := &group{name: "eng", providers: []string{"open*"}} + if !grp.allowsProvider("openai") { + t.Error("want glob match") + } + if grp.allowsProvider("anthropic") { + t.Error("want no match") + } +} + +func TestGroup_AllowsMCP_Glob(t *testing.T) { + grp := &group{name: "eng", mcpServers: []string{"search-*"}} + if !grp.allowsMCP("search-web") { + t.Error("want glob match") + } + if grp.allowsMCP("db-admin") { + t.Error("want no match") + } +} + +func TestGroup_AllowsAgent_Glob(t *testing.T) { + grp := &group{name: "eng", agents: []string{"triage"}} + if !grp.allowsAgent("triage") { + t.Error("want exact match") + } + if grp.allowsAgent("other") { + t.Error("want no match") + } +} + +func TestGroup_AllowsModel_EmptyListAllowsAnything(t *testing.T) { + grp := &group{name: "eng"} + if !grp.allowsModel("anything-goes") { + t.Fatal("empty models list should allow anything") + } +} + +func TestAuthStore_ReplaceFileUsers_AddsAndSwapsUsers(t *testing.T) { + a, err := newAuthStore(testAuthCfg()) // inline user "a" / sk-secret in group eng + if err != nil { + t.Fatal(err) + } + + if err := a.replaceFileUsers([]*UserConfig{{Name: "f1", Group: "eng", APIKey: "sk-file1"}}); err != nil { + t.Fatalf("replaceFileUsers: %v", err) + } + + inlineReq := httptest.NewRequest("POST", "/v1/chat/completions", nil) + inlineReq.Header.Set("Authorization", "Bearer sk-secret") + if _, _, ok := a.identify(inlineReq); !ok { + t.Fatal("want inline user to remain identifiable after replaceFileUsers") + } + + fileReq := httptest.NewRequest("POST", "/v1/chat/completions", nil) + fileReq.Header.Set("Authorization", "Bearer sk-file1") + u, _, ok := a.identify(fileReq) + if !ok || u.name != "f1" { + t.Fatalf("want file user f1 identifiable, got %v ok=%v", u, ok) + } + + // Second call swaps the file-sourced set wholesale: f1 must disappear, + // f2 must appear, inline user must still work. + if err := a.replaceFileUsers([]*UserConfig{{Name: "f2", Group: "eng", APIKey: "sk-file2"}}); err != nil { + t.Fatalf("replaceFileUsers (second): %v", err) + } + if _, _, ok := a.identify(fileReq); ok { + t.Fatal("want f1 to no longer be identifiable after being replaced") + } + if _, _, ok := a.identify(inlineReq); !ok { + t.Fatal("want inline user to remain identifiable after second replaceFileUsers") + } + f2Req := httptest.NewRequest("POST", "/v1/chat/completions", nil) + f2Req.Header.Set("Authorization", "Bearer sk-file2") + if _, _, ok := a.identify(f2Req); !ok { + t.Fatal("want file user f2 identifiable after second replaceFileUsers") + } +} + +func TestAuthStore_ReplaceFileUsers_DuplicateWithInline_ReturnsErrorAndLeavesStoreUnchanged(t *testing.T) { + a, err := newAuthStore(testAuthCfg()) // inline user "a" / sk-secret + if err != nil { + t.Fatal(err) + } + err = a.replaceFileUsers([]*UserConfig{{Name: "dup", Group: "eng", APIKey: "sk-secret"}}) + if err == nil { + t.Fatal("want error for file user colliding with inline user's API key") + } + + // Store must still be usable and unchanged by the rejected replace. + r := httptest.NewRequest("POST", "/v1/chat/completions", nil) + r.Header.Set("Authorization", "Bearer sk-secret") + u, _, ok := a.identify(r) + if !ok || u.name != "a" { + t.Fatalf("want original inline user still resolvable after rejected replace, got %v ok=%v", u, ok) + } +} + +func TestAuthStore_ReplaceFileUsers_UnknownGroup_ReturnsError(t *testing.T) { + a, err := newAuthStore(testAuthCfg()) + if err != nil { + t.Fatal(err) + } + err = a.replaceFileUsers([]*UserConfig{{Name: "f1", Group: "nonexistent", APIKey: "sk-file1"}}) + if err == nil { + t.Fatal("want error for file user referencing unknown group") + } +} diff --git a/llmgateway.go b/llmgateway.go index 96915dd..a4197ab 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -101,9 +101,10 @@ func CreateConfig() *Config { type Gateway struct { next http.Handler cfg *Config + auth *authStore name string - // wired in later tasks: auth *authStore, limiter *limiter, - // registry *modelRegistry, adapters map[string]providerAdapter + // wired in later tasks: limiter *limiter, registry *modelRegistry, + // adapters map[string]providerAdapter } // New creates the middleware. NOTE: no tail call — Yaegi zeroes @@ -120,7 +121,11 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str if config == nil || len(config.Providers) == 0 { return nil, errors.New("llmgateway: at least one provider must be configured") } - return &Gateway{next: next, name: name, cfg: config}, nil + auth, err := newAuthStore(config) + if err != nil { + return nil, err + } + return &Gateway{next: next, name: name, cfg: config, auth: auth}, nil } // ServeHTTP is the internal router. It grows in later tasks.