driver(external): implement bidirectional streaming for GuestAgentConn()

Signed-off-by: Ansuman Sahoo <[email protected]>

Author: unsuman

pkg/driver/external/client/methods.go

@@ -253,7 +253,7 @@ func (d *DriverClient) ForwardGuestAgent() bool {
 func (d *DriverClient) GuestAgentConn(ctx context.Context) (net.Conn, error) {
 	d.logger.Debug("Getting guest agent connection")
 
-	connStream, err := d.DriverSvc.GuestAgentConn(ctx, &emptypb.Empty{})
+	connStream, err := d.DriverSvc.GuestAgentConn(ctx)
 	if err != nil {
 		d.logger.Errorf("Failed to get guest agent connection: %v", err)
 		return nil, err

pkg/driver/external/driver.pb.go

@@ -24,7 +24,7 @@ service Driver {
     rpc Unregister(google.protobuf.Empty) returns (google.protobuf.Empty);
 
     rpc ForwardGuestAgent(google.protobuf.Empty) returns (ForwardGuestAgentResponse);
-    rpc GuestAgentConn(google.protobuf.Empty) returns (stream GuestAgentConnResponse);
+    rpc GuestAgentConn(stream GuestAgentConnStream) returns (stream GuestAgentConnStream);
 
     rpc SetConfig(SetConfigRequest) returns (google.protobuf.Empty);
 
@@ -73,6 +73,6 @@ message ForwardGuestAgentResponse {
     bool should_forward = 1;
 }
 
-message GuestAgentConnResponse {
+message GuestAgentConnStream {
     bytes net_conn = 1;
 }

pkg/driver/external/driver.proto

@@ -64,19 +64,31 @@ func (s *DriverServer) SetConfig(ctx context.Context, req *pb.SetConfigRequest)
 	return &emptypb.Empty{}, nil
 }
 
-// NOTE: If this doesn't work, try with bidirectional streaming(like in buildkit) and see if that helps.
-func (s *DriverServer) GuestAgentConn(empty *emptypb.Empty, stream pb.Driver_GuestAgentConnServer) error {
+func (s *DriverServer) GuestAgentConn(stream pb.Driver_GuestAgentConnServer) error {
 	s.logger.Debug("Received GuestAgentConn request")
 	conn, err := s.driver.GuestAgentConn(stream.Context())
 	if err != nil {
 		s.logger.Errorf("GuestAgentConn failed: %v", err)
 		return err
 	}
 
-	defer conn.Close()
+	go func() {
+		for {
+			msg, err := stream.Recv()
+			if err != nil {
+				return
+			}
+			if len(msg.NetConn) > 0 {
+				_, err = conn.Write(msg.NetConn)
+				if err != nil {
+					s.logger.Errorf("Error writing to connection: %v", err)
+					return
+				}
+			}
+		}
+	}()
 
 	buf := make([]byte, 32*1<<10)
-
 	for {
 		n, err := conn.Read(buf)
 		if err != nil {
@@ -86,7 +98,7 @@ func (s *DriverServer) GuestAgentConn(empty *emptypb.Empty, stream pb.Driver_Gue
 			return status.Errorf(codes.Internal, "error reading: %v", err)
 		}
 
-		msg := &pb.GuestAgentConnResponse{NetConn: buf[:n]}
+		msg := &pb.GuestAgentConnStream{NetConn: buf[:n]}
 		if err := stream.Send(msg); err != nil {
 			return err
 		}

pkg/driver/external/driver_grpc.pb.go

@@ -59,7 +59,7 @@ const (
 	HostAgentSock           = "ha.sock"
 	HostAgentStdoutLog      = "ha.stdout.log"
 	HostAgentStderrLog      = "ha.stderr.log"
-	ExternalDriverStderrLog = "ed.stderr.log"
+	ExternalDriverStderrLog = "driver.stderr.log"
 	VzIdentifier            = "vz-identifier"
 	VzEfi                   = "vz-efi"           // efi variable store
 	QemuEfiCodeFD           = "qemu-efi-code.fd" // efi code; not always created