c8379c3 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

fix(auth): add better atomic refresh for hyper

c8379c3Kieran Klukas5 months ago

internal/config/atomicwrite.gonew file+38

compare with main

unifiedsplit

@@ -0,0 +1,38 @@

1+package config

3+import (

4+ "os"

5+ "path/filepath"

6+)

8+// atomicWriteFile writes data to a file atomically by writing to a unique

9+// temporary file in the same directory and renaming it into place. This

10+// prevents concurrent readers from observing a partially-written file.

11+func atomicWriteFile(path string, data []byte, perm os.FileMode) error {

12+ path = filepath.Clean(path)

13+ dir := filepath.Dir(path)

14+ f, err := os.CreateTemp(dir, filepath.Base(path)+".*.tmp")

15+ if err != nil {

16+ return err

17+ }

18+ tmp := f.Name()

19+ if _, err := f.Write(data); err != nil {

20+ f.Close()

21+ os.Remove(tmp)

22+ return err

23+ }

24+ if err := f.Chmod(perm); err != nil {

25+ f.Close()

26+ os.Remove(tmp)

27+ return err

28+ }

29+ if err := f.Close(); err != nil {

30+ os.Remove(tmp)

31+ return err

32+ }

33+ if err := os.Rename(tmp, path); err != nil {

34+ os.Remove(tmp)

35+ return err

36+ }

37+ return nil

38+}

internal/config/atomicwrite_test.gonew file+43

@@ -0,0 +1,43 @@

1+package config

3+import (

4+ "os"

5+ "path/filepath"

6+ "runtime"

7+ "testing"

9+ "github.com/stretchr/testify/require"

10+)

11+

12+func TestAtomicWriteFile(t *testing.T) {

13+ t.Parallel()

14+ dir := t.TempDir()

15+ path := filepath.Join(dir, "test.json")

16+

17+ require.NoError(t, atomicWriteFile(path, []byte(`{"key":"value"}`), 0o600))

18+

19+ data, err := os.ReadFile(path)

20+ require.NoError(t, err)

21+ require.Equal(t, `{"key":"value"}`, string(data))

22+

23+ // No temp files should linger.

24+ entries, err := os.ReadDir(dir)

25+ require.NoError(t, err)

26+ require.Len(t, entries, 1)

27+ require.Equal(t, "test.json", entries[0].Name())

28+}

29+

30+func TestAtomicWriteFile_PermissionsApplied(t *testing.T) {

31+ if runtime.GOOS == "windows" {

32+ t.Skip("Windows does not support Unix file permissions")

33+ }

34+ t.Parallel()

35+ dir := t.TempDir()

36+ path := filepath.Join(dir, "test.json")

37+

38+ require.NoError(t, atomicWriteFile(path, []byte(`{}`), 0o600))

39+

40+ info, err := os.Stat(path)

41+ require.NoError(t, err)

42+ require.Equal(t, os.FileMode(0o600), info.Mode().Perm())

43+}

internal/config/store.go+29−11

