feat(finalmask): add TCP and UDP source IP filtering

go.mod

@@ -9,7 +9,7 @@ require (
 	github.com/google/go-cmp v0.7.0
 	github.com/google/uuid v1.6.0
 	github.com/gorilla/websocket v1.5.3
-	github.com/klauspost/cpuid/v2 v2.4.1
+	github.com/klauspost/cpuid/v2 v2.4.0
 	github.com/libp2p/go-nat v1.0.1-0.20250821073202-01afc089f138
 	github.com/miekg/dns v1.1.73
 	github.com/pelletier/go-toml v1.9.5

go.sum

@@ -33,8 +33,8 @@ github.com/juju/ratelimit v1.0.2 h1:sRxmtRiajbvrcLQT7S+JbqU0ntsb9W2yhSdNN8tWfaI=
 github.com/juju/ratelimit v1.0.2/go.mod h1:qapgC/Gy+xNh9UxzV13HGGl/6UXNN+ct+vwSgWNm/qk=
 github.com/klauspost/compress v1.20.0 h1:a3C1ke2ohxFymNlb2HWAHjDeKCI90scRskErZkR0ezA=
 github.com/klauspost/compress v1.20.0/go.mod h1:LUdAzn7YLVvxLpc7y3V1m40wESHTgc1422pwwBSKYuI=
-github.com/klauspost/cpuid/v2 v2.4.1 h1:I9KxadAkvo+QjtcEE2f2tyfEvmIla/T0m+14PrZNl0U=
-github.com/klauspost/cpuid/v2 v2.4.1/go.mod h1:19jmZ9mjzoF//ddRSUsv0zfBTJWh3QJh9FNxZTMrGxU=
+github.com/klauspost/cpuid/v2 v2.4.0 h1:S6Hrbc7+ywsr0r+RLapfGBHfyefhCTwEh3A0tV913Dw=
+github.com/klauspost/cpuid/v2 v2.4.0/go.mod h1:19jmZ9mjzoF//ddRSUsv0zfBTJWh3QJh9FNxZTMrGxU=
 github.com/koron/go-ssdp v0.0.4 h1:1IDwrghSKYM7yLf7XCzbByg2sJ/JcNOZRXS2jczTwz0=
 github.com/koron/go-ssdp v0.0.4/go.mod h1:oDXq+E5IL5q0U8uSBcoAXzTzInwy5lEgC91HoKtbmZk=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=

infra/conf/transport_finalmask.go

@@ -15,11 +15,13 @@ 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/fragment"
 	"github.com/xtls/xray-core/transport/internet/finalmask/header/custom"
+	"github.com/xtls/xray-core/transport/internet/finalmask/ipfilter"
 	"github.com/xtls/xray-core/transport/internet/finalmask/mkcp/aes128gcm"
 	"github.com/xtls/xray-core/transport/internet/finalmask/mkcp/header"
 	"github.com/xtls/xray-core/transport/internet/finalmask/mkcp/original"
@@ -74,6 +76,7 @@ var (
 	customVarNamePattern = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_]*$`)
 
 	tcpmaskLoader = NewJSONConfigLoader(ConfigCreatorCache{
+		"ip-filter":     func() interface{} { return new(IPFilter) },
 		"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{
+		"ip-filter":     func() interface{} { return new(IPFilter) },
 		"header-custom": func() interface{} { return new(HeaderCustomUDP) },
 		"mkcp-legacy":   func() interface{} { return new(MkcpLegacy) },
 		"noise":         func() interface{} { return new(NoiseMask) },
@@ -93,6 +97,18 @@ var (
 	}, "type", "settings")
 )
 
+type IPFilter struct {
+	BlockedIPs []string `json:"blockedIPs"`
+}
+
+func (c *IPFilter) Build() (proto.Message, error) {
+	rules, err := geodata.ParseIPRules(c.BlockedIPs)
+	if err != nil {
+		return nil, errors.New(`invalid "blockedIPs"`).Base(err)
+	}
+	return &ipfilter.Config{BlockedIps: rules}, nil
+}
+
 type TCPItem struct {
 	Delay     Int32Range       `json:"delay"`
 	Rand      int32            `json:"rand"`

infra/conf/transport_finalmask_ipfilter_test.go

@@ -0,0 +1,57 @@
+package conf_test
+
+import (
+	"encoding/json"
+	"fmt"
+	"net"
+	"strings"
+	"testing"
+
+	. "github.com/xtls/xray-core/infra/conf"
+	"github.com/xtls/xray-core/transport/internet"
+	"github.com/xtls/xray-core/transport/internet/finalmask/ipfilter"
+)
+
+func TestFinalMaskIPFilter(t *testing.T) {
+	for _, network := range []string{"tcp", "udp"} {
+		t.Run(network, func(t *testing.T) {
+			input := fmt.Sprintf(`{"finalmask":{"%s":[{"type":"ip-filter","settings":{"blockedIPs":["192.0.2.0/24"]}}]}}`, network)
+			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.(*ipfilter.Config)
+			if !ok {
+				t.Fatalf("unexpected mask type: %T", instance)
+			}
+			if len(filter.BlockedIps) != 1 {
+				t.Fatalf("expected one blocked IP rule, got %d", len(filter.BlockedIps))
+			}
+			cidr := filter.BlockedIps[0].GetCustom().GetCidr()
+			if cidr.GetPrefix() != 24 || !net.IP(cidr.GetIp()).Equal(net.IPv4(192, 0, 2, 0)) {
+				t.Fatalf("unexpected blocked IP rule: %v", cidr)
+			}
+			if _, err := internet.ToMemoryStreamConfig(built); err != nil {
+				t.Fatalf("failed to initialize runtime finalmask: %v", err)
+			}
+		})
+	}
+	if _, err := (&IPFilter{BlockedIPs: []string{"192.0.2.0/33"}}).Build(); err == nil || !strings.Contains(err.Error(), "blockedIPs") {
+		t.Errorf("expected invalid blockedIPs, got %v", err)
+	}
+}

transport/internet/finalmask/ipfilter/config.go

@@ -0,0 +1,49 @@
+package ipfilter
+
+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) matcher() (geodata.IPMatcher, error) {
+	if len(c.BlockedIps) == 0 {
+		return nil, nil
+	}
+	return geodata.IPReg.BuildIPMatcher(c.BlockedIps)
+}
+
+func (c *Config) Listen(inner net.Listener) (net.Listener, error) {
+	matcher, err := c.matcher()
+	if err != nil {
+		return nil, err
+	}
+	if matcher == nil {
+		return inner, nil
+	}
+	return &listener{Listener: inner, matcher: matcher}, 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) {
+	matcher, err := c.matcher()
+	if err != nil {
+		return nil, err
+	}
+	if matcher == nil {
+		return conn, nil
+	}
+	return &packetConn{PacketConn: conn, matcher: matcher}, nil
+}

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

@@ -0,0 +1,127 @@
+// Code generated by protoc-gen-go. DO NOT EDIT.
+// versions:
+// 	protoc-gen-go v1.36.11
+// 	protoc        v6.33.5
+// source: transport/internet/finalmask/ipfilter/config.proto
+
+package ipfilter
+
+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"`
+	BlockedIps    []*geodata.IPRule      `protobuf:"bytes,1,rep,name=blocked_ips,json=blockedIps,proto3" json:"blocked_ips,omitempty"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
+}
+
+func (x *Config) Reset() {
+	*x = Config{}
+	mi := &file_transport_internet_finalmask_ipfilter_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_ipfilter_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_ipfilter_config_proto_rawDescGZIP(), []int{0}
+}
+
+func (x *Config) GetBlockedIps() []*geodata.IPRule {
+	if x != nil {
+		return x.BlockedIps
+	}
+	return nil
+}
+
+var File_transport_internet_finalmask_ipfilter_config_proto protoreflect.FileDescriptor
+
+const file_transport_internet_finalmask_ipfilter_config_proto_rawDesc = "" +
+	"\n" +
+	"2transport/internet/finalmask/ipfilter/config.proto\x12*xray.transport.internet.finalmask.ipfilter\x1a\x1bcommon/geodata/geodat.proto\"F\n" +
+	"\x06Config\x12<\n" +
+	"\vblocked_ips\x18\x01 \x03(\v2\x1b.xray.common.geodata.IPRuleR\n" +
+	"blockedIpsB\xa0\x01\n" +
+	".com.xray.transport.internet.finalmask.ipfilterP\x01Z?github.com/xtls/xray-core/transport/internet/finalmask/ipfilter\xaa\x02*Xray.Transport.Internet.Finalmask.Ipfilterb\x06proto3"
+
+var (
+	file_transport_internet_finalmask_ipfilter_config_proto_rawDescOnce sync.Once
+	file_transport_internet_finalmask_ipfilter_config_proto_rawDescData []byte
+)
+
+func file_transport_internet_finalmask_ipfilter_config_proto_rawDescGZIP() []byte {
+	file_transport_internet_finalmask_ipfilter_config_proto_rawDescOnce.Do(func() {
+		file_transport_internet_finalmask_ipfilter_config_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_transport_internet_finalmask_ipfilter_config_proto_rawDesc), len(file_transport_internet_finalmask_ipfilter_config_proto_rawDesc)))
+	})
+	return file_transport_internet_finalmask_ipfilter_config_proto_rawDescData
+}
+
+var file_transport_internet_finalmask_ipfilter_config_proto_msgTypes = make([]protoimpl.MessageInfo, 1)
+var file_transport_internet_finalmask_ipfilter_config_proto_goTypes = []any{
+	(*Config)(nil),         // 0: xray.transport.internet.finalmask.ipfilter.Config
+	(*geodata.IPRule)(nil), // 1: xray.common.geodata.IPRule
+}
+var file_transport_internet_finalmask_ipfilter_config_proto_depIdxs = []int32{
+	1, // 0: xray.transport.internet.finalmask.ipfilter.Config.blocked_ips:type_name -> xray.common.geodata.IPRule
+	1, // [1:1] is the sub-list for method output_type
+	1, // [1:1] is the sub-list for method input_type
+	1, // [1:1] is the sub-list for extension type_name
+	1, // [1:1] is the sub-list for extension extendee
+	0, // [0:1] is the sub-list for field type_name
+}
+
+func init() { file_transport_internet_finalmask_ipfilter_config_proto_init() }
+func file_transport_internet_finalmask_ipfilter_config_proto_init() {
+	if File_transport_internet_finalmask_ipfilter_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_ipfilter_config_proto_rawDesc), len(file_transport_internet_finalmask_ipfilter_config_proto_rawDesc)),
+			NumEnums:      0,
+			NumMessages:   1,
+			NumExtensions: 0,
+			NumServices:   0,
+		},
+		GoTypes:           file_transport_internet_finalmask_ipfilter_config_proto_goTypes,
+		DependencyIndexes: file_transport_internet_finalmask_ipfilter_config_proto_depIdxs,
+		MessageInfos:      file_transport_internet_finalmask_ipfilter_config_proto_msgTypes,
+	}.Build()
+	File_transport_internet_finalmask_ipfilter_config_proto = out.File
+	file_transport_internet_finalmask_ipfilter_config_proto_goTypes = nil
+	file_transport_internet_finalmask_ipfilter_config_proto_depIdxs = nil
+}

transport/internet/finalmask/ipfilter/config.proto

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

transport/internet/finalmask/ipfilter/ip_filter.go

@@ -0,0 +1,58 @@
+package ipfilter
+
+import (
+	"github.com/xtls/xray-core/common/geodata"
+	"github.com/xtls/xray-core/common/net"
+)
+
+type listener struct {
+	net.Listener
+	matcher geodata.IPMatcher
+}
+
+// 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 matchAddress(l.matcher, conn.RemoteAddr()) {
+			conn.Close()
+			continue
+		}
+		return conn, nil
+	}
+}
+
+type packetConn struct {
+	net.PacketConn
+	matcher geodata.IPMatcher
+}
+
+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 matchAddress(c.matcher, addr) {
+			if err != nil {
+				return 0, nil, err
+			}
+			continue
+		}
+		return n, addr, err
+	}
+}
+
+func matchAddress(matcher geodata.IPMatcher, addr net.Addr) bool {
+	switch addr := addr.(type) {
+	case *net.TCPAddr:
+		return matcher.Match(addr.IP)
+	case *net.UDPAddr:
+		return matcher.Match(addr.IP)
+	case *net.IPAddr:
+		return matcher.Match(addr.IP)
+	default:
+		return true
+	}
+}

transport/internet/finalmask/ipfilter/ip_filter_test.go

@@ -0,0 +1,251 @@
+package ipfilter
+
+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) {
+	blocked4 := sourceConn("192.0.2.1")
+	blocked6 := sourceConn("2001:db8::1")
+	unknown := &testConn{addr: &net.UnixAddr{Name: "ipfilter.sock", Net: "unix"}}
+	missing := new(testConn)
+	allowed := sourceConn("198.51.100.1")
+	raw := &testListener{conns: []net.Conn{blocked4, blocked6, unknown, missing, allowed}}
+	spy := new(recordMask)
+	fm := finalmask.NewFinalMask([]finalmask.TCPMask{spy, testConfig(t, "192.0.2.0/24", "2001:db8::/32")}, 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()
+	got, err := l.Accept()
+	if err != nil || got != allowed {
+		t.Fatalf("Accept() = (%v, %v), want unblocked connection", got, err)
+	}
+	// Fake sockets panic on Read or Write. Rejection must not require handshake data.
+	if !blocked4.closed || !blocked6.closed || !unknown.closed || !missing.closed || allowed.closed || raw.closed {
+		t.Fatal("filter did not close only the blocked connections")
+	}
+	if len(spy.seen) != 1 || spy.seen[0] != allowed {
+		t.Fatal("blocked connections 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) {
+	allowed := &net.UDPAddr{IP: net.ParseIP("198.51.100.1"), Port: 1234}
+	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}},
+		{data: []byte("unknown"), addr: &net.UnixAddr{Name: "ipfilter.sock", Net: "unixgram"}},
+		{data: []byte("missing address")},
+		{addr: allowed},
+		{data: []byte("allowed"), addr: allowed},
+		{data: []byte("blocked again"), addr: &net.UDPAddr{IP: net.ParseIP("192.0.2.2"), Port: 1234}},
+	}}
+	fm := finalmask.NewFinalMask(nil, []finalmask.UDPMask{testConfig(t, "192.0.2.0/24", "2001:db8::/32")}, 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 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)
+	}
+	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) {
+	rawListener := new(testListener)
+	rawPacket := new(testPacketConn)
+	empty := new(Config)
+	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 _, 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, true},
+	} {
+		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 := testConfig(t, "192.0.2.0/24").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 testConfig(t *testing.T, entries ...string) *Config {
+	t.Helper()
+	rules, err := geodata.ParseIPRules(entries)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return &Config{BlockedIps: 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
+}