prober: add helper function to check all IPs for a DNS hostname

This allows us to check all IP addresses (and address families) for a
given DNS hostname while dynamically discovering new IPs and removing
old ones as they're no longer valid.

Also add a testable example that demonstrates how to use it.

Alternative to #11610
Updates tailscale/corp#16367

Signed-off-by: Andrew Dunham <andrew@du.nham.ca>
Change-Id: I6d6f39bafc30e6dfcf6708185d09faee2a374599

prober/dns_example_test.go

@@ -0,0 +1,98 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package prober_test
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"log"
+	"net"
+	"net/netip"
+	"os"
+	"os/signal"
+	"time"
+
+	"tailscale.com/prober"
+	"tailscale.com/types/logger"
+)
+
+const (
+	every30s = 30 * time.Second
+)
+
+var (
+	hostname = flag.String("hostname", "tailscale.com", "hostname to probe")
+	oneshot  = flag.Bool("oneshot", true, "run probes once and exit")
+	verbose  = flag.Bool("verbose", false, "enable verbose logging")
+)
+
+// This example demonstrates how to use ForEachAddr to create a TLS probe for
+// each IP address in the DNS record of a given hostname.
+func ExampleForEachAddr() {
+	flag.Parse()
+
+	p := prober.New().WithSpread(true)
+	if *oneshot {
+		p = p.WithOnce(true)
+	}
+
+	// This function is called every time we discover a new IP address to check.
+	makeTLSProbe := func(addr netip.Addr) *prober.Probe {
+		pf := prober.TLSWithIP(*hostname, netip.AddrPortFrom(addr, 443))
+		if *verbose {
+			logger := logger.WithPrefix(log.Printf, fmt.Sprintf("[tls %s]: ", addr))
+			pf = probeLogWrapper(logger, pf)
+		}
+
+		return p.Run(fmt.Sprintf("website/%s/tls", addr), every30s, nil, pf)
+	}
+
+	// Determine whether to use IPv4 or IPv6 based on whether we can create
+	// an IPv6 listening socket on localhost.
+	sock, err := net.Listen("tcp", "[::1]:0")
+	supportsIPv6 := err == nil
+	if sock != nil {
+		sock.Close()
+	}
+
+	networks := []string{"ip4"}
+	if supportsIPv6 {
+		networks = append(networks, "ip6")
+	}
+
+	var vlogf logger.Logf = logger.Discard
+	if *verbose {
+		vlogf = log.Printf
+	}
+
+	// This is the outer probe that resolves the hostname and creates a new
+	// TLS probe for each IP.
+	p.Run("website/dns", every30s, nil, prober.ForEachAddr(*hostname, makeTLSProbe, prober.ForEachAddrOpts{
+		Logf:     vlogf,
+		Networks: networks,
+	}))
+
+	defer log.Printf("done")
+
+	// Wait until all probes have run if we're running in oneshot mode.
+	if *oneshot {
+		p.Wait()
+		return
+	}
+
+	// Otherwise, wait until we get a signal.
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, os.Interrupt)
+	<-sigCh
+}
+
+func probeLogWrapper(logf logger.Logf, pf prober.ProbeFunc) prober.ProbeFunc {
+	return func(ctx context.Context) error {
+		logf("starting probe")
+		err := pf(ctx)
+		logf("probe finished with %v", err)
+		return err
+	}
+}