From eacc8c648d39c0b0422c02f66d1eb355a4c6c738 Mon Sep 17 00:00:00 2001
From: cthuang <chungting@cloudflare.com>
Date: Thu, 29 Sep 2022 15:42:30 +0100
Subject: [PATCH] TUN-6812: Drop IP packets if ICMP proxy is not initialized

---
 connection/quic.go    | 21 +++++++--------------
 packet/router.go      | 30 +++++++++++++++---------------
 packet/router_test.go |  4 ++--
 3 files changed, 24 insertions(+), 31 deletions(-)

diff --git a/connection/quic.go b/connection/quic.go
index 95829a6e..3004b929 100644
--- a/connection/quic.go
+++ b/connection/quic.go
@@ -50,7 +50,7 @@ type QUICConnection struct {
 	// sessionManager tracks active sessions. It receives datagrams from quic connection via datagramMuxer
 	sessionManager datagramsession.Manager
 	// datagramMuxer mux/demux datagrams from quic connection
-	datagramMuxer        quicpogs.BaseDatagramMuxer
+	datagramMuxer        *quicpogs.DatagramMuxerV2
 	packetRouter         *packet.Router
 	controlStreamHandler ControlStreamHandler
 	connOptions          *tunnelpogs.ConnectionOptions
@@ -75,11 +75,7 @@ func NewQUICConnection(
 	sessionDemuxChan := make(chan *packet.Session, demuxChanCapacity)
 	datagramMuxer := quicpogs.NewDatagramMuxerV2(session, logger, sessionDemuxChan)
 	sessionManager := datagramsession.NewManager(logger, datagramMuxer.SendToSession, sessionDemuxChan)
-
-	var pr *packet.Router
-	if packetRouterConfig != nil {
-		pr = packet.NewRouter(packetRouterConfig, datagramMuxer, &returnPipe{muxer: datagramMuxer}, logger)
-	}
+	packetRouter := packet.NewRouter(packetRouterConfig, datagramMuxer, &returnPipe{muxer: datagramMuxer}, logger)
 
 	return &QUICConnection{
 		session:              session,
@@ -87,7 +83,7 @@ func NewQUICConnection(
 		logger:               logger,
 		sessionManager:       sessionManager,
 		datagramMuxer:        datagramMuxer,
-		packetRouter:         pr,
+		packetRouter:         packetRouter,
 		controlStreamHandler: controlStreamHandler,
 		connOptions:          connOptions,
 	}, nil
@@ -123,17 +119,14 @@ func (q *QUICConnection) Serve(ctx context.Context) error {
 		defer cancel()
 		return q.sessionManager.Serve(ctx)
 	})
-
 	errGroup.Go(func() error {
 		defer cancel()
 		return q.datagramMuxer.ServeReceive(ctx)
 	})
-	if q.packetRouter != nil {
-		errGroup.Go(func() error {
-			defer cancel()
-			return q.packetRouter.Serve(ctx)
-		})
-	}
+	errGroup.Go(func() error {
+		defer cancel()
+		return q.packetRouter.Serve(ctx)
+	})
 
 	return errGroup.Wait()
 }
diff --git a/packet/router.go b/packet/router.go
index 2f0fa37b..3fc451fc 100644
--- a/packet/router.go
+++ b/packet/router.go
@@ -23,12 +23,10 @@ type Upstream interface {
 
 // Router routes packets between Upstream and ICMPRouter. Currently it rejects all other type of ICMP packets
 type Router struct {
-	upstream   Upstream
-	returnPipe FunnelUniPipe
-	icmpRouter ICMPRouter
-	ipv4Src    netip.Addr
-	ipv6Src    netip.Addr
-	logger     *zerolog.Logger
+	upstream     Upstream
+	returnPipe   FunnelUniPipe
+	globalConfig *GlobalRouterConfig
+	logger       *zerolog.Logger
 }
 
 // GlobalRouterConfig is the configuration shared by all instance of Router.
@@ -41,12 +39,10 @@ type GlobalRouterConfig struct {
 
 func NewRouter(globalConfig *GlobalRouterConfig, upstream Upstream, returnPipe FunnelUniPipe, logger *zerolog.Logger) *Router {
 	return &Router{
-		upstream:   upstream,
-		returnPipe: returnPipe,
-		icmpRouter: globalConfig.ICMPRouter,
-		ipv4Src:    globalConfig.IPv4Src,
-		ipv6Src:    globalConfig.IPv6Src,
-		logger:     logger,
+		upstream:     upstream,
+		returnPipe:   returnPipe,
+		globalConfig: globalConfig,
+		logger:       logger,
 	}
 }
 
@@ -58,6 +54,10 @@ func (r *Router) Serve(ctx context.Context) error {
 		if err != nil {
 			return err
 		}
+		// Drop packets if ICMPRouter wasn't created
+		if r.globalConfig == nil {
+			continue
+		}
 		icmpPacket, err := icmpDecoder.Decode(rawPacket)
 		if err != nil {
 			r.logger.Err(err).Msg("Failed to decode ICMP packet from quic datagram")
@@ -72,7 +72,7 @@ func (r *Router) Serve(ctx context.Context) error {
 		}
 		icmpPacket.TTL--
 
-		if err := r.icmpRouter.Request(icmpPacket, r.returnPipe); err != nil {
+		if err := r.globalConfig.ICMPRouter.Request(icmpPacket, r.returnPipe); err != nil {
 			r.logger.Err(err).
 				Str("src", icmpPacket.Src.String()).
 				Str("dst", icmpPacket.Dst.String()).
@@ -86,9 +86,9 @@ func (r *Router) Serve(ctx context.Context) error {
 func (r *Router) sendTTLExceedMsg(pk *ICMP, rawPacket RawPacket, encoder *Encoder) error {
 	var srcIP netip.Addr
 	if pk.Dst.Is4() {
-		srcIP = r.ipv4Src
+		srcIP = r.globalConfig.IPv4Src
 	} else {
-		srcIP = r.ipv6Src
+		srcIP = r.globalConfig.IPv6Src
 	}
 	ttlExceedPacket := NewICMPTTLExceedPacket(pk.IP, rawPacket, srcIP)
 
diff --git a/packet/router_test.go b/packet/router_test.go
index 8b009c1c..c1056450 100644
--- a/packet/router_test.go
+++ b/packet/router_test.go
@@ -56,7 +56,7 @@ func TestRouterReturnTTLExceed(t *testing.T) {
 			},
 		},
 	}
-	assertTTLExceed(t, &pk, router.ipv4Src, upstream, returnPipe)
+	assertTTLExceed(t, &pk, router.globalConfig.IPv4Src, upstream, returnPipe)
 	pk = ICMP{
 		IP: &IP{
 			Src:      netip.MustParseAddr("fd51:2391:523:f4ee::1"),
@@ -74,7 +74,7 @@ func TestRouterReturnTTLExceed(t *testing.T) {
 			},
 		},
 	}
-	assertTTLExceed(t, &pk, router.ipv6Src, upstream, returnPipe)
+	assertTTLExceed(t, &pk, router.globalConfig.IPv6Src, upstream, returnPipe)
 
 	cancel()
 	<-routerStopped