godns/hosts/hosts_file.go

158 lines
2.8 KiB
Go

package hosts
import (
"bufio"
"errors"
"github.com/fsnotify/fsnotify"
"github.com/miekg/dns"
"github.com/ryanuber/go-glob"
log "github.com/sirupsen/logrus"
"meow.tf/joker/godns/utils"
"os"
"regexp"
"strings"
"sync"
"time"
)
var (
errUnsupportedType = errors.New("unsupported type")
errRecordNotFound = errors.New("record not found")
errUnsupportedOperation = errors.New("unsupported operation")
)
type FileHosts struct {
Provider
file string
hosts map[string]Host
mu sync.RWMutex
ttl time.Duration
}
func NewFileProvider(file string, ttl time.Duration) Provider {
fp := &FileHosts{
file: file,
hosts: make(map[string]Host),
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) (*Host, error) {
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, errUnsupportedType
}
f.mu.RLock()
defer f.mu.RUnlock()
domain = strings.ToLower(domain)
if host, ok := f.hosts[domain]; ok {
return &host, nil
}
if idx := strings.Index(domain, "."); idx != -1 {
wildcard := "*." + domain[strings.Index(domain, ".") + 1:]
if host, ok := f.hosts[wildcard]; ok {
return &host, nil
}
}
for hostname, host := range f.hosts {
if glob.Glob(hostname, domain) {
return &host, nil
}
}
return nil, errRecordNotFound
}
func (f *FileHosts) Set(domain string, host *Host) error {
return errUnsupportedOperation
}
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 line == "" || line[0] == '#' {
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)] = Host{Values: []string{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]Host)
}