RSDK-14208 - automatically configure ports for tunneling

cli/client.go

@@ -3986,25 +3986,6 @@ func RobotsPartTunnelAction(ctx context.Context, cmd *cli.Command, args robotsPa
 }
 
 func tunnelTraffic(ctx context.Context, cmd *cli.Command, robotClient *client.RobotClient, local, dest int) error {
-	// don't block tunnel attempt if ListTunnels fails in any way - it may be unimplemented.
-	// TODO: early return if ListTunnels fails.
-	if tunnels, err := robotClient.ListTunnels(ctx); err == nil {
-		allowed := false
-		for _, t := range tunnels {
-			if t.Port == dest {
-				allowed = true
-				break
-			}
-		}
-		if !allowed {
-			return errors.Errorf(
-				"tunneling to destination port %v not allowed. "+
-					"Please ensure the traffic_tunnel_endpoints configuration is set correctly on the machine.",
-				dest,
-			)
-		}
-	}
-
 	li, err := net.Listen("tcp", net.JoinHostPort("localhost", strconv.Itoa(local)))
 	if err != nil {
 		return fmt.Errorf("failed to create listener %w", err)
@@ -4072,9 +4053,126 @@ func (c *viamClient) robotPartTunnel(ctx context.Context, cmd *cli.Command, args
 	if err != nil {
 		return err
 	}
+
+	if err := c.ensureTunnelPortAllowed(ctx, cmd, robotClient, args); err != nil {
+		return err
+	}
+
 	return tunnelTraffic(ctx, cmd, robotClient, args.LocalPort, args.DestinationPort)
 }
 
+// tunnelLister lists the tunnel endpoints a machine part is currently configured to
+// allow. Satisfied by *client.RobotClient; declared as an interface so the auto-add
+// flow can be tested without a live machine.
+type tunnelLister interface {
+	ListTunnels(ctx context.Context) ([]rconfig.TrafficTunnelEndpoint, error)
+}
+
+// tunnelPortAllowed reports whether the given destination port is present in the
+// machine part's configured tunnel endpoints. The second return value is false if
+// the tunnel list could not be retrieved (e.g. ListTunnels is unimplemented), in
+// which case callers should proceed without gating on the result.
+func tunnelPortAllowed(ctx context.Context, lister tunnelLister, dest int) (allowed, known bool) {
+	tunnels, err := lister.ListTunnels(ctx)
+	if err != nil {
+		return false, false
+	}
+	for _, t := range tunnels {
+		if t.Port == dest {
+			return true, true
+		}
+	}
+	return false, true
+}
+
+// ensureTunnelPortAllowed checks whether the destination port is configured as a
+// tunnel endpoint on the machine part. If it is not, and the user has permission to
+// edit the machine's config, the port is added to traffic_tunnel_endpoints
+// automatically and we wait for the machine to pick up the new config before
+// returning.
+func (c *viamClient) ensureTunnelPortAllowed(
+	ctx context.Context, cmd *cli.Command, lister tunnelLister, args robotsPartTunnelArgs,
+) error {
+	dest := args.DestinationPort
+
+	allowed, known := tunnelPortAllowed(ctx, lister, dest)
+	// If we couldn't read the tunnel list (e.g. ListTunnels is unimplemented on an
+	// older machine, or a transient error), don't gate on it. The auto-add flow below
+	// relies on ListTunnels to poll for the config taking effect, so without it we
+	// can't confirm an edit worked - and editing the machine config off an unreliable
+	// signal risks breaking a tunnel attempt that would otherwise have succeeded. Fall
+	// back to the previous behavior of just attempting the tunnel.
+	if allowed || !known {
+		return nil
+	}
+
+	infof(cmd.Root().Writer,
+		"destination port %v is not configured for tunneling; attempting to add it to the machine config...", dest)
+
+	part, err := c.robotPart(ctx, args.Organization, args.Location, args.Machine, args.Part)
+	if err != nil {
+		return errors.Wrapf(err,
+			"tunneling to destination port %v not allowed, and failed to look up the machine part to add it", dest)
+	}
+
+	config := part.RobotConfig.AsMap()
+	network, _ := config["network"].(map[string]any)
+	if network == nil {
+		network = map[string]any{}
+	}
+	endpoints, _ := network["traffic_tunnel_endpoints"].([]any)
+	endpoints = append(endpoints, map[string]any{"port": dest})
+	network["traffic_tunnel_endpoints"] = endpoints
+	config["network"] = network
+
+	pbConfig, err := protoutils.StructToStructPb(config)
+	if err != nil {
+		return err
+	}
+
+	if _, err := c.client.UpdateRobotPart(ctx, &apppb.UpdateRobotPartRequest{
+		Id:          part.Id,
+		Name:        part.Name,
+		RobotConfig: pbConfig,
+	}); err != nil {
+		if s, ok := status.FromError(err); ok && s.Code() == codes.PermissionDenied {
+			return errors.Errorf(
+				"tunneling to destination port %v not allowed. You do not have permission to edit this machine's config, "+
+					"so it could not be added automatically. Please ensure the traffic_tunnel_endpoints configuration is "+
+					"set correctly on the machine.",
+				dest,
+			)
+		}
+		return errors.Wrapf(err, "failed to add destination port %v to the machine's tunnel endpoints", dest)
+	}
+
+	infof(cmd.Root().Writer,
+		"added destination port %v to the machine config; waiting for the machine to apply the new config...", dest)
+
+	// wait for the machine to pick up the updated config and start allowing tunnels to
+	// the new port.
+	const (
+		tunnelConfigTimeout      = time.Minute
+		tunnelConfigPollInterval = 3 * time.Second
+	)
+	timeoutCtx, cancel := context.WithTimeout(ctx, tunnelConfigTimeout)
+	defer cancel()
+	for {
+		if allowed, _ := tunnelPortAllowed(timeoutCtx, lister, dest); allowed {
+			return nil
+		}
+		select {
+		case <-timeoutCtx.Done():
+			return errors.Errorf(
+				"destination port %v was added to the machine config, but the machine did not apply it within %s. "+
+					"Please try running the tunnel command again shortly.",
+				dest, tunnelConfigTimeout,
+			)
+		case <-time.After(tunnelConfigPollInterval):
+		}
+	}
+}
+
 // checkUpdateResponse holds the values used to hold release information.
 type getLatestReleaseResponse struct {
 	Name       string `json:"name"`

cli/client_test.go

@@ -1711,11 +1711,6 @@ func TestTunnelE2ECLI(t *testing.T) {
 	//nolint:dogsled
 	cCtx, _, _, _ := setup(nil, nil, nil, nil, "token")
 
-	// error early if tunnel not listed
-	err = tunnelTraffic(ctx, cCtx, rc, sourcePort, 1)
-	test.That(t, err, test.ShouldNotBeNil)
-	test.That(t, err.Error(), test.ShouldContainSubstring, "not allowed")
-
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
@@ -1751,6 +1746,166 @@ func TestTunnelE2ECLI(t *testing.T) {
 	wg.Wait()
 }
 
+// fakeTunnelLister is a tunnelLister test double. Before `reloadAfter` ListTunnels
+// calls it reports no configured ports (simulating a machine that has not yet picked
+// up a config change); after that it reports `ports`. If `err` is set, every call
+// fails with it.
+type fakeTunnelLister struct {
+	calls       int
+	reloadAfter int
+	ports       []int
+	err         error
+}
+
+func (f *fakeTunnelLister) ListTunnels(ctx context.Context) ([]robotconfig.TrafficTunnelEndpoint, error) {
+	f.calls++
+	if f.err != nil {
+		return nil, f.err
+	}
+	if f.calls <= f.reloadAfter {
+		return nil, nil
+	}
+	var out []robotconfig.TrafficTunnelEndpoint
+	for _, p := range f.ports {
+		out = append(out, robotconfig.TrafficTunnelEndpoint{Port: p})
+	}
+	return out, nil
+}
+
+func TestEnsureTunnelPortAllowed(t *testing.T) {
+	const (
+		partID       = "part-id"
+		partName     = "part-name"
+		existingPort = 9999
+		destPort     = 2222
+	)
+
+	// robotConfig builds a part config with the given tunnel endpoint ports already set.
+	robotConfig := func(ports ...int) *structpb.Struct {
+		endpoints := make([]any, 0, len(ports))
+		for _, p := range ports {
+			endpoints = append(endpoints, map[string]any{"port": float64(p)})
+		}
+		s, err := structpb.NewStruct(map[string]any{
+			"network": map[string]any{
+				"traffic_tunnel_endpoints": endpoints,
+			},
+		})
+		test.That(t, err, test.ShouldBeNil)
+		return s
+	}
+
+	// tunnelArgs targets destPort by part ID. An erroring ListOrganizations forces
+	// robotPart to fall back to a direct GetRobotPart lookup by ID.
+	tunnelArgs := robotsPartTunnelArgs{Part: partID, DestinationPort: destPort}
+	listOrganizationsFunc := func(ctx context.Context, in *apppb.ListOrganizationsRequest,
+		opts ...grpc.CallOption,
+	) (*apppb.ListOrganizationsResponse, error) {
+		return nil, errors.New("not used in this test")
+	}
+	getRobotPartFunc := func(ctx context.Context, in *apppb.GetRobotPartRequest,
+		opts ...grpc.CallOption,
+	) (*apppb.GetRobotPartResponse, error) {
+		return &apppb.GetRobotPartResponse{
+			Part: &apppb.RobotPart{Id: partID, Name: partName, RobotConfig: robotConfig(existingPort)},
+		}, nil
+	}
+
+	t.Run("port already allowed is a no-op", func(t *testing.T) {
+		var updateCalled bool
+		asc := &inject.AppServiceClient{
+			ListOrganizationsFunc: listOrganizationsFunc,
+			GetRobotPartFunc:      getRobotPartFunc,
+			UpdateRobotPartFunc: func(ctx context.Context, in *apppb.UpdateRobotPartRequest,
+				opts ...grpc.CallOption,
+			) (*apppb.UpdateRobotPartResponse, error) {
+				updateCalled = true
+				return &apppb.UpdateRobotPartResponse{}, nil
+			},
+		}
+		cCtx, ac, _, _ := setup(asc, nil, nil, nil, "token")
+		lister := &fakeTunnelLister{ports: []int{destPort}}
+
+		err := ac.ensureTunnelPortAllowed(context.Background(), cCtx, lister, tunnelArgs)
+		test.That(t, err, test.ShouldBeNil)
+		test.That(t, updateCalled, test.ShouldBeFalse)
+	})
+
+	t.Run("ListTunnels error proceeds without editing config", func(t *testing.T) {
+		var updateCalled bool
+		asc := &inject.AppServiceClient{
+			ListOrganizationsFunc: listOrganizationsFunc,
+			GetRobotPartFunc:      getRobotPartFunc,
+			UpdateRobotPartFunc: func(ctx context.Context, in *apppb.UpdateRobotPartRequest,
+				opts ...grpc.CallOption,
+			) (*apppb.UpdateRobotPartResponse, error) {
+				updateCalled = true
+				return &apppb.UpdateRobotPartResponse{}, nil
+			},
+		}
+		cCtx, ac, _, _ := setup(asc, nil, nil, nil, "token")
+		lister := &fakeTunnelLister{err: errors.New("unimplemented")}
+
+		err := ac.ensureTunnelPortAllowed(context.Background(), cCtx, lister, tunnelArgs)
+		test.That(t, err, test.ShouldBeNil)
+		test.That(t, updateCalled, test.ShouldBeFalse)
+	})
+
+	t.Run("auto-adds missing port and waits for reload", func(t *testing.T) {
+		var capturedReq *apppb.UpdateRobotPartRequest
+		asc := &inject.AppServiceClient{
+			ListOrganizationsFunc: listOrganizationsFunc,
+			GetRobotPartFunc:      getRobotPartFunc,
+			UpdateRobotPartFunc: func(ctx context.Context, in *apppb.UpdateRobotPartRequest,
+				opts ...grpc.CallOption,
+			) (*apppb.UpdateRobotPartResponse, error) {
+				capturedReq = in
+				return &apppb.UpdateRobotPartResponse{}, nil
+			},
+		}
+		cCtx, ac, _, _ := setup(asc, nil, nil, nil, "token")
+		// Not allowed on the initial check; allowed once the machine "reloads".
+		lister := &fakeTunnelLister{reloadAfter: 1, ports: []int{existingPort, destPort}}
+
+		err := ac.ensureTunnelPortAllowed(context.Background(), cCtx, lister, tunnelArgs)
+		test.That(t, err, test.ShouldBeNil)
+
+		// The update should preserve the existing port and append the new one.
+		test.That(t, capturedReq, test.ShouldNotBeNil)
+		test.That(t, capturedReq.Id, test.ShouldEqual, partID)
+		test.That(t, capturedReq.Name, test.ShouldEqual, partName)
+		network, ok := capturedReq.RobotConfig.AsMap()["network"].(map[string]any)
+		test.That(t, ok, test.ShouldBeTrue)
+		endpoints, ok := network["traffic_tunnel_endpoints"].([]any)
+		test.That(t, ok, test.ShouldBeTrue)
+		var ports []int
+		for _, e := range endpoints {
+			ports = append(ports, int(e.(map[string]any)["port"].(float64)))
+		}
+		test.That(t, ports, test.ShouldContain, existingPort)
+		test.That(t, ports, test.ShouldContain, destPort)
+	})
+
+	t.Run("permission denied surfaces a friendly error", func(t *testing.T) {
+		asc := &inject.AppServiceClient{
+			ListOrganizationsFunc: listOrganizationsFunc,
+			GetRobotPartFunc:      getRobotPartFunc,
+			UpdateRobotPartFunc: func(ctx context.Context, in *apppb.UpdateRobotPartRequest,
+				opts ...grpc.CallOption,
+			) (*apppb.UpdateRobotPartResponse, error) {
+				return nil, status.Error(codes.PermissionDenied, "nope")
+			},
+		}
+		cCtx, ac, _, _ := setup(asc, nil, nil, nil, "token")
+		lister := &fakeTunnelLister{} // never reports the port
+
+		err := ac.ensureTunnelPortAllowed(context.Background(), cCtx, lister, tunnelArgs)
+		test.That(t, err, test.ShouldNotBeNil)
+		test.That(t, err.Error(), test.ShouldContainSubstring, "permission")
+		test.That(t, err.Error(), test.ShouldContainSubstring, "not allowed")
+	})
+}
+
 func TestAPIToGRPCServiceName(t *testing.T) {
 	tests := []struct {
 		name     string