godns/hosts/hosts_redis.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)
}