godns/hosts_file.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)
}