godns/hosts/hosts_file.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) List() (HostMap, error) {
	return nil, errUnsupportedOperation
}

func (f *FileHosts) Get(queryType uint16, domain string) (*Host, error) {
	log.WithFields(log.Fields{
		"queryType": dns.TypeToString[queryType],
		"question": domain,
	}).Debug("Checking file provider")

	// 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
}

var (
	hostRegexp = regexp.MustCompile("^(.*?)\\s+(.*)$")
)

func (f *FileHosts) Refresh() {
	buf, err := os.Open(f.file)

	if err != nil {
		log.WithFields(log.Fields{
			"file": f.file,
			"error": err,
		}).Warn("Hosts update from file failed")
		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.WithFields(log.Fields{
		"file": f.file,
		"count": len(f.hosts),
	}).Debug("Updated hosts records")
}

func (f *FileHosts) clear() {
	f.hosts = make(map[string]Host)
}