godns/hosts.go

173 lines
3.1 KiB
Go

package main
import (
"bufio"
"net"
"os"
"regexp"
"strings"
"time"
"github.com/hoisie/redis"
)
type Hosts struct {
fileHosts *FileHosts
redisHosts *RedisHosts
}
func NewHosts(hs HostsSettings, rs RedisSettings) Hosts {
fileHosts := &FileHosts{BaseHosts{make(map[string]string)}, hs.HostsFile}
var redisHosts *RedisHosts
if hs.RedisEnable {
rc := &redis.Client{Addr: rs.Addr(), Db: rs.DB, Password: rs.Password}
redisHosts = &RedisHosts{BaseHosts{make(map[string]string)}, rc, hs.RedisKey}
}
hosts := Hosts{fileHosts, redisHosts}
hosts.refresh()
return hosts
}
/*
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.Get(domain); !ok {
if h.redisHosts != nil {
sip, ok = h.redisHosts.Get(domain)
}
}
if sip == "" {
return nil, false
}
switch family {
case _IP4Query:
ip = net.ParseIP(sip).To4()
case _IP6Query:
ip = net.ParseIP(sip).To16()
default:
return nil, false
}
return ip, (ip != nil)
}
/*
Update hosts records from /etc/hosts file and redis per minute
*/
func (h *Hosts) refresh() {
ticker := time.NewTicker(time.Minute)
go func() {
for {
h.fileHosts.Refresh()
if h.redisHosts != nil {
h.redisHosts.Refresh()
}
<-ticker.C
}
}()
}
type BaseHosts struct {
hosts map[string]string
}
func (h *BaseHosts) Get(domain string) (ip string, ok bool) {
ip, ok = h.hosts[domain]
if ok {
return
}
for host, ip := range h.hosts {
if strings.HasPrefix(host, "*.") {
upperLevelDomain := strings.Split(host, "*.")[1]
if strings.HasSuffix(domain, upperLevelDomain) {
return ip, true
}
}
}
return
}
type RedisHosts struct {
BaseHosts
redis *redis.Client
key string
}
func (r *RedisHosts) Set(domain, ip string) (bool, error) {
return r.redis.Hset(r.key, domain, []byte(ip))
}
func (r *RedisHosts) Refresh() {
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")
}
}
type FileHosts struct {
BaseHosts
file string
}
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()
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
}
f.hosts[domain] = ip
}
logger.Debug("update hosts records from %s", f.file)
}
func (f *FileHosts) isDomain(domain string) bool {
if f.isIP(domain) {
return false
}
match, _ := regexp.MatchString(`^([a-zA-Z0-9\*]([a-zA-Z0-9\-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,6}$`, domain)
return match
}
func (f *FileHosts) isIP(ip string) bool {
return (net.ParseIP(ip) != nil)
}