Skip to content

Commit b5ed21e

Browse files
authored
Merge pull request #1263 from cipherstash/fix/go-client-cancel-isolation
fix(golang): a call whose context ends fails on its own, and leaves the client open
2 parents eca3b4f + 94f995d commit b5ed21e

23 files changed

Lines changed: 1209 additions & 130 deletions

‎languages/golang/CHANGELOG.md‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,18 @@ listed under **Breaking**. Releases are tagged `languages/golang/vX.Y.Z`.
8181

8282
### Changed
8383

84+
- **A call whose context ends fails on its own and leaves the client
85+
open.** A cancelled or timed-out call used to close the `encrypt.Client`,
86+
or the `auth` profile behind its strategy, for every later caller, and a
87+
request whose deadline passed while it waited for the client closed it
88+
before running. Now a call waiting for the client leaves the queue with
89+
the context's error; a ZeroKMS request in flight is cancelled, and the
90+
error matches both the context's error and `ErrTransport`; a device
91+
session's refresh in flight finishes in the background, for up to 10
92+
seconds, so a rotated refresh token is never lost; and the client stays
93+
open. A client or profile still closes itself (`ErrState`) when the
94+
guest traps, or when a call runs on for 30 seconds after its context
95+
ended.
8496
- **stashgen refuses a match index on an index-only field**
8597
(`stash:"name,index=match"` with no `encrypt`). A match query's rows are
8698
checked by decrypting them, and such a field stores no ciphertext.
Lines changed: 303 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,303 @@
1+
package auth
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"errors"
7+
"fmt"
8+
"io"
9+
"net/http"
10+
"net/http/httptest"
11+
"os"
12+
"path/filepath"
13+
"sync"
14+
"sync/atomic"
15+
"testing"
16+
"time"
17+
18+
"github.com/cipherstash/stack/languages/golang/internal/guest"
19+
"github.com/tetratelabs/wazero/api"
20+
)
21+
22+
// A profile serves every strategy, and through them every encrypt client,
23+
// in a process: a request cancelled while its token is fetched must fail
24+
// on its own and leave the profile open for the next.
25+
26+
// tokenServer answers the access-key exchange, holding each request until
27+
// the client gives up while hold is set.
28+
func tokenServer(t *testing.T, hold *atomic.Bool) *httptest.Server {
29+
t.Helper()
30+
var jwt string
31+
// The server notices a client that hung up a moment after the client
32+
// has gone: it reads the closed connection on a goroutine of its own.
33+
// So the test's end waits for each held request to see its
34+
// cancellation, for a bounded time, before the server is closed (which
35+
// would end the request's context itself and hide a request the client
36+
// never cancelled).
37+
var held sync.WaitGroup
38+
done := make(chan struct{})
39+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
40+
if hold.Load() {
41+
held.Add(1)
42+
defer held.Done()
43+
// The server notices a client that hung up only once the body
44+
// is read.
45+
_, _ = io.Copy(io.Discard, r.Body)
46+
select {
47+
case <-r.Context().Done():
48+
return
49+
case <-done:
50+
}
51+
select {
52+
case <-r.Context().Done():
53+
case <-time.After(5 * time.Second):
54+
t.Error("a held request was never cancelled")
55+
}
56+
return
57+
}
58+
fmt.Fprintf(w, `{"accessToken":%q,"expiry":%d}`, jwt, time.Now().Add(time.Hour).Unix())
59+
}))
60+
t.Cleanup(server.Close)
61+
t.Cleanup(func() {
62+
close(done)
63+
held.Wait()
64+
})
65+
jwt = testJWT(t, server.URL)
66+
return server
67+
}
68+
69+
func TestATokenFetchCancelledFailsOnlyItself(t *testing.T) {
70+
guestOrSkip(t)
71+
var hold atomic.Bool
72+
server := tokenServer(t, &hold)
73+
profile, err := Open(context.Background(), t.TempDir())
74+
if err != nil {
75+
t.Fatal(err)
76+
}
77+
defer profile.Close()
78+
strategy, err := profile.AccessKey(context.Background(), testCRN, "CSAKtestKeyId.testKeySecret", WithBaseURL(server.URL))
79+
if err != nil {
80+
t.Fatal(err)
81+
}
82+
defer strategy.Close()
83+
84+
hold.Store(true)
85+
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
86+
defer cancel()
87+
if _, err := strategy.Token(ctx); !errors.Is(err, context.DeadlineExceeded) {
88+
t.Fatalf("cancelled Token: %v, want the deadline", err)
89+
}
90+
hold.Store(false)
91+
// The next Token waits for the abandoned one to finish.
92+
if token, err := strategy.Token(context.Background()); err != nil || token == "" {
93+
t.Fatalf("the next Token: %q, %v", token, err)
94+
}
95+
if profile.root.closed || profile.root.inst.module.IsClosed() {
96+
t.Fatal("the profile was closed")
97+
}
98+
}
99+
100+
// An access-key exchange rotates nothing, so it ends with its caller: the
101+
// held request is cancelled at the caller's deadline, not run on for the
102+
// refresh grace, and Close does not wait for it.
103+
func TestAnAccessKeyExchangeEndsWithItsCaller(t *testing.T) {
104+
guestOrSkip(t)
105+
var hold atomic.Bool
106+
server := tokenServer(t, &hold)
107+
profile, err := Open(context.Background(), t.TempDir())
108+
if err != nil {
109+
t.Fatal(err)
110+
}
111+
defer profile.Close()
112+
strategy, err := profile.AccessKey(context.Background(), testCRN, "CSAKtestKeyId.testKeySecret", WithBaseURL(server.URL))
113+
if err != nil {
114+
t.Fatal(err)
115+
}
116+
117+
hold.Store(true)
118+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
119+
defer cancel()
120+
if _, err := strategy.Token(ctx); !errors.Is(err, context.DeadlineExceeded) {
121+
t.Fatalf("cancelled Token: %v, want the deadline", err)
122+
}
123+
// The grace is 10 seconds; a Close that waited for it would take that.
124+
start := time.Now()
125+
if err := strategy.Close(); err != nil {
126+
t.Fatal(err)
127+
}
128+
if waited := time.Since(start); waited > 2*time.Second {
129+
t.Fatalf("Close took %v: the exchange ran on past its caller", waited)
130+
}
131+
}
132+
133+
// A device session's refresh is the request that rotates a credential: the
134+
// service answers with a new refresh token and retires the old one. A
135+
// caller that gives up returns at once, but the refresh runs on, and what
136+
// it brings is saved. The next Token is served from it with no second
137+
// request, and the profile on disk holds the rotated refresh token.
138+
func TestARefreshAnsweredAfterItsCallerGaveUpIsKept(t *testing.T) {
139+
guestOrSkip(t)
140+
dir, _ := expiredDeviceProfile(t)
141+
var requests atomic.Int32
142+
answered := make(chan struct{})
143+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
144+
if requests.Add(1) > 1 {
145+
// The old refresh token is retired once the service answers.
146+
w.WriteHeader(http.StatusBadRequest)
147+
fmt.Fprint(w, `{"error":"invalid_grant"}`)
148+
return
149+
}
150+
_, _ = io.Copy(io.Discard, r.Body)
151+
time.Sleep(150 * time.Millisecond)
152+
fmt.Fprint(w, `{"access_token":"fresh","token_type":"Bearer","expires_in":3600,"refresh_token":"refresh-2"}`)
153+
close(answered)
154+
}))
155+
t.Cleanup(server.Close)
156+
profile, err := Open(context.Background(), dir)
157+
if err != nil {
158+
t.Fatal(err)
159+
}
160+
defer profile.Close()
161+
ws, err := profile.WorkspaceStore(context.Background(), wsA)
162+
if err != nil {
163+
t.Fatal(err)
164+
}
165+
strategy, err := ws.DeviceSession(context.Background(), WithBaseURL(server.URL))
166+
if err != nil {
167+
t.Fatal(err)
168+
}
169+
defer strategy.Close()
170+
171+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
172+
defer cancel()
173+
if _, err := strategy.Token(ctx); !errors.Is(err, context.DeadlineExceeded) {
174+
t.Fatalf("cancelled Token: %v, want the deadline", err)
175+
}
176+
select {
177+
case <-answered:
178+
t.Fatal("the cancelled Token waited for the service's answer")
179+
default:
180+
}
181+
if token, err := strategy.Token(context.Background()); err != nil || token != "fresh" {
182+
t.Fatalf("the next Token: %q, %v, want the refreshed token", token, err)
183+
}
184+
if n := requests.Load(); n != 1 {
185+
t.Fatalf("%d refresh requests, want 1: the late answer was dropped", n)
186+
}
187+
data, err := os.ReadFile(filepath.Join(dir, "workspaces", wsA, "auth.json"))
188+
if err != nil {
189+
t.Fatal(err)
190+
}
191+
var stored map[string]any
192+
if err := json.Unmarshal(data, &stored); err != nil {
193+
t.Fatal(err)
194+
}
195+
if stored["refresh_token"] != "refresh-2" {
196+
t.Fatalf("stored refresh token = %v, want the rotated one", stored["refresh_token"])
197+
}
198+
}
199+
200+
// A refresh runs on after its caller gives up, so the guest call under it
201+
// can succeed after the caller's context ended. It still fails with the
202+
// context's error; what the refresh brought is kept for the next call.
203+
func TestARefreshThatSucceedsAfterItsContextEndsFailsWithIt(t *testing.T) {
204+
guestOrSkip(t)
205+
dir, _ := expiredDeviceProfile(t)
206+
ctx, cancel := context.WithCancel(context.Background())
207+
defer cancel()
208+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
209+
_, _ = io.Copy(io.Discard, r.Body)
210+
cancel()
211+
fmt.Fprint(w, `{"access_token":"fresh","token_type":"Bearer","expires_in":3600,"refresh_token":"refresh-2"}`)
212+
}))
213+
t.Cleanup(server.Close)
214+
profile, err := Open(context.Background(), dir)
215+
if err != nil {
216+
t.Fatal(err)
217+
}
218+
defer profile.Close()
219+
ws, err := profile.WorkspaceStore(context.Background(), wsA)
220+
if err != nil {
221+
t.Fatal(err)
222+
}
223+
strategy, err := ws.DeviceSession(context.Background(), WithBaseURL(server.URL))
224+
if err != nil {
225+
t.Fatal(err)
226+
}
227+
defer strategy.Close()
228+
229+
out, err := strategy.store.callArgs(withRefresh(ctx), func(i *instance) api.Function { return i.authRefresh }, guest.BufArg([]byte(strategy.handle)))
230+
if !errors.Is(err, context.Canceled) || out != nil {
231+
t.Fatalf("refresh: %q, %v, want no token and context.Canceled", out, err)
232+
}
233+
if token, err := strategy.Token(context.Background()); err != nil || token != "fresh" {
234+
t.Fatalf("the next Token: %q, %v, want the refreshed token", token, err)
235+
}
236+
}
237+
238+
func TestATokenRequestWhoseContextEndsInTheQueueNeverReachesTheGuest(t *testing.T) {
239+
guestOrSkip(t)
240+
var hold atomic.Bool
241+
server := tokenServer(t, &hold)
242+
profile, err := Open(context.Background(), t.TempDir())
243+
if err != nil {
244+
t.Fatal(err)
245+
}
246+
defer profile.Close()
247+
strategy, err := profile.AccessKey(context.Background(), testCRN, "CSAKtestKeyId.testKeySecret", WithBaseURL(server.URL))
248+
if err != nil {
249+
t.Fatal(err)
250+
}
251+
defer strategy.Close()
252+
253+
for name, lock := range map[string]interface {
254+
Acquire(context.Context) error
255+
Unlock()
256+
}{"the strategy's": &strategy.mu, "the profile's": &profile.root.mu} {
257+
t.Run(name, func(t *testing.T) {
258+
if err := lock.Acquire(context.Background()); err != nil {
259+
t.Fatal(err)
260+
}
261+
released := make(chan struct{})
262+
time.AfterFunc(200*time.Millisecond, func() { lock.Unlock(); close(released) })
263+
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
264+
defer cancel()
265+
if _, err := strategy.Token(ctx); !errors.Is(err, context.DeadlineExceeded) {
266+
t.Fatalf("queued Token: %v, want the deadline", err)
267+
}
268+
<-released
269+
})
270+
}
271+
if token, err := strategy.Token(context.Background()); err != nil || token == "" {
272+
t.Fatalf("the next Token: %q, %v", token, err)
273+
}
274+
}
275+
276+
// The backstop: a call still running long after its caller gave up closes
277+
// the profile. The caller has returned by then; the next call is ErrState.
278+
func TestATokenFetchThatRunsOnAfterItsContextEndsClosesTheProfile(t *testing.T) {
279+
guestOrSkip(t)
280+
profile, err := Open(context.Background(), t.TempDir(), WithRoundTripper(roundTripFunc(func(*http.Request) (*http.Response, error) {
281+
// A RoundTripper that ignores the request's context.
282+
time.Sleep(300 * time.Millisecond)
283+
return nil, errors.New("too late")
284+
})))
285+
if err != nil {
286+
t.Fatal(err)
287+
}
288+
defer profile.Close()
289+
strategy, err := profile.AccessKey(context.Background(), testCRN, "CSAKtestKeyId.testKeySecret", WithBaseURL("https://cts.example.com"))
290+
if err != nil {
291+
t.Fatal(err)
292+
}
293+
defer strategy.Close()
294+
profile.root.stuck = 20 * time.Millisecond
295+
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
296+
defer cancel()
297+
if _, err := strategy.Token(ctx); !errors.Is(err, context.DeadlineExceeded) {
298+
t.Fatalf("stuck Token: %v, want the deadline", err)
299+
}
300+
if _, err := strategy.Token(context.Background()); !errors.Is(err, ErrState) {
301+
t.Fatalf("a Token after the close: %v, want ErrState", err)
302+
}
303+
}

