Unified --usb/--udp/--iface target flags for all CLI commands

cmd/picomap/main.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"strings"
 	"sync"
 	"time"
 
@@ -16,9 +17,103 @@ import (
 	"github.com/theater/picomap/lib/uf2"
 )
 
+type target struct {
+	name   string
+	client *client.Client
+}
+
+type targetFlags struct {
+	usb   string
+	udp   string
+	iface string
+}
+
+func addTargetFlags(fs *flag.FlagSet) *targetFlags {
+	tf := &targetFlags{}
+	fs.StringVar(&tf.usb, "usb", "", "comma-separated USB serial devices")
+	fs.StringVar(&tf.udp, "udp", "", "comma-separated UDP IP addresses")
+	fs.StringVar(&tf.iface, "iface", "", "network interface for multicast discovery")
+	return tf
+}
+
+func (tf *targetFlags) connect(timeout time.Duration) ([]target, error) {
+	var targets []target
+
+	if tf.usb != "" {
+		for _, dev := range strings.Split(tf.usb, ",") {
+			dev = strings.TrimSpace(dev)
+			c, err := client.NewSerial(dev, timeout)
+			if err != nil {
+				return nil, fmt.Errorf("usb %s: %w", dev, err)
+			}
+			targets = append(targets, target{dev, c})
+		}
+	}
+
+	if tf.udp != "" {
+		for _, addr := range strings.Split(tf.udp, ",") {
+			addr = strings.TrimSpace(addr)
+			c, err := client.NewUDP(addr, tf.iface, timeout)
+			if err != nil {
+				return nil, fmt.Errorf("udp %s: %w", addr, err)
+			}
+			targets = append(targets, target{addr, c})
+		}
+	}
+
+	if tf.iface != "" && tf.udp == "" {
+		bcast, err := client.InterfaceBroadcast(tf.iface)
+		if err != nil {
+			return nil, err
+		}
+		c, err := client.NewUDP(bcast, tf.iface, timeout)
+		if err != nil {
+			return nil, err
+		}
+		infos, err := c.InfoAll()
+		c.Close()
+		if err != nil {
+			return nil, fmt.Errorf("discovery: %w", err)
+		}
+		for _, r := range infos {
+			uc, err := client.NewUDP(r.From, tf.iface, timeout)
+			if err != nil {
+				return nil, fmt.Errorf("udp %s: %w", r.From, err)
+			}
+			targets = append(targets, target{r.From, uc})
+		}
+	}
+
+	if tf.usb == "" && tf.udp == "" && tf.iface == "" {
+		devs, err := client.ListSerial()
+		if err != nil {
+			return nil, err
+		}
+		for _, dev := range devs {
+			c, err := client.NewSerial(dev, timeout)
+			if err != nil {
+				slog.Warn("connect error", "dev", dev, "err", err)
+				continue
+			}
+			targets = append(targets, target{dev, c})
+		}
+	}
+
+	if len(targets) == 0 {
+		return nil, fmt.Errorf("no devices found")
+	}
+	return targets, nil
+}
+
+func closeTargets(targets []target) {
+	for _, t := range targets {
+		t.client.Close()
+	}
+}
+
 func main() {
 	if len(os.Args) < 2 {
-		slog.Error("usage: picomap <info|load|test> [args...]")
+		slog.Error("usage: picomap <info|load|log|test> [args...]")
 		os.Exit(1)
 	}
 	cmd := os.Args[1]
@@ -44,16 +139,9 @@ func main() {
 	}
 }
 
