godns/hosts.go

125 lines
2.3 KiB
Go
Raw Normal View History

2013-07-26 04:06:16 +00:00
package main
import (
2013-07-26 10:54:19 +00:00
"bufio"
2013-07-26 04:06:16 +00:00
"github.com/hoisie/redis"
2013-07-26 10:54:19 +00:00
"os"
"regexp"
"strings"
2013-07-26 04:06:16 +00:00
)
2013-07-26 10:54:19 +00:00
type Hosts struct {
FileHosts map[string]string
RedisHosts *RedisHosts
2013-07-26 04:06:16 +00:00
}
2013-07-26 10:54:19 +00:00
func NewHosts(hs HostsSettings, rs RedisSettings) Hosts {
fileHosts := &FileHosts{hs.HostsFile}
redis := &redis.Client{Addr: rs.Addr(), Db: rs.DB, Password: rs.Password}
redisHosts := &RedisHosts{redis, hs.RedisKey}
hosts := Hosts{fileHosts.GetAll(), redisHosts}
return hosts
}
/*
1. Resolve hosts file only one times
2. Request redis on every get called, may lead to performance lose serious
3. Match local /etc/hosts file first, remote redis records second
*/
func (h *Hosts) Get(domain string) (ip string, ok bool) {
if ip, ok = h.FileHosts[domain]; ok {
return
}
if ip, ok = h.RedisHosts.Get(domain); ok {
return
}
return "", 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 {
var hosts = make(map[string]string)
r.redis.Hgetall(r.key, hosts)
return hosts
}
func (r *RedisHosts) Get(domain string) (ip string, ok bool) {
b, err := r.redis.Hget(r.key, domain)
return string(b), err == nil
2013-07-26 04:06:16 +00:00
}
2013-07-26 10:54:19 +00:00
func (r *RedisHosts) Set(domain, ip string) (bool, error) {
return r.redis.Hset(r.key, domain, []byte(ip))
}
type FileHosts struct {
file string
}
func (f *FileHosts) GetAll() map[string]string {
2013-07-26 04:06:16 +00:00
var hosts = make(map[string]string)
2013-07-26 10:54:19 +00:00
buf, err := os.Open(f.file)
if err != nil {
panic("Can't open " + f.file)
}
scanner := bufio.NewScanner(buf)
2013-07-26 04:06:16 +00:00
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]
2013-07-26 10:54:19 +00:00
if !f.isDomain(domain) || !f.isIP(ip) {
2013-07-26 04:06:16 +00:00
continue
}
hosts[domain] = ip
}
return hosts
}
2013-07-26 10:54:19 +00:00
func (f *FileHosts) isDomain(domain string) bool {
2013-07-26 04:06:16 +00:00
match, _ := regexp.MatchString("^[a-z]", domain)
return match
}
2013-07-26 10:54:19 +00:00
func (f *FileHosts) isIP(ip string) bool {
2013-07-26 04:06:16 +00:00
match, _ := regexp.MatchString("^[1-9]", ip)
return match
}