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) }