godns/hosts_file.go

136 lines
2.3 KiB
Go

package main
import (
"bufio"
"github.com/fsnotify/fsnotify"
"github.com/ryanuber/go-glob"
"os"
"regexp"
"strings"
"sync"
)
type FileHosts struct {
HostProvider
file string
hosts map[string]string
mu sync.RWMutex
}
func NewFileProvider(file string) HostProvider {
fp := &FileHosts{
file: file,
hosts: make(map[string]string),
}
watcher, err := fsnotify.NewWatcher()
// Use fsnotify to notify us of host file changes
if err == nil {
watcher.Add(file)
go func() {
for {
e := <- watcher.Events
if e.Op == fsnotify.Write {
fp.Refresh()
}
}
}()
}
return fp
}
func (f *FileHosts) Get(domain string) ([]string, bool) {
logger.Debug("Checking file provider for %s", domain)
f.mu.RLock()
defer f.mu.RUnlock()
domain = strings.ToLower(domain)
if ip, ok := f.hosts[domain]; ok {
return strings.Split(ip, ","), true
}
if idx := strings.Index(domain, "."); idx != -1 {
wildcard := "*." + domain[strings.Index(domain, ".") + 1:]
if ip, ok := f.hosts[wildcard]; ok {
return strings.Split(ip, ","), true
}
}
for host, ip := range f.hosts {
if glob.Glob(host, domain) {
return strings.Split(ip, ","), true
}
}
return nil, false
}
var (
hostRegexp = regexp.MustCompile("^(.*?)\\s+(.*)$")
)
func (f *FileHosts) Refresh() {
buf, err := os.Open(f.file)
if err != nil {
logger.Warn("Update hosts records from file failed %s", err)
return
}
defer buf.Close()
f.mu.Lock()
defer f.mu.Unlock()
f.clear()
scanner := bufio.NewScanner(buf)
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if strings.HasPrefix(line, "#") || line == "" {
continue
}
m := hostRegexp.FindStringSubmatch(line)
if m == nil {
continue
}
ip := m[1]
if !isIP(ip) {
continue
}
domains := strings.Split(m[2], " ")
// Would have multiple columns of domain in line.
// Such as "127.0.0.1 localhost localhost.domain" on linux.
// The domains may not strict standard, like "local" so don't check with f.isDomain(domain).
for _, domain := range domains {
domain = strings.TrimSpace(domain)
if domain == "" {
continue
}
f.hosts[strings.ToLower(domain)] = ip
}
}
logger.Debug("update hosts records from %s, total %d records.", f.file, len(f.hosts))
}
func (f *FileHosts) clear() {
f.hosts = make(map[string]string)
}