godns/hosts_redis.go

129 lines
2.7 KiB
Go
Raw Permalink Normal View History

package main
import (
"github.com/hoisie/redis"
2019-09-26 04:43:17 +00:00
"github.com/ryanuber/go-glob"
"strings"
2019-09-26 04:43:17 +00:00
"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,
2019-09-26 04:43:17 +00:00
key: key,
hosts: make(map[string]string),
}
2018-08-05 08:48:26 +00:00
// Force an initial refresh
rh.Refresh()
// Use pubsub to listen for key update events
go func() {
keyspaceEvent := "__keyspace@0__:" + key
2018-08-05 08:48:26 +00:00
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 {
2019-09-26 04:43:17 +00:00
msg := <-messages
if msg.Channel == "godns:update" {
2018-08-05 08:48:26 +00:00
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 {
2018-08-05 08:48:26 +00:00
logger.Warn("Record %s does not exist, but was updated", recordName)
continue
}
2018-08-05 08:48:26 +00:00
logger.Debug("Record %s was updated to %s", recordName, string(b))
rh.mu.Lock()
rh.hosts[recordName] = string(b)
2018-08-05 08:48:26 +00:00
rh.mu.Unlock()
2018-07-03 05:35:20 +00:00
} else if msg.Channel == "godns:remove_record" {
2018-08-05 08:48:26 +00:00
logger.Debug("Record %s was removed", msg.Message)
2018-09-01 01:27:26 +00:00
recordName := string(msg.Message)
2018-08-05 08:48:26 +00:00
rh.mu.Lock()
2018-09-01 01:27:26 +00:00
delete(rh.hosts, recordName)
2018-08-05 08:48:26 +00:00
rh.mu.Unlock()
} else if msg.Channel == keyspaceEvent {
2018-08-05 08:48:26 +00:00
logger.Debug("Refreshing redis records due to update")
rh.Refresh()
}
}
}()
return rh
}
func (r *RedisHosts) Get(domain string) ([]string, bool) {
2018-08-05 04:16:15 +00:00
logger.Debug("Checking redis provider for %s", domain)
r.mu.RLock()
defer r.mu.RUnlock()
domain = strings.ToLower(domain)
2018-09-01 01:27:26 +00:00
if ip, ok := r.hosts[domain]; ok {
return strings.Split(ip, ","), true
}
2018-09-01 01:27:26 +00:00
if idx := strings.Index(domain, "."); idx != -1 {
2019-09-26 04:43:17 +00:00
wildcard := "*." + domain[strings.Index(domain, ".")+1:]
2018-09-01 01:27:26 +00:00
if ip, ok := r.hosts[wildcard]; ok {
return strings.Split(ip, ","), true
}
}
for host, ip := range r.hosts {
2018-09-01 01:27:26 +00:00
if glob.Glob(host, domain) {
return strings.Split(ip, ","), true
}
}
2018-09-01 01:27:26 +00:00
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)
}