From 44a8db5d63d8876d6885088b4665ec9666232579 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Fri, 16 Aug 2024 17:33:08 +0800
Subject: [PATCH] fix missing context ID in datagrams

For UDP payloads, the context ID is 0, see section 5 of RFC 9298.
---
 conn.go       | 17 ++++++++++++++---
 proxy.go      | 17 +++++++++++++++--
 proxy_test.go |  2 +-
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/conn.go b/conn.go
index 30d60f1..539c304 100644
--- a/conn.go
+++ b/conn.go
@@ -3,6 +3,7 @@ package masque
 import (
 	"context"
 	"errors"
+	"fmt"
 	"io"
 	"log"
 	"net"
@@ -77,16 +78,26 @@ start:
 		}
 		return 0, nil, os.ErrDeadlineExceeded
 	}
+	contextID, n, err := quicvarint.Parse(data)
+	if err != nil {
+		return 0, nil, fmt.Errorf("masque: malformed datagram: %w", err)
+	}
+	if contextID != 0 {
+		// Drop this datagram. We currently only support proxying of UDP payloads.
+		goto start
+	}
 	// If b is too small, additional bytes are discarded.
 	// This mirrors the behavior of large UDP datagrams received on a UDP socket (on Linux).
-	n = copy(b, data)
-	return n, c.remoteAddr, nil
+	return copy(b, data[n:]), c.remoteAddr, nil
 }
 
 // WriteTo sends a UDP datagram to the target.
 // The net.Addr parameter is ignored.
 func (c *proxiedConn) WriteTo(p []byte, _ net.Addr) (n int, err error) {
-	return len(p), c.str.SendDatagram(p)
+	data := make([]byte, 0, len(contextIDZero)+len(p))
+	data = append(data, contextIDZero...)
+	data = append(data, p...)
+	return len(p), c.str.SendDatagram(data)
 }
 
 func (c *proxiedConn) Close() error {
diff --git a/proxy.go b/proxy.go
index 482cd1d..94c0fda 100644
--- a/proxy.go
+++ b/proxy.go
@@ -19,6 +19,8 @@ const (
 	uriTemplateTargetPort = "target_port"
 )
 
+var contextIDZero = quicvarint.Append([]byte{}, 0)
+
 type proxyEntry struct {
 	str  http3.Stream
 	conn *net.UDPConn
@@ -126,7 +128,15 @@ func (s *Proxy) proxyConnSend(conn *net.UDPConn, str http3.Stream) error {
 		if err != nil {
 			return err
 		}
-		if _, err := conn.Write(data); err != nil {
+		contextID, n, err := quicvarint.Parse(data)
+		if err != nil {
+			return err
+		}
+		if contextID != 0 {
+			// Drop this datagram. We currently only support proxying of UDP payloads.
+			continue
+		}
+		if _, err := conn.Write(data[n:]); err != nil {
 			return err
 		}
 	}
@@ -139,7 +149,10 @@ func (s *Proxy) proxyConnReceive(conn *net.UDPConn, str http3.Stream) error {
 		if err != nil {
 			return err
 		}
-		if err := str.SendDatagram(b[:n]); err != nil {
+		data := make([]byte, 0, len(contextIDZero)+n)
+		data = append(data, contextIDZero...)
+		data = append(data, b[:n]...)
+		if err := str.SendDatagram(data); err != nil {
 			return err
 		}
 	}
diff --git a/proxy_test.go b/proxy_test.go
index 6dc75c0..547632d 100644
--- a/proxy_test.go
+++ b/proxy_test.go
@@ -56,7 +56,7 @@ func TestProxyCloseProxiedConn(t *testing.T) {
 	done := make(chan struct{})
 	str := NewMockStream(gomock.NewController(t))
 	str.EXPECT().ReceiveDatagram(gomock.Any()).DoAndReturn(func(context.Context) ([]byte, error) {
-		return []byte("foo"), nil
+		return append(contextIDZero, []byte("foo")...), nil
 	})
 	// This datagram is received after the connection is closed.
 	// We expect that it won't get sent on.