wdevs/amcs
5
0
Fork 0
Code Issues 10 Pull Requests Actions Packages Projects Releases 1 Wiki Activity

feat(app): add lightweight status access tracking

Browse Source
This commit is contained in:
Jack O'Neill
2026-04-04 14:16:02 +02:00
parent 7c41a3e846
commit 50870dd369
8 changed files with 405 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

81
internal/auth/access_tracker.go Normal file
View File

@@ -0,0 +1,81 @@
package auth
import (
"sort"
"sync"
"time"
)
type AccessSnapshot struct {
KeyID string
LastPath string
RemoteAddr string
UserAgent string
RequestCount int
LastAccessedAt time.Time
}
type AccessTracker struct {
mu sync.RWMutex
entries map[string]AccessSnapshot
}
func NewAccessTracker() *AccessTracker {
return &AccessTracker{entries: make(map[string]AccessSnapshot)}
}
func (t *AccessTracker) Record(keyID, path, remoteAddr, userAgent string, now time.Time) {
if t == nil || keyID == "" {
return
}
t.mu.Lock()
defer t.mu.Unlock()
entry := t.entries[keyID]
entry.KeyID = keyID
entry.LastPath = path
entry.RemoteAddr = remoteAddr
entry.UserAgent = userAgent
entry.LastAccessedAt = now.UTC()
entry.RequestCount++
t.entries[keyID] = entry
}
func (t *AccessTracker) Snapshot() []AccessSnapshot {
if t == nil {
return nil
}
t.mu.RLock()
defer t.mu.RUnlock()
items := make([]AccessSnapshot, 0, len(t.entries))
for _, entry := range t.entries {
items = append(items, entry)
}
sort.Slice(items, func(i, j int) bool {
return items[i].LastAccessedAt.After(items[j].LastAccessedAt)
})
return items
}
func (t *AccessTracker) ConnectedCount(now time.Time, window time.Duration) int {
if t == nil {
return 0
}
cutoff := now.UTC().Add(-window)
t.mu.RLock()
defer t.mu.RUnlock()
count := 0
for _, entry := range t.entries {
if !entry.LastAccessedAt.Before(cutoff) {
count++
}
}
return count
}

45
internal/auth/access_tracker_test.go Normal file
View File

@@ -0,0 +1,45 @@
package auth
import (
"testing"
"time"
)
func TestAccessTrackerRecordAndSnapshot(t *testing.T) {
tracker := NewAccessTracker()
older := time.Date(2026, 4, 4, 10, 0, 0, 0, time.UTC)
newer := older.Add(2 * time.Minute)
tracker.Record("client-a", "/files", "10.0.0.1:1234", "agent-a", older)
tracker.Record("client-b", "/mcp", "10.0.0.2:1234", "agent-b", newer)
tracker.Record("client-a", "/files/1", "10.0.0.1:1234", "agent-a2", newer.Add(30*time.Second))
snap := tracker.Snapshot()
if len(snap) != 2 {
t.Fatalf("len(snapshot) = %d, want 2", len(snap))
}
if snap[0].KeyID != "client-a" {
t.Fatalf("snapshot[0].KeyID = %q, want client-a", snap[0].KeyID)
}
if snap[0].RequestCount != 2 {
t.Fatalf("snapshot[0].RequestCount = %d, want 2", snap[0].RequestCount)
}
if snap[0].LastPath != "/files/1" {
t.Fatalf("snapshot[0].LastPath = %q, want /files/1", snap[0].LastPath)
}
if snap[0].UserAgent != "agent-a2" {
t.Fatalf("snapshot[0].UserAgent = %q, want agent-a2", snap[0].UserAgent)
}
}
func TestAccessTrackerConnectedCount(t *testing.T) {
tracker := NewAccessTracker()
now := time.Date(2026, 4, 4, 12, 0, 0, 0, time.UTC)
tracker.Record("recent", "/mcp", "", "", now.Add(-2*time.Minute))
tracker.Record("stale", "/mcp", "", "", now.Add(-11*time.Minute))
if got := tracker.ConnectedCount(now, 10*time.Minute); got != 1 {
t.Fatalf("ConnectedCount() = %d, want 1", got)
}
}

10
internal/auth/keyring_test.go
View File

