diff --git a/README.md b/README.md index c1e2a57..4aa9339 100644 --- a/README.md +++ b/README.md @@ -305,6 +305,39 @@ AgentGate is deterministic. Agents never see tokens, calls go only to configured 6. **OAuth state encrypted** with AES-256-GCM and expires after 10 minutes. 7. **OAuth refresh.** Configured OAuth providers refresh tokens expiring within 5 minutes before dispatch. A failed refresh returns `token_expired` without calling the SaaS API. +### Vault cache and refresh + +Vault reads cache decrypted access and refresh tokens in process memory. +The default TTL is `60s`, with an LRU limit of 10,000 user-service pairs. +An entry expires earlier when its token reaches 30 seconds before expiry. +Tokens without an expiry use the TTL alone. +Tokens already within that 30-second window aren't cached. + +Token strings cannot be reliably zeroed. Keep the cache bounded. +Delete drops cached references. Go may retain string bytes until garbage collection. +Expired entries aren't reused. Lookup or eviction removes their cached references. +Put writes SQLite before caching the new token. Failed writes invalidate the local entry. +All writes must pass through the decorated store. + +Concurrent refreshes share one result per store and user-service pair within a process. +Current waiters share refresh failures. Later calls retry. +Canceling a waiter doesn't cancel the shared refresh. +Each shared refresh has a separate `15s` deadline. OAuth HTTP exchanges also have a `15s` client timeout. +Refresh functions receive that context and must honor cancellation. +Successful exchanges still persist their results after the refresh deadline. That deadline doesn't cancel persistence retries. +Shared refresh panics return a fixed error without panic values. The flight is cleared so later requests can retry. +Timeouts and canceled waits return HTTP `504` with receipt error code `vault_refresh_timeout`, not `token_expired`. +Refresh writes compare the stored refresh token inside a SQLite transaction. +Deleted credentials stay deleted. A concurrent re-link keeps its replacement token and scopes. +On a comparison mismatch, refresh returns the current credential only when it remains valid. +Rotated refresh tokens get up to 3 conditional write attempts. +If every write fails, the current request still receives the new access token. +The failure is logged without token values. Later requests may require re-linking because the rotation wasn't saved. +Replicas maintain separate caches and refresh independently. +Rotating refresh tokens still need cross-process coordination when replicas share credentials. +External writes or deletes can leave stale local credentials until cache expiry. +Set `AGENTGATE_VAULT_CACHE_TTL=0` to disable plaintext caching. Refresh sharing remains enabled. + ## Configuration ### Environment Variables @@ -312,6 +345,7 @@ AgentGate is deterministic. Agents never see tokens, calls go only to configured | Variable | Description | Required | |----------|-------------|----------| | `AGENTGATE_VAULT_KEY` | 32-byte vault key. Also derives receipt encryption keys and the API-key HMAC pepper. | Yes | +| `AGENTGATE_VAULT_CACHE_TTL` | Decrypted-token cache duration. Default `60s`; `0` disables caching. Negative or invalid durations stop startup. | No | | `AGENTGATE_ADMIN_SECRET` | Secret for admin API access | Yes | | `AGENTGATE_PUBLIC_URL` | Base URL used to build the OAuth callback (default `http://localhost:8080`) | No | | `AGENTGATE_API_KEY_CACHE_TTL` | Verified-key cache duration. Default `30s`; `0` disables caching. | No | diff --git a/cmd/agentgw/main.go b/cmd/agentgw/main.go index 7d31a0b..9c2daed 100644 --- a/cmd/agentgw/main.go +++ b/cmd/agentgw/main.go @@ -96,11 +96,21 @@ func main() { os.Exit(1) } - vaultStore, err := vault.NewSQLiteStore(database, masterKey) + sqliteVault, err := vault.NewSQLiteStore(database, masterKey) if err != nil { logger.Error("init vault", "error", err) os.Exit(1) } + vaultCacheTTL, err := loadVaultCacheTTL() + if err != nil { + logger.Error("load vault cache TTL", "error", err) + os.Exit(1) + } + vaultStore, err := vault.NewCachedStore(sqliteVault, vault.WithCacheTTL(vaultCacheTTL)) + if err != nil { + logger.Error("init vault cache", "error", err) + os.Exit(1) + } pepper, err := auth.DeriveKeyPepper(masterKey) if err != nil { @@ -275,6 +285,18 @@ func loadAPIKeyCacheTTL() (time.Duration, error) { return ttl, nil } +func loadVaultCacheTTL() (time.Duration, error) { + value, set := os.LookupEnv("AGENTGATE_VAULT_CACHE_TTL") + if !set { + return vault.DefaultCacheTTL, nil + } + ttl, err := time.ParseDuration(value) + if err != nil || ttl < 0 { + return 0, fmt.Errorf("AGENTGATE_VAULT_CACHE_TTL must be a non-negative duration") + } + return ttl, nil +} + // loadMasterKey reads the 32-byte master secret for the vault. // It also derives the signer's encryption key and the API-key HMAC pepper. func loadMasterKey() ([]byte, error) { diff --git a/cmd/agentgw/vault_cache_ttl_test.go b/cmd/agentgw/vault_cache_ttl_test.go new file mode 100644 index 0000000..bd5967c --- /dev/null +++ b/cmd/agentgw/vault_cache_ttl_test.go @@ -0,0 +1,41 @@ +package main + +import ( + "os" + "testing" + "time" +) + +func TestLoadVaultCacheTTL(t *testing.T) { + for _, tc := range []struct { + value string + want time.Duration + fail bool + }{ + {"60s", time.Minute, false}, + {"5m", 5 * time.Minute, false}, + {"0", 0, false}, + {"0s", 0, false}, + {"-1s", 0, true}, + {"invalid", 0, true}, + {"", 0, true}, + } { + t.Run(tc.value, func(t *testing.T) { + t.Setenv("AGENTGATE_VAULT_CACHE_TTL", tc.value) + got, err := loadVaultCacheTTL() + if (err != nil) != tc.fail || got != tc.want { + t.Fatalf("TTL = %v, err = %v", got, err) + } + }) + } + t.Run("unset", func(t *testing.T) { + t.Setenv("AGENTGATE_VAULT_CACHE_TTL", "") + if err := os.Unsetenv("AGENTGATE_VAULT_CACHE_TTL"); err != nil { + t.Fatal(err) + } + got, err := loadVaultCacheTTL() + if err != nil || got != time.Minute { + t.Fatalf("default TTL = %v, err = %v", got, err) + } + }) +} diff --git a/internal/gateway/gateway.go b/internal/gateway/gateway.go index 4b50bf8..a6f4954 100644 --- a/internal/gateway/gateway.go +++ b/internal/gateway/gateway.go @@ -457,8 +457,16 @@ func (s *Server) executeAttempt(ctx context.Context, att *attempt) *outcome { } } if refresher := s.cfg.Refreshers[att.req.Service]; refresher != nil { - tok, err = vault.GetOrRefresh(s.cfg.Vault, att.req.OnBehalfOf, att.req.Service, refresher) + tok, err = vault.GetOrRefresh(ctx, s.cfg.Vault, att.req.OnBehalfOf, att.req.Service, refresher) if err != nil { + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return &outcome{ + status: http.StatusGatewayTimeout, + body: ErrorResponse{Error: "vault refresh timed out or was canceled", Code: "vault_refresh_timeout"}, + policyDecision: "allow", + errorCode: "vault_refresh_timeout", + } + } return &outcome{ status: http.StatusForbidden, body: ErrorResponse{Error: "token expired — user must re-authenticate", Code: "token_expired"}, diff --git a/internal/gateway/gateway_test.go b/internal/gateway/gateway_test.go index b12002b..a4277fb 100644 --- a/internal/gateway/gateway_test.go +++ b/internal/gateway/gateway_test.go @@ -468,7 +468,7 @@ services: }}, Receipts: &fakeReceipts{}, Refreshers: map[string]vault.RefreshFunc{ - "github": func(refreshToken string) (string, string, time.Duration, error) { + "github": func(_ context.Context, refreshToken string) (string, string, time.Duration, error) { if refreshToken != "refresh-token" { t.Errorf("refresh token = %q, want stored refresh token", refreshToken) } @@ -508,7 +508,7 @@ func TestAct_RefreshFailureBlocksDispatchAndIsReceipted(t *testing.T) { ExpiresAt: time.Now().Add(-time.Minute), }) srv.cfg.Refreshers = map[string]vault.RefreshFunc{ - "github": func(string) (string, string, time.Duration, error) { + "github": func(context.Context, string) (string, string, time.Duration, error) { return "", "", 0, errors.New("provider rejected refresh token") }, } diff --git a/internal/gateway/path_test.go b/internal/gateway/path_test.go index 0c88f6a..d099203 100644 --- a/internal/gateway/path_test.go +++ b/internal/gateway/path_test.go @@ -1,6 +1,7 @@ package gateway import ( + "context" "encoding/json" "fmt" "io/fs" @@ -333,7 +334,7 @@ func TestActInvalidPathParamsBeforeVault(t *testing.T) { srv.cfg.Vault = counted var refreshCalls atomic.Int32 srv.cfg.Refreshers = map[string]vault.RefreshFunc{ - "github": func(string) (string, string, time.Duration, error) { + "github": func(context.Context, string) (string, string, time.Duration, error) { refreshCalls.Add(1) return "new-token", "new-refresh", time.Hour, nil }, diff --git a/internal/gateway/refresh_cancellation_test.go b/internal/gateway/refresh_cancellation_test.go new file mode 100644 index 0000000..6346a29 --- /dev/null +++ b/internal/gateway/refresh_cancellation_test.go @@ -0,0 +1,78 @@ +package gateway + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/Clawdlinux/agentgate/internal/vault" +) + +func TestActRefreshCancellationIsNotTokenExpired(t *testing.T) { + for _, deadline := range []bool{false, true} { + t.Run(map[bool]string{false: "canceled", true: "deadline"}[deadline], func(t *testing.T) { + srv, store, receipts := testSetup(t) + _ = store.Put("user-1", "github", vault.Token{ + AccessToken: "old", RefreshToken: "refresh", ExpiresAt: time.Now().Add(-time.Minute), + }) + started, release := make(chan struct{}), make(chan struct{}) + refresh := func(context.Context, string) (string, string, time.Duration, error) { + close(started) + <-release + return "new", "rotated", time.Hour, nil + } + srv.cfg.Refreshers = map[string]vault.RefreshFunc{"github": refresh} + ctx, cancel := context.WithCancel(context.Background()) + if deadline { + cancel() + ctx, cancel = context.WithTimeout(context.Background(), 100*time.Millisecond) + } + defer cancel() + req := httptest.NewRequest("POST", "/v1/act", strings.NewReader( + `{"service":"github","action":"list_repos","on_behalf_of":"user-1"}`)).WithContext(ctx) + req.Header.Set("Authorization", "Bearer test-key") + w := httptest.NewRecorder() + done := make(chan struct{}) + go func() { + defer close(done) + srv.ServeHTTP(w, req) + }() + select { + case <-started: + case <-time.After(time.Second): + close(release) + t.Fatal("refresh never started") + } + if !deadline { + cancel() + } + select { + case <-done: + case <-time.After(time.Second): + close(release) + t.Fatal("gateway didn't honor waiter cancellation") + } + close(release) + if _, err := vault.GetOrRefresh(context.Background(), store, "user-1", "github", refresh); err != nil { + t.Fatal(err) + } + var body ErrorResponse + if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil { + t.Fatal(err) + } + if w.Code != http.StatusGatewayTimeout || body.Code != "vault_refresh_timeout" { + t.Fatalf("status = %d, body = %s", w.Code, w.Body) + } + if body.Error != "vault refresh timed out or was canceled" { + t.Fatalf("unexpected cancellation message: %q", body.Error) + } + if receipts.count() != 1 || receipts.last().Error != "vault_refresh_timeout" { + t.Fatalf("receipt = %+v", receipts.last()) + } + }) + } +} diff --git a/internal/oauth/handler.go b/internal/oauth/handler.go index d8c0ad2..1c84aef 100644 --- a/internal/oauth/handler.go +++ b/internal/oauth/handler.go @@ -1,8 +1,10 @@ package oauth import ( + "context" "encoding/json" "fmt" + "html/template" "io" "log/slog" "net/http" @@ -33,6 +35,12 @@ type CallbackHandler struct { logger *slog.Logger } +var linkedPage = template.Must(template.New("linked").Parse(` +
+Your {{.}} account has been connected. You can close this tab.
+`)) + // NewCallbackHandler creates a new OAuth callback handler. func NewCallbackHandler(providers map[string]*Provider, v vault.Store, encKey []byte, callbackBase string, logger *slog.Logger) *CallbackHandler { if logger == nil { @@ -118,7 +126,7 @@ func (h *CallbackHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { } // Exchange code for tokens. - tok, err := h.exchangeCode(provider, code, service) + tok, err := h.exchangeCode(r.Context(), provider, code, service) if err != nil { h.logger.Error("oauth code exchange failed", "service", service, "error", err) http.Error(w, "token exchange failed", http.StatusInternalServerError) @@ -135,14 +143,12 @@ func (h *CallbackHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.logger.Info("oauth linked", "user", userID, "service", service) w.Header().Set("Content-Type", "text/html") - fmt.Fprintf(w, ` - -Your %s account has been connected. You can close this tab.
-`, service) + if err := linkedPage.Execute(w, service); err != nil { + h.logger.Error("render oauth callback", "error", err) + } } -func (h *CallbackHandler) exchangeCode(provider *Provider, code, service string) (*vault.Token, error) { +func (h *CallbackHandler) exchangeCode(ctx context.Context, provider *Provider, code, service string) (*vault.Token, error) { data := url.Values{ "grant_type": {"authorization_code"}, "code": {code}, @@ -150,7 +156,7 @@ func (h *CallbackHandler) exchangeCode(provider *Provider, code, service string) "client_secret": {provider.ClientSecret}, "redirect_uri": {h.callbackBase + "/auth/callback/" + service}, } - tok, _, err := exchangeToken(provider, data) + tok, _, err := exchangeToken(ctx, provider, data) return tok, err } @@ -158,8 +164,8 @@ func (h *CallbackHandler) exchangeCode(provider *Provider, code, service string) // provider. It returns an error when the provider omits an expiry because the // caller must not persist a token that will immediately be treated as expired. func NewRefreshFunc(provider *Provider) vault.RefreshFunc { - return func(refreshToken string) (string, string, time.Duration, error) { - tok, expiresIn, err := exchangeToken(provider, url.Values{ + return func(ctx context.Context, refreshToken string) (string, string, time.Duration, error) { + tok, expiresIn, err := exchangeToken(ctx, provider, url.Values{ "grant_type": {"refresh_token"}, "refresh_token": {refreshToken}, "client_id": {provider.ClientID}, @@ -175,14 +181,24 @@ func NewRefreshFunc(provider *Provider) vault.RefreshFunc { } } -func exchangeToken(provider *Provider, data url.Values) (*vault.Token, time.Duration, error) { - resp, err := http.PostForm(provider.TokenURL, data) +var tokenHTTPClient = &http.Client{Timeout: vault.DefaultRefreshTimeout} + +func exchangeToken(ctx context.Context, provider *Provider, data url.Values) (*vault.Token, time.Duration, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.TokenURL, strings.NewReader(data.Encode())) + if err != nil { + return nil, 0, fmt.Errorf("oauth: create exchange request: %w", err) + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + resp, err := tokenHTTPClient.Do(req) if err != nil { return nil, 0, fmt.Errorf("oauth: exchange: %w", err) } defer resp.Body.Close() - body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if err != nil { + return nil, 0, fmt.Errorf("oauth: read token response: %w", err) + } if resp.StatusCode != 200 { return nil, 0, fmt.Errorf("oauth: exchange returned %d: %s", resp.StatusCode, string(body)) diff --git a/internal/oauth/linked_page_test.go b/internal/oauth/linked_page_test.go new file mode 100644 index 0000000..4579b93 --- /dev/null +++ b/internal/oauth/linked_page_test.go @@ -0,0 +1,17 @@ +package oauth + +import ( + "bytes" + "strings" + "testing" +) + +func TestLinkedPageEscapesService(t *testing.T) { + var body bytes.Buffer + if err := linkedPage.Execute(&body, ""); err != nil { + t.Fatal(err) + } + if strings.Contains(body.String(), "