Finalmask: Add filter for inbound (TCP & UDP, source IP filtering) (#7206)

https://github.com/XTLS/Xray-core/pull/7206#issuecomment-6312634417

infra/conf/transport_finalmask.go

@@ -15,9 +15,11 @@ import (
 
 	googleuuid "github.com/google/uuid"
 	"github.com/xtls/xray-core/common/errors"
+	"github.com/xtls/xray-core/common/geodata"
 	xlua "github.com/xtls/xray-core/common/lua"
 	"github.com/xtls/xray-core/common/net"
 	"github.com/xtls/xray-core/transport/internet"
+	"github.com/xtls/xray-core/transport/internet/finalmask/filter"
 	"github.com/xtls/xray-core/transport/internet/finalmask/fragment"
 	"github.com/xtls/xray-core/transport/internet/finalmask/header/custom"
 	"github.com/xtls/xray-core/transport/internet/finalmask/mkcp/aes128gcm"
@@ -74,6 +76,7 @@ var (
 	customVarNamePattern = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_]*$`)
 
 	tcpmaskLoader = NewJSONConfigLoader(ConfigCreatorCache{
+		"filter":        func() interface{} { return new(Filter) },
 		"header-custom": func() interface{} { return new(HeaderCustomTCP) },
 		"fragment":      func() interface{} { return new(FragmentMask) },
 		"sudoku":        func() interface{} { return new(Sudoku) },
@@ -81,6 +84,7 @@ var (
 	}, "type", "settings")
 
 	udpmaskLoader = NewJSONConfigLoader(ConfigCreatorCache{
+		"filter":        func() interface{} { return new(Filter) },
 		"header-custom": func() interface{} { return new(HeaderCustomUDP) },
 		"mkcp-legacy":   func() interface{} { return new(MkcpLegacy) },
 		"noise":         func() interface{} { return new(NoiseMask) },
@@ -93,6 +97,23 @@ var (
 	}, "type", "settings")
 )
 
+type Filter struct {
+	BlockedSourceIPs []string `json:"blockedSourceIPs"`
+	AllowedSourceIPs []string `json:"allowedSourceIPs"`
+}
+
+func (c *Filter) Build() (proto.Message, error) {
+	rules, err := geodata.ParseIPRules(c.BlockedSourceIPs)
+	if err != nil {
+		return nil, errors.New(`invalid "blockedSourceIPs"`).Base(err)
+	}
+	allowed, err := geodata.ParseIPRules(c.AllowedSourceIPs)
+	if err != nil {
+		return nil, errors.New(`invalid "allowedSourceIPs"`).Base(err)
+	}
+	return &filter.Config{BlockedSourceIps: rules, AllowedSourceIps: allowed}, nil
+}
+
 type TCPItem struct {
 	Delay     Int32Range       `json:"delay"`
 	Rand      int32            `json:"rand"`

infra/conf/transport_finalmask_filter_test.go

@@ -0,0 +1,79 @@
+package conf_test
+
+import (
+	"encoding/json"
+	"fmt"
+	"net"
+	"strings"
+	"testing"
+
+	"github.com/xtls/xray-core/common/geodata"
+	. "github.com/xtls/xray-core/infra/conf"
+	"github.com/xtls/xray-core/transport/internet"
+	"github.com/xtls/xray-core/transport/internet/finalmask/filter"
+)
+
+func TestFinalMaskFilter(t *testing.T) {
+	for _, network := range []string{"tcp", "udp"} {
+		for _, test := range []struct {
+			name, settings   string
+			blocked, allowed int
+		}{
+			{"blacklist", `{"blockedSourceIPs":["192.0.2.0/24"]}`, 1, 0},
+			{"whitelist", `{"allowedSourceIPs":["192.0.2.0/24"]}`, 0, 1},
+			{"both", `{"blockedSourceIPs":["192.0.2.0/24"],"allowedSourceIPs":["192.0.2.0/24"]}`, 1, 1},
+			{"empty_whitelist", `{"allowedSourceIPs":[]}`, 0, 0},
+			{"empty_blacklist", `{"blockedSourceIPs":[]}`, 0, 0},
+		} {
+			t.Run(network+"/"+test.name, func(t *testing.T) {
+				input := fmt.Sprintf(`{"finalmask":{"%s":[{"type":"filter","settings":%s}]}}`, network, test.settings)
+				var config StreamConfig
+				if err := json.Unmarshal([]byte(input), &config); err != nil {
+					t.Fatal(err)
+				}
+				built, err := config.Build()
+				if err != nil {
+					t.Fatal(err)
+				}
+				masks := built.Tcpmasks
+				if network == "udp" {
+					masks = built.Udpmasks
+				}
+				if len(masks) != 1 {
+					t.Fatalf("expected one mask, got %d", len(masks))
+				}
+				instance, err := masks[0].GetInstance()
+				if err != nil {
+					t.Fatal(err)
+				}
+				filter, ok := instance.(*filter.Config)
+				if !ok {
+					t.Fatalf("unexpected mask type: %T", instance)
+				}
+				if len(filter.BlockedSourceIps) != test.blocked || len(filter.AllowedSourceIps) != test.allowed {
+					t.Fatalf("unexpected filter policy: %v", filter)
+				}
+				for _, rules := range [][]*geodata.IPRule{filter.BlockedSourceIps, filter.AllowedSourceIps} {
+					for _, rule := range rules {
+						cidr := rule.GetCustom().GetCidr()
+						if cidr.GetPrefix() != 24 || !net.IP(cidr.GetIp()).Equal(net.IPv4(192, 0, 2, 0)) {
+							t.Fatalf("unexpected IP rule: %v", cidr)
+						}
+					}
+				}
+				if _, err := internet.ToMemoryStreamConfig(built); err != nil {
+					t.Fatalf("failed to initialize runtime finalmask: %v", err)
+				}
+			})
+		}
+	}
+	for _, field := range []string{"blockedSourceIPs", "allowedSourceIPs"} {
+		var config Filter
+		if err := json.Unmarshal([]byte(fmt.Sprintf(`{"%s":["192.0.2.0/33"]}`, field)), &config); err != nil {
+			t.Fatal(err)
+		}
+		if _, err := config.Build(); err == nil || !strings.Contains(err.Error(), field) {
+			t.Errorf("expected invalid %s, got %v", field, err)
+		}
+	}
+}

transport/internet/finalmask/filter/config.go

@@ -0,0 +1,63 @@
+package filter
+
+import (
+	"github.com/xtls/xray-core/common/geodata"
+	"github.com/xtls/xray-core/common/net"
+	"github.com/xtls/xray-core/transport/internet/finalmask"
+)
+
+func (c *Config) buildFilter() (*filter, error) {
+	if len(c.BlockedSourceIps) == 0 && len(c.AllowedSourceIps) == 0 {
+		return nil, nil
+	}
+	f := new(filter)
+	var err error
+	if len(c.BlockedSourceIps) > 0 {
+		f.blocked, err = geodata.IPReg.BuildIPMatcher(c.BlockedSourceIps)
+		if err != nil {
+			return nil, err
+		}
+	}
+	if len(c.AllowedSourceIps) > 0 {
+		f.allowed, err = geodata.IPReg.BuildIPMatcher(c.AllowedSourceIps)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return f, nil
+}
+
+func (c *Config) Listen(inner net.Listener) (net.Listener, error) {
+	f, err := c.buildFilter()
+	if err != nil {
+		return nil, err
+	}
+	if f == nil {
+		return inner, nil
+	}
+	return &listener{Listener: inner, filter: f}, nil
+}
+
+func (c *Config) WrapConnClient(conn net.Conn, dest *net.Destination, dialer *finalmask.Dialer) (net.Conn, error) {
+	return conn, nil
+}
+
+func (c *Config) WrapConnServer(conn net.Conn) (net.Conn, error) {
+	// TCP filtering is installed by Listen, before per-connection masks.
+	return conn, nil
+}
+
+func (c *Config) WrapPacketConnClient(conn net.PacketConn, dest *net.Destination, dialer *finalmask.Dialer) (net.PacketConn, error) {
+	return conn, nil
+}
+
+func (c *Config) WrapPacketConnServer(conn net.PacketConn, addr net.Addr, lc *finalmask.ListenConfig) (net.PacketConn, error) {
+	f, err := c.buildFilter()
+	if err != nil {
+		return nil, err
+	}
+	if f == nil {
+		return conn, nil
+	}
+	return &packetConn{PacketConn: conn, filter: f}, nil
+}

transport/internet/finalmask/filter/config.pb.go

@@ -0,0 +1,136 @@
+// Code generated by protoc-gen-go. DO NOT EDIT.
+// versions:
+// 	protoc-gen-go v1.36.11
+// 	protoc        v6.33.5
+// source: transport/internet/finalmask/filter/config.proto
+
+package filter
+
+import (
+	geodata "github.com/xtls/xray-core/common/geodata"
+	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
+	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
+	reflect "reflect"
+	sync "sync"
+	unsafe "unsafe"
+)
+
+const (
+	// Verify that this generated code is sufficiently up-to-date.
+	_ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion)
+	// Verify that runtime/protoimpl is sufficiently up-to-date.
+	_ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20)
+)
+
+type Config struct {
+	state            protoimpl.MessageState `protogen:"open.v1"`
+	BlockedSourceIps []*geodata.IPRule      `protobuf:"bytes,1,rep,name=blocked_source_ips,json=blockedSourceIps,proto3" json:"blocked_source_ips,omitempty"`
+	AllowedSourceIps []*geodata.IPRule      `protobuf:"bytes,2,rep,name=allowed_source_ips,json=allowedSourceIps,proto3" json:"allowed_source_ips,omitempty"`
+	unknownFields    protoimpl.UnknownFields
+	sizeCache        protoimpl.SizeCache
+}
+
+func (x *Config) Reset() {
+	*x = Config{}
+	mi := &file_transport_internet_finalmask_filter_config_proto_msgTypes[0]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *Config) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*Config) ProtoMessage() {}
+
+func (x *Config) ProtoReflect() protoreflect.Message {
+	mi := &file_transport_internet_finalmask_filter_config_proto_msgTypes[0]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use Config.ProtoReflect.Descriptor instead.
+func (*Config) Descriptor() ([]byte, []int) {
+	return file_transport_internet_finalmask_filter_config_proto_rawDescGZIP(), []int{0}
+}
+
+func (x *Config) GetBlockedSourceIps() []*geodata.IPRule {
+	if x != nil {
+		return x.BlockedSourceIps
+	}
+	return nil
+}
+
+func (x *Config) GetAllowedSourceIps() []*geodata.IPRule {
+	if x != nil {
+		return x.AllowedSourceIps
+	}
+	return nil
+}
+
+var File_transport_internet_finalmask_filter_config_proto protoreflect.FileDescriptor
+
+const file_transport_internet_finalmask_filter_config_proto_rawDesc = "" +
+	"\n" +
+	"0transport/internet/finalmask/filter/config.proto\x12(xray.transport.internet.finalmask.filter\x1a\x1bcommon/geodata/geodat.proto\"\x9e\x01\n" +
+	"\x06Config\x12I\n" +
+	"\x12blocked_source_ips\x18\x01 \x03(\v2\x1b.xray.common.geodata.IPRuleR\x10blockedSourceIps\x12I\n" +
+	"\x12allowed_source_ips\x18\x02 \x03(\v2\x1b.xray.common.geodata.IPRuleR\x10allowedSourceIpsB\x9a\x01\n" +
+	",com.xray.transport.internet.finalmask.filterP\x01Z=github.com/xtls/xray-core/transport/internet/finalmask/filter\xaa\x02(Xray.Transport.Internet.Finalmask.Filterb\x06proto3"
+
+var (
+	file_transport_internet_finalmask_filter_config_proto_rawDescOnce sync.Once
+	file_transport_internet_finalmask_filter_config_proto_rawDescData []byte
+)
+
+func file_transport_internet_finalmask_filter_config_proto_rawDescGZIP() []byte {
+	file_transport_internet_finalmask_filter_config_proto_rawDescOnce.Do(func() {
+		file_transport_internet_finalmask_filter_config_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_transport_internet_finalmask_filter_config_proto_rawDesc), len(file_transport_internet_finalmask_filter_config_proto_rawDesc)))
+	})
+	return file_transport_internet_finalmask_filter_config_proto_rawDescData
+}
+
+var file_transport_internet_finalmask_filter_config_proto_msgTypes = make([]protoimpl.MessageInfo, 1)
+var file_transport_internet_finalmask_filter_config_proto_goTypes = []any{
+	(*Config)(nil),         // 0: xray.transport.internet.finalmask.filter.Config
+	(*geodata.IPRule)(nil), // 1: xray.common.geodata.IPRule
+}
+var file_transport_internet_finalmask_filter_config_proto_depIdxs = []int32{
+	1, // 0: xray.transport.internet.finalmask.filter.Config.blocked_source_ips:type_name -> xray.common.geodata.IPRule
+	1, // 1: xray.transport.internet.finalmask.filter.Config.allowed_source_ips:type_name -> xray.common.geodata.IPRule
+	2, // [2:2] is the sub-list for method output_type
+	2, // [2:2] is the sub-list for method input_type
+	2, // [2:2] is the sub-list for extension type_name
+	2, // [2:2] is the sub-list for extension extendee
+	0, // [0:2] is the sub-list for field type_name
+}
+
+func init() { file_transport_internet_finalmask_filter_config_proto_init() }
+func file_transport_internet_finalmask_filter_config_proto_init() {
+	if File_transport_internet_finalmask_filter_config_proto != nil {
+		return
+	}
+	type x struct{}
+	out := protoimpl.TypeBuilder{
+		File: protoimpl.DescBuilder{
+			GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
+			RawDescriptor: unsafe.Slice(unsafe.StringData(file_transport_internet_finalmask_filter_config_proto_rawDesc), len(file_transport_internet_finalmask_filter_config_proto_rawDesc)),
+			NumEnums:      0,
+			NumMessages:   1,
+			NumExtensions: 0,
+			NumServices:   0,
+		},
+		GoTypes:           file_transport_internet_finalmask_filter_config_proto_goTypes,
+		DependencyIndexes: file_transport_internet_finalmask_filter_config_proto_depIdxs,
+		MessageInfos:      file_transport_internet_finalmask_filter_config_proto_msgTypes,
+	}.Build()
+	File_transport_internet_finalmask_filter_config_proto = out.File
+	file_transport_internet_finalmask_filter_config_proto_goTypes = nil
+	file_transport_internet_finalmask_filter_config_proto_depIdxs = nil
+}

transport/internet/finalmask/filter/config.proto

@@ -0,0 +1,14 @@
+syntax = "proto3";
+
+package xray.transport.internet.finalmask.filter;
+option csharp_namespace = "Xray.Transport.Internet.Finalmask.Filter";
+option go_package = "github.com/xtls/xray-core/transport/internet/finalmask/filter";
+option java_package = "com.xray.transport.internet.finalmask.filter";
+option java_multiple_files = true;
+
+import "common/geodata/geodat.proto";
+
+message Config {
+  repeated xray.common.geodata.IPRule blocked_source_ips = 1;
+  repeated xray.common.geodata.IPRule allowed_source_ips = 2;
+}

transport/internet/finalmask/filter/filter.go

@@ -0,0 +1,69 @@
+package filter
+
+import (
+	"github.com/xtls/xray-core/common/geodata"
+	"github.com/xtls/xray-core/common/net"
+)
+
+type listener struct {
+	net.Listener
+	filter *filter
+}
+
+// Accept filters connections after the TCP handshake and before higher-level processing.
+func (l *listener) Accept() (net.Conn, error) {
+	for {
+		conn, err := l.Listener.Accept()
+		if err != nil {
+			return nil, err
+		}
+		if l.filter.rejects(conn.RemoteAddr()) {
+			conn.Close()
+			continue
+		}
+		return conn, nil
+	}
+}
+
+type packetConn struct {
+	net.PacketConn
+	filter *filter
+}
+
+func (c *packetConn) ReadFrom(p []byte) (int, net.Addr, error) {
+	for {
+		n, addr, err := c.PacketConn.ReadFrom(p)
+		// Reads may return both data and an error; filter the data first.
+		if c.filter.rejects(addr) {
+			if err != nil {
+				return 0, nil, err
+			}
+			continue
+		}
+		return n, addr, err
+	}
+}
+
+type filter struct {
+	blocked geodata.IPMatcher
+	allowed geodata.IPMatcher
+}
+
+func (f *filter) rejects(addr net.Addr) bool {
+	var ip net.IP
+	switch addr := addr.(type) {
+	case *net.TCPAddr:
+		ip = addr.IP
+	case *net.UDPAddr:
+		ip = addr.IP
+	case *net.IPAddr:
+		ip = addr.IP
+	default:
+		// Unknown sources cannot satisfy an allowlist.
+		return f.allowed != nil
+	}
+	if f.blocked != nil && f.blocked.Match(ip) {
+		return true
+	}
+	return f.allowed != nil && !f.allowed.Match(ip)
+}

transport/internet/finalmask/filter/filter_test.go

@@ -0,0 +1,311 @@
+package filter
+
+import (
+	"context"
+	"errors"
+	"io"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/pires/go-proxyproto"
+	"github.com/xtls/xray-core/common/geodata"
+	xnet "github.com/xtls/xray-core/common/net"
+	"github.com/xtls/xray-core/transport/internet/finalmask"
+)
+
+func TestTCPFilterBeforeOtherMasks(t *testing.T) {
+	for _, policy := range testPolicies(t) {
+		t.Run(policy.name, func(t *testing.T) {
+			blocked4 := sourceConn("192.0.2.1")
+			blocked6 := sourceConn("2001:db8::1")
+			unknown := &testConn{addr: &net.UnixAddr{Name: "filter.sock", Net: "unix"}}
+			missing := new(testConn)
+			allowed := sourceConn("198.51.100.1")
+			allowed6 := sourceConn("2001:db8:1::1")
+			raw := &testListener{conns: []net.Conn{blocked4, blocked6, unknown, missing, allowed, allowed6}}
+			spy := new(recordMask)
+			fm := finalmask.NewFinalMask([]finalmask.TCPMask{spy, policy.config}, nil, nil,
+				func(context.Context, net.Addr) (net.Listener, error) { return raw, nil }, nil, nil)
+			l, err := fm.Listen(context.Background(), &net.TCPAddr{})
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer l.Close()
+			whitelist := len(policy.config.AllowedSourceIps) > 0
+			want := []net.Conn{allowed, allowed6}
+			if !whitelist {
+				want = append([]net.Conn{unknown, missing}, want...)
+			}
+			for _, expected := range want {
+				if got, err := l.Accept(); err != nil || got != expected {
+					t.Fatalf("Accept() = (%v, %v), want %v", got, err, expected)
+				}
+			}
+			// Fake sockets panic on Read or Write. Rejection must not require handshake data.
+			if !blocked4.closed || !blocked6.closed || unknown.closed != whitelist || missing.closed != whitelist || allowed.closed || allowed6.closed || raw.closed {
+				t.Fatal("filter did not close only the blocked connections")
+			}
+			if len(spy.seen) != len(want) {
+				t.Fatal("blocked connections reached the next mask")
+			}
+			for i, conn := range spy.seen {
+				if conn != want[i] {
+					t.Fatal("unexpected connection reached the next mask")
+				}
+			}
+			if conn, err := l.Accept(); conn != nil || !errors.Is(err, io.EOF) {
+				t.Fatalf("listener error was not propagated: (%v, %v)", conn, err)
+			}
+		})
+	}
+}
+
+func TestTCPFilterUsesProxyHeader(t *testing.T) {
+	for _, test := range []struct {
+		name, peer, reported string
+		blocked              bool
+	}{
+		{"blocked_header", "198.51.100.1", "192.0.2.1", true},
+		{"allowed_header", "192.0.2.1", "198.51.100.1", false},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			server, client := net.Pipe()
+			defer server.Close()
+			defer client.Close()
+			if err := server.SetDeadline(time.Now().Add(5 * time.Second)); err != nil {
+				t.Fatal(err)
+			}
+			peer := sourceConn(test.peer)
+			peer.Conn = server
+			conn := proxyproto.NewConn(peer)
+			next := sourceConn("198.51.100.2")
+			l, err := testConfig(t, "192.0.2.0/24").Listen(&testListener{conns: []net.Conn{conn, next}})
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer l.Close()
+			written := make(chan error, 1)
+			go func() {
+				_, err := client.Write([]byte("PROXY TCP4 " + test.reported + " 203.0.113.1 1234 443\r\n"))
+				written <- err
+			}()
+			got, err := l.Accept()
+			if err != nil {
+				t.Fatal(err)
+			}
+			if err := <-written; err != nil {
+				t.Fatal(err)
+			}
+			var want net.Conn = conn
+			if test.blocked {
+				want = next
+			}
+			if got != want || peer.closed != test.blocked {
+				t.Fatal("filter did not follow the address supplied by PROXY protocol")
+			}
+		})
+	}
+}
+
+func TestUDPFilter(t *testing.T) {
+	for _, policy := range testPolicies(t) {
+		t.Run(policy.name, func(t *testing.T) {
+			allowed := &net.UDPAddr{IP: net.ParseIP("198.51.100.1"), Port: 1234}
+			allowed6 := &net.UDPAddr{IP: net.ParseIP("2001:db8:1::1"), Port: 1234}
+			unknown := datagram{data: []byte("unknown"), addr: &net.UnixAddr{Name: "filter.sock", Net: "unixgram"}}
+			missing := datagram{data: []byte("missing address")}
+			raw := &testPacketConn{reads: []datagram{
+				{data: []byte("blocked"), addr: &net.UDPAddr{IP: net.ParseIP("192.0.2.1"), Port: 1234}},
+				{addr: &net.UDPAddr{IP: net.ParseIP("2001:db8::1"), Port: 1234}},
+				unknown,
+				missing,
+				{addr: allowed},
+				{data: []byte("allowed"), addr: allowed},
+				{data: []byte("allowed6"), addr: allowed6},
+				{data: []byte("blocked again"), addr: &net.UDPAddr{IP: net.ParseIP("192.0.2.2"), Port: 1234}},
+			}}
+			fm := finalmask.NewFinalMask(nil, []finalmask.UDPMask{policy.config}, nil, nil, nil,
+				func(context.Context, net.Addr) (net.PacketConn, error) { return raw, nil })
+			conn, err := fm.ListenPacket(context.Background(), &net.UDPAddr{})
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer conn.Close()
+			buf := make([]byte, 64)
+			if len(policy.config.AllowedSourceIps) == 0 {
+				for _, expected := range []datagram{unknown, missing} {
+					n, addr, err := conn.ReadFrom(buf)
+					if err != nil || addr != expected.addr || string(buf[:n]) != string(expected.data) {
+						t.Fatalf("unknown source was not allowed by blacklist: (%q, %v, %v)", buf[:n], addr, err)
+					}
+				}
+			}
+			if n, addr, err := conn.ReadFrom(buf); n != 0 || addr != allowed || err != nil {
+				t.Fatalf("zero-length allowed datagram was lost: (%d, %v, %v)", n, addr, err)
+			}
+			n, addr, err := conn.ReadFrom(buf)
+			if err != nil || addr != allowed || string(buf[:n]) != "allowed" {
+				t.Fatalf("ReadFrom() = (%q, %v, %v)", buf[:n], addr, err)
+			}
+			n, addr, err = conn.ReadFrom(buf)
+			if err != nil || addr != allowed6 || string(buf[:n]) != "allowed6" {
+				t.Fatalf("IPv6 ReadFrom() = (%q, %v, %v)", buf[:n], addr, err)
+			}
+			if _, _, err := conn.ReadFrom(buf); !errors.Is(err, io.EOF) {
+				t.Fatalf("underlying read error was not propagated: %v", err)
+			}
+			if raw.closed {
+				t.Fatal("dropping a packet closed the shared UDP socket")
+			}
+		})
+	}
+}
+
+func TestEmptyFilter(t *testing.T) {
+	for _, empty := range []*Config{{}, {AllowedSourceIps: []*geodata.IPRule{}}} {
+		rawListener := new(testListener)
+		rawPacket := new(testPacketConn)
+		if l, err := empty.Listen(rawListener); err != nil || l != rawListener {
+			t.Fatal("an empty TCP filter should preserve the listener")
+		}
+		if conn, err := empty.WrapPacketConnServer(rawPacket, nil, nil); err != nil || conn != rawPacket {
+			t.Fatal("an empty UDP filter should preserve the socket")
+		}
+	}
+}
+
+func TestUDPFilterReadError(t *testing.T) {
+	for _, policy := range testPolicies(t) {
+		t.Run(policy.name, func(t *testing.T) {
+			for _, test := range []struct {
+				name    string
+				source  net.Addr
+				blocked bool
+			}{
+				{"blocked", &net.UDPAddr{IP: net.ParseIP("192.0.2.1")}, true},
+				{"allowed", &net.UDPAddr{IP: net.ParseIP("198.51.100.1")}, false},
+				{"missing address", nil, len(policy.config.AllowedSourceIps) > 0},
+				{"unknown address", &net.UnixAddr{Name: "filter.sock", Net: "unixgram"}, len(policy.config.AllowedSourceIps) > 0},
+			} {
+				t.Run(test.name, func(t *testing.T) {
+					source := test.source
+					raw := &testPacketConn{reads: []datagram{{data: []byte("partial"), addr: source, err: io.ErrShortBuffer}}}
+					conn, err := policy.config.WrapPacketConnServer(raw, nil, nil)
+					if err != nil {
+						t.Fatal(err)
+					}
+					buf := make([]byte, 32)
+					n, addr, err := conn.ReadFrom(buf)
+					if !errors.Is(err, io.ErrShortBuffer) {
+						t.Fatalf("read error was not preserved: %v", err)
+					}
+					if test.blocked {
+						if n != 0 || addr != nil {
+							t.Fatal("partial data from a blocked source was returned with an error")
+						}
+					} else if n != 7 || addr != source || string(buf[:n]) != "partial" {
+						t.Fatal("partial data from an unblocked source was changed")
+					}
+				})
+			}
+		})
+	}
+}
+
+func testPolicies(t *testing.T) []struct {
+	name   string
+	config *Config
+} {
+	t.Helper()
+	blocked := testConfig(t, "192.0.2.0/24", "2001:db8::/48").BlockedSourceIps
+	allowed := testConfig(t, "198.51.100.0/24", "2001:db8:1::/48").BlockedSourceIps
+	broader := testConfig(t, "192.0.2.0/24", "198.51.100.0/24", "2001:db8::/32").BlockedSourceIps
+	return []struct {
+		name   string
+		config *Config
+	}{
+		{"blacklist", &Config{BlockedSourceIps: blocked}},
+		{"whitelist", &Config{AllowedSourceIps: allowed}},
+		{"both", &Config{BlockedSourceIps: blocked, AllowedSourceIps: broader}},
+	}
+}
+
+func testConfig(t *testing.T, entries ...string) *Config {
+	t.Helper()
+	rules, err := geodata.ParseIPRules(entries)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return &Config{BlockedSourceIps: rules}
+}
+
+type testConn struct {
+	net.Conn
+	addr   net.Addr
+	closed bool
+}
+
+func sourceConn(ip string) *testConn {
+	return &testConn{addr: &net.TCPAddr{IP: net.ParseIP(ip), Port: 1234}}
+}
+
+func (c *testConn) RemoteAddr() net.Addr { return c.addr }
+func (c *testConn) Close() error {
+	c.closed = true
+	if c.Conn != nil {
+		return c.Conn.Close()
+	}
+	return nil
+}
+
+type testListener struct {
+	net.Listener
+	conns  []net.Conn
+	closed bool
+}
+
+func (l *testListener) Accept() (net.Conn, error) {
+	if len(l.conns) == 0 {
+		return nil, io.EOF
+	}
+	conn := l.conns[0]
+	l.conns = l.conns[1:]
+	return conn, nil
+}
+
+func (l *testListener) Close() error { l.closed = true; return nil }
+
+type datagram struct {
+	data []byte
+	addr net.Addr
+	err  error
+}
+
+type testPacketConn struct {
+	net.PacketConn
+	reads  []datagram
+	closed bool
+}
+
+func (c *testPacketConn) ReadFrom(p []byte) (int, net.Addr, error) {
+	if len(c.reads) == 0 {
+		return 0, nil, io.EOF
+	}
+	packet := c.reads[0]
+	c.reads = c.reads[1:]
+	return copy(p, packet.data), packet.addr, packet.err
+}
+
+func (c *testPacketConn) Close() error { c.closed = true; return nil }
+
+type recordMask struct{ seen []net.Conn }
+
+func (m *recordMask) WrapConnClient(c net.Conn, _ *xnet.Destination, _ *finalmask.Dialer) (net.Conn, error) {
+	return c, nil
+}
+
+func (m *recordMask) WrapConnServer(c net.Conn) (net.Conn, error) {
+	m.seen = append(m.seen, c)
+	return c, nil
+}