feat: Add source port rule

Author: layou233

adapter/metadata.go

@@ -21,7 +21,7 @@ type Metadata struct {
 	ConnectionID        string
 	ServiceName         string
 	SniffedProtocol     Protocol
-	SourceIP            netip.Addr
+	SourceAddress       netip.AddrPort
 	DestinationHostname string
 	DestinationPort     uint16
 	Minecraft           *MinecraftMetadata

route/router.go

@@ -142,11 +142,11 @@ func (r *Router) HandleConnection(conn net.Conn, metadata *adapter.Metadata) {
 		}
 		r.access.RUnlock()
 		err = bufio.CopyConn(destinationConn, cachedConn)
+		logger := r.logger.Warn().Str("id", metadata.ConnectionID).Str("outbound", outbound.Name())
 		if err != nil {
-			r.logger.Warn().Str("id", metadata.ConnectionID).Str("outbound", outbound.Name()).Err(err).Msg("Handled connection")
-		} else {
-			r.logger.Info().Str("id", metadata.ConnectionID).Str("outbound", outbound.Name()).Msg("Handled connection")
+			logger = logger.Err(err)
 		}
+		logger.Msg("Handled outbound connection")
 		cachedConn.Close()
 		return
 	}

route/rule.go

@@ -39,6 +39,8 @@ func NewRule(logger *log.Logger, config *config.Rule, listMap map[string]set.Str
 		return NewSourceIPVersionRule(config)
 	case "SourceIP":
 		return NewSourceIPRule(config, listMap)
+	case "SourcePort":
+		return NewSourcePortRule(config)
 	case "MinecraftHostname":
 		return NewMinecraftHostnameRule(config, listMap)
 	case "MinecraftPlayerName":

route/rule_ip.go

@@ -83,7 +83,7 @@ func (r *RuleSourceIP) Config() *config.Rule {
 }
 
 func (r *RuleSourceIP) Match(metadata *adapter.Metadata) (match bool) {
-	match = r.set.Contains(metadata.SourceIP.WithZone(""))
+	match = r.set.Contains(metadata.SourceAddress.Addr().WithZone(""))
 	if r.config.Invert {
 		match = !match
 	}

route/rule_ipversion.go

@@ -32,7 +32,7 @@ func (r *RuleSourceIPVersion) Config() *config.Rule {
 }
 
 func (r *RuleSourceIPVersion) Match(metadata *adapter.Metadata) (match bool) {
-	if metadata.SourceIP.Is4() {
+	if metadata.SourceAddress.Addr().Is4() {
 		match = r.version == 4
 	} else {
 		match = r.version == 6

route/rule_port.go

@@ -0,0 +1,44 @@
+package route
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/layou233/zbproxy/v3/adapter"
+	"github.com/layou233/zbproxy/v3/config"
+)
+
+type RuleSourcePort struct {
+	config *config.Rule
+	ports  map[uint16]struct{}
+}
+
+var _ Rule = (*RuleSourcePort)(nil)
+
+func NewSourcePortRule(config *config.Rule) (*RuleSourcePort, error) {
+	var ports []uint16
+	err := json.Unmarshal(config.Parameter, &ports)
+	if err != nil {
+		return nil, fmt.Errorf("bad port list [%v]: %w", config.Parameter, err)
+	}
+	portsMap := make(map[uint16]struct{}, len(ports))
+	for _, port := range ports {
+		portsMap[port] = struct{}{}
+	}
+	return &RuleSourcePort{
+		config: config,
+		ports:  portsMap,
+	}, nil
+}
+
+func (r *RuleSourcePort) Config() *config.Rule {
+	return r.config
+}
+
+func (r *RuleSourcePort) Match(metadata *adapter.Metadata) (match bool) {
+	_, match = r.ports[metadata.SourceAddress.Port()]
+	if r.config.Invert {
+		match = !match
+	}
+	return
+}

service/service.go

@@ -50,8 +50,8 @@ func (s *Service) listenLoop() {
 		if err != nil {
 			return
 		}
-		ip := conn.RemoteAddr().(*net.TCPAddr).IP
-		ipString := ip.String()
+		tcpAddress := conn.RemoteAddr().(*net.TCPAddr)
+		ipString := tcpAddress.IP.String()
 		if s.ipAccessLists != nil &&
 			!access.Check(s.ipAccessLists, s.config.IPAccess.Mode, ipString) {
 			conn.SetLinger(0)
@@ -63,7 +63,7 @@ func (s *Service) listenLoop() {
 			ServiceName:         s.config.Name,
 			DestinationHostname: s.config.TargetAddress,
 			DestinationPort:     s.config.TargetPort,
-			SourceIP:            common.MustOK(netip.AddrFromSlice(ip)).Unmap(),
+			SourceAddress:       netip.AddrPortFrom(common.MustOK(netip.AddrFromSlice(tcpAddress.IP)).Unmap(), uint16(tcpAddress.Port)),
 		}
 		metadata.GenerateID()
 		s.logger.Info().Str("id", metadata.ConnectionID).Str("service", s.config.Name).