-type deviceResult struct {
-	dev  string
-	info *client.ResponseInfo
-	err  error
-}
-
-func printInfo(via string, from string, info *client.ResponseInfo) {
+func printInfo(via string, info *client.ResponseInfo) {
 	slog.Info("device",
 		"via", via,
-		"from", from,
 		"board_id", hex.EncodeToString(info.BoardID[:]),
 		"mac", net.HardwareAddr(info.MAC[:]).String(),
 		"ip", net.IP(info.IP[:]).String(),
@@ -62,74 +150,52 @@ func printInfo(via string, from string, info *client.ResponseInfo) {
 
 func cmdInfo(args []string) error {
 	fs := flag.NewFlagSet("info", flag.ExitOnError)
-	udpAddr := fs.String("udp", "", "connect via UDP to this IP address")
-	iface := fs.String("iface", "", "bind to this network interface (for broadcast)")
+	tf := addTargetFlags(fs)
 	fs.Parse(args)
 
-	if *udpAddr == "" && *iface != "" {
-		bcast, err := client.InterfaceBroadcast(*iface)
-		if err != nil {
-			return err
-		}
-		*udpAddr = bcast
-	}
-
-	if *udpAddr != "" {
-		c, err := client.NewUDP(*udpAddr, *iface, 500*time.Millisecond)
-		if err != nil {
-			return err
-		}
-		defer c.Close()
-		infos, err := c.InfoAll()
-		if err != nil {
-			return err
-		}
-		if len(infos) == 0 {
-			return fmt.Errorf("no devices responded")
-		}
-		for _, r := range infos {
-			printInfo(*udpAddr, r.From, r.Value)
-		}
-		return nil
-	}
-
-	devs, err := client.ListSerial()
+	targets, err := tf.connect(500*time.Millisecond)
 	if err != nil {
 		return err
 	}
-	if len(devs) == 0 {
-		return fmt.Errorf("no devices found")
-	}
+	defer closeTargets(targets)
 
-	results := make([]deviceResult, len(devs))
-	var wg sync.WaitGroup
-	for i, dev := range devs {
-		results[i].dev = dev
-		wg.Go(func() {
-			c, err := client.NewSerial(dev, 500*time.Millisecond)
-			if err != nil {
-				results[i].err = err
-				return
-			}
-			info, err := c.Info()
-			c.Close()
-			if err != nil {
-				results[i].err = err
-				return
-			}
-			results[i].info = info
-		})
-	}
-	wg.Wait()
-
-	for _, r := range results {
-		if r.err != nil {
-			slog.Error("device error", "dev", r.dev, "err", r.err)
+	for _, t := range targets {
+		info, err := t.client.Info()
+		if err != nil {
+			slog.Error("info error", "via", t.name, "err", err)
 			continue
 		}
-		printInfo(r.dev, r.dev, r.info)
+		printInfo(t.name, info)
 	}
+	return nil
+}
 
+func cmdLog(args []string) error {
+	fs := flag.NewFlagSet("log", flag.ExitOnError)
+	tf := addTargetFlags(fs)
+	fs.Parse(args)
+
+	targets, err := tf.connect(500*time.Millisecond)
+	if err != nil {
+		return err
+	}
+	defer closeTargets(targets)
+
+	for _, t := range targets {
+		log := slog.With("dev", t.name)
+		resp, err := t.client.Log()
+		if err != nil {
+			log.Error("log error", "err", err)
+			continue
+		}
+		if len(resp.Entries) == 0 {
+			log.Info("no debug messages")
+			continue
+		}
+		for _, e := range resp.Entries {
+			log.Info("dlog", "t_us", e.TimestampUS, "msg", e.Message)
+		}
+	}
 	return nil
 }
 
@@ -138,13 +204,6 @@ func boardSerial(id [8]byte) string {
 		id[0], id[1], id[2], id[3], id[4], id[5], id[6], id[7])
 }
 
-type deviceInfo struct {
-	dev    string
-	serial string
-	uf2    string
-	name   string
-}
-
 func buildFirmware(buildDir string) error {
 	slog.Info("configuring")
 	cmake := exec.Command("cmake", "-S", filepath.Join(filepath.Dir(buildDir)), "-B", buildDir)
@@ -164,45 +223,14 @@ func buildFirmware(buildDir string) error {
 	return nil
 }
 
-func cmdLog(_ []string) error {
-	devs, err := client.ListSerial()
-	if err != nil {
-		return err
-	}
-	if len(devs) == 0 {
-		return fmt.Errorf("no devices found")
-	}
-	for _, dev := range devs {
-		log := slog.With("dev", dev)
-		c, err := client.NewSerial(dev, 500*time.Millisecond)
-		if err != nil {
-			log.Error("connect error", "err", err)
-			continue
-		}
-		resp, err := c.Log()
-		c.Close()
-		if err != nil {
-			log.Error("log error", "err", err)
-			continue
-		}
-		if len(resp.Entries) == 0 {
-			log.Info("no debug messages")
-			continue
-		}
-		for _, e := range resp.Entries {
-			log.Info("dlog", "t_us", e.TimestampUS, "msg", e.Message)
-		}
-	}
-	return nil
-}
-
 func cmdLoad(args []string) error {
 	fs := flag.NewFlagSet("load", flag.ExitOnError)
+	tf := addTargetFlags(fs)
 	dryRun := fs.Bool("dry-run", false, "parse UF2 and log operations without flashing")
 	fs.Parse(args)
-	target := "all"
+	loadTarget := "all"
 	if fs.NArg() > 0 {
-		target = fs.Arg(0)
+		loadTarget = fs.Arg(0)
 	}
 
 	wd, err := os.Getwd()
@@ -215,78 +243,66 @@ func cmdLoad(args []string) error {
 		return err
 	}
 
-	allTargets := []struct {
+	type firmwareTarget struct {
 		name string
 		uf2  string
-	}{
+	}
+	allTargets := []firmwareTarget{
 		{"picomap", filepath.Join(buildDir, "picomap.uf2")},
 		{"picomap_test", filepath.Join(buildDir, "picomap_test.uf2")},
 	}
 
-	var targets []struct {
-		name string
-		uf2  string
-	}
-	switch target {
+	var fwTargets []firmwareTarget
+	switch loadTarget {
 	case "all":
-		targets = allTargets
+		fwTargets = allTargets
 	case "picomap":
-		targets = allTargets[:1]
+		fwTargets = allTargets[:1]
 	case "picomap_test":
-		targets = allTargets[1:]
+		fwTargets = allTargets[1:]
 	default:
-		return fmt.Errorf("unknown target %q", target)
+		return fmt.Errorf("unknown target %q", loadTarget)
 	}
 
-	devs, err := client.ListSerial()
+	targets, err := tf.connect(5*time.Second)
 	if err != nil {
 		return err
 	}
-	if len(devs) < len(targets) {
-		return fmt.Errorf("need %d device(s), found %d", len(targets), len(devs))
+	defer closeTargets(targets)
+
+	if len(targets) < len(fwTargets) {
+		return fmt.Errorf("need %d device(s), found %d", len(fwTargets), len(targets))
 	}
 
-	devices := make([]deviceInfo, len(targets))
-	errs := make([]error, len(targets))
-
-	var wg sync.WaitGroup
-	for i := range targets {
-		log := slog.With("dev", devs[i])
-		wg.Go(func() {
-			c, err := client.NewSerial(devs[i], 500*time.Millisecond)
-			if err != nil {
-				errs[i] = err
-				return
-			}
-			info, err := c.Info()
-			c.Close()
-			if err != nil {
-				errs[i] = err
-				return
-			}
-			devices[i] = deviceInfo{
-				dev:    devs[i],
-				serial: boardSerial(info.BoardID),
-				uf2:    targets[i].uf2,
-				name:   targets[i].name,
-			}
-			log.Info("got info", "serial", devices[i].serial, "firmware", info.FirmwareName)
-		})
+	type deviceInfo struct {
+		target target
+		serial string
+		fw     firmwareTarget
 	}
-	wg.Wait()
-	for i, err := range errs {
+
+	devices := make([]deviceInfo, len(fwTargets))
+	for i, t := range targets[:len(fwTargets)] {
+		info, err := t.client.Info()
 		if err != nil {
-			return fmt.Errorf("[%s] info: %w", devs[i], err)
+			return fmt.Errorf("[%s] info: %w", t.name, err)
 		}
+		devices[i] = deviceInfo{
+			target: t,
+			serial: boardSerial(info.BoardID),
+			fw:     fwTargets[i],
+		}
+		slog.Info("got info", "dev", t.name, "serial", devices[i].serial, "firmware", info.FirmwareName)
 	}
 
+	errs := make([]error, len(devices))
+	var wg sync.WaitGroup
 	for i := range devices {
 		log := slog.With("serial", devices[i].serial)
 		wg.Go(func() {
-			log.Info("flashing", "uf2", devices[i].name)
-			errs[i] = flashDevice(devices[i].dev, devices[i].uf2, *dryRun, log)
+			log.Info("flashing", "uf2", devices[i].fw.name)
+			errs[i] = flashDevice(devices[i].target.client, devices[i].fw.uf2, *dryRun, log)
 			if errs[i] == nil {
-				log.Info("flashed", "uf2", devices[i].name)
+				log.Info("flashed", "uf2", devices[i].fw.name)
 			}
 		})
 	}
@@ -301,7 +317,7 @@ func cmdLoad(args []string) error {
 	return nil
 }
 
-func flashDevice(dev, uf2Path string, dryRun bool, log *slog.Logger) error {
+func flashDevice(c *client.Client, uf2Path string, dryRun bool, log *slog.Logger) error {
 	blocks, err := uf2.Parse(uf2Path)
 	if err != nil {
 		return fmt.Errorf("parse uf2: %w", err)
@@ -324,12 +340,6 @@ func flashDevice(dev, uf2Path string, dryRun bool, log *slog.Logger) error {
 		return nil
 	}
 
-	c, err := client.NewSerial(dev, 5*time.Second)
-	if err != nil {
-		return err
-	}
-	defer c.Close()
-
 	erased := make(map[uint32]bool)
 	for _, b := range blocks {
 		sector := b.Addr &^ (sectorSize - 1)
@@ -349,29 +359,33 @@ func flashDevice(dev, uf2Path string, dryRun bool, log *slog.Logger) error {
 	return nil
 }
 
-func findTestDevice() (string, error) {
-	devs, err := client.ListSerial()
+func findTestDevice(args []string) (target, error) {
+	fs := flag.NewFlagSet("test", flag.ExitOnError)
+	tf := addTargetFlags(fs)
+	fs.Parse(args)
+
+	targets, err := tf.connect(10*time.Second)
 	if err != nil {
-		return "", err
+		return target{}, err
 	}
-	for _, dev := range devs {
-		log := slog.With("dev", dev)
-		c, err := client.NewSerial(dev, 500*time.Millisecond)
+
+	for _, t := range targets {
+		info, err := t.client.Info()
 		if err != nil {
-			log.Warn("connect error", "err", err)
-			continue
-		}
-		info, err := c.Info()
-		c.Close()
-		if err != nil {
-			log.Warn("info error", "err", err)
+			t.client.Close()
 			continue
 		}
 		if info.FirmwareName == "picomap_test" {
-			return dev, nil
+			for _, other := range targets {
+				if other.name != t.name {
+					other.client.Close()
+				}
+			}
+			return t, nil
 		}
+		t.client.Close()
 	}
-	return "", fmt.Errorf("no picomap_test device found")
+	return target{}, fmt.Errorf("no picomap_test device found")
 }
 
 func cmdTestGroup(args []string) error {
@@ -390,49 +404,40 @@ func cmdTestGroup(args []string) error {
 	}
 }
 
-func cmdTestList(_ []string) error {
-	dev, err := findTestDevice()
+func cmdTestList(args []string) error {
+	t, err := findTestDevice(args)
 	if err != nil {
 		return err
 	}
-	c, err := client.NewSerial(dev, 10*time.Second)
-	if err != nil {
-		return err
-	}
-	defer c.Close()
+	defer t.client.Close()
 
-	result, err := c.ListTests()
+	result, err := t.client.ListTests()
 	if err != nil {
 		return fmt.Errorf("remote: %w", err)
 	}
 	for _, name := range result.Names {
-		slog.Info("test", "dev", dev, "name", name)
+		slog.Info("test", "dev", t.name, "name", name)
 	}
 	return nil
 }
 
-func cmdTestAll(_ []string) error {
-	dev, err := findTestDevice()
+func cmdTestAll(args []string) error {
+	t, err := findTestDevice(args)
 	if err != nil {
 		return err
 	}
+	defer t.client.Close()
 
-	c, err := client.NewSerial(dev, 10*time.Second)
-	if err != nil {
-		return err
-	}
-	defer c.Close()
-
-	list, err := c.ListTests()
+	list, err := t.client.ListTests()
 	if err != nil {
 		return fmt.Errorf("remote: %w", err)
 	}
 
-	log := slog.With("dev", dev)
+	log := slog.With("dev", t.name)
 	failed := 0
 	for _, name := range list.Names {
 		log.Info("running test", "name", name)
-		result, err := c.Test(name)
+		result, err := t.client.Test(name)
 		if err != nil {
 			log.Error("error", "name", name, "err", err)
 			failed++
@@ -462,21 +467,16 @@ func cmdTestRun(args []string) error {
 	}
 	name := args[0]
 
-	dev, err := findTestDevice()
+	t, err := findTestDevice(args[1:])
 	if err != nil {
 		return err
 	}
+	defer t.client.Close()
 
-	log := slog.With("dev", dev)
+	log := slog.With("dev", t.name)
 	log.Info("running test", "name", name)
 
-	c, err := client.NewSerial(dev, 10*time.Second)
-	if err != nil {
-		return err
-	}
-	defer c.Close()
-
-	result, err := c.Test(name)
+	result, err := t.client.Test(name)
 	if err != nil {
 		return fmt.Errorf("remote: %w", err)
 	}