diff --git a/agent/internal/client/commander.go b/agent/internal/client/commander.go index 681755d..340975d 100644 --- a/agent/internal/client/commander.go +++ b/agent/internal/client/commander.go @@ -13,26 +13,30 @@ import ( "golang.org/x/sync/errgroup" "google.golang.org/grpc" "google.golang.org/grpc/credentials" + "google.golang.org/grpc/metadata" ) type CommanderClient struct { cmder *commander.Commander wg *sync.WaitGroup + id string } func New( cmder *commander.Commander, wg *sync.WaitGroup, + id string, ) CommanderClient { - return CommanderClient{cmder, wg} + return CommanderClient{cmder, wg, id} } + func (self *CommanderClient) HandleCommands(ctx context.Context, srvAddr string, tc credentials.TransportCredentials) error { cli, err := grpc.NewClient(srvAddr, grpc.WithTransportCredentials(tc)) if err != nil { return fmt.Errorf("Failed to connect to gRPC: %w", err) } ccli := proto.NewCommanderClient(cli) - bidi, err := ccli.Stream(ctx) + bidi, err := ccli.Stream(metadata.NewOutgoingContext(ctx, metadata.MD{"agentid": []string{self.id}})) if err != nil { return err } @@ -56,7 +60,6 @@ func (self *CommanderClient) recv(bidi grpc.BidiStreamingClient[proto.FinishedCo } self.wg.Go(func() { func() error { - fc, err := self.cmder.Execute(msg) if err != nil { return err