This commit is contained in:
@@ -13,26 +13,30 @@ import (
|
|||||||
"golang.org/x/sync/errgroup"
|
"golang.org/x/sync/errgroup"
|
||||||
"google.golang.org/grpc"
|
"google.golang.org/grpc"
|
||||||
"google.golang.org/grpc/credentials"
|
"google.golang.org/grpc/credentials"
|
||||||
|
"google.golang.org/grpc/metadata"
|
||||||
)
|
)
|
||||||
|
|
||||||
type CommanderClient struct {
|
type CommanderClient struct {
|
||||||
cmder *commander.Commander
|
cmder *commander.Commander
|
||||||
wg *sync.WaitGroup
|
wg *sync.WaitGroup
|
||||||
|
id string
|
||||||
}
|
}
|
||||||
|
|
||||||
func New(
|
func New(
|
||||||
cmder *commander.Commander,
|
cmder *commander.Commander,
|
||||||
wg *sync.WaitGroup,
|
wg *sync.WaitGroup,
|
||||||
|
id string,
|
||||||
) CommanderClient {
|
) CommanderClient {
|
||||||
return CommanderClient{cmder, wg}
|
return CommanderClient{cmder, wg, id}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (self *CommanderClient) HandleCommands(ctx context.Context, srvAddr string, tc credentials.TransportCredentials) error {
|
func (self *CommanderClient) HandleCommands(ctx context.Context, srvAddr string, tc credentials.TransportCredentials) error {
|
||||||
cli, err := grpc.NewClient(srvAddr, grpc.WithTransportCredentials(tc))
|
cli, err := grpc.NewClient(srvAddr, grpc.WithTransportCredentials(tc))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("Failed to connect to gRPC: %w", err)
|
return fmt.Errorf("Failed to connect to gRPC: %w", err)
|
||||||
}
|
}
|
||||||
ccli := proto.NewCommanderClient(cli)
|
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 {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
@@ -56,7 +60,6 @@ func (self *CommanderClient) recv(bidi grpc.BidiStreamingClient[proto.FinishedCo
|
|||||||
}
|
}
|
||||||
self.wg.Go(func() {
|
self.wg.Go(func() {
|
||||||
func() error {
|
func() error {
|
||||||
|
|
||||||
fc, err := self.cmder.Execute(msg)
|
fc, err := self.cmder.Execute(msg)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
|
|||||||
Reference in New Issue
Block a user