godns/hosts.go

package main

import (
	"bufio"
	"errors"
	"net"
	"os"
	"regexp"
	"strings"

	"github.com/hoisie/redis"
)

type Hosts struct {
	FileHosts  map[string]string
	RedisHosts *RedisHosts
}

func NewHosts(hs HostsSettings, rs RedisSettings) Hosts {
	fileHosts := &FileHosts{hs.HostsFile}

	var rc *redis.Client
	if hs.RedisEnable {
		rc = &redis.Client{Addr: rs.Addr(), Db: rs.DB, Password: rs.Password}
	} else {
		rc = nil
	}
	redisHosts := &RedisHosts{rc, hs.RedisKey}

	hosts := Hosts{fileHosts.GetAll(), redisHosts}
	return hosts

}

/*
1. Resolve hosts file only one times
2. Request redis on every query called, not found performance lose serious yet.
3. Match local /etc/hosts file first, remote redis records second
*/

func (h *Hosts) Get(domain string, family int) (ip net.IP, ok bool) {
	var sip string

	if sip, ok = h.FileHosts[domain]; !ok {
		if sip, ok = h.RedisHosts.Get(domain); !ok {
			return nil, false
		}
	}

	switch family {
	case _IP4Query:
		ip = net.ParseIP(sip).To4()
		return ip, (ip != nil)
	case _IP6Query:
		ip = net.ParseIP(sip).To16()
		return ip, (ip != nil)
	}
	return nil, false
}

func (h *Hosts) GetAll() map[string]string {

	m := make(map[string]string)
	for domain, ip := range h.RedisHosts.GetAll() {
		m[domain] = ip
	}
	for domain, ip := range h.FileHosts {
		m[domain] = ip
	}
	return m
}

type RedisHosts struct {
	redis *redis.Client
	key   string
}

func (r *RedisHosts) GetAll() map[string]string {
	if r.redis == nil {
		return map[string]string{}
	}
	var hosts = make(map[string]string)
	r.redis.Hgetall(r.key, hosts)
	return hosts
}

func (r *RedisHosts) Get(domain string) (ip string, ok bool) {
	if r.redis == nil {
		return "", false
	}
	b, err := r.redis.Hget(r.key, domain)
	return string(b), err == nil
}

func (r *RedisHosts) Set(domain, ip string) (bool, error) {
	if r.redis == nil {
		return false, errors.New("Redis not enabled")
	}
	return r.redis.Hset(r.key, domain, []byte(ip))
}

type FileHosts struct {
	file string
}

func (f *FileHosts) GetAll() map[string]string {
	var hosts = make(map[string]string)

	buf, err := os.Open(f.file)
	if err != nil {
		panic("Can't open " + f.file)
	}

	scanner := bufio.NewScanner(buf)
	for scanner.Scan() {

		line := scanner.Text()
		line = strings.TrimSpace(line)

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

		sli := strings.Split(line, " ")
		if len(sli) == 1 {
			sli = strings.Split(line, "\t")
		}

		if len(sli) < 2 {
			continue
		}

		domain := sli[len(sli)-1]
		ip := sli[0]
		if !f.isDomain(domain) || !f.isIP(ip) {
			continue
		}

		hosts[domain] = ip
	}
	return hosts
}

func (f *FileHosts) isDomain(domain string) bool {
	if f.isIP(domain) {
		return false
	}
	match, _ := regexp.MatchString("^[a-zA-Z0-9][a-zA-Z0-9-]", domain)
	return match
}

func (f *FileHosts) isIP(ip string) bool {
	return (net.ParseIP(ip) != nil)
}