Implement client side ACL for SOCKS5 TCP

pkg/acl/engine.go

@@ -2,14 +2,23 @@ package acl
 
 import (
 	"bufio"
+	lru "github.com/hashicorp/golang-lru"
 	"net"
 	"os"
 	"strings"
 )
 
+const entryCacheSize = 1024
+
 type Engine struct {
 	DefaultAction Action
 	Entries       []Entry
+	Cache         *lru.ARCCache
+}
+
+type cacheEntry struct {
+	Action Action
+	Arg    string
 }
 
 func LoadFromFile(filename string) (*Engine, error) {
@@ -32,20 +41,48 @@ func LoadFromFile(filename string) (*Engine, error) {
 		}
 		entries = append(entries, entry)
 	}
+	cache, err := lru.NewARC(entryCacheSize)
+	if err != nil {
+		return nil, err
+	}
 	return &Engine{
 		DefaultAction: ActionProxy,
 		Entries:       entries,
+		Cache:         cache,
 	}, nil
 }
 
 func (e *Engine) Lookup(domain string, ip net.IP) (Action, string) {
-	if len(domain) == 0 && ip == nil {
+	if len(domain) > 0 {
+		// Domain
+		if v, ok := e.Cache.Get(domain); ok {
+			// Cache hit
+			ce := v.(cacheEntry)
+			return ce.Action, ce.Arg
+		}
+		ips, _ := net.LookupIP(domain)
+		for _, entry := range e.Entries {
+			if entry.MatchDomain(domain) || (len(ips) > 0 && entry.MatchIPs(ips)) {
+				e.Cache.Add(domain, cacheEntry{entry.Action, entry.ActionArg})
+				return entry.Action, entry.ActionArg
+			}
+		}
+		return e.DefaultAction, ""
+	} else if ip != nil {
+		// IP
+		if v, ok := e.Cache.Get(ip.String()); ok {
+			// Cache hit
+			ce := v.(cacheEntry)
+			return ce.Action, ce.Arg
+		}
+		for _, entry := range e.Entries {
+			if entry.MatchIP(ip) {
+				e.Cache.Add(ip.String(), cacheEntry{entry.Action, entry.ActionArg})
+				return entry.Action, entry.ActionArg
+			}
+		}
+		return e.DefaultAction, ""
+	} else {
 		return e.DefaultAction, ""
 	}
-	for _, entry := range e.Entries {
-		if entry.Match(domain, ip) {
-			return entry.Action, entry.ActionArg
-		}
-	}
-	return e.DefaultAction, ""
 }