A simple, go-based DNS resolver/caching server
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

128 lines
2.7 KiB

package main
import (
"github.com/hoisie/redis"
"github.com/ryanuber/go-glob"
"strings"
"sync"
)
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),
}
// Force an initial refresh
rh.Refresh()
// Use pubsub to listen for key update events
go func() {
keyspaceEvent := "__keyspace@0__:" + key
sub := make(chan string, 3)
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" {
logger.Debug("Refreshing redis records due to update")
rh.Refresh()
} else if msg.Channel == "godns:update_record" {
recordName := string(msg.Message)
b, err := rc.Hget(key, recordName)
if err != nil {
logger.Warn("Record %s does not exist, but was updated", recordName)
continue
}
logger.Debug("Record %s was updated to %s", recordName, string(b))
rh.mu.Lock()
rh.hosts[recordName] = string(b)
rh.mu.Unlock()
} else if msg.Channel == "godns:remove_record" {
logger.Debug("Record %s was removed", msg.Message)
recordName := string(msg.Message)
rh.mu.Lock()
delete(rh.hosts, recordName)
rh.mu.Unlock()
} else if msg.Channel == keyspaceEvent {
logger.Debug("Refreshing redis records due to update")
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)
if ip, ok := r.hosts[domain]; ok {
return strings.Split(ip, ","), true
}
if idx := strings.Index(domain, "."); idx != -1 {
wildcard := "*." + domain[strings.Index(domain, ".")+1:]
if ip, ok := r.hosts[wildcard]; ok {
return strings.Split(ip, ","), true
}
}
for host, ip := range r.hosts {
if glob.Glob(host, domain) {
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)
}