From 00d68d9465220fc12ec48b87b455c8297e16b178 Mon Sep 17 00:00:00 2001 From: Pasha Sviderski Date: Fri, 28 Aug 2026 20:37:50 +1000 Subject: [PATCH] refactor(client): make ProxySingleMachineContext and ProxyMachinesContext package functions as well --- cmd/uc/ps.go | 2 +- pkg/client/client.go | 27 +++++++++++-- pkg/client/client_test.go | 79 +++++++++++++++++++++++++++++++++++++++ pkg/client/container.go | 8 ++-- pkg/client/image.go | 2 +- pkg/client/logs.go | 4 +- pkg/client/machine.go | 2 +- pkg/client/volume.go | 6 +-- 8 files changed, 114 insertions(+), 16 deletions(-) create mode 100644 pkg/client/client_test.go diff --git a/cmd/uc/ps.go b/cmd/uc/ps.go index 00f663b9..249f6e4b 100644 --- a/cmd/uc/ps.go +++ b/cmd/uc/ps.go @@ -187,7 +187,7 @@ func printContainers(containers []containerInfo) error { } func collectContainers(ctx context.Context, cli *client.Client) ([]containerInfo, error) { - listCtx := cli.ProxyMachinesContext(ctx, nil) + listCtx := client.ProxyMachinesContext(ctx, nil) // List all service containers across all machines in the cluster. machineContainers, err := cli.Docker.ListServiceContainers( diff --git a/pkg/client/client.go b/pkg/client/client.go index a0faf69c..f8926522 100644 --- a/pkg/client/client.go +++ b/pkg/client/client.go @@ -78,8 +78,8 @@ func (cli *Client) progressOut() *streams.Out { // ProxyMachinesContext returns a new context that proxies gRPC requests to the specified machines. // If namesOrIDs is nil or empty, all machines are included. // This triggers One2Many proxying, which always injects metadata into the response. -func (cli *Client) ProxyMachinesContext(ctx context.Context, namesOrIDs []string) context.Context { - md := metadata.New(nil) +func ProxyMachinesContext(ctx context.Context, namesOrIDs []string) context.Context { + md := outgoingMetadataWithoutProxyTargets(ctx) if len(namesOrIDs) == 0 { md.Append("machines", "*") } else { @@ -92,7 +92,26 @@ func (cli *Client) ProxyMachinesContext(ctx context.Context, namesOrIDs []string // ProxySingleMachineContext returns a new context that proxies gRPC requests to a single specified machine. // This triggers One2One proxying, which does NOT inject metadata into the response. // Use this for requests that expect a single response message without metadata wrapper. -func (cli *Client) ProxySingleMachineContext(ctx context.Context, nameOrID string) context.Context { - md := metadata.Pairs("machine", nameOrID) +func ProxySingleMachineContext(ctx context.Context, nameOrID string) context.Context { + md := outgoingMetadataWithoutProxyTargets(ctx) + md.Set("machine", nameOrID) return metadata.NewOutgoingContext(ctx, md) } + +func outgoingMetadataWithoutProxyTargets(ctx context.Context) metadata.MD { + md, _ := metadata.FromOutgoingContext(ctx) + md = md.Copy() + md.Delete("machine") + md.Delete("machines") + return md +} + +// ProxyMachinesContext returns a new context that proxies gRPC requests to the specified machines. +func (cli *Client) ProxyMachinesContext(ctx context.Context, namesOrIDs []string) context.Context { + return ProxyMachinesContext(ctx, namesOrIDs) +} + +// ProxySingleMachineContext returns a new context that proxies gRPC requests to a single specified machine. +func (cli *Client) ProxySingleMachineContext(ctx context.Context, nameOrID string) context.Context { + return ProxySingleMachineContext(ctx, nameOrID) +} diff --git a/pkg/client/client_test.go b/pkg/client/client_test.go new file mode 100644 index 00000000..6aa875ea --- /dev/null +++ b/pkg/client/client_test.go @@ -0,0 +1,79 @@ +package client + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + "google.golang.org/grpc/metadata" +) + +func TestProxySingleMachineContext(t *testing.T) { + original := metadata.Pairs( + "authorization", "token", + "machine", "old-machine", + "machines", "old-machine-a", + "machines", "old-machine-b", + ) + ctx := metadata.NewOutgoingContext(context.Background(), original) + + proxyCtx := ProxySingleMachineContext(ctx, "new-machine") + + md, ok := metadata.FromOutgoingContext(proxyCtx) + require.True(t, ok) + require.Equal(t, metadata.Pairs( + "authorization", "token", + "machine", "new-machine", + ), md) + require.Equal(t, metadata.Pairs( + "authorization", "token", + "machine", "old-machine", + "machines", "old-machine-a", + "machines", "old-machine-b", + ), original) +} + +func TestProxyMachinesContext(t *testing.T) { + tests := []struct { + name string + machines []string + want []string + }{ + { + name: "specified machines", + machines: []string{"machine-a", "machine-b"}, + want: []string{"machine-a", "machine-b"}, + }, + { + name: "all machines", + want: []string{"*"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + original := metadata.Pairs( + "authorization", "token", + "machine", "old-machine", + "machines", "old-machine-a", + "machines", "old-machine-b", + ) + ctx := metadata.NewOutgoingContext(context.Background(), original) + + proxyCtx := ProxyMachinesContext(ctx, tt.machines) + + md, ok := metadata.FromOutgoingContext(proxyCtx) + require.True(t, ok) + require.Equal(t, metadata.MD{ + "authorization": {"token"}, + "machines": tt.want, + }, md) + require.Equal(t, metadata.Pairs( + "authorization", "token", + "machine", "old-machine", + "machines", "old-machine-a", + "machines", "old-machine-b", + ), original) + }) + } +} diff --git a/pkg/client/container.go b/pkg/client/container.go index ae582b8c..ccf5eb9b 100644 --- a/pkg/client/container.go +++ b/pkg/client/container.go @@ -74,7 +74,7 @@ func (cli *Client) createServiceContainerWithPull( resp.Name = containerName // Proxy Docker gRPC requests to the selected machine. - ctx = cli.ProxySingleMachineContext(ctx, machine.Machine.Id) + ctx = ProxySingleMachineContext(ctx, machine.Machine.Id) pw := progress.ContextWriter(ctx) eventID := cliprogress.NewContainerEventID(ctx, containerName, machine.Machine.Name) @@ -277,7 +277,7 @@ func (cli *Client) resolveContainerOperation( eventID := cliprogress.ContainerEventID(ctx, ctr.Container.ServiceSpec.Name, ctr.Container.ID, ctr.MachineName) return containerOperationContext{ - ctx: cli.ProxySingleMachineContext(ctx, ctr.MachineID), + ctx: ProxySingleMachineContext(ctx, ctr.MachineID), containerID: ctr.Container.ID, eventID: eventID, }, nil @@ -375,7 +375,7 @@ func (cli *Client) ExecContainer( } // Proxy Docker gRPC requests to the machine hosting the container - ctx = cli.ProxySingleMachineContext(ctx, machine.Machine.Id) + ctx = ProxySingleMachineContext(ctx, machine.Machine.Id) // Execute the command in the container exitCode, err := cli.Docker.ExecContainer(ctx, machinedocker.ExecConfig{ @@ -452,7 +452,7 @@ func (cli *Client) WaitContainerHealthy( } // For containers with a health check, wait until Docker reports healthy or unhealthy. - mctx := cli.ProxySingleMachineContext(ctx, machine.Machine.Id) + mctx := ProxySingleMachineContext(ctx, machine.Machine.Id) mctx, cancel := context.WithTimeout(mctx, healthcheckTimeout(mc.Container.Config.Healthcheck)) defer cancel() ticker := time.NewTicker(1 * time.Second) diff --git a/pkg/client/image.go b/pkg/client/image.go index 5a20da7e..2b05678e 100644 --- a/pkg/client/image.go +++ b/pkg/client/image.go @@ -59,7 +59,7 @@ func (cli *Client) InspectRemoteImage(ctx context.Context, id string) ([]api.Mac // it lists images on all machines. func (cli *Client) ListImages(ctx context.Context, filter api.ImageFilter) ([]api.MachineImages, error) { // Broadcast the image list request to the specified machines or all machines if none specified. - listCtx := cli.ProxyMachinesContext(ctx, filter.Machines) + listCtx := ProxyMachinesContext(ctx, filter.Machines) opts := image.ListOptions{Manifests: true} if filter.Name != "" { diff --git a/pkg/client/logs.go b/pkg/client/logs.go index 0d7874c1..a11d0757 100644 --- a/pkg/client/logs.go +++ b/pkg/client/logs.go @@ -102,7 +102,7 @@ func (cli *Client) ServiceLogs( func (cli *Client) ContainerLogs( ctx context.Context, machineNameOrID string, containerID string, opts api.ServiceLogsOptions, ) (<-chan api.LogEntry, error) { - proxyCtx := cli.ProxySingleMachineContext(ctx, machineNameOrID) + proxyCtx := ProxySingleMachineContext(ctx, machineNameOrID) req := &pb.LogsRequest{ Id: containerID, @@ -198,7 +198,7 @@ func (cli *Client) MachineLogs( func (cli *Client) systemServiceLogs( ctx context.Context, machineID, service string, opts api.ServiceLogsOptions, ) (<-chan api.LogEntry, error) { - proxyCtx := cli.ProxySingleMachineContext(ctx, machineID) + proxyCtx := ProxySingleMachineContext(ctx, machineID) req := &pb.LogsRequest{ Id: service, diff --git a/pkg/client/machine.go b/pkg/client/machine.go index 313217fe..ef339c03 100644 --- a/pkg/client/machine.go +++ b/pkg/client/machine.go @@ -78,7 +78,7 @@ func (cli *Client) ListMachines(ctx context.Context, filter *api.MachineFilter) func (cli *Client) UpdateMachine( ctx context.Context, nameOrID string, req *pb.UpdateMachineRequest, ) (*pb.MachineInfo, error) { - ctx = cli.ProxySingleMachineContext(ctx, nameOrID) + ctx = ProxySingleMachineContext(ctx, nameOrID) resp, err := cli.MachineClient.UpdateMachine(ctx, req) if err != nil { if s, ok := status.FromError(err); ok && s.Code() == codes.NotFound { diff --git a/pkg/client/volume.go b/pkg/client/volume.go index abe0e90f..f8b8de3a 100644 --- a/pkg/client/volume.go +++ b/pkg/client/volume.go @@ -27,7 +27,7 @@ func (cli *Client) CreateVolume( return resp, fmt.Errorf("inspect machine '%s': %w", machineNameOrID, err) } // Proxy Docker gRPC requests to the selected machine. - ctx = cli.ProxySingleMachineContext(ctx, machine.Machine.Id) + ctx = ProxySingleMachineContext(ctx, machine.Machine.Id) pw := progress.ContextWriter(ctx) eventID := cliprogress.VolumeEventID(opts.Name, machine.Machine.Name) @@ -56,7 +56,7 @@ func (cli *Client) ListVolumes(ctx context.Context, filter *api.VolumeFilter) ([ proxyMachines = filter.Machines } - listCtx := cli.ProxyMachinesContext(ctx, proxyMachines) + listCtx := ProxyMachinesContext(ctx, proxyMachines) machineVolumes, err := cli.Docker.ListVolumes(listCtx, volume.ListOptions{}) if err != nil { return nil, err @@ -107,7 +107,7 @@ func (cli *Client) RemoveVolume(ctx context.Context, machineNameOrID, volumeName return fmt.Errorf("inspect machine '%s': %w", machineNameOrID, err) } // Proxy Docker gRPC requests to the selected machine. - ctx = cli.ProxySingleMachineContext(ctx, machine.Machine.Id) + ctx = ProxySingleMachineContext(ctx, machine.Machine.Id) pw := progress.ContextWriter(ctx) eventID := cliprogress.VolumeEventID(volumeName, machine.Machine.Name)