From 52d81fd8148d311a97acdb45cadd53fcd0023fdc Mon Sep 17 00:00:00 2001 From: Ethan Rodkin Date: Tue, 7 Jul 2026 14:34:57 -0400 Subject: [PATCH] automatically configure ports for tunneling --- cli/client.go | 136 +++++++++++++++++++++++++++++++------ cli/client_test.go | 165 +++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 277 insertions(+), 24 deletions(-) diff --git a/cli/client.go b/cli/client.go index ef03f596a37..d20b711c1ba 100644 --- a/cli/client.go +++ b/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"` diff --git a/cli/client_test.go b/cli/client_test.go index 6b80232a0f2..85df07196e9 100644 --- a/cli/client_test.go +++ b/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