go/src/infra/cros/recovery/internal/components/cros/androidtools/adb_tools.go

// Copyright 2026 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package androidtools

import (
	"context"
	"strings"
	"time"

	"go.chromium.org/chromiumos/config/go/test/api"
	"go.chromium.org/luci/common/errors"

	"go.chromium.org/infra/cros/recovery/internal/components/cft"
	"go.chromium.org/infra/cros/recovery/internal/env"
	"go.chromium.org/infra/cros/recovery/internal/log"
	"go.chromium.org/infra/cros/recovery/tlw"
)

// ADBRun executes a generic adb command.
//
// This function can run any adb command. For specific common commands like shell, root,
// connect, disconnect, more specific wrapper functions are available (e.g., ADBShell, ADBRoot).
// It abstracts the client creation and error handling.
//
// Args:
//
//	ctx: Context for the operation.
//	dut: The device under test.
//	timeout: The maximum time to allow for the command execution.
//	command: The adb command to execute (e.g., "devices", "reboot").
//	args: Additional arguments for the adb command.
//
// Returns:
//
//	Response: The result of the command execution.
//	error: An error if the command failed to execute or the exit code was non-zero.
func ADBRun(ctx context.Context, dut *tlw.Dut, timeout time.Duration, command string, args ...string) (Response, error) {
	client, err := adbClient(ctx, dut)
	if err != nil {
		return emptyResponse(), errors.WrapIf(err, "adb run")
	}
	res, err := adbRun(ctx, client, timeout, command, args...)
	if err != nil {
		return res, errors.WrapIf(err, "adb run")
	}
	if res.GetExitCode() != 0 {
		return res, errors.Fmt("adb run: failed with exitcode: %d", res.GetExitCode())
	}
	return res, nil
}

// ADBRoot attempts to restart the adb daemon with root permissions on the device.
//
// Equivalent to running "adb -s <serial> root".
//
// Args:
//
//	ctx: Context for the operation.
//	dut: The device under test.
//	timeout: The maximum time to allow for the command execution.
//
// Returns:
//
//	Response: The result of the command execution.
//	error: An error if the command failed to execute or the exit code was non-zero.
func ADBRoot(ctx context.Context, dut *tlw.Dut, timeout time.Duration) (Response, error) {
	client, err := adbClient(ctx, dut)
	if err != nil {
		return emptyResponse(), errors.WrapIf(err, "adb root")
	}
	args := []string{ADBSerial(dut), rootCommand}
	res, err := adbRun(ctx, client, timeout, serialParam, args...)
	if err != nil {
		return res, errors.WrapIf(err, "adb root")
	}
	if res.GetExitCode() != 0 {
		return res, errors.Fmt("adb root: failed with exitcode: %d", res.GetExitCode())
	}
	return res, nil
}

// ADBConnect connects to the device, usually over TCP/IP.
//
// Equivalent to running "adb connect <serial>".
//
// Args:
//
//	ctx: Context for the operation.
//	dut: The device under test.
//	timeout: The maximum time to allow for the command execution.
//
// Returns:
//
//	Response: The result of the command execution.
//	error: An error if the command failed to execute or the exit code was non-zero.
func ADBConnect(ctx context.Context, dut *tlw.Dut, timeout time.Duration) (Response, error) {
	client, err := adbClient(ctx, dut)
	if err != nil {
		return emptyResponse(), errors.WrapIf(err, "adb connect")
	}
	args := []string{ADBSerial(dut)}
	res, err := adbRun(ctx, client, timeout, connectCommand, args...)
	if err != nil {
		return res, errors.WrapIf(err, "adb connect")
	}
	if res.GetExitCode() != 0 {
		return res, errors.Fmt("adb connect: failed with exitcode: %d", res.GetExitCode())
	}
	return res, nil
}

// ADBDisconnect disconnects from the device, usually over TCP/IP.
//
// Equivalent to running "adb disconnect <serial>".
//
// Args:
//
//	ctx: Context for the operation.
//	dut: The device under test.
//	timeout: The maximum time to allow for the command execution.
//
// Returns:
//
//	Response: The result of the command execution.
//	error: An error if the command failed to execute or the exit code was non-zero.
func ADBDisconnect(ctx context.Context, dut *tlw.Dut, timeout time.Duration) (Response, error) {
	client, err := adbClient(ctx, dut)
	if err != nil {
		return emptyResponse(), errors.WrapIf(err, "adb disconnect")
	}
	args := []string{ADBSerial(dut)}
	res, err := adbRun(ctx, client, timeout, disconnectCommand, args...)
	if err != nil {
		return res, errors.WrapIf(err, "adb disconnect")
	}
	if res.GetExitCode() != 0 {
		return res, errors.Fmt("adb disconnect: failed with exitcode: %d", res.GetExitCode())
	}
	return res, nil
}

// ADBIsConnected checks if the device is currently listed in the output of "adb devices".
//
// Args:
//
//	ctx: Context for the operation.
//	dut: The device under test.
//	timeout: The maximum time to allow for the "adb devices" command.
//
// Returns:
//
//	bool: True if the device is listed, false otherwise.
func ADBIsConnected(ctx context.Context, dut *tlw.Dut, timeout time.Duration) bool {
	res, err := adbDevices(ctx, dut, timeout)
	if err != nil {
		log.Debugf(ctx, "Failed to read adb devices, assume device isn't listed yet: %s", err)
		return false
	}
	deviceName := ADBSerial(dut)
	if out := string(res.GetStdout()); out != "" && strings.Contains(out, deviceName) {
		log.Debugf(ctx, "Device is listed.")
		return true
	}
	log.Debugf(ctx, "Device isn't listed yet: %s", err)
	return false
}