‎languages/golang/auth/doc.go‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,18 @@
3939
// [OAuth2TokenSource] adapts an existing golang.org/x/oauth2.TokenSource
4040
// into the OIDC provider interface.
4141
//
42+
// A call whose context ends, [Strategy.Token] included, fails with the
43+
// context's error and leaves the store and its strategies open: a request
44+
// cancelled while its token is exchanged or refreshed does not stop the
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.
53+
//
4254
// # Why a second guest
4355
//
4456
// The crypto guest behind encrypt has no filesystem and no

‎languages/golang/auth/guest.go‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,9 @@ func newInstance(ctx context.Context, wasm []byte, hostDir string, policy guest.
111111
return nil, fmt.Errorf("%w: %s: %w", ErrNoProfile, hostDir, err)
112112
}
113113
}
114+
// A call never runs under its caller's context (see guest.Detach), so
115+
// this closes the module only for a call that ran on long after its
116+
// caller gave up: ProfileStore.callArgs.
114117
config := wazero.NewRuntimeConfig().
115118
WithCompilationCache(compilationCache()).
116119
WithCloseOnContextDone(true)
@@ -183,8 +186,8 @@ func newInstance(ctx context.Context, wasm []byte, hostDir string, policy guest.
183186

184187
// release runs the guest's shutdown — every buffer it still holds wiped —
185188
// and closes the runtime, which frees the linear memory through the
186-
// allocator's wipe, then the directory handle the mount holds. A module an
187-
// interrupted call or a trap already closed cannot run sa_shutdown; the
189+
// allocator's wipe, then the directory handle the mount holds. A module a
190+
// stuck call or a trap already closed cannot run sa_shutdown; the
188191
// runtime close still wipes and frees its memory.
189192
func (inst *instance) release() error {
190193
ctx := context.Background()

0 commit comments

Comments
 (0)