guestagent: close the gRPC connection when a write fails

If the gRPC server fails to write, it stops sending but keeps reading.
The connection then stays half-open: requests reach the guest, but no
response ever comes back, and the host agent never reconnects.

Close the connection on the first write error, so that the read side
fails too and the host agent reconnects.

Fixes #5545

Signed-off-by: Sarthak <sarthakshreshtha345@gmail.com>

pkg/guestagent/api/server/server.go

@@ -29,6 +29,7 @@ func StartServer(ctx context.Context, lis net.Listener, guest *GuestServer) erro
 		grpc.KeepaliveParams(keepalive.ServerParameters{Time: 0, Timeout: 0, MaxConnectionIdle: 0}),
 	)
 	api.RegisterGuestServiceServer(server, guest)
+	lis = &closeOnWriteErrorListener{Listener: lis}
 	go func() {
 		<-ctx.Done()
 		logrus.Debug("Stopping the gRPC server")
@@ -44,6 +45,39 @@ func StartServer(ctx context.Context, lis net.Listener, guest *GuestServer) erro
 	return err
 }
 
+// closeOnWriteErrorListener wraps accepted connections with closeOnWriteErrorConn.
+type closeOnWriteErrorListener struct {
+	net.Listener
+}
+
+func (l *closeOnWriteErrorListener) Accept() (net.Conn, error) {
+	conn, err := l.Listener.Accept()
+	if err != nil {
+		return nil, err
+	}
+	return &closeOnWriteErrorConn{Conn: conn}, nil
+}
+
+// closeOnWriteErrorConn closes the connection when a write fails.
+//
+// When the gRPC server fails to write, it stops sending but keeps reading.
+// If the read side still works, the connection stays half-open forever:
+// requests reach the guest, but no response is sent back to the host.
+// Closing the connection makes the reader fail too, so the host agent reconnects.
+// https://github.com/lima-vm/lima/issues/5545
+type closeOnWriteErrorConn struct {
+	net.Conn
+}
+
+func (c *closeOnWriteErrorConn) Write(b []byte) (int, error) {
+	n, err := c.Conn.Write(b)
+	if err != nil {
+		logrus.WithError(err).Warn("Closing the gRPC connection after a write error")
+		_ = c.Conn.Close()
+	}
+	return n, err
+}
+
 type GuestServer struct {
 	api.UnimplementedGuestServiceServer
 	Agent   guestagent.Agent

pkg/guestagent/api/server/server_test.go

@@ -0,0 +1,62 @@
+// SPDX-FileCopyrightText: Copyright The Lima Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package server
+
+import (
+	"errors"
+	"net"
+	"testing"
+
+	"gotest.tools/v3/assert"
+)
+
+type fakeConn struct {
+	net.Conn
+	writeErr error
+	closed   bool
+}
+
+func (c *fakeConn) Write(b []byte) (int, error) {
+	if c.writeErr != nil {
+		return 0, c.writeErr
+	}
+	return len(b), nil
+}
+
+func (c *fakeConn) Close() error {
+	c.closed = true
+	return nil
+}
+
+type fakeListener struct {
+	net.Listener
+	conn net.Conn
+}
+
+func (l *fakeListener) Accept() (net.Conn, error) {
+	return l.conn, nil
+}
+
+func TestCloseOnWriteErrorListener(t *testing.T) {
+	t.Run("write succeeds", func(t *testing.T) {
+		inner := &fakeConn{}
+		lis := &closeOnWriteErrorListener{Listener: &fakeListener{conn: inner}}
+		conn, err := lis.Accept()
+		assert.NilError(t, err)
+		n, err := conn.Write([]byte("hello"))
+		assert.NilError(t, err)
+		assert.Equal(t, n, 5)
+		assert.Assert(t, !inner.closed)
+	})
+	t.Run("write fails", func(t *testing.T) {
+		writeErr := errors.New("write failed")
+		inner := &fakeConn{writeErr: writeErr}
+		lis := &closeOnWriteErrorListener{Listener: &fakeListener{conn: inner}}
+		conn, err := lis.Accept()
+		assert.NilError(t, err)
+		_, err = conn.Write([]byte("hello"))
+		assert.ErrorIs(t, err, writeErr)
+		assert.Assert(t, inner.closed, "connection must be closed after a write error")
+	})
+}