From baed5f4eeae487a187fa98b0f859a589a5c1a3d7 Mon Sep 17 00:00:00 2001
From: cthuang <chungting@cloudflare.com>
Date: Mon, 30 May 2022 13:38:15 +0100
Subject: [PATCH] TUN-6308: Add debug logs to see if packets are sent/received
 from edge

---
 connection/quic.go          |  4 ++--
 ingress/origin_udp_proxy.go | 11 ++++++++---
 quic/datagram.go            | 19 +++++++++++--------
 quic/datagram_test.go       |  7 +++++--
 4 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/connection/quic.go b/connection/quic.go
index 2c04edf3..4e2f4681 100644
--- a/connection/quic.go
+++ b/connection/quic.go
@@ -58,7 +58,7 @@ func NewQUICConnection(
 		return nil, fmt.Errorf("failed to dial to edge: %w", err)
 	}
 
-	datagramMuxer, err := quicpogs.NewDatagramMuxer(session)
+	datagramMuxer, err := quicpogs.NewDatagramMuxer(session, logger)
 	if err != nil {
 		return nil, err
 	}
@@ -217,7 +217,7 @@ func (q *QUICConnection) RegisterUdpSession(ctx context.Context, sessionID uuid.
 
 	go q.serveUDPSession(session, closeAfterIdleHint)
 
-	q.logger.Debug().Msgf("Registered session %v, %v, %v", sessionID, dstIP, dstPort)
+	q.logger.Debug().Str("sessionID", sessionID.String()).Str("src", originProxy.LocalAddr().String()).Str("dst", fmt.Sprintf("%s:%d", dstIP, dstPort)).Msgf("Registered session")
 	return nil
 }
 
diff --git a/ingress/origin_udp_proxy.go b/ingress/origin_udp_proxy.go
index 3f2b86c7..836489be 100644
--- a/ingress/origin_udp_proxy.go
+++ b/ingress/origin_udp_proxy.go
@@ -6,11 +6,16 @@ import (
 	"net"
 )
 
-type UDPProxy struct {
+type UDPProxy interface {
 	io.ReadWriteCloser
+	LocalAddr() net.Addr
 }
 
-func DialUDP(dstIP net.IP, dstPort uint16) (*UDPProxy, error) {
+type udpProxy struct {
+	*net.UDPConn
+}
+
+func DialUDP(dstIP net.IP, dstPort uint16) (UDPProxy, error) {
 	dstAddr := &net.UDPAddr{
 		IP:   dstIP,
 		Port: int(dstPort),
@@ -23,5 +28,5 @@ func DialUDP(dstIP net.IP, dstPort uint16) (*UDPProxy, error) {
 		return nil, fmt.Errorf("unable to create UDP proxy to origin (%v:%v): %w", dstIP, dstPort, err)
 	}
 
-	return &UDPProxy{udpConn}, nil
+	return &udpProxy{udpConn}, nil
 }
diff --git a/quic/datagram.go b/quic/datagram.go
index 6b3e2bfd..e9b305af 100644
--- a/quic/datagram.go
+++ b/quic/datagram.go
@@ -6,6 +6,7 @@ import (
 	"github.com/google/uuid"
 	"github.com/lucas-clemente/quic-go"
 	"github.com/pkg/errors"
+	"github.com/rs/zerolog"
 )
 
 const (
@@ -13,18 +14,14 @@ const (
 )
 
 type DatagramMuxer struct {
-	ID      uuid.UUID
 	session quic.Session
+	logger  *zerolog.Logger
 }
 
-func NewDatagramMuxer(quicSession quic.Session) (*DatagramMuxer, error) {
-	muxerID, err := uuid.NewRandom()
-	if err != nil {
-		return nil, err
-	}
+func NewDatagramMuxer(quicSession quic.Session, logger *zerolog.Logger) (*DatagramMuxer, error) {
 	return &DatagramMuxer{
-		ID:      muxerID,
 		session: quicSession,
+		logger:  logger,
 	}, nil
 }
 
@@ -42,6 +39,7 @@ func (dm *DatagramMuxer) SendTo(sessionID uuid.UUID, payload []byte) error {
 	if err := dm.session.SendMessage(msgWithID); err != nil {
 		return errors.Wrap(err, "Failed to send datagram back to edge")
 	}
+	dm.logger.Debug().Str("sessionID", sessionID.String()).Int("bytes", len(payload)).Msg("Send datagram back to edge")
 	return nil
 }
 
@@ -53,7 +51,12 @@ func (dm *DatagramMuxer) ReceiveFrom() (uuid.UUID, []byte, error) {
 	if err != nil {
 		return uuid.Nil, nil, err
 	}
-	return extractSessionID(msg)
+	sessionID, payload, err := extractSessionID(msg)
+	if err != nil {
+		return uuid.Nil, nil, err
+	}
+	dm.logger.Debug().Str("sessionID", sessionID.String()).Int("bytes", len(payload)).Msg("Received datagram from edge")
+	return sessionID, payload, nil
 }
 
 // Maximum application payload to send to / receive from QUIC datagram frame
diff --git a/quic/datagram_test.go b/quic/datagram_test.go
index b8745fa5..3f3eef1a 100644
--- a/quic/datagram_test.go
+++ b/quic/datagram_test.go
@@ -14,6 +14,7 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/lucas-clemente/quic-go"
+	"github.com/rs/zerolog"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/sync/errgroup"
 )
@@ -71,7 +72,8 @@ func TestMaxDatagramPayload(t *testing.T) {
 			return err
 		}
 
-		muxer, err := NewDatagramMuxer(quicSession)
+		logger := zerolog.Nop()
+		muxer, err := NewDatagramMuxer(quicSession, &logger)
 		if err != nil {
 			return err
 		}
@@ -96,7 +98,8 @@ func TestMaxDatagramPayload(t *testing.T) {
 		quicSession, err := quic.DialAddrEarly(quicListener.Addr().String(), tlsClientConfig, quicConfig)
 		require.NoError(t, err)
 
-		muxer, err := NewDatagramMuxer(quicSession)
+		logger := zerolog.Nop()
+		muxer, err := NewDatagramMuxer(quicSession, &logger)
 		if err != nil {
 			return err
 		}