feat: ACL country support

This commit is contained in:
Toby
2022-01-09 18:11:52 -08:00
parent c3b76a5b44
commit 89452dd9c5
9 changed files with 136 additions and 29 deletions

View File

@@ -3,6 +3,7 @@ package acl
import (
"bufio"
lru "github.com/hashicorp/golang-lru"
"github.com/oschwald/geoip2-golang"
"github.com/tobyxdd/hysteria/pkg/transport"
"net"
"os"
@@ -16,6 +17,7 @@ type Engine struct {
Entries []Entry
Cache *lru.ARCCache
Transport transport.Transport
GeoIPReader *geoip2.Reader
}
type cacheEntry struct {
@@ -23,7 +25,7 @@ type cacheEntry struct {
Arg string
}
func LoadFromFile(filename string, transport transport.Transport) (*Engine, error) {
func LoadFromFile(filename string, transport transport.Transport, geoIPLoadFunc func() (*geoip2.Reader, error)) (*Engine, error) {
f, err := os.Open(filename)
if err != nil {
return nil, err
@@ -31,6 +33,7 @@ func LoadFromFile(filename string, transport transport.Transport) (*Engine, erro
defer f.Close()
scanner := bufio.NewScanner(f)
entries := make([]Entry, 0, 1024)
var geoIPReader *geoip2.Reader
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if len(line) == 0 || strings.HasPrefix(line, "#") {
@@ -41,6 +44,12 @@ func LoadFromFile(filename string, transport transport.Transport) (*Engine, erro
if err != nil {
return nil, err
}
if len(entry.Country) > 0 && geoIPReader == nil {
geoIPReader, err = geoIPLoadFunc() // lazy load GeoIP reader only when needed
if err != nil {
return nil, err
}
}
entries = append(entries, entry)
}
cache, err := lru.NewARC(entryCacheSize)
@@ -52,6 +61,7 @@ func LoadFromFile(filename string, transport transport.Transport) (*Engine, erro
Entries: entries,
Cache: cache,
Transport: transport,
GeoIPReader: geoIPReader,
}, nil
}
@@ -66,7 +76,7 @@ func (e *Engine) ResolveAndMatch(host string) (Action, string, *net.IPAddr, erro
return ce.Action, ce.Arg, ipAddr, err
}
for _, entry := range e.Entries {
if entry.MatchDomain(host) || (ipAddr != nil && entry.MatchIP(ipAddr.IP)) {
if entry.MatchDomain(host) || (ipAddr != nil && entry.MatchIP(ipAddr.IP, e.GeoIPReader)) {
e.Cache.Add(host, cacheEntry{entry.Action, entry.ActionArg})
return entry.Action, entry.ActionArg, ipAddr, err
}
@@ -84,7 +94,7 @@ func (e *Engine) ResolveAndMatch(host string) (Action, string, *net.IPAddr, erro
}, nil
}
for _, entry := range e.Entries {
if entry.MatchIP(ip) {
if entry.MatchIP(ip, e.GeoIPReader) {
e.Cache.Add(ip.String(), cacheEntry{entry.Action, entry.ActionArg})
return entry.Action, entry.ActionArg, &net.IPAddr{
IP: ip,