119 lines
2.3 KiB
Go
119 lines
2.3 KiB
Go
package main
|
|
|
|
import (
|
|
"github.com/hoisie/redis"
|
|
"sync"
|
|
"strings"
|
|
"golang.org/x/net/publicsuffix"
|
|
)
|
|
|
|
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() {
|
|
keyspaceEvent := "__keyspace@0__:" + key
|
|
|
|
sub := make(chan string, 2)
|
|
sub <- keyspaceEvent
|
|
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.mu.RLock()
|
|
rh.hosts[recordName] = string(b)
|
|
rh.mu.RUnlock()
|
|
} else if msg.Channel == "godns:remove_record" {
|
|
rh.mu.RLock()
|
|
delete(rh.hosts, string(msg.Message))
|
|
rh.mu.RUnlock()
|
|
} else if msg.Channel == keyspaceEvent {
|
|
rh.Refresh()
|
|
}
|
|
}
|
|
}()
|
|
|
|
return rh
|
|
}
|
|
|
|
func (r *RedisHosts) Get(domain string) ([]string, bool) {
|
|
logger.Debug("Checking redis provider for %s", domain)
|
|
|
|
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)
|
|
}
|