From 7a5ddd7ceca5256eed09684b1d6f81330820ce8f Mon Sep 17 00:00:00 2001
From: Nir Soffer <nirsof@gmail.com>
Date: Wed, 4 Dec 2024 01:02:59 +0200
Subject: [PATCH] portfwd: Fix handling the last UDP  message

PacketCon.ReadFrom documentation says[1]:

> Callers should always process the n > 0 bytes returned before
> considering the error err.

But we handled err first, and dropped the last read bytes. Fixed by
sending a message if n > 0, and handling the error after the send.

I'm not sure what how this bug affects the system. It was found while
reviewing #2969.

stream.Recv() does is not documented, but the trivial implementation
ensure that we get nil message on any error. Add comment to make it more
clear.

[1] https://pkg.go.dev/net#PacketConn

Fixes #2970

Thanks: Tamir Duberstein <tamird@gmail.com>
Signed-off-by: Nir Soffer <nirsof@gmail.com>
---
 pkg/portfwd/client.go | 26 ++++++++++++++++----------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/pkg/portfwd/client.go b/pkg/portfwd/client.go
index 1c3d496882b..3285fe83447 100644
--- a/pkg/portfwd/client.go
+++ b/pkg/portfwd/client.go
@@ -58,6 +58,20 @@ func HandleUDPConnection(ctx context.Context, client *guestagentclient.GuestAgen
 		buf := make([]byte, 65507)
 		for {
 			n, addr, err := conn.ReadFrom(buf)
+			// We must handle n > 0 bytes before considering the error.
+			// https://pkg.go.dev/net#PacketConn
+			if n > 0 {
+				msg := &api.TunnelMessage{
+					Id:            id + "-" + addr.String(),
+					Protocol:      "udp",
+					GuestAddr:     guestAddr,
+					Data:          buf[:n],
+					UdpTargetAddr: addr.String(),
+				}
+				if err := stream.Send(msg); err != nil {
+					return err
+				}
+			}
 			if err != nil {
 				// https://pkg.go.dev/net#PacketConn does not mention io.EOF semantics.
 				if errors.Is(err, io.EOF) {
@@ -65,21 +79,12 @@ func HandleUDPConnection(ctx context.Context, client *guestagentclient.GuestAgen
 				}
 				return err
 			}
-			msg := &api.TunnelMessage{
-				Id:            id + "-" + addr.String(),
-				Protocol:      "udp",
-				GuestAddr:     guestAddr,
-				Data:          buf[:n],
-				UdpTargetAddr: addr.String(),
-			}
-			if err := stream.Send(msg); err != nil {
-				return err
-			}
 		}
 	})
 
 	g.Go(func() error {
 		for {
+			// Not documented: when err != nil, in is always nil.
 			in, err := stream.Recv()
 			if err != nil {
 				if errors.Is(err, io.EOF) {
@@ -128,6 +133,7 @@ func (g GrpcClientRW) Write(p []byte) (n int, err error) {
 }
 
 func (g GrpcClientRW) Read(p []byte) (n int, err error) {
+	// Not documented: when err != nil, in is always nil.
 	in, err := g.stream.Recv()
 	if err != nil {
 		if errors.Is(err, io.EOF) {