From 549a9aadc31a5aaec0f530b513e0b72591c9ceb6 Mon Sep 17 00:00:00 2001
From: layou233 <admin@launium.com>
Date: Sun, 21 Jul 2024 14:56:06 +0800
Subject: [PATCH] feat: New service name rule

---
 adapter/metadata.go     |  3 +-
 route/rule.go           |  2 ++
 route/rule_logic.go     |  4 +--
 route/rule_minecraft.go |  4 +--
 route/rule_service.go   | 63 +++++++++++++++++++++++++++++++++++++++++
 service/service.go      | 10 ++++---
 6 files changed, 77 insertions(+), 9 deletions(-)
 create mode 100644 route/rule_service.go

diff --git a/adapter/metadata.go b/adapter/metadata.go
index c0a42fa..bce6263 100644
--- a/adapter/metadata.go
+++ b/adapter/metadata.go
@@ -19,6 +19,7 @@ type Protocol = uint8
 
 type Metadata struct {
 	ConnectionID        string
+	ServiceName         string
 	SniffedProtocol     Protocol
 	SourceIP            netip.Addr
 	DestinationHostname string
@@ -30,7 +31,7 @@ type Metadata struct {
 
 func (m *Metadata) GenerateID() {
 	id := int64(fastrand.Int31())
-	idColor := fastrand.Intn(len(color.List))
+	idColor := fastrand.Int31n(int32(len(color.List)))
 	m.ConnectionID = color.Apply(color.List[idColor], "["+strconv.FormatInt(id, 10)+"]")
 }
 
diff --git a/route/rule.go b/route/rule.go
index 851adf6..ced5f0b 100644
--- a/route/rule.go
+++ b/route/rule.go
@@ -33,6 +33,8 @@ func NewRule(logger *log.Logger, config *config.Rule, listMap map[string]set.Str
 		return NewLogicalAndRule(logger, config, listMap, ruleRegistry)
 	case "or":
 		return NewLogicalOrRule(logger, config, listMap, ruleRegistry)
+	case "ServiceName":
+		return NewServiceNameRule(config, listMap)
 	case "SourceIPVersion":
 		return NewSourceIPVersionRule(config)
 	case "SourceIP":
diff --git a/route/rule_logic.go b/route/rule_logic.go
index 9af9dd0..92261bf 100644
--- a/route/rule_logic.go
+++ b/route/rule_logic.go
@@ -55,7 +55,7 @@ func NewLogicalAndRule(logger *log.Logger, newConfig *config.Rule, listMap map[s
 	return &RuleLogicalAnd{logicRule}, nil
 }
 
-func (r RuleLogicalAnd) Match(metadata *adapter.Metadata) (match bool) {
+func (r *RuleLogicalAnd) Match(metadata *adapter.Metadata) (match bool) {
 	match = true
 	for _, rule := range r.rules {
 		if !rule.Match(metadata) {
@@ -83,7 +83,7 @@ func NewLogicalOrRule(logger *log.Logger, newConfig *config.Rule, listMap map[st
 	return &RuleLogicalOr{logicRule}, nil
 }
 
-func (r RuleLogicalOr) Match(metadata *adapter.Metadata) (match bool) {
+func (r *RuleLogicalOr) Match(metadata *adapter.Metadata) (match bool) {
 	for _, rule := range r.rules {
 		if rule.Match(metadata) {
 			match = true
diff --git a/route/rule_minecraft.go b/route/rule_minecraft.go
index 3357bcf..865ac7c 100644
--- a/route/rule_minecraft.go
+++ b/route/rule_minecraft.go
@@ -45,11 +45,11 @@ func NewMinecraftPlayerNameRule(newConfig *config.Rule, listMap map[string]set.S
 	}, nil
 }
 
-func (r RuleMinecraftPlayerName) Config() *config.Rule {
+func (r *RuleMinecraftPlayerName) Config() *config.Rule {
 	return r.config
 }
 
-func (r RuleMinecraftPlayerName) Match(metadata *adapter.Metadata) (match bool) {
+func (r *RuleMinecraftPlayerName) Match(metadata *adapter.Metadata) (match bool) {
 	if metadata.Minecraft != nil {
 		for _, nameSet := range r.sets {
 			match = nameSet.Has(metadata.Minecraft.PlayerName)
diff --git a/route/rule_service.go b/route/rule_service.go
new file mode 100644
index 0000000..af4fbc1
--- /dev/null
+++ b/route/rule_service.go
@@ -0,0 +1,63 @@
+package route
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+
+	"github.com/layou233/zbproxy/v3/adapter"
+	"github.com/layou233/zbproxy/v3/common/jsonx"
+	"github.com/layou233/zbproxy/v3/common/set"
+	"github.com/layou233/zbproxy/v3/config"
+)
+
+type RuleServiceName struct {
+	sets   []set.StringSet
+	config *config.Rule
+}
+
+var _ Rule = (*RuleServiceName)(nil)
+
+func NewServiceNameRule(newConfig *config.Rule, listMap map[string]set.StringSet) (Rule, error) {
+	var serviceList jsonx.Listable[string]
+	err := json.Unmarshal(newConfig.Parameter, &serviceList)
+	if err != nil {
+		return nil, fmt.Errorf("bad service name list %v: %w", newConfig.Parameter, err)
+	}
+	sets := []set.StringSet{
+		{}, // new set for individual names
+	}
+	for _, i := range serviceList {
+		if strings.HasPrefix(i, parameterListPrefix) {
+			i = strings.TrimPrefix(i, parameterListPrefix)
+			nameSet, found := listMap[i]
+			if !found {
+				return nil, fmt.Errorf("list [%v] is not found", i)
+			}
+			sets = append(sets, nameSet)
+		} else {
+			sets[0].Add(i)
+		}
+	}
+	return &RuleServiceName{
+		sets:   sets,
+		config: newConfig,
+	}, nil
+}
+
+func (r RuleServiceName) Config() *config.Rule {
+	return r.config
+}
+
+func (r RuleServiceName) Match(metadata *adapter.Metadata) (match bool) {
+	for _, nameSet := range r.sets {
+		match = nameSet.Has(metadata.ServiceName)
+		if match {
+			break
+		}
+	}
+	if r.config.Invert {
+		match = !match
+	}
+	return
+}
diff --git a/service/service.go b/service/service.go
index 2283398..727fd2d 100644
--- a/service/service.go
+++ b/service/service.go
@@ -59,11 +59,13 @@ func (s *Service) listenLoop() {
 			s.logger.Warn().Str("service", s.config.Name).Str("ip", ipString).Msg("Rejected by access control")
 			continue
 		}
-		metadata := &adapter.Metadata{}
+		metadata := &adapter.Metadata{
+			ServiceName:         s.config.Name,
+			DestinationHostname: s.config.TargetAddress,
+			DestinationPort:     s.config.TargetPort,
+			SourceIP:            common.MustOK(netip.AddrFromSlice(ip)).Unmap(),
+		}
 		metadata.GenerateID()
-		metadata.DestinationHostname = s.config.TargetAddress
-		metadata.DestinationPort = s.config.TargetPort
-		metadata.SourceIP = common.MustOK(netip.AddrFromSlice(ip)).Unmap()
 		s.logger.Info().Str("id", metadata.ConnectionID).Str("service", s.config.Name).
 			Str("ip", ipString).Msg("New inbound connection")
 		if s.legacyOutbound != nil {