You've already forked agentic-coding-workflow
123 lines
3.6 KiB
Go
123 lines
3.6 KiB
Go
package security
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/go-chi/chi/v5"
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/yan1h/agent-coding-workflow/internal/jobs"
|
|
)
|
|
|
|
type fakeRotator struct {
|
|
active int
|
|
rotated bool
|
|
}
|
|
|
|
func (f *fakeRotator) RotateActive(_ context.Context, _, _ string) (int, error) {
|
|
f.active++
|
|
f.rotated = true
|
|
return f.active, nil
|
|
}
|
|
func (f *fakeRotator) ActiveVersion(context.Context) (int, error) { return f.active, nil }
|
|
|
|
type fakeEnqueuer struct{ enqueued jobs.JobType }
|
|
|
|
func (f *fakeEnqueuer) Enqueue(_ context.Context, typ jobs.JobType, _ json.RawMessage, _ time.Time, _ int32) (*jobs.Job, error) {
|
|
f.enqueued = typ
|
|
return &jobs.Job{ID: uuid.New(), Type: typ}, nil
|
|
}
|
|
|
|
type fakeStatus struct{ stale map[string]int }
|
|
|
|
func (f *fakeStatus) Tables() []string { return []string{"workspaces"} }
|
|
func (f *fakeStatus) StaleCount(_ context.Context, t string, _ int) (int, error) {
|
|
return f.stale[t], nil
|
|
}
|
|
|
|
// fakeResolver maps a fixed token to a fixed user id.
|
|
type fakeResolver struct{ uid uuid.UUID }
|
|
|
|
func (f fakeResolver) ResolveSession(_ context.Context, token string) (uuid.UUID, bool, error) {
|
|
if token == "good" {
|
|
return f.uid, true, nil
|
|
}
|
|
return uuid.Nil, false, nil
|
|
}
|
|
|
|
type fakeAdmin struct{ admins map[uuid.UUID]bool }
|
|
|
|
func (f fakeAdmin) IsAdmin(_ context.Context, id uuid.UUID) (bool, error) {
|
|
return f.admins[id], nil
|
|
}
|
|
|
|
func mountHandler(t *testing.T, admin bool) (*chi.Mux, *fakeRotator, *fakeEnqueuer) {
|
|
t.Helper()
|
|
uid := uuid.New()
|
|
rot := &fakeRotator{active: 1}
|
|
enq := &fakeEnqueuer{}
|
|
st := &fakeStatus{stale: map[string]int{"workspaces": 2}}
|
|
h := NewHandler(rot, enq, st, fakeResolver{uid: uid}, fakeAdmin{admins: map[uuid.UUID]bool{uid: admin}}, nil)
|
|
r := chi.NewRouter()
|
|
h.Mount(r)
|
|
return r, rot, enq
|
|
}
|
|
|
|
func TestRotateKey_AdminEnqueuesReencrypt(t *testing.T) {
|
|
r, rot, enq := mountHandler(t, true)
|
|
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/security/rotate-key", nil)
|
|
req.Header.Set("Authorization", "Bearer good")
|
|
rr := httptest.NewRecorder()
|
|
r.ServeHTTP(rr, req)
|
|
|
|
require.Equal(t, http.StatusAccepted, rr.Code)
|
|
require.True(t, rot.rotated)
|
|
require.Equal(t, jobs.TypeSecretReencrypt, enq.enqueued)
|
|
|
|
var resp rotateKeyResp
|
|
require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &resp))
|
|
require.Equal(t, 2, resp.ActiveVersion)
|
|
require.NotEmpty(t, resp.JobID)
|
|
}
|
|
|
|
func TestRotateKey_NonAdminForbidden(t *testing.T) {
|
|
r, rot, enq := mountHandler(t, false)
|
|
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/security/rotate-key", nil)
|
|
req.Header.Set("Authorization", "Bearer good")
|
|
rr := httptest.NewRecorder()
|
|
r.ServeHTTP(rr, req)
|
|
|
|
require.Equal(t, http.StatusForbidden, rr.Code)
|
|
require.False(t, rot.rotated)
|
|
require.Empty(t, enq.enqueued)
|
|
}
|
|
|
|
func TestRotateKey_Unauthenticated(t *testing.T) {
|
|
r, _, _ := mountHandler(t, true)
|
|
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/security/rotate-key", nil)
|
|
req.Header.Set("Authorization", "Bearer bad")
|
|
rr := httptest.NewRecorder()
|
|
r.ServeHTTP(rr, req)
|
|
require.Equal(t, http.StatusUnauthorized, rr.Code)
|
|
}
|
|
|
|
func TestReencryptStatus_ReportsStale(t *testing.T) {
|
|
r, _, _ := mountHandler(t, true)
|
|
req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/security/reencrypt-status", nil)
|
|
req.Header.Set("Authorization", "Bearer good")
|
|
rr := httptest.NewRecorder()
|
|
r.ServeHTTP(rr, req)
|
|
|
|
require.Equal(t, http.StatusOK, rr.Code)
|
|
var resp reencryptStatusResp
|
|
require.NoError(t, json.Unmarshal(rr.Body.Bytes(), &resp))
|
|
require.Equal(t, 2, resp.Stale["workspaces"])
|
|
require.False(t, resp.Complete)
|
|
}
|