Skip to content

Commit c81fe3c

Browse files
perf(vault): cached credentials and safe singleflight refresh (#46)
* perf(vault): cache tokens and coalesce refreshes Prevent concurrent OAuth refreshes from invalidating rotating tokens. Keep decrypted tokens in a bounded, expiry-aware local LRU. Signed-off-by: Shreyansh Sancheti <43677304+shreyanshjain7174@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix(vault): bound refreshes and guard writes Apply a separate 15s deadline to shared refreshes. Use conditional writes so deletes and re-links beat stale tokens. Retry rotated-token writes without logging credentials. Escape OAuth callback HTML to satisfy the expanded Semgrep scan. Signed-off-by: Shreyansh Sancheti <43677304+shreyanshjain7174@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix(vault): recover shared refresh panics Always release waiters and clear flights after a panic. Use fixed errors and logs so panic values cannot expose credentials. Persist successful rotations even after the exchange deadline. Signed-off-by: Shreyansh Sancheti <43677304+shreyanshjain7174@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Signed-off-by: Shreyansh Sancheti <43677304+shreyanshjain7174@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent 033e77c commit c81fe3c

21 files changed

Lines changed: 1733 additions & 49 deletions

‎README.md‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -305,13 +305,47 @@ AgentGate is deterministic. Agents never see tokens, calls go only to configured
305305
6. **OAuth state encrypted** with AES-256-GCM and expires after 10 minutes.
306306
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.
307307

308+
### Vault cache and refresh
309+
310+
Vault reads cache decrypted access and refresh tokens in process memory.
311+
The default TTL is `60s`, with an LRU limit of 10,000 user-service pairs.
312+
An entry expires earlier when its token reaches 30 seconds before expiry.
313+
Tokens without an expiry use the TTL alone.
314+
Tokens already within that 30-second window aren't cached.
315+
316+
Token strings cannot be reliably zeroed. Keep the cache bounded.
317+
Delete drops cached references. Go may retain string bytes until garbage collection.
318+
Expired entries aren't reused. Lookup or eviction removes their cached references.
319+
Put writes SQLite before caching the new token. Failed writes invalidate the local entry.
320+
All writes must pass through the decorated store.
321+
322+
Concurrent refreshes share one result per store and user-service pair within a process.
323+
Current waiters share refresh failures. Later calls retry.
324+
Canceling a waiter doesn't cancel the shared refresh.
325+
Each shared refresh has a separate `15s` deadline. OAuth HTTP exchanges also have a `15s` client timeout.
326+
Refresh functions receive that context and must honor cancellation.
327+
Successful exchanges still persist their results after the refresh deadline. That deadline doesn't cancel persistence retries.
328+
Shared refresh panics return a fixed error without panic values. The flight is cleared so later requests can retry.
329+
Timeouts and canceled waits return HTTP `504` with receipt error code `vault_refresh_timeout`, not `token_expired`.
330+
Refresh writes compare the stored refresh token inside a SQLite transaction.
331+
Deleted credentials stay deleted. A concurrent re-link keeps its replacement token and scopes.
332+
On a comparison mismatch, refresh returns the current credential only when it remains valid.
333+
Rotated refresh tokens get up to 3 conditional write attempts.
334+
If every write fails, the current request still receives the new access token.
335+
The failure is logged without token values. Later requests may require re-linking because the rotation wasn't saved.
336+
Replicas maintain separate caches and refresh independently.
337+
Rotating refresh tokens still need cross-process coordination when replicas share credentials.
338+
External writes or deletes can leave stale local credentials until cache expiry.
339+
Set `AGENTGATE_VAULT_CACHE_TTL=0` to disable plaintext caching. Refresh sharing remains enabled.
340+
308341
## Configuration
309342

310343
### Environment Variables
311344

312345
| Variable | Description | Required |
313346
|----------|-------------|----------|
314347
| `AGENTGATE_VAULT_KEY` | 32-byte vault key. Also derives receipt encryption keys and the API-key HMAC pepper. | Yes |
348+
| `AGENTGATE_VAULT_CACHE_TTL` | Decrypted-token cache duration. Default `60s`; `0` disables caching. Negative or invalid durations stop startup. | No |
315349
| `AGENTGATE_ADMIN_SECRET` | Secret for admin API access | Yes |
316350
| `AGENTGATE_PUBLIC_URL` | Base URL used to build the OAuth callback (default `http://localhost:8080`) | No |
317351
| `AGENTGATE_API_KEY_CACHE_TTL` | Verified-key cache duration. Default `30s`; `0` disables caching. | No |

‎cmd/agentgw/main.go‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,11 +96,21 @@ func main() {
9696
os.Exit(1)
9797
}
9898

99-
vaultStore, err := vault.NewSQLiteStore(database, masterKey)
99+
sqliteVault, err := vault.NewSQLiteStore(database, masterKey)
100100
if err != nil {
101101
logger.Error("init vault", "error", err)
102102
os.Exit(1)
103103
}
104+
vaultCacheTTL, err := loadVaultCacheTTL()
105+
if err != nil {
106+
logger.Error("load vault cache TTL", "error", err)
107+
os.Exit(1)
108+
}
109+
vaultStore, err := vault.NewCachedStore(sqliteVault, vault.WithCacheTTL(vaultCacheTTL))
110+
if err != nil {
111+
logger.Error("init vault cache", "error", err)
112+
os.Exit(1)
113+
}
104114

105115
pepper, err := auth.DeriveKeyPepper(masterKey)
106116
if err != nil {
@@ -275,6 +285,18 @@ func loadAPIKeyCacheTTL() (time.Duration, error) {
275285
return ttl, nil
276286
}
277287

288+
func loadVaultCacheTTL() (time.Duration, error) {
289+
value, set := os.LookupEnv("AGENTGATE_VAULT_CACHE_TTL")
290+
if !set {
291+
return vault.DefaultCacheTTL, nil
292+
}
293+
ttl, err := time.ParseDuration(value)
294+
if err != nil || ttl < 0 {
295+
return 0, fmt.Errorf("AGENTGATE_VAULT_CACHE_TTL must be a non-negative duration")
296+
}
297+
return ttl, nil
298+
}
299+
278300
// loadMasterKey reads the 32-byte master secret for the vault.
279301
// It also derives the signer's encryption key and the API-key HMAC pepper.
280302
func loadMasterKey() ([]byte, error) {
Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
package main
2+
3+
import (
4+
"os"
5+
"testing"
6+
"time"
7+
)
8+
9+
func TestLoadVaultCacheTTL(t *testing.T) {
10+
for _, tc := range []struct {
11+
value string
12+
want time.Duration
13+
fail bool
14+
}{
15+
{"60s", time.Minute, false},
16+
{"5m", 5 * time.Minute, false},
17+
{"0", 0, false},
18+
{"0s", 0, false},
19+
{"-1s", 0, true},
20+
{"invalid", 0, true},
21+
{"", 0, true},
22+
} {
23+
t.Run(tc.value, func(t *testing.T) {
24+
t.Setenv("AGENTGATE_VAULT_CACHE_TTL", tc.value)
25+
got, err := loadVaultCacheTTL()
26+
if (err != nil) != tc.fail || got != tc.want {
27+
t.Fatalf("TTL = %v, err = %v", got, err)
28+
}
29+
})
30+
}
31+
t.Run("unset", func(t *testing.T) {
32+
t.Setenv("AGENTGATE_VAULT_CACHE_TTL", "")
33+
if err := os.Unsetenv("AGENTGATE_VAULT_CACHE_TTL"); err != nil {
34+
t.Fatal(err)
35+
}
36+
got, err := loadVaultCacheTTL()
37+
if err != nil || got != time.Minute {
38+
t.Fatalf("default TTL = %v, err = %v", got, err)
39+
}
40+
})
41+
}

‎internal/gateway/gateway.go‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -457,8 +457,16 @@ func (s *Server) executeAttempt(ctx context.Context, att *attempt) *outcome {
457457
}
458458
}
459459
if refresher := s.cfg.Refreshers[att.req.Service]; refresher != nil {
460-
tok, err = vault.GetOrRefresh(s.cfg.Vault, att.req.OnBehalfOf, att.req.Service, refresher)
460+
tok, err = vault.GetOrRefresh(ctx, s.cfg.Vault, att.req.OnBehalfOf, att.req.Service, refresher)
461461
if err != nil {
462+
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
463+
return &outcome{
464+
status: http.StatusGatewayTimeout,
465+
body: ErrorResponse{Error: "vault refresh timed out or was canceled", Code: "vault_refresh_timeout"},
466+
policyDecision: "allow",
467+
errorCode: "vault_refresh_timeout",
468+
}
469+
}
462470
return &outcome{
463471
status: http.StatusForbidden,
464472
body: ErrorResponse{Error: "token expired — user must re-authenticate", Code: "token_expired"},

‎internal/gateway/gateway_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -468,7 +468,7 @@ services:
468468
}},
469469
Receipts: &fakeReceipts{},
470470
Refreshers: map[string]vault.RefreshFunc{
471-
"github": func(refreshToken string) (string, string, time.Duration, error) {
471+
"github": func(_ context.Context, refreshToken string) (string, string, time.Duration, error) {
472472
if refreshToken != "refresh-token" {
473473
t.Errorf("refresh token = %q, want stored refresh token", refreshToken)
474474
}
@@ -508,7 +508,7 @@ func TestAct_RefreshFailureBlocksDispatchAndIsReceipted(t *testing.T) {
508508
ExpiresAt: time.Now().Add(-time.Minute),
509509
})
510510
srv.cfg.Refreshers = map[string]vault.RefreshFunc{
511-
"github": func(string) (string, string, time.Duration, error) {
511+
"github": func(context.Context, string) (string, string, time.Duration, error) {
512512
return "", "", 0, errors.New("provider rejected refresh token")
513513
},
514514
}

‎internal/gateway/path_test.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package gateway
22

33
import (
4+
"context"
45
"encoding/json"
56
"fmt"
67
"io/fs"
@@ -333,7 +334,7 @@ func TestActInvalidPathParamsBeforeVault(t *testing.T) {
333334
srv.cfg.Vault = counted
334335
var refreshCalls atomic.Int32
335336
srv.cfg.Refreshers = map[string]vault.RefreshFunc{
336-
"github": func(string) (string, string, time.Duration, error) {
337+
"github": func(context.Context, string) (string, string, time.Duration, error) {
337338
refreshCalls.Add(1)
338339
return "new-token", "new-refresh", time.Hour, nil
339340
},
Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
package gateway
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"net/http"
7+
"net/http/httptest"
8+
"strings"
9+
"testing"
10+
"time"
11+
12+
"github.com/Clawdlinux/agentgate/internal/vault"
13+
)
14+
15+
func TestActRefreshCancellationIsNotTokenExpired(t *testing.T) {
16+
for _, deadline := range []bool{false, true} {
17+
t.Run(map[bool]string{false: "canceled", true: "deadline"}[deadline], func(t *testing.T) {
18+
srv, store, receipts := testSetup(t)
19+
_ = store.Put("user-1", "github", vault.Token{
20+
AccessToken: "old", RefreshToken: "refresh", ExpiresAt: time.Now().Add(-time.Minute),
21+
})
22+
started, release := make(chan struct{}), make(chan struct{})
23+
refresh := func(context.Context, string) (string, string, time.Duration, error) {
24+
close(started)
25+
<-release
26+
return "new", "rotated", time.Hour, nil
27+
}
28+
srv.cfg.Refreshers = map[string]vault.RefreshFunc{"github": refresh}
29+
ctx, cancel := context.WithCancel(context.Background())
30+
if deadline {
31+
cancel()
32+
ctx, cancel = context.WithTimeout(context.Background(), 100*time.Millisecond)
33+
}
34+
defer cancel()
35+
req := httptest.NewRequest("POST", "/v1/act", strings.NewReader(
36+
`{"service":"github","action":"list_repos","on_behalf_of":"user-1"}`)).WithContext(ctx)
37+
req.Header.Set("Authorization", "Bearer test-key")
38+
w := httptest.NewRecorder()
39+
done := make(chan struct{})
40+
go func() {
41+
defer close(done)
42+
srv.ServeHTTP(w, req)
43+
}()
44+
select {
45+
case <-started:
46+
case <-time.After(time.Second):
47+
close(release)
48+
t.Fatal("refresh never started")
49+
}
50+
if !deadline {
51+
cancel()
52+
}
53+
select {
54+
case <-done:
55+
case <-time.After(time.Second):
56+
close(release)
57+
t.Fatal("gateway didn't honor waiter cancellation")
58+
}
59+
close(release)
60+
if _, err := vault.GetOrRefresh(context.Background(), store, "user-1", "github", refresh); err != nil {
61+
t.Fatal(err)
62+
}
63+
var body ErrorResponse
64+
if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
65+
t.Fatal(err)
66+
}
67+
if w.Code != http.StatusGatewayTimeout || body.Code != "vault_refresh_timeout" {
68+
t.Fatalf("status = %d, body = %s", w.Code, w.Body)
69+
}
70+
if body.Error != "vault refresh timed out or was canceled" {
71+
t.Fatalf("unexpected cancellation message: %q", body.Error)
72+
}
73+
if receipts.count() != 1 || receipts.last().Error != "vault_refresh_timeout" {
74+
t.Fatalf("receipt = %+v", receipts.last())
75+
}
76+
})
77+
}
78+
}

‎internal/oauth/handler.go‎

Lines changed: 29 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
package oauth
22

33
import (
4+
"context"
45
"encoding/json"
56
"fmt"
7+
"html/template"
68
"io"
79
"log/slog"
810
"net/http"
@@ -33,6 +35,12 @@ type CallbackHandler struct {
3335
logger *slog.Logger
3436
}
3537

38+
var linkedPage = template.Must(template.New("linked").Parse(`<!DOCTYPE html>
39+
<html><body>
40+
<h2>Account linked!</h2>
41+
<p>Your {{.}} account has been connected. You can close this tab.</p>
42+
</body></html>`))
43+
3644
// NewCallbackHandler creates a new OAuth callback handler.
3745
func NewCallbackHandler(providers map[string]*Provider, v vault.Store, encKey []byte, callbackBase string, logger *slog.Logger) *CallbackHandler {
3846
if logger == nil {
@@ -118,7 +126,7 @@ func (h *CallbackHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
118126
}
119127

120128
// Exchange code for tokens.
121-
tok, err := h.exchangeCode(provider, code, service)
129+
tok, err := h.exchangeCode(r.Context(), provider, code, service)
122130
if err != nil {
123131
h.logger.Error("oauth code exchange failed", "service", service, "error", err)
124132
http.Error(w, "token exchange failed", http.StatusInternalServerError)
@@ -135,31 +143,29 @@ func (h *CallbackHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
135143
h.logger.Info("oauth linked", "user", userID, "service", service)
136144

137145
w.Header().Set("Content-Type", "text/html")
138-
fmt.Fprintf(w, `<!DOCTYPE html>
139-
<html><body>
140-
<h2>Account linked!</h2>
141-
<p>Your %s account has been connected. You can close this tab.</p>
142-
</body></html>`, service)
146+
if err := linkedPage.Execute(w, service); err != nil {
147+
h.logger.Error("render oauth callback", "error", err)
148+
}
143149
}
144150

145-
func (h *CallbackHandler) exchangeCode(provider *Provider, code, service string) (*vault.Token, error) {
151+
func (h *CallbackHandler) exchangeCode(ctx context.Context, provider *Provider, code, service string) (*vault.Token, error) {
146152
data := url.Values{
147153
"grant_type": {"authorization_code"},
148154
"code": {code},
149155
"client_id": {provider.ClientID},
150156
"client_secret": {provider.ClientSecret},
151157
"redirect_uri": {h.callbackBase + "/auth/callback/" + service},
152158
}
153-
tok, _, err := exchangeToken(provider, data)
159+
tok, _, err := exchangeToken(ctx, provider, data)
154160
return tok, err
155161
}
156162

157163
// NewRefreshFunc returns the refresh-token exchange for one configured OAuth
158164
// provider. It returns an error when the provider omits an expiry because the
159165
// caller must not persist a token that will immediately be treated as expired.
160166
func NewRefreshFunc(provider *Provider) vault.RefreshFunc {
161-
return func(refreshToken string) (string, string, time.Duration, error) {
162-
tok, expiresIn, err := exchangeToken(provider, url.Values{
167+
return func(ctx context.Context, refreshToken string) (string, string, time.Duration, error) {
168+
tok, expiresIn, err := exchangeToken(ctx, provider, url.Values{
163169
"grant_type": {"refresh_token"},
164170
"refresh_token": {refreshToken},
165171
"client_id": {provider.ClientID},
@@ -175,14 +181,24 @@ func NewRefreshFunc(provider *Provider) vault.RefreshFunc {
175181
}
176182
}
177183

178-
func exchangeToken(provider *Provider, data url.Values) (*vault.Token, time.Duration, error) {
179-
resp, err := http.PostForm(provider.TokenURL, data)
184+
var tokenHTTPClient = &http.Client{Timeout: vault.DefaultRefreshTimeout}
185+
186+
func exchangeToken(ctx context.Context, provider *Provider, data url.Values) (*vault.Token, time.Duration, error) {
187+
req, err := http.NewRequestWithContext(ctx, http.MethodPost, provider.TokenURL, strings.NewReader(data.Encode()))
188+
if err != nil {
189+
return nil, 0, fmt.Errorf("oauth: create exchange request: %w", err)
190+
}
191+
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
192+
resp, err := tokenHTTPClient.Do(req)
180193
if err != nil {
181194
return nil, 0, fmt.Errorf("oauth: exchange: %w", err)
182195
}
183196
defer resp.Body.Close()
184197

185-
body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
198+
body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
199+
if err != nil {
200+
return nil, 0, fmt.Errorf("oauth: read token response: %w", err)
201+
}
186202

187203
if resp.StatusCode != 200 {
188204
return nil, 0, fmt.Errorf("oauth: exchange returned %d: %s", resp.StatusCode, string(body))

‎internal/oauth/linked_page_test.go‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
package oauth
2+
3+
import (
4+
"bytes"
5+
"strings"
6+
"testing"
7+
)
8+
9+
func TestLinkedPageEscapesService(t *testing.T) {
10+
var body bytes.Buffer
11+
if err := linkedPage.Execute(&body, "<script>alert(1)</script>"); err != nil {
12+
t.Fatal(err)
13+
}
14+
if strings.Contains(body.String(), "<script>") || !strings.Contains(body.String(), "&lt;script&gt;") {
15+
t.Fatal("callback page did not escape the service name")
16+
}
17+
}

0 commit comments

Comments
 (0)