From 150b1162316cb63bbbf5c3c01f3a2b494a56e72d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Wed, 1 Nov 2023 20:57:42 +0800
Subject: [PATCH] Add multicast filter

---
 stack_gvisor.go        |  1 +
 stack_gvisor_filter.go | 52 ++++++++++++++++++++++++++++++++++++++++++
 stack_system.go        |  7 ++++++
 3 files changed, 60 insertions(+)
 create mode 100644 stack_gvisor_filter.go

diff --git a/stack_gvisor.go b/stack_gvisor.go
index 51c9179..616598c 100644
--- a/stack_gvisor.go
+++ b/stack_gvisor.go
@@ -70,6 +70,7 @@ func (t *GVisor) Start() error {
 	if err != nil {
 		return err
 	}
+	linkEndpoint = &LinkEndpointFilter{linkEndpoint, t.tun.CreateVectorisedWriter()}
 	ipStack, err := newGVisorStack(linkEndpoint)
 	if err != nil {
 		return err
diff --git a/stack_gvisor_filter.go b/stack_gvisor_filter.go
new file mode 100644
index 0000000..7f943ae
--- /dev/null
+++ b/stack_gvisor_filter.go
@@ -0,0 +1,52 @@
+//go:build with_gvisor
+
+package tun
+
+import (
+	"github.com/sagernet/gvisor/pkg/tcpip"
+	"github.com/sagernet/gvisor/pkg/tcpip/header"
+	"github.com/sagernet/gvisor/pkg/tcpip/stack"
+	"github.com/sagernet/sing/common/bufio"
+	N "github.com/sagernet/sing/common/network"
+)
+
+var _ stack.LinkEndpoint = (*LinkEndpointFilter)(nil)
+
+type LinkEndpointFilter struct {
+	stack.LinkEndpoint
+	Writer N.VectorisedWriter
+}
+
+func (w *LinkEndpointFilter) Attach(dispatcher stack.NetworkDispatcher) {
+	w.LinkEndpoint.Attach(&networkDispatcherFilter{dispatcher, w.Writer})
+}
+
+var _ stack.NetworkDispatcher = (*networkDispatcherFilter)(nil)
+
+type networkDispatcherFilter struct {
+	stack.NetworkDispatcher
+	writer N.VectorisedWriter
+}
+
+func (w *networkDispatcherFilter) DeliverNetworkPacket(protocol tcpip.NetworkProtocolNumber, pkt stack.PacketBufferPtr) {
+	var network header.Network
+	if protocol == header.IPv4ProtocolNumber {
+		if headerPackets, loaded := pkt.Data().PullUp(header.IPv4MinimumSize); loaded {
+			network = header.IPv4(headerPackets)
+		}
+	} else {
+		if headerPackets, loaded := pkt.Data().PullUp(header.IPv6MinimumSize); loaded {
+			network = header.IPv6(headerPackets)
+		}
+	}
+	if network == nil {
+		w.NetworkDispatcher.DeliverNetworkPacket(protocol, pkt)
+		return
+	}
+	destination := AddrFromAddress(network.DestinationAddress())
+	if destination.IsMulticast() || !destination.IsGlobalUnicast() {
+		_, _ = bufio.WriteVectorised(w.writer, pkt.AsSlices())
+		return
+	}
+	w.NetworkDispatcher.DeliverNetworkPacket(protocol, pkt)
+}
diff --git a/stack_system.go b/stack_system.go
index e8a64ca..e2b87cc 100644
--- a/stack_system.go
+++ b/stack_system.go
@@ -233,6 +233,10 @@ func (s *System) acceptLoop(listener net.Listener) {
 }
 
 func (s *System) processIPv4(packet clashtcpip.IPv4Packet) error {
+	destination := packet.DestinationIP()
+	if destination.IsMulticast() || !destination.IsGlobalUnicast() {
+		return common.Error(s.tun.Write(packet))
+	}
 	switch packet.Protocol() {
 	case clashtcpip.TCP:
 		return s.processIPv4TCP(packet, packet.Payload())
@@ -246,6 +250,9 @@ func (s *System) processIPv4(packet clashtcpip.IPv4Packet) error {
 }
 
 func (s *System) processIPv6(packet clashtcpip.IPv6Packet) error {
+	if !packet.DestinationIP().IsGlobalUnicast() {
+		return common.Error(s.tun.Write(packet))
+	}
 	switch packet.Protocol() {
 	case clashtcpip.TCP:
 		return s.processIPv6TCP(packet, packet.Payload())