Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
196 changes: 102 additions & 94 deletions internal/cmd/initcmd/initcmd.go
Original file line number Diff line number Diff line change
Expand Up @@ -5386,6 +5386,19 @@ func initCredentialEntryDeferredByDecision(decisions map[initCredentialDecisionK

type initReviewerCredentialCleanupGroup = credentials.ReviewerCredentialCleanupGroup

type initApplicationPlan struct {
path string
cfg config.File
credentialPlan []initCredentialPlanEntry
writes map[string]map[string]string
overwriteRefs map[string]bool
backendFlagSet bool
overwrite bool
primaryStore initStore
primaryResolved credentials.ResolvedSecretsStore
profileName string
}

func openInitStoreForEntry(deps initDeps, opts *root.Options, flagSet bool, cfg config.File, entry initCredentialPlanEntry) (initStore, error) {
backend := ""
if opts != nil {
Expand Down Expand Up @@ -6129,54 +6142,14 @@ func initCredentialOptionalKeySummary(entry initCredentialPlanEntry) string {
}

func applyInteractiveInitSessionPlan(opts *root.Options, deps initDeps, plan initSessionPlan) error {
if err := config.Validate(plan.cfg); err != nil {
if errors.Is(err, config.ErrInvalid) || errors.Is(err, config.ErrProfileNotFound) {
return cmderr.Config(err)
}
return err
}
touchedCredentialRefs := map[string]struct{}{}
if len(plan.writes) > 0 {
groups, err := credentials.GroupWritesByStore(plan.credentialPlan, plan.writes, plan.overwriteRefs)
if err != nil {
return cmderr.Config(err)
}
for _, group := range groups {
store, err := openInitStoreForEntry(deps, opts, plan.backendFlagSet, plan.cfg, initCredentialPlanEntry{SecretsStore: group.Resolved})
if err != nil {
if config.IsConfigSelection(err) {
return cmderr.Config(err)
}
return cmderr.Credential(err)
}
if err := preflightNoOverwrite(store, group.Writes, group.OverwriteRefs); err != nil {
_ = store.Close()
return cmderr.Credential(err)
}
if _, err := writeBundles(store, group.Writes, false, group.OverwriteRefs); err != nil {
_ = store.Close()
return cmderr.Credential(err)
}
for ref := range group.Writes {
touchedCredentialRefs[ref] = struct{}{}
}
_ = store.Close()
}
}
cleanupGroups, err := credentials.GroupStaleReviewerCredentialCleanupsByStore(plan.cfg, plan.credentialPlan)
if err != nil {
return cmderr.Config(err)
}
if err := deps.saveConfig(plan.path, plan.cfg); err != nil {
if errors.Is(err, config.ErrInvalid) || errors.Is(err, config.ErrProfileNotFound) {
return cmderr.Config(err)
}
if len(touchedCredentialRefs) > 0 {
return fmt.Errorf("init updated credentials but failed to save config; credential names needing cleanup: %v: %w", sortedCredentialRefSet(touchedCredentialRefs), err)
}
return err
}
if err := applyStaleReviewerCredentialCleanups(opts, deps, plan.backendFlagSet, plan.cfg, cleanupGroups); err != nil {
if err := applyInitApplicationPlan(opts, deps, initApplicationPlan{
path: plan.path,
cfg: plan.cfg,
credentialPlan: plan.credentialPlan,
writes: plan.writes,
overwriteRefs: plan.overwriteRefs,
backendFlagSet: plan.backendFlagSet,
}); err != nil {
return err
}
if err := writeInteractiveInitSessionSummary(opts.Stdout, plan); err != nil {
Expand Down Expand Up @@ -6330,13 +6303,79 @@ func applyStaleReviewerCredentialCleanups(opts *root.Options, deps initDeps, bac
return nil
}

func applyInitPlan(opts *root.Options, flags initOptions, deps initDeps, plan initPlan) error {
var store initStore
var primaryResolved credentials.ResolvedSecretsStore
func applyInitApplicationPlan(opts *root.Options, deps initDeps, plan initApplicationPlan) error {
if err := config.Validate(plan.cfg); err != nil {
if errors.Is(err, config.ErrInvalid) || errors.Is(err, config.ErrProfileNotFound) {
return cmderr.Config(err)
}
return err
}
var writeGroups []credentials.WriteGroup
if len(plan.writes) > 0 {
var err error
writeGroups, err = credentials.GroupWritesByStore(plan.credentialPlan, plan.writes, plan.overwriteRefs)
if err != nil {
return cmderr.Config(err)
}
}
cleanupGroups, err := credentials.GroupStaleReviewerCredentialCleanupsByStore(plan.cfg, plan.credentialPlan)
if err != nil {
return cmderr.Config(err)
}
touchedCredentialRefs := map[string]struct{}{}
for _, group := range writeGroups {
store := plan.primaryStore
closeStore := false
if store == nil || !sameInitCredentialStore(group.Resolved, plan.primaryResolved) {
store, err = openInitStoreForEntry(deps, opts, plan.backendFlagSet, plan.cfg, initCredentialPlanEntry{SecretsStore: group.Resolved})
if err != nil {
if config.IsConfigSelection(err) {
return cmderr.Config(err)
}
return cmderr.Credential(err)
}
closeStore = true
}
if !plan.overwrite {
if err := preflightNoOverwrite(store, group.Writes, group.OverwriteRefs); err != nil {
if closeStore {
_ = store.Close()
}
return cmderr.Credential(err)
}
}
if _, err := writeBundles(store, group.Writes, plan.overwrite, group.OverwriteRefs); err != nil {
if closeStore {
_ = store.Close()
}
return cmderr.Credential(err)
}
if closeStore {
_ = store.Close()
}
for ref := range group.Writes {
touchedCredentialRefs[ref] = struct{}{}
}
}
if err := deps.saveConfig(plan.path, plan.cfg); err != nil {
if errors.Is(err, config.ErrInvalid) || errors.Is(err, config.ErrProfileNotFound) {
return cmderr.Config(err)
}
if len(touchedCredentialRefs) > 0 {
refs := sortedCredentialRefSet(touchedCredentialRefs)
if plan.profileName != "" {
return fmt.Errorf("init updated credentials but failed to save config for profile %q; credential names needing cleanup: %v: %w", plan.profileName, refs, err)
}
return fmt.Errorf("init updated credentials but failed to save config; credential names needing cleanup: %v: %w", refs, err)
}
return err
}
return applyStaleReviewerCredentialCleanups(opts, deps, plan.backendFlagSet, plan.cfg, cleanupGroups)
}

func applyInitPlan(opts *root.Options, flags initOptions, deps initDeps, plan initPlan) error {
var store initStore
var primaryResolved credentials.ResolvedSecretsStore
needsPrimaryStore := len(plan.writes) > 0 || (plan.profile.LLM.Auth == config.LLMAuthAPIKey && !plan.allowDeferredLLM)
openPrimaryStore := func() (initStore, credentials.ResolvedSecretsStore, error) {
if store != nil {
Expand Down Expand Up @@ -6391,49 +6430,18 @@ func applyInitPlan(opts *root.Options, flags initOptions, deps initDeps, plan in
return exitcode.Usage(fmt.Errorf("api_key LLM auth requires --llm-api-key-stdin or --llm-api-key-from-env"))
}
}
touchedCredentialRefs := map[string]struct{}{}
if len(plan.writes) > 0 {
groups, err := credentials.GroupWritesByStore(plan.credentialPlan, plan.writes, plan.overwriteRefs)
if err != nil {
return cmderr.Config(err)
}
for _, group := range groups {
groupStore := store
if groupStore == nil || !sameInitCredentialStore(group.Resolved, primaryResolved) {
groupStore, err = openInitStoreForEntry(deps, opts, plan.backendFlagSet, plan.cfg, initCredentialPlanEntry{SecretsStore: group.Resolved})
if err != nil {
if config.IsConfigSelection(err) {
return cmderr.Config(err)
}
return cmderr.Credential(err)
}
if groupStore != store {
defer groupStore.Close()
}
}
if !flags.overwrite {
if err := preflightNoOverwrite(groupStore, group.Writes, group.OverwriteRefs); err != nil {
return cmderr.Credential(err)
}
}
if _, err := writeBundles(groupStore, group.Writes, flags.overwrite, group.OverwriteRefs); err != nil {
return cmderr.Credential(err)
}
for ref := range group.Writes {
touchedCredentialRefs[ref] = struct{}{}
}
}
}
if err := deps.saveConfig(plan.path, plan.cfg); err != nil {
if errors.Is(err, config.ErrInvalid) || errors.Is(err, config.ErrProfileNotFound) {
return cmderr.Config(err)
}
if len(touchedCredentialRefs) > 0 {
return fmt.Errorf("init updated credentials but failed to save config for profile %q; credential names needing cleanup: %v: %w", plan.profileName, sortedCredentialRefSet(touchedCredentialRefs), err)
}
return err
}
if err := applyStaleReviewerCredentialCleanups(opts, deps, plan.backendFlagSet, plan.cfg, cleanupGroups); err != nil {
if err := applyInitApplicationPlan(opts, deps, initApplicationPlan{
path: plan.path,
cfg: plan.cfg,
credentialPlan: plan.credentialPlan,
writes: plan.writes,
overwriteRefs: plan.overwriteRefs,
backendFlagSet: plan.backendFlagSet,
overwrite: flags.overwrite,
primaryStore: store,
primaryResolved: primaryResolved,
profileName: plan.profileName,
}); err != nil {
return err
}
if _, err := fmt.Fprintf(opts.Stdout, "Initialized profile %s\n", plan.profileName); err != nil {
Expand Down
80 changes: 80 additions & 0 deletions internal/cmd/initcmd/initcmd_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -17908,6 +17908,10 @@ func TestApplyInteractiveInitSessionPlanSaveFailureDoesNotDeleteStaleReviewerKey
credentials.GitHubAppPrivateKeyKey: "private-key",
},
})
store.deleteFunc = func(string, string) error {
t.Fatal("Delete called despite config save failure")
return nil
}
profile := basicProfile("work")
profile.ReviewerCredentials = &config.ReviewerCredentials{
AuthMode: config.GitAuthModeGitHubApp,
Expand Down Expand Up @@ -18083,6 +18087,78 @@ func TestApplyInitPlanDeletesStaleReviewerKeyWithoutWrites(t *testing.T) {
}
}

func TestApplyInitPlanCleanupFailureLeavesSavedConfigAndReportsPartialCleanup(t *testing.T) {
store := newFakeInitStore(map[string]map[string]string{
"work-reviewer": {
credentials.GitTokenKey: "old-reviewer-pat",
credentials.GitHubAppIDKey: "12345",
credentials.GitHubAppPrivateKeyKey: "private-key",
},
})
store.deleteFunc = func(profile, key string) error {
if key == credentials.GitHubAppIDKey {
return errors.New("cleanup failed")
}
delete(store.bundles[profile], key)
return nil
}
profile := basicProfile("work")
profile.ReviewerCredentials = &config.ReviewerCredentials{
AuthMode: config.GitAuthModeGitHubApp,
GitHubApp: &config.GitHubAppConfig{AppID: "12345"},
Credential: config.CredentialLocation{Store: config.LocalOSCredentialStoreID, Name: "codereview/work-reviewer"},
}
cfg := config.File{Profiles: map[string]config.Profile{"work": profile}}
resolved, err := credentials.ResolveSecretsStoreForProfile(cfg, profile)
if err != nil {
t.Fatalf("ResolveSecretsStoreForProfile: %v", err)
}
plan := initPlan{
path: filepath.Join(t.TempDir(), "config.yml"),
cfg: cfg,
profileName: "work",
profile: profile,
credentialPlan: []initCredentialPlanEntry{{
Ref: config.CredentialRef{
Purpose: "reviewer_credentials",
Ref: "codereview/work-reviewer",
Mode: string(config.GitAuthModeGitHubApp),
},
PreviousRef: &config.CredentialRef{
Purpose: "reviewer_credentials",
Ref: "codereview/work-reviewer",
Mode: string(config.GitAuthModePAT),
},
SecretsStore: resolved,
KeySpecs: []credentials.KeySpec{
{Key: credentials.GitHubAppPrivateKeyKey, Required: true},
},
State: initCredentialPlanStateKeepExisting,
}},
}
var saved bool
err = applyInitPlan(&root.Options{Stdout: &bytes.Buffer{}, Stderr: &bytes.Buffer{}}, initOptions{}, initDeps{
openStore: func(string, bool, config.File) (initStore, error) { return store, nil },
saveConfig: func(string, config.File) error {
saved = true
return nil
},
}, plan)
wantErr := "init saved config before failing to delete stale reviewer keys on ; stale credential names needing manual cleanup: [codereview/work-reviewer]: cleanup failed"
if err == nil || err.Error() != wantErr {
t.Fatalf("applyInitPlan error = %v, want %q", err, wantErr)
}
if !saved {
t.Fatal("config was not saved before cleanup failure")
}
if ok, err := store.Exists("work-reviewer", credentials.GitTokenKey); err != nil || ok {
t.Fatalf("git_token exists = %t, err = %v; want partial cleanup retained", ok, err)
}
if ok, err := store.Exists("work-reviewer", credentials.GitHubAppIDKey); err != nil || !ok {
t.Fatalf("github_app_id exists = %t, err = %v; want failed key retained", ok, err)
}
}

func TestStaleReviewerCleanupKeepsKeysRequiredByAnotherActiveModeWithoutWrites(t *testing.T) {
store := newFakeInitStore(map[string]map[string]string{
"shared-reviewer": {
Expand Down Expand Up @@ -18207,6 +18283,7 @@ type fakeInitStore struct {
existsFunc func(string, string) (bool, error)
listBundleFunc func(string) ([]string, error)
setBundleFunc func(string, map[string]string, ...credstore.SetOpt) (credstore.Result, error)
deleteFunc func(string, string) error
}

func newFakeInitStore(bundles map[string]map[string]string) *fakeInitStore {
Expand Down Expand Up @@ -18265,6 +18342,9 @@ func (s *fakeInitStore) SetBundle(profile string, kv map[string]string, opts ...
}

func (s *fakeInitStore) Delete(profile, key string) error {
if s.deleteFunc != nil {
return s.deleteFunc(profile, key)
}
if _, ok := s.bundles[profile][key]; !ok {
return credstore.ErrNotFound
}
Expand Down
Loading