Skip to content

Commit 19eed21

Browse files
johnstcnjohnstcn
committed
retry fetching agent ssh key
1 parent b1ba280 commit 19eed21

File tree

3 files changed

+90
-21
lines changed

3 files changed

+90
-21
lines changed

‎envbuilder.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -194,7 +194,7 @@ func Run(ctx context.Context, options Options) error {
194194
CABundle: caBundle,
195195
}
196196

197-
cloneOpts.RepoAuth = SetupRepoAuth(&options)
197+
cloneOpts.RepoAuth = SetupRepoAuth(ctx, &options)
198198
if options.GitHTTPProxyURL != "" {
199199
cloneOpts.ProxyOptions = transport.ProxyOptions{
200200
URL: options.GitHTTPProxyURL,

‎git.go

+39-4
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,13 @@ import (
77
"fmt"
88
"io"
99
"net"
10+
"net/http"
1011
"net/url"
1112
"os"
1213
"strings"
1314
"time"
1415

16+
"github.com/cenkalti/backoff/v4"
1517
"github.com/coder/coder/v2/codersdk"
1618
"github.com/coder/coder/v2/codersdk/agentsdk"
1719
"github.com/go-git/go-billy/v5"
@@ -186,7 +188,7 @@ func LogHostKeyCallback(log LoggerFunc) gossh.HostKeyCallback {
186188
// If SSH_KNOWN_HOSTS is not set, the SSH auth method will be configured
187189
// to accept and log all host keys. Otherwise, host key checking will be
188190
// performed as usual.
189-
func SetupRepoAuth(options *Options) transport.AuthMethod {
191+
func SetupRepoAuth(ctx context.Context, options *Options) transport.AuthMethod {
190192
if options.GitURL == "" {
191193
options.Logger(codersdk.LogLevelInfo, "#1: ❔ No Git URL supplied!")
192194
return nil
@@ -231,14 +233,14 @@ func SetupRepoAuth(options *Options) transport.AuthMethod {
231233
// an SSH key from Coder!
232234
if signer == nil && options.CoderAgentURL != "" && options.CoderAgentToken != "" {
233235
options.Logger(codersdk.LogLevelInfo, "#1: 🔑 Fetching key from %s!", options.CoderAgentURL)
234-
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
236+
fetchCtx, cancel := context.WithCancel(ctx)
235237
defer cancel()
236-
s, err := FetchCoderSSHKey(ctx, options.CoderAgentURL, options.CoderAgentToken)
238+
s, err := FetchCoderSSHKeyRetry(fetchCtx, options.Logger, options.CoderAgentURL, options.CoderAgentToken)
237239
if err == nil {
238240
signer = s
239241
options.Logger(codersdk.LogLevelInfo, "#1: 🔑 Fetched %s key %s !", signer.PublicKey().Type(), keyFingerprint(signer)[:8])
240242
} else {
241-
options.Logger(codersdk.LogLevelInfo, "#1: ❌ Failed to fetch SSH key: %w", options.CoderAgentURL, err)
243+
options.Logger(codersdk.LogLevelInfo, "#1: ❌ Failed to fetch SSH key from %s: %w", options.CoderAgentURL, err)
242244
}
243245
}
244246

@@ -276,6 +278,39 @@ func SetupRepoAuth(options *Options) transport.AuthMethod {
276278
return auth
277279
}
278280

281+
// FetchCoderSSHKeyRetry wraps FetchCoderSSHKey in backoff.Retry.
282+
// Retries are attempted if Coder responds with a 401 Unauthorized.
283+
// This indicates that the workspace build has not yet completed.
284+
// It will retry for up to 1 minute with exponential backoff.
285+
// Any other error is considered a permanent failure.
286+
func FetchCoderSSHKeyRetry(ctx context.Context, log LoggerFunc, coderURL, agentToken string) (gossh.Signer, error) {
287+
ctx, cancel := context.WithCancel(ctx)
288+
defer cancel()
289+
290+
signerChan := make(chan gossh.Signer, 1)
291+
eb := backoff.NewExponentialBackOff()
292+
eb.MaxElapsedTime = 0
293+
eb.MaxInterval = time.Minute
294+
bkoff := backoff.WithContext(eb, ctx)
295+
err := backoff.Retry(func() error {
296+
s, err := FetchCoderSSHKey(ctx, coderURL, agentToken)
297+
if err != nil {
298+
var sdkErr *codersdk.Error
299+
if errors.As(err, &sdkErr) && sdkErr.StatusCode() == http.StatusUnauthorized {
300+
// Retry, as this may just mean that the workspace build has not yet
301+
// completed.
302+
log(codersdk.LogLevelInfo, "#1: 🕐 Backing off as the workspace build has not yet completed...")
303+
return err
304+
}
305+
close(signerChan)
306+
return backoff.Permanent(err)
307+
}
308+
signerChan <- s
309+
return nil
310+
}, bkoff)
311+
return <-signerChan, err
312+
}
313+
279314
// FetchCoderSSHKey fetches the user's Git SSH key from Coder using the supplied
280315
// Coder URL and agent token.
281316
func FetchCoderSSHKey(ctx context.Context, coderURL string, agentToken string) (gossh.Signer, error) {

‎git_test.go

+50-16
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
"os"
1313
"path/filepath"
1414
"regexp"
15+
"sync/atomic"
1516
"testing"
1617

1718
"github.com/coder/coder/v2/codersdk"
@@ -268,11 +269,12 @@ func TestCloneRepoSSH(t *testing.T) {
268269
// nolint:paralleltest // t.Setenv for SSH_AUTH_SOCK
269270
func TestSetupRepoAuth(t *testing.T) {
270271
t.Setenv("SSH_AUTH_SOCK", "")
272+
ctx := context.Background()
271273
t.Run("Empty", func(t *testing.T) {
272274
opts := &envbuilder.Options{
273275
Logger: testLog(t),
274276
}
275-
auth := envbuilder.SetupRepoAuth(opts)
277+
auth := envbuilder.SetupRepoAuth(ctx, opts)
276278
require.Nil(t, auth)
277279
})
278280

@@ -281,7 +283,7 @@ func TestSetupRepoAuth(t *testing.T) {
281283
GitURL: "http://host.tld/repo",
282284
Logger: testLog(t),
283285
}
284-
auth := envbuilder.SetupRepoAuth(opts)
286+
auth := envbuilder.SetupRepoAuth(ctx, opts)
285287
require.Nil(t, auth)
286288
})
287289

@@ -292,7 +294,7 @@ func TestSetupRepoAuth(t *testing.T) {
292294
GitPassword: "pass",
293295
Logger: testLog(t),
294296
}
295-
auth := envbuilder.SetupRepoAuth(opts)
297+
auth := envbuilder.SetupRepoAuth(ctx, opts)
296298
ba, ok := auth.(*githttp.BasicAuth)
297299
require.True(t, ok)
298300
require.Equal(t, opts.GitUsername, ba.Username)
@@ -306,7 +308,7 @@ func TestSetupRepoAuth(t *testing.T) {
306308
GitPassword: "pass",
307309
Logger: testLog(t),
308310
}
309-
auth := envbuilder.SetupRepoAuth(opts)
311+
auth := envbuilder.SetupRepoAuth(ctx, opts)
310312
ba, ok := auth.(*githttp.BasicAuth)
311313
require.True(t, ok)
312314
require.Equal(t, opts.GitUsername, ba.Username)
@@ -320,7 +322,7 @@ func TestSetupRepoAuth(t *testing.T) {
320322
GitSSHPrivateKeyPath: kPath,
321323
Logger: testLog(t),
322324
}
323-
auth := envbuilder.SetupRepoAuth(opts)
325+
auth := envbuilder.SetupRepoAuth(ctx, opts)
324326
_, ok := auth.(*gitssh.PublicKeys)
325327
require.True(t, ok)
326328
})
@@ -332,7 +334,7 @@ func TestSetupRepoAuth(t *testing.T) {
332334
GitSSHPrivateKeyPath: kPath,
333335
Logger: testLog(t),
334336
}
335-
auth := envbuilder.SetupRepoAuth(opts)
337+
auth := envbuilder.SetupRepoAuth(ctx, opts)
336338
_, ok := auth.(*gitssh.PublicKeys)
337339
require.True(t, ok)
338340
})
@@ -345,7 +347,7 @@ func TestSetupRepoAuth(t *testing.T) {
345347
GitSSHPrivateKeyPath: kPath,
346348
Logger: testLog(t),
347349
}
348-
auth := envbuilder.SetupRepoAuth(opts)
350+
auth := envbuilder.SetupRepoAuth(ctx, opts)
349351
_, ok := auth.(*gitssh.PublicKeys)
350352
require.True(t, ok)
351353
})
@@ -358,7 +360,7 @@ func TestSetupRepoAuth(t *testing.T) {
358360
GitUsername: "user",
359361
Logger: testLog(t),
360362
}
361-
auth := envbuilder.SetupRepoAuth(opts)
363+
auth := envbuilder.SetupRepoAuth(ctx, opts)
362364
_, ok := auth.(*gitssh.PublicKeys)
363365
require.True(t, ok)
364366
})
@@ -370,7 +372,7 @@ func TestSetupRepoAuth(t *testing.T) {
370372
GitSSHPrivateKeyPath: kPath,
371373
Logger: testLog(t),
372374
}
373-
auth := envbuilder.SetupRepoAuth(opts)
375+
auth := envbuilder.SetupRepoAuth(ctx, opts)
374376
pk, ok := auth.(*gitssh.PublicKeys)
375377
require.True(t, ok)
376378
require.NotNil(t, pk.Signer)
@@ -384,7 +386,7 @@ func TestSetupRepoAuth(t *testing.T) {
384386
GitURL: "ssh://[email protected]:repo/path",
385387
Logger: testLog(t),
386388
}
387-
auth := envbuilder.SetupRepoAuth(opts)
389+
auth := envbuilder.SetupRepoAuth(ctx, opts)
388390
require.Nil(t, auth) // TODO: actually test SSH_AUTH_SOCK
389391
})
390392

