Skip to content

Commit 255d775

Browse files
SantiagoDePoloniaSantiagoDePolonia
authored
feat(auth): add managed API key support (#190)
* feat(auth): add managed API key support * fix(auth): address managed key review feedback * fix(auth): tighten postgres auth key handling
1 parent a7609be commit 255d775

29 files changed

Lines changed: 2039 additions & 85 deletions

‎internal/admin/dashboard/templates/index.html‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -975,6 +975,7 @@ <h2>Audit Logs</h2>
975975
<span class="provider-badge audit-alias-badge" x-show="entry.alias_used">alias</span>
976976
<span class="provider-badge mono" x-show="entry.alias_used && entry.resolved_model" x-text="'resolved: ' + entry.resolved_model"></span>
977977
<span class="provider-badge mono" x-text="'request_id: ' + (entry.request_id || '-')"></span>
978+
<span class="provider-badge mono" x-show="entry.auth_key_id" x-text="'auth_key_id: ' + entry.auth_key_id"></span>
978979
<span class="provider-badge mono" x-show="entry.client_ip" x-text="'ip: ' + entry.client_ip"></span>
979980
<span class="provider-badge" x-show="entry.stream">stream</span>
980981
<span class="provider-badge" x-show="entry.error_type" x-text="entry.error_type"></span>

‎internal/admin/handler.go‎

Lines changed: 118 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ package admin
44
import (
55
"context"
66
"errors"
7+
"log/slog"
78
"net/http"
89
"net/url"
910
"slices"
@@ -15,6 +16,7 @@ import (
1516

1617
"gomodel/internal/aliases"
1718
"gomodel/internal/auditlog"
19+
"gomodel/internal/authkeys"
1820
"gomodel/internal/core"
1921
"gomodel/internal/executionplans"
2022
"gomodel/internal/guardrails"
@@ -27,6 +29,7 @@ type Handler struct {
2729
usageReader usage.UsageReader
2830
auditReader auditlog.Reader
2931
registry *providers.ModelRegistry
32+
authKeys *authkeys.Service
3033
aliases *aliases.Service
3134
plans *executionplans.Service
3235
guardrails *guardrails.Registry
@@ -69,6 +72,13 @@ func WithAliases(service *aliases.Service) Option {
6972
}
7073
}
7174

75+
// WithAuthKeys enables managed auth key administration endpoints.
76+
func WithAuthKeys(service *authkeys.Service) Option {
77+
return func(h *Handler) {
78+
h.authKeys = service
79+
}
80+
}
81+
7282
// WithExecutionPlans enables execution-plan administration endpoints.
7383
func WithExecutionPlans(service *executionplans.Service) Option {
7484
return func(h *Handler) {
@@ -613,6 +623,12 @@ type createExecutionPlanRequest struct {
613623
Payload executionplans.Payload `json:"plan_payload"`
614624
}
615625

626+
type createAuthKeyRequest struct {
627+
Name string `json:"name"`
628+
Description string `json:"description,omitempty"`
629+
ExpiresAt *time.Time `json:"expires_at,omitempty"`
630+
}
631+
616632
func featureUnavailableError(message string) error {
617633
return core.NewInvalidRequestErrorWithStatus(http.StatusServiceUnavailable, message, nil).
618634
WithCode("feature_unavailable")
@@ -622,6 +638,10 @@ func (h *Handler) aliasesUnavailableError() error {
622638
return featureUnavailableError("aliases feature is unavailable")
623639
}
624640

641+
func (h *Handler) authKeysUnavailableError() error {
642+
return featureUnavailableError("auth keys feature is unavailable")
643+
}
644+
625645
func (h *Handler) executionPlansUnavailableError() error {
626646
return featureUnavailableError("execution plans feature is unavailable")
627647
}
@@ -646,6 +666,98 @@ func executionPlanWriteError(err error) error {
646666
return err
647667
}
648668

669+
func authKeyWriteError(err error) error {
670+
if err == nil {
671+
return nil
672+
}
673+
if authkeys.IsValidationError(err) {
674+
return core.NewInvalidRequestError(err.Error(), err)
675+
}
676+
return err
677+
}
678+
679+
func deactivateByID(
680+
c *echo.Context,
681+
unavailableErr error,
682+
idLabel string,
683+
notFoundErr error,
684+
notFoundMessage string,
685+
deactivate func(context.Context, string) error,
686+
writeError func(error) error,
687+
) error {
688+
if unavailableErr != nil {
689+
return handleError(c, unavailableErr)
690+
}
691+
692+
id := strings.TrimSpace(c.Param("id"))
693+
if id == "" {
694+
return handleError(c, core.NewInvalidRequestError(idLabel+" id is required", nil))
695+
}
696+
697+
if err := deactivate(c.Request().Context(), id); err != nil {
698+
if errors.Is(err, notFoundErr) {
699+
return handleError(c, core.NewNotFoundError(notFoundMessage+id))
700+
}
701+
return handleError(c, writeError(err))
702+
}
703+
return c.NoContent(http.StatusNoContent)
704+
}
705+
706+
// ListAuthKeys handles GET /admin/api/v1/auth-keys
707+
func (h *Handler) ListAuthKeys(c *echo.Context) error {
708+
if h.authKeys == nil {
709+
return handleError(c, h.authKeysUnavailableError())
710+
}
711+
views := h.authKeys.ListViews()
712+
if views == nil {
713+
views = []authkeys.View{}
714+
}
715+
return c.JSON(http.StatusOK, views)
716+
}
717+
718+
// CreateAuthKey handles POST /admin/api/v1/auth-keys
719+
func (h *Handler) CreateAuthKey(c *echo.Context) error {
720+
if h.authKeys == nil {
721+
return handleError(c, h.authKeysUnavailableError())
722+
}
723+
724+
var req createAuthKeyRequest
725+
if err := c.Bind(&req); err != nil {
726+
return handleError(c, core.NewInvalidRequestError("invalid request body: "+err.Error(), err))
727+
}
728+
729+
issued, err := h.authKeys.Create(c.Request().Context(), authkeys.CreateInput{
730+
Name: req.Name,
731+
Description: req.Description,
732+
ExpiresAt: req.ExpiresAt,
733+
})
734+
if err != nil {
735+
return handleError(c, authKeyWriteError(err))
736+
}
737+
if issued == nil {
738+
requestID := strings.TrimSpace(core.GetRequestID(c.Request().Context()))
739+
slog.Error("auth key service returned nil issued key", "request_id", requestID, "path", c.Request().URL.Path)
740+
return c.JSON(http.StatusInternalServerError, (&core.GatewayError{
741+
Type: core.ErrorType("internal_error"),
742+
Message: "auth key creation failed unexpectedly",
743+
StatusCode: http.StatusInternalServerError,
744+
}).WithCode("auth_key_issue_failed").ToJSON())
745+
}
746+
return c.JSON(http.StatusCreated, issued)
747+
}
748+
749+
// DeactivateAuthKey handles POST /admin/api/v1/auth-keys/:id/deactivate
750+
func (h *Handler) DeactivateAuthKey(c *echo.Context) error {
751+
var unavailableErr error
752+
var deactivate func(context.Context, string) error
753+
if h.authKeys == nil {
754+
unavailableErr = h.authKeysUnavailableError()
755+
} else {
756+
deactivate = h.authKeys.Deactivate
757+
}
758+
return deactivateByID(c, unavailableErr, "auth key", authkeys.ErrNotFound, "auth key not found: ", deactivate, authKeyWriteError)
759+
}
760+
649761
// ListAliases handles GET /admin/api/v1/aliases
650762
func (h *Handler) ListAliases(c *echo.Context) error {
651763
if h.aliases == nil {
@@ -806,22 +918,14 @@ func (h *Handler) CreateExecutionPlan(c *echo.Context) error {
806918

807919
// DeactivateExecutionPlan handles POST /admin/api/v1/execution-plans/:id/deactivate
808920
func (h *Handler) DeactivateExecutionPlan(c *echo.Context) error {
921+
var unavailableErr error
922+
var deactivate func(context.Context, string) error
809923
if h.plans == nil {
810-
return handleError(c, h.executionPlansUnavailableError())
811-
}
812-
813-
id := strings.TrimSpace(c.Param("id"))
814-
if id == "" {
815-
return handleError(c, core.NewInvalidRequestError("execution plan id is required", nil))
816-
}
817-
818-
if err := h.plans.Deactivate(c.Request().Context(), id); err != nil {
819-
if errors.Is(err, executionplans.ErrNotFound) {
820-
return handleError(c, core.NewNotFoundError("workflow not found: "+id))
821-
}
822-
return handleError(c, executionPlanWriteError(err))
924+
unavailableErr = h.executionPlansUnavailableError()
925+
} else {
926+
deactivate = h.plans.Deactivate
823927
}
824-
return c.NoContent(http.StatusNoContent)
928+
return deactivateByID(c, unavailableErr, "execution plan", executionplans.ErrNotFound, "workflow not found: ", deactivate, executionPlanWriteError)
825929
}
826930

827931
func (h *Handler) validateExecutionPlanGuardrails(payload executionplans.Payload) error {

‎internal/admin/handler_authkeys_test.go‎

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
1+
package admin
2+
3+
import (
4+
"bytes"
5+
"context"
6+
"encoding/json"
7+
"net/http"
8+
"net/http/httptest"
9+
"testing"
10+
"time"
11+
12+
"github.com/labstack/echo/v5"
13+
14+
"gomodel/internal/authkeys"
15+
)
16+
17+
type authKeyTestStore struct {
18+
keys map[string]authkeys.AuthKey
19+
}
20+
21+
func newAuthKeyTestStore(keys ...authkeys.AuthKey) *authKeyTestStore {
22+
store := &authKeyTestStore{keys: make(map[string]authkeys.AuthKey, len(keys))}
23+
for _, key := range keys {
24+
store.keys[key.ID] = key
25+
}
26+
return store
27+
}
28+
29+
func (s *authKeyTestStore) List(_ context.Context) ([]authkeys.AuthKey, error) {
30+
result := make([]authkeys.AuthKey, 0, len(s.keys))
31+
for _, key := range s.keys {
32+
result = append(result, key)
33+
}
34+
return result, nil
35+
}
36+
37+
func (s *authKeyTestStore) Create(_ context.Context, key authkeys.AuthKey) error {
38+
s.keys[key.ID] = key
39+
return nil
40+
}
41+
42+
func (s *authKeyTestStore) Deactivate(_ context.Context, id string, now time.Time) error {
43+
key, ok := s.keys[id]
44+
if !ok {
45+
return authkeys.ErrNotFound
46+
}
47+
key.Enabled = false
48+
key.UpdatedAt = now.UTC()
49+
if key.DeactivatedAt == nil {
50+
deactivatedAt := now.UTC()
51+
key.DeactivatedAt = &deactivatedAt
52+
}
53+
s.keys[id] = key
54+
return nil
55+
}
56+
57+
func (s *authKeyTestStore) Close() error { return nil }
58+
59+
func newAuthKeyHandler(t *testing.T, store authkeys.Store) *Handler {
60+
t.Helper()
61+
service, err := authkeys.NewService(store)
62+
if err != nil {
63+
t.Fatalf("NewService() error = %v", err)
64+
}
65+
if err := service.Refresh(context.Background()); err != nil {
66+
t.Fatalf("Refresh() error = %v", err)
67+
}
68+
return NewHandler(nil, nil, WithAuthKeys(service))
69+
}
70+
71+
func TestAuthKeyEndpointsReturn503WhenServiceUnavailable(t *testing.T) {
72+
h := NewHandler(nil, nil)
73+
e := echo.New()
74+
75+
listCtx, listRec := newHandlerContext("/admin/api/v1/auth-keys")
76+
if err := h.ListAuthKeys(listCtx); err != nil {
77+
t.Fatalf("ListAuthKeys() error = %v", err)
78+
}
79+
if listRec.Code != http.StatusServiceUnavailable {
80+
t.Fatalf("ListAuthKeys() status = %d, want 503", listRec.Code)
81+
}
82+
83+
createReq := httptest.NewRequest(http.MethodPost, "/admin/api/v1/auth-keys", bytes.NewBufferString(`{"name":"primary"}`))
84+
createReq.Header.Set("Content-Type", "application/json")
85+
createRec := httptest.NewRecorder()
86+
createCtx := e.NewContext(createReq, createRec)
87+
if err := h.CreateAuthKey(createCtx); err != nil {
88+
t.Fatalf("CreateAuthKey() error = %v", err)
89+
}
90+
if createRec.Code != http.StatusServiceUnavailable {
91+
t.Fatalf("CreateAuthKey() status = %d, want 503", createRec.Code)
92+
}
93+
94+
deactivateReq := httptest.NewRequest(http.MethodPost, "/admin/api/v1/auth-keys/test-key/deactivate", nil)
95+
deactivateRec := httptest.NewRecorder()
96+
deactivateCtx := e.NewContext(deactivateReq, deactivateRec)
97+
deactivateCtx.SetPathValues(echo.PathValues{{Name: "id", Value: "test-key"}})
98+
if err := h.DeactivateAuthKey(deactivateCtx); err != nil {
99+
t.Fatalf("DeactivateAuthKey() error = %v", err)
100+
}
101+
if deactivateRec.Code != http.StatusServiceUnavailable {
102+
t.Fatalf("DeactivateAuthKey() status = %d, want 503", deactivateRec.Code)
103+
}
104+
}
105+
106+
func TestCreateListAndDeactivateAuthKey(t *testing.T) {
107+
h := newAuthKeyHandler(t, newAuthKeyTestStore())
108+
e := echo.New()
109+
110+
createReq := httptest.NewRequest(http.MethodPost, "/admin/api/v1/auth-keys", bytes.NewBufferString(`{"name":"primary","description":"prod key"}`))
111+
createReq.Header.Set("Content-Type", "application/json")
112+
createRec := httptest.NewRecorder()
113+
createCtx := e.NewContext(createReq, createRec)
114+
115+
if err := h.CreateAuthKey(createCtx); err != nil {
116+
t.Fatalf("CreateAuthKey() error = %v", err)
117+
}
118+
if createRec.Code != http.StatusCreated {
119+
t.Fatalf("CreateAuthKey() status = %d, want 201", createRec.Code)
120+
}
121+
122+
var issued authkeys.IssuedKey
123+
if err := json.Unmarshal(createRec.Body.Bytes(), &issued); err != nil {
124+
t.Fatalf("unmarshal create response: %v", err)
125+
}
126+
if issued.Value == "" || issued.ID == "" {
127+
t.Fatalf("issued response = %#v, want id and value", issued)
128+
}
129+
130+
listCtx, listRec := newHandlerContext("/admin/api/v1/auth-keys")
131+
if err := h.ListAuthKeys(listCtx); err != nil {
132+
t.Fatalf("ListAuthKeys() error = %v", err)
133+
}
134+
if listRec.Code != http.StatusOK {
135+
t.Fatalf("ListAuthKeys() status = %d, want 200", listRec.Code)
136+
}
137+
138+
var views []authkeys.View
139+
if err := json.Unmarshal(listRec.Body.Bytes(), &views); err != nil {
140+
t.Fatalf("unmarshal list response: %v", err)
141+
}
142+
if len(views) != 1 || !views[0].Active {
143+
t.Fatalf("list response = %#v, want one active key", views)
144+
}
145+
146+
deactivateReq := httptest.NewRequest(http.MethodPost, "/admin/api/v1/auth-keys/"+issued.ID+"/deactivate", nil)
147+
deactivateRec := httptest.NewRecorder()
148+
deactivateCtx := e.NewContext(deactivateReq, deactivateRec)
149+
deactivateCtx.SetPathValues(echo.PathValues{{Name: "id", Value: issued.ID}})
150+
151+
if err := h.DeactivateAuthKey(deactivateCtx); err != nil {
152+
t.Fatalf("DeactivateAuthKey() error = %v", err)
153+
}
154+
if deactivateRec.Code != http.StatusNoContent {
155+
t.Fatalf("DeactivateAuthKey() status = %d, want 204", deactivateRec.Code)
156+
}
157+
158+
listCtx, listRec = newHandlerContext("/admin/api/v1/auth-keys")
159+
if err := h.ListAuthKeys(listCtx); err != nil {
160+
t.Fatalf("ListAuthKeys() error after deactivate = %v", err)
161+
}
162+
if err := json.Unmarshal(listRec.Body.Bytes(), &views); err != nil {
163+
t.Fatalf("unmarshal list response after deactivate: %v", err)
164+
}
165+
if len(views) != 1 || views[0].Active {
166+
t.Fatalf("list response after deactivate = %#v, want one inactive key", views)
167+
}
168+
}

0 commit comments

Comments
 (0)