go/src/infra/cros/cmd/provision/common-utils/cross-over/full_os_image_state.go - infra/infra.git - Git at Google

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

 // DLC constants and helpers
 package cross_over

 import (
 	"context"
 	"errors"
 	"fmt"
 	"log"
 	"regexp"
 	"strings"
 	"time"

 	"golang.org/x/crypto/ssh"
 	"google.golang.org/protobuf/types/known/anypb"

 	"go.chromium.org/chromiumos/config/go/test/api"

 	"go.chromium.org/infra/cros/cmd/cft/common/adb"
 	commonutils "go.chromium.org/infra/cros/cmd/provision/common-utils"
 )

 const (
 	// sshMsgIgnore is the SSH global message sent to ping the host.
 	// See RFC 4253 11.2, "Ignored Data Message".
 	sshMsgIgnore              = "SSH_MSG_IGNORE"
 	pingTimeout               = time.Second
 	pingRetryDelay            = 2 * time.Second
 	pingRetryCount            = 20
 	networkWaitTime           = 1 * time.Minute
 	powerResetDelay           = 10 * time.Second
 	installTimeout            = 12 * time.Minute
 	networkErrMsg             = "Network is unreachable"
 	cpuUartCapture            = "cpu_uart_capture"
 	engImg                    = "-eng"
 	userImg                   = "-user"
 	androidBuild              = "android-build"
 	checkConnectionRetryCount = 3
 )

 var launchTargetPatterns = []*regexp.Regexp{
 	regexp.MustCompile(`build_details/P?[0-9]+/([a-z_\-]+)`),
 	regexp.MustCompile(`artifacts_list/P?[0-9]+/([a-z_\-]+)`)}

 var buildIdPatterns = []*regexp.Regexp{
 	regexp.MustCompile(`build_details/(P?[0-9]+)/`),
 	regexp.MustCompile(`artifacts_list/(P?[0-9]+)/`)}

 // FullOSImageState copies the full Android/Cros image onto DUT disk.
 type FullOSImageState struct {
 	params *CrossOverParameters
 }

 func NewFullOSImageState(params *CrossOverParameters) commonutils.ServiceState {
 	return &FullOSImageState{
 		params: params,
 	}
 }

 // reboot restarts the DUT using the servo.
 func reboot(ctx context.Context, log *log.Logger, params *CrossOverParameters) error {
 	defer setPDRole(ctx, log, "src", params)

 	if err := servoSpecialColdResetModeForC2D2(ctx, log, params); err != nil {
 		return fmt.Errorf("INFRA: unable to update cold reset mode: %w", err)
 	}
 	if err := callServodRetry(ctx, log, "power_state", "off", params); err != nil {
 		return fmt.Errorf("INFRA: unable to set turn off USB: %w", err)
 	}
 	log.Printf("\nWaiting for the power state to turn off.")
 	// TODO: Instead of fixed wait time, use a poll based status checker. Refer WaitForPowerStates and GetECSystemPowerState in firmware code.
 	time.Sleep(waitForPowerOff)
 	if err := callServodRetry(ctx, log, "power_state", "rec", params); err != nil {
 		// A hard failure isn't necessary at this point, as subsequent checks will identify servod call failures and system state problems.
 		log.Printf("Letting passthrough from failed rec mode boot: %s", err)
 	}
 	return nil
 }

 func (s FullOSImageState) Execute(ctx context.Context, log *log.Logger) (*anypb.Any, api.InstallResponse_Status, error) {
 	log.Println("Executing " + s.Name())
 	dutAddress := fmt.Sprintf("%s:%v", s.params.Dut.GetChromeos().GetSsh().GetAddress(), s.params.Dut.GetChromeos().GetSsh().GetPort())

 	// Override to adb, should split this out to own state.
 	useAdb := false

 	// Should use props, but this is prior to connection being established.
 	if strings.Contains(s.params.TargetImagePath.GetPath(), "bluey") {
 		// Could have inlined into initial setting, but this is simple enough w/ logging.
 		log.Printf("Using adb for connection from USB boot")
 		useAdb = true
 	}

 	var client *ssh.Client

 	// f calls the checks connection function and returns an error on a nil client. This
 	// will be used as the UDF for the generic retry function.
 	f := func() error {
 		fail := false

 		if useAdb {
 			fail = !checkADB(log, s.params.Dut.GetChromeos().GetSsh().GetAddress(), bootWaitRetryCount*bootWaitRetryInterval)
 		} else {
 			client = checkSSH(log, dutAddress, bootWaitRetryCount, bootWaitRetryInterval)
 			fail = client == nil
 		}

 		if fail {
 			setPDRole(ctx, log, "src", s.params)
 			err := reboot(ctx, log, s.params)
 			if err != nil {
 				return err
 			}

 			return fmt.Errorf("cannot connect onto DUT booted from USB common provision image")
 		}

 		return nil
 	}

 	// Attempt to connect to the DUT post boot from USB. Retry
 	// checkConnectionRetryCount times using the common retry function.
 	if err := commonutils.Retry(log, checkConnectionRetryCount, "checkConnection", f); err != nil {
 		setPDRole(ctx, log, "src", s.params)
 		return commonutils.WrapStringInAny(s.errStatus(err.Error())), api.InstallResponse_STATUS_PRE_PROVISION_USB_BOOT_FAILURE, errors.New(err.Error())

 	}
 	if !useAdb {
 		defer client.Close()
 	}
 	setPDRole(ctx, log, "src", s.params)

 	cacheServer := s.params.Dut.GetCacheServer().GetAddress()
 	ipFunc := func() {
 		log.Println("IP info")
 		var session *ssh.Session
 		session, err := client.NewSession()
 		if err != nil {
 			log.Println("Failed to create session: ", err)
 			return
 		}

 		// Jumble the command, we don't care what fails or not, any output, simply print it.
 		const ipCmd = "set -e; ip rule show; ip route show table all; route; ifconfig eth0; echo 'All commands succeeded'"
 		if output, err := session.CombinedOutput(ipCmd); err != nil {
 			log.Println("Warning: failed to run IP info commands: ", ipCmd, err)
 			log.Println(string(output))
 		} else {
 			log.Println(string(output))
 		}
 	}
 	networkWaitFunc := func() error {
 		shorterCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
 		defer cancel()
 		networkCheckCmd := fmt.Sprintf("curl --max-time 5 -s -o/dev/null http://%s:%v/check_health", cacheServer.GetAddress(), cacheServer.GetPort())
 		for {
 			var session *ssh.Session
 			session, err := client.NewSession()
 			if err != nil {
 				log.Println("Failed to create session: ", err)
 				return err
 			}

 			select {
 			case <-shorterCtx.Done():
 				return errors.New("Failed while waiting for network connection.")
 			default:
 				if err := session.Run(networkCheckCmd); err != nil {
 					ipFunc()
 					// Retry w/ delay.
 					time.Sleep(time.Second)
 					continue
 				}
 				return nil
 			}
 		}
 	}
 	if useAdb {
 		// We can likely skip this for now, until necessary.
 		log.Println("Using adb, skipping network wait")
 	} else {
 		if err := networkWaitFunc(); err != nil {
 			log.Println(err)
 		}
 	}

 	installCommand := ""
 	targetImgPath := s.params.TargetImagePath.GetPath()
 	if strings.Contains(targetImgPath, androidBuild) {
 		buildId, err := targetBuild(targetImgPath)
 		if err != nil {
 			return commonutils.WrapStringInAny(s.errStatus("INFRA: unable to parse image path from target")), api.InstallResponse_STATUS_INVALID_REQUEST, err
 		}
 		launchTarget, err := extractTargetLaunch(targetImgPath)
 		if err != nil {
 			return commonutils.WrapStringInAny(s.errStatus("INFRA: unable to parse image path from target")), api.InstallResponse_STATUS_INVALID_REQUEST, err
 		}
 		if strings.Contains(launchTarget, engImg) {
 			// Uart logs are only enabled in -eng img.
 			callServodSET(ctx, log, cpuUartCapture, on, s.params)
 		}

 		installCommand = fmt.Sprintf("al-install android-build/builds/%s/%s/attempts/latest/artifacts/android-desktop_image.bin.gz %s", buildId, launchTarget, cacheServer.GetAddress())

 		log.Printf("LaunchTarget: %s\n", launchTarget)
 		if strings.HasSuffix(launchTarget, userImg) {
 			userImgType := "user"
 			ramdiskType := "vendor_boot-debug.img"
 			found := 0

 			for _, arg := range s.params.ExecutionMetadata.GetArgs() {
 				if arg.GetFlag() == "useSignedImage" && arg.GetValue() == "true" {
 					userImgType = "signed-user"
 					found = found + 1
 				}
 				if arg.GetFlag() == "useTestRamdisk" && arg.GetValue() == "true" {
 					ramdiskType = "vendor_boot-test-harness.img"
 					found = found + 1
 				}
 				if found == 2 {
 					break
 				}
 			}

 			// User images require the debug vendor boot to be flashed in the
 			// provision script.
 			// Need the "" for 3rd argument to be backwards compatible with partners.
 			// For pre-v10.0.0 images, it will default to "user" image type (previously "swap-vendor-boot").
 			installCommand = fmt.Sprintf("%s \"\" %s %s", installCommand, userImgType, ramdiskType)
 		}

 	} else {
 		installCommand = fmt.Sprintf("cros-install %s %s", targetImgPath[5:], cacheServer.GetAddress())
 	}

 	if useAdb {
 		if _, err := adb.AdbShellCmd(
 			strings.Split(installCommand, " "),
 			s.params.Dut.GetChromeos().GetSsh().GetAddress(),
 			log, 3, int(installTimeout.Seconds())); err != nil {
 			log.Println("Failed to run installCommand over adb: ", installCommand, err)
 			return commonutils.WrapStringInAny(s.errStatus("FAILED TO RUN CROSS OVER INSTALL COMMAND")), api.InstallResponse_STATUS_DOWNLOADING_IMAGE_FAILED, err
 		}
 	} else if err := installCommandWithRetry(client, installCommand, log); err != nil {
 		log.Println("Failed to run installCommand: ", installCommand, err)
 		collectUARTLogs(ctx, log, s.params)
 		return commonutils.WrapStringInAny(s.errStatus("FAILED TO RUN CROSS OVER INSTALL COMMAND")), api.InstallResponse_STATUS_DOWNLOADING_IMAGE_FAILED, err
 	}
 	errStatus, responseStatus, err := s.handleReboot(ctx, client, log)
 	if err != nil {
 		collectUARTLogs(ctx, log, s.params)
 	}
 	return errStatus, responseStatus, err
 }

 func (s FullOSImageState) errStatus(curErr string) string {
 	if s.params.PrevError != "" {
 		return fmt.Sprintf("Flash and OTA failed: %s", s.params.PrevError)
 	}
 	return curErr
 }

 func (s FullOSImageState) Next() commonutils.ServiceState {
 	return NewValidateState(s.params)
 }

 func (s FullOSImageState) Name() string {
 	return "Full OS Image State"
 }

 func (s FullOSImageState) handleReboot(ctx context.Context, sshClient *ssh.Client, log *log.Logger) (*anypb.Any, api.InstallResponse_Status, error) {
 	if err := callServodRetry(ctx, log, "power_state", "off", s.params); err != nil {
 		return commonutils.WrapStringInAny("INFRA: unable to set turn off power via servo"), api.InstallResponse_STATUS_PROVISIONING_FAILED, err
 	}
 	if err := sshFailWait(sshClient, log); err != nil {
 		return commonutils.WrapStringInAny("INFRA: dut did not power off post provision image installation"), api.InstallResponse_STATUS_PROVISIONING_FAILED, err
 	}
 	log.Println("Device successfully powered off.")
 	if err := callServodRetry(ctx, log, "image_usbkey_direction", "servo_sees_usbkey", s.params); err != nil {
 		return commonutils.WrapStringInAny("INFRA: unable to set usb direction via servo"), api.InstallResponse_STATUS_PROVISIONING_FAILED, err
 	}
 	log.Printf("\nWaiting for the image_usbkey_direction.")
 	// TODO: Instead of fixed wait time, use a poll based status checker. Refer WaitForPowerStates and GetECSystemPowerState in firmware code.
 	time.Sleep(10 * time.Second)

 	if err := callServodRetry(ctx, log, "power_state", "on", s.params); err != nil {
 		// reset might be needed if power on failed: b/381982169
 		log.Println("Warning: unable to set turn on power via servo ", err)
 		log.Println("Waiting for 10 seconds before trying power resetting.")
 		time.Sleep(powerResetDelay)
 		if err := callServodRetry(ctx, log, "power_state", "reset", s.params); err != nil {
 			return commonutils.WrapStringInAny("INFRA: unable to reset power via servo"), api.InstallResponse_STATUS_PROVISIONING_FAILED, err
 		}
 	}
 	return nil, api.InstallResponse_STATUS_SUCCESS, nil
 }

 func installCommandWithRetry(client *ssh.Client, installCommand string, log *log.Logger) error {
 	log.Println("Running command on host now ", installCommand)
 	retryCount := 3
 	for ; retryCount >= 0; retryCount-- {
 		var session *ssh.Session
 		session, err := client.NewSession()
 		if err != nil {
 			log.Println("Failed to create session: ", err)
 			return err
 		}
 		output, err := session.CombinedOutput(installCommand)
 		log.Println(string(output))
 		if err == nil {
 			return nil
 		}
 		if !strings.Contains(string(output), networkErrMsg) {
 			return err
 		}
 		log.Println("Got error " + networkErrMsg + " while running install command. Will Retry...")
 	}
 	return fmt.Errorf("Failed to run installCommand")
 }

 // sshFailWait return nil, if the device cannot be sshed within 40 seconds.
 // If the device is alive even after 40 seconds, it returns error.
 func sshFailWait(sshClient *ssh.Client, log *log.Logger) error {
 	for attemp := 1; attemp < pingRetryCount; attemp++ {
 		log.Println("pinging device to check if its offline...")
 		if err := ping(sshClient, pingTimeout); err != nil {
 			return nil
 		}
 		time.Sleep(pingRetryDelay)
 	}
 	return errors.New("dut failed to become unreachable")
 }

 func ping(sshClient *ssh.Client, timeout time.Duration) error {
 	ch := make(chan error, 1)
 	go func() {
 		_, _, err := sshClient.SendRequest(sshMsgIgnore, true, []byte{})
 		ch <- err
 	}()

 	select {
 	case err := <-ch:
 		return err
 	case <-time.After(timeout):
 		return errors.New("timed out")
 	}
 }

 func extractTargetLaunch(path string) (string, error) {
 	for _, re := range launchTargetPatterns {
 		matches := re.FindStringSubmatch(path)
 		if len(matches) == 2 {
 			return matches[1], nil
 		}
 	}
 	return "", fmt.Errorf("could not extract launch from %s", path)
 }

 func targetBuild(imagePath string) (string, error) {
 	for _, re := range buildIdPatterns {
 		matches := re.FindStringSubmatch(imagePath)
 		if len(matches) == 2 {
 			return matches[1], nil
 		}
 	}
 	return "", fmt.Errorf("could not extract buildId from %s", imagePath)
 }

 func (s FullOSImageState) Retry() bool {
 	return false
 }
	// Copyright 2024 The Chromium Authors
	// Use of this source code is governed by a BSD-style license that can be
	// found in the LICENSE file.

	// DLC constants and helpers
	package cross_over

	import (
	"context"
	"errors"
	"fmt"
	"log"
	"regexp"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
	"google.golang.org/protobuf/types/known/anypb"

	"go.chromium.org/chromiumos/config/go/test/api"

	"go.chromium.org/infra/cros/cmd/cft/common/adb"
	commonutils "go.chromium.org/infra/cros/cmd/provision/common-utils"
	)

	const (
	// sshMsgIgnore is the SSH global message sent to ping the host.
	// See RFC 4253 11.2, "Ignored Data Message".
	sshMsgIgnore = "SSH_MSG_IGNORE"
	pingTimeout = time.Second
	pingRetryDelay = 2 * time.Second
	pingRetryCount = 20
	networkWaitTime = 1 * time.Minute
	powerResetDelay = 10 * time.Second
	installTimeout = 12 * time.Minute
	networkErrMsg = "Network is unreachable"
	cpuUartCapture = "cpu_uart_capture"
	engImg = "-eng"
	userImg = "-user"
	androidBuild = "android-build"
	checkConnectionRetryCount = 3
	)

	var launchTargetPatterns = []*regexp.Regexp{
	regexp.MustCompile(`build_details/P?[0-9]+/([a-z_\-]+)`),
	regexp.MustCompile(`artifacts_list/P?[0-9]+/([a-z_\-]+)`)}

	var buildIdPatterns = []*regexp.Regexp{
	regexp.MustCompile(`build_details/(P?[0-9]+)/`),
	regexp.MustCompile(`artifacts_list/(P?[0-9]+)/`)}

	// FullOSImageState copies the full Android/Cros image onto DUT disk.
	type FullOSImageState struct {
	params *CrossOverParameters
	}

	func NewFullOSImageState(params *CrossOverParameters) commonutils.ServiceState {
	return &FullOSImageState{
	params: params,
	}
	}

	// reboot restarts the DUT using the servo.
	func reboot(ctx context.Context, log log.Logger, params CrossOverParameters) error {
	defer setPDRole(ctx, log, "src", params)

	if err := servoSpecialColdResetModeForC2D2(ctx, log, params); err != nil {
	return fmt.Errorf("INFRA: unable to update cold reset mode: %w", err)
	}
	if err := callServodRetry(ctx, log, "power_state", "off", params); err != nil {
	return fmt.Errorf("INFRA: unable to set turn off USB: %w", err)
	}
	log.Printf("\nWaiting for the power state to turn off.")
	// TODO: Instead of fixed wait time, use a poll based status checker. Refer WaitForPowerStates and GetECSystemPowerState in firmware code.
	time.Sleep(waitForPowerOff)
	if err := callServodRetry(ctx, log, "power_state", "rec", params); err != nil {
	// A hard failure isn't necessary at this point, as subsequent checks will identify servod call failures and system state problems.
	log.Printf("Letting passthrough from failed rec mode boot: %s", err)
	}
	return nil
	}

	func (s FullOSImageState) Execute(ctx context.Context, log log.Logger) (anypb.Any, api.InstallResponse_Status, error) {
	log.Println("Executing " + s.Name())
	dutAddress := fmt.Sprintf("%s:%v", s.params.Dut.GetChromeos().GetSsh().GetAddress(), s.params.Dut.GetChromeos().GetSsh().GetPort())

	// Override to adb, should split this out to own state.
	useAdb := false

	// Should use props, but this is prior to connection being established.
	if strings.Contains(s.params.TargetImagePath.GetPath(), "bluey") {
	// Could have inlined into initial setting, but this is simple enough w/ logging.
	log.Printf("Using adb for connection from USB boot")
	useAdb = true
	}

	var client *ssh.Client

	// f calls the checks connection function and returns an error on a nil client. This
	// will be used as the UDF for the generic retry function.
	f := func() error {
	fail := false

	if useAdb {
	fail = !checkADB(log, s.params.Dut.GetChromeos().GetSsh().GetAddress(), bootWaitRetryCount*bootWaitRetryInterval)
	} else {
	client = checkSSH(log, dutAddress, bootWaitRetryCount, bootWaitRetryInterval)
	fail = client == nil
	}

	if fail {
	setPDRole(ctx, log, "src", s.params)
	err := reboot(ctx, log, s.params)
	if err != nil {
	return err
	}

	return fmt.Errorf("cannot connect onto DUT booted from USB common provision image")
	}

	return nil
	}

	// Attempt to connect to the DUT post boot from USB. Retry
	// checkConnectionRetryCount times using the common retry function.
	if err := commonutils.Retry(log, checkConnectionRetryCount, "checkConnection", f); err != nil {
	setPDRole(ctx, log, "src", s.params)
	return commonutils.WrapStringInAny(s.errStatus(err.Error())), api.InstallResponse_STATUS_PRE_PROVISION_USB_BOOT_FAILURE, errors.New(err.Error())

	}
	if !useAdb {
	defer client.Close()
	}
	setPDRole(ctx, log, "src", s.params)

	cacheServer := s.params.Dut.GetCacheServer().GetAddress()
	ipFunc := func() {
	log.Println("IP info")
	var session *ssh.Session
	session, err := client.NewSession()
	if err != nil {
	log.Println("Failed to create session: ", err)
	return
	}

	// Jumble the command, we don't care what fails or not, any output, simply print it.
	const ipCmd = "set -e; ip rule show; ip route show table all; route; ifconfig eth0; echo 'All commands succeeded'"
	if output, err := session.CombinedOutput(ipCmd); err != nil {
	log.Println("Warning: failed to run IP info commands: ", ipCmd, err)
	log.Println(string(output))
	} else {
	log.Println(string(output))
	}
	}
	networkWaitFunc := func() error {
	shorterCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()
	networkCheckCmd := fmt.Sprintf("curl --max-time 5 -s -o/dev/null http://%s:%v/check_health", cacheServer.GetAddress(), cacheServer.GetPort())
	for {
	var session *ssh.Session
	session, err := client.NewSession()
	if err != nil {
	log.Println("Failed to create session: ", err)
	return err
	}

	select {
	case <-shorterCtx.Done():
	return errors.New("Failed while waiting for network connection.")
	default:
	if err := session.Run(networkCheckCmd); err != nil {
	ipFunc()
	// Retry w/ delay.
	time.Sleep(time.Second)
	continue
	}
	return nil
	}
	}
	}
	if useAdb {
	// We can likely skip this for now, until necessary.
	log.Println("Using adb, skipping network wait")
	} else {
	if err := networkWaitFunc(); err != nil {
	log.Println(err)
	}
	}

	installCommand := ""
	targetImgPath := s.params.TargetImagePath.GetPath()
	if strings.Contains(targetImgPath, androidBuild) {
	buildId, err := targetBuild(targetImgPath)
	if err != nil {
	return commonutils.WrapStringInAny(s.errStatus("INFRA: unable to parse image path from target")), api.InstallResponse_STATUS_INVALID_REQUEST, err
	}
	launchTarget, err := extractTargetLaunch(targetImgPath)
	if err != nil {
	return commonutils.WrapStringInAny(s.errStatus("INFRA: unable to parse image path from target")), api.InstallResponse_STATUS_INVALID_REQUEST, err
	}
	if strings.Contains(launchTarget, engImg) {
	// Uart logs are only enabled in -eng img.
	callServodSET(ctx, log, cpuUartCapture, on, s.params)
	}

	installCommand = fmt.Sprintf("al-install android-build/builds/%s/%s/attempts/latest/artifacts/android-desktop_image.bin.gz %s", buildId, launchTarget, cacheServer.GetAddress())

	log.Printf("LaunchTarget: %s\n", launchTarget)
	if strings.HasSuffix(launchTarget, userImg) {
	userImgType := "user"
	ramdiskType := "vendor_boot-debug.img"
	found := 0

	for _, arg := range s.params.ExecutionMetadata.GetArgs() {
	if arg.GetFlag() == "useSignedImage" && arg.GetValue() == "true" {
	userImgType = "signed-user"
	found = found + 1
	}
	if arg.GetFlag() == "useTestRamdisk" && arg.GetValue() == "true" {
	ramdiskType = "vendor_boot-test-harness.img"
	found = found + 1
	}
	if found == 2 {
	break
	}
	}

	// User images require the debug vendor boot to be flashed in the
	// provision script.
	// Need the "" for 3rd argument to be backwards compatible with partners.
	// For pre-v10.0.0 images, it will default to "user" image type (previously "swap-vendor-boot").
	installCommand = fmt.Sprintf("%s \"\" %s %s", installCommand, userImgType, ramdiskType)
	}

	} else {
	installCommand = fmt.Sprintf("cros-install %s %s", targetImgPath[5:], cacheServer.GetAddress())
	}

	if useAdb {
	if _, err := adb.AdbShellCmd(
	strings.Split(installCommand, " "),
	s.params.Dut.GetChromeos().GetSsh().GetAddress(),
	log, 3, int(installTimeout.Seconds())); err != nil {
	log.Println("Failed to run installCommand over adb: ", installCommand, err)
	return commonutils.WrapStringInAny(s.errStatus("FAILED TO RUN CROSS OVER INSTALL COMMAND")), api.InstallResponse_STATUS_DOWNLOADING_IMAGE_FAILED, err
	}
	} else if err := installCommandWithRetry(client, installCommand, log); err != nil {
	log.Println("Failed to run installCommand: ", installCommand, err)
	collectUARTLogs(ctx, log, s.params)
	return commonutils.WrapStringInAny(s.errStatus("FAILED TO RUN CROSS OVER INSTALL COMMAND")), api.InstallResponse_STATUS_DOWNLOADING_IMAGE_FAILED, err
	}
	errStatus, responseStatus, err := s.handleReboot(ctx, client, log)
	if err != nil {
	collectUARTLogs(ctx, log, s.params)
	}
	return errStatus, responseStatus, err
	}

	func (s FullOSImageState) errStatus(curErr string) string {
	if s.params.PrevError != "" {
	return fmt.Sprintf("Flash and OTA failed: %s", s.params.PrevError)
	}
	return curErr
	}

	func (s FullOSImageState) Next() commonutils.ServiceState {
	return NewValidateState(s.params)
	}

	func (s FullOSImageState) Name() string {
	return "Full OS Image State"
	}

	func (s FullOSImageState) handleReboot(ctx context.Context, sshClient ssh.Client, log log.Logger) (*anypb.Any, api.InstallResponse_Status, error) {
	if err := callServodRetry(ctx, log, "power_state", "off", s.params); err != nil {
	return commonutils.WrapStringInAny("INFRA: unable to set turn off power via servo"), api.InstallResponse_STATUS_PROVISIONING_FAILED, err
	}
	if err := sshFailWait(sshClient, log); err != nil {
	return commonutils.WrapStringInAny("INFRA: dut did not power off post provision image installation"), api.InstallResponse_STATUS_PROVISIONING_FAILED, err
	}
	log.Println("Device successfully powered off.")
	if err := callServodRetry(ctx, log, "image_usbkey_direction", "servo_sees_usbkey", s.params); err != nil {
	return commonutils.WrapStringInAny("INFRA: unable to set usb direction via servo"), api.InstallResponse_STATUS_PROVISIONING_FAILED, err
	}
	log.Printf("\nWaiting for the image_usbkey_direction.")
	// TODO: Instead of fixed wait time, use a poll based status checker. Refer WaitForPowerStates and GetECSystemPowerState in firmware code.
	time.Sleep(10 * time.Second)

	if err := callServodRetry(ctx, log, "power_state", "on", s.params); err != nil {
	// reset might be needed if power on failed: b/381982169
	log.Println("Warning: unable to set turn on power via servo ", err)
	log.Println("Waiting for 10 seconds before trying power resetting.")
	time.Sleep(powerResetDelay)
	if err := callServodRetry(ctx, log, "power_state", "reset", s.params); err != nil {
	return commonutils.WrapStringInAny("INFRA: unable to reset power via servo"), api.InstallResponse_STATUS_PROVISIONING_FAILED, err
	}
	}
	return nil, api.InstallResponse_STATUS_SUCCESS, nil
	}

	func installCommandWithRetry(client ssh.Client, installCommand string, log log.Logger) error {
	log.Println("Running command on host now ", installCommand)
	retryCount := 3
	for ; retryCount >= 0; retryCount-- {
	var session *ssh.Session
	session, err := client.NewSession()
	if err != nil {
	log.Println("Failed to create session: ", err)
	return err
	}
	output, err := session.CombinedOutput(installCommand)
	log.Println(string(output))
	if err == nil {
	return nil
	}
	if !strings.Contains(string(output), networkErrMsg) {
	return err
	}
	log.Println("Got error " + networkErrMsg + " while running install command. Will Retry...")
	}
	return fmt.Errorf("Failed to run installCommand")
	}

	// sshFailWait return nil, if the device cannot be sshed within 40 seconds.
	// If the device is alive even after 40 seconds, it returns error.
	func sshFailWait(sshClient ssh.Client, log log.Logger) error {
	for attemp := 1; attemp < pingRetryCount; attemp++ {
	log.Println("pinging device to check if its offline...")
	if err := ping(sshClient, pingTimeout); err != nil {
	return nil
	}
	time.Sleep(pingRetryDelay)
	}
	return errors.New("dut failed to become unreachable")
	}

	func ping(sshClient *ssh.Client, timeout time.Duration) error {
	ch := make(chan error, 1)
	go func() {
	_, _, err := sshClient.SendRequest(sshMsgIgnore, true, []byte{})
	ch <- err
	}()

	select {
	case err := <-ch:
	return err
	case <-time.After(timeout):
	return errors.New("timed out")
	}
	}

	func extractTargetLaunch(path string) (string, error) {
	for _, re := range launchTargetPatterns {
	matches := re.FindStringSubmatch(path)
	if len(matches) == 2 {
	return matches[1], nil
	}
	}
	return "", fmt.Errorf("could not extract launch from %s", path)
	}

	func targetBuild(imagePath string) (string, error) {
	for _, re := range buildIdPatterns {
	matches := re.FindStringSubmatch(imagePath)
	if len(matches) == 2 {
	return matches[1], nil
	}
	}
	return "", fmt.Errorf("could not extract buildId from %s", imagePath)
	}

	func (s FullOSImageState) Retry() bool {
	return false
	}