rio/internal/dns/dns.go

package dns

import (
	"net"
	"strings"
	"time"

	"github.com/patrickmn/go-cache"
)

const (
	// TXT record name that lookupRepoTXT will try to lookup.
	TxtRepoRecord = "_rio-pages."

	// The key that the TXT record will have to start with, e.g.
	// "repo=some-random-repo".
	TxtRepoKey = "repo="
)

var (
	// Cache for CNAME resolution results.
	cnameCache = cache.New(1*time.Hour, 1*time.Hour)

	// Cache for TXT resolution results.
	txtRepoCache = cache.New(1*time.Hour, 1*time.Hour)
)

// Query the domain for the a repository redirect.
// Returns the new repository name or "", if we could not
// resolve a repository redirect.
func LookupRepoTXT(domain string) (string, error) {
	repoLookup, found := txtRepoCache.Get(domain)
	if found {
		return repoLookup.(string), nil
	}

	txts, err := net.LookupTXT("_rio-pages." + domain)
	if err != nil {
		return "", err
	}

	repo := ""
	for _, txt := range txts {
		if !strings.HasPrefix(txt, TxtRepoKey) {
			continue
		}

		repo = strings.TrimPrefix(txt, TxtRepoKey)
		break
	}

	txtRepoCache.Set(domain, repo, cache.DefaultExpiration)
	return repo, nil
}

// Query the domain for a CNAME record. Returns the resolved
// CNAME or "", if no CNAME could be queried.
func LookupCNAME(domain string) (string, error) {
	cname, found := cnameCache.Get(domain)
	if found {
		return cname.(string), nil
	}

	query, err := net.LookupCNAME(domain)
	if err == nil {
		if query[len(query)-1] == '.' {
			query = query[:len(query)-1]
		}
		cnameCache.Set(domain, query, cache.DefaultExpiration)
		return query, nil
	}

	return "", err
}