godns/hosts.go

package main

import (
	"bufio"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/hoisie/redis"
	"golang.org/x/net/publicsuffix"
	"github.com/fsnotify/fsnotify"
)

type Hosts struct {
	providers []HostProvider
	refreshInterval time.Duration
}

type HostProvider interface {
	Get(domain string) ([]string, bool)
	Refresh()
}

func NewHosts(hs HostsSettings, rs RedisSettings) Hosts {
	providers := []HostProvider{
		NewFileProvider(hs.HostsFile),
	}

	if hs.RedisEnable {
		rc := &redis.Client{Addr: rs.Addr(), Db: rs.DB, Password: rs.Password}

		providers = append(providers, NewRedisProvider(rc, hs.RedisKey))
	}

	return Hosts{providers, time.Second * time.Duration(hs.RefreshInterval)}
}

func (h *Hosts) refresh() {
	ticker := time.NewTicker(h.refreshInterval)

	go func() {
		for {
			// Force a refresh every refreshInterval
			for _, provider := range h.providers {
				provider.Refresh()
			}

			<-ticker.C
		}
	}()
}

/*
Match local /etc/hosts file first, remote redis records second
*/
func (h *Hosts) Get(domain string, family int) ([]net.IP, bool) {
	var sips []string
	var ok bool
	var ip net.IP
	var ips []net.IP

	for _, provider := range h.providers {
		sips, ok = provider.Get(domain)

		if ok {
			break
		}
	}

	if sips == nil {
		return nil, false
	}

	for _, sip := range sips {
		switch family {
		case _IP4Query:
			ip = net.ParseIP(sip).To4()
		case _IP6Query:
			ip = net.ParseIP(sip).To16()
		default:
			continue
		}
		if ip != nil {
			ips = append(ips, ip)
		}
	}

	return ips, ips != nil
}

type RedisHosts struct {
	HostProvider

	redis *redis.Client
	key   string
	hosts map[string]string
	mu    sync.RWMutex
}

func NewRedisProvider(rc *redis.Client, key string) HostProvider {
	rh := &RedisHosts{
		redis: rc,
		key: key,
		hosts: make(map[string]string),
	}

	// Use pubsub to listen for key update events
	go func() {
		sub := make(chan string, 2)
		sub <- "godns:update"
		sub <- "godns:update_record"
		messages := make(chan redis.Message, 0)
		go rc.Subscribe(sub, nil, nil, nil, messages)

		for {
			msg := <- messages

			if msg.Channel == "godns:update" {
				rh.Refresh()
			} else if msg.Channel == "godns:update_record" {
				recordName := string(msg.Message)

				b, err := rc.Hget(key, recordName)

				if err != nil {
					continue
				}

				rh.hosts[recordName] = string(b)
			}
		}
	}()

	return rh
}

func (r *RedisHosts) Get(domain string) ([]string, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()

	domain = strings.ToLower(domain)
	ip, ok := r.hosts[domain]
	if ok {
		return strings.Split(ip, ","), true
	}

	sld, err := publicsuffix.EffectiveTLDPlusOne(domain)
	if err != nil {
		return nil, false
	}

	for host, ip := range r.hosts {
		if strings.HasPrefix(host, "*.") {
			old, err := publicsuffix.EffectiveTLDPlusOne(host)
			if err != nil {
				continue
			}
			if sld == old {
				return strings.Split(ip, ","), true
			}
		}
	}
	return nil, false
}

func (r *RedisHosts) Set(domain, ip string) (bool, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.redis.Hset(r.key, strings.ToLower(domain), []byte(ip))
}

func (r *RedisHosts) Refresh() {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.clear()
	err := r.redis.Hgetall(r.key, r.hosts)
	if err != nil {
		logger.Warn("Update hosts records from redis failed %s", err)
	} else {
		logger.Debug("Update hosts records from redis")
	}
}

func (r *RedisHosts) clear() {
	r.hosts = make(map[string]string)
}

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) {
	f.mu.RLock()
	defer f.mu.RUnlock()
	domain = strings.ToLower(domain)
	ip, ok := f.hosts[domain]
	if ok {
		return []string{ip}, true
	}

	sld, err := publicsuffix.EffectiveTLDPlusOne(domain)
	if err != nil {
		return nil, false
	}

	for host, ip := range f.hosts {
		if strings.HasPrefix(host, "*.") {
			old, err := publicsuffix.EffectiveTLDPlusOne(host)
			if err != nil {
				continue
			}
			if sld == old {
				return []string{ip}, true
			}
		}
	}

	return nil, false
}

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 := scanner.Text()
		line = strings.TrimSpace(line)
		line = strings.Replace(line, "\t", " ", -1)

		if strings.HasPrefix(line, "#") || line == "" {
			continue
		}

		sli := strings.Split(line, " ")

		if len(sli) < 2 {
			continue
		}

		ip := sli[0]
		if !isIP(ip) {
			continue
		}

		// 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 i := 1; i <= len(sli)-1; i++ {
			domain := strings.TrimSpace(sli[i])
			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)
}