From aa13dc0ef4a919b0951ea54932154983261ed127 Mon Sep 17 00:00:00 2001
From: Balaji Vijayakumar <kuttibalaji.v6@gmail.com>
Date: Mon, 19 Feb 2024 11:55:32 +0530
Subject: [PATCH] Fix for grpc frame too large and context closed warning

Signed-off-by: Balaji Vijayakumar <kuttibalaji.v6@gmail.com>
---
 pkg/guestagent/api/client/client.go      |  4 +-
 pkg/guestagent/api/client/credentials.go | 51 ++++++++++++++++++++++++
 pkg/hostagent/hostagent.go               |  8 +++-
 pkg/vz/network_darwin.go                 |  1 -
 4 files changed, 59 insertions(+), 5 deletions(-)
 create mode 100644 pkg/guestagent/api/client/credentials.go

diff --git a/pkg/guestagent/api/client/client.go b/pkg/guestagent/api/client/client.go
index 2b98907b2ee..2312723dfb5 100644
--- a/pkg/guestagent/api/client/client.go
+++ b/pkg/guestagent/api/client/client.go
@@ -4,8 +4,6 @@ import (
 	"context"
 	"net"
 
-	"google.golang.org/grpc/credentials/insecure"
-
 	"github.com/lima-vm/lima/pkg/guestagent/api"
 	"google.golang.org/grpc"
 	"google.golang.org/protobuf/types/known/emptypb"
@@ -20,7 +18,7 @@ func NewGuestAgentClient(dialFn func(ctx context.Context) (net.Conn, error)) (*G
 		grpc.WithContextDialer(func(ctx context.Context, target string) (net.Conn, error) {
 			return dialFn(ctx)
 		}),
-		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithTransportCredentials(NewCredentials()),
 	}
 
 	clientConn, err := grpc.Dial("", opts...)
diff --git a/pkg/guestagent/api/client/credentials.go b/pkg/guestagent/api/client/credentials.go
new file mode 100644
index 00000000000..76470f14da3
--- /dev/null
+++ b/pkg/guestagent/api/client/credentials.go
@@ -0,0 +1,51 @@
+package client
+
+import (
+	"context"
+	"net"
+
+	"google.golang.org/grpc/credentials"
+)
+
+// NewCredentials returns a lima credential implementing credentials.TransportCredentials.
+func NewCredentials() credentials.TransportCredentials {
+	return &secureTC{
+		info: credentials.ProtocolInfo{
+			SecurityProtocol: "local",
+		},
+	}
+}
+
+// secureTC is the credentials required to establish a lima guest connection.
+type secureTC struct {
+	info credentials.ProtocolInfo
+}
+
+func (c *secureTC) Info() credentials.ProtocolInfo {
+	return c.info
+}
+
+func (*secureTC) ClientHandshake(_ context.Context, _ string, conn net.Conn) (net.Conn, credentials.AuthInfo, error) {
+	return conn, info{credentials.CommonAuthInfo{SecurityLevel: credentials.PrivacyAndIntegrity}}, nil
+}
+
+func (*secureTC) ServerHandshake(conn net.Conn) (net.Conn, credentials.AuthInfo, error) {
+	return conn, info{credentials.CommonAuthInfo{SecurityLevel: credentials.PrivacyAndIntegrity}}, nil
+}
+
+func (c *secureTC) Clone() credentials.TransportCredentials {
+	return &secureTC{info: c.info}
+}
+
+func (c *secureTC) OverrideServerName(serverNameOverride string) error {
+	c.info.ServerName = serverNameOverride
+	return nil
+}
+
+type info struct {
+	credentials.CommonAuthInfo
+}
+
+func (info) AuthType() string {
+	return "local"
+}
diff --git a/pkg/hostagent/hostagent.go b/pkg/hostagent/hostagent.go
index dcc7321fbb2..6247ecbea83 100644
--- a/pkg/hostagent/hostagent.go
+++ b/pkg/hostagent/hostagent.go
@@ -16,6 +16,9 @@ import (
 	"sync"
 	"time"
 
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
 	"github.com/lima-vm/lima/pkg/driver"
 	"github.com/lima-vm/lima/pkg/driverutil"
 	"github.com/lima-vm/lima/pkg/networks"
@@ -595,7 +598,7 @@ func (a *HostAgent) watchGuestAgentEvents(ctx context.Context) {
 		if err == nil {
 			if err := a.processGuestAgentEvents(ctx, client); err != nil {
 				if !errors.Is(err, context.Canceled) {
-					logrus.WithError(err).Warn("connection to the guest agent was closed unexpectedly", err)
+					logrus.WithError(err).Warn("guest agent events closed unexpectedly", err)
 				}
 			}
 		} else {
@@ -658,6 +661,9 @@ func (a *HostAgent) processGuestAgentEvents(ctx context.Context, client *guestag
 	}
 
 	if err := client.Events(ctx, onEvent); err != nil {
+		if status.Code(err) == codes.Canceled {
+			return context.Canceled
+		}
 		return err
 	}
 	return io.EOF
diff --git a/pkg/vz/network_darwin.go b/pkg/vz/network_darwin.go
index fc5611f2d1d..e114335ddc0 100644
--- a/pkg/vz/network_darwin.go
+++ b/pkg/vz/network_darwin.go
@@ -79,7 +79,6 @@ func (v *QEMUPacketConn) Read(b []byte) (n int, err error) {
 	size := binary.BigEndian.Uint32(header)
 	reader := io.LimitReader(v.unixConn, int64(size))
 	_, err = reader.Read(b)
-
 	if err != nil {
 		logrus.Errorln("Failed to read packet", err)
 	}