RSDK-9863 Add ListTunnels wrappers (viamrobotics#4821)

Author: benjirewis

go.mod

@@ -75,7 +75,7 @@ require (
 	go.uber.org/atomic v1.11.0
 	go.uber.org/multierr v1.11.0
 	go.uber.org/zap v1.27.0
-	go.viam.com/api v0.1.396
+	go.viam.com/api v0.1.397
 	go.viam.com/test v1.2.4
 	go.viam.com/utils v0.1.130
 	goji.io v2.0.2+incompatible

go.sum

@@ -1513,8 +1513,8 @@ go.uber.org/zap v1.18.1/go.mod h1:xg/QME4nWcxGxrpdeYfq7UvYrLh66cuVKdrbD1XF/NI=
 go.uber.org/zap v1.23.0/go.mod h1:D+nX8jyLsMHMYrln8A0rJjFt/T/9/bGgIhAqxv5URuY=
 go.uber.org/zap v1.27.0 h1:aJMhYGrd5QSmlpLMr2MftRKl7t8J8PTZPA732ud/XR8=
 go.uber.org/zap v1.27.0/go.mod h1:GB2qFLM7cTU87MWRP2mPIjqfIDnGu+VIO4V/SdhGo2E=
-go.viam.com/api v0.1.396 h1:HqX+zt24BJFPcIQZk1fttDR6goiRCTt1pcVZyBOi7n4=
-go.viam.com/api v0.1.396/go.mod h1:drvlBWaiHFxPziz5jayHvibez1qG7lylcNCC1LF8onU=
+go.viam.com/api v0.1.397 h1:BGHQ71TmGDtUOWSmY66STA8/hcbUFV+okWAtgYVK7SA=
+go.viam.com/api v0.1.397/go.mod h1:drvlBWaiHFxPziz5jayHvibez1qG7lylcNCC1LF8onU=
 go.viam.com/test v1.2.4 h1:JYgZhsuGAQ8sL9jWkziAXN9VJJiKbjoi9BsO33TW3ug=
 go.viam.com/test v1.2.4/go.mod h1:zI2xzosHdqXAJ/kFqcN+OIF78kQuTV2nIhGZ8EzvaJI=
 go.viam.com/utils v0.1.130 h1:5vggHNK/ar/YHXjf9z5WuihX+n9jYFnA2CEFOYi99d8=

robot/client/client.go

@@ -1306,6 +1306,30 @@ func (rc *RobotClient) Tunnel(ctx context.Context, conn io.ReadWriteCloser, dest
 	return errors.Join(err, readerSenderErr, recvWriterErr)
 }
 
+// ListTunnels lists all available tunnels configured on the robot.
+func (rc *RobotClient) ListTunnels(ctx context.Context) ([]config.TrafficTunnelEndpoint, error) {
+	var ttes []config.TrafficTunnelEndpoint
+
+	resp, err := rc.client.ListTunnels(ctx, &pb.ListTunnelsRequest{})
+	if err != nil {
+		return ttes, err
+	}
+
+	for _, protoTTE := range resp.Tunnels {
+		if protoTTE == nil {
+			continue
+		}
+
+		tte := config.TrafficTunnelEndpoint{
+			Port:              int(protoTTE.Port),
+			ConnectionTimeout: protoTTE.ConnectionTimeout.AsDuration(),
+		}
+		ttes = append(ttes, tte)
+	}
+
+	return ttes, nil
+}
+
 // SetPeerConnection is only to be called internally from modules.
 func (rc *RobotClient) SetPeerConnection(pc *webrtc.PeerConnection) {
 	rc.mu.Lock()
@@ -1366,8 +1390,3 @@ func streamClientInterceptor() googlegrpc.StreamClientInterceptor {
 		return streamer(ctx, desc, cc, method, opts...)
 	}
 }
-
-// ListTunnels is a no-op (returns nil) for robot clients.
-func (rc *RobotClient) ListTunnels() []config.TrafficTunnelEndpoint {
-	return nil
-}

robot/client/client_test.go

@@ -2286,3 +2286,49 @@ func TestVersion(t *testing.T) {
 	test.That(t, err, test.ShouldBeNil)
 	test.That(t, md, test.ShouldResemble, version)
 }
+
+func TestListTunnels(t *testing.T) {
+	logger := logging.NewTestLogger(t)
+	listener, err := net.Listen("tcp", "localhost:0")
+	test.That(t, err, test.ShouldBeNil)
+	gServer := grpc.NewServer()
+
+	expectedTTEs := []config.TrafficTunnelEndpoint{
+		{
+			Port:              9090,
+			ConnectionTimeout: 20 * time.Second,
+		},
+		{
+			Port:              27017,
+			ConnectionTimeout: 40 * time.Millisecond,
+		},
+		{
+			Port: 23654,
+		},
+	}
+	injectRobot := &inject.Robot{
+		ResourceNamesFunc:   func() []resource.Name { return nil },
+		ResourceRPCAPIsFunc: func() []resource.RPCAPI { return nil },
+		MachineStatusFunc: func(ctx context.Context) (robot.MachineStatus, error) {
+			return robot.MachineStatus{State: robot.StateRunning}, nil
+		},
+		ListTunnelsFunc: func(ctx context.Context) ([]config.TrafficTunnelEndpoint, error) {
+			return expectedTTEs, nil
+		},
+	}
+
+	pb.RegisterRobotServiceServer(gServer, server.New(injectRobot))
+
+	go gServer.Serve(listener)
+	defer gServer.Stop()
+
+	client, err := New(context.Background(), listener.Addr().String(), logger)
+	test.That(t, err, test.ShouldBeNil)
+	defer func() {
+		test.That(t, client.Close(context.Background()), test.ShouldBeNil)
+	}()
+
+	ttes, err := client.ListTunnels(context.Background())
+	test.That(t, err, test.ShouldBeNil)
+	test.That(t, ttes, test.ShouldResemble, expectedTTEs)
+}

robot/impl/local_robot.go

@@ -1590,10 +1590,10 @@ func (r *localRobot) RestartAllowed() bool {
 }
 
 // ListTunnels returns information on available traffic tunnels.
-func (r *localRobot) ListTunnels() []config.TrafficTunnelEndpoint {
+func (r *localRobot) ListTunnels(_ context.Context) ([]config.TrafficTunnelEndpoint, error) {
 	cfg := r.Config()
 	if cfg != nil {
-		return cfg.Network.NetworkConfigData.TrafficTunnelEndpoints
+		return cfg.Network.NetworkConfigData.TrafficTunnelEndpoints, nil
 	}
-	return nil
+	return nil, nil
 }

robot/impl/local_robot_test.go

@@ -4640,6 +4640,7 @@ func TestListTunnels(t *testing.T) {
 
 	r := setupLocalRobot(t, ctx, cfg, logger)
 
-	ttes := r.ListTunnels()
+	ttes, err := r.ListTunnels(ctx)
+	test.That(t, err, test.ShouldBeNil)
 	test.That(t, ttes, test.ShouldResemble, trafficTunnelEndpoints)
 }

robot/impl/resource_manager_test.go

@@ -2037,8 +2037,8 @@ func (rr *dummyRobot) Version(ctx context.Context) (robot.VersionResponse, error
 }
 
 // ListTunnels returns information on available traffic tunnels.
-func (rr *dummyRobot) ListTunnels() []config.TrafficTunnelEndpoint {
-	return nil
+func (rr *dummyRobot) ListTunnels(_ context.Context) ([]config.TrafficTunnelEndpoint, error) {
+	return nil, nil
 }
 
 // managerForDummyRobot integrates all parts from a given robot except for its remotes.

robot/robot.go

@@ -157,7 +157,7 @@ type Robot interface {
 	Version(ctx context.Context) (VersionResponse, error)
 
 	// ListTunnels returns information on available traffic tunnels.
-	ListTunnels() []config.TrafficTunnelEndpoint
+	ListTunnels(ctx context.Context) ([]config.TrafficTunnelEndpoint, error)
 }
 
 // A LocalRobot is a Robot that can have its parts modified.

robot/server/server.go

@@ -135,6 +135,25 @@ func (s *Server) Tunnel(srv pb.RobotService_TunnelServer) error {
 	return errors.Join(err, readerSenderErr, recvWriterErr)
 }
 
+// ListTunnels lists all available tunnels on the server.
+func (s *Server) ListTunnels(ctx context.Context, req *pb.ListTunnelsRequest) (*pb.ListTunnelsResponse, error) {
+	res := &pb.ListTunnelsResponse{}
+
+	ttes, err := s.robot.ListTunnels(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, tte := range ttes {
+		res.Tunnels = append(res.Tunnels, &pb.Tunnel{
+			Port:              uint32(tte.Port),
+			ConnectionTimeout: durationpb.New(tte.ConnectionTimeout),
+		})
+	}
+
+	return res, nil
+}
+
 // GetOperations lists all running operations.
 func (s *Server) GetOperations(ctx context.Context, req *pb.GetOperationsRequest) (*pb.GetOperationsResponse, error) {
 	me := operation.Get(ctx)

testutils/inject/robot.go

@@ -51,6 +51,7 @@ type Robot struct {
 	CloudMetadataFunc       func(ctx context.Context) (cloud.Metadata, error)
 	MachineStatusFunc       func(ctx context.Context) (robot.MachineStatus, error)
 	ShutdownFunc            func(ctx context.Context) error
+	ListTunnelsFunc         func(ctx context.Context) ([]config.TrafficTunnelEndpoint, error)
 
 	ops        *operation.Manager
 	SessMgr    session.Manager
@@ -316,6 +317,16 @@ func (r *Robot) Shutdown(ctx context.Context) error {
 	return r.ShutdownFunc(ctx)
 }
 
+// ListTunnels calls the injected ListTunnels or the real one.
+func (r *Robot) ListTunnels(ctx context.Context) ([]config.TrafficTunnelEndpoint, error) {
+	r.Mu.RLock()
+	defer r.Mu.RUnlock()
+	if r.ListTunnelsFunc == nil {
+		return r.LocalRobot.ListTunnels(ctx)
+	}
+	return r.ListTunnelsFunc(ctx)
+}
+
 type noopSessionManager struct{}
 
 func (m noopSessionManager) Start(ctx context.Context, ownerID string) (*session.Session, error) {