ACL protocol & port support

This commit is contained in:
Toby
2022-05-11 17:26:39 -07:00
parent 7a0977023e
commit e9974b0398
7 changed files with 422 additions and 176 deletions

View File

@@ -20,7 +20,13 @@ type Engine struct {
GeoIPReader *geoip2.Reader
}
type cacheEntry struct {
type cacheKey struct {
Host string
Port uint16
IsUDP bool
}
type cacheValue struct {
Action Action
Arg string
}
@@ -44,7 +50,7 @@ func LoadFromFile(filename string, resolveIPAddr func(string) (*net.IPAddr, erro
if err != nil {
return nil, err
}
if len(entry.Country) > 0 && geoIPReader == nil {
if _, ok := entry.Matcher.(*countryMatcher); ok && geoIPReader == nil {
geoIPReader, err = geoIPLoadFunc() // lazy load GeoIP reader only when needed
if err != nil {
return nil, err
@@ -66,44 +72,69 @@ func LoadFromFile(filename string, resolveIPAddr func(string) (*net.IPAddr, erro
}
// action, arg, isDomain, resolvedIP, error
func (e *Engine) ResolveAndMatch(host string) (Action, string, bool, *net.IPAddr, error) {
func (e *Engine) ResolveAndMatch(host string, port uint16, isUDP bool) (Action, string, bool, *net.IPAddr, error) {
ip, zone := utils.ParseIPZone(host)
if ip == nil {
// Domain
ipAddr, err := e.ResolveIPAddr(host)
if v, ok := e.Cache.Get(host); ok {
if v, ok := e.Cache.Get(cacheKey{host, port, isUDP}); ok {
// Cache hit
ce := v.(cacheEntry)
ce := v.(cacheValue)
return ce.Action, ce.Arg, true, ipAddr, err
}
for _, entry := range e.Entries {
if entry.MatchDomain(host) || (ipAddr != nil && entry.MatchIP(ipAddr.IP, e.GeoIPReader)) {
e.Cache.Add(host, cacheEntry{entry.Action, entry.ActionArg})
mReq := MatchRequest{
Domain: host,
Port: port,
DB: e.GeoIPReader,
}
if ipAddr != nil {
mReq.IP = ipAddr.IP
}
if isUDP {
mReq.Protocol = ProtocolUDP
} else {
mReq.Protocol = ProtocolTCP
}
if entry.Match(mReq) {
e.Cache.Add(cacheKey{host, port, isUDP},
cacheValue{entry.Action, entry.ActionArg})
return entry.Action, entry.ActionArg, true, ipAddr, err
}
}
e.Cache.Add(host, cacheEntry{e.DefaultAction, ""})
e.Cache.Add(cacheKey{host, port, isUDP}, cacheValue{e.DefaultAction, ""})
return e.DefaultAction, "", true, ipAddr, err
} else {
// IP
if v, ok := e.Cache.Get(ip.String()); ok {
if v, ok := e.Cache.Get(cacheKey{ip.String(), port, isUDP}); ok {
// Cache hit
ce := v.(cacheEntry)
ce := v.(cacheValue)
return ce.Action, ce.Arg, false, &net.IPAddr{
IP: ip,
Zone: zone,
}, nil
}
for _, entry := range e.Entries {
if entry.MatchIP(ip, e.GeoIPReader) {
e.Cache.Add(ip.String(), cacheEntry{entry.Action, entry.ActionArg})
mReq := MatchRequest{
IP: ip,
Port: port,
DB: e.GeoIPReader,
}
if isUDP {
mReq.Protocol = ProtocolUDP
} else {
mReq.Protocol = ProtocolTCP
}
if entry.Match(mReq) {
e.Cache.Add(cacheKey{ip.String(), port, isUDP},
cacheValue{entry.Action, entry.ActionArg})
return entry.Action, entry.ActionArg, false, &net.IPAddr{
IP: ip,
Zone: zone,
}, nil
}
}
e.Cache.Add(ip.String(), cacheEntry{e.DefaultAction, ""})
e.Cache.Add(cacheKey{ip.String(), port, isUDP}, cacheValue{e.DefaultAction, ""})
return e.DefaultAction, "", false, &net.IPAddr{
IP: ip,
Zone: zone,