@@ -415,19 +417,51 @@ func TestSetupRepoAuth(t *testing.T) {
415417
GitURL: "ssh://[email protected]:repo/path",
416418
Logger: testLog(t),
417419
}
418-
auth := envbuilder.SetupRepoAuth(opts)
420+
auth := envbuilder.SetupRepoAuth(ctx, opts)
419421
pk, ok := auth.(*gitssh.PublicKeys)
420422
require.True(t, ok)
421423
require.NotNil(t, pk.Signer)
422424
require.Equal(t, actualSigner, pk.Signer)
423425
})
424426

425-
t.Run("SSH/CoderForbidden", func(t *testing.T) {
427+
t.Run("SSH/CoderRetry", func(t *testing.T) {
426428
token := uuid.NewString()
429+
actualSigner, err := gossh.ParsePrivateKey([]byte(testKey))
430+
require.NoError(t, err)
431+
var count atomic.Int64
432+
// Return 401 initially, but eventually 200.
427433
handler := func(w http.ResponseWriter, r *http.Request) {
428-
hdr := r.Header.Get(codersdk.SessionTokenHeader)
429-
assert.Equal(t, hdr, token)
430-
w.WriteHeader(http.StatusForbidden)
434+
c := count.Add(1)
435+
if c < 3 {
436+
hdr := r.Header.Get(codersdk.SessionTokenHeader)
437+
assert.Equal(t, hdr, token)
438+
w.WriteHeader(http.StatusUnauthorized)
439+
return
440+
}
441+
_ = json.NewEncoder(w).Encode(&agentsdk.GitSSHKey{
442+
PublicKey: string(actualSigner.PublicKey().Marshal()),
443+
PrivateKey: string(testKey),
444+
})
445+
}
446+
srv := httptest.NewServer(http.HandlerFunc(handler))
447+
opts := &envbuilder.Options{
448+
CoderAgentURL: srv.URL,
449+
CoderAgentToken: token,
450+
GitURL: "ssh://[email protected]:repo/path",
451+
Logger: testLog(t),
452+
}
453+
auth := envbuilder.SetupRepoAuth(ctx, opts)
454+
pk, ok := auth.(*gitssh.PublicKeys)
455+
require.True(t, ok)
456+
require.NotNil(t, pk.Signer)
457+
require.Equal(t, actualSigner, pk.Signer)
458+
})
459+
460+
t.Run("SSH/NotCoder", func(t *testing.T) {
461+
token := uuid.NewString()
462+
handler := func(w http.ResponseWriter, r *http.Request) {
463+
w.WriteHeader(http.StatusTeapot)
464+
_, _ = w.Write([]byte("I'm a teapot!"))
431465
}
432466
srv := httptest.NewServer(http.HandlerFunc(handler))
433467
opts := &envbuilder.Options{
@@ -436,7 +470,7 @@ func TestSetupRepoAuth(t *testing.T) {
436470
GitURL: "ssh://[email protected]:repo/path",
437471
Logger: testLog(t),
438472
}
439-
auth := envbuilder.SetupRepoAuth(opts)
473+
auth := envbuilder.SetupRepoAuth(ctx, opts)
440474
require.Nil(t, auth)
441475
})
442476
}

0 commit comments

Comments
 (0)