artemis_control/internal/adb/discover.go

package adb

import (
	"fmt"
	"sort"
	"strings"
	"sync"
	"time"

	"net"
)

// The wireless-debugging port changes on every phone reboot or toggle of
// wireless debugging, so it is never hardcoded. Android opens two TCP ports:
// a pairing port (useless here) and the connect port. `adb connect` succeeds on
// the connect port and is refused on the pairing port.

// DefaultScanRange is the range Android picks wireless-debugging ports from.
const (
	DefaultScanFrom = 30000
	DefaultScanTo   = 65535
)

// Connect asks the adb server to connect to host:port.
func (r *Runner) Connect(address string) (string, error) {
	out, err := r.run(15*time.Second, "connect", address)
	msg := strings.TrimSpace(string(out))
	if err != nil {
		return msg, err
	}
	// adb connect can exit 0 while reporting failure in its output.
	lower := strings.ToLower(msg)
	if strings.Contains(lower, "failed") || strings.Contains(lower, "cannot") {
		return msg, fmt.Errorf("adb connect %s: %s", address, msg)
	}
	return msg, nil
}

// DeviceOnline reports whether serial is present and in state "device".
func (r *Runner) DeviceOnline(serial string) bool {
	devices, err := r.Devices()
	if err != nil {
		return false
	}
	for _, d := range devices {
		if d.Serial == serial && d.State == "device" {
			return true
		}
	}
	return false
}

// ScanPorts does a concurrent TCP connect scan of host:from-to and returns the
// open ports in ascending order.
func ScanPorts(host string, from, to int, timeout time.Duration) []int {
	if timeout <= 0 {
		timeout = 500 * time.Millisecond
	}
	const workers = 512

	type result struct {
		port int
		open bool
	}

	jobs := make(chan int)
	results := make(chan result)
	var wg sync.WaitGroup

	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for port := range jobs {
				addr := net.JoinHostPort(host, fmt.Sprintf("%d", port))
				conn, err := net.DialTimeout("tcp", addr, timeout)
				if err != nil {
					results <- result{port: port, open: false}
					continue
				}
				conn.Close()
				results <- result{port: port, open: true}
			}
		}()
	}

	go func() {
		for p := from; p <= to; p++ {
			jobs <- p
		}
		close(jobs)
	}()

	go func() {
		wg.Wait()
		close(results)
	}()

	var open []int
	for res := range results {
		if res.open {
			open = append(open, res.port)
		}
	}
	sort.Ints(open)
	return open
}

// ConnectAuto brings the phone online. It reuses an already-connected device if
// one is present, otherwise it scans host for open ports and tries
// `adb connect` on each one until a port comes up in state "device".
//
// It returns the serial it connected to.
func (r *Runner) ConnectAuto(host string) (string, error) {
	// Already connected? Find a device whose serial is host:port.
	for _, d := range mustDevices(r) {
		if d.State == "device" && strings.HasPrefix(d.Serial, host+":") {
			r.Serial = d.Serial
			return d.Serial, nil
		}
	}

	ports := ScanPorts(host, DefaultScanFrom, DefaultScanTo, 0)
	if len(ports) == 0 {
		return "", fmt.Errorf("no open TCP ports on %s in %d-%d; is wireless debugging on?",
			host, DefaultScanFrom, DefaultScanTo)
	}

	var errs []string
	for _, port := range ports {
		serial := fmt.Sprintf("%s:%d", host, port)
		if _, err := r.Connect(serial); err != nil {
			errs = append(errs, fmt.Sprintf("%d: %v", port, err))
			continue
		}
		// Give the adb server a moment to move the device to "device".
		deadline := time.Now().Add(5 * time.Second)
		for time.Now().Before(deadline) {
			if r.DeviceOnline(serial) {
				r.Serial = serial
				return serial, nil
			}
			time.Sleep(250 * time.Millisecond)
		}
		errs = append(errs, fmt.Sprintf("%d: connect accepted but device never came online", port))
	}
	return "", fmt.Errorf("open ports on %s were %v but none accepted adb connect: %s",
		host, ports, strings.Join(errs, "; "))
}

func mustDevices(r *Runner) []Device {
	devices, err := r.Devices()
	if err != nil {
		return nil
	}
	return devices
}