Skip to content

Commit a885a4d

Browse files
coderdancoderdan
committed
fix(golang): finish a device session's refresh in the background after its caller gives up, so a rotated refresh token is kept
The token service rotates the refresh token when it answers a device session's refresh. A refresh cancelled with its caller after the service answered dropped the new token, and the next refresh sent one the service had retired, which it rejects: the strategy could not refresh again without a new login. This is the failure stack-auth's native refresh now avoids by finishing the refresh in a task after its caller leaves (#1172); the Go auth guest runs that refresh inline in WebAssembly, so the Go host has to do the same. Strategy.Token now runs its work on a goroutine that holds the strategy's lock until it is done, and returns to a caller whose context ends at once, without making a copy of a token nobody will read. The refresh call's context is marked, and only its requests run under guest.Detach of the caller's context with a grace of 10 seconds: the refresh finishes, the guest saves what came back, and the next Token reads it. Past the grace it fails as a transport failure, well inside the 30-second backstop, so the profile stays open. Close waits for such a refresh, so its rotated token is saved. Every other request to the token service, an access-key or OIDC exchange, rotates nothing and is still cancelled with its caller, as is a ZeroKMS request, so a cancelled NewClient or a shutdown does not wait on them. Refs #1169.
1 parent 4fa1ee2 commit a885a4d

9 files changed

Lines changed: 235 additions & 37 deletions

File tree

‎languages/golang/CHANGELOG.md‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -82,11 +82,13 @@ listed under **Breaking**. Releases are tagged `languages/golang/vX.Y.Z`.
8282
or the `auth` profile behind its strategy, for every later caller, and a
8383
request whose deadline passed while it waited for the client closed it
8484
before running. Now a call waiting for the client leaves the queue with
85-
the context's error; a ZeroKMS request or a token fetch in flight is
86-
cancelled, and the error matches both the context's error and
87-
`ErrTransport`; and the client stays open. A client or profile still
88-
closes itself (`ErrState`) when the guest traps, or when a call runs on
89-
for 30 seconds after its context ended.
85+
the context's error; a ZeroKMS request in flight is cancelled, and the
86+
error matches both the context's error and `ErrTransport`; a device
87+
session's refresh in flight finishes in the background, for up to 10
88+
seconds, so a rotated refresh token is never lost; and the client stays
89+
open. A client or profile still closes itself (`ErrState`) when the
90+
guest traps, or when a call runs on for 30 seconds after its context
91+
ended.
9092
- **stashgen refuses a match index on an index-only field**
9193
(`stash:"name,index=match"` with no `encrypt`). A match query's rows are
9294
checked by decrypting them, and such a field stores no ciphertext.

‎languages/golang/auth/cancel_test.go‎

Lines changed: 110 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,14 @@ package auth
22

33
import (
44
"context"
5+
"encoding/json"
56
"errors"
67
"fmt"
78
"io"
89
"net/http"
910
"net/http/httptest"
10-
"strings"
11+
"os"
12+
"path/filepath"
1113
"sync/atomic"
1214
"testing"
1315
"time"
@@ -66,13 +68,114 @@ func TestATokenFetchCancelledFailsOnlyItself(t *testing.T) {
6668
if _, err := strategy.Token(ctx); !errors.Is(err, context.DeadlineExceeded) {
6769
t.Fatalf("cancelled Token: %v, want the deadline", err)
6870
}
69-
if profile.root.closed || profile.root.inst.module.IsClosed() {
70-
t.Fatal("the profile was closed")
71-
}
7271
hold.Store(false)
72+
// The next Token waits for the abandoned one to finish.
7373
if token, err := strategy.Token(context.Background()); err != nil || token == "" {
7474
t.Fatalf("the next Token: %q, %v", token, err)
7575
}
76+
if profile.root.closed || profile.root.inst.module.IsClosed() {
77+
t.Fatal("the profile was closed")
78+
}
79+
}
80+
81+
// An access-key exchange rotates nothing, so it ends with its caller: the
82+
// held request is cancelled at the caller's deadline, not run on for the
83+
// refresh grace, and Close does not wait for it.
84+
func TestAnAccessKeyExchangeEndsWithItsCaller(t *testing.T) {
85+
guestOrSkip(t)
86+
var hold atomic.Bool
87+
server := tokenServer(t, &hold)
88+
profile, err := Open(context.Background(), t.TempDir())
89+
if err != nil {
90+
t.Fatal(err)
91+
}
92+
defer profile.Close()
93+
strategy, err := profile.AccessKey(context.Background(), testCRN, "CSAKtestKeyId.testKeySecret", WithBaseURL(server.URL))
94+
if err != nil {
95+
t.Fatal(err)
96+
}
97+
98+
hold.Store(true)
99+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
100+
defer cancel()
101+
if _, err := strategy.Token(ctx); !errors.Is(err, context.DeadlineExceeded) {
102+
t.Fatalf("cancelled Token: %v, want the deadline", err)
103+
}
104+
// The grace is 10 seconds; a Close that waited for it would take that.
105+
start := time.Now()
106+
if err := strategy.Close(); err != nil {
107+
t.Fatal(err)
108+
}
109+
if waited := time.Since(start); waited > 2*time.Second {
110+
t.Fatalf("Close took %v: the exchange ran on past its caller", waited)
111+
}
112+
}
113+
114+
// A device session's refresh is the request that rotates a credential: the
115+
// service answers with a new refresh token and retires the old one. A
116+
// caller that gives up returns at once, but the refresh runs on, and what
117+
// it brings is saved. The next Token is served from it with no second
118+
// request, and the profile on disk holds the rotated refresh token.
119+
func TestARefreshAnsweredAfterItsCallerGaveUpIsKept(t *testing.T) {
120+
guestOrSkip(t)
121+
dir, _ := expiredDeviceProfile(t)
122+
var requests atomic.Int32
123+
answered := make(chan struct{})
124+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
125+
if requests.Add(1) > 1 {
126+
// The old refresh token is retired once the service answers.
127+
w.WriteHeader(http.StatusBadRequest)
128+
fmt.Fprint(w, `{"error":"invalid_grant"}`)
129+
return
130+
}
131+
_, _ = io.Copy(io.Discard, r.Body)
132+
time.Sleep(150 * time.Millisecond)
133+
fmt.Fprint(w, `{"access_token":"fresh","token_type":"Bearer","expires_in":3600,"refresh_token":"refresh-2"}`)
134+
close(answered)
135+
}))
136+
t.Cleanup(server.Close)
137+
profile, err := Open(context.Background(), dir)
138+
if err != nil {
139+
t.Fatal(err)
140+
}
141+
defer profile.Close()
142+
ws, err := profile.WorkspaceStore(context.Background(), wsA)
143+
if err != nil {
144+
t.Fatal(err)
145+
}
146+
strategy, err := ws.DeviceSession(context.Background(), WithBaseURL(server.URL))
147+
if err != nil {
148+
t.Fatal(err)
149+
}
150+
defer strategy.Close()
151+
152+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
153+
defer cancel()
154+
if _, err := strategy.Token(ctx); !errors.Is(err, context.DeadlineExceeded) {
155+
t.Fatalf("cancelled Token: %v, want the deadline", err)
156+
}
157+
select {
158+
case <-answered:
159+
t.Fatal("the cancelled Token waited for the service's answer")
160+
default:
161+
}
162+
if token, err := strategy.Token(context.Background()); err != nil || token != "fresh" {
163+
t.Fatalf("the next Token: %q, %v, want the refreshed token", token, err)
164+
}
165+
if n := requests.Load(); n != 1 {
166+
t.Fatalf("%d refresh requests, want 1: the late answer was dropped", n)
167+
}
168+
data, err := os.ReadFile(filepath.Join(dir, "workspaces", wsA, "auth.json"))
169+
if err != nil {
170+
t.Fatal(err)
171+
}
172+
var stored map[string]any
173+
if err := json.Unmarshal(data, &stored); err != nil {
174+
t.Fatal(err)
175+
}
176+
if stored["refresh_token"] != "refresh-2" {
177+
t.Fatalf("stored refresh token = %v, want the rotated one", stored["refresh_token"])
178+
}
76179
}
77180

