godns/hosts_redis.go

package main

import (
	"github.com/hoisie/redis"
	"sync"
	"strings"
		"github.com/ryanuber/go-glob"
)

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

	// Force an initial refresh
	rh.Refresh()

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

		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 {
			msg := <- messages

			if msg.Channel == "godns:update" {
				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 {
					logger.Warn("Record %s does not exist, but was updated", recordName)
					continue
				}

				logger.Debug("Record %s was updated to %s", recordName, string(b))

				rh.mu.Lock()
				rh.hosts[recordName] = string(b)
				rh.mu.Unlock()
			} else if msg.Channel == "godns:remove_record" {
				logger.Debug("Record %s was removed", msg.Message)

				recordName := string(msg.Message)

				rh.mu.Lock()
				delete(rh.hosts, recordName)
				rh.mu.Unlock()
			} else if msg.Channel == keyspaceEvent {
				logger.Debug("Refreshing redis records due to update")
				rh.Refresh()
			}
		}
	}()

	return rh
}

func (r *RedisHosts) Get(domain string) ([]string, bool) {
	logger.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, ","), 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, ","), true
		}
	}

	for host, ip := range r.hosts {
		if glob.Glob(host, domain) {
			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)
}