@@ -39,7 +39,7 @@ func TestMiddlewareAllowsHeaderAuthAndSetsContext(t *testing.T) {
t.Fatalf("NewKeyring() error = %v", err)
}
handler := Middleware(config.AuthConfig{HeaderName: "x-brain-key"}, keyring, nil, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
handler := Middleware(config.AuthConfig{HeaderName: "x-brain-key"}, keyring, nil, nil, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
keyID, ok := KeyIDFromContext(r.Context())
if !ok || keyID != "client-a" {
t.Fatalf("KeyIDFromContext() = (%q, %v), want (client-a, true)", keyID, ok)
@@ -63,7 +63,7 @@ func TestMiddlewareAllowsBearerAuthAndSetsContext(t *testing.T) {
t.Fatalf("NewKeyring() error = %v", err)
}
handler := Middleware(config.AuthConfig{HeaderName: "x-brain-key"}, keyring, nil, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
handler := Middleware(config.AuthConfig{HeaderName: "x-brain-key"}, keyring, nil, nil, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
keyID, ok := KeyIDFromContext(r.Context())
if !ok || keyID != "client-a" {
t.Fatalf("KeyIDFromContext() = (%q, %v), want (client-a, true)", keyID, ok)
@@ -90,7 +90,7 @@ func TestMiddlewarePrefersExplicitHeaderOverBearerAuth(t *testing.T) {
t.Fatalf("NewKeyring() error = %v", err)
}
handler := Middleware(config.AuthConfig{HeaderName: "x-brain-key"}, keyring, nil, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
handler := Middleware(config.AuthConfig{HeaderName: "x-brain-key"}, keyring, nil, nil, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
keyID, ok := KeyIDFromContext(r.Context())
if !ok || keyID != "client-a" {
t.Fatalf("KeyIDFromContext() = (%q, %v), want (client-a, true)", keyID, ok)
@@ -119,7 +119,7 @@ func TestMiddlewareAllowsQueryParamWhenEnabled(t *testing.T) {
HeaderName: "x-brain-key",
QueryParam: "key",
AllowQueryParam: true,
}, keyring, nil, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
}, keyring, nil, nil, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusNoContent)
}))
@@ -138,7 +138,7 @@ func TestMiddlewareRejectsMissingOrInvalidKey(t *testing.T) {
t.Fatalf("NewKeyring() error = %v", err)
}
handler := Middleware(config.AuthConfig{HeaderName: "x-brain-key"}, keyring, nil, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
handler := Middleware(config.AuthConfig{HeaderName: "x-brain-key"}, keyring, nil, nil, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
t.Fatal("next handler should not be called")
}))

13
internal/auth/middleware.go
View File

@@ -6,6 +6,7 @@ import (
"log/slog"
"net/http"
"strings"
"time"
"git.warky.dev/wdevs/amcs/internal/config"
)
@@ -14,11 +15,16 @@ type contextKey string
const keyIDContextKey contextKey = "auth.key_id"
func Middleware(cfg config.AuthConfig, keyring *Keyring, oauthRegistry *OAuthRegistry, tokenStore *TokenStore, log *slog.Logger) func(http.Handler) http.Handler {
func Middleware(cfg config.AuthConfig, keyring *Keyring, oauthRegistry *OAuthRegistry, tokenStore *TokenStore, tracker *AccessTracker, log *slog.Logger) func(http.Handler) http.Handler {
headerName := cfg.HeaderName
if headerName == "" {
headerName = "x-brain-key"
}
recordAccess := func(r *http.Request, keyID string) {
if tracker != nil {
tracker.Record(keyID, r.URL.Path, r.RemoteAddr, r.UserAgent(), time.Now())
}
}
return func(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// 1. Custom header → keyring only.
@@ -30,6 +36,7 @@ func Middleware(cfg config.AuthConfig, keyring *Keyring, oauthRegistry *OAuthReg
http.Error(w, "invalid API key", http.StatusUnauthorized)
return
}
recordAccess(r, keyID)
next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID)))
return
}
@@ -39,12 +46,14 @@ func Middleware(cfg config.AuthConfig, keyring *Keyring, oauthRegistry *OAuthReg
if bearer := extractBearer(r); bearer != "" {
if tokenStore != nil {
if keyID, ok := tokenStore.Lookup(bearer); ok {
recordAccess(r, keyID)
next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, 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)))
return
}
@@ -66,6 +75,7 @@ func Middleware(cfg config.AuthConfig, keyring *Keyring, oauthRegistry *OAuthReg
http.Error(w, "invalid OAuth client credentials", http.StatusUnauthorized)
return
}
recordAccess(r, keyID)
next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID)))
return
}
@@ -79,6 +89,7 @@ func Middleware(cfg config.AuthConfig, keyring *Keyring, oauthRegistry *OAuthReg
http.Error(w, "invalid API key", http.StatusUnauthorized)
return
}
recordAccess(r, keyID)
next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID)))
return
}

4
internal/auth/oauth_registry_test.go
View File

@@ -42,7 +42,7 @@ func TestMiddlewareAllowsOAuthBasicAuthAndSetsContext(t *testing.T) {
t.Fatalf("NewOAuthRegistry() error = %v", err)
}
handler := Middleware(config.AuthConfig{}, nil, oauthRegistry, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
handler := Middleware(config.AuthConfig{}, nil, oauthRegistry, nil, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
keyID, ok := KeyIDFromContext(r.Context())
if !ok || keyID != "oauth-client" {
t.Fatalf("KeyIDFromContext() = (%q, %v), want (oauth-client, true)", keyID, ok)
@@ -70,7 +70,7 @@ func TestMiddlewareRejectsOAuthMissingOrInvalidCredentials(t *testing.T) {
t.Fatalf("NewOAuthRegistry() error = %v", err)
}
handler := Middleware(config.AuthConfig{}, nil, oauthRegistry, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
handler := Middleware(config.AuthConfig{}, nil, oauthRegistry, nil, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
t.Fatal("next handler should not be called")
}))