158 lines
2.8 KiB
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)
|
|
}
|