TUN-6813: Only proxy ICMP packets when warp-routing is enabled

packet/router.go

@@ -23,10 +23,11 @@ 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
-	globalConfig *GlobalRouterConfig
-	logger       *zerolog.Logger
+	upstream               Upstream
+	returnPipe             FunnelUniPipe
+	globalConfig           *GlobalRouterConfig
+	logger                 *zerolog.Logger
+	checkRouterEnabledFunc func() bool
 }
 
 // GlobalRouterConfig is the configuration shared by all instance of Router.
@@ -37,12 +38,13 @@ type GlobalRouterConfig struct {
 	Zone       string
 }
 
-func NewRouter(globalConfig *GlobalRouterConfig, upstream Upstream, returnPipe FunnelUniPipe, logger *zerolog.Logger) *Router {
+func NewRouter(globalConfig *GlobalRouterConfig, upstream Upstream, returnPipe FunnelUniPipe, logger *zerolog.Logger, checkRouterEnabledFunc func() bool) *Router {
 	return &Router{
-		upstream:     upstream,
-		returnPipe:   returnPipe,
-		globalConfig: globalConfig,
-		logger:       logger,
+		upstream:               upstream,
+		returnPipe:             returnPipe,
+		globalConfig:           globalConfig,
+		logger:                 logger,
+		checkRouterEnabledFunc: checkRouterEnabledFunc,
 	}
 }
 
@@ -54,10 +56,16 @@ 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
 		}
+
+		if enabled := r.checkRouterEnabledFunc(); !enabled {
+			continue
+		}
+
 		icmpPacket, err := icmpDecoder.Decode(rawPacket)
 		if err != nil {
 			r.logger.Err(err).Msg("Failed to decode ICMP packet from quic datagram")

packet/router_test.go

@@ -5,7 +5,9 @@ import (
 	"context"
 	"fmt"
 	"net/netip"
+	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/google/gopacket/layers"
 	"github.com/rs/zerolog"
@@ -16,7 +18,12 @@ import (
 )
 
 var (
-	noopLogger = zerolog.Nop()
+	noopLogger   = zerolog.Nop()
+	packetConfig = &GlobalRouterConfig{
+		ICMPRouter: &mockICMPRouter{},
+		IPv4Src:    netip.MustParseAddr("172.16.0.1"),
+		IPv6Src:    netip.MustParseAddr("fd51:2391:523:f4ee::1"),
+	}
 )
 
 func TestRouterReturnTTLExceed(t *testing.T) {
@@ -26,12 +33,9 @@ func TestRouterReturnTTLExceed(t *testing.T) {
 	returnPipe := &mockFunnelUniPipe{
 		uniPipe: make(chan RawPacket),
 	}
-	packetConfig := &GlobalRouterConfig{
-		ICMPRouter: &mockICMPRouter{},
-		IPv4Src:    netip.MustParseAddr("172.16.0.1"),
-		IPv6Src:    netip.MustParseAddr("fd51:2391:523:f4ee::1"),
-	}
-	router := NewRouter(packetConfig, upstream, returnPipe, &noopLogger)
+	routerEnabled := &routerEnabledChecker{}
+	routerEnabled.set(true)
+	router := NewRouter(packetConfig, upstream, returnPipe, &noopLogger, routerEnabled.isEnabled)
 	ctx, cancel := context.WithCancel(context.Background())
 	routerStopped := make(chan struct{})
 	go func() {
@@ -80,12 +84,71 @@ func TestRouterReturnTTLExceed(t *testing.T) {
 	<-routerStopped
 }
 
+func TestRouterCheckEnabled(t *testing.T) {
+	upstream := &mockUpstream{
+		source: make(chan RawPacket),
+	}
+	returnPipe := &mockFunnelUniPipe{
+		uniPipe: make(chan RawPacket),
+	}
+	routerEnabled := &routerEnabledChecker{}
+	router := NewRouter(packetConfig, upstream, returnPipe, &noopLogger, routerEnabled.isEnabled)
+	ctx, cancel := context.WithCancel(context.Background())
+	routerStopped := make(chan struct{})
+	go func() {
+		router.Serve(ctx)
+		close(routerStopped)
+	}()
+
+	pk := ICMP{
+		IP: &IP{
+			Src:      netip.MustParseAddr("192.168.1.1"),
+			Dst:      netip.MustParseAddr("10.0.0.1"),
+			Protocol: layers.IPProtocolICMPv4,
+			TTL:      1,
+		},
+		Message: &icmp.Message{
+			Type: ipv4.ICMPTypeEcho,
+			Code: 0,
+			Body: &icmp.Echo{
+				ID:   12481,
+				Seq:  8036,
+				Data: []byte(t.Name()),
+			},
+		},
+	}
+
+	// router is disabled
+	require.NoError(t, upstream.send(&pk))
+	select {
+	case <-time.After(time.Millisecond * 10):
+	case <-returnPipe.uniPipe:
+		t.Error("Unexpected reply when router is disabled")
+	}
+	routerEnabled.set(true)
+	// router is enabled, expects reply
+	require.NoError(t, upstream.send(&pk))
+	<-returnPipe.uniPipe
+
+	routerEnabled.set(false)
+	// router is disabled
+	require.NoError(t, upstream.send(&pk))
+	select {
+	case <-time.After(time.Millisecond * 10):
+	case <-returnPipe.uniPipe:
+		t.Error("Unexpected reply when router is disabled")
+	}
+
+	cancel()
+	<-routerStopped
+}
+
 func assertTTLExceed(t *testing.T, originalPacket *ICMP, expectedSrc netip.Addr, upstream *mockUpstream, returnPipe *mockFunnelUniPipe) {
 	encoder := NewEncoder()
 	rawPacket, err := encoder.Encode(originalPacket)
 	require.NoError(t, err)
-
 	upstream.source <- rawPacket
+
 	resp := <-returnPipe.uniPipe
 	decoder := NewICMPDecoder()
 	decoded, err := decoder.Decode(resp)
@@ -111,6 +174,16 @@ type mockUpstream struct {
 	source chan RawPacket
 }
 
+func (ms *mockUpstream) send(pk Packet) error {
+	encoder := NewEncoder()
+	rawPacket, err := encoder.Encode(pk)
+	if err != nil {
+		return err
+	}
+	ms.source <- rawPacket
+	return nil
+}
+
 func (ms *mockUpstream) ReceivePacket(ctx context.Context) (RawPacket, error) {
 	select {
 	case <-ctx.Done():
@@ -129,3 +202,22 @@ func (mir mockICMPRouter) Serve(ctx context.Context) error {
 func (mir mockICMPRouter) Request(pk *ICMP, responder FunnelUniPipe) error {
 	return fmt.Errorf("Request not implemented by mockICMPRouter")
 }
+
+type routerEnabledChecker struct {
+	enabled uint32
+}
+
+func (rec *routerEnabledChecker) isEnabled() bool {
+	if atomic.LoadUint32(&rec.enabled) == 0 {
+		return false
+	}
+	return true
+}
+
+func (rec *routerEnabledChecker) set(enabled bool) {
+	if enabled {
+		atomic.StoreUint32(&rec.enabled, 1)
+	} else {
+		atomic.StoreUint32(&rec.enabled, 0)
+	}
+}