godns/hosts/hosts_redis.go

88 lines
1.6 KiB
Go

package hosts
import (
"github.com/go-redis/redis/v7"
"github.com/ryanuber/go-glob"
"meow.tf/joker/godns/log"
"strings"
"sync"
"time"
)
type RedisHosts struct {
Provider
redis *redis.Client
key string
hosts map[string]string
mu sync.RWMutex
ttl time.Duration
}
func NewRedisProvider(rc *redis.Client, key string, ttl time.Duration) Provider {
rh := &RedisHosts{
redis: rc,
key: key,
hosts: make(map[string]string),
ttl: ttl,
}
// Force an initial refresh
rh.Refresh()
return rh
}
func (r *RedisHosts) Get(domain string) ([]string, time.Duration, bool) {
log.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, ","), r.ttl, 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, ","), r.ttl, true
}
}
for host, ip := range r.hosts {
if glob.Glob(host, domain) {
return strings.Split(ip, ","), r.ttl, true
}
}
return nil, time.Duration(0), 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)).Result()
}
func (r *RedisHosts) Refresh() {
r.mu.Lock()
defer r.mu.Unlock()
r.clear()
var err error
r.hosts, err = r.redis.HGetAll(r.key).Result()
if err != nil {
log.Warn("Update hosts records from redis failed %s", err)
} else {
log.Debug("Update hosts records from redis")
}
}
func (r *RedisHosts) clear() {
r.hosts = make(map[string]string)
}