Skip to content

Commit 6d2a609

Browse files
patelspratikpatelspratik
committed
feat(BRE2-814): allocate ssh port explicitly
1 parent 9266638 commit 6d2a609

9 files changed

Lines changed: 201 additions & 20 deletions

File tree

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ go 1.24.0
44

55
require (
66
buf.build/gen/go/brevdev/devplane/connectrpc/go v1.19.1-20260228021043-887d38e1b474.2
7-
buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.11-20260305201249-af5106090c8a.1
7+
buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.11-20260309172248-8105d701fdce.1
88
connectrpc.com/connect v1.19.1
99
github.com/NVIDIA/go-nvml v0.13.0-1
1010
github.com/alessio/shellescape v1.4.1

‎go.sum‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
buf.build/gen/go/brevdev/devplane/connectrpc/go v1.19.1-20260228021043-887d38e1b474.2 h1:Sq0kIa/xKzScbJcqB5EbPVhOL0QYHPr3araQaupL2lk=
22
buf.build/gen/go/brevdev/devplane/connectrpc/go v1.19.1-20260228021043-887d38e1b474.2/go.mod h1:Yh34p9aADmWsKv2umYlMpnCZuBmNBE9N+HImgRriJXM=
3-
buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.11-20260305201249-af5106090c8a.1 h1:d+kY4OSI/WV2eBuO5G7ezCQ8RiLtDOIrUbtmZOKi5Kw=
4-
buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.11-20260305201249-af5106090c8a.1/go.mod h1:V/y7Wxg0QvU4XPVwqErF5NHLobUT1QEyfgrGuQIxdPo=
3+
buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.11-20260309172248-8105d701fdce.1 h1:lWdcuXsXpMfPOer4yawjwomVbtSAnGgFAWYF8ggK9g4=
4+
buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.11-20260309172248-8105d701fdce.1/go.mod h1:V/y7Wxg0QvU4XPVwqErF5NHLobUT1QEyfgrGuQIxdPo=
55
buf.build/gen/go/brevdev/protoc-gen-gotag/protocolbuffers/go v1.36.11-20220906235457-8b4922735da5.1 h1:6amhprQmCKJ4wgJ6ngkh32d9V+dQcOLUZ/SfHdOnYgo=
66
buf.build/gen/go/brevdev/protoc-gen-gotag/protocolbuffers/go v1.36.11-20220906235457-8b4922735da5.1/go.mod h1:O+pnSHMru/naTMrm4tmpBoH3wz6PHa+R75HR7Mv8X2g=
77
cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=

‎pkg/cmd/enablessh/enablessh.go‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,11 @@ func enableSSH(
103103
return fmt.Errorf("SSH port: %w", err)
104104
}
105105

106-
if err := register.GrantSSHAccessToNode(ctx, t, deps.nodeClients, tokenProvider, reg, brevUser, u, port); err != nil {
106+
if err := register.OpenSSHPort(ctx, t, deps.nodeClients, tokenProvider, reg, port); err != nil {
107+
return fmt.Errorf("enable SSH failed: %w", err)
108+
}
109+
110+
if err := register.GrantSSHAccessToNode(ctx, t, deps.nodeClients, tokenProvider, reg, brevUser, u); err != nil {
107111
return fmt.Errorf("enable SSH failed: %w", err)
108112
}
109113

‎pkg/cmd/grantssh/grantssh.go‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -131,12 +131,7 @@ func runGrantSSH(ctx context.Context, t *terminal.Terminal, s GrantSSHStore, dep
131131
t.Vprintf(" Linux user: %s\n", linuxUser)
132132
t.Vprint("")
133133

134-
port, err := register.PromptSSHPort(t)
135-
if err != nil {
136-
return fmt.Errorf("SSH port: %w", err)
137-
}
138-
139-
if err := register.GrantSSHAccessToNode(ctx, t, deps.nodeClients, s, reg, selectedUser, osUser, port); err != nil {
134+
if err := register.GrantSSHAccessToNode(ctx, t, deps.nodeClients, s, reg, selectedUser, osUser); err != nil {
140135
return fmt.Errorf("grant SSH failed: %w", err)
141136
}
142137

‎pkg/cmd/grantssh/grantssh_test.go‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,17 @@ type fakeNodeService struct {
111111
grantSSHFn func(*nodev1.GrantNodeSSHAccessRequest) (*nodev1.GrantNodeSSHAccessResponse, error)
112112
}
113113

114+
func (f *fakeNodeService) OpenPort(_ context.Context, req *connect.Request[nodev1.OpenPortRequest]) (*connect.Response[nodev1.OpenPortResponse], error) {
115+
return connect.NewResponse(&nodev1.OpenPortResponse{
116+
Port: &nodev1.Port{
117+
PortId: "port_ssh",
118+
Protocol: req.Msg.GetProtocol(),
119+
PortNumber: req.Msg.GetPortNumber(),
120+
ServerPort: req.Msg.GetPortNumber(),
121+
},
122+
}), nil
123+
}
124+
114125
func (f *fakeNodeService) GrantNodeSSHAccess(_ context.Context, req *connect.Request[nodev1.GrantNodeSSHAccessRequest]) (*connect.Response[nodev1.GrantNodeSSHAccessResponse], error) {
115126
resp, err := f.grantSSHFn(req.Msg)
116127
if err != nil {
@@ -271,9 +282,6 @@ func Test_runGrantSSH_HappyPath(t *testing.T) {
271282
deps, server := testGrantSSHDeps(t, svc, regStore)
272283
defer server.Close()
273284

274-
register.SetTestSSHPort(22)
275-
defer register.ClearTestSSHPort()
276-
277285
term := terminal.New()
278286
err := runGrantSSH(context.Background(), term, store, deps)
279287
if err != nil {
@@ -326,9 +334,6 @@ func Test_runGrantSSH_RPCFailure(t *testing.T) {
326334
deps, server := testGrantSSHDeps(t, svc, regStore)
327335
defer server.Close()
328336

329-
register.SetTestSSHPort(22)
330-
defer register.ClearTestSSHPort()
331-
332337
term := terminal.New()
333338
err := runGrantSSH(context.Background(), term, store, deps)
334339
if err == nil {

‎pkg/cmd/register/register.go‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -333,7 +333,11 @@ func grantSSHAccess(ctx context.Context, t *terminal.Terminal, deps registerDeps
333333
return fmt.Errorf("SSH port: %w", err)
334334
}
335335

336-
err = GrantSSHAccessToNode(ctx, t, deps.nodeClients, tokenProvider, reg, brevUser, osUser, port)
336+
if err := OpenSSHPort(ctx, t, deps.nodeClients, tokenProvider, reg, port); err != nil {
337+
return fmt.Errorf("allocate SSH port failed: %w", err)
338+
}
339+
340+
err = GrantSSHAccessToNode(ctx, t, deps.nodeClients, tokenProvider, reg, brevUser, osUser)
337341
if err != nil {
338342
return fmt.Errorf("grant SSH failed: %w", err)
339343
}

‎pkg/cmd/register/register_test.go‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -913,3 +913,134 @@ func Test_runRegister_NoNameAlreadyRegistered(t *testing.T) {
913913
t.Error("expected registration to still exist")
914914
}
915915
}
916+
917+
func Test_runRegister_OpenSSHPort(t *testing.T) { // nolint:funlen // test
918+
tests := []struct {
919+
name string
920+
port int32
921+
openFn func(*nodev1.OpenPortRequest) (*nodev1.OpenPortResponse, error)
922+
verify func(t *testing.T, openReq *nodev1.OpenPortRequest, grantReq *nodev1.GrantNodeSSHAccessRequest, reg *mockRegistrationStore, err error)
923+
}{
924+
{
925+
name: "SendsCorrectArgs",
926+
port: 2222,
927+
openFn: func(req *nodev1.OpenPortRequest) (*nodev1.OpenPortResponse, error) {
928+
return &nodev1.OpenPortResponse{
929+
Port: &nodev1.Port{
930+
PortId: "port_ssh",
931+
Protocol: req.GetProtocol(),
932+
PortNumber: req.GetPortNumber(),
933+
},
934+
}, nil
935+
},
936+
verify: func(t *testing.T, openReq *nodev1.OpenPortRequest, _ *nodev1.GrantNodeSSHAccessRequest, _ *mockRegistrationStore, err error) {
937+
t.Helper()
938+
if err != nil {
939+
t.Fatalf("runRegister failed: %v", err)
940+
}
941+
if openReq == nil {
942+
t.Fatal("expected OpenPort to be called")
943+
}
944+
if openReq.GetExternalNodeId() != "unode_abc" {
945+
t.Errorf("expected node ID unode_abc, got %s", openReq.GetExternalNodeId())
946+
}
947+
if openReq.GetProtocol() != nodev1.PortProtocol_PORT_PROTOCOL_SSH {
948+
t.Errorf("expected PORT_PROTOCOL_SSH, got %s", openReq.GetProtocol())
949+
}
950+
if openReq.GetPortNumber() != 2222 {
951+
t.Errorf("expected port 2222, got %d", openReq.GetPortNumber())
952+
}
953+
},
954+
},
955+
{
956+
name: "FailureIsSoftError",
957+
port: 22,
958+
openFn: func(_ *nodev1.OpenPortRequest) (*nodev1.OpenPortResponse, error) {
959+
return nil, connect.NewError(connect.CodeInternal, fmt.Errorf("skybridge unavailable"))
960+
},
961+
verify: func(t *testing.T, _ *nodev1.OpenPortRequest, _ *nodev1.GrantNodeSSHAccessRequest, regStore *mockRegistrationStore, err error) {
962+
t.Helper()
963+
if err != nil {
964+
t.Fatalf("registration should succeed even when OpenSSHPort fails (soft error), got: %v", err)
965+
}
966+
exists, _ := regStore.Exists()
967+
if !exists {
968+
t.Error("expected registration to still exist after OpenSSHPort failure")
969+
}
970+
},
971+
},
972+
{
973+
name: "GrantRequestHasNoPort",
974+
port: 22,
975+
verify: func(t *testing.T, _ *nodev1.OpenPortRequest, grantReq *nodev1.GrantNodeSSHAccessRequest, _ *mockRegistrationStore, err error) {
976+
t.Helper()
977+
if err != nil {
978+
t.Fatalf("runRegister failed: %v", err)
979+
}
980+
if grantReq == nil {
981+
t.Fatal("expected GrantNodeSSHAccess to be called")
982+
}
983+
if grantReq.GetExternalNodeId() != "unode_abc" {
984+
t.Errorf("expected node ID unode_abc, got %s", grantReq.GetExternalNodeId())
985+
}
986+
if grantReq.GetUserId() != "user_1" {
987+
t.Errorf("expected user ID user_1, got %s", grantReq.GetUserId())
988+
}
989+
},
990+
},
991+
}
992+
993+
for _, tt := range tests {
994+
t.Run(tt.name, func(t *testing.T) {
995+
regStore := &mockRegistrationStore{}
996+
store := &mockRegisterStore{
997+
user: &entity.User{ID: "user_1"},
998+
org: &entity.Organization{ID: "org_123", Name: "TestOrg"},
999+
token: "tok",
1000+
}
1001+
1002+
var gotOpenReq *nodev1.OpenPortRequest
1003+
var gotGrantReq *nodev1.GrantNodeSSHAccessRequest
1004+
svc := &fakeNodeService{
1005+
addNodeFn: func(req *nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) {
1006+
return &nodev1.AddNodeResponse{
1007+
ExternalNode: &nodev1.ExternalNode{
1008+
ExternalNodeId: "unode_abc",
1009+
OrganizationId: "org_123",
1010+
Name: req.GetName(),
1011+
DeviceId: req.GetDeviceId(),
1012+
ConnectivityInfo: &nodev1.ConnectivityInfo{
1013+
RegistrationCommand: "netbird up --key abc",
1014+
},
1015+
},
1016+
}, nil
1017+
},
1018+
openPortFn: func(req *nodev1.OpenPortRequest) (*nodev1.OpenPortResponse, error) {
1019+
gotOpenReq = req
1020+
if tt.openFn != nil {
1021+
return tt.openFn(req)
1022+
}
1023+
return &nodev1.OpenPortResponse{
1024+
Port: &nodev1.Port{PortId: "port_ssh", Protocol: req.GetProtocol(), PortNumber: req.GetPortNumber()},
1025+
}, nil
1026+
},
1027+
grantNodeSSHAccessFn: func(_ *nodev1.GrantNodeSSHAccessRequest) (*nodev1.GrantNodeSSHAccessResponse, error) {
1028+
return &nodev1.GrantNodeSSHAccessResponse{}, nil
1029+
},
1030+
}
1031+
1032+
deps, server := testRegisterDeps(t, svc, regStore)
1033+
defer server.Close()
1034+
1035+
deps.prompter = mockConfirmer{confirm: true}
1036+
1037+
SetTestSSHPort(tt.port)
1038+
defer ClearTestSSHPort()
1039+
1040+
term := terminal.New()
1041+
err := runRegister(context.Background(), term, store, "my-spark", "", deps)
1042+
1043+
tt.verify(t, gotOpenReq, gotGrantReq, regStore, err)
1044+
})
1045+
}
1046+
}

‎pkg/cmd/register/rpcclient_test.go‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,7 @@ type fakeNodeService struct {
179179
addNodeFn func(*nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error)
180180
removeNodeFn func(*nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error)
181181
getNodeFn func(*nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error)
182+
openPortFn func(*nodev1.OpenPortRequest) (*nodev1.OpenPortResponse, error)
182183
grantNodeSSHAccessFn func(*nodev1.GrantNodeSSHAccessRequest) (*nodev1.GrantNodeSSHAccessResponse, error)
183184
}
184185

@@ -209,6 +210,24 @@ func (f *fakeNodeService) GetNode(_ context.Context, req *connect.Request[nodev1
209210
return connect.NewResponse(resp), nil
210211
}
211212

213+
func (f *fakeNodeService) OpenPort(_ context.Context, req *connect.Request[nodev1.OpenPortRequest]) (*connect.Response[nodev1.OpenPortResponse], error) {
214+
if f.openPortFn == nil {
215+
return connect.NewResponse(&nodev1.OpenPortResponse{
216+
Port: &nodev1.Port{
217+
PortId: "port_ssh",
218+
Protocol: req.Msg.GetProtocol(),
219+
PortNumber: req.Msg.GetPortNumber(),
220+
ServerPort: req.Msg.GetPortNumber(),
221+
},
222+
}), nil
223+
}
224+
resp, err := f.openPortFn(req.Msg)
225+
if err != nil {
226+
return nil, err
227+
}
228+
return connect.NewResponse(resp), nil
229+
}
230+
212231
func (f *fakeNodeService) GrantNodeSSHAccess(_ context.Context, req *connect.Request[nodev1.GrantNodeSSHAccessRequest]) (*connect.Response[nodev1.GrantNodeSSHAccessResponse], error) {
213232
if f.grantNodeSSHAccessFn == nil {
214233
return nil, connect.NewError(connect.CodeUnimplemented, nil)

‎pkg/cmd/register/sshkeys.go‎

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -131,9 +131,34 @@ func RemoveAuthorizedKeyLine(u *user.User, line string) error {
131131
return nil
132132
}
133133

134+
// OpenSSHPort calls the OpenPort RPC to allocate an SSH port on the node.
135+
// This must be called before GrantSSHAccessToNode when enabling SSH for the
136+
// first time on a device. The call is idempotent — if the port is already
137+
// open, the server returns the existing allocation.
138+
func OpenSSHPort(
139+
ctx context.Context,
140+
t *terminal.Terminal,
141+
nodeClients externalnode.NodeClientFactory,
142+
tokenProvider externalnode.TokenProvider,
143+
reg *DeviceRegistration,
144+
port int32,
145+
) error {
146+
client := nodeClients.NewNodeClient(tokenProvider, config.GlobalConfig.GetBrevPublicAPIURL())
147+
_, err := client.OpenPort(ctx, connect.NewRequest(&nodev1.OpenPortRequest{
148+
ExternalNodeId: reg.ExternalNodeID,
149+
Protocol: nodev1.PortProtocol_PORT_PROTOCOL_SSH,
150+
PortNumber: port,
151+
}))
152+
if err != nil {
153+
return fmt.Errorf("failed to allocate SSH port: %w", err)
154+
}
155+
t.Vprintf(" SSH port %d allocated.\n", port)
156+
return nil
157+
}
158+
134159
// GrantSSHAccessToNode installs the user's public key in authorized_keys and
135160
// calls GrantNodeSSHAccess to record access server-side. If the RPC fails,
136-
// the installed key is rolled back. port is the target SSH port (e.g. 22).
161+
// the installed key is rolled back.
137162
func GrantSSHAccessToNode(
138163
ctx context.Context,
139164
t *terminal.Terminal,
@@ -142,7 +167,6 @@ func GrantSSHAccessToNode(
142167
reg *DeviceRegistration,
143168
targetUser *entity.User,
144169
osUser *user.User,
145-
port int32,
146170
) error {
147171
if targetUser.PublicKey != "" {
148172
if added, err := InstallAuthorizedKey(osUser, targetUser.PublicKey, targetUser.ID); err != nil {
@@ -165,7 +189,6 @@ func GrantSSHAccessToNode(
165189
ExternalNodeId: reg.ExternalNodeID,
166190
UserId: targetUser.ID,
167191
LinuxUser: osUser.Username,
168-
Port: port,
169192
}))
170193
if err != nil {
171194
// Retryable error

0 commit comments

Comments
 (0)