feat: discover from DNS SRV records (#138)

* feat: discover from DNS SRV records

Signed-off-by: Carlos A Becker <caarlos0@users.noreply.github.com>

* chore: lint issue

Signed-off-by: Carlos A Becker <caarlos0@users.noreply.github.com>

* feat: txt records

Signed-off-by: Carlos A Becker <caarlos0@users.noreply.github.com>

Signed-off-by: Carlos A Becker <caarlos0@users.noreply.github.com>

README.md

@@ -200,6 +200,23 @@ Otherwise, the services will simply be added to the list.
 
 Run `wishlist --help` to see all the options.
 
+### SRV records
+
+Wishlist can also find nodes from DNS SRV records, on one or more domains.
+
+Run `wishlist --srv.domain {your domain}` to get started. You can repeat the
+flag for multiple domains.
+
+By default, Wishlist will set the name of the endpoint to the SRV target.
+You can, however, customize that with a TXT record in the following format:
+
+```
+wishlist.name full.address:22=thename
+```
+
+So, in this case, a SRV record pointing to `full.address` on port `22` will get
+the name `thename`.
+
 ### Using the binary
 
 ```sh

cmd/wishlist/main.go

@@ -17,6 +17,7 @@ import (
 	"github.com/charmbracelet/wish/activeterm"
 	lm "github.com/charmbracelet/wish/logging"
 	"github.com/charmbracelet/wishlist"
+	"github.com/charmbracelet/wishlist/srv"
 	"github.com/charmbracelet/wishlist/sshconfig"
 	"github.com/charmbracelet/wishlist/zeroconf"
 	"github.com/hashicorp/go-multierror"
@@ -123,6 +124,7 @@ var serverCmd = &cobra.Command{
 
 var (
 	configFile      string
+	srvDomains      []string
 	zeroconfEnabled bool
 	zeroconfDomain  string
 	zeroconfTimeout time.Duration
@@ -134,6 +136,7 @@ func init() {
 	rootCmd.PersistentFlags().BoolVar(&zeroconfEnabled, "zeroconf.enabled", false, "Whether to enable zeroconf service discovery (Avahi/Bonjour/mDNS)")
 	rootCmd.PersistentFlags().StringVar(&zeroconfDomain, "zeroconf.domain", "", "Domain to use with zeroconf service discovery")
 	rootCmd.PersistentFlags().DurationVar(&zeroconfTimeout, "zeroconf.timeout", time.Second, "How long should zeroconf keep searching for hosts")
+	rootCmd.PersistentFlags().StringSliceVar(&srvDomains, "srv.domain", nil, "SRV domains to discover endpoints")
 	rootCmd.AddCommand(serverCmd, manCmd)
 }
 
@@ -180,6 +183,13 @@ func getConfig(configFile string) (wishlist.Config, error) {
 		}
 		seed = endpoints
 	}
+	for _, domain := range srvDomains {
+		endpoints, err := srv.Endpoints(domain)
+		if err != nil {
+			return wishlist.Config{}, err //nolint: wrapcheck
+		}
+		seed = append(seed, endpoints...)
+	}
 	for _, path := range append([]string{configFile}, userConfigPaths()...) {
 		if path == "" {
 			continue

srv/srv.go

@@ -0,0 +1,54 @@
+package srv
+
+import (
+	"fmt"
+	"log"
+	"net"
+	"strings"
+
+	"github.com/charmbracelet/wishlist"
+)
+
+// Endpoints returns the _ssh._tcp SRV records on the given domain as
+// Wishlist endpoints.
+func Endpoints(domain string) ([]*wishlist.Endpoint, error) {
+	log.Printf("discovering _ssh._tcp SRV records on domain %q...", domain)
+	_, srvs, err := net.LookupSRV("ssh", "tcp", domain)
+	if err != nil {
+		return nil, fmt.Errorf("srv: could not resolve %s: %w", domain, err)
+	}
+	txts, err := net.LookupTXT(domain)
+	if err != nil {
+		return nil, fmt.Errorf("srv: could not resolve %s: %w", domain, err)
+	}
+	return fromRecords(srvs, txts), nil
+}
+
+const txtPrefix = "wishlist.name "
+
+func fromRecords(srvs []*net.SRV, txts []string) []*wishlist.Endpoint {
+	result := make([]*wishlist.Endpoint, 0, len(srvs))
+	for _, entry := range srvs {
+		hostname := strings.TrimSuffix(entry.Target, ".")
+		name := hostname
+		port := fmt.Sprintf("%d", entry.Port)
+		address := net.JoinHostPort(hostname, port)
+		for _, txt := range txts {
+			if !strings.HasPrefix(txt, txtPrefix) {
+				continue
+			}
+			txtAddr, txtName, ok := strings.Cut(strings.TrimPrefix(txt, txtPrefix), "=")
+			if !ok {
+				continue
+			}
+			if txtAddr == address {
+				name = txtName
+			}
+		}
+		result = append(result, &wishlist.Endpoint{
+			Name:    name,
+			Address: address,
+		})
+	}
+	return result
+}

srv/srv_test.go

@@ -0,0 +1,66 @@
+package srv
+
+import (
+	"net"
+	"testing"
+
+	"github.com/charmbracelet/wishlist"
+	"github.com/stretchr/testify/require"
+)
+
+func TestFromRecords(t *testing.T) {
+	t.Run("no txt records", func(t *testing.T) {
+		require.ElementsMatch(t, []*wishlist.Endpoint{
+			{
+				Name:    "foo.local",
+				Address: "foo.local:2222",
+			},
+			{
+				Name:    "foo.bar",
+				Address: "foo.bar:22",
+			},
+		}, fromRecords([]*net.SRV{
+			{
+				Target:   "foo.bar",
+				Port:     22,
+				Priority: 10,
+				Weight:   10,
+			},
+			{
+				Target:   "foo.local",
+				Port:     2222,
+				Priority: 10,
+				Weight:   10,
+			},
+		}, nil))
+	})
+
+	t.Run("with txt records", func(t *testing.T) {
+		require.ElementsMatch(t, []*wishlist.Endpoint{
+			{
+				Name:    "local-foo",
+				Address: "foo.local:2222",
+			},
+			{
+				Name:    "foobar",
+				Address: "foo.bar:22",
+			},
+		}, fromRecords([]*net.SRV{
+			{
+				Target:   "foo.bar",
+				Port:     22,
+				Priority: 10,
+				Weight:   10,
+			},
+			{
+				Target:   "foo.local",
+				Port:     2222,
+				Priority: 10,
+				Weight:   10,
+			},
+		}, []string{
+			"wishlist.name foo.bar:22=foobar",
+			"wishlist.name foo.local:2222=local-foo",
+		}))
+	})
+}

zeroconf/zeroconf.go

@@ -17,7 +17,7 @@ const service = "_ssh._tcp"
 
 // Endpoints returns the found endpoints from zeroconf.
 func Endpoints(domain string, timeout time.Duration) ([]*wishlist.Endpoint, error) {
-	log.Printf("getting %s from zeroconf on domain %q...", service, domain)
+	log.Printf("discovering %s from zeroconf on domain %q...", service, domain)
 	r, err := zeroconf.NewResolver()
 	if err != nil {
 		return nil, fmt.Errorf("zeroconf: could not create resolver: %w", err)