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)
+		}
+	}
+}