diff --git a/infra/conf/transport_finalmask.go b/infra/conf/transport_finalmask.go index eee3d2e43..6f3292313 100644 --- a/infra/conf/transport_finalmask.go +++ b/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"` diff --git a/infra/conf/transport_finalmask_filter_test.go b/infra/conf/transport_finalmask_filter_test.go new file mode 100644 index 000000000..63618cf50 --- /dev/null +++ b/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) + } + } +} diff --git a/transport/internet/finalmask/filter/config.go b/transport/internet/finalmask/filter/config.go new file mode 100644 index 000000000..f13b10579 --- /dev/null +++ b/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 +} diff --git a/transport/internet/finalmask/filter/config.pb.go b/transport/internet/finalmask/filter/config.pb.go new file mode 100644 index 000000000..d3f081dc7 --- /dev/null +++ b/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 +} diff --git a/transport/internet/finalmask/filter/config.proto b/transport/internet/finalmask/filter/config.proto new file mode 100644 index 000000000..b7b1f6af4 --- /dev/null +++ b/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; +} diff --git a/transport/internet/finalmask/filter/filter.go b/transport/internet/finalmask/filter/filter.go new file mode 100644 index 000000000..a9065f6b6 --- /dev/null +++ b/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) +} diff --git a/transport/internet/finalmask/filter/filter_test.go b/transport/internet/finalmask/filter/filter_test.go new file mode 100644 index 000000000..4780c1f3a --- /dev/null +++ b/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 +}