78181
func TestATokenRequestWhoseContextEndsInTheQueueNeverReachesTheGuest(t *testing.T) {
@@ -114,7 +217,7 @@ func TestATokenRequestWhoseContextEndsInTheQueueNeverReachesTheGuest(t *testing.
114217
}
115218

116219
// The backstop: a call still running long after its caller gave up closes
117-
// the profile, and says so.
220+
// the profile. The caller has returned by then; the next call is ErrState.
118221
func TestATokenFetchThatRunsOnAfterItsContextEndsClosesTheProfile(t *testing.T) {
119222
guestOrSkip(t)
120223
profile, err := Open(context.Background(), t.TempDir(), WithRoundTripper(roundTripFunc(func(*http.Request) (*http.Response, error) {
@@ -134,8 +237,8 @@ func TestATokenFetchThatRunsOnAfterItsContextEndsClosesTheProfile(t *testing.T)
134237
profile.root.stuck = 20 * time.Millisecond
135238
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
136239
defer cancel()
137-
if _, err := strategy.Token(ctx); err == nil || !strings.Contains(err.Error(), "the profile is closed") {
138-
t.Fatalf("stuck Token: %v, want the profile closed", err)
240+
if _, err := strategy.Token(ctx); !errors.Is(err, context.DeadlineExceeded) {
241+
t.Fatalf("stuck Token: %v, want the deadline", err)
139242
}
140243
if _, err := strategy.Token(context.Background()); !errors.Is(err, ErrState) {
141244
t.Fatalf("a Token after the close: %v, want ErrState", err)

‎languages/golang/auth/doc.go‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,9 +42,14 @@
4242
// A call whose context ends, [Strategy.Token] included, fails with the
4343
// context's error and leaves the store and its strategies open: a request
4444
// cancelled while its token is exchanged or refreshed does not stop the
45-
// encrypt clients that read tokens from the same store. The store closes
46-
// itself, and every later call is [ErrState], only when its guest traps or
47-
// a call runs on for 30 seconds after its context ended.
45+
// encrypt clients that read tokens from the same store. A request to the
46+
// token service is cancelled with its caller, except a device session's
47+
// refresh: the service rotates the refresh token when it answers one, so
48+
// the caller returns at once while the refresh runs on, for up to 10
49+
// seconds, and the store keeps what it sends back for the next call.
50+
// [Strategy.Close] waits for such a refresh. The store closes itself, and
51+
// every later call is [ErrState], only when its guest traps or a call runs
52+
// on for 30 seconds after its context ended.
4853
//
4954
// # Why a second guest
5055
//

‎languages/golang/auth/store.go‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,9 @@ func RequireLockedMemory() Option {
5757
// A call whose context ends fails with the context's error, and only that
5858
// call: a request cancelled while its token is fetched or refreshed never
5959
// closes the profile, or the strategies and encrypt clients that read
60-
// tokens from it. The profile closes itself, and every later call is
60+
// tokens from it. A device session's refresh already sent finishes in the
61+
// background, for up to 10 seconds, so a rotated refresh token is never
62+
// lost (see refreshGrace). The profile closes itself, and every later call is
6163
// [ErrState], only when the guest's state is unknown: it trapped, or it
6264
// went on running long after its caller gave up, which is a bug.
6365
type ProfileStore struct {
@@ -235,7 +237,8 @@ func (s *ProfileStore) call(ctx context.Context, fn export, args ...string) ([]b
235237
//
236238
// The caller's context bounds the wait for the lock, as encrypt's
237239
// Client.call does, and the guest runs under guest.Detach, so a caller
238-
// that gives up cancels its token request without closing the profile.
240+
// that gives up never closes the profile. A refresh request runs on for
241+
// refreshGrace, so the guest keeps the rotated token the service sends.
239242
func (s *ProfileStore) callArgs(ctx context.Context, fn export, args ...guest.Arg) ([]byte, error) {
240243
r := s.root
241244
if err := r.mu.Acquire(ctx); err != nil {

‎languages/golang/auth/strategy.go‎

Lines changed: 43 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -186,27 +186,58 @@ func (s *Strategy) Token(ctx context.Context) (string, error) {
186186
if err := s.mu.Acquire(ctx); err != nil {
187187
return "", err
188188
}
189-
defer s.mu.Unlock()
190189
if s.closed {
190+
s.mu.Unlock()
191191
return "", ErrState
192192
}
193-
ctx, status := withAuthHTTPStatus(ctx)
194-
token, err := s.token(ctx)
195-
return token, status.wrap(err)
193+
// The work runs on its own goroutine, which holds the lock until it is
194+
// done, so a caller that gives up returns at once. A device session's
195+
// refresh already sent finishes there, and the guest keeps what the
196+
// service sends (see refreshGrace): a rotated refresh token is never
197+
// lost. The token it brings is dropped; the next Token reads it. Any
198+
// other request ends with the caller.
199+
type result struct {
200+
token string
201+
err error
202+
}
203+
done := make(chan result, 1)
204+
go func() {
205+
defer s.mu.Unlock()
206+
ctx, status := withAuthHTTPStatus(ctx)
207+
token, err := s.token(ctx)
208+
done <- result{token, status.wrap(err)}
209+
}()
210+
select {
211+
case r := <-done:
212+
return r.token, r.err
213+
case <-ctx.Done():
214+
// A call that finished as the context ended still answers.
215+
select {
216+
case r := <-done:
217+
return r.token, r.err
218+
default:
219+
return "", ctx.Err()
220+
}
221+
}
196222
}
197223

198224
// token is Token under the strategy's lock: a read, then a locked refresh
199225
// for a device session that needs one.
200226
func (s *Strategy) token(ctx context.Context) (string, error) {
201-
call := func(export func(*instance) api.Function) (string, error) {
227+
call := func(ctx context.Context, export func(*instance) api.Function) (string, error) {
202228
out, err := s.store.callArgs(ctx, export, guest.BufArg([]byte(s.handle)))
203229
if err != nil {
204230
return "", err
205231
}
206232
defer guest.Wipe(out)
233+
// A caller that has left gets no copy: a string cannot be wiped,
234+
// and the next Token reads the guest's.
235+
if err := ctx.Err(); err != nil {
236+
return "", err
237+
}
207238
return string(out), nil
208239
}
209-
token, err := call(func(i *instance) api.Function { return i.authToken })
240+
token, err := call(ctx, func(i *instance) api.Function { return i.authToken })
210241
if !s.device || !errors.Is(err, guest.ErrAuthRefreshRequired) {
211242
return token, err
212243
}
@@ -216,7 +247,9 @@ func (s *Strategy) token(ctx context.Context) (string, error) {
216247
}
217248
err = withRefreshLock(ctx, path, func() error {
218249
var err error
219-
token, err = call(func(i *instance) api.Function { return i.authRefresh })
250+
// The one request that rotates a credential: it runs on after its
251+
// caller leaves (see refreshGrace).
252+
token, err = call(withRefresh(ctx), func(i *instance) api.Function { return i.authRefresh })
220253
return err
221254
})
222255
return token, err
@@ -232,6 +265,9 @@ func (s *Strategy) token(ctx context.Context) (string, error) {
232265
func (s *Strategy) MemoryLockError() error { return s.store.MemoryLockError() }
233266

234267
// Close drops the strategy's cached token and unregisters its provider.
268+
// It waits for a call in flight, including a device session's refresh that
269+
// a cancelled Token left running: that refresh finishes, for up to 10
270+
// seconds, so the rotated refresh token it brings is saved rather than lost.
235271
func (s *Strategy) Close() error {
236272
s.mu.Lock()
237273
defer s.mu.Unlock()

‎languages/golang/auth/transport.go‎

Lines changed: 43 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"sort"
1111
"strings"
1212
"sync"
13+
"time"
1314

1415
"github.com/cipherstash/stack/languages/golang/internal/guest"
1516
"github.com/tetratelabs/wazero"
@@ -70,7 +71,10 @@ func (s *authHTTPStatus) wrap(err error) error {
7071
}
7172

7273
type authTransport struct {
73-
rt http.RoundTripper
74+
rt http.RoundTripper
75+
// grace is how long a refresh request runs on after its caller gives
76+
// up: refreshGrace, shortened by tests.
77+
grace time.Duration
7478
mu sync.Mutex
7579
providers map[uint32]OIDCProvider
7680
nextProvider uint32
@@ -80,7 +84,7 @@ func newAuthTransport(rt http.RoundTripper) *authTransport {
8084
if rt == nil {
8185
rt = http.DefaultTransport
8286
}
83-
return &authTransport{rt: rt, providers: make(map[uint32]OIDCProvider)}
87+
return &authTransport{rt: rt, grace: refreshGrace, providers: make(map[uint32]OIDCProvider)}
8488
}
8589

8690
func (t *authTransport) register(provider OIDCProvider) uint32 {
@@ -101,6 +105,40 @@ func (t *authTransport) unregister(id uint32) {
101105
t.mu.Unlock()
102106
}
103107

108+
// refreshGrace is how long a device session's refresh request runs on after
109+
// its caller's context ends. The token service rotates the refresh token
110+
// when it answers a refresh, so a refresh abandoned after the service
111+
// answered would lose the new token, and the next refresh would send one
112+
// the service has retired. So a refresh is not cancelled with its caller:
113+
// the caller returns at once (Strategy.Token), the request finishes, and
114+
// the guest saves what came back. A refresh still running this long after
115+
// its caller gave up fails as a transport failure, well inside
116+
// guest.StuckAfter, so the profile stays open.
117+
//
118+
// Every other request to the token service (an access-key exchange, an
119+
// OIDC exchange) rotates nothing, and is cancelled with its caller.
120+
const refreshGrace = 10 * time.Second
121+
122+
// refreshKey marks the context of a device session's refresh call
123+
// (withRefresh), whose requests get refreshGrace.
124+
type refreshKey struct{}
125+
126+
// withRefresh marks ctx as a refresh call's.
127+
func withRefresh(ctx context.Context) context.Context {
128+
return context.WithValue(ctx, refreshKey{}, true)
129+
}
130+
131+
// requestContext is the context a request to the token service runs under:
132+
// the caller's, so a caller that gives up cancels it, except for a refresh,
133+
// which runs on for the grace (see refreshGrace). done must be called when
134+
// the request is finished with.
135+
func (t *authTransport) requestContext(ctx context.Context) (context.Context, func()) {
136+
if refresh, _ := ctx.Value(refreshKey{}).(bool); refresh {
137+
return guest.Detach(guest.Caller(ctx), t.grace)
138+
}
139+
return guest.Caller(ctx), func() {}
140+
}
141+
104142
func (t *authTransport) instantiate(ctx context.Context, r wazero.Runtime) error {
105143
_, err := r.NewHostModuleBuilder("cipherstash_transport").
106144
NewFunctionBuilder().WithFunc(t.send).Export("transport_send").
@@ -124,8 +162,9 @@ func (t *authTransport) send(ctx context.Context, m api.Module,
124162
return t.placeResponse(ctx, m, respHeadersPtrOut, respHeadersLenOut, respBodyPtrOut, respBodyLenOut, -1, nil, []byte("request buffer out of range"))
125163
}
126164
requestBody := newAuthRequestBody(body)
127-
// The caller's context, not the call's: see guest.Detach.
128-
req, err := http.NewRequestWithContext(guest.Caller(ctx), string(method), string(url), requestBody)
165+
reqCtx, done := t.requestContext(ctx)
166+
defer done()
167+
req, err := http.NewRequestWithContext(reqCtx, string(method), string(url), requestBody)
129168
if err != nil {
130169
_ = requestBody.Close()
131170
return t.placeResponse(ctx, m, respHeadersPtrOut, respHeadersLenOut, respBodyPtrOut, respBodyLenOut, -1, nil, []byte(err.Error()))

‎languages/golang/encrypt/README.md‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -123,9 +123,13 @@ request giving up never closes the shared client for the others.
123123

124124
- A call waiting for the client (it runs one call at a time) leaves the
125125
queue when its `ctx` ends, and never starts.
126-
- A call waiting on ZeroKMS, or on its auth strategy for a token, has that
127-
request cancelled. Its error matches both `context.Canceled` or
128-
`context.DeadlineExceeded` and `ErrTransport`.
126+
- A call waiting on ZeroKMS has that request cancelled. Its error matches
127+
both `context.Canceled` or `context.DeadlineExceeded` and `ErrTransport`.
128+
- A call waiting on its auth strategy for a token returns too, with the
129+
context's error. If the strategy is refreshing a device session, the
130+
refresh finishes in the background, for up to 10 seconds. The service
131+
rotates the refresh token when it answers, and dropping that answer
132+
would leave the strategy unable to refresh.
129133

130134
A client closes itself, and every later call returns `ErrState`, only when
131135
the guest's state is unknown: the guest trapped, or a call ran on for 30

0 commit comments

Comments
 (0)