From 35893ae75a2bd9a0862b21bb4d60287888e47c0f Mon Sep 17 00:00:00 2001
From: Shreyansh Sancheti <43677304+shreyanshjain7174@users.noreply.github.com>
Date: Sun, 18 Oct 2026 05:08:07 +0530
Subject: [PATCH 1/3] 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>
---
README.md | 23 ++
cmd/agentgw/main.go | 24 +-
cmd/agentgw/vault_cache_ttl_test.go | 41 ++++
internal/gateway/gateway.go | 2 +-
internal/vault/benchmark_test.go | 41 ++++
internal/vault/cached_store.go | 174 +++++++++++++++
internal/vault/cached_store_test.go | 334 ++++++++++++++++++++++++++++
internal/vault/refresh.go | 69 +++++-
internal/vault/refresh_test.go | 256 ++++++++++++++++++++-
internal/vault/sqlite_store_test.go | 2 +-
10 files changed, 952 insertions(+), 14 deletions(-)
create mode 100644 cmd/agentgw/vault_cache_ttl_test.go
create mode 100644 internal/vault/benchmark_test.go
create mode 100644 internal/vault/cached_store.go
create mode 100644 internal/vault/cached_store_test.go
diff --git a/README.md b/README.md
index c1e2a57..08046e7 100644
--- a/README.md
+++ b/README.md
@@ -305,6 +305,28 @@ 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.
+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 +334,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..28930f0 100644
--- a/internal/gateway/gateway.go
+++ b/internal/gateway/gateway.go
@@ -457,7 +457,7 @@ 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 {
return &outcome{
status: http.StatusForbidden,
diff --git a/internal/vault/benchmark_test.go b/internal/vault/benchmark_test.go
new file mode 100644
index 0000000..626fcb7
--- /dev/null
+++ b/internal/vault/benchmark_test.go
@@ -0,0 +1,41 @@
+package vault
+
+import (
+ "testing"
+ "time"
+)
+
+func BenchmarkVaultGet(b *testing.B) {
+ store, err := NewSQLiteStore(setupTestDB(b), testKey())
+ if err != nil {
+ b.Fatal(err)
+ }
+ if err := store.Put("user", "service", Token{
+ AccessToken: "access-token",
+ RefreshToken: "refresh-token",
+ ExpiresAt: time.Now().Add(time.Hour),
+ Scopes: []string{"read", "write"},
+ }); err != nil {
+ b.Fatal(err)
+ }
+ b.Run("SQLite", func(b *testing.B) {
+ b.ReportAllocs()
+ for b.Loop() {
+ if _, err := store.Get("user", "service"); err != nil {
+ b.Fatal(err)
+ }
+ }
+ })
+ cached, err := NewCachedStore(store)
+ if err != nil {
+ b.Fatal(err)
+ }
+ b.Run("Cached", func(b *testing.B) {
+ b.ReportAllocs()
+ for b.Loop() {
+ if _, err := cached.Get("user", "service"); err != nil {
+ b.Fatal(err)
+ }
+ }
+ })
+}
diff --git a/internal/vault/cached_store.go b/internal/vault/cached_store.go
new file mode 100644
index 0000000..0531ebd
--- /dev/null
+++ b/internal/vault/cached_store.go
@@ -0,0 +1,174 @@
+package vault
+
+import (
+ "container/list"
+ "fmt"
+ "slices"
+ "sync"
+ "time"
+)
+
+const (
+ DefaultCacheTTL = time.Minute
+ DefaultCacheSize = 10_000
+ cacheExpirySkew = 30 * time.Second
+)
+
+type credentialKey struct {
+ user, service string
+}
+
+type cacheEntry struct {
+ key credentialKey
+ token Token
+ expiresAt time.Time
+}
+
+// CachedStore keeps a bounded LRU of decrypted tokens in process memory.
+// Strings cannot be reliably zeroed. Delete drops cached references.
+// All writes must use this decorator to keep its cache coherent.
+type CachedStore struct {
+ inner Store
+ ttl time.Duration
+ size int
+ now func() time.Time
+
+ // The gate prevents an in-flight read from caching a token after a write.
+ gate sync.RWMutex
+ mu sync.Mutex
+ entries map[credentialKey]*list.Element
+ lru list.List
+}
+
+var _ Store = (*CachedStore)(nil)
+
+type CacheOption func(*CachedStore) error
+
+// WithCacheTTL sets the cache lifetime. Zero disables caching.
+func WithCacheTTL(ttl time.Duration) CacheOption {
+ return func(s *CachedStore) error {
+ if ttl < 0 {
+ return fmt.Errorf("vault: cache TTL must not be negative")
+ }
+ s.ttl = ttl
+ return nil
+ }
+}
+
+// WithCacheSize sets the maximum number of cached token pairs.
+func WithCacheSize(size int) CacheOption {
+ return func(s *CachedStore) error {
+ if size <= 0 {
+ return fmt.Errorf("vault: cache size must be positive")
+ }
+ s.size = size
+ return nil
+ }
+}
+
+// NewCachedStore wraps a store with a 60-second, 10,000-entry cache by default.
+func NewCachedStore(inner Store, opts ...CacheOption) (*CachedStore, error) {
+ if inner == nil {
+ return nil, fmt.Errorf("vault: cache requires a store")
+ }
+ s := &CachedStore{
+ inner: inner, ttl: DefaultCacheTTL, size: DefaultCacheSize,
+ now: time.Now, entries: make(map[credentialKey]*list.Element),
+ }
+ for _, opt := range opts {
+ if err := opt(s); err != nil {
+ return nil, err
+ }
+ }
+ return s, nil
+}
+
+func cloneToken(tok Token) Token {
+ tok.Scopes = slices.Clone(tok.Scopes)
+ return tok
+}
+
+func (s *CachedStore) remove(key credentialKey) {
+ if element, ok := s.entries[key]; ok {
+ delete(s.entries, key)
+ s.lru.Remove(element)
+ }
+}
+
+func (s *CachedStore) remember(key credentialKey, tok Token) {
+ if s.ttl == 0 {
+ return
+ }
+ now := s.now()
+ deadline := now.Add(s.ttl)
+ if !tok.ExpiresAt.IsZero() && tok.ExpiresAt.Add(-cacheExpirySkew).Before(deadline) {
+ deadline = tok.ExpiresAt.Add(-cacheExpirySkew)
+ }
+ if !deadline.After(now) {
+ return
+ }
+ s.remove(key)
+ s.entries[key] = s.lru.PushFront(cacheEntry{key: key, token: cloneToken(tok), expiresAt: deadline})
+ if s.lru.Len() > s.size {
+ s.remove(s.lru.Back().Value.(cacheEntry).key)
+ }
+}
+
+func (s *CachedStore) Get(user, service string) (Token, error) {
+ if s.ttl == 0 {
+ return s.inner.Get(user, service)
+ }
+ s.gate.RLock()
+ defer s.gate.RUnlock()
+ key := credentialKey{user, service}
+ s.mu.Lock()
+ if element, ok := s.entries[key]; ok {
+ entry := element.Value.(cacheEntry)
+ if s.now().Before(entry.expiresAt) {
+ s.lru.MoveToFront(element)
+ tok := cloneToken(entry.token)
+ s.mu.Unlock()
+ return tok, nil
+ }
+ s.remove(key)
+ }
+ s.mu.Unlock()
+
+ tok, err := s.inner.Get(user, service)
+ if err != nil {
+ return Token{}, err
+ }
+ s.mu.Lock()
+ s.remember(key, tok)
+ s.mu.Unlock()
+ return tok, nil
+}
+
+func (s *CachedStore) Put(user, service string, tok Token) error {
+ s.gate.Lock()
+ defer s.gate.Unlock()
+ key := credentialKey{user, service}
+ s.mu.Lock()
+ s.remove(key)
+ s.mu.Unlock()
+ if err := s.inner.Put(user, service, tok); err != nil {
+ return err
+ }
+ s.mu.Lock()
+ s.remember(key, tok)
+ s.mu.Unlock()
+ return nil
+}
+
+func (s *CachedStore) Delete(user, service string) error {
+ s.gate.Lock()
+ defer s.gate.Unlock()
+ s.mu.Lock()
+ s.remove(credentialKey{user, service})
+ s.mu.Unlock()
+ return s.inner.Delete(user, service)
+}
+
+func (s *CachedStore) ListServices(user string) ([]string, error) {
+ return s.inner.ListServices(user)
+}
diff --git a/internal/vault/cached_store_test.go b/internal/vault/cached_store_test.go
new file mode 100644
index 0000000..367cc1c
--- /dev/null
+++ b/internal/vault/cached_store_test.go
@@ -0,0 +1,334 @@
+package vault
+
+import (
+ "errors"
+ "sync"
+ "testing"
+ "time"
+)
+
+type tokenPair struct {
+ user, service string
+}
+
+type countingStore struct {
+ mu sync.Mutex
+ tokens map[tokenPair]Token
+ gets, puts int
+ putErr error
+ deleteErr error
+ getStarted chan struct{}
+ getContinue chan struct{}
+}
+
+func newCountingStore() *countingStore {
+ return &countingStore{tokens: make(map[tokenPair]Token)}
+}
+
+func (s *countingStore) Get(user, service string) (Token, error) {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ s.gets++
+ tok, ok := s.tokens[tokenPair{user, service}]
+ if s.getStarted != nil {
+ close(s.getStarted)
+ <-s.getContinue
+ s.getStarted = nil
+ }
+ if !ok {
+ return Token{}, ErrTokenNotFound
+ }
+ return tok, nil
+}
+
+func (s *countingStore) Put(user, service string, tok Token) error {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ s.puts++
+ if s.putErr != nil {
+ return s.putErr
+ }
+ s.tokens[tokenPair{user, service}] = tok
+ return nil
+}
+
+func (s *countingStore) Delete(user, service string) error {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ if s.deleteErr != nil {
+ return s.deleteErr
+ }
+ delete(s.tokens, tokenPair{user, service})
+ return nil
+}
+
+func (s *countingStore) ListServices(user string) ([]string, error) {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ var services []string
+ for key := range s.tokens {
+ if key.user == user {
+ services = append(services, key.service)
+ }
+ }
+ return services, nil
+}
+
+func newTestCache(t *testing.T, inner Store, opts ...CacheOption) *CachedStore {
+ t.Helper()
+ store, err := NewCachedStore(inner, opts...)
+ if err != nil {
+ t.Fatal(err)
+ }
+ return store
+}
+
+func TestCachedStoreHitAndKeyIsolation(t *testing.T) {
+ inner := newCountingStore()
+ _ = inner.Put("a:b", "c", Token{AccessToken: "first"})
+ _ = inner.Put("a", "b:c", Token{AccessToken: "second"})
+ store := newTestCache(t, inner)
+ for range 3 {
+ for _, tc := range []struct{ user, service, want string }{
+ {"a:b", "c", "first"},
+ {"a", "b:c", "second"},
+ } {
+ tok, err := store.Get(tc.user, tc.service)
+ if err != nil || tok.AccessToken != tc.want {
+ t.Fatalf("Get = %v, %v", tok, err)
+ }
+ }
+ }
+ if inner.gets != 2 {
+ t.Fatalf("inner Gets = %d, want 2", inner.gets)
+ }
+}
+
+func TestCachedStoreLifetime(t *testing.T) {
+ for _, tc := range []struct {
+ name string
+ expiry time.Duration
+ life time.Duration
+ }{
+ {"ttl", time.Hour, time.Minute},
+ {"skew", 50 * time.Second, 20 * time.Second},
+ {"no-expiry", 0, time.Minute},
+ {"within-skew", 20 * time.Second, 0},
+ {"expired", -time.Second, 0},
+ } {
+ t.Run(tc.name, func(t *testing.T) {
+ now := time.Date(2026, 10, 18, 0, 0, 0, 0, time.UTC)
+ inner := newCountingStore()
+ tok := Token{AccessToken: "token"}
+ if tc.expiry != 0 {
+ tok.ExpiresAt = now.Add(tc.expiry)
+ }
+ _ = inner.Put("u", "s", tok)
+ store := newTestCache(t, inner)
+ store.now = func() time.Time { return now }
+ _, _ = store.Get("u", "s")
+ if tc.life > 0 {
+ now = now.Add(tc.life - time.Nanosecond)
+ _, _ = store.Get("u", "s")
+ if inner.gets != 1 {
+ t.Fatal("cache expired early")
+ }
+ now = now.Add(time.Nanosecond)
+ }
+ _, _ = store.Get("u", "s")
+ if inner.gets != 2 {
+ t.Fatalf("inner Gets = %d, want 2", inner.gets)
+ }
+ })
+ }
+}
+
+func TestCachedStoreWriteThrough(t *testing.T) {
+ inner := newCountingStore()
+ store := newTestCache(t, inner)
+ if err := store.Put("u", "s", Token{AccessToken: "old"}); err != nil {
+ t.Fatal(err)
+ }
+ if err := store.Put("u", "s", Token{AccessToken: "new"}); err != nil {
+ t.Fatal(err)
+ }
+ tok, err := store.Get("u", "s")
+ if err != nil || tok.AccessToken != "new" || inner.gets != 0 {
+ t.Fatalf("write-through Get = %v, %v; inner Gets = %d", tok, err, inner.gets)
+ }
+ services, err := store.ListServices("u")
+ if err != nil || len(services) != 1 || services[0] != "s" {
+ t.Fatalf("ListServices = %v, %v", services, err)
+ }
+ if err := store.Delete("u", "s"); err != nil {
+ t.Fatal(err)
+ }
+ if _, err := store.Get("u", "s"); !errors.Is(err, ErrTokenNotFound) {
+ t.Fatalf("Get after Delete = %v", err)
+ }
+ if len(store.entries) != 0 {
+ t.Fatal("Delete retained plaintext")
+ }
+}
+
+func TestCachedStoreWriteFailureInvalidates(t *testing.T) {
+ for _, operation := range []string{"put", "delete"} {
+ t.Run(operation, func(t *testing.T) {
+ inner := newCountingStore()
+ store := newTestCache(t, inner)
+ _ = store.Put("u", "s", Token{AccessToken: "old"})
+ want := errors.New("database unavailable")
+ var err error
+ if operation == "put" {
+ inner.putErr = want
+ err = store.Put("u", "s", Token{AccessToken: "new"})
+ } else {
+ inner.deleteErr = want
+ err = store.Delete("u", "s")
+ }
+ if !errors.Is(err, want) || len(store.entries) != 0 {
+ t.Fatalf("write = %v; entries = %d", err, len(store.entries))
+ }
+ tok, err := store.Get("u", "s")
+ if err != nil || tok.AccessToken != "old" || inner.gets != 1 {
+ t.Fatalf("Get = %v, %v; inner Gets = %d", tok, err, inner.gets)
+ }
+ })
+ }
+}
+
+func TestCachedStoreLRU(t *testing.T) {
+ inner := newCountingStore()
+ for _, service := range []string{"a", "b", "c"} {
+ _ = inner.Put("u", service, Token{AccessToken: service})
+ }
+ store := newTestCache(t, inner, WithCacheSize(2))
+ for _, service := range []string{"a", "b", "a", "c", "a"} {
+ _, _ = store.Get("u", service)
+ }
+ if inner.gets != 3 || len(store.entries) != 2 {
+ t.Fatalf("Gets = %d, entries = %d", inner.gets, len(store.entries))
+ }
+ _, _ = store.Get("u", "b")
+ if inner.gets != 4 || len(store.entries) != 2 {
+ t.Fatal("least recently used entry was not evicted")
+ }
+}
+
+func TestCachedStoreDisabledAndErrors(t *testing.T) {
+ inner := newCountingStore()
+ store := newTestCache(t, inner, WithCacheTTL(0))
+ _ = store.Put("u", "s", Token{AccessToken: "token"})
+ for range 2 {
+ _, _ = store.Get("u", "s")
+ }
+ if inner.gets != 2 || len(store.entries) != 0 {
+ t.Fatal("disabled cache retained a token")
+ }
+ store = newTestCache(t, inner)
+ for range 2 {
+ if _, err := store.Get("u", "missing"); !errors.Is(err, ErrTokenNotFound) {
+ t.Fatalf("Get = %v", err)
+ }
+ }
+ if inner.gets != 4 {
+ t.Fatal("Get errors were cached")
+ }
+ for _, opt := range []CacheOption{WithCacheTTL(-time.Second), WithCacheSize(0), WithCacheSize(-1)} {
+ if _, err := NewCachedStore(inner, opt); err == nil {
+ t.Fatal("invalid option accepted")
+ }
+ }
+ if _, err := NewCachedStore(nil); err == nil {
+ t.Fatal("nil store accepted")
+ }
+}
+
+func TestCachedStoreScopesAreCopied(t *testing.T) {
+ inner := newCountingStore()
+ store := newTestCache(t, inner)
+ tok := Token{AccessToken: "token", Scopes: []string{"read"}}
+ _ = store.Put("u", "s", tok)
+ tok.Scopes[0] = "changed"
+ got, _ := store.Get("u", "s")
+ if got.Scopes[0] != "read" {
+ t.Fatal("Put retained caller scopes")
+ }
+ got.Scopes[0] = "changed-again"
+ got, _ = store.Get("u", "s")
+ if got.Scopes[0] != "read" {
+ t.Fatal("Get exposed cached scopes")
+ }
+}
+
+func TestCachedStoreConcurrentReadWrite(t *testing.T) {
+ inner := newCountingStore()
+ store := newTestCache(t, inner)
+ var wg sync.WaitGroup
+ for range 20 {
+ wg.Go(func() {
+ for range 20 {
+ _ = store.Put("u", "s", Token{AccessToken: "token", Scopes: []string{"read"}})
+ tok, _ := store.Get("u", "s")
+ if len(tok.Scopes) > 0 {
+ tok.Scopes[0] = "private"
+ }
+ _ = store.Delete("u", "s")
+ }
+ })
+ }
+ wg.Wait()
+}
+
+func TestCachedStoreReadCannotRepopulateAfterWrite(t *testing.T) {
+ for _, operation := range []string{"put", "delete", "failed-put", "failed-delete"} {
+ t.Run(operation, func(t *testing.T) {
+ inner := newCountingStore()
+ _ = inner.Put("u", "s", Token{AccessToken: "old"})
+ inner.getStarted = make(chan struct{})
+ inner.getContinue = make(chan struct{})
+ if operation == "failed-put" {
+ inner.putErr = errors.New("write failed")
+ }
+ if operation == "failed-delete" {
+ inner.deleteErr = errors.New("delete failed")
+ }
+ store := newTestCache(t, inner)
+ readDone := make(chan struct{})
+ go func() {
+ defer close(readDone)
+ _, _ = store.Get("u", "s")
+ }()
+ awaitSignal(t, inner.getStarted)
+ writeDone := make(chan struct{})
+ go func() {
+ defer close(writeDone)
+ if operation == "put" || operation == "failed-put" {
+ _ = store.Put("u", "s", Token{AccessToken: "new"})
+ } else {
+ _ = store.Delete("u", "s")
+ }
+ }()
+ close(inner.getContinue)
+ awaitSignal(t, readDone)
+ awaitSignal(t, writeDone)
+ if operation != "put" && len(store.entries) != 0 {
+ t.Fatal("read repopulated the cache after invalidation")
+ }
+ tok, err := store.Get("u", "s")
+ if operation == "delete" {
+ if !errors.Is(err, ErrTokenNotFound) {
+ t.Fatalf("Get after Delete = %v", err)
+ }
+ } else {
+ want := "old"
+ if operation == "put" {
+ want = "new"
+ }
+ if err != nil || tok.AccessToken != want {
+ t.Fatalf("Get = %v, %v; want %s", tok, err, want)
+ }
+ }
+ })
+ }
+}
diff --git a/internal/vault/refresh.go b/internal/vault/refresh.go
index efadbc2..fcd936c 100644
--- a/internal/vault/refresh.go
+++ b/internal/vault/refresh.go
@@ -1,7 +1,9 @@
package vault
import (
+ "context"
"fmt"
+ "sync"
"time"
)
@@ -9,10 +11,63 @@ import (
// refresh token and returns new access/refresh tokens and expiry.
type RefreshFunc func(refreshToken string) (newAccess, newRefresh string, expiresIn time.Duration, err error)
-// GetOrRefresh retrieves a token and auto-refreshes if it expires within
-// the given buffer period (default 5 minutes). If refresh fails, returns
-// an error indicating the user must re-link.
-func GetOrRefresh(store Store, userID, service string, refreshFn RefreshFunc) (Token, error) {
+type refreshKey struct {
+ store Store
+ credentialKey
+}
+
+type refreshCall struct {
+ done chan struct{}
+ token Token
+ err error
+}
+
+var refreshFlights = struct {
+ sync.Mutex
+ calls map[refreshKey]*refreshCall
+}{calls: make(map[refreshKey]*refreshCall)}
+
+// GetOrRefresh refreshes tokens expiring within 5 minutes.
+// Concurrent callers sharing a store and token pair share one refresh.
+// Cancellation stops only the caller's wait, not the shared refresh.
+func GetOrRefresh(ctx context.Context, store Store, userID, service string, refreshFn RefreshFunc) (Token, error) {
+ if err := ctx.Err(); err != nil {
+ return Token{}, err
+ }
+ tok, err := store.Get(userID, service)
+ if err != nil {
+ return Token{}, err
+ }
+ if tok.ExpiresAt.IsZero() || time.Until(tok.ExpiresAt) > 5*time.Minute {
+ return tok, nil
+ }
+
+ key := refreshKey{store: store, credentialKey: credentialKey{userID, service}}
+ refreshFlights.Lock()
+ call, exists := refreshFlights.calls[key]
+ if !exists {
+ call = &refreshCall{done: make(chan struct{})}
+ refreshFlights.calls[key] = call
+ go func() {
+ call.token, call.err = refreshToken(store, userID, service, refreshFn)
+ refreshFlights.Lock()
+ delete(refreshFlights.calls, key)
+ close(call.done)
+ refreshFlights.Unlock()
+ }()
+ }
+ refreshFlights.Unlock()
+
+ select {
+ case <-ctx.Done():
+ return Token{}, ctx.Err()
+ case <-call.done:
+ return cloneToken(call.token), call.err
+ }
+}
+
+func refreshToken(store Store, userID, service string, refreshFn RefreshFunc) (Token, error) {
+ // Re-read after joining the flight. Another refresh may have finished after the caller's first read.
tok, err := store.Get(userID, service)
if err != nil {
return Token{}, err
@@ -23,9 +78,9 @@ func GetOrRefresh(store Store, userID, service string, refreshFn RefreshFunc) (T
return tok, nil
}
- // Token is expired or expiring soon — attempt refresh.
+ // Token is expired or expiring soon. Attempt refresh.
if tok.RefreshToken == "" {
- return Token{}, fmt.Errorf("vault: token for %s:%s expired and no refresh token available — user must re-link", userID, service)
+ return Token{}, fmt.Errorf("vault: token for %s:%s expired and no refresh token available. User must re-link", userID, service)
}
if refreshFn == nil {
@@ -34,7 +89,7 @@ func GetOrRefresh(store Store, userID, service string, refreshFn RefreshFunc) (T
newAccess, newRefresh, expiresIn, err := refreshFn(tok.RefreshToken)
if err != nil {
- return Token{}, fmt.Errorf("vault: refresh failed for %s:%s — user must re-link: %w", userID, service, err)
+ return Token{}, fmt.Errorf("vault: refresh failed for %s:%s. User must re-link: %w", userID, service, err)
}
tok.AccessToken = newAccess
diff --git a/internal/vault/refresh_test.go b/internal/vault/refresh_test.go
index 9008209..00bf07d 100644
--- a/internal/vault/refresh_test.go
+++ b/internal/vault/refresh_test.go
@@ -1,7 +1,10 @@
package vault
import (
+ "context"
"errors"
+ "sync"
+ "sync/atomic"
"testing"
"time"
)
@@ -16,7 +19,7 @@ func TestGetOrRefresh_FreshToken(t *testing.T) {
})
refreshCalled := false
- tok, err := GetOrRefresh(store, "user-1", "github", func(rt string) (string, string, time.Duration, error) {
+ tok, err := GetOrRefresh(context.Background(), store, "user-1", "github", func(rt string) (string, string, time.Duration, error) {
refreshCalled = true
return "", "", 0, nil
})
@@ -41,7 +44,7 @@ func TestGetOrRefresh_ExpiredToken(t *testing.T) {
ExpiresAt: time.Now().Add(-time.Hour),
})
- tok, err := GetOrRefresh(store, "user-1", "github", func(rt string) (string, string, time.Duration, error) {
+ tok, err := GetOrRefresh(context.Background(), store, "user-1", "github", func(rt string) (string, string, time.Duration, error) {
if rt != "refresh-123" {
t.Fatalf("unexpected refresh token: %s", rt)
}
@@ -71,7 +74,7 @@ func TestGetOrRefresh_RefreshFails(t *testing.T) {
ExpiresAt: time.Now().Add(-time.Hour),
})
- _, err := GetOrRefresh(store, "user-1", "github", func(rt string) (string, string, time.Duration, error) {
+ _, err := GetOrRefresh(context.Background(), store, "user-1", "github", func(rt string) (string, string, time.Duration, error) {
return "", "", 0, errors.New("refresh denied")
})
if err == nil {
@@ -88,8 +91,253 @@ func TestGetOrRefresh_NoRefreshToken(t *testing.T) {
ExpiresAt: time.Now().Add(-time.Hour),
})
- _, err := GetOrRefresh(store, "user-1", "github", nil)
+ _, err := GetOrRefresh(context.Background(), store, "user-1", "github", nil)
if err == nil {
t.Fatal("expected error when no refresh token")
}
}
+
+// Done signals that a caller has joined the flight and started waiting.
+type waitingContext struct {
+ context.Context
+ joined chan<- struct{}
+ once sync.Once
+}
+
+func (c *waitingContext) Done() <-chan struct{} {
+ c.once.Do(func() { c.joined <- struct{}{} })
+ return c.Context.Done()
+}
+
+func awaitSignal(t *testing.T, ch <-chan struct{}) {
+ t.Helper()
+ select {
+ case <-ch:
+ case <-time.After(5 * time.Second):
+ t.Fatal("timed out waiting for refresh")
+ }
+}
+
+func TestGetOrRefreshSingleflight(t *testing.T) {
+ for _, fail := range []bool{false, true} {
+ t.Run(map[bool]string{false: "success", true: "failure"}[fail], func(t *testing.T) {
+ inner := newCountingStore()
+ store := newTestCache(t, inner)
+ _ = store.Put("u", "s", Token{
+ AccessToken: "old", RefreshToken: "refresh", ExpiresAt: time.Now().Add(time.Minute),
+ Scopes: []string{"read"},
+ })
+ joined := make(chan struct{}, 50)
+ release := make(chan struct{})
+ defer func() {
+ select {
+ case <-release:
+ default:
+ close(release)
+ }
+ }()
+ wantErr := errors.New("provider denied refresh")
+ var refreshes atomic.Int32
+ refresh := func(rt string) (string, string, time.Duration, error) {
+ refreshes.Add(1)
+ <-release
+ if fail {
+ return "", "", 0, wantErr
+ }
+ if rt != "refresh" {
+ return "", "", 0, errors.New("wrong refresh token")
+ }
+ return "new", "rotated", time.Hour, nil
+ }
+ var wg sync.WaitGroup
+ for range 50 {
+ wg.Go(func() {
+ ctx := &waitingContext{Context: context.Background(), joined: joined}
+ tok, err := GetOrRefresh(ctx, store, "u", "s", refresh)
+ if fail {
+ if !errors.Is(err, wantErr) {
+ t.Errorf("refresh error = %v", err)
+ }
+ } else if err != nil || tok.AccessToken != "new" || tok.RefreshToken != "rotated" {
+ t.Errorf("refresh = %v, %v", tok, err)
+ } else {
+ tok.Scopes[0] = "private"
+ }
+ })
+ }
+ for range 50 {
+ awaitSignal(t, joined)
+ }
+ close(release)
+ wg.Wait()
+ if refreshes.Load() != 1 {
+ t.Fatalf("refreshes = %d, want 1", refreshes.Load())
+ }
+ wantPuts := 2
+ if fail {
+ wantPuts = 1
+ }
+ if inner.puts != wantPuts {
+ t.Fatalf("Puts = %d, want %d", inner.puts, wantPuts)
+ }
+ _, err := GetOrRefresh(context.Background(), store, "u", "s", refresh)
+ if fail {
+ if !errors.Is(err, wantErr) || refreshes.Load() != 2 {
+ t.Fatalf("retry = %v; refreshes = %d", err, refreshes.Load())
+ }
+ } else if err != nil || refreshes.Load() != 1 {
+ t.Fatalf("fresh Get = %v; refreshes = %d", err, refreshes.Load())
+ }
+ })
+ }
+}
+
+func TestGetOrRefreshCancellation(t *testing.T) {
+ inner := newCountingStore()
+ store := newTestCache(t, inner)
+ _ = store.Put("u", "s", Token{
+ AccessToken: "old", RefreshToken: "refresh", ExpiresAt: time.Now().Add(-time.Minute),
+ })
+ started := make(chan struct{})
+ release := make(chan struct{})
+ defer close(release)
+ refresh := func(string) (string, string, time.Duration, error) {
+ close(started)
+ <-release
+ return "new", "", time.Hour, nil
+ }
+ leaderCtx, cancelLeader := context.WithCancel(context.Background())
+ leaderDone := make(chan error, 1)
+ go func() {
+ _, err := GetOrRefresh(leaderCtx, store, "u", "s", refresh)
+ leaderDone <- err
+ }()
+ awaitSignal(t, started)
+ joined := make(chan struct{}, 2)
+ ctx, cancel := context.WithCancel(context.Background())
+ defer cancel()
+ waiterDone := make(chan error, 1)
+ go func() {
+ _, err := GetOrRefresh(&waitingContext{Context: ctx, joined: joined}, store, "u", "s", refresh)
+ waiterDone <- err
+ }()
+ survivorDone := make(chan error, 1)
+ go func() {
+ tok, err := GetOrRefresh(&waitingContext{Context: context.Background(), joined: joined}, store, "u", "s", refresh)
+ if err == nil && (tok.AccessToken != "new" || tok.RefreshToken != "refresh") {
+ err = errors.New("incorrect refreshed token")
+ }
+ survivorDone <- err
+ }()
+ awaitSignal(t, joined)
+ awaitSignal(t, joined)
+ cancel()
+ cancelLeader()
+ for _, done := range []<-chan error{waiterDone, leaderDone} {
+ select {
+ case err := <-done:
+ if !errors.Is(err, context.Canceled) {
+ t.Fatalf("canceled caller = %v", err)
+ }
+ case <-time.After(time.Second):
+ t.Fatal("cancellation blocked on refresh")
+ }
+ }
+ release <- struct{}{}
+ select {
+ case err := <-survivorDone:
+ if err != nil {
+ t.Fatal(err)
+ }
+ case <-time.After(5 * time.Second):
+ t.Fatal("refresh did not continue for surviving caller")
+ }
+}
+
+func TestGetOrRefreshMissingAndCanceled(t *testing.T) {
+ store := newCountingStore()
+ if _, err := GetOrRefresh(context.Background(), store, "u", "s", nil); !errors.Is(err, ErrTokenNotFound) {
+ t.Fatalf("missing token = %v", err)
+ }
+ ctx, cancel := context.WithCancel(context.Background())
+ cancel()
+ if _, err := GetOrRefresh(ctx, store, "u", "s", nil); !errors.Is(err, context.Canceled) {
+ t.Fatalf("canceled caller = %v", err)
+ }
+ if store.gets != 1 {
+ t.Fatal("canceled caller touched the store")
+ }
+}
+
+func TestGetOrRefreshPutFailure(t *testing.T) {
+ inner := newCountingStore()
+ store := newTestCache(t, inner)
+ _ = store.Put("u", "s", Token{
+ AccessToken: "old", RefreshToken: "refresh", ExpiresAt: time.Now().Add(-time.Minute),
+ })
+ want := errors.New("database write failed")
+ inner.putErr = want
+ refreshes := 0
+ refresh := func(string) (string, string, time.Duration, error) {
+ refreshes++
+ return "new", "", time.Hour, nil
+ }
+ for range 2 {
+ if _, err := GetOrRefresh(context.Background(), store, "u", "s", refresh); !errors.Is(err, want) {
+ t.Fatalf("Put error = %v", err)
+ }
+ if len(store.entries) != 0 {
+ t.Fatal("failed write cached refreshed credentials")
+ }
+ }
+ if refreshes != 2 {
+ t.Fatalf("refreshes = %d, want 2", refreshes)
+ }
+}
+
+func TestGetOrRefreshIndependentPairsAndStores(t *testing.T) {
+ blocked := newCountingStore()
+ for _, service := range []string{"a", "b"} {
+ _ = blocked.Put("u", service, Token{
+ AccessToken: "old", RefreshToken: "refresh", ExpiresAt: time.Now().Add(-time.Minute),
+ })
+ }
+ other := newCountingStore()
+ _ = other.Put("u", "a", Token{
+ AccessToken: "other", RefreshToken: "other-refresh", ExpiresAt: time.Now().Add(-time.Minute),
+ })
+ started, release := make(chan struct{}), make(chan struct{})
+ done := make(chan error, 1)
+ go func() {
+ _, err := GetOrRefresh(context.Background(), blocked, "u", "a", func(string) (string, string, time.Duration, error) {
+ close(started)
+ <-release
+ return "new", "", time.Hour, nil
+ })
+ done <- err
+ }()
+ awaitSignal(t, started)
+ defer func() {
+ close(release)
+ if err := <-done; err != nil {
+ t.Error(err)
+ }
+ }()
+ ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+ defer cancel()
+ for _, tc := range []struct {
+ store Store
+ service string
+ want string
+ }{
+ {blocked, "b", "new-b"},
+ {other, "a", "new-other"},
+ } {
+ tok, err := GetOrRefresh(ctx, tc.store, "u", tc.service, func(string) (string, string, time.Duration, error) {
+ return tc.want, "", time.Hour, nil
+ })
+ if err != nil || tok.AccessToken != tc.want {
+ t.Fatalf("independent refresh = %v, %v", tok, err)
+ }
+ }
+}
diff --git a/internal/vault/sqlite_store_test.go b/internal/vault/sqlite_store_test.go
index 2ffb3ac..639dd9a 100644
--- a/internal/vault/sqlite_store_test.go
+++ b/internal/vault/sqlite_store_test.go
@@ -9,7 +9,7 @@ import (
_ "modernc.org/sqlite"
)
-func setupTestDB(t *testing.T) *sql.DB {
+func setupTestDB(t testing.TB) *sql.DB {
t.Helper()
db, err := sql.Open("sqlite", ":memory:")
if err != nil {
From c18cd082cefb4b00082d2f5a9fff72461e31bbec Mon Sep 17 00:00:00 2001
From: Shreyansh Sancheti <43677304+shreyanshjain7174@users.noreply.github.com>
Date: Sun, 18 Oct 2026 05:20:05 +0530
Subject: [PATCH 2/3] 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>
---
README.md | 9 +
internal/gateway/gateway.go | 8 +
internal/gateway/gateway_test.go | 4 +-
internal/gateway/path_test.go | 3 +-
internal/gateway/refresh_cancellation_test.go | 78 ++++++
internal/oauth/handler.go | 42 +++-
internal/oauth/linked_page_test.go | 17 ++
internal/oauth/refresh_cancellation_test.go | 62 +++++
internal/oauth/refresh_test.go | 3 +-
internal/vault/cached_store.go | 17 ++
internal/vault/cached_store_test.go | 24 ++
internal/vault/refresh.go | 60 ++++-
internal/vault/refresh_safety_test.go | 222 ++++++++++++++++++
internal/vault/refresh_test.go | 24 +-
internal/vault/sqlite_store.go | 57 +++++
internal/vault/vault.go | 51 +++-
16 files changed, 635 insertions(+), 46 deletions(-)
create mode 100644 internal/gateway/refresh_cancellation_test.go
create mode 100644 internal/oauth/linked_page_test.go
create mode 100644 internal/oauth/refresh_cancellation_test.go
create mode 100644 internal/vault/refresh_safety_test.go
diff --git a/README.md b/README.md
index 08046e7..e1121b1 100644
--- a/README.md
+++ b/README.md
@@ -322,6 +322,15 @@ 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.
+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.
diff --git a/internal/gateway/gateway.go b/internal/gateway/gateway.go
index 28930f0..a6f4954 100644
--- a/internal/gateway/gateway.go
+++ b/internal/gateway/gateway.go
@@ -459,6 +459,14 @@ func (s *Server) executeAttempt(ctx context.Context, att *attempt) *outcome {
if refresher := s.cfg.Refreshers[att.req.Service]; refresher != nil {
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(`
+
+Account linked!
+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, `
-
-Account linked!
-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(), "