From 76db34a2f919b92699f1e4c0ca6706fedb6240f7 Mon Sep 17 00:00:00 2001 From: Pratik Patel Date: Wed, 19 Aug 2026 16:54:02 -0700 Subject: [PATCH 1/6] simplify register --- pkg/cmd/register/register.go | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/pkg/cmd/register/register.go b/pkg/cmd/register/register.go index 2ad88b43..41692d55 100644 --- a/pkg/cmd/register/register.go +++ b/pkg/cmd/register/register.go @@ -235,12 +235,7 @@ func runRegister(ctx context.Context, t *terminal.Terminal, s RegisterStore, opt // Determine if SSH access should be enabled enableSSH := false sshPortForGrant := int32(0) - if opts.interactive { - enableSSH = deps.prompter.ConfirmYesNo("Would you like to enable SSH access to this device?") - if enableSSH { - sshPortForGrant = 0 // prompt for port - } - } else if opts.sshPort != 0 { + if opts.sshPort != 0 { enableSSH = true sshPortForGrant = opts.sshPort } From 9175303c4fc47a2960e4b8282e55a124bb2067e8 Mon Sep 17 00:00:00 2001 From: Pratik Patel Date: Fri, 21 Aug 2026 09:16:51 -0700 Subject: [PATCH 2/6] flags --- pkg/cmd/cmd.go | 12 +- pkg/cmd/deregister/deregister_test.go | 7 + pkg/cmd/grantssh/grantssh_test.go | 7 + pkg/cmd/register/device_registration_store.go | 53 ++ .../device_registration_store_test.go | 66 ++ pkg/cmd/register/register.go | 332 +++++--- pkg/cmd/register/register_test.go | 716 +++++++++++------- pkg/cmd/revokessh/revokessh_test.go | 7 + 8 files changed, 828 insertions(+), 372 deletions(-) diff --git a/pkg/cmd/cmd.go b/pkg/cmd/cmd.go index 831a2aa9..00461a6d 100644 --- a/pkg/cmd/cmd.go +++ b/pkg/cmd/cmd.go @@ -265,12 +265,18 @@ func NewBrevCommand() *cobra.Command { //nolint:funlen,gocognit,gocyclo // defin fmt.Printf("%v\n", err) } - createCmdTree(cmds, t, loginCmdStore, noLoginCmdStore, loginAuth, externalNodeCmdStore) + // accessKeyAuth seeds the in-memory session store (used by register's API + // calls) and durably persists a supplied --access-key via the file-backed + // loginAuth (the same path as `brev login --api-key`), so other brev + // commands authenticate with the key afterwards. + accessKeyAuth := register.NewAccessKeyAuthenticator(memAuthStore.MemoryAuthStore, loginAuth) + + createCmdTree(cmds, t, loginCmdStore, noLoginCmdStore, loginAuth, externalNodeCmdStore, accessKeyAuth) return cmds } -func createCmdTree(cmd *cobra.Command, t *terminal.Terminal, loginCmdStore *store.AuthHTTPStore, noLoginCmdStore *store.AuthHTTPStore, loginAuth *auth.LoginAuth, externalNodeCmdStore *store.AuthHTTPStore) { //nolint:funlen // define brev command +func createCmdTree(cmd *cobra.Command, t *terminal.Terminal, loginCmdStore *store.AuthHTTPStore, noLoginCmdStore *store.AuthHTTPStore, loginAuth *auth.LoginAuth, externalNodeCmdStore *store.AuthHTTPStore, accessKeyAuth register.AccessKeyAuthenticator) { //nolint:funlen // define brev command cmd.AddCommand(set.NewCmdSet(t, loginCmdStore, noLoginCmdStore)) cmd.AddCommand(ls.NewCmdLs(t, loginCmdStore, noLoginCmdStore)) cmd.AddCommand(org.NewCmdOrg(t, loginCmdStore, noLoginCmdStore)) @@ -318,7 +324,7 @@ func createCmdTree(cmd *cobra.Command, t *terminal.Terminal, loginCmdStore *stor cmd.AddCommand(reset.NewCmdReset(t, loginCmdStore, noLoginCmdStore)) cmd.AddCommand(profile.NewCmdProfile(t, loginCmdStore, noLoginCmdStore)) cmd.AddCommand(refresh.NewCmdRefresh(t, loginCmdStore)) - cmd.AddCommand(register.NewCmdRegister(t, externalNodeCmdStore)) + cmd.AddCommand(register.NewCmdRegister(t, externalNodeCmdStore, accessKeyAuth)) cmd.AddCommand(deregister.NewCmdDeregister(t, externalNodeCmdStore)) cmd.AddCommand(upgrade.NewCmdUpgrade(t, noLoginCmdStore)) cmd.AddCommand(enablessh.NewCmdEnableSSH(t, externalNodeCmdStore)) diff --git a/pkg/cmd/deregister/deregister_test.go b/pkg/cmd/deregister/deregister_test.go index 95c5ac99..ea96a929 100644 --- a/pkg/cmd/deregister/deregister_test.go +++ b/pkg/cmd/deregister/deregister_test.go @@ -64,6 +64,13 @@ func (m *mockRegistrationStore) Load() (*register.DeviceRegistration, error) { return m.reg, nil } +func (m *mockRegistrationStore) LoadAny() (*register.DeviceRegistration, error) { + if m.reg == nil { + return nil, fmt.Errorf("no registration") + } + return m.reg, nil +} + func (m *mockRegistrationStore) Delete() error { m.reg = nil return nil diff --git a/pkg/cmd/grantssh/grantssh_test.go b/pkg/cmd/grantssh/grantssh_test.go index 10b3cbab..ef4dc367 100644 --- a/pkg/cmd/grantssh/grantssh_test.go +++ b/pkg/cmd/grantssh/grantssh_test.go @@ -51,6 +51,13 @@ func (m *mockRegistrationStore) Load() (*register.DeviceRegistration, error) { return m.reg, nil } +func (m *mockRegistrationStore) LoadAny() (*register.DeviceRegistration, error) { + if m.reg == nil { + return nil, fmt.Errorf("no registration") + } + return m.reg, nil +} + func (m *mockRegistrationStore) Delete() error { m.reg = nil return nil diff --git a/pkg/cmd/register/device_registration_store.go b/pkg/cmd/register/device_registration_store.go index 315dfb99..0dd0235d 100644 --- a/pkg/cmd/register/device_registration_store.go +++ b/pkg/cmd/register/device_registration_store.go @@ -20,6 +20,18 @@ const ( globalRegistrationDir = "/etc/brev" ) +// Registration status values written to DeviceRegistration.Status. +const ( + // RegistrationStatusPending marks an intent record written before AddNode + // has been confirmed. It captures the chosen device ID (plus name/org and + // hardware) so a retry can resume with the same ID: AddNode is idempotent on + // device_id, so reusing it does not create a duplicate backend node. + RegistrationStatusPending = "pending" + // RegistrationStatusRegistered marks a fully confirmed registration: AddNode + // succeeded and the external node ID has been persisted locally. + RegistrationStatusRegistered = "registered" +) + // DeviceRegistration is the persistent identity file for a registered device. // Fields align with the AddNodeResponse from dev-plane. type DeviceRegistration struct { @@ -30,12 +42,27 @@ type DeviceRegistration struct { DeviceID string `json:"device_id"` RegisteredAt string `json:"registered_at"` HardwareProfile HardwareProfile `json:"hardware_profile"` + // RegistrationToken is a UI-supplied token captured at registration time and + // persisted so a retry reuses the same token. It is not yet sent to the + // backend (AddNode has no token field today); wired in for future use. + RegistrationToken string `json:"registration_token,omitempty"` + // Status is RegistrationStatusPending (intent written, AddNode not yet + // confirmed) or RegistrationStatusRegistered (confirmed). Empty is treated + // as registered for backwards compatibility with files written by older + // CLI versions that predate the Status field. + Status string `json:"status,omitempty"` } // RegistrationStore defines the contract for persisting device registration data. type RegistrationStore interface { Save(reg *DeviceRegistration) error + // Load returns a fully-confirmed registration (one with an external node + // ID). It errors for pending (in-progress) records. Load() (*DeviceRegistration, error) + // LoadAny returns the registration record regardless of whether it has + // been confirmed, erroring only if the file is missing or unparseable. + // Used by the register flow to detect and resume pending attempts. + LoadAny() (*DeviceRegistration, error) Delete() error Exists() (bool, error) } @@ -87,6 +114,32 @@ func (s *FileRegistrationStore) Load() (*DeviceRegistration, error) { return nil, breverrors.WrapAndTrace(err) } if reg.ExternalNodeID == "" || reg.OrgID == "" { + if reg.Status == RegistrationStatusPending { + return nil, breverrors.New("device registration is incomplete; re-run 'brev register' to finish") + } + return nil, breverrors.New("malformed registration") + } + return ®, nil +} + +// LoadAny reads the registration file and returns the parsed DeviceRegistration +// regardless of whether registration has been confirmed. It errors only if the +// file is missing or unparseable. Use Load() when a fully-confirmed registration +// is required (e.g., enable-ssh, deregister). +func (s *FileRegistrationStore) LoadAny() (*DeviceRegistration, error) { + path := s.path() + exists, err := s.Exists() + if err != nil { + return nil, breverrors.WrapAndTrace(err) + } + if !exists { + return nil, breverrors.New("device registration not found, run 'brev register' first") + } + var reg DeviceRegistration + if err := files.ReadJSON(files.AppFs, path, ®); err != nil { + return nil, breverrors.WrapAndTrace(err) + } + if reg.OrgID == "" && reg.DeviceID == "" { return nil, breverrors.New("malformed registration") } return ®, nil diff --git a/pkg/cmd/register/device_registration_store_test.go b/pkg/cmd/register/device_registration_store_test.go index 39d7b1a2..e8622568 100644 --- a/pkg/cmd/register/device_registration_store_test.go +++ b/pkg/cmd/register/device_registration_store_test.go @@ -1,6 +1,7 @@ package register import ( + "strings" "testing" "github.com/brevdev/brev-cli/pkg/files" @@ -197,3 +198,68 @@ func Test_DeleteRegistration_FailsWhenMissing(t *testing.T) { t.Error("expected error deleting missing registration") } } + +// Test_LoadAny_ReturnsPendingRecord verifies LoadAny returns a pending (intent) +// record that has no external node ID, which the strict Load() rejects. +func Test_LoadAny_ReturnsPendingRecord(t *testing.T) { + cleanup := setupTestFs(t) + defer cleanup() + + store := NewFileRegistrationStore() + + pending := &DeviceRegistration{ + DisplayName: "My Spark", + OrgID: "org_xyz", + DeviceID: "device-uuid-123", + Status: RegistrationStatusPending, + } + if err := store.Save(pending); err != nil { + t.Fatalf("Save failed: %v", err) + } + + loaded, err := store.LoadAny() + if err != nil { + t.Fatalf("LoadAny failed: %v", err) + } + if loaded.DeviceID != "device-uuid-123" { + t.Errorf("DeviceID mismatch: got %s, want device-uuid-123", loaded.DeviceID) + } + if loaded.Status != RegistrationStatusPending { + t.Errorf("Status mismatch: got %q, want %q", loaded.Status, RegistrationStatusPending) + } + if loaded.ExternalNodeID != "" { + t.Errorf("pending record should have no ExternalNodeID, got %q", loaded.ExternalNodeID) + } + + // The strict Load() must reject the pending record. + if _, err := store.Load(); err == nil { + t.Error("expected Load() to error on a pending record") + } +} + +// Test_Load_PendingRecordErrorMessage verifies Load() reports a clear, actionable +// message for a pending (in-progress) registration. +func Test_Load_PendingRecordErrorMessage(t *testing.T) { + cleanup := setupTestFs(t) + defer cleanup() + + store := NewFileRegistrationStore() + + pending := &DeviceRegistration{ + DisplayName: "My Spark", + OrgID: "org_xyz", + DeviceID: "device-uuid-123", + Status: RegistrationStatusPending, + } + if err := store.Save(pending); err != nil { + t.Fatalf("Save failed: %v", err) + } + + _, err := store.Load() + if err == nil { + t.Fatal("expected Load() to error on a pending record") + } + if !strings.Contains(err.Error(), "incomplete") { + t.Errorf("expected 'incomplete' in error, got: %v", err) + } +} diff --git a/pkg/cmd/register/register.go b/pkg/cmd/register/register.go index 41692d55..a7e970f8 100644 --- a/pkg/cmd/register/register.go +++ b/pkg/cmd/register/register.go @@ -5,7 +5,7 @@ import ( "context" "errors" "fmt" - "os/user" + "os" "strings" "time" @@ -13,12 +13,14 @@ import ( "connectrpc.com/connect" "github.com/google/uuid" + "github.com/brevdev/brev-cli/pkg/auth" "github.com/brevdev/brev-cli/pkg/config" "github.com/brevdev/brev-cli/pkg/entity" breverrors "github.com/brevdev/brev-cli/pkg/errors" "github.com/brevdev/brev-cli/pkg/externalnode" "github.com/brevdev/brev-cli/pkg/externalnode/helpers" "github.com/brevdev/brev-cli/pkg/names" + "github.com/brevdev/brev-cli/pkg/store" "github.com/brevdev/brev-cli/pkg/sudo" "github.com/brevdev/brev-cli/pkg/terminal" @@ -61,6 +63,7 @@ type registerDeps struct { nodeClients externalnode.NodeClientFactory hardwareProfiler HardwareProfiler registrationStore RegistrationStore + accessKeyAuth AccessKeyAuthenticator } func defaultRegisterDeps() registerDeps { @@ -75,31 +78,110 @@ func defaultRegisterDeps() registerDeps { nodeClients: DefaultNodeClientFactory{}, hardwareProfiler: &SystemHardwareProfiler{}, registrationStore: NewFileRegistrationStore(), + accessKeyAuth: noopAccessKeyAuth{}, + } +} + +// For now the access key is a Brev API key (the bak- prefix); in the future +// this may be a more specific machine key. +type AccessKeyAuthenticator interface { + SeedSession(apiKey string) error + Persist(apiKey, orgID string) error +} + +// accessKeyAuthenticator is the production AccessKeyAuthenticator. It seeds the +// in-memory session store (used by external-node commands) and persists the +// key to the file-backed credential store (used by other brev commands). +type accessKeyAuthenticator struct { + session *store.MemoryAuthStore + durable *auth.LoginAuth +} + +// NewAccessKeyAuthenticator returns the production AccessKeyAuthenticator. +// session is the in-memory store backing the register command's auth (so the +// current command authenticates with the key); durable is the file-backed auth +// used by 'brev login --api-key' (so the key survives for other commands). +func NewAccessKeyAuthenticator(session *store.MemoryAuthStore, durable *auth.LoginAuth) AccessKeyAuthenticator { + return accessKeyAuthenticator{session: session, durable: durable} +} + +func (a accessKeyAuthenticator) SeedSession(apiKey string) error { + a.session.Seed(entity.AuthTokens{APIKey: apiKey}) + return nil +} + +func (a accessKeyAuthenticator) Persist(apiKey, orgID string) error { + if err := a.durable.LoginWithAPIKey(apiKey, orgID); err != nil { + return fmt.Errorf("failed to save access key: %w", err) + } + return nil +} + +// noopAccessKeyAuth is a no-op AccessKeyAuthenticator used as the default in +// defaultRegisterDeps (tests and any path that doesn't supply an access key). +// Production wiring overrides it via NewCmdRegister. +type noopAccessKeyAuth struct{} + +func (noopAccessKeyAuth) SeedSession(string) error { return nil } +func (noopAccessKeyAuth) Persist(string, string) error { return nil } + +// accessKeyEnvVar is the environment variable consulted when --access-key is +// not set, enabling headless/CI registration without modifying the command. +const accessKeyEnvVar = "BREV_ACCESS_KEY" + +// resolveAccessKey returns the access key from the --access-key flag, falling +// back to the BREV_ACCESS_KEY environment variable. Returns "" when neither is +// set, in which case the caller falls back to the interactive login-link flow. +func resolveAccessKey(flagValue string) string { + if v := strings.TrimSpace(flagValue); v != "" { + return v + } + return strings.TrimSpace(os.Getenv(accessKeyEnvVar)) +} + +// persistAccessKey durably saves the access key for other brev commands. It is +// best-effort: the in-memory session was already seeded, so the current command +// authenticates with the key regardless; a failure here only means future +// commands won't pick it up automatically. +func persistAccessKey(t *terminal.Terminal, deps registerDeps, apiKey, orgID string) { + if err := deps.accessKeyAuth.Persist(apiKey, orgID); err != nil { + t.Vprintf(" %s\n", t.Yellow(fmt.Sprintf("Warning: failed to save access key for future commands: %v", err))) } } var ( registerLong = `Register your device with NVIDIA Brev -This command sets up network connectivity and registers this machine with Brev. +This command registers this machine with Brev and brings up the Brev tunnel. +It no longer enables SSH as part of registration; run 'brev enable-ssh' afterwards +if you want to SSH to this device. Two modes are supported: - • Interactive (default): run 'brev register' with no flags and follow prompts for device name, org, and options. - • Non-interactive: use any of --name, --org, or --ssh-port. No prompts; --name and --org are required. Use for scripts/CI.` + • Interactive (default): run 'brev register' with no flags and follow prompts for device name and org. + • Non-interactive: use --name and --org. No prompts; both are required. Use for scripts/CI. + +Headless auth (credential chain): pass --access-key (a Brev API key) or set +the BREV_ACCESS_KEY environment variable to authenticate without the login +link; the key is saved for other brev commands, like 'brev login --api-key'. +If neither is set, the interactive login-link flow is used.` registerExample = ` # Interactive (prompts for device name, org, confirmations) brev register - # Non-interactive (any flag implies no prompts; --name and --org required) + # Non-interactive (--name and --org required) brev register --name my-node --org my-org - brev register --name my-node --org my-org --ssh-port 22` + + # Enable SSH access to this device after registering + brev enable-ssh` ) -func NewCmdRegister(t *terminal.Terminal, store RegisterStore) *cobra.Command { +func NewCmdRegister(t *terminal.Terminal, store RegisterStore, accessKeyAuth AccessKeyAuthenticator) *cobra.Command { var orgFlag string var nameFlag string - var sshPort int + var sshPort int // deprecated; accepted for backwards compatibility, no longer acted on var approveFlag bool + var registrationToken string + var accessKey string cmd := &cobra.Command{ Annotations: map[string]string{"configuration": ""}, @@ -112,13 +194,16 @@ func NewCmdRegister(t *terminal.Terminal, store RegisterStore) *cobra.Command { RunE: func(cmd *cobra.Command, args []string) error { interactive := nameFlag == "" && orgFlag == "" && sshPort == 0 opts := registerOpts{ - interactive: interactive, - name: nameFlag, - orgName: orgFlag, - sshPort: int32(sshPort), - skipConfirm: approveFlag, + interactive: interactive, + name: nameFlag, + orgName: orgFlag, + skipConfirm: approveFlag, + registrationToken: registrationToken, + accessKey: accessKey, } - return runRegister(cmd.Context(), t, store, opts, defaultRegisterDeps()) + deps := defaultRegisterDeps() + deps.accessKeyAuth = accessKeyAuth + return runRegister(cmd.Context(), t, store, opts, deps) }, } @@ -126,17 +211,28 @@ func NewCmdRegister(t *terminal.Terminal, store RegisterStore) *cobra.Command { cmd.Flags().StringVarP(&nameFlag, "name", "n", "", "device name (required when using non-interactive mode)") cmd.Flags().IntVarP(&sshPort, "ssh-port", "p", 0, "SSH port (if ssh access is desired)") cmd.Flags().BoolVar(&approveFlag, "approve", false, "skip all confirmation prompts (assume yes)") + cmd.Flags().StringVar(®istrationToken, "registration-token", "", "UI-supplied registration token (used when registering via the UI flow)") + cmd.Flags().StringVar(&accessKey, "access-key", "", "durable headless auth key (Brev API key); falls back to the "+accessKeyEnvVar+" env var, then the login-link flow") + + // --ssh-port is deprecated: registration no longer enables SSH. The flag is + // kept (and still accepted) for backwards compatibility with existing + // scripts. Use 'brev enable-ssh' after registration to enable SSH access. + // MarkDeprecated hides the flag from help and prints the deprecation notice + // when either --ssh-port or its -p shorthand is supplied. + _ = cmd.Flags().MarkDeprecated("ssh-port", "use 'brev enable-ssh' after registration to enable SSH access") return cmd } -// registerOpts carries mode and inputs: when interactive, name/orgName/sshPort are from prompts; otherwise from flags. +// registerOpts carries mode and inputs: when interactive, name/orgName are from +// prompts; otherwise from flags. type registerOpts struct { - interactive bool - name string - orgName string - sshPort int32 - skipConfirm bool + interactive bool + name string + orgName string + skipConfirm bool + registrationToken string + accessKey string } // runRegister runs a single registration flow; the only difference by mode is whether we prompt or use opts. @@ -155,18 +251,50 @@ func runRegister(ctx context.Context, t *terminal.Terminal, s RegisterStore, opt } } - // Run through the login flow - brevUser, err := s.GetCurrentUser() - if err != nil { + // Credential chain: --access-key flag -> BREV_ACCESS_KEY env -> interactive + // login-link flow (handled by the store's auth when no key is supplied). + // When an access key is present it is seeded into the in-memory session auth + // now, before any API call (e.g. GetCurrentUser, org resolution), so those + // calls authenticate with it. The org ID isn't known yet, so the durable, + // org-scoped save happens after the org is resolved (see persistAccessKey). + accessKey := resolveAccessKey(opts.accessKey) + if accessKey != "" { + if !auth.IsBrevAPIKey(accessKey) { + return breverrors.NewValidationError(fmt.Sprintf("access key must be a Brev API key (expected %s prefix); see 'brev login --api-key'", auth.BrevAPIKeyPrefix)) + } + if err := deps.accessKeyAuth.SeedSession(accessKey); err != nil { + return fmt.Errorf("failed to establish access key: %w", err) + } + t.Vprintf(" %s\n", t.Green("Authenticating with access key.")) + } + + // Verify the user is logged in before performing any local side effects. + if _, err := s.GetCurrentUser(); err != nil { return breverrors.WrapAndTrace(err) } - // Check if the device is already registered - alreadyRegistered, err := deps.registrationStore.Exists() + // Check for an existing registration (confirmed or in-progress). + exists, err := deps.registrationStore.Exists() if err != nil { return breverrors.WrapAndTrace(err) } - if alreadyRegistered { + if exists { + reg, err := deps.registrationStore.LoadAny() + if err != nil { + return breverrors.WrapAndTrace(err) + } + // The existing record (pending or confirmed) already carries the org ID, + // so durably persist the access key here if one was supplied. + if accessKey != "" { + persistAccessKey(t, deps, accessKey, reg.OrgID) + } + if reg.Status == RegistrationStatusPending { + // An earlier attempt wrote an intent record but never confirmed. + // Resume it, reusing the same device ID so AddNode is idempotent + // (a repeat call returns the existing node rather than a duplicate) + // and rectifying backend state into local CLI state. + return resumeRegistration(ctx, t, s, deps, reg) + } return checkExistingRegistration(ctx, t, s, deps) } @@ -199,6 +327,12 @@ func runRegister(ctx context.Context, t *terminal.Terminal, s RegisterStore, opt return err } + // Org is resolved now, so durably persist the access key (org-scoped) for + // other brev commands, mirroring 'brev login --api-key'. + if accessKey != "" { + persistAccessKey(t, deps, accessKey, org.ID) + } + t.Vprint("") t.Vprint(t.White("══════════════════════════════════════════════════")) t.Vprint(t.White(" Registering your device with Brev")) @@ -226,37 +360,32 @@ func runRegister(ctx context.Context, t *terminal.Terminal, s RegisterStore, opt } } - // Perform the registration steps - reg, err := runRegisterSteps(ctx, t, s, name, org, deps) + // Perform the registration steps. A stable device ID is generated here so a + // retry (which loads the pending record written by runRegisterSteps) reuses + // the same ID: AddNode is idempotent on device_id, so this does not create a + // duplicate backend node. + deviceID := uuid.New().String() + reg, err := runRegisterSteps(ctx, t, s, name, org, deps, deviceID, opts.registrationToken) if err != nil { return err } - // Determine if SSH access should be enabled - enableSSH := false - sshPortForGrant := int32(0) - if opts.sshPort != 0 { - enableSSH = true - sshPortForGrant = opts.sshPort - } - - // Grant SSH access if requested - if enableSSH { - osUser, err := user.Current() - if err != nil { - return fmt.Errorf("failed to determine current Linux user: %w", err) - } - if err := grantSSHAccessWithPort(ctx, t, deps, s, reg, brevUser, osUser, sshPortForGrant, opts.interactive, opts.skipConfirm); err != nil { - t.Vprintf(" %s\n", t.Yellow(fmt.Sprintf("Warning: %v", err))) - } - } - + suggestEnableSSH(t, reg) return nil } -// runRegisterSteps performs netbird install, hardware profile, AddNode, save registration, and runSetup. -// It does not prompt or enable SSH. Used by both flag-driven and prompt-driven flows. -func runRegisterSteps(ctx context.Context, t *terminal.Terminal, s RegisterStore, name string, org *entity.Organization, deps registerDeps) (*DeviceRegistration, error) { +// runRegisterSteps performs the registration steps: install the Brev tunnel, +// collect a hardware profile, write a pending intent record, call AddNode, +// persist the confirmed registration, and run the local setup command. +// +// deviceID is reused across retries so AddNode is idempotent: a second call with +// the same device_id returns the existing node rather than creating a duplicate. +// The pending intent record is written before AddNode so that a crash or +// timeout after the node was created backend-side can be resumed (see +// resumeRegistration) without orphaning the backend node. +// +// It does not prompt or enable SSH. Used by both fresh and resume flows. +func runRegisterSteps(ctx context.Context, t *terminal.Terminal, s RegisterStore, name string, org *entity.Organization, deps registerDeps, deviceID, registrationToken string) (*DeviceRegistration, error) { t.Vprint("") t.Vprint(t.Yellow("[Step 1/5] Downloading and installing Brev tunnel...")) @@ -279,7 +408,23 @@ func runRegisterSteps(ctx context.Context, t *terminal.Terminal, s RegisterStore t.Vprint("") t.Vprint(t.Yellow("[Step 3/5] Registering device with Brev...")) - deviceID := uuid.New().String() + + // Write a pending intent record before AddNode so a crash/timeout can be + // resumed with the same device ID (AddNode is idempotent on device_id). + pending := &DeviceRegistration{ + DisplayName: name, + OrgID: org.ID, + OrgName: org.Name, + DeviceID: deviceID, + RegistrationToken: registrationToken, + HardwareProfile: *hwProfile, + Status: RegistrationStatusPending, + RegisteredAt: time.Now().UTC().Format(time.RFC3339), + } + if err := deps.registrationStore.Save(pending); err != nil { + return nil, fmt.Errorf("failed to write pending registration: %w", err) + } + client := deps.nodeClients.NewNodeClient(s, config.GlobalConfig.GetBrevPublicAPIURL()) addResp, err := client.AddNode(ctx, connect.NewRequest(&nodev1.AddNodeRequest{ OrganizationId: org.ID, @@ -299,13 +444,15 @@ func runRegisterSteps(ctx context.Context, t *terminal.Terminal, s RegisterStore node := addResp.Msg.GetExternalNode() reg := &DeviceRegistration{ - ExternalNodeID: node.GetExternalNodeId(), - DisplayName: name, - OrgID: org.ID, - OrgName: org.Name, - DeviceID: deviceID, - RegisteredAt: time.Now().UTC().Format(time.RFC3339), - HardwareProfile: *hwProfile, + ExternalNodeID: node.GetExternalNodeId(), + DisplayName: name, + OrgID: org.ID, + OrgName: org.Name, + DeviceID: deviceID, + RegistrationToken: registrationToken, + RegisteredAt: time.Now().UTC().Format(time.RFC3339), + HardwareProfile: *hwProfile, + Status: RegistrationStatusRegistered, } t.Vprint("") @@ -419,58 +566,41 @@ func runSetup(node *nodev1.ExternalNode, t *terminal.Terminal, deps registerDeps } } -// grantSSHAccessWithPort enables SSH: shows confirm table, uses port or prompts if port is 0, then allocates port and grants access. -func grantSSHAccessWithPort(ctx context.Context, t *terminal.Terminal, deps registerDeps, tokenProvider externalnode.TokenProvider, reg *DeviceRegistration, brevUser *entity.User, osUser *user.User, port int32, interactive bool, skipConfirm bool) error { - brevUserName := brevUser.Username - if brevUserName == "" { - brevUserName = brevUser.Email - } - if brevUserName == "" { - brevUserName = brevUser.ID - } - +// resumeRegistration resumes an in-progress (pending) registration, reusing the +// same device ID so AddNode is idempotent. This rectifies backend state (the +// node created by a previous, partially-successful attempt) into local CLI +// state: a repeat AddNode with the same device ID returns the existing node, +// which is then persisted locally as confirmed. +func resumeRegistration(ctx context.Context, t *terminal.Terminal, s RegisterStore, deps registerDeps, pending *DeviceRegistration) error { t.Vprint("") t.Vprint(t.White("══════════════════════════════════════════════════")) - t.Vprint(t.White(" Enabling SSH access on this device")) + t.Vprint(t.White(" Resuming incomplete registration")) t.Vprint(t.White("══════════════════════════════════════════════════")) t.Vprint("") - if interactive && !skipConfirm { - t.Vprint(t.Green(" Please confirm before continuing:")) - t.Vprint("") - } - t.Vprintf(" %s %s\n", t.Green(fmt.Sprintf("%-14s", "Device:")), t.BoldBlue(reg.DisplayName+" ("+reg.ExternalNodeID+")")) - t.Vprintf(" %s %s\n", t.Green(fmt.Sprintf("%-14s", "Organization:")), t.BoldBlue(reg.OrgName+" ("+reg.OrgID+")")) - t.Vprintf(" %s %s\n", t.Green(fmt.Sprintf("%-14s", "Brev user:")), t.BoldBlue(brevUserName+" ("+brevUser.ID+")")) - t.Vprintf(" %s %s\n", t.Green(fmt.Sprintf("%-14s", "Linux user:")), t.BoldBlue(osUser.Username)) - - var err error - if port == 0 { - t.Vprint("") - port, err = PromptSSHPort(t) - if err != nil { - return fmt.Errorf("invalid SSH port: %w", err) - } - } else { - t.Vprintf(" %s %s\n", t.Green(fmt.Sprintf("%-14s", "SSH port:")), t.BoldBlue(fmt.Sprintf("%d", port))) - } + t.Vprintf(" %s %s\n", t.Green(fmt.Sprintf("%-14s", "Device:")), t.BoldBlue(pending.DisplayName)) + t.Vprintf(" %s %s\n", t.Green(fmt.Sprintf("%-14s", "Organization:")), t.BoldBlue(pending.OrgName+" ("+pending.OrgID+")")) + t.Vprintf(" %s %s\n", t.Green(fmt.Sprintf("%-14s", "Device ID:")), t.BoldBlue(pending.DeviceID)) t.Vprint("") + t.Vprint(" A previous registration attempt did not finish. Reusing the same") + t.Vprint(" device ID so no duplicate node is created on Brev.") - return grantSSHAccess(ctx, t, deps, tokenProvider, reg, brevUser, osUser, port) -} - -func grantSSHAccess(ctx context.Context, t *terminal.Terminal, deps registerDeps, tokenProvider externalnode.TokenProvider, reg *DeviceRegistration, brevUser *entity.User, osUser *user.User, port int32) error { - brevPortID, err := OpenSSHPort(ctx, t, deps.nodeClients, tokenProvider, reg, port) + org := &entity.Organization{ID: pending.OrgID, Name: pending.OrgName} + reg, err := runRegisterSteps(ctx, t, s, pending.DisplayName, org, deps, pending.DeviceID, pending.RegistrationToken) if err != nil { - return fmt.Errorf("allocate SSH port failed: %w", err) + return err } - err = SetupAndRegisterNodeSSHAccess(ctx, t, deps.nodeClients, tokenProvider, reg, brevUser, osUser.Username, brevPortID) - if err != nil { - return fmt.Errorf("grant SSH failed: %w", err) - } + suggestEnableSSH(t, reg) + return nil +} +// suggestEnableSSH prints a hint that SSH access can be enabled separately now +// that registration is complete. Registration no longer enables SSH itself. +func suggestEnableSSH(t *terminal.Terminal, reg *DeviceRegistration) { t.Vprint("") - t.Vprint(t.Green(fmt.Sprintf("SSH access enabled. You can now SSH to this device via: brev shell %s", reg.DisplayName))) - t.Vprint("") - return nil + if reg != nil && reg.DisplayName != "" { + t.Vprintf(" %s\n", t.Green(fmt.Sprintf("To enable SSH access to %s, run: brev enable-ssh", reg.DisplayName))) + } else { + t.Vprintf(" %s\n", t.Green("To enable SSH access to this device, run: brev enable-ssh")) + } } diff --git a/pkg/cmd/register/register_test.go b/pkg/cmd/register/register_test.go index d98b1a92..b9cdaafe 100644 --- a/pkg/cmd/register/register_test.go +++ b/pkg/cmd/register/register_test.go @@ -2,6 +2,7 @@ package register import ( "context" + "errors" "fmt" "net/http/httptest" "strings" @@ -11,7 +12,9 @@ import ( nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" "connectrpc.com/connect" + "github.com/brevdev/brev-cli/pkg/auth" "github.com/brevdev/brev-cli/pkg/entity" + breverrors "github.com/brevdev/brev-cli/pkg/errors" "github.com/brevdev/brev-cli/pkg/externalnode" "github.com/brevdev/brev-cli/pkg/sudo" "github.com/brevdev/brev-cli/pkg/terminal" @@ -79,6 +82,13 @@ func (m *mockRegistrationStore) Load() (*DeviceRegistration, error) { return m.reg, nil } +func (m *mockRegistrationStore) LoadAny() (*DeviceRegistration, error) { + if m.reg == nil { + return nil, fmt.Errorf("no registration") + } + return m.reg, nil +} + func (m *mockRegistrationStore) Delete() error { m.reg = nil return nil @@ -141,6 +151,32 @@ func (m mockNodeClientFactory) NewNodeClient(provider externalnode.TokenProvider return NewNodeServiceClient(provider, m.serverURL) } +// accessKeyPersistCall records a single Persist invocation. +type accessKeyPersistCall struct { + apiKey string + orgID string +} + +// mockAccessKeyAuth records SeedSession/Persist calls so tests can assert the +// credential chain wired up the access key (and which org ID it was persisted +// against) without touching real auth stores. +type mockAccessKeyAuth struct { + seedCalls []string + persistCalls []accessKeyPersistCall + seedErr error + persistErr error +} + +func (m *mockAccessKeyAuth) SeedSession(apiKey string) error { + m.seedCalls = append(m.seedCalls, apiKey) + return m.seedErr +} + +func (m *mockAccessKeyAuth) Persist(apiKey, orgID string) error { + m.persistCalls = append(m.persistCalls, accessKeyPersistCall{apiKey: apiKey, orgID: orgID}) + return m.persistErr +} + // testHardwareProfile returns a realistic HardwareProfile for use in tests. func testHardwareProfile() *HardwareProfile { cpuCount := int32(2) @@ -181,6 +217,7 @@ func testRegisterDeps(t *testing.T, svc *fakeNodeService, regStore RegistrationS profile: testHardwareProfile(), }, registrationStore: regStore, + accessKeyAuth: noopAccessKeyAuth{}, }, server } @@ -223,11 +260,8 @@ func Test_runRegister_HappyPath(t *testing.T) { deps.setupRunner = setupRunner - SetTestSSHPort(22) - defer ClearTestSSHPort() - term := terminal.New() - opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg", sshPort: 22} + opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg"} err := runRegister(context.Background(), term, store, opts, deps) if err != nil { t.Fatalf("runRegister failed: %v", err) @@ -295,7 +329,7 @@ func Test_runRegister_UserCancels(t *testing.T) { }) term := terminal.New() - opts := registerOpts{interactive: true, name: "", orgName: "", sshPort: 0} + opts := registerOpts{interactive: true, name: "", orgName: ""} err := runRegister(context.Background(), term, store, opts, deps) if err == nil { t.Fatal("expected error when user declines sudo gate") @@ -383,7 +417,7 @@ func Test_runRegister_AlreadyRegistered(t *testing.T) { term := terminal.New() // Pass the same name as the existing registration so we go through // the checkExistingRegistration path (not the different-name path). - opts := registerOpts{interactive: false, name: "Existing", orgName: "TestOrg", sshPort: 22} + opts := registerOpts{interactive: false, name: "Existing", orgName: "TestOrg"} err := runRegister(context.Background(), term, store, opts, deps) if err != nil { t.Fatalf("expected nil error, got: %v", err) @@ -412,7 +446,7 @@ func Test_runRegister_NoOrganization(t *testing.T) { defer server.Close() term := terminal.New() - opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg", sshPort: 22} + opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg"} err := runRegister(context.Background(), term, store, opts, deps) if err == nil { t.Fatal("expected error when no org exists") @@ -451,11 +485,8 @@ func Test_runRegister_WithOrgFlag(t *testing.T) { defer server.Close() deps.setupRunner = setupRunner - SetTestSSHPort(22) - defer ClearTestSSHPort() - term := terminal.New() - opts := registerOpts{interactive: false, name: "my-spark", orgName: "SpecificOrg", sshPort: 22} + opts := registerOpts{interactive: false, name: "my-spark", orgName: "SpecificOrg"} err := runRegister(context.Background(), term, store, opts, deps) if err != nil { t.Fatalf("runRegister with --org failed: %v", err) @@ -489,7 +520,7 @@ func Test_runRegister_WithOrgFlag_NotFound(t *testing.T) { defer server.Close() term := terminal.New() - opts := registerOpts{interactive: false, name: "my-spark", orgName: "NonexistentOrg", sshPort: 22} + opts := registerOpts{interactive: false, name: "my-spark", orgName: "NonexistentOrg"} err := runRegister(context.Background(), term, store, opts, deps) if err == nil { t.Fatal("expected error when org not found") @@ -519,19 +550,30 @@ func Test_runRegister_AddNodeFails(t *testing.T) { defer server.Close() term := terminal.New() - opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg", sshPort: 22} + opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg"} err := runRegister(context.Background(), term, store, opts, deps) if err == nil { t.Fatal("expected error when AddNode fails") } - // Registration should not exist on failure - exists, err := regStore.Exists() - if err != nil { - t.Fatalf("Exists error: %v", err) + // AddNode failed, but a pending intent record is written before AddNode so a + // retry can resume with the same device ID. The record must be pending (no + // external node ID yet) and carry a device ID for the retry to reuse. + reg, loadErr := regStore.LoadAny() + if loadErr != nil { + t.Fatalf("expected pending intent record after AddNode failure, got: %v", loadErr) } - if exists { - t.Error("registration should not exist after AddNode failure") + if reg.Status != RegistrationStatusPending { + t.Errorf("expected pending status, got %q", reg.Status) + } + if reg.ExternalNodeID != "" { + t.Errorf("expected no ExternalNodeID on AddNode failure, got %q", reg.ExternalNodeID) + } + if reg.DeviceID == "" { + t.Error("expected pending record to carry a device ID for retry") + } + if reg.DisplayName != "my-spark" || reg.OrgID != "org_123" { + t.Errorf("pending record mismatch: name=%q org=%q", reg.DisplayName, reg.OrgID) } } @@ -566,11 +608,8 @@ func Test_runRegister_NoSetupCommand(t *testing.T) { deps.setupRunner = setupRunner - SetTestSSHPort(22) - defer ClearTestSSHPort() - term := terminal.New() - opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg", sshPort: 22} + opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg"} err := runRegister(context.Background(), term, store, opts, deps) if err != nil { t.Fatalf("runRegister failed: %v", err) @@ -662,109 +701,6 @@ Peers count: 0/0 Connected` } } -func Test_runRegister_GrantSSH_retries_on_connection_error_then_succeeds(t *testing.T) { - regStore := &mockRegistrationStore{} - - store := &mockRegisterStore{ - user: &entity.User{ID: "user_1"}, - org: &entity.Organization{ID: "org_123", Name: "TestOrg"}, - token: "tok", - } - - var grantCalls int - svc := &fakeNodeService{ - addNodeFn: func(req *nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) { - return &nodev1.AddNodeResponse{ - ExternalNode: &nodev1.ExternalNode{ - ExternalNodeId: "unode_abc", - OrganizationId: "org_123", - Name: req.GetName(), - DeviceId: req.GetDeviceId(), - ConnectivityInfo: &nodev1.ConnectivityInfo{ - RegistrationCommand: "netbird up --key abc", - }, - }, - }, nil - }, - grantNodeSSHAccessFn: func(_ *nodev1.GrantNodeSSHAccessRequest) (*nodev1.GrantNodeSSHAccessResponse, error) { - grantCalls++ - if grantCalls < 2 { - return nil, connect.NewError(connect.CodeInternal, nil) - } - return &nodev1.GrantNodeSSHAccessResponse{}, nil - }, - } - - deps, server := testRegisterDeps(t, svc, regStore) - defer server.Close() - - deps.prompter = mockConfirmer{confirm: true} - - SetTestSSHPort(22) - defer ClearTestSSHPort() - - term := terminal.New() - opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg", sshPort: 22} - err := runRegister(context.Background(), term, store, opts, deps) - if err != nil { - t.Fatalf("runRegister failed: %v", err) - } - - if grantCalls != 2 { - t.Errorf("expected GrantNodeSSHAccess to be called 2 times (retry once), got %d", grantCalls) - } -} - -func Test_runRegister_GrantSSH_no_retry_on_permanent_error(t *testing.T) { - regStore := &mockRegistrationStore{} - - store := &mockRegisterStore{ - user: &entity.User{ID: "user_1"}, - org: &entity.Organization{ID: "org_123", Name: "TestOrg"}, - token: "tok", - } - - var grantCalls int - svc := &fakeNodeService{ - addNodeFn: func(req *nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) { - return &nodev1.AddNodeResponse{ - ExternalNode: &nodev1.ExternalNode{ - ExternalNodeId: "unode_abc", - OrganizationId: "org_123", - Name: req.GetName(), - DeviceId: req.GetDeviceId(), - ConnectivityInfo: &nodev1.ConnectivityInfo{ - RegistrationCommand: "netbird up --key abc", - }, - }, - }, nil - }, - grantNodeSSHAccessFn: func(_ *nodev1.GrantNodeSSHAccessRequest) (*nodev1.GrantNodeSSHAccessResponse, error) { - grantCalls++ - return nil, connect.NewError(connect.CodePermissionDenied, nil) - }, - } - - deps, server := testRegisterDeps(t, svc, regStore) - defer server.Close() - - deps.prompter = mockConfirmer{confirm: true} - - SetTestSSHPort(22) - defer ClearTestSSHPort() - - term := terminal.New() - opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg", sshPort: 22} - err := runRegister(context.Background(), term, store, opts, deps) - if err != nil { - t.Fatalf("runRegister should not fail the overall flow when SSH grant fails: %v", err) - } - - if grantCalls != 1 { - t.Errorf("expected GrantNodeSSHAccess to be called once (no retry on permanent error), got %d", grantCalls) - } -} - func Test_runRegister_NameValidation(t *testing.T) { tests := []struct { name string @@ -811,12 +747,9 @@ func Test_runRegister_NameValidation(t *testing.T) { deps, server := testRegisterDeps(t, svc, regStore) defer server.Close() - SetTestSSHPort(22) - defer ClearTestSSHPort() - term := terminal.New() var err error - opts := registerOpts{interactive: false, name: tt.input, orgName: "TestOrg", sshPort: 22} + opts := registerOpts{interactive: false, name: tt.input, orgName: "TestOrg"} err = runRegister(context.Background(), term, store, opts, deps) if tt.wantErr { if err == nil { @@ -848,7 +781,7 @@ func Test_runRegister_PlatformIncompatible(t *testing.T) { deps.platform = mockPlatform{compatible: false} term := terminal.New() - opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg", sshPort: 22} + opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg"} err := runRegister(context.Background(), term, store, opts, deps) if err == nil { t.Fatal("expected error when platform is incompatible") @@ -874,7 +807,7 @@ func Test_runRegister_HardwareProfilerFailure(t *testing.T) { deps.hardwareProfiler = &mockHardwareProfiler{err: fmt.Errorf("nvml init failed")} term := terminal.New() - opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg", sshPort: 22} + opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg"} err := runRegister(context.Background(), term, store, opts, deps) if err == nil { t.Fatal("expected error when hardware profiler fails") @@ -900,7 +833,7 @@ func Test_runRegister_NetBirdInstallFailure(t *testing.T) { deps.netbird = mockNetBirdManager{err: fmt.Errorf("install failed")} term := terminal.New() - opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg", sshPort: 22} + opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg"} err := runRegister(context.Background(), term, store, opts, deps) if err == nil { t.Fatal("expected error when NetBird install fails") @@ -925,7 +858,7 @@ func Test_runRegister_NoNameNotRegistered(t *testing.T) { defer server.Close() term := terminal.New() - opts := registerOpts{interactive: false, name: "", orgName: "", sshPort: 22} + opts := registerOpts{interactive: false, name: "", orgName: ""} err := runRegister(context.Background(), term, store, opts, deps) if err == nil { t.Fatal("expected error when no name/org in non-interactive mode") @@ -967,7 +900,7 @@ func Test_runRegister_NoNameAlreadyRegistered(t *testing.T) { defer server.Close() term := terminal.New() - opts := registerOpts{interactive: false, name: "Existing", orgName: "TestOrg", sshPort: 22} + opts := registerOpts{interactive: false, name: "Existing", orgName: "TestOrg"} err := runRegister(context.Background(), term, store, opts, deps) if err != nil { t.Fatalf("expected nil error when already registered with no name, got: %v", err) @@ -980,152 +913,399 @@ func Test_runRegister_NoNameAlreadyRegistered(t *testing.T) { } } -func Test_runRegister_OpenSSHPort(t *testing.T) { // nolint:funlen, gocyclo, gocognit // test - tests := []struct { - name string - port int32 - openFn func(*nodev1.OpenPortRequest) (*nodev1.OpenPortResponse, error) - verify func(t *testing.T, openReq *nodev1.OpenPortRequest, grantReq *nodev1.GrantNodeSSHAccessRequest, reg *mockRegistrationStore, err error) - }{ - { - name: "SendsCorrectArgs", - port: 2222, - openFn: func(req *nodev1.OpenPortRequest) (*nodev1.OpenPortResponse, error) { - return &nodev1.OpenPortResponse{ - Port: &nodev1.Port{ - PortId: "port_ssh", - Protocol: req.GetProtocol(), - PortNumber: req.GetPortNumber(), +// Test_runRegister_ResumesPendingRegistration verifies that a pending (intent) +// record left by an interrupted attempt is resumed: the same device ID is reused +// (AddNode is idempotent on device_id), and the final record is marked +// registered with the backend-assigned external node ID. The captured +// registration token is preserved across the resume. +func Test_runRegister_ResumesPendingRegistration(t *testing.T) { + const pendingDeviceID = "device-uuid-pending" + pending := &DeviceRegistration{ + DisplayName: "My Spark", + OrgID: "org_123", + OrgName: "TestOrg", + DeviceID: pendingDeviceID, + RegistrationToken: "ui-token-pending", + Status: RegistrationStatusPending, + } + regStore := &mockRegistrationStore{reg: pending} + + store := &mockRegisterStore{ + user: &entity.User{ID: "user_1"}, + org: &entity.Organization{ID: "org_123", Name: "TestOrg"}, + token: "tok", + } + + var addNodeDeviceIDs []string + var addNodeCalls int + svc := &fakeNodeService{ + addNodeFn: func(req *nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) { + addNodeCalls++ + addNodeDeviceIDs = append(addNodeDeviceIDs, req.GetDeviceId()) + return &nodev1.AddNodeResponse{ + ExternalNode: &nodev1.ExternalNode{ + ExternalNodeId: "unode_abc", + OrganizationId: req.GetOrganizationId(), + Name: req.GetName(), + DeviceId: req.GetDeviceId(), + ConnectivityInfo: &nodev1.ConnectivityInfo{ + RegistrationCommand: "netbird up --key abc", }, - }, nil - }, - verify: func(t *testing.T, openReq *nodev1.OpenPortRequest, _ *nodev1.GrantNodeSSHAccessRequest, _ *mockRegistrationStore, err error) { - t.Helper() - if err != nil { - t.Fatalf("runRegister failed: %v", err) - } - if openReq == nil { - t.Fatal("expected OpenPort to be called") - } - if openReq.GetExternalNodeId() != "unode_abc" { - t.Errorf("expected node ID unode_abc, got %s", openReq.GetExternalNodeId()) - } - if openReq.GetProtocol() != nodev1.PortProtocol_PORT_PROTOCOL_TCP { - t.Errorf("expected PORT_PROTOCOL_TCP, got %s", openReq.GetProtocol()) - } - if openReq.GetPortNumber() != 2222 { - t.Errorf("expected port 2222, got %d", openReq.GetPortNumber()) - } - }, + }, + }, nil }, - { - name: "FailureIsSoftError", - port: 22, - openFn: func(_ *nodev1.OpenPortRequest) (*nodev1.OpenPortResponse, error) { - return nil, connect.NewError(connect.CodeInternal, fmt.Errorf("skybridge unavailable")) - }, - verify: func(t *testing.T, _ *nodev1.OpenPortRequest, _ *nodev1.GrantNodeSSHAccessRequest, regStore *mockRegistrationStore, err error) { - t.Helper() - if err != nil { - t.Fatalf("registration should succeed even when OpenSSHPort fails (soft error), got: %v", err) - } - exists, _ := regStore.Exists() - if !exists { - t.Error("expected registration to still exist after OpenSSHPort failure") - } - }, + } + + deps, server := testRegisterDeps(t, svc, regStore) + defer server.Close() + + term := terminal.New() + // On resume, opts (name/org/interactive) are ignored in favor of the pending + // record. Use interactive mode to skip the non-interactive --name/--org + // validation that runs before the Exists check. + err := runRegister(context.Background(), term, store, registerOpts{interactive: true}, deps) + if err != nil { + t.Fatalf("runRegister failed: %v", err) + } + + if addNodeCalls != 1 { + t.Fatalf("expected AddNode to be called once, got %d", addNodeCalls) + } + if len(addNodeDeviceIDs) != 1 || addNodeDeviceIDs[0] != pendingDeviceID { + t.Errorf("expected AddNode to reuse device ID %q, got %v", pendingDeviceID, addNodeDeviceIDs) + } + + reg, loadErr := regStore.LoadAny() + if loadErr != nil { + t.Fatalf("LoadAny failed: %v", loadErr) + } + if reg.Status != RegistrationStatusRegistered { + t.Errorf("expected status %q after resume, got %q", RegistrationStatusRegistered, reg.Status) + } + if reg.ExternalNodeID != "unode_abc" { + t.Errorf("expected ExternalNodeID unode_abc, got %q", reg.ExternalNodeID) + } + if reg.DeviceID != pendingDeviceID { + t.Errorf("expected device ID to remain %q, got %q", pendingDeviceID, reg.DeviceID) + } + if reg.RegistrationToken != "ui-token-pending" { + t.Errorf("expected registration token to be preserved, got %q", reg.RegistrationToken) + } +} + +// Test_runRegister_PersistsRegistrationToken verifies the --registration-token +// flag value is threaded through registration and persisted to local state (it +// is not yet sent to the backend). +func Test_runRegister_PersistsRegistrationToken(t *testing.T) { + regStore := &mockRegistrationStore{} + store := &mockRegisterStore{ + user: &entity.User{ID: "user_1"}, + org: &entity.Organization{ID: "org_123", Name: "TestOrg"}, + token: "tok", + } + svc := &fakeNodeService{ + addNodeFn: func(req *nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) { + return &nodev1.AddNodeResponse{ + ExternalNode: &nodev1.ExternalNode{ + ExternalNodeId: "unode_abc", + OrganizationId: req.GetOrganizationId(), + Name: req.GetName(), + DeviceId: req.GetDeviceId(), + }, + }, nil }, - { - name: "InvalidPortNoAPICall", - port: 99999, - verify: func(t *testing.T, openReq *nodev1.OpenPortRequest, _ *nodev1.GrantNodeSSHAccessRequest, regStore *mockRegistrationStore, err error) { - t.Helper() - if err != nil { - t.Fatalf("registration should succeed even when SSH port is invalid (soft error), got: %v", err) - } - if openReq != nil { - t.Error("expected OpenPort NOT to be called for invalid port") - } - exists, _ := regStore.Exists() - if !exists { - t.Error("expected registration to still exist after invalid port") - } + } + + deps, server := testRegisterDeps(t, svc, regStore) + defer server.Close() + + term := terminal.New() + opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg", registrationToken: "ui-token-xyz"} + err := runRegister(context.Background(), term, store, opts, deps) + if err != nil { + t.Fatalf("runRegister failed: %v", err) + } + + reg, loadErr := regStore.LoadAny() + if loadErr != nil { + t.Fatalf("LoadAny failed: %v", loadErr) + } + if reg.RegistrationToken != "ui-token-xyz" { + t.Errorf("expected RegistrationToken to be persisted, got %q", reg.RegistrationToken) + } + if reg.Status != RegistrationStatusRegistered { + t.Errorf("expected registered status, got %q", reg.Status) + } +} + +// --- Access key credential chain --- + +const testAccessKey = auth.BrevAPIKeyPrefix + "test-key" + +// accessKeyAddNodeFn returns an AddNode handler that echoes the request back as +// a registered node, suitable for the access-key happy-path tests. +func accessKeyAddNodeFn(t *testing.T) func(*nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) { + t.Helper() + return func(req *nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) { + if req.GetOrganizationId() != "org_123" { + t.Errorf("unexpected org: %s", req.GetOrganizationId()) + } + return &nodev1.AddNodeResponse{ + ExternalNode: &nodev1.ExternalNode{ + ExternalNodeId: "unode_abc", + OrganizationId: req.GetOrganizationId(), + Name: req.GetName(), + DeviceId: req.GetDeviceId(), }, + }, nil + } +} + +// ensureNoAccessKeyEnv unsets BREV_ACCESS_KEY for the test so the env fallback +// doesn't leak in (or out) of the credential-chain tests. +func ensureNoAccessKeyEnv(t *testing.T) { + t.Helper() + t.Setenv(accessKeyEnvVar, "") +} + +// Test_resolveAccessKey verifies the flag -> env -> empty fallback order. +func Test_resolveAccessKey(t *testing.T) { + ensureNoAccessKeyEnv(t) + + if got := resolveAccessKey(""); got != "" { + t.Errorf("expected empty when no flag and no env, got %q", got) + } + + t.Setenv(accessKeyEnvVar, auth.BrevAPIKeyPrefix+"env-key") + if got := resolveAccessKey(""); got != auth.BrevAPIKeyPrefix+"env-key" { + t.Errorf("expected env value, got %q", got) + } + + // Flag wins over env. + if got := resolveAccessKey(auth.BrevAPIKeyPrefix + "flag-key"); got != auth.BrevAPIKeyPrefix+"flag-key" { + t.Errorf("expected flag value to win, got %q", got) + } + + // Whitespace-only flag falls through to env. + if got := resolveAccessKey(" "); got != auth.BrevAPIKeyPrefix+"env-key" { + t.Errorf("expected env value when flag is blank, got %q", got) + } +} + +// Test_runRegister_AccessKeyFlag_SeedsAndPersists verifies that --access-key +// seeds the session (before any API call) and is durably persisted with the +// resolved org ID once the org is known, then registration proceeds. +func Test_runRegister_AccessKeyFlag_SeedsAndPersists(t *testing.T) { + ensureNoAccessKeyEnv(t) + + regStore := &mockRegistrationStore{} + store := &mockRegisterStore{ + user: &entity.User{ID: "user_1"}, + org: &entity.Organization{ID: "org_123", Name: "TestOrg"}, + token: "tok", + } + svc := &fakeNodeService{addNodeFn: accessKeyAddNodeFn(t)} + + deps, server := testRegisterDeps(t, svc, regStore) + defer server.Close() + accessKeyAuth := &mockAccessKeyAuth{} + deps.accessKeyAuth = accessKeyAuth + + term := terminal.New() + opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg", accessKey: testAccessKey} + if err := runRegister(context.Background(), term, store, opts, deps); err != nil { + t.Fatalf("runRegister failed: %v", err) + } + + if len(accessKeyAuth.seedCalls) != 1 || accessKeyAuth.seedCalls[0] != testAccessKey { + t.Errorf("expected SeedSession called once with the access key, got %v", accessKeyAuth.seedCalls) + } + if len(accessKeyAuth.persistCalls) != 1 { + t.Fatalf("expected Persist called once, got %d", len(accessKeyAuth.persistCalls)) + } + if accessKeyAuth.persistCalls[0].apiKey != testAccessKey { + t.Errorf("expected persist api key %q, got %q", testAccessKey, accessKeyAuth.persistCalls[0].apiKey) + } + if accessKeyAuth.persistCalls[0].orgID != "org_123" { + t.Errorf("expected persist org ID org_123 (resolved), got %q", accessKeyAuth.persistCalls[0].orgID) + } +} + +// Test_runRegister_AccessKeyEnv_Fallback verifies the BREV_ACCESS_KEY env var is +// used when --access-key is not supplied. +func Test_runRegister_AccessKeyEnv_Fallback(t *testing.T) { + ensureNoAccessKeyEnv(t) + t.Setenv(accessKeyEnvVar, testAccessKey) + + regStore := &mockRegistrationStore{} + store := &mockRegisterStore{ + user: &entity.User{ID: "user_1"}, + org: &entity.Organization{ID: "org_123", Name: "TestOrg"}, + token: "tok", + } + svc := &fakeNodeService{addNodeFn: accessKeyAddNodeFn(t)} + + deps, server := testRegisterDeps(t, svc, regStore) + defer server.Close() + accessKeyAuth := &mockAccessKeyAuth{} + deps.accessKeyAuth = accessKeyAuth + + term := terminal.New() + // No accessKey in opts; should fall back to the env var. + opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg"} + if err := runRegister(context.Background(), term, store, opts, deps); err != nil { + t.Fatalf("runRegister failed: %v", err) + } + + if len(accessKeyAuth.seedCalls) != 1 || accessKeyAuth.seedCalls[0] != testAccessKey { + t.Errorf("expected SeedSession called with env access key, got %v", accessKeyAuth.seedCalls) + } + if len(accessKeyAuth.persistCalls) != 1 || accessKeyAuth.persistCalls[0].apiKey != testAccessKey { + t.Errorf("expected Persist called with env access key, got %v", accessKeyAuth.persistCalls) + } +} + +// Test_runRegister_AccessKeyFlag_PrecedenceOverEnv verifies the flag wins when +// both --access-key and BREV_ACCESS_KEY are set. +func Test_runRegister_AccessKeyFlag_PrecedenceOverEnv(t *testing.T) { + ensureNoAccessKeyEnv(t) + t.Setenv(accessKeyEnvVar, auth.BrevAPIKeyPrefix+"env-key") + + regStore := &mockRegistrationStore{} + store := &mockRegisterStore{ + user: &entity.User{ID: "user_1"}, + org: &entity.Organization{ID: "org_123", Name: "TestOrg"}, + token: "tok", + } + svc := &fakeNodeService{addNodeFn: accessKeyAddNodeFn(t)} + + deps, server := testRegisterDeps(t, svc, regStore) + defer server.Close() + accessKeyAuth := &mockAccessKeyAuth{} + deps.accessKeyAuth = accessKeyAuth + + term := terminal.New() + opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg", accessKey: testAccessKey} + if err := runRegister(context.Background(), term, store, opts, deps); err != nil { + t.Fatalf("runRegister failed: %v", err) + } + + if len(accessKeyAuth.seedCalls) != 1 || accessKeyAuth.seedCalls[0] != testAccessKey { + t.Errorf("expected flag value to win over env, got %v", accessKeyAuth.seedCalls) + } +} + +// Test_runRegister_AccessKeyInvalid verifies a non-API-key value is rejected +// before any registration side effect. +func Test_runRegister_AccessKeyInvalid(t *testing.T) { + ensureNoAccessKeyEnv(t) + + regStore := &mockRegistrationStore{} + store := &mockRegisterStore{ + user: &entity.User{ID: "user_1"}, + org: &entity.Organization{ID: "org_123", Name: "TestOrg"}, + token: "tok", + } + svc := &fakeNodeService{addNodeFn: accessKeyAddNodeFn(t)} + + deps, server := testRegisterDeps(t, svc, regStore) + defer server.Close() + accessKeyAuth := &mockAccessKeyAuth{} + deps.accessKeyAuth = accessKeyAuth + + term := terminal.New() + opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg", accessKey: "not-a-brev-key"} + err := runRegister(context.Background(), term, store, opts, deps) + if err == nil { + t.Fatal("expected error for invalid access key") + } + var ve breverrors.ValidationError + if !errors.As(err, &ve) { + t.Errorf("expected a ValidationError, got %T: %v", err, err) + } + if !strings.Contains(err.Error(), auth.BrevAPIKeyPrefix) { + t.Errorf("expected error to mention the %s prefix, got: %v", auth.BrevAPIKeyPrefix, err) + } + if len(accessKeyAuth.seedCalls) != 0 { + t.Errorf("expected no SeedSession call on validation failure, got %v", accessKeyAuth.seedCalls) + } +} + +// Test_runRegister_AccessKey_PersistsOnAlreadyRegistered verifies that when an +// access key is supplied against an already-registered machine, it is still +// persisted (using the existing record's org ID) for other brev commands. +func Test_runRegister_AccessKey_PersistsOnAlreadyRegistered(t *testing.T) { + ensureNoAccessKeyEnv(t) + + regStore := &mockRegistrationStore{ + reg: &DeviceRegistration{ + ExternalNodeID: "unode_existing", + DisplayName: "Existing", + OrgID: "org_123", + Status: RegistrationStatusRegistered, }, - { - name: "GrantRequestHasNoPort", - port: 22, - verify: func(t *testing.T, _ *nodev1.OpenPortRequest, grantReq *nodev1.GrantNodeSSHAccessRequest, _ *mockRegistrationStore, err error) { - t.Helper() - if err != nil { - t.Fatalf("runRegister failed: %v", err) - } - if grantReq == nil { - t.Fatal("expected GrantNodeSSHAccess to be called") - } - if grantReq.GetExternalNodeId() != "unode_abc" { - t.Errorf("expected node ID unode_abc, got %s", grantReq.GetExternalNodeId()) - } - if grantReq.GetUserId() != "user_1" { - t.Errorf("expected user ID user_1, got %s", grantReq.GetUserId()) - } - }, + } + store := &mockRegisterStore{ + user: &entity.User{ID: "user_1"}, + org: &entity.Organization{ID: "org_123", Name: "TestOrg"}, + token: "tok", + } + svc := &fakeNodeService{ + getNodeFn: func(req *nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) { + return &nodev1.GetNodeResponse{ + ExternalNode: &nodev1.ExternalNode{ + ExternalNodeId: req.GetExternalNodeId(), + ConnectivityInfo: &nodev1.ConnectivityInfo{ + Status: nodev1.NetworkMemberStatus_NETWORK_MEMBER_STATUS_CONNECTED, + }, + }, + }, nil }, } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - regStore := &mockRegistrationStore{} - store := &mockRegisterStore{ - user: &entity.User{ID: "user_1"}, - org: &entity.Organization{ID: "org_123", Name: "TestOrg"}, - token: "tok", - } + deps, server := testRegisterDeps(t, svc, regStore) + defer server.Close() + accessKeyAuth := &mockAccessKeyAuth{} + deps.accessKeyAuth = accessKeyAuth - var gotOpenReq *nodev1.OpenPortRequest - var gotGrantReq *nodev1.GrantNodeSSHAccessRequest - svc := &fakeNodeService{ - addNodeFn: func(req *nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) { - return &nodev1.AddNodeResponse{ - ExternalNode: &nodev1.ExternalNode{ - ExternalNodeId: "unode_abc", - OrganizationId: "org_123", - Name: req.GetName(), - DeviceId: req.GetDeviceId(), - ConnectivityInfo: &nodev1.ConnectivityInfo{ - RegistrationCommand: "netbird up --key abc", - }, - }, - }, nil - }, - openPortFn: func(req *nodev1.OpenPortRequest) (*nodev1.OpenPortResponse, error) { - gotOpenReq = req - if tt.openFn != nil { - return tt.openFn(req) - } - return &nodev1.OpenPortResponse{ - Port: &nodev1.Port{PortId: "port_ssh", Protocol: req.GetProtocol(), PortNumber: req.GetPortNumber()}, - }, nil - }, - grantNodeSSHAccessFn: func(req *nodev1.GrantNodeSSHAccessRequest) (*nodev1.GrantNodeSSHAccessResponse, error) { - gotGrantReq = req - return &nodev1.GrantNodeSSHAccessResponse{}, nil - }, - } + term := terminal.New() + opts := registerOpts{interactive: false, name: "Existing", orgName: "TestOrg", accessKey: testAccessKey} + if err := runRegister(context.Background(), term, store, opts, deps); err != nil { + t.Fatalf("runRegister failed: %v", err) + } - deps, server := testRegisterDeps(t, svc, regStore) - defer server.Close() + if len(accessKeyAuth.seedCalls) != 1 || accessKeyAuth.seedCalls[0] != testAccessKey { + t.Errorf("expected SeedSession called with the access key, got %v", accessKeyAuth.seedCalls) + } + if len(accessKeyAuth.persistCalls) != 1 || accessKeyAuth.persistCalls[0].orgID != "org_123" { + t.Errorf("expected Persist called with existing org ID org_123, got %v", accessKeyAuth.persistCalls) + } +} - deps.prompter = mockConfirmer{confirm: true} +// Test_runRegister_AccessKey_PersistFailureIsSoft verifies a durable-save +// failure does not fail registration (the session was already seeded). +func Test_runRegister_AccessKey_PersistFailureIsSoft(t *testing.T) { + ensureNoAccessKeyEnv(t) - SetTestSSHPort(tt.port) - defer ClearTestSSHPort() + regStore := &mockRegistrationStore{} + store := &mockRegisterStore{ + user: &entity.User{ID: "user_1"}, + org: &entity.Organization{ID: "org_123", Name: "TestOrg"}, + token: "tok", + } + svc := &fakeNodeService{addNodeFn: accessKeyAddNodeFn(t)} - term := terminal.New() - opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg", sshPort: tt.port} - err := runRegister(context.Background(), term, store, opts, deps) + deps, server := testRegisterDeps(t, svc, regStore) + defer server.Close() + accessKeyAuth := &mockAccessKeyAuth{persistErr: fmt.Errorf("disk full")} + deps.accessKeyAuth = accessKeyAuth - tt.verify(t, gotOpenReq, gotGrantReq, regStore, err) - }) + term := terminal.New() + opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg", accessKey: testAccessKey} + if err := runRegister(context.Background(), term, store, opts, deps); err != nil { + t.Fatalf("registration should succeed even when durable save fails, got: %v", err) + } + if len(accessKeyAuth.seedCalls) != 1 { + t.Errorf("expected session to be seeded regardless, got %v", accessKeyAuth.seedCalls) } } diff --git a/pkg/cmd/revokessh/revokessh_test.go b/pkg/cmd/revokessh/revokessh_test.go index 447c53cb..9eec15d1 100644 --- a/pkg/cmd/revokessh/revokessh_test.go +++ b/pkg/cmd/revokessh/revokessh_test.go @@ -50,6 +50,13 @@ func (m *mockRegistrationStore) Load() (*register.DeviceRegistration, error) { return m.reg, nil } +func (m *mockRegistrationStore) LoadAny() (*register.DeviceRegistration, error) { + if m.reg == nil { + return nil, fmt.Errorf("no registration") + } + return m.reg, nil +} + func (m *mockRegistrationStore) Delete() error { m.reg = nil return nil From f4616e12c0d148194fa6291e080fa966babe726b Mon Sep 17 00:00:00 2001 From: Pratik Patel Date: Fri, 21 Aug 2026 09:44:11 -0700 Subject: [PATCH 3/6] simplify comments --- pkg/cmd/cmd.go | 4 - pkg/cmd/register/device_registration_store.go | 29 ++----- pkg/cmd/register/register.go | 84 +++++-------------- 3 files changed, 29 insertions(+), 88 deletions(-) diff --git a/pkg/cmd/cmd.go b/pkg/cmd/cmd.go index 00461a6d..7b0541de 100644 --- a/pkg/cmd/cmd.go +++ b/pkg/cmd/cmd.go @@ -265,10 +265,6 @@ func NewBrevCommand() *cobra.Command { //nolint:funlen,gocognit,gocyclo // defin fmt.Printf("%v\n", err) } - // accessKeyAuth seeds the in-memory session store (used by register's API - // calls) and durably persists a supplied --access-key via the file-backed - // loginAuth (the same path as `brev login --api-key`), so other brev - // commands authenticate with the key afterwards. accessKeyAuth := register.NewAccessKeyAuthenticator(memAuthStore.MemoryAuthStore, loginAuth) createCmdTree(cmds, t, loginCmdStore, noLoginCmdStore, loginAuth, externalNodeCmdStore, accessKeyAuth) diff --git a/pkg/cmd/register/device_registration_store.go b/pkg/cmd/register/device_registration_store.go index 0dd0235d..3c0daaf8 100644 --- a/pkg/cmd/register/device_registration_store.go +++ b/pkg/cmd/register/device_registration_store.go @@ -22,13 +22,10 @@ const ( // Registration status values written to DeviceRegistration.Status. const ( - // RegistrationStatusPending marks an intent record written before AddNode - // has been confirmed. It captures the chosen device ID (plus name/org and - // hardware) so a retry can resume with the same ID: AddNode is idempotent on - // device_id, so reusing it does not create a duplicate backend node. + // RegistrationStatusPending marks an intent record written before AddNode is + // confirmed; it lets a retry reuse the same device ID. RegistrationStatusPending = "pending" - // RegistrationStatusRegistered marks a fully confirmed registration: AddNode - // succeeded and the external node ID has been persisted locally. + // RegistrationStatusRegistered marks a confirmed registration (AddNode succeeded, external node ID persisted). RegistrationStatusRegistered = "registered" ) @@ -42,26 +39,18 @@ type DeviceRegistration struct { DeviceID string `json:"device_id"` RegisteredAt string `json:"registered_at"` HardwareProfile HardwareProfile `json:"hardware_profile"` - // RegistrationToken is a UI-supplied token captured at registration time and - // persisted so a retry reuses the same token. It is not yet sent to the - // backend (AddNode has no token field today); wired in for future use. + // UI-supplied token persisted for retry reuse; not yet sent to the backend (no AddNode token field yet). RegistrationToken string `json:"registration_token,omitempty"` - // Status is RegistrationStatusPending (intent written, AddNode not yet - // confirmed) or RegistrationStatusRegistered (confirmed). Empty is treated - // as registered for backwards compatibility with files written by older - // CLI versions that predate the Status field. + // Status is pending or registered; empty is treated as registered (backwards compat with older files). Status string `json:"status,omitempty"` } // RegistrationStore defines the contract for persisting device registration data. type RegistrationStore interface { Save(reg *DeviceRegistration) error - // Load returns a fully-confirmed registration (one with an external node - // ID). It errors for pending (in-progress) records. + // Load returns a confirmed registration only; it errors for pending records. Load() (*DeviceRegistration, error) - // LoadAny returns the registration record regardless of whether it has - // been confirmed, erroring only if the file is missing or unparseable. - // Used by the register flow to detect and resume pending attempts. + // LoadAny returns the record regardless of status (pending or confirmed); used to resume pending attempts. LoadAny() (*DeviceRegistration, error) Delete() error Exists() (bool, error) @@ -122,10 +111,6 @@ func (s *FileRegistrationStore) Load() (*DeviceRegistration, error) { return ®, nil } -// LoadAny reads the registration file and returns the parsed DeviceRegistration -// regardless of whether registration has been confirmed. It errors only if the -// file is missing or unparseable. Use Load() when a fully-confirmed registration -// is required (e.g., enable-ssh, deregister). func (s *FileRegistrationStore) LoadAny() (*DeviceRegistration, error) { path := s.path() exists, err := s.Exists() diff --git a/pkg/cmd/register/register.go b/pkg/cmd/register/register.go index a7e970f8..79484fc8 100644 --- a/pkg/cmd/register/register.go +++ b/pkg/cmd/register/register.go @@ -89,18 +89,13 @@ type AccessKeyAuthenticator interface { Persist(apiKey, orgID string) error } -// accessKeyAuthenticator is the production AccessKeyAuthenticator. It seeds the -// in-memory session store (used by external-node commands) and persists the -// key to the file-backed credential store (used by other brev commands). +// accessKeyAuthenticator seeds the in-memory session (for this command) and +// persists to the file-backed store (for other commands). type accessKeyAuthenticator struct { session *store.MemoryAuthStore durable *auth.LoginAuth } -// NewAccessKeyAuthenticator returns the production AccessKeyAuthenticator. -// session is the in-memory store backing the register command's auth (so the -// current command authenticates with the key); durable is the file-backed auth -// used by 'brev login --api-key' (so the key survives for other commands). func NewAccessKeyAuthenticator(session *store.MemoryAuthStore, durable *auth.LoginAuth) AccessKeyAuthenticator { return accessKeyAuthenticator{session: session, durable: durable} } @@ -117,21 +112,14 @@ func (a accessKeyAuthenticator) Persist(apiKey, orgID string) error { return nil } -// noopAccessKeyAuth is a no-op AccessKeyAuthenticator used as the default in -// defaultRegisterDeps (tests and any path that doesn't supply an access key). -// Production wiring overrides it via NewCmdRegister. +// noopAccessKeyAuth is the default for tests/no-key paths; production injects a real one via NewCmdRegister. type noopAccessKeyAuth struct{} func (noopAccessKeyAuth) SeedSession(string) error { return nil } func (noopAccessKeyAuth) Persist(string, string) error { return nil } -// accessKeyEnvVar is the environment variable consulted when --access-key is -// not set, enabling headless/CI registration without modifying the command. const accessKeyEnvVar = "BREV_ACCESS_KEY" -// resolveAccessKey returns the access key from the --access-key flag, falling -// back to the BREV_ACCESS_KEY environment variable. Returns "" when neither is -// set, in which case the caller falls back to the interactive login-link flow. func resolveAccessKey(flagValue string) string { if v := strings.TrimSpace(flagValue); v != "" { return v @@ -139,10 +127,8 @@ func resolveAccessKey(flagValue string) string { return strings.TrimSpace(os.Getenv(accessKeyEnvVar)) } -// persistAccessKey durably saves the access key for other brev commands. It is -// best-effort: the in-memory session was already seeded, so the current command -// authenticates with the key regardless; a failure here only means future -// commands won't pick it up automatically. +// persistAccessKey durably saves the access key. Best-effort: the session is +// already seeded, so a failure only affects future commands, not this one. func persistAccessKey(t *terminal.Terminal, deps registerDeps, apiKey, orgID string) { if err := deps.accessKeyAuth.Persist(apiKey, orgID); err != nil { t.Vprintf(" %s\n", t.Yellow(fmt.Sprintf("Warning: failed to save access key for future commands: %v", err))) @@ -214,11 +200,7 @@ func NewCmdRegister(t *terminal.Terminal, store RegisterStore, accessKeyAuth Acc cmd.Flags().StringVar(®istrationToken, "registration-token", "", "UI-supplied registration token (used when registering via the UI flow)") cmd.Flags().StringVar(&accessKey, "access-key", "", "durable headless auth key (Brev API key); falls back to the "+accessKeyEnvVar+" env var, then the login-link flow") - // --ssh-port is deprecated: registration no longer enables SSH. The flag is - // kept (and still accepted) for backwards compatibility with existing - // scripts. Use 'brev enable-ssh' after registration to enable SSH access. - // MarkDeprecated hides the flag from help and prints the deprecation notice - // when either --ssh-port or its -p shorthand is supplied. + // --ssh-port is deprecated (registration no longer enables SSH) but kept for backwards compatibility. _ = cmd.Flags().MarkDeprecated("ssh-port", "use 'brev enable-ssh' after registration to enable SSH access") return cmd @@ -251,12 +233,8 @@ func runRegister(ctx context.Context, t *terminal.Terminal, s RegisterStore, opt } } - // Credential chain: --access-key flag -> BREV_ACCESS_KEY env -> interactive - // login-link flow (handled by the store's auth when no key is supplied). - // When an access key is present it is seeded into the in-memory session auth - // now, before any API call (e.g. GetCurrentUser, org resolution), so those - // calls authenticate with it. The org ID isn't known yet, so the durable, - // org-scoped save happens after the org is resolved (see persistAccessKey). + // Seed the access key before any API call; the durable (org-scoped) save is + // deferred until the org is resolved below. accessKey := resolveAccessKey(opts.accessKey) if accessKey != "" { if !auth.IsBrevAPIKey(accessKey) { @@ -283,16 +261,12 @@ func runRegister(ctx context.Context, t *terminal.Terminal, s RegisterStore, opt if err != nil { return breverrors.WrapAndTrace(err) } - // The existing record (pending or confirmed) already carries the org ID, - // so durably persist the access key here if one was supplied. + // The existing record already carries the org ID, so we can persist the access key now. if accessKey != "" { persistAccessKey(t, deps, accessKey, reg.OrgID) } if reg.Status == RegistrationStatusPending { - // An earlier attempt wrote an intent record but never confirmed. - // Resume it, reusing the same device ID so AddNode is idempotent - // (a repeat call returns the existing node rather than a duplicate) - // and rectifying backend state into local CLI state. + // Resume: reuse the same device ID (AddNode is idempotent) to rectify backend state. return resumeRegistration(ctx, t, s, deps, reg) } return checkExistingRegistration(ctx, t, s, deps) @@ -327,8 +301,7 @@ func runRegister(ctx context.Context, t *terminal.Terminal, s RegisterStore, opt return err } - // Org is resolved now, so durably persist the access key (org-scoped) for - // other brev commands, mirroring 'brev login --api-key'. + // Org is resolved; persist the access key for other commands (mirrors 'brev login --api-key'). if accessKey != "" { persistAccessKey(t, deps, accessKey, org.ID) } @@ -360,10 +333,7 @@ func runRegister(ctx context.Context, t *terminal.Terminal, s RegisterStore, opt } } - // Perform the registration steps. A stable device ID is generated here so a - // retry (which loads the pending record written by runRegisterSteps) reuses - // the same ID: AddNode is idempotent on device_id, so this does not create a - // duplicate backend node. + // Generate the device ID here so a retry reuses it (AddNode is idempotent on device_id). deviceID := uuid.New().String() reg, err := runRegisterSteps(ctx, t, s, name, org, deps, deviceID, opts.registrationToken) if err != nil { @@ -374,17 +344,10 @@ func runRegister(ctx context.Context, t *terminal.Terminal, s RegisterStore, opt return nil } -// runRegisterSteps performs the registration steps: install the Brev tunnel, -// collect a hardware profile, write a pending intent record, call AddNode, -// persist the confirmed registration, and run the local setup command. -// -// deviceID is reused across retries so AddNode is idempotent: a second call with -// the same device_id returns the existing node rather than creating a duplicate. -// The pending intent record is written before AddNode so that a crash or -// timeout after the node was created backend-side can be resumed (see -// resumeRegistration) without orphaning the backend node. -// -// It does not prompt or enable SSH. Used by both fresh and resume flows. +// runRegisterSteps runs the registration steps (tunnel, hardware profile, +// AddNode, persist, setup). It writes a pending intent record before AddNode so +// a crash/timeout after the node exists can be resumed (see resumeRegistration) +// without orphaning the backend node; deviceID reuse makes AddNode idempotent. func runRegisterSteps(ctx context.Context, t *terminal.Terminal, s RegisterStore, name string, org *entity.Organization, deps registerDeps, deviceID, registrationToken string) (*DeviceRegistration, error) { t.Vprint("") @@ -409,8 +372,7 @@ func runRegisterSteps(ctx context.Context, t *terminal.Terminal, s RegisterStore t.Vprint("") t.Vprint(t.Yellow("[Step 3/5] Registering device with Brev...")) - // Write a pending intent record before AddNode so a crash/timeout can be - // resumed with the same device ID (AddNode is idempotent on device_id). + // Write pending intent before AddNode (see resumeRegistration). pending := &DeviceRegistration{ DisplayName: name, OrgID: org.ID, @@ -566,11 +528,10 @@ func runSetup(node *nodev1.ExternalNode, t *terminal.Terminal, deps registerDeps } } -// resumeRegistration resumes an in-progress (pending) registration, reusing the -// same device ID so AddNode is idempotent. This rectifies backend state (the -// node created by a previous, partially-successful attempt) into local CLI -// state: a repeat AddNode with the same device ID returns the existing node, -// which is then persisted locally as confirmed. +// resumeRegistration resumes a pending registration, reusing the same device ID. +// AddNode is idempotent on device_id, so this recovers when AddNode succeeded +// backend-side but the CLI never confirmed: the repeat call returns the existing +// node, which is then persisted locally. func resumeRegistration(ctx context.Context, t *terminal.Terminal, s RegisterStore, deps registerDeps, pending *DeviceRegistration) error { t.Vprint("") t.Vprint(t.White("══════════════════════════════════════════════════")) @@ -594,8 +555,7 @@ func resumeRegistration(ctx context.Context, t *terminal.Terminal, s RegisterSto return nil } -// suggestEnableSSH prints a hint that SSH access can be enabled separately now -// that registration is complete. Registration no longer enables SSH itself. +// suggestEnableSSH prints a hint to run 'brev enable-ssh'; registration no longer enables SSH itself. func suggestEnableSSH(t *terminal.Terminal, reg *DeviceRegistration) { t.Vprint("") if reg != nil && reg.DisplayName != "" { From d75c89cb9f5af52c110aeee37b3d992ddf990a99 Mon Sep 17 00:00:00 2001 From: Pratik Patel Date: Fri, 21 Aug 2026 12:24:50 -0700 Subject: [PATCH 4/6] clean --- pkg/cmd/deregister/deregister.go | 2 +- pkg/cmd/deregister/deregister_test.go | 9 +-- pkg/cmd/enablessh/enablessh.go | 2 +- pkg/cmd/grantssh/grantssh_test.go | 9 +-- pkg/cmd/register/device_registration_store.go | 62 ++++++------------- .../device_registration_store_test.go | 30 ++++----- pkg/cmd/register/register.go | 14 ++--- pkg/cmd/register/register_test.go | 52 +++------------- pkg/cmd/register/sshkeys.go | 2 +- pkg/cmd/revokessh/revokessh_test.go | 9 +-- 10 files changed, 51 insertions(+), 140 deletions(-) diff --git a/pkg/cmd/deregister/deregister.go b/pkg/cmd/deregister/deregister.go index efd9090b..fec58249 100644 --- a/pkg/cmd/deregister/deregister.go +++ b/pkg/cmd/deregister/deregister.go @@ -106,7 +106,7 @@ func runDeregister(ctx context.Context, t *terminal.Terminal, s DeregisterStore, return fmt.Errorf("sudo issue: %w", err) } - reg, err := deps.registrationStore.Load() + reg, err := deps.registrationStore.Load(true) // deregister should still work for pending registrations if err != nil { return err //nolint:wrapcheck // do not present stack trace for this error } diff --git a/pkg/cmd/deregister/deregister_test.go b/pkg/cmd/deregister/deregister_test.go index ea96a929..0e702e39 100644 --- a/pkg/cmd/deregister/deregister_test.go +++ b/pkg/cmd/deregister/deregister_test.go @@ -57,14 +57,7 @@ func (m *mockRegistrationStore) Save(reg *register.DeviceRegistration) error { return nil } -func (m *mockRegistrationStore) Load() (*register.DeviceRegistration, error) { - if m.reg == nil { - return nil, fmt.Errorf("no registration") - } - return m.reg, nil -} - -func (m *mockRegistrationStore) LoadAny() (*register.DeviceRegistration, error) { +func (m *mockRegistrationStore) Load(bool) (*register.DeviceRegistration, error) { if m.reg == nil { return nil, fmt.Errorf("no registration") } diff --git a/pkg/cmd/enablessh/enablessh.go b/pkg/cmd/enablessh/enablessh.go index 9788b0e6..6b6ac548 100644 --- a/pkg/cmd/enablessh/enablessh.go +++ b/pkg/cmd/enablessh/enablessh.go @@ -66,7 +66,7 @@ func runEnableSSH(ctx context.Context, t *terminal.Terminal, s EnableSSHStore, d return fmt.Errorf("brev enable-ssh is only supported on Linux") } - reg, err := deps.registrationStore.Load() + reg, err := deps.registrationStore.Load(false) if err != nil { return fmt.Errorf("failed to read registration file: %w", err) } diff --git a/pkg/cmd/grantssh/grantssh_test.go b/pkg/cmd/grantssh/grantssh_test.go index ef4dc367..e1e6441a 100644 --- a/pkg/cmd/grantssh/grantssh_test.go +++ b/pkg/cmd/grantssh/grantssh_test.go @@ -44,14 +44,7 @@ func (m *mockRegistrationStore) Save(reg *register.DeviceRegistration) error { return nil } -func (m *mockRegistrationStore) Load() (*register.DeviceRegistration, error) { - if m.reg == nil { - return nil, fmt.Errorf("no registration") - } - return m.reg, nil -} - -func (m *mockRegistrationStore) LoadAny() (*register.DeviceRegistration, error) { +func (m *mockRegistrationStore) Load(bool) (*register.DeviceRegistration, error) { if m.reg == nil { return nil, fmt.Errorf("no registration") } diff --git a/pkg/cmd/register/device_registration_store.go b/pkg/cmd/register/device_registration_store.go index 3c0daaf8..0838b098 100644 --- a/pkg/cmd/register/device_registration_store.go +++ b/pkg/cmd/register/device_registration_store.go @@ -20,47 +20,35 @@ const ( globalRegistrationDir = "/etc/brev" ) -// Registration status values written to DeviceRegistration.Status. const ( - // RegistrationStatusPending marks an intent record written before AddNode is - // confirmed; it lets a retry reuse the same device ID. - RegistrationStatusPending = "pending" - // RegistrationStatusRegistered marks a confirmed registration (AddNode succeeded, external node ID persisted). + RegistrationStatusPending = "pending" // used for retries RegistrationStatusRegistered = "registered" ) // DeviceRegistration is the persistent identity file for a registered device. // Fields align with the AddNodeResponse from dev-plane. type DeviceRegistration struct { - ExternalNodeID string `json:"external_node_id"` - DisplayName string `json:"display_name"` - OrgID string `json:"org_id"` - OrgName string `json:"org_name"` - DeviceID string `json:"device_id"` - RegisteredAt string `json:"registered_at"` - HardwareProfile HardwareProfile `json:"hardware_profile"` - // UI-supplied token persisted for retry reuse; not yet sent to the backend (no AddNode token field yet). - RegistrationToken string `json:"registration_token,omitempty"` - // Status is pending or registered; empty is treated as registered (backwards compat with older files). - Status string `json:"status,omitempty"` + ExternalNodeID string `json:"external_node_id"` + DisplayName string `json:"display_name"` + OrgID string `json:"org_id"` + OrgName string `json:"org_name"` + DeviceID string `json:"device_id"` + RegisteredAt string `json:"registered_at"` + HardwareProfile HardwareProfile `json:"hardware_profile"` + RegistrationToken string `json:"registration_token,omitempty"` + Status string `json:"status,omitempty"` } // RegistrationStore defines the contract for persisting device registration data. type RegistrationStore interface { Save(reg *DeviceRegistration) error - // Load returns a confirmed registration only; it errors for pending records. - Load() (*DeviceRegistration, error) - // LoadAny returns the record regardless of status (pending or confirmed); used to resume pending attempts. - LoadAny() (*DeviceRegistration, error) + Load(includeAll bool) (*DeviceRegistration, error) Delete() error Exists() (bool, error) } -// FileRegistrationStore implements RegistrationStore using the global /etc/brev/ path. type FileRegistrationStore struct{} -// NewFileRegistrationStore returns a FileRegistrationStore that reads/writes -// from /etc/brev/device_registration.json. func NewFileRegistrationStore() *FileRegistrationStore { return &FileRegistrationStore{} } @@ -88,8 +76,7 @@ func (s *FileRegistrationStore) Save(reg *DeviceRegistration) error { return sudoWriteFile(path, data) } -// Load reads the registration file and returns the parsed DeviceRegistration -func (s *FileRegistrationStore) Load() (*DeviceRegistration, error) { +func (s *FileRegistrationStore) Load(includeAll bool) (*DeviceRegistration, error) { path := s.path() exists, err := s.Exists() if !exists { @@ -102,6 +89,12 @@ func (s *FileRegistrationStore) Load() (*DeviceRegistration, error) { if err := files.ReadJSON(files.AppFs, path, ®); err != nil { return nil, breverrors.WrapAndTrace(err) } + if includeAll { + if reg.OrgID == "" && reg.DeviceID == "" { + return nil, breverrors.New("malformed registration") + } + return ®, nil + } if reg.ExternalNodeID == "" || reg.OrgID == "" { if reg.Status == RegistrationStatusPending { return nil, breverrors.New("device registration is incomplete; re-run 'brev register' to finish") @@ -111,25 +104,6 @@ func (s *FileRegistrationStore) Load() (*DeviceRegistration, error) { return ®, nil } -func (s *FileRegistrationStore) LoadAny() (*DeviceRegistration, error) { - path := s.path() - exists, err := s.Exists() - if err != nil { - return nil, breverrors.WrapAndTrace(err) - } - if !exists { - return nil, breverrors.New("device registration not found, run 'brev register' first") - } - var reg DeviceRegistration - if err := files.ReadJSON(files.AppFs, path, ®); err != nil { - return nil, breverrors.WrapAndTrace(err) - } - if reg.OrgID == "" && reg.DeviceID == "" { - return nil, breverrors.New("malformed registration") - } - return ®, nil -} - func (s *FileRegistrationStore) Delete() error { path := s.path() err := files.DeleteFile(files.AppFs, path) diff --git a/pkg/cmd/register/device_registration_store_test.go b/pkg/cmd/register/device_registration_store_test.go index e8622568..c5e53fb5 100644 --- a/pkg/cmd/register/device_registration_store_test.go +++ b/pkg/cmd/register/device_registration_store_test.go @@ -43,7 +43,7 @@ func Test_SaveAndLoadRegistration_RoundTrip(t *testing.T) { t.Fatalf("Save failed: %v", err) } - loaded, err := store.Load() + loaded, err := store.Load(false) if err != nil { t.Fatalf("Load failed: %v", err) } @@ -139,7 +139,7 @@ func Test_LoadRegistration_FailsWhenMissing(t *testing.T) { store := NewFileRegistrationStore() - _, err := store.Load() + _, err := store.Load(false) if err == nil { t.Error("expected error loading missing registration") } @@ -160,7 +160,7 @@ func Test_LoadRegistration_RejectsMissingExternalNodeID(t *testing.T) { t.Fatalf("Save failed: %v", err) } - _, err := store.Load() + _, err := store.Load(false) if err == nil { t.Fatal("expected error loading registration with empty ExternalNodeID") } @@ -181,7 +181,7 @@ func Test_LoadRegistration_RejectsMissingOrgID(t *testing.T) { t.Fatalf("Save failed: %v", err) } - _, err := store.Load() + _, err := store.Load(false) if err == nil { t.Fatal("expected error loading registration with empty OrgID") } @@ -199,9 +199,9 @@ func Test_DeleteRegistration_FailsWhenMissing(t *testing.T) { } } -// Test_LoadAny_ReturnsPendingRecord verifies LoadAny returns a pending (intent) -// record that has no external node ID, which the strict Load() rejects. -func Test_LoadAny_ReturnsPendingRecord(t *testing.T) { +// Test_Load_IncludeAllReturnsPendingRecord verifies Load(true) returns a pending +// (intent) record that has no external node ID, which the strict Load(false) rejects. +func Test_Load_IncludeAllReturnsPendingRecord(t *testing.T) { cleanup := setupTestFs(t) defer cleanup() @@ -217,9 +217,9 @@ func Test_LoadAny_ReturnsPendingRecord(t *testing.T) { t.Fatalf("Save failed: %v", err) } - loaded, err := store.LoadAny() + loaded, err := store.Load(true) if err != nil { - t.Fatalf("LoadAny failed: %v", err) + t.Fatalf("Load failed: %v", err) } if loaded.DeviceID != "device-uuid-123" { t.Errorf("DeviceID mismatch: got %s, want device-uuid-123", loaded.DeviceID) @@ -231,14 +231,14 @@ func Test_LoadAny_ReturnsPendingRecord(t *testing.T) { t.Errorf("pending record should have no ExternalNodeID, got %q", loaded.ExternalNodeID) } - // The strict Load() must reject the pending record. - if _, err := store.Load(); err == nil { - t.Error("expected Load() to error on a pending record") + // The strict Load(false) must reject the pending record. + if _, err := store.Load(false); err == nil { + t.Error("expected Load(false) to error on a pending record") } } -// Test_Load_PendingRecordErrorMessage verifies Load() reports a clear, actionable -// message for a pending (in-progress) registration. +// Test_Load_PendingRecordErrorMessage verifies Load(false) reports a clear, +// actionable message for a pending (in-progress) registration. func Test_Load_PendingRecordErrorMessage(t *testing.T) { cleanup := setupTestFs(t) defer cleanup() @@ -255,7 +255,7 @@ func Test_Load_PendingRecordErrorMessage(t *testing.T) { t.Fatalf("Save failed: %v", err) } - _, err := store.Load() + _, err := store.Load(false) if err == nil { t.Fatal("expected Load() to error on a pending record") } diff --git a/pkg/cmd/register/register.go b/pkg/cmd/register/register.go index 79484fc8..9c57c514 100644 --- a/pkg/cmd/register/register.go +++ b/pkg/cmd/register/register.go @@ -89,8 +89,6 @@ type AccessKeyAuthenticator interface { Persist(apiKey, orgID string) error } -// accessKeyAuthenticator seeds the in-memory session (for this command) and -// persists to the file-backed store (for other commands). type accessKeyAuthenticator struct { session *store.MemoryAuthStore durable *auth.LoginAuth @@ -112,7 +110,6 @@ func (a accessKeyAuthenticator) Persist(apiKey, orgID string) error { return nil } -// noopAccessKeyAuth is the default for tests/no-key paths; production injects a real one via NewCmdRegister. type noopAccessKeyAuth struct{} func (noopAccessKeyAuth) SeedSession(string) error { return nil } @@ -127,10 +124,9 @@ func resolveAccessKey(flagValue string) string { return strings.TrimSpace(os.Getenv(accessKeyEnvVar)) } -// persistAccessKey durably saves the access key. Best-effort: the session is -// already seeded, so a failure only affects future commands, not this one. func persistAccessKey(t *terminal.Terminal, deps registerDeps, apiKey, orgID string) { if err := deps.accessKeyAuth.Persist(apiKey, orgID); err != nil { + // Best-effort: the session is already seeded, so a failure only affects future commands, not this one. t.Vprintf(" %s\n", t.Yellow(fmt.Sprintf("Warning: failed to save access key for future commands: %v", err))) } } @@ -198,9 +194,7 @@ func NewCmdRegister(t *terminal.Terminal, store RegisterStore, accessKeyAuth Acc cmd.Flags().IntVarP(&sshPort, "ssh-port", "p", 0, "SSH port (if ssh access is desired)") cmd.Flags().BoolVar(&approveFlag, "approve", false, "skip all confirmation prompts (assume yes)") cmd.Flags().StringVar(®istrationToken, "registration-token", "", "UI-supplied registration token (used when registering via the UI flow)") - cmd.Flags().StringVar(&accessKey, "access-key", "", "durable headless auth key (Brev API key); falls back to the "+accessKeyEnvVar+" env var, then the login-link flow") - - // --ssh-port is deprecated (registration no longer enables SSH) but kept for backwards compatibility. + cmd.Flags().StringVar(&accessKey, "access-key", "", "API key; falls back to the "+accessKeyEnvVar+" env var, then the login-link flow") _ = cmd.Flags().MarkDeprecated("ssh-port", "use 'brev enable-ssh' after registration to enable SSH access") return cmd @@ -257,7 +251,7 @@ func runRegister(ctx context.Context, t *terminal.Terminal, s RegisterStore, opt return breverrors.WrapAndTrace(err) } if exists { - reg, err := deps.registrationStore.LoadAny() + reg, err := deps.registrationStore.Load(true) if err != nil { return breverrors.WrapAndTrace(err) } @@ -457,7 +451,7 @@ func resolveOrg(s RegisterStore, orgName string) (*entity.Organization, error) { // local netbird service is running, starting it if necessary. Returns nil if // the node is healthy, or an error describing what's wrong. func checkExistingRegistration(ctx context.Context, t *terminal.Terminal, s RegisterStore, deps registerDeps) error { - reg, loadErr := deps.registrationStore.Load() + reg, loadErr := deps.registrationStore.Load(false) if loadErr != nil { return fmt.Errorf("this machine is already registered but the registration file could not be read: %w", loadErr) } diff --git a/pkg/cmd/register/register_test.go b/pkg/cmd/register/register_test.go index b9cdaafe..6fba5a73 100644 --- a/pkg/cmd/register/register_test.go +++ b/pkg/cmd/register/register_test.go @@ -75,14 +75,7 @@ func (m *mockRegistrationStore) Save(reg *DeviceRegistration) error { return nil } -func (m *mockRegistrationStore) Load() (*DeviceRegistration, error) { - if m.reg == nil { - return nil, fmt.Errorf("no registration") - } - return m.reg, nil -} - -func (m *mockRegistrationStore) LoadAny() (*DeviceRegistration, error) { +func (m *mockRegistrationStore) Load(bool) (*DeviceRegistration, error) { if m.reg == nil { return nil, fmt.Errorf("no registration") } @@ -276,7 +269,7 @@ func Test_runRegister_HappyPath(t *testing.T) { t.Fatal("expected registration to exist after successful register") } - reg, err := regStore.Load() + reg, err := regStore.Load(false) if err != nil { t.Fatalf("Load failed: %v", err) } @@ -496,7 +489,7 @@ func Test_runRegister_WithOrgFlag(t *testing.T) { t.Errorf("expected org_456, got %s", capturedOrgID) } - reg, err := regStore.Load() + reg, err := regStore.Load(false) if err != nil { t.Fatalf("Load failed: %v", err) } @@ -559,7 +552,7 @@ func Test_runRegister_AddNodeFails(t *testing.T) { // AddNode failed, but a pending intent record is written before AddNode so a // retry can resume with the same device ID. The record must be pending (no // external node ID yet) and carry a device ID for the retry to reuse. - reg, loadErr := regStore.LoadAny() + reg, loadErr := regStore.Load(true) if loadErr != nil { t.Fatalf("expected pending intent record after AddNode failure, got: %v", loadErr) } @@ -913,11 +906,6 @@ func Test_runRegister_NoNameAlreadyRegistered(t *testing.T) { } } -// Test_runRegister_ResumesPendingRegistration verifies that a pending (intent) -// record left by an interrupted attempt is resumed: the same device ID is reused -// (AddNode is idempotent on device_id), and the final record is marked -// registered with the backend-assigned external node ID. The captured -// registration token is preserved across the resume. func Test_runRegister_ResumesPendingRegistration(t *testing.T) { const pendingDeviceID = "device-uuid-pending" pending := &DeviceRegistration{ @@ -975,9 +963,9 @@ func Test_runRegister_ResumesPendingRegistration(t *testing.T) { t.Errorf("expected AddNode to reuse device ID %q, got %v", pendingDeviceID, addNodeDeviceIDs) } - reg, loadErr := regStore.LoadAny() + reg, loadErr := regStore.Load(true) if loadErr != nil { - t.Fatalf("LoadAny failed: %v", loadErr) + t.Fatalf("Load failed: %v", loadErr) } if reg.Status != RegistrationStatusRegistered { t.Errorf("expected status %q after resume, got %q", RegistrationStatusRegistered, reg.Status) @@ -993,9 +981,6 @@ func Test_runRegister_ResumesPendingRegistration(t *testing.T) { } } -// Test_runRegister_PersistsRegistrationToken verifies the --registration-token -// flag value is threaded through registration and persisted to local state (it -// is not yet sent to the backend). func Test_runRegister_PersistsRegistrationToken(t *testing.T) { regStore := &mockRegistrationStore{} store := &mockRegisterStore{ @@ -1026,9 +1011,9 @@ func Test_runRegister_PersistsRegistrationToken(t *testing.T) { t.Fatalf("runRegister failed: %v", err) } - reg, loadErr := regStore.LoadAny() + reg, loadErr := regStore.Load(true) if loadErr != nil { - t.Fatalf("LoadAny failed: %v", loadErr) + t.Fatalf("Load failed: %v", loadErr) } if reg.RegistrationToken != "ui-token-xyz" { t.Errorf("expected RegistrationToken to be persisted, got %q", reg.RegistrationToken) @@ -1038,12 +1023,8 @@ func Test_runRegister_PersistsRegistrationToken(t *testing.T) { } } -// --- Access key credential chain --- - const testAccessKey = auth.BrevAPIKeyPrefix + "test-key" -// accessKeyAddNodeFn returns an AddNode handler that echoes the request back as -// a registered node, suitable for the access-key happy-path tests. func accessKeyAddNodeFn(t *testing.T) func(*nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) { t.Helper() return func(req *nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) { @@ -1061,14 +1042,11 @@ func accessKeyAddNodeFn(t *testing.T) func(*nodev1.AddNodeRequest) (*nodev1.AddN } } -// ensureNoAccessKeyEnv unsets BREV_ACCESS_KEY for the test so the env fallback -// doesn't leak in (or out) of the credential-chain tests. func ensureNoAccessKeyEnv(t *testing.T) { t.Helper() t.Setenv(accessKeyEnvVar, "") } -// Test_resolveAccessKey verifies the flag -> env -> empty fallback order. func Test_resolveAccessKey(t *testing.T) { ensureNoAccessKeyEnv(t) @@ -1092,9 +1070,6 @@ func Test_resolveAccessKey(t *testing.T) { } } -// Test_runRegister_AccessKeyFlag_SeedsAndPersists verifies that --access-key -// seeds the session (before any API call) and is durably persisted with the -// resolved org ID once the org is known, then registration proceeds. func Test_runRegister_AccessKeyFlag_SeedsAndPersists(t *testing.T) { ensureNoAccessKeyEnv(t) @@ -1131,8 +1106,6 @@ func Test_runRegister_AccessKeyFlag_SeedsAndPersists(t *testing.T) { } } -// Test_runRegister_AccessKeyEnv_Fallback verifies the BREV_ACCESS_KEY env var is -// used when --access-key is not supplied. func Test_runRegister_AccessKeyEnv_Fallback(t *testing.T) { ensureNoAccessKeyEnv(t) t.Setenv(accessKeyEnvVar, testAccessKey) @@ -1165,8 +1138,6 @@ func Test_runRegister_AccessKeyEnv_Fallback(t *testing.T) { } } -// Test_runRegister_AccessKeyFlag_PrecedenceOverEnv verifies the flag wins when -// both --access-key and BREV_ACCESS_KEY are set. func Test_runRegister_AccessKeyFlag_PrecedenceOverEnv(t *testing.T) { ensureNoAccessKeyEnv(t) t.Setenv(accessKeyEnvVar, auth.BrevAPIKeyPrefix+"env-key") @@ -1195,8 +1166,6 @@ func Test_runRegister_AccessKeyFlag_PrecedenceOverEnv(t *testing.T) { } } -// Test_runRegister_AccessKeyInvalid verifies a non-API-key value is rejected -// before any registration side effect. func Test_runRegister_AccessKeyInvalid(t *testing.T) { ensureNoAccessKeyEnv(t) @@ -1231,9 +1200,6 @@ func Test_runRegister_AccessKeyInvalid(t *testing.T) { } } -// Test_runRegister_AccessKey_PersistsOnAlreadyRegistered verifies that when an -// access key is supplied against an already-registered machine, it is still -// persisted (using the existing record's org ID) for other brev commands. func Test_runRegister_AccessKey_PersistsOnAlreadyRegistered(t *testing.T) { ensureNoAccessKeyEnv(t) @@ -1282,8 +1248,6 @@ func Test_runRegister_AccessKey_PersistsOnAlreadyRegistered(t *testing.T) { } } -// Test_runRegister_AccessKey_PersistFailureIsSoft verifies a durable-save -// failure does not fail registration (the session was already seeded). func Test_runRegister_AccessKey_PersistFailureIsSoft(t *testing.T) { ensureNoAccessKeyEnv(t) diff --git a/pkg/cmd/register/sshkeys.go b/pkg/cmd/register/sshkeys.go index 1188766d..2ca9bb91 100644 --- a/pkg/cmd/register/sshkeys.go +++ b/pkg/cmd/register/sshkeys.go @@ -30,7 +30,7 @@ func SelectNodeFromList(ctx context.Context, t *terminal.Terminal, prompter term return nil, fmt.Errorf("no nodes found in organization") } var thisNodeID string - if reg, err := registrationStore.Load(); err == nil && reg != nil { + if reg, err := registrationStore.Load(false); err == nil && reg != nil { thisNodeID = reg.ExternalNodeID } t.Vprint("") diff --git a/pkg/cmd/revokessh/revokessh_test.go b/pkg/cmd/revokessh/revokessh_test.go index 9eec15d1..71165928 100644 --- a/pkg/cmd/revokessh/revokessh_test.go +++ b/pkg/cmd/revokessh/revokessh_test.go @@ -43,14 +43,7 @@ func (m *mockRegistrationStore) Save(reg *register.DeviceRegistration) error { return nil } -func (m *mockRegistrationStore) Load() (*register.DeviceRegistration, error) { - if m.reg == nil { - return nil, fmt.Errorf("no registration") - } - return m.reg, nil -} - -func (m *mockRegistrationStore) LoadAny() (*register.DeviceRegistration, error) { +func (m *mockRegistrationStore) Load(bool) (*register.DeviceRegistration, error) { if m.reg == nil { return nil, fmt.Errorf("no registration") } From 11630ffae72c3435151a3b81397258a30359b414 Mon Sep 17 00:00:00 2001 From: Pratik Patel Date: Fri, 21 Aug 2026 13:21:25 -0700 Subject: [PATCH 5/6] idempotency --- pkg/cmd/deregister/deregister.go | 31 +++++++-- pkg/cmd/deregister/deregister_test.go | 92 +++++++++++++++++++++++++++ 2 files changed, 116 insertions(+), 7 deletions(-) diff --git a/pkg/cmd/deregister/deregister.go b/pkg/cmd/deregister/deregister.go index fec58249..3e6db2d9 100644 --- a/pkg/cmd/deregister/deregister.go +++ b/pkg/cmd/deregister/deregister.go @@ -3,6 +3,7 @@ package deregister import ( "context" + "errors" "fmt" "os/user" @@ -158,14 +159,30 @@ func runDeregister(ctx context.Context, t *terminal.Terminal, s DeregisterStore, } t.Vprint(t.Yellow("[Step 1/4] Removing node from Brev...")) - client := deps.nodeClients.NewNodeClient(s, config.GlobalConfig.GetBrevPublicAPIURL()) - _, err = client.RemoveNode(ctx, connect.NewRequest(&nodev1.RemoveNodeRequest{ - ExternalNodeId: reg.ExternalNodeID, - })) - if err != nil { - return fmt.Errorf("failed to deregister node: %w", err) + if reg.ExternalNodeID == "" { + // Pending registration: AddNode was never confirmed, so the CLI has no + // external node ID to remove. Skip RemoveNode and clean up local state + // so the user isn't stuck with an undeletable pending record. + t.Vprintf(" %s\n", t.Yellow("No registered node to remove (pending registration); cleaning up local state.")) + } else { + client := deps.nodeClients.NewNodeClient(s, config.GlobalConfig.GetBrevPublicAPIURL()) + _, err = client.RemoveNode(ctx, connect.NewRequest(&nodev1.RemoveNodeRequest{ + ExternalNodeId: reg.ExternalNodeID, + })) + if err != nil { + // NotFound means the node is already gone (e.g. a previous attempt + // succeeded before a transient error). Treat as success so local + // cleanup proceeds and a retry isn't blocked. + var connectErr *connect.Error + if errors.As(err, &connectErr) && connectErr.Code() == connect.CodeNotFound { + t.Vprintf(" %s\n", t.Yellow("Node not found on Brev (already removed); continuing.")) + } else { + return fmt.Errorf("failed to deregister node: %w", err) + } + } else { + t.Vprintf("%s Node removed from Brev.\n", t.Green(" ✓")) + } } - t.Vprintf("%s Node removed from Brev.\n", t.Green(" ✓")) t.Vprint("") t.Vprint(t.Yellow("[Step 2/4] Removing Brev SSH keys...")) diff --git a/pkg/cmd/deregister/deregister_test.go b/pkg/cmd/deregister/deregister_test.go index 0e702e39..fdf2c223 100644 --- a/pkg/cmd/deregister/deregister_test.go +++ b/pkg/cmd/deregister/deregister_test.go @@ -291,6 +291,98 @@ func Test_runDeregister_RemoveNodeFails(t *testing.T) { } } +// Test_runDeregister_RemoveNodeNotFound_ProceedsCleanup verifies that a NotFound +// from RemoveNode is treated as success: the node is already gone (e.g. a +// previous attempt succeeded before a transient error), so local cleanup must +// proceed and not leave the user stuck with an undeletable registration file. +func Test_runDeregister_RemoveNodeNotFound_ProceedsCleanup(t *testing.T) { + regStore := &mockRegistrationStore{ + reg: ®ister.DeviceRegistration{ + ExternalNodeID: "unode_abc", + DisplayName: "My Spark", + OrgID: "org_123", + }, + } + + store := &mockDeregisterStore{ + user: &entity.User{ID: "user_1"}, + token: "tok", + } + + svc := &fakeNodeService{ + removeNodeFn: func(_ *nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error) { + return nil, connect.NewError(connect.CodeNotFound, nil) + }, + } + + deps, server := testDeregisterDeps(t, svc, regStore) + defer server.Close() + + term := terminal.New() + err := runDeregister(context.Background(), term, store, deps, false) + if err != nil { + t.Fatalf("NotFound should be treated as success (node already gone), got: %v", err) + } + + exists, err := regStore.Exists() + if err != nil { + t.Fatalf("Exists error: %v", err) + } + if exists { + t.Error("expected local registration to be deleted even when RemoveNode returns NotFound") + } +} + +// Test_runDeregister_PendingRegistration_SkipsRemoveNodeAndCleansUp verifies that +// deregistering a partial (pending) registration — which has no ExternalNodeID — +// skips the RemoveNode call (there's no known backend node ID to remove) and +// still cleans up local state, so the user isn't stuck with an undeletable +// pending record. +func Test_runDeregister_PendingRegistration_SkipsRemoveNodeAndCleansUp(t *testing.T) { + regStore := &mockRegistrationStore{ + reg: ®ister.DeviceRegistration{ + DisplayName: "My Spark", + OrgID: "org_123", + DeviceID: "dev-uuid-pending", + Status: register.RegistrationStatusPending, + }, + } + + store := &mockDeregisterStore{ + user: &entity.User{ID: "user_1"}, + token: "tok", + } + + var removeCalled bool + svc := &fakeNodeService{ + removeNodeFn: func(req *nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error) { + removeCalled = true + return nil, fmt.Errorf("RemoveNode should not be called with empty ID %q", req.GetExternalNodeId()) + }, + } + + deps, server := testDeregisterDeps(t, svc, regStore) + defer server.Close() + + term := terminal.New() + err := runDeregister(context.Background(), term, store, deps, false) + if err != nil { + t.Fatalf("deregister of a pending registration should succeed, got: %v", err) + } + + if removeCalled { + t.Error("RemoveNode should not be called for a pending registration (no ExternalNodeID)") + } + + exists, err := regStore.Exists() + if err != nil { + t.Fatalf("Exists error: %v", err) + } + if exists { + t.Error("expected local pending registration to be deleted") + } +} + func Test_runDeregister_AlwaysUninstallsNetbird(t *testing.T) { regStore := &mockRegistrationStore{ reg: ®ister.DeviceRegistration{ From 587657ba1ab9f972a6b369bd9ba5866b14f48b76 Mon Sep 17 00:00:00 2001 From: Pratik Patel Date: Fri, 21 Aug 2026 14:01:15 -0700 Subject: [PATCH 6/6] clean --- pkg/cmd/deregister/deregister_test.go | 10 ---------- pkg/cmd/register/device_registration_store_test.go | 5 ----- pkg/cmd/register/register_test.go | 13 ++----------- 3 files changed, 2 insertions(+), 26 deletions(-) diff --git a/pkg/cmd/deregister/deregister_test.go b/pkg/cmd/deregister/deregister_test.go index fdf2c223..a18c093d 100644 --- a/pkg/cmd/deregister/deregister_test.go +++ b/pkg/cmd/deregister/deregister_test.go @@ -281,7 +281,6 @@ func Test_runDeregister_RemoveNodeFails(t *testing.T) { t.Fatal("expected error when RemoveNode fails") } - // Registration should still exist (server-side removal failed) exists, err := regStore.Exists() if err != nil { t.Fatalf("Exists error: %v", err) @@ -291,10 +290,6 @@ func Test_runDeregister_RemoveNodeFails(t *testing.T) { } } -// Test_runDeregister_RemoveNodeNotFound_ProceedsCleanup verifies that a NotFound -// from RemoveNode is treated as success: the node is already gone (e.g. a -// previous attempt succeeded before a transient error), so local cleanup must -// proceed and not leave the user stuck with an undeletable registration file. func Test_runDeregister_RemoveNodeNotFound_ProceedsCleanup(t *testing.T) { regStore := &mockRegistrationStore{ reg: ®ister.DeviceRegistration{ @@ -333,11 +328,6 @@ func Test_runDeregister_RemoveNodeNotFound_ProceedsCleanup(t *testing.T) { } } -// Test_runDeregister_PendingRegistration_SkipsRemoveNodeAndCleansUp verifies that -// deregistering a partial (pending) registration — which has no ExternalNodeID — -// skips the RemoveNode call (there's no known backend node ID to remove) and -// still cleans up local state, so the user isn't stuck with an undeletable -// pending record. func Test_runDeregister_PendingRegistration_SkipsRemoveNodeAndCleansUp(t *testing.T) { regStore := &mockRegistrationStore{ reg: ®ister.DeviceRegistration{ diff --git a/pkg/cmd/register/device_registration_store_test.go b/pkg/cmd/register/device_registration_store_test.go index c5e53fb5..2e1b085f 100644 --- a/pkg/cmd/register/device_registration_store_test.go +++ b/pkg/cmd/register/device_registration_store_test.go @@ -199,8 +199,6 @@ func Test_DeleteRegistration_FailsWhenMissing(t *testing.T) { } } -// Test_Load_IncludeAllReturnsPendingRecord verifies Load(true) returns a pending -// (intent) record that has no external node ID, which the strict Load(false) rejects. func Test_Load_IncludeAllReturnsPendingRecord(t *testing.T) { cleanup := setupTestFs(t) defer cleanup() @@ -231,14 +229,11 @@ func Test_Load_IncludeAllReturnsPendingRecord(t *testing.T) { t.Errorf("pending record should have no ExternalNodeID, got %q", loaded.ExternalNodeID) } - // The strict Load(false) must reject the pending record. if _, err := store.Load(false); err == nil { t.Error("expected Load(false) to error on a pending record") } } -// Test_Load_PendingRecordErrorMessage verifies Load(false) reports a clear, -// actionable message for a pending (in-progress) registration. func Test_Load_PendingRecordErrorMessage(t *testing.T) { cleanup := setupTestFs(t) defer cleanup() diff --git a/pkg/cmd/register/register_test.go b/pkg/cmd/register/register_test.go index 6fba5a73..1af5bdff 100644 --- a/pkg/cmd/register/register_test.go +++ b/pkg/cmd/register/register_test.go @@ -144,15 +144,11 @@ func (m mockNodeClientFactory) NewNodeClient(provider externalnode.TokenProvider return NewNodeServiceClient(provider, m.serverURL) } -// accessKeyPersistCall records a single Persist invocation. type accessKeyPersistCall struct { apiKey string orgID string } -// mockAccessKeyAuth records SeedSession/Persist calls so tests can assert the -// credential chain wired up the access key (and which org ID it was persisted -// against) without touching real auth stores. type mockAccessKeyAuth struct { seedCalls []string persistCalls []accessKeyPersistCall @@ -549,9 +545,6 @@ func Test_runRegister_AddNodeFails(t *testing.T) { t.Fatal("expected error when AddNode fails") } - // AddNode failed, but a pending intent record is written before AddNode so a - // retry can resume with the same device ID. The record must be pending (no - // external node ID yet) and carry a device ID for the retry to reuse. reg, loadErr := regStore.Load(true) if loadErr != nil { t.Fatalf("expected pending intent record after AddNode failure, got: %v", loadErr) @@ -948,9 +941,8 @@ func Test_runRegister_ResumesPendingRegistration(t *testing.T) { defer server.Close() term := terminal.New() - // On resume, opts (name/org/interactive) are ignored in favor of the pending - // record. Use interactive mode to skip the non-interactive --name/--org - // validation that runs before the Exists check. + // Use interactive mode so non-interactive --name/--org validation doesn't run + // before the Exists check; on resume the pending record's values are used. err := runRegister(context.Background(), term, store, registerOpts{interactive: true}, deps) if err != nil { t.Fatalf("runRegister failed: %v", err) @@ -1124,7 +1116,6 @@ func Test_runRegister_AccessKeyEnv_Fallback(t *testing.T) { deps.accessKeyAuth = accessKeyAuth term := terminal.New() - // No accessKey in opts; should fall back to the env var. opts := registerOpts{interactive: false, name: "my-spark", orgName: "TestOrg"} if err := runRegister(context.Background(), term, store, opts, deps); err != nil { t.Fatalf("runRegister failed: %v", err)