updates

server/internal/grpc/pb/keymanager.pb.go

@@ -45,12 +45,87 @@ type UploadKeyResponse struct {
 	KeyId string `json:"key_id"`
 }
 
+// CommandStream message types
+
+type ServerCommand struct {
+	CommandId   string          `json:"command_id"`
+	GenerateKey *GenerateKeyCmd `json:"generate_key,omitempty"`
+}
+
+type GenerateKeyCmd struct {
+	Label string `json:"label"`
+}
+
+type AgentMessage struct {
+	ServerId   string         `json:"server_id"`
+	AgentToken string         `json:"agent_token"`
+	Ready      *AgentReady    `json:"ready,omitempty"`
+	Result     *CommandResult `json:"result,omitempty"`
+}
+
+type AgentReady struct{}
+
+type CommandResult struct {
+	CommandId string `json:"command_id"`
+	Success   bool   `json:"success"`
+	Message   string `json:"message"`
+}
+
+// CommandStream server-side interface
+
+type KeyManager_CommandStreamServer interface {
+	Send(*ServerCommand) error
+	Recv() (*AgentMessage, error)
+	grpc.ServerStream
+}
+
+type keyManagerCommandStreamServer struct {
+	grpc.ServerStream
+}
+
+func (s *keyManagerCommandStreamServer) Send(m *ServerCommand) error {
+	return s.ServerStream.SendMsg(m)
+}
+
+func (s *keyManagerCommandStreamServer) Recv() (*AgentMessage, error) {
+	m := new(AgentMessage)
+	if err := s.ServerStream.RecvMsg(m); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
+// CommandStream client-side interface
+
+type KeyManager_CommandStreamClient interface {
+	Send(*AgentMessage) error
+	Recv() (*ServerCommand, error)
+	grpc.ClientStream
+}
+
+type keyManagerCommandStreamClient struct {
+	grpc.ClientStream
+}
+
+func (c *keyManagerCommandStreamClient) Send(m *AgentMessage) error {
+	return c.ClientStream.SendMsg(m)
+}
+
+func (c *keyManagerCommandStreamClient) Recv() (*ServerCommand, error) {
+	m := new(ServerCommand)
+	if err := c.ClientStream.RecvMsg(m); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
 // Server interface
 
 type KeyManagerServer interface {
 	Register(context.Context, *RegisterRequest) (*RegisterResponse, error)
 	SyncKeys(context.Context, *SyncRequest) (*SyncResponse, error)
 	UploadGeneratedKey(context.Context, *UploadKeyRequest) (*UploadKeyResponse, error)
+	CommandStream(KeyManager_CommandStreamServer) error
 }
 
 type UnimplementedKeyManagerServer struct{}
@@ -67,12 +142,17 @@ func (UnimplementedKeyManagerServer) UploadGeneratedKey(context.Context, *Upload
 	return nil, status.Errorf(codes.Unimplemented, "method UploadGeneratedKey not implemented")
 }
 
+func (UnimplementedKeyManagerServer) CommandStream(KeyManager_CommandStreamServer) error {
+	return status.Errorf(codes.Unimplemented, "method CommandStream not implemented")
+}
+
 // Client interface
 
 type KeyManagerClient interface {
 	Register(ctx context.Context, in *RegisterRequest, opts ...grpc.CallOption) (*RegisterResponse, error)
 	SyncKeys(ctx context.Context, in *SyncRequest, opts ...grpc.CallOption) (*SyncResponse, error)
 	UploadGeneratedKey(ctx context.Context, in *UploadKeyRequest, opts ...grpc.CallOption) (*UploadKeyResponse, error)
+	CommandStream(ctx context.Context, opts ...grpc.CallOption) (KeyManager_CommandStreamClient, error)
 }
 
 type keyManagerClient struct {
@@ -107,6 +187,14 @@ func (c *keyManagerClient) UploadGeneratedKey(ctx context.Context, in *UploadKey
 	return out, nil
 }
 
+func (c *keyManagerClient) CommandStream(ctx context.Context, opts ...grpc.CallOption) (KeyManager_CommandStreamClient, error) {
+	stream, err := c.cc.NewStream(ctx, &KeyManager_ServiceDesc.Streams[0], "/keymanager.v1.KeyManager/CommandStream", opts...)
+	if err != nil {
+		return nil, err
+	}
+	return &keyManagerCommandStreamClient{stream}, nil
+}
+
 // Server registration
 
 func RegisterKeyManagerServer(s grpc.ServiceRegistrar, srv KeyManagerServer) {
@@ -121,7 +209,14 @@ var KeyManager_ServiceDesc = grpc.ServiceDesc{
 		{MethodName: "SyncKeys", Handler: _KeyManager_SyncKeys_Handler},
 		{MethodName: "UploadGeneratedKey", Handler: _KeyManager_UploadGeneratedKey_Handler},
 	},
-	Streams:  []grpc.StreamDesc{},
+	Streams: []grpc.StreamDesc{
+		{
+			StreamName:    "CommandStream",
+			Handler:       _KeyManager_CommandStream_Handler,
+			ServerStreams: true,
+			ClientStreams: true,
+		},
+	},
 	Metadata: "keymanager/v1/keymanager.proto",
 }
 
@@ -169,3 +264,7 @@ func _KeyManager_UploadGeneratedKey_Handler(srv interface{}, ctx context.Context
 	}
 	return interceptor(ctx, in, info, handler)
 }
+
+func _KeyManager_CommandStream_Handler(srv interface{}, stream grpc.ServerStream) error {
+	return srv.(KeyManagerServer).CommandStream(&keyManagerCommandStreamServer{stream})
+}

server/internal/grpc/server.go

@@ -67,6 +67,59 @@ func (s *keyManagerServer) UploadGeneratedKey(ctx context.Context, req *pb.Uploa
 	return &pb.UploadKeyResponse{KeyId: key.KeyID}, nil
 }
 
+func (s *keyManagerServer) CommandStream(stream pb.KeyManager_CommandStreamServer) error {
+	// First message authenticates the agent and signals readiness.
+	msg, err := stream.Recv()
+	if err != nil {
+		return status.Errorf(codes.InvalidArgument, "expected initial auth message: %v", err)
+	}
+
+	srv, err := services.ValidateAgentToken(msg.ServerId, msg.AgentToken)
+	if err != nil {
+		return status.Errorf(codes.Unauthenticated, "invalid agent token")
+	}
+
+	if err := services.UpdateServerLastSeen(srv.ServerID); err != nil {
+		log.Printf("update last seen %s: %v", srv.ServerID, err)
+	}
+
+	ch := services.Dispatcher.Connect(srv.ServerID)
+	defer services.Dispatcher.Disconnect(srv.ServerID)
+
+	log.Printf("agent %s connected command stream", srv.ServerID)
+	defer log.Printf("agent %s disconnected command stream", srv.ServerID)
+
+	// Drain inbound results in the background so client Send calls never block.
+	// UploadGeneratedKey handles the real storage; these are just confirmation logs.
+	go func() {
+		for {
+			m, err := stream.Recv()
+			if err != nil {
+				return
+			}
+			if m.Result != nil {
+				r := m.Result
+				log.Printf("agent %s cmd %s: success=%v %s", srv.ServerID, r.CommandId, r.Success, r.Message)
+			}
+		}
+	}()
+
+	ctx := stream.Context()
+	for {
+		select {
+		case <-ctx.Done():
+			return nil
+		case cmd, ok := <-ch:
+			if !ok {
+				return nil
+			}
+			if err := stream.Send(cmd); err != nil {
+				return err
+			}
+		}
+	}
+}
+
 func StartGRPC(port int) error {
 	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
 	if err != nil {