prologic-saltyim/resolv.go

package saltyim

import (
	"context"
	"errors"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/likexian/doh-go"
	"github.com/likexian/doh-go/dns"
	log "github.com/sirupsen/logrus"
)

type Resolver interface {
	LookupSRV(service, proto, domain string) (string, error)
}

var (
	_ Resolver = (*StandardResolver)(nil)

	ErrSRVRecordNotFound = errors.New("error: No SRV records found")

	resolver Resolver
)

func init() {
	SetResolver(&StandardResolver{})
}

func SetResolver(r Resolver) {
	resolver = r
}

type StandardResolver struct{}

func (r *StandardResolver) LookupSRV(service, proto, domain string) (string, error) {
	log.Debugf("Using StandardResolver, looking up SRV _%s._%s.%s", service, proto, domain)
	_, records, err := net.LookupSRV(service, proto, domain)
	if err != nil {
		return "", fmt.Errorf("error looking up _%s._%s.%s : %w", service, proto, domain, err)
	}
	if len(records) == 0 {
		return "", ErrSRVRecordNotFound
	}
	return strings.TrimSuffix(records[0].Target, "."), nil
}

type DNSOverHTTPResolver struct{}

func (r *DNSOverHTTPResolver) LookupSRV(service, proto, domain string) (string, error) {
	name := fmt.Sprintf("_%s._%s.%s", service, proto, domain)

	log.Debugf("Using DNSOverHTTPResolver, looking up SRV %s", name)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	c := doh.Use(doh.CloudflareProvider, doh.DNSPodProvider)

	res, err := c.Query(ctx, dns.Domain(name), dns.Type("SRV"))
	if err != nil {
		return "", fmt.Errorf("error looking up _%s._%s.%s : %w", service, proto, domain, err)
	}
	defer c.Close()

	if len(res.Answer) == 0 {
		return "", ErrSRVRecordNotFound
	}

	data := res.Answer[0].Data
	fields := strings.Split(data, " ")
	if len(fields) != 4 {
		return "", fmt.Errorf("invalid SRV records found expected 4 fields got %d: %q", len(fields), data)
	}
	return strings.TrimSuffix(fields[3], "."), nil
}