godns/hosts_redis.go

package main

import (
	"github.com/hoisie/redis"
	"sync"
	"strings"
	"golang.org/x/net/publicsuffix"
)

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),
	}

	// Use pubsub to listen for key update events
	go func() {
		keyspaceEvent := "__keyspace@0__:" + key

		sub := make(chan string, 2)
		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" {
				rh.Refresh()
			} else if msg.Channel == "godns:update_record" {
				recordName := string(msg.Message)

				b, err := rc.Hget(key, recordName)

				if err != nil {
					continue
				}

				rh.hosts[recordName] = string(b)
			} else if msg.Channel == keyspaceEvent {
				rh.Refresh()
			}
		}
	}()

	return rh
}

func (r *RedisHosts) Get(domain string) ([]string, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()

	domain = strings.ToLower(domain)
	ip, ok := r.hosts[domain]
	if ok {
		return strings.Split(ip, ","), true
	}

	sld, err := publicsuffix.EffectiveTLDPlusOne(domain)
	if err != nil {
		return nil, false
	}

	for host, ip := range r.hosts {
		if strings.HasPrefix(host, "*.") {
			old, err := publicsuffix.EffectiveTLDPlusOne(host)
			if err != nil {
				continue
			}
			if sld == old {
				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)
}