refractor: create limiter conn

SagerNet · Nov 16, 2023 · 630cdb0 · 630cdb0
1 parent ef6b563
commit 630cdb0
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 23 deletions.
diff --git a/limiter/builder.go b/limiter/builder.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net"
 
+	"github.com/sagernet/sing-box/adapter"
 	"github.com/sagernet/sing-box/common/humanize"
 	"github.com/sagernet/sing-box/log"
 	"github.com/sagernet/sing-box/option"
@@ -86,22 +87,23 @@ func (m *defaultManager) createLimiter(ctx context.Context, option option.Limite
 	return
 }
 
-func (m *defaultManager) LoadLimiters(tags []string, user, inbound string) (limiters []*limiter) {
-	for _, tag := range tags {
-		if v, ok := m.mp[limiterKey{prefixTag, tag}]; ok {
-			limiters = append(limiters, v)
+func (m *defaultManager) NewConnWithLimiters(ctx context.Context, conn net.Conn, metadata *adapter.InboundContext, rule adapter.Rule) net.Conn {
+	var limiters []*limiter
+	if rule != nil {
+		for _, tag := range rule.Limiters() {
+			if v, ok := m.mp[limiterKey{prefixTag, tag}]; ok {
+				limiters = append(limiters, v)
+			}
 		}
 	}
-	if v, ok := m.mp[limiterKey{prefixUser, user}]; ok {
-		limiters = append(limiters, v)
-	}
-	if v, ok := m.mp[limiterKey{prefixInbound, inbound}]; ok {
-		limiters = append(limiters, v)
+	if metadata != nil {
+		if v, ok := m.mp[limiterKey{prefixUser, metadata.User}]; ok {
+			limiters = append(limiters, v)
+		}
+		if v, ok := m.mp[limiterKey{prefixInbound, metadata.Inbound}]; ok {
+			limiters = append(limiters, v)
+		}
 	}
-	return
-}
-
-func (m *defaultManager) NewConnWithLimiters(ctx context.Context, conn net.Conn, limiters []*limiter) net.Conn {
 	for _, limiter := range limiters {
 		conn = &connWithLimiter{Conn: conn, limiter: limiter, ctx: ctx}
 	}

diff --git a/limiter/manager.go b/limiter/manager.go
@@ -3,9 +3,10 @@ package limiter
 import (
 	"context"
 	"net"
+
+	"github.com/sagernet/sing-box/adapter"
 )
 
 type Manager interface {
-	LoadLimiters(tags []string, user, inbound string) []*limiter
-	NewConnWithLimiters(ctx context.Context, conn net.Conn, limiters []*limiter) net.Conn
+	NewConnWithLimiters(ctx context.Context, conn net.Conn, metadata *adapter.InboundContext, rule adapter.Rule) net.Conn
 }
diff --git a/route/router.go b/route/router.go
@@ -708,14 +708,7 @@ func (r *Router) RouteConnection(ctx context.Context, conn net.Conn, metadata ad
 	}
 
 	if r.limiterManager != nil {
-		var limiterTags []string
-		if matchedRule != nil {
-			limiterTags = matchedRule.Limiters()
-		}
-		limiters := r.limiterManager.LoadLimiters(limiterTags, metadata.User, metadata.Inbound)
-		if len(limiters) > 0 {
-			conn = r.limiterManager.NewConnWithLimiters(ctx, conn, limiters)
-		}
+		conn = r.limiterManager.NewConnWithLimiters(ctx, conn, &metadata, matchedRule)
 	}
 
 	if r.clashServer != nil {