restarting gRPC conn

Author: sean-breen

internal/command/command_plugin.go

@@ -93,6 +93,15 @@ func (cp *CommandPlugin) processResourceUpdate(ctx context.Context, msg *bus.Mes
 	if resource, ok := msg.Data.(*mpi.Resource); ok {
 		if !cp.commandService.IsConnected() && cp.config.IsFeatureEnabled(pkgConfig.FeatureConnection) {
 			cp.createConnection(ctx, resource)
+
+			newConn, err := cp.conn.Restart(ctx)
+			if err != nil {
+				slog.ErrorContext(ctx, "Failed to restart connection", "error", err)
+				return
+			}
+			cp.conn = newConn
+			cp.createConnection(ctx, resource)
+
 		} else {
 			statusErr := cp.commandService.UpdateDataPlaneStatus(ctx, resource)
 			if statusErr != nil {

internal/config/config.go

@@ -302,6 +302,12 @@ func registerCommandFlags(fs *flag.FlagSet) {
 		DefCommandAuthTokenKey,
 		"The token used in the authentication handshake with the command server endpoint for command and control.",
 	)
+	fs.String(
+		CommandAuthTokenPathKey,
+		DefCommandAuthTokenPathKey,
+		"The path to the file containing the token used in the authentication handshake with "+
+			"the command server endpoint for command and control.",
+	)
 	fs.String(
 		CommandTLSCertKey,
 		DefCommandTLSCertKey,
@@ -816,9 +822,10 @@ func resolveCommand() *Command {
 		},
 	}
 
-	if viperInstance.IsSet(CommandAuthTokenKey) {
+	if viperInstance.IsSet(CommandAuthTokenKey) || viperInstance.IsSet(CommandAuthTokenPathKey) {
 		command.Auth = &AuthConfig{
-			Token: viperInstance.GetString(CommandAuthTokenKey),
+			Token:     viperInstance.GetString(CommandAuthTokenKey),
+			TokenPath: viperInstance.GetString(CommandAuthTokenPathKey),
 		}
 	}
 

internal/grpc/grpc.go

@@ -41,6 +41,7 @@ type (
 		CommandServiceClient() mpi.CommandServiceClient
 		FileServiceClient() mpi.FileServiceClient
 		Close(ctx context.Context) error
+		Restart(ctx context.Context) (*GrpcConnection, error)
 	}
 
 	GrpcConnection struct {
@@ -130,6 +131,22 @@ func (gc *GrpcConnection) Close(ctx context.Context) error {
 	return nil
 }
 
+func (gc *GrpcConnection) Restart(ctx context.Context) (*GrpcConnection, error) {
+	slog.InfoContext(ctx, "Restarting grpc connection")
+	err := gc.Close(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	slog.InfoContext(ctx, "Creating grpc connection")
+	newConn, err := NewGrpcConnection(ctx, gc.config)
+	if err != nil {
+		return nil, err
+	}
+
+	return newConn, nil
+}
+
 func (w *wrappedStream) RecvMsg(message any) error {
 	err := w.ClientStream.RecvMsg(message)
 	if err == nil {
@@ -249,7 +266,7 @@ func addPerRPCCredentials(agentConfig *config.Config, resourceID string, opts []
 		}
 	}
 
-	slog.Debug("Adding token to RPC credentials")
+	slog.Debug("Adding token to RPC credentials", "token", token)
 	opts = append(opts,
 		grpc.WithPerRPCCredentials(
 			&PerRPCCredentials{