// adbDevices runs the "adb devices" command.
func adbDevices(ctx context.Context, dut *tlw.Dut, timeout time.Duration) (Response, error) {
	client, err := adbClient(ctx, dut)
	if err != nil {
		return emptyResponse(), errors.WrapIf(err, "adb devices")
	}
	res, err := adbRun(ctx, client, timeout, devicesCommand)
	if err != nil {
		return res, errors.WrapIf(err, "adb devices")
	}
	if res.GetExitCode() != 0 {
		return res, errors.Fmt("adb devices: failed with exitcode: %d", res.GetExitCode())
	}
	return res, nil
}

// ADBShell executes a shell command on the device via adb.
//
// Equivalent to running "adb -s <serial> shell <command> <args...>".
//
// Args:
//
//	ctx: Context for the operation.
//	dut: The device under test.
//	timeout: The maximum time to allow for the command execution.
//	command: The shell command to execute on the device.
//	args: Additional arguments for the shell command.
//
// Returns:
//
//	Response: The result of the command execution.
//	error: An error if the command failed to execute or the exit code was non-zero.
func ADBShell(ctx context.Context, dut *tlw.Dut, timeout time.Duration, command string, args ...string) (Response, error) {
	client, err := adbClient(ctx, dut)
	if err != nil {
		return emptyResponse(), errors.WrapIf(err, "adb shell")
	}
	// Example adb -s device1:5555 shell xyz
	// Where `-s` is the command and rest will be args.
	newArgs := []string{ADBSerial(dut), shellCommand, command}
	newArgs = append(newArgs, args...)
	res, err := adbRun(ctx, client, timeout, serialParam, newArgs...)
	if err != nil {
		return res, errors.WrapIf(err, "adb shell")
	}
	if res.GetExitCode() != 0 {
		return res, errors.Fmt("exec adb command: failed with exitcode: %d", res.GetExitCode())
	}
	return res, nil
}

// adbRun runs raw command by base-adb service and return result.
//
// Error is provided only if the service fails to execute command or timeout.
// Command execution is always represented as exec-code in response.
func adbRun(ctx context.Context, client api.ToolsServiceClient, timeout time.Duration, command string, args ...string) (Response, error) {
	if command == "" {
		return nil, errors.Fmt("adb run: command is empty")
	}
	fullCmd := "adb " + command
	if len(args) > 0 {
		fullCmd += " " + strings.Join(args, " ")
	}
	log.Debugf(ctx, "Prepare to run adb command: %q", fullCmd)
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	var err error
	var response Response
	if client != nil {
		response, err = client.ADBRun(ctx, &api.ADBRunRequest{
			Command: command,
			Params:  args,
		})
	} else {
		response, err = localRun(ctx, env.ADBPath(), command, args...)
	}
	if response != nil {
		printResponse(ctx, response)
	}
	if err != nil {
		err = errors.Fmt("adb run: failed execute command %q, finished with error: %w", fullCmd, err)
	}
	return response, errors.WrapIf(err, "adb run: command %q", fullCmd)
}

func adbClient(ctx context.Context, dut *tlw.Dut) (api.ToolsServiceClient, error) {
	if UsesContainer(ctx) {
		return cft.ToolClientFromScope(ctx, dut)
	}
	// For local test we do not provide a client.
	return nil, nil
}

// ADBDeviceState returns the state of the device.
//
// Returns one of: "device", "offline", "bootloader", "unauthorized", "unknown".
func ADBDeviceState(ctx context.Context, dut *tlw.Dut, timeout time.Duration) (string, error) {
	client, err := adbClient(ctx, dut)
	if err != nil {
		return "", errors.WrapIf(err, "adb get-state")
	}
	args := []string{ADBSerial(dut), "get-state"}
	res, err := adbRun(ctx, client, timeout, serialParam, args...)
	if err != nil {
		return "", errors.WrapIf(err, "adb get-state")
	}
	return parseGetStateOutput(res.GetStdout(), res.GetExitCode(), res.GetStderr())
}

func parseGetStateOutput(stdout []byte, exitCode int32, stderr []byte) (string, error) {
	if exitCode != 0 {
		stderrStr := string(stderr)
		if strings.Contains(stderrStr, "not found") {
			return "unknown", nil
		}
		return "", errors.Fmt("adb get-state: failed with exitcode: %d, stderr: %s", exitCode, stderrStr)
	}
	return strings.TrimSpace(string(stdout)), nil
}

// ADBDeviceVisible checks if the device is visible via ADB.
// It returns nil if the device is visible (state is not offline or unknown),
// and an error otherwise.
func ADBDeviceVisible(ctx context.Context, dut *tlw.Dut, timeout time.Duration) error {
	state, err := ADBDeviceState(ctx, dut, timeout)
	if err != nil {
		return errors.Fmt("adb device visible: %w", err)
	}
	log.Debugf(ctx, "ADB device state: %q", state)
	if state == "offline" || state == "unknown" {
		return errors.Reason("adb device visible: device is in %q state", state)
	}
	return nil
}