94 lines
1.8 KiB
Go
94 lines
1.8 KiB
Go
package hosts
|
|
|
|
import (
|
|
"github.com/go-redis/redis/v7"
|
|
"github.com/miekg/dns"
|
|
"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(queryType uint16, domain string) ([]string, time.Duration, bool) {
|
|
log.Debug("Checking redis provider for %s", domain)
|
|
|
|
// Don't support queries other than A/AAAA for now
|
|
if queryType != dns.TypeA || queryType != dns.TypeAAAA {
|
|
return nil, zeroDuration, false
|
|
}
|
|
|
|
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(t, 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)
|
|
}
|