Skip to content

Commit 92efa6c

Browse files
SantiagoDePoloniaSantiagoDePolonia
committed
feat(auth): add managed API key support
1 parent 64320e7 commit 92efa6c

29 files changed

Lines changed: 1784 additions & 84 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: 111 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515

1616
"gomodel/internal/aliases"
1717
"gomodel/internal/auditlog"
18+
"gomodel/internal/authkeys"
1819
"gomodel/internal/core"
1920
"gomodel/internal/executionplans"
2021
"gomodel/internal/guardrails"
@@ -27,6 +28,7 @@ type Handler struct {
2728
usageReader usage.UsageReader
2829
auditReader auditlog.Reader
2930
registry *providers.ModelRegistry
31+
authKeys *authkeys.Service
3032
aliases *aliases.Service
3133
plans *executionplans.Service
3234
guardrails *guardrails.Registry
@@ -69,6 +71,13 @@ func WithAliases(service *aliases.Service) Option {
6971
}
7072
}
7173

74+
// WithAuthKeys enables managed auth key administration endpoints.
75+
func WithAuthKeys(service *authkeys.Service) Option {
76+
return func(h *Handler) {
77+
h.authKeys = service
78+
}
79+
}
80+
7281
// WithExecutionPlans enables execution-plan administration endpoints.
7382
func WithExecutionPlans(service *executionplans.Service) Option {
7483
return func(h *Handler) {
@@ -613,6 +622,12 @@ type createExecutionPlanRequest struct {
613622
Payload executionplans.Payload `json:"plan_payload"`
614623
}
615624

625+
type createAuthKeyRequest struct {
626+
Name string `json:"name"`
627+
Description string `json:"description,omitempty"`
628+
ExpiresAt *time.Time `json:"expires_at,omitempty"`
629+
}
630+
616631
func featureUnavailableError(message string) error {
617632
return core.NewInvalidRequestErrorWithStatus(http.StatusServiceUnavailable, message, nil).
618633
WithCode("feature_unavailable")
@@ -622,6 +637,10 @@ func (h *Handler) aliasesUnavailableError() error {
622637
return featureUnavailableError("aliases feature is unavailable")
623638
}
624639

640+
func (h *Handler) authKeysUnavailableError() error {
641+
return featureUnavailableError("auth keys feature is unavailable")
642+
}
643+
625644
func (h *Handler) executionPlansUnavailableError() error {
626645
return featureUnavailableError("execution plans feature is unavailable")
627646
}
@@ -646,6 +665,92 @@ func executionPlanWriteError(err error) error {
646665
return err
647666
}
648667

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

807912
// DeactivateExecutionPlan handles POST /admin/api/v1/execution-plans/:id/deactivate
808913
func (h *Handler) DeactivateExecutionPlan(c *echo.Context) error {
914+
var unavailableErr error
915+
var deactivate func(context.Context, string) error
809916
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))
917+
unavailableErr = h.executionPlansUnavailableError()
918+
} else {
919+
deactivate = h.plans.Deactivate
823920
}
824-
return c.NoContent(http.StatusNoContent)
921+
return deactivateByID(c, unavailableErr, "execution plan", executionplans.ErrNotFound, "workflow not found: ", deactivate, executionPlanWriteError)
825922
}
826923

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

‎internal/admin/handler_authkeys_test.go‎

Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
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+
95+
func TestCreateListAndDeactivateAuthKey(t *testing.T) {
96+
h := newAuthKeyHandler(t, newAuthKeyTestStore())
97+
e := echo.New()
98+
99+
createReq := httptest.NewRequest(http.MethodPost, "/admin/api/v1/auth-keys", bytes.NewBufferString(`{"name":"primary","description":"prod key"}`))
100+
createReq.Header.Set("Content-Type", "application/json")
101+
createRec := httptest.NewRecorder()
102+
createCtx := e.NewContext(createReq, createRec)
103+
104+
if err := h.CreateAuthKey(createCtx); err != nil {
105+
t.Fatalf("CreateAuthKey() error = %v", err)
106+
}
107+
if createRec.Code != http.StatusCreated {
108+
t.Fatalf("CreateAuthKey() status = %d, want 201", createRec.Code)
109+
}
110+
111+
var issued authkeys.IssuedKey
112+
if err := json.Unmarshal(createRec.Body.Bytes(), &issued); err != nil {
113+
t.Fatalf("unmarshal create response: %v", err)
114+
}
115+
if issued.Value == "" || issued.ID == "" {
116+
t.Fatalf("issued response = %#v, want id and value", issued)
117+
}
118+
119+
listCtx, listRec := newHandlerContext("/admin/api/v1/auth-keys")
120+
if err := h.ListAuthKeys(listCtx); err != nil {
121+
t.Fatalf("ListAuthKeys() error = %v", err)
122+
}
123+
if listRec.Code != http.StatusOK {
124+
t.Fatalf("ListAuthKeys() status = %d, want 200", listRec.Code)
125+
}
126+
127+
var views []authkeys.View
128+
if err := json.Unmarshal(listRec.Body.Bytes(), &views); err != nil {
129+
t.Fatalf("unmarshal list response: %v", err)
130+
}
131+
if len(views) != 1 || !views[0].Active {
132+
t.Fatalf("list response = %#v, want one active key", views)
133+
}
134+
135+
deactivateReq := httptest.NewRequest(http.MethodPost, "/admin/api/v1/auth-keys/"+issued.ID+"/deactivate", nil)
136+
deactivateRec := httptest.NewRecorder()
137+
deactivateCtx := e.NewContext(deactivateReq, deactivateRec)
138+
deactivateCtx.SetPathValues(echo.PathValues{{Name: "id", Value: issued.ID}})
139+
140+
if err := h.DeactivateAuthKey(deactivateCtx); err != nil {
141+
t.Fatalf("DeactivateAuthKey() error = %v", err)
142+
}
143+
if deactivateRec.Code != http.StatusNoContent {
144+
t.Fatalf("DeactivateAuthKey() status = %d, want 204", deactivateRec.Code)
145+
}
146+
147+
listCtx, listRec = newHandlerContext("/admin/api/v1/auth-keys")
148+
if err := h.ListAuthKeys(listCtx); err != nil {
149+
t.Fatalf("ListAuthKeys() error after deactivate = %v", err)
150+
}
151+
if err := json.Unmarshal(listRec.Body.Bytes(), &views); err != nil {
152+
t.Fatalf("unmarshal list response after deactivate: %v", err)
153+
}
154+
if len(views) != 1 || views[0].Active {
155+
t.Fatalf("list response after deactivate = %#v, want one inactive key", views)
156+
}
157+
}

0 commit comments

Comments
 (0)