rio/internal/dns/dns.go

package dns

import (
	"net"
	"errors"
	"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="

	TxtCNAMERecord = "_rio-cname."

	TxtCNAMEKey = "cname="
)

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
	}

	cname, err := lookupCNAME(domain)
	if err == nil {
		cnameCache.Set(domain, cname, cache.DefaultExpiration)
		return cname.(string), nil
	}

	altCname, err := lookupCNAMETxt(domain)
	if err == nil {
		cnameCache.Set(domain, altCname, cache.DefaultExpiration)
		return altCname, nil
	}

	return "", err
}

// Lookup the CNAME by trying to find a CNAME RR. Contrary to net.LookupCNAME,
// this method fails if we find no CNAME RR.
func lookupCNAME(domain string) (string, error) {
	query, err := net.LookupCNAME(domain)
	if err == nil {
		if query[len(query)-1] == '.' {
			query = query[:len(query)-1]
		}

		// Fail if we have no CNAME RR.
		if query == domain {
			return "", errors.New("CNAME is equal to domain")
		}

		return query, nil
	}

	return "", err
}

// Performs an alternative CNAME lookup by looking for a special TXT record.
func lookupCNAMETxt(domain string) (string, error) {
	txts, err := net.LookupTXT(TxtCNAMERecord+domain)
	if err == nil {
		for _, txt := range txts {
			if !strings.HasPrefix(txt, TxtCNAMEKey) {
				continue
			}

			return strings.TrimPrefix(txt, TxtCNAMEKey), nil
		}
	}
	
	return "", err
}