@@ -158,7 +158,7 @@ func (s *ConfigStore) SetConfigFields(scope Scope, kv map[string]any) error {

158 if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {

159 return fmt.Errorf("failed to create config directory %q: %w", path, err)

160 }

161− if err := os.WriteFile(path, []byte(newValue), 0o600); err != nil {

161+ if err := atomicWriteFile(path, []byte(newValue), 0o600); err != nil {

162 return fmt.Errorf("failed to write config file: %w", err)

163 }

164

@@ -193,7 +193,7 @@ func (s *ConfigStore) RemoveConfigField(scope Scope, key string) error {

193 if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {

194 return fmt.Errorf("failed to create config directory %q: %w", path, err)

195 }

196− if err := os.WriteFile(path, []byte(newValue), 0o600); err != nil {

196+ if err := atomicWriteFile(path, []byte(newValue), 0o600); err != nil {

197 return fmt.Errorf("failed to write config file: %w", err)

198 }

199

@@ -302,7 +302,10 @@ func (s *ConfigStore) SetProviderAPIKey(scope Scope, providerID string, apiKey a

302// RefreshOAuthToken refreshes the OAuth token for the given provider.

303// Before making an external refresh request, it checks the config file on

304// disk to see if another Crush session has already refreshed the token. If

305−// a newer token is found, it is used instead of refreshing.

305+// a newer, unexpired token is found, it is used instead of refreshing. If

306+// the exchange fails (e.g. because another session already rotated the

307+// refresh token), the disk is re-checked to recover the other session's

308+// token.

306func (s *ConfigStore) RefreshOAuthToken(ctx context.Context, scope Scope, providerID string) error {

309func (s *ConfigStore) RefreshOAuthToken(ctx context.Context, scope Scope, providerID string) error {

307 providerConfig, exists := s.config.Providers.Get(providerID)

310 providerConfig, exists := s.config.Providers.Get(providerID)

308 if !exists {

311 if !exists {

@@ -318,15 +321,9 @@ func (s *ConfigStore) RefreshOAuthToken(ctx context.Context, scope Scope, provid

318 newToken, err := s.loadTokenFromDisk(scope, providerID)

321 newToken, err := s.loadTokenFromDisk(scope, providerID)

319 if err != nil {

322 if err != nil {

320 slog.Warn("Failed to read token from config file, proceeding with refresh", "provider", providerID, "error", err)

323 slog.Warn("Failed to read token from config file, proceeding with refresh", "provider", providerID, "error", err)

321− } else if newToken != nil && newToken.AccessToken != providerConfig.OAuthToken.AccessToken {

324+ } else if newToken != nil && !newToken.IsExpired() && newToken.AccessToken != providerConfig.OAuthToken.AccessToken {

322 slog.Info("Using token refreshed by another session", "provider", providerID)

325 slog.Info("Using token refreshed by another session", "provider", providerID)

323− providerConfig.OAuthToken = newToken

326+ return s.applyToken(providerConfig, newToken, providerID)

324− providerConfig.APIKey = newToken.AccessToken

325− if providerID == string(catwalk.InferenceProviderCopilot) {

326− providerConfig.SetupGitHubCopilot()

327− }

328− s.config.Providers.Set(providerID, providerConfig)

329− return nil

330 }

327 }

331

328

332 var refreshedToken *oauth.Token

329 var refreshedToken *oauth.Token

@@ -340,6 +337,16 @@ func (s *ConfigStore) RefreshOAuthToken(ctx context.Context, scope Scope, provid

340 return fmt.Errorf("OAuth refresh not supported for provider %s", providerID)

337 return fmt.Errorf("OAuth refresh not supported for provider %s", providerID)

341 }

338 }

342 if refreshErr != nil {

339 if refreshErr != nil {

340+ // The exchange may have failed because another session already

341+ // rotated the refresh token. Re-read the config file and use the

342+ // other session's token if available.

343+ if diskToken, diskErr := s.loadTokenFromDisk(scope, providerID); diskErr == nil &&

344+ diskToken != nil &&

345+ !diskToken.IsExpired() &&

346+ diskToken.AccessToken != providerConfig.OAuthToken.AccessToken {

347+ slog.Info("Using token refreshed by another session after exchange failure", "provider", providerID)

348+ return s.applyToken(providerConfig, diskToken, providerID)

349+ }

343 return fmt.Errorf("failed to refresh OAuth token for provider %s: %w", providerID, refreshErr)

350 return fmt.Errorf("failed to refresh OAuth token for provider %s: %w", providerID, refreshErr)

344 }

351 }

345

352

@@ -364,6 +371,17 @@ func (s *ConfigStore) RefreshOAuthToken(ctx context.Context, scope Scope, provid

364 return nil

371 return nil

365}

372}

366

373

374+// applyToken updates the in-memory provider config with the given token.

375+func (s *ConfigStore) applyToken(providerConfig ProviderConfig, token *oauth.Token, providerID string) error {

376+ providerConfig.OAuthToken = token

377+ providerConfig.APIKey = token.AccessToken

378+ if providerID == string(catwalk.InferenceProviderCopilot) {

379+ providerConfig.SetupGitHubCopilot()

380+ }

381+ s.config.Providers.Set(providerID, providerConfig)

382+ return nil

383+}

384+

367// loadTokenFromDisk reads the OAuth token for the given provider from the

385// loadTokenFromDisk reads the OAuth token for the given provider from the

368// config file on disk. Returns nil if the token is not found or matches the

386// config file on disk. Returns nil if the token is not found or matches the

369// current in-memory token.

387// current in-memory token.