diff --git a/pkg/cmd/cmd.go b/pkg/cmd/cmd.go index 831a2aa9..7b0541de 100644 --- a/pkg/cmd/cmd.go +++ b/pkg/cmd/cmd.go @@ -265,12 +265,14 @@ func NewBrevCommand() *cobra.Command { //nolint:funlen,gocognit,gocyclo // defin fmt.Printf("%v\n", err) } - createCmdTree(cmds, t, loginCmdStore, noLoginCmdStore, loginAuth, externalNodeCmdStore) + 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 +320,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.go b/pkg/cmd/deregister/deregister.go index efd9090b..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" @@ -106,7 +107,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 } @@ -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 95c5ac99..a18c093d 100644 --- a/pkg/cmd/deregister/deregister_test.go +++ b/pkg/cmd/deregister/deregister_test.go @@ -57,7 +57,7 @@ func (m *mockRegistrationStore) Save(reg *register.DeviceRegistration) error { return nil } -func (m *mockRegistrationStore) Load() (*register.DeviceRegistration, error) { +func (m *mockRegistrationStore) Load(bool) (*register.DeviceRegistration, error) { if m.reg == nil { return nil, fmt.Errorf("no registration") } @@ -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,6 +290,89 @@ func Test_runDeregister_RemoveNodeFails(t *testing.T) { } } +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") + } +} + +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{ 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 10b3cbab..e1e6441a 100644 --- a/pkg/cmd/grantssh/grantssh_test.go +++ b/pkg/cmd/grantssh/grantssh_test.go @@ -44,7 +44,7 @@ func (m *mockRegistrationStore) Save(reg *register.DeviceRegistration) error { return nil } -func (m *mockRegistrationStore) Load() (*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 315dfb99..0838b098 100644 --- a/pkg/cmd/register/device_registration_store.go +++ b/pkg/cmd/register/device_registration_store.go @@ -20,31 +20,35 @@ const ( globalRegistrationDir = "/etc/brev" ) +const ( + 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"` + 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() (*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{} } @@ -72,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 { @@ -86,7 +89,16 @@ 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") + } 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..2e1b085f 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" @@ -42,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) } @@ -138,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") } @@ -159,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") } @@ -180,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") } @@ -197,3 +198,63 @@ func Test_DeleteRegistration_FailsWhenMissing(t *testing.T) { t.Error("expected error deleting missing registration") } } + +func Test_Load_IncludeAllReturnsPendingRecord(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.Load(true) + if err != nil { + t.Fatalf("Load 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) + } + + if _, err := store.Load(false); err == nil { + t.Error("expected Load(false) to error on a pending record") + } +} + +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(false) + 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 2ad88b43..9c57c514 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,92 @@ 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 +} + +type accessKeyAuthenticator struct { + session *store.MemoryAuthStore + durable *auth.LoginAuth +} + +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 +} + +type noopAccessKeyAuth struct{} + +func (noopAccessKeyAuth) SeedSession(string) error { return nil } +func (noopAccessKeyAuth) Persist(string, string) error { return nil } + +const accessKeyEnvVar = "BREV_ACCESS_KEY" + +func resolveAccessKey(flagValue string) string { + if v := strings.TrimSpace(flagValue); v != "" { + return v + } + return strings.TrimSpace(os.Getenv(accessKeyEnvVar)) +} + +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))) } } 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 +176,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 +193,22 @@ 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", "", "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 } -// 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 +227,42 @@ func runRegister(ctx context.Context, t *terminal.Terminal, s RegisterStore, opt } } - // Run through the login flow - brevUser, err := s.GetCurrentUser() - if err != nil { + // 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) { + 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.Load(true) + if err != nil { + return breverrors.WrapAndTrace(err) + } + // 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 { + // 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) } @@ -199,6 +295,11 @@ func runRegister(ctx context.Context, t *terminal.Terminal, s RegisterStore, opt return err } + // Org is resolved; persist the access key for other commands (mirrors '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,42 +327,22 @@ func runRegister(ctx context.Context, t *terminal.Terminal, s RegisterStore, opt } } - // Perform the registration steps - reg, err := runRegisterSteps(ctx, t, s, name, org, deps) + // 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 { return err } - // 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 { - 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 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("") t.Vprint(t.Yellow("[Step 1/5] Downloading and installing Brev tunnel...")) @@ -284,7 +365,22 @@ 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 pending intent before AddNode (see resumeRegistration). + 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, @@ -304,13 +400,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("") @@ -353,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) } @@ -424,58 +522,39 @@ 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 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("══════════════════════════════════════════════════")) - 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 to run 'brev enable-ssh'; 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..1af5bdff 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" @@ -72,7 +75,7 @@ func (m *mockRegistrationStore) Save(reg *DeviceRegistration) error { return nil } -func (m *mockRegistrationStore) Load() (*DeviceRegistration, error) { +func (m *mockRegistrationStore) Load(bool) (*DeviceRegistration, error) { if m.reg == nil { return nil, fmt.Errorf("no registration") } @@ -141,6 +144,28 @@ func (m mockNodeClientFactory) NewNodeClient(provider externalnode.TokenProvider return NewNodeServiceClient(provider, m.serverURL) } +type accessKeyPersistCall struct { + apiKey string + orgID string +} + +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 +206,7 @@ func testRegisterDeps(t *testing.T, svc *fakeNodeService, regStore RegistrationS profile: testHardwareProfile(), }, registrationStore: regStore, + accessKeyAuth: noopAccessKeyAuth{}, }, server } @@ -223,11 +249,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) @@ -242,7 +265,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) } @@ -295,7 +318,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 +406,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 +435,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 +474,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) @@ -465,7 +485,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) } @@ -489,7 +509,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 +539,27 @@ 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) + reg, loadErr := regStore.Load(true) + 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 +594,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 +687,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 +733,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 +767,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 +793,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 +819,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 +844,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 +886,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 +899,368 @@ 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(), +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() + // 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) + } + + 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.Load(true) + if loadErr != nil { + t.Fatalf("Load 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) + } +} + +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.Load(true) + if loadErr != nil { + t.Fatalf("Load 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) + } +} + +const testAccessKey = auth.BrevAPIKeyPrefix + "test-key" + +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 + } +} + +func ensureNoAccessKeyEnv(t *testing.T) { + t.Helper() + t.Setenv(accessKeyEnvVar, "") +} + +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) + } +} + +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) + } +} + +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() + 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) + } +} + +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) + } +} + +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) + } +} + +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} +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/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 447c53cb..71165928 100644 --- a/pkg/cmd/revokessh/revokessh_test.go +++ b/pkg/cmd/revokessh/revokessh_test.go @@ -43,7 +43,7 @@ func (m *mockRegistrationStore) Save(reg *register.DeviceRegistration) error { return nil } -func (m *mockRegistrationStore) Load() (*register.DeviceRegistration, error) { +func (m *mockRegistrationStore) Load(bool) (*register.DeviceRegistration, error) { if m.reg == nil { return nil, fmt.Errorf("no registration") }