godns/hosts/hosts_file.go

152 lines
2.8 KiB
Go

package hosts
import (
"bufio"
"errors"
"github.com/fsnotify/fsnotify"
"github.com/miekg/dns"
"github.com/ryanuber/go-glob"
"meow.tf/joker/godns/log"
"meow.tf/joker/godns/utils"
"os"
"regexp"
"strings"
"sync"
"time"
)
type FileHosts struct {
Provider
file string
hosts map[string]string
mu sync.RWMutex
ttl time.Duration
}
func NewFileProvider(file string, ttl time.Duration) Provider {
fp := &FileHosts{
file: file,
hosts: make(map[string]string),
ttl: ttl,
}
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(queryType uint16, domain string) ([]string, time.Duration, bool) {
log.Debug("Checking file provider for %s : %s", queryType, domain)
// Does not support CNAME/TXT/etc
if queryType != dns.TypeA && queryType != dns.TypeAAAA {
return nil, zeroDuration, false
}
f.mu.RLock()
defer f.mu.RUnlock()
domain = strings.ToLower(domain)
if ip, ok := f.hosts[domain]; ok {
return strings.Split(ip, ","), f.ttl, 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, ","), f.ttl, true
}
}
for host, ip := range f.hosts {
if glob.Glob(host, domain) {
return strings.Split(ip, ","), f.ttl, true
}
}
return nil, time.Duration(0), false
}
func (f *FileHosts) Set(t, domain, value string) (bool, error) {
return false, errors.New("file provider does not support setting values")
}
var (
hostRegexp = regexp.MustCompile("^(.*?)\\s+(.*)$")
)
func (f *FileHosts) Refresh() {
buf, err := os.Open(f.file)
if err != nil {
log.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 !utils.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
}
}
log.Debug("update hosts records from %s, total %d records.", f.file, len(f.hosts))
}
func (f *FileHosts) clear() {
f.hosts = make(map[string]string)
}