Skip to content

improve register resiliency and idempotency #463

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
patelspratik wants to merge 1 commit into from
Open

improve register resiliency and idempotency #463

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions pkg/auth/auth.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,7 @@ type Auth struct {

const BrevAPIKeyPrefix = "bak-"

const MissingAPIKeyOrgIDMessage = "api key auth requires an org id; run brev login --api-key <api-key> --org-id <org-id>"
const MissingAPIKeyOrgIDMessage = "org id missing, please login again; run 'brev login --api-key <api-key>'"

Comment on lines +105 to 106

@drewmalin drewmalin Aug 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks like the old error message was more correct here (specifically the guidance to use --org-id <org-id>)?

type APIKeyAuthStore interface {
GetAuthTokens() (*entity.AuthTokens, error)
Expand Down Expand Up @@ -217,7 +217,6 @@ func (t Auth) GetFreshAccessTokenOrNil() (string, error) {
return apiKey, nil
}

// should always at least have access token?
if tokens.AccessToken == "" {
breverrors.GetDefaultErrorReporter().ReportMessage("access token is an empty string but shouldn't be")
}
Expand Down
12 changes: 3 additions & 9 deletions pkg/auth/auth_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,15 +28,6 @@ func TestIsAccessTokenValid(t *testing.T) {
if !assert.False(t, res) {
return
}

// expiredToken := "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCIsImtpZCI6ImdLTXBESXlRc0ZXSF9zYWdiT2oyViJ9.eyJpc3MiOiJodHRwczovL2JyZXZkZXYudXMuYXV0aDAuY29tLyIsInN1YiI6Imdvb2dsZS1vYXV0aDJ8MTAxNzY0NjMwNTEwODYxNDk5MTgwIiwiYXVkIjpbImh0dHBzOi8vYnJldmRldi51cy5hdXRoMC5jb20vYXBpL3YyLyIsImh0dHBzOi8vYnJldmRldi51cy5hdXRoMC5jb20vdXNlcmluZm8iXSwiaWF0IjoxNjM4NTYyMzY4LCJleHAiOjE2Mzg2NDg3NjgsImF6cCI6IkphcUpSTEVzZGF0NXc3VGIwV3FtVHh6SWVxd3FlcG1rIiwic2NvcGUiOiJvcGVuaWQgcHJvZmlsZSBlbWFpbCBvZmZsaW5lX2FjY2VzcyJ9.YCiO-som26ehT91qGAX5ZfrtVg4eYwamnlMRoCuUljXmg8Nf-ArDyoG32CqZkQ6YJ5XnzrVX9bVk5ZNHP_AFSE9SJvYL6MchoN09nR84WTbevRBCtZedIZUk5ULg6rWo5mszGr-S2gi08od4iTzXtKySPx1JnT60muRj_k9VV3MyixqvngEz5NvmFDdA8glGes5_iOuiBidmjOJzi_CVfKJ9s48BhlxzciSXFC0_DUBnT9OThjYjUP-22ohOuWwJWomRUv6gMSq78hJOALc330LwvmEsLdzlP7a3otIYM43hTtAVJ9QEL6M08GKqm3PdikzTxiGdfuQUhgMDlXygbQ"
// res, err := isAccessTokenValid(expiredToken)
// if !assert.Nil(t, err) {
// return
// }
// if !assert.False(t, res) {
// return
// }
}

type MockAuthStore struct {
Expand All @@ -48,6 +39,7 @@ type MockAuthStore struct {
func (m *MockAuthStore) SaveAuthTokens(tokens entity.AuthTokens) error {
m.saved = tokens
m.didSave = true
m.authTokens = &tokens // write-then-read consistent (mirrors a real store)
return nil
}

Expand Down Expand Up @@ -130,6 +122,8 @@ func TestIsAPIKeyAuthStore_LegacyCredentialsAreNotAPIKeyAuth(t *testing.T) {
assert.False(t, s.getAccessTokenCalled)
}

// When the env key matches the persisted key, its persisted org is valid.

type cliAuthStore struct {
tokens *entity.AuthTokens
user *entity.User
Expand Down
62 changes: 51 additions & 11 deletions pkg/cmd/deregister/deregister.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package deregister

import (
"context"
"errors"
"fmt"
"os/user"

Expand All @@ -20,18 +21,15 @@ import (
"github.com/spf13/cobra"
)

// DeregisterStore defines the store methods needed by the deregister command.
type DeregisterStore interface {
GetCurrentUser() (*entity.User, error)
GetAccessToken() (string, error)
}

// SSHKeyRemover removes Brev-managed SSH keys and returns the lines removed.
type SSHKeyRemover interface {
RemoveBrevKeys(u *user.User) ([]string, error)
}

// brevSSHKeyRemover delegates to register.RemoveBrevAuthorizedKeys.
type brevSSHKeyRemover struct{}

func (brevSSHKeyRemover) RemoveBrevKeys(u *user.User) ([]string, error) {
Expand Down Expand Up @@ -97,6 +95,53 @@ func NewCmdDeregister(t *terminal.Terminal, store DeregisterStore) *cobra.Comman
return cmd
}

func removeNodeFromBrev(ctx context.Context, t *terminal.Terminal, s DeregisterStore, deps deregisterDeps, reg *register.DeviceRegistration) error {
externalNodeID := reg.ExternalNodeID
if externalNodeID == "" && reg.DeviceID != "" {
lookedUp, lookupErr := findNodeByDeviceID(ctx, s, deps, reg.OrgID, reg.DeviceID)
if lookupErr != nil {
t.Vprintf(" %s\n", t.Yellow(fmt.Sprintf("Could not look up pending node by device ID: %v", lookupErr)))
}
if lookedUp != "" {
externalNodeID = lookedUp
}
}
if externalNodeID == "" {
t.Vprintf(" %s\n", t.Yellow("No registered node to remove (pending registration); cleaning up local state."))

@drewmalin drewmalin Aug 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks like this will happen if:

  • the node truly wasn't there, or
  • we had some failure in looking up the node

For the second case, I wonder if we actually shouldn't clean up local state so that we can retry the deregister or understand that a new registration shouldn't be allowed (or maybe we just abort the deregister altogether).

return nil
}
client := deps.nodeClients.NewNodeClient(s, config.GlobalConfig.GetBrevPublicAPIURL())
_, err := client.RemoveNode(ctx, connect.NewRequest(&nodev1.RemoveNodeRequest{
ExternalNodeId: externalNodeID,
}))
if err != nil {
var connectErr *connect.Error
if errors.As(err, &connectErr) && connectErr.Code() == connect.CodeNotFound {
t.Vprintf(" %s\n", t.Yellow("Node not found on Brev; continuing."))
return nil
}
return fmt.Errorf("failed to deregister node: %w", err)
}
t.Vprintf("%s Node removed from Brev.\n", t.Green(" ✓"))
return nil
}

func findNodeByDeviceID(ctx context.Context, s externalnode.TokenProvider, deps deregisterDeps, orgID, deviceID string) (string, error) {
client := deps.nodeClients.NewNodeClient(s, config.GlobalConfig.GetBrevPublicAPIURL())
resp, err := client.ListNodes(ctx, connect.NewRequest(&nodev1.ListNodesRequest{

@drewmalin drewmalin Aug 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Paginate?

OrganizationId: orgID,
}))
if err != nil {
return "", fmt.Errorf("failed to list nodes: %w", err)
}
for _, n := range resp.Msg.GetItems() {
if n.GetDeviceId() == deviceID {
return n.GetExternalNodeId(), nil
}
}
return "", nil
}

func runDeregister(ctx context.Context, t *terminal.Terminal, s DeregisterStore, deps deregisterDeps, skipConfirm bool) error { //nolint:funlen,gocyclo // deregistration flow
if !deps.platform.IsCompatible() {
return fmt.Errorf("brev deregister is only supported on Linux")
Expand All @@ -106,7 +151,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
}
Expand Down Expand Up @@ -158,14 +203,9 @@ 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 err := removeNodeFromBrev(ctx, t, s, deps, reg); err != nil {
return err
}
t.Vprintf("%s Node removed from Brev.\n", t.Green(" ✓"))
t.Vprint("")

t.Vprint(t.Yellow("[Step 2/4] Removing Brev SSH keys..."))
Expand Down
141 changes: 139 additions & 2 deletions pkg/cmd/deregister/deregister_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ func (m *mockDeregisterStore) GetAccessToken() (string, error) { return m.token,
type fakeNodeService struct {
nodev1connect.UnimplementedExternalNodeServiceHandler
removeNodeFn func(*nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error)
listNodesFn func(*nodev1.ListNodesRequest) (*nodev1.ListNodesResponse, error)
}

func (f *fakeNodeService) RemoveNode(_ context.Context, req *connect.Request[nodev1.RemoveNodeRequest]) (*connect.Response[nodev1.RemoveNodeResponse], error) {
Expand All @@ -47,6 +48,14 @@ func (f *fakeNodeService) RemoveNode(_ context.Context, req *connect.Request[nod
return connect.NewResponse(resp), nil
}

func (f *fakeNodeService) ListNodes(_ context.Context, req *connect.Request[nodev1.ListNodesRequest]) (*connect.Response[nodev1.ListNodesResponse], error) {
resp, err := f.listNodesFn(req.Msg)
if err != nil {
return nil, err
}
return connect.NewResponse(resp), nil
}

// mockRegistrationStore satisfies register.RegistrationStore for deregister tests.
type mockRegistrationStore struct {
reg *register.DeviceRegistration
Expand All @@ -57,7 +66,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")
}
Expand Down Expand Up @@ -281,7 +290,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)
Expand All @@ -291,6 +299,135 @@ func Test_runDeregister_RemoveNodeFails(t *testing.T) {
}
}

func Test_runDeregister_RemoveNodeNotFound_ProceedsCleanup(t *testing.T) {
regStore := &mockRegistrationStore{
reg: &register.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(t *testing.T) {
const deviceID = "dev-uuid-pending"
reg := &register.DeviceRegistration{
DisplayName: "My Spark",
OrgID: "org_123",
DeviceID: deviceID,
Status: register.RegistrationStatusPending,
}

tests := []struct {
name string
listNodesFn func(*nodev1.ListNodesRequest) (*nodev1.ListNodesResponse, error)
removeNodeFn func(*nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error)
wantRemovedID string // empty = RemoveNode must not be called
}{
{
name: "no backend node matches: cleans up locally without RemoveNode",
listNodesFn: func(*nodev1.ListNodesRequest) (*nodev1.ListNodesResponse, error) {
return &nodev1.ListNodesResponse{}, nil
},
removeNodeFn: func(req *nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error) {
return nil, fmt.Errorf("RemoveNode should not be called with empty ID %q", req.GetExternalNodeId())
},
},
{
name: "backend node recovered by device ID: removed",
listNodesFn: func(req *nodev1.ListNodesRequest) (*nodev1.ListNodesResponse, error) {
return &nodev1.ListNodesResponse{
Items: []*nodev1.ExternalNode{
{ExternalNodeId: "unode_other", DeviceId: "dev-different"},
{ExternalNodeId: "unode_recovered", DeviceId: deviceID},
},
}, nil
},
removeNodeFn: func(req *nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error) {
return &nodev1.RemoveNodeResponse{}, nil
},
wantRemovedID: "unode_recovered",
},
{
name: "ListNodes failure is non-fatal: local state still deleted",
listNodesFn: func(*nodev1.ListNodesRequest) (*nodev1.ListNodesResponse, error) {
return nil, connect.NewError(connect.CodeInternal, nil)
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
regStore := &mockRegistrationStore{reg: reg}
store := &mockDeregisterStore{user: &entity.User{ID: "user_1"}, token: "tok"}

var gotOrgID string
var removedNodeID string
svc := &fakeNodeService{
listNodesFn: func(req *nodev1.ListNodesRequest) (*nodev1.ListNodesResponse, error) {
gotOrgID = req.GetOrganizationId()
return tt.listNodesFn(req)
},
removeNodeFn: func(req *nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error) {
removedNodeID = req.GetExternalNodeId()
return tt.removeNodeFn(req)
},
}

deps, server := testDeregisterDeps(t, svc, regStore)
defer server.Close()

term := terminal.New()
if err := runDeregister(context.Background(), term, store, deps, false); err != nil {
t.Fatalf("deregister failed: %v", err)
}

if gotOrgID != "org_123" {
t.Errorf("expected ListNodes scoped to org_123, got %q", gotOrgID)
}
if tt.wantRemovedID == "" {
if removedNodeID != "" {
t.Errorf("RemoveNode should not be called, got %q", removedNodeID)
}
} else if removedNodeID != tt.wantRemovedID {
t.Errorf("expected RemoveNode called with %q, got %q", tt.wantRemovedID, removedNodeID)
}

exists, _ := regStore.Exists()
if exists {
t.Error("expected local registration to be deleted")
}
})
}
}

func Test_runDeregister_AlwaysUninstallsNetbird(t *testing.T) {
regStore := &mockRegistrationStore{
reg: &register.DeviceRegistration{
Expand Down
3 changes: 1 addition & 2 deletions pkg/cmd/enablessh/enablessh.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down Expand Up @@ -96,7 +96,6 @@ func enableSSH(
linuxUsername := linuxUser.Username

checkSSHDaemon(t)

t.Vprint("")
t.Vprint(t.Green("Enabling SSH access on this device"))
t.Vprint("")
Expand Down
7 changes: 3 additions & 4 deletions pkg/cmd/grantssh/grantssh.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -183,12 +183,11 @@ func runGrantSSH(ctx context.Context, t *terminal.Terminal, s GrantSSHStore, opt
return err
}
linuxUserOptions := uniqueLinuxUsersFromNodeSSHAccess(node)
if len(linuxUserOptions) > 0 {
t.Vprint("")
linuxUser = deps.prompter.Select("Select Linux user on the node", linuxUserOptions)
} else {
if len(linuxUserOptions) == 0 {
return fmt.Errorf("no Linux users on this node yet; run with --linux-user to specify one (e.g. after enable-ssh on the node)")
}
t.Vprint("")
linuxUser = deps.prompter.Select("Select Linux user on the node", linuxUserOptions)
} else {
selectedUser, err = findUserByIDOrEmail(orgMembers, opts.userIDOrEmail)
if err != nil {
Expand Down
2 changes: 1 addition & 1 deletion pkg/cmd/grantssh/grantssh_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
}
Expand Down
2 changes: 1 addition & 1 deletion pkg/cmd/ls/ls_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -214,7 +214,7 @@ func TestRunLs_APIKeyRequiresCredentialOrg(t *testing.T) {
if err == nil {
t.Fatal("expected missing API key org error, got nil")
}
if !strings.Contains(err.Error(), "api key auth requires an org id") {
if !strings.Contains(err.Error(), "org id missing") {
t.Fatalf("expected API key org validation error, got %v", err)
}
if s.workspaceOrgID != "" {
Expand Down
Loading
Loading