driver(external): revamp grpchijack and GuestAgentConn()

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

Author: unsuman

cmd/limactl/main_darwin.go

@@ -6,4 +6,4 @@
 package main
 
 // Import vz driver to register it in the registry on darwin.
-import _ "github.com/lima-vm/lima/pkg/driver/vz"
+// import _ "github.com/lima-vm/lima/pkg/driver/vz"

pkg/driver/external/client/client.go

@@ -8,13 +8,11 @@ import (
 	"io"
 	"math"
 	"net"
-	"time"
 
 	pb "github.com/lima-vm/lima/pkg/driver/external"
 	"github.com/sirupsen/logrus"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
-	"google.golang.org/grpc/keepalive"
 )
 
 type DriverClient struct {
@@ -36,18 +34,15 @@ func NewDriverClient(stdin io.WriteCloser, stdout io.ReadCloser, logger *logrus.
 			return pipeConn, nil
 		}),
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
-		grpc.WithKeepaliveParams(keepalive.ClientParameters{
-			Time:                10 * time.Second,
-			Timeout:             20 * time.Second,
-			PermitWithoutStream: true,
-		}),
 	}
 
 	// conn, err := grpc.NewClient("pipe", opts...)
 	// if err != nil {
 	// 	logger.Errorf("failed to create gRPC driver client connection: %v", err)
 	// 	return nil, err
 	// }
+	// -> ERRO[2025-06-04T21:32:54+05:30] Failed to set config: rpc error: code =
+	// Unavailable desc = name resolver error: produced zero addresses
 
 	conn, err := grpc.Dial("pipe", opts...)
 	if err != nil {

pkg/driver/external/client/grpchijack/hijack.go

@@ -0,0 +1,135 @@
+// SPDX-FileCopyrightText: Copyright The Lima Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package grpchijack
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"sync"
+	"time"
+
+	pb "github.com/lima-vm/lima/pkg/driver/external"
+	"google.golang.org/grpc"
+)
+
+type streamConn struct {
+	stream     pb.Driver_GuestAgentConnClient
+	readBuf    []byte
+	lastBuf    []byte
+	readMu     sync.Mutex
+	writeMu    sync.Mutex
+	closeCh    chan struct{}
+	closedOnce sync.Once
+	closed     bool
+}
+
+func StreamToConn(stream pb.Driver_GuestAgentConnClient) *streamConn {
+	return &streamConn{
+		stream:  stream,
+		closeCh: make(chan struct{}),
+	}
+}
+
+func (c *streamConn) Read(b []byte) (int, error) {
+	c.readMu.Lock()
+	defer c.readMu.Unlock()
+
+	if c.closed {
+		return 0, io.EOF
+	}
+
+	if c.lastBuf != nil {
+		n := copy(b, c.lastBuf)
+		c.lastBuf = c.lastBuf[n:]
+		if len(c.lastBuf) == 0 {
+			c.lastBuf = nil
+		}
+		return n, nil
+	}
+
+	msg := new(pb.BytesMessage)
+	msg, err := c.stream.Recv()
+	if err != nil {
+		c.closed = true
+		if err == io.EOF {
+			return 0, io.EOF
+		}
+		return 0, fmt.Errorf("stream receive error: %w", err)
+	}
+
+	n := copy(b, msg.Data)
+	if n < len(msg.Data) {
+		c.readBuf = make([]byte, len(msg.Data)-n)
+		copy(c.readBuf, msg.Data[n:])
+	}
+
+	return n, nil
+}
+
+func (c *streamConn) Write(b []byte) (int, error) {
+	c.writeMu.Lock()
+	defer c.writeMu.Unlock()
+
+	if c.closed {
+		return 0, errors.New("connection closed")
+	}
+
+	err := c.stream.Send(&pb.BytesMessage{Data: b})
+	if err != nil {
+		c.closed = true
+		return 0, fmt.Errorf("stream send error: %w", err)
+	}
+
+	return len(b), nil
+}
+
+func (c *streamConn) Close() error {
+	c.closedOnce.Do(func() {
+		defer func() {
+			close(c.closeCh)
+		}()
+
+		if cs, ok := c.stream.(grpc.ClientStream); ok {
+			c.writeMu.Lock()
+			err := cs.CloseSend()
+			c.writeMu.Unlock()
+			if err != nil {
+				return
+			}
+		}
+
+		c.readMu.Lock()
+		for {
+			m := new(pb.BytesMessage)
+			m.Data = c.readBuf
+			err := c.stream.RecvMsg(m)
+			if err != nil {
+				if !errors.Is(err, io.EOF) {
+					c.readMu.Unlock()
+					return
+				}
+				err = nil
+				break
+			}
+			c.readBuf = m.Data[:cap(m.Data)]
+			c.lastBuf = append(c.lastBuf, c.readBuf...)
+		}
+		c.readMu.Unlock()
+	})
+	return nil
+}
+
+func (c *streamConn) LocalAddr() net.Addr  { return &grpcAddr{} }
+func (c *streamConn) RemoteAddr() net.Addr { return &grpcAddr{} }
+
+func (c *streamConn) SetDeadline(t time.Time) error      { return nil }
+func (c *streamConn) SetReadDeadline(t time.Time) error  { return nil }
+func (c *streamConn) SetWriteDeadline(t time.Time) error { return nil }
+
+type grpcAddr struct{}
+
+func (grpcAddr) Network() string { return "grpc" }
+func (grpcAddr) String() string  { return "grpc-stream" }

pkg/driver/external/client/methods.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/lima-vm/lima/pkg/driver"
 	pb "github.com/lima-vm/lima/pkg/driver/external"
+	"github.com/lima-vm/lima/pkg/driver/external/client/grpchijack"
 	"github.com/lima-vm/lima/pkg/store"
 )
 
@@ -259,7 +260,7 @@ func (d *DriverClient) GuestAgentConn(ctx context.Context) (net.Conn, error) {
 		return nil, err
 	}
 
-	return streamToConn(connStream), nil
+	return grpchijack.StreamToConn(connStream), nil
 }
 
 func (d *DriverClient) GetInfo() driver.Info {

pkg/driver/external/client/pipe.go

@@ -12,9 +12,10 @@ import (
 type PipeConn struct {
 	Reader io.Reader
 	Writer io.Writer
+	Closer io.Closer
 }
 
-func newPipeConn(writer io.WriteCloser, reader io.ReadCloser) *PipeConn {
+func newPipeConn(writer io.WriteCloser, reader io.ReadCloser) net.Conn {
 	return &PipeConn{
 		Reader: reader,
 		Writer: writer,
@@ -30,16 +31,7 @@ func (p *PipeConn) Write(b []byte) (n int, err error) {
 }
 
 func (p *PipeConn) Close() error {
-	var err error
-	if closer, ok := p.Reader.(io.Closer); ok {
-		err = closer.Close()
-	}
-	if closer, ok := p.Writer.(io.Closer); ok {
-		if closeErr := closer.Close(); closeErr != nil && err == nil {
-			err = closeErr
-		}
-	}
-	return err
+	return p.Closer.Close()
 }
 
 func (p *PipeConn) LocalAddr() net.Addr {