src/fastboot/cmd.c

/*
 * Copyright (C) 2021 Google Inc.
 *
 * See file CREDITS for list of people who contributed to this
 * project.
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation; either version 2 of
 * the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but without any warranty; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 */

#include "boot/avb.h"
#include "boot/bootbyte.h"
#include "boot/zbi.h"
#include "config.h"
#include "fastboot/cmd.h"
#include "fastboot/disk.h"
#include "fastboot/fastboot.h"
#include "fastboot/tcp.h"
#include "fastboot/vars.h"
#include "fuchsia-sdk/pkg/avb/libavb/libavb.h"
#include "fuchsia-sdk/pkg/avb/libavb_atx/libavb_atx.h"
#include "fuchsia-sdk/pkg/zircon_boot/zircon_vboot.h"
#include "gpt_misc.h"
#include <stdlib.h>
#include <libpayload.h>

static int parse_hex(const char *str, size_t len, uint32_t *ret)
{
	int valid = 0;
	uint32_t result = 0;
	for (size_t i = 0; i < len; i++) {
		char c = str[i];
		int units = 0;
		if (c >= '0' && c <= '9') {
			units = c - '0';
		} else if (c >= 'A' && c <= 'F') {
			units = 10 + (c - 'A');
		} else if (c >= 'a' && c <= 'f') {
			units = 10 + (c - 'a');
		} else {
			break;
		}

		result *= 0x10;
		result += units;
		valid++;
	}

	*ret = result;

	return valid;
}

static void fastboot_cmd_continue(fastboot_session_t *fb, const char *arg,
				  uint64_t arg_len)
{
	fastboot_okay(fb, "Continuing boot");
	fb->state = FINISHED;
}

static void fastboot_cmd_download(fastboot_session_t *fb, const char *arg,
				  uint64_t arg_len)
{
	uint32_t size = 0;
	int digits = parse_hex(arg, arg_len, &size);
	if (digits != arg_len) {
		fastboot_fail(fb, "Invalid argument");
		return;
	}

	if (size > FASTBOOT_MAX_DOWNLOAD_SIZE) {
		fastboot_fail(fb, "File too big");
		return;
	}

	fastboot_data(fb, size);
}

static void fastboot_cmd_flash(fastboot_session_t *fb, const char *arg,
			       uint64_t arg_len)
{
	if (!fb->has_download) {
		fastboot_fail(fb, "No data staged to flash");
		return;
	}

	uint64_t data_len;
	void *data = fastboot_get_download_buffer(fb, &data_len);

	if (!strncmp(arg, "bootloader", arg_len)) {
		fastboot_write_bootloader(fb, data, data_len);
		return;
	} else if (!strncmp(arg, "gpt", arg_len)) {
		fastboot_write_gpt(fb, data, data_len);
		return;
	}

	struct fastboot_disk disk;
	if (!fastboot_disk_init(&disk)) {
		fastboot_fail(fb, "Failed to init disk");
		return;
	}

	fastboot_write(fb, &disk, arg, arg_len, data, (uint32_t)data_len);
	fastboot_disk_destroy(&disk);
}

static size_t upload_size = 0;

void fastboot_set_upload_size(size_t size) { upload_size = size; }

static void fastboot_cmd_upload(fastboot_session_t *fb, const char *arg,
				uint64_t arg_len)
{
	if (upload_size == 0) {
		fastboot_fail(fb, "No data has been staged");
		return;
	}

	char response[FASTBOOT_MSG_MAX + 1];
	snprintf(response, FASTBOOT_MSG_MAX, "DATA%08zx", upload_size);
	fastboot_send(response, strlen(response));

	// Send the data
	printf("Uploading %zu bytes of data...\n", upload_size);
	fastboot_send(fastboot_get_upload_buffer(), upload_size);
	printf("Done\n");

	// Reset upload size.
	fastboot_set_upload_size(0);
	fastboot_succeed(fb);
}

static void fastboot_cmd_oem_stage_partition(fastboot_session_t *fb,
					     const char *arg, uint64_t arg_len)
{
	struct fastboot_disk disk;
	if (!fastboot_disk_init(&disk)) {
		fastboot_fail(fb, "Failed to init disk");
		return;
	}

	if (!strncmp(arg, "gpt", arg_len))
		fastboot_stage_gpt(fb, &disk);
	else
		fastboot_fail(fb, "Partition is unsupported for staging");

	fastboot_disk_destroy(&disk);
}

static void fastboot_cmd_erase(fastboot_session_t *fb, const char *arg,
			       uint64_t arg_len)
{
	struct fastboot_disk disk;
	if (!fastboot_disk_init(&disk)) {
		fastboot_fail(fb, "Failed to init disk");
		return;
	}

	fastboot_erase(fb, &disk, arg, arg_len);
	fastboot_disk_destroy(&disk);
}

static void fastboot_cmd_oem_add_staged_bootloader_file(fastboot_session_t *fb,
							const char *arg,
							uint64_t arg_len)
{
	if (!fb->has_download) {
		fastboot_fail(fb, "Nothing staged");
		return;
	}

	uint64_t data_len;
	const void *data = fastboot_get_download_buffer(fb, &data_len);
	if (zbi_stage_bootloader_file(arg, arg_len, data, data_len)) {
		fastboot_fail(fb, "Failed to stage bootloader file");
		return;
	}

	fastboot_succeed(fb);
}

// `fastboot oem get-kernels` returns a list of slot letter:kernel mapping.
// This is useful for us because the partition tables are more flexible than on
// more traditional devices, so there could be many kernels.
// For instance, installing Fuchsia on a Chromebook will result in the device
// having five kernel partitions:
// KERN-A
// KERN-B
// zircon-a
// zircon-b
// zircon-r
//
// We map them to fastboot slots by having the first slot be the first partition
// we found.
struct get_kernels_ctx {
	fastboot_session_t *fb;
	int cur_kernel;
};
static bool get_kernels_cb(void *ctx, int index, GptEntry *e,
			   char *partition_name)
{
	if (get_slot_for_partition_name(e, partition_name) == 0)
		return false;

	struct get_kernels_ctx *gk = (struct get_kernels_ctx *)ctx;
	char fastboot_slot = gk->cur_kernel + 'a';
	fastboot_info(gk->fb, "%c:%s:prio=%d", fastboot_slot, partition_name,
		      GetEntryPriority(e));
	gk->cur_kernel++;
	return false;
}
static void fastboot_cmd_oem_get_kernels(fastboot_session_t *fb,
					 const char *arg, uint64_t arg_len)
{
	struct fastboot_disk disk;
	if (!fastboot_disk_init(&disk)) {
		fastboot_fail(fb, "Failed to init disk");
		return;
	}
	struct get_kernels_ctx ctx = {
		.fb = fb,
		.cur_kernel = 0,
	};
	fastboot_disk_foreach_partition(&disk, get_kernels_cb, &ctx);
	fastboot_disk_destroy(&disk);
	fastboot_succeed(fb);
}

static void fastboot_cmd_reboot(fastboot_session_t *fb, const char *arg,
				uint64_t arg_len)
{
	fastboot_succeed(fb);
	fb->state = REBOOT;
}

static void fastboot_cmd_reboot_recovery(fastboot_session_t *fb,
					 const char *arg, uint64_t arg_len)
{
	bootbyte_set_mode(BOOT_MODE_RECOVERY);
	fastboot_succeed(fb);
	fb->state = REBOOT;
}

static void fastboot_cmd_reboot_bootloader(fastboot_session_t *fb,
					   const char *arg, uint64_t arg_len)
{
	bootbyte_set_mode(BOOT_MODE_BOOTLOADER);
	fastboot_succeed(fb);
	fb->state = REBOOT;
}

static void fastboot_cmd_set_active(fastboot_session_t *fb, const char *arg,
				    uint64_t arg_len)
{
	struct fastboot_disk disk;
	if (!fastboot_disk_init(&disk)) {
		fastboot_fail(fb, "Failed to init disk");
		return;
	}
	GptEntry *slot = fastboot_get_kernel_for_slot(&disk, arg[0]);
	if (slot == NULL) {
		fastboot_fail(fb, "Could not find slot");
		goto out;
	}

	fastboot_slots_disable_all(&disk);
	GptUpdateKernelWithEntry(disk.gpt, slot, GPT_UPDATE_ENTRY_ACTIVE);

	// Keep zircon_r successful. Otherwise it's left invalid and
	// `dm reboot-recovery` and `fastboot recovery` won't work.
	GptEntry *zircon_r_slot =
		fastboot_find_partition(&disk, "zircon_r", strlen("zircon_r"));
	if (zircon_r_slot == NULL) {
		fastboot_fail(fb, "Failed to find zircon_r slot");
		goto out;
	}

	SetEntrySuccessful(zircon_r_slot, 1);
	GptModified(disk.gpt);

	fastboot_succeed(fb);
out:
	fastboot_disk_destroy(&disk);
}

static void fastboot_vx_get_unlock_challenge(fastboot_session_t *fb,
					     const char *arg, uint64_t arg_len)
{
	if (sizeof(AvbAtxUnlockChallenge) > fastboot_get_upload_buffer_size()) {
		fastboot_fail(fb, "Upload buffer is too small");
		return;
	}

	struct fastboot_disk disk;
	if (!fastboot_disk_init(&disk)) {
		fastboot_fail(fb, "Failed to init disk");
		return;
	}

	struct zircon_boot_context priv;
	ZirconBootOps zircon_boot_ops;
	setup_zircon_boot_context_and_ops(disk.disk, disk.gpt, &priv,
					  &zircon_boot_ops);

	AvbAtxUnlockChallenge unlock_challenge;
	if (!ZirconVbootGenerateUnlockChallenge(&zircon_boot_ops,
						&unlock_challenge)) {
		fastboot_fail(fb, "Failed to generate unlock challenge");
		fastboot_disk_destroy(&disk);
		return;
	}

	memcpy(fastboot_get_upload_buffer(), &unlock_challenge,
	       sizeof(unlock_challenge));
	fastboot_set_upload_size(sizeof(unlock_challenge));
	fastboot_disk_destroy(&disk);
	fastboot_succeed(fb);
}

static void fastboot_vx_unlock(fastboot_session_t *fb, const char *arg,
			       uint64_t arg_len)
{
	uint64_t data_len;
	void *download_buffer = fastboot_get_download_buffer(fb, &data_len);
	if (data_len != sizeof(AvbAtxUnlockCredential)) {
		fastboot_fail(fb, "Invalid unlock credential size");
		return;
	}

	static AvbAtxUnlockCredential unlock_credential;
	memcpy(&unlock_credential, download_buffer, data_len);

	struct fastboot_disk disk;
	if (!fastboot_disk_init(&disk)) {
		fastboot_fail(fb, "Failed to init disk");
		return;
	}

	struct zircon_boot_context priv;
	ZirconBootOps zircon_boot_ops;
	setup_zircon_boot_context_and_ops(disk.disk, disk.gpt, &priv,
					  &zircon_boot_ops);

	if (!ZirconVbootValidateUnlockCredential(&zircon_boot_ops,
						 &unlock_credential)) {
		fastboot_fail(fb, "Invalid unlock credential");
		return;
	}

	unlock_device();
	fastboot_succeed(fb);
}

static void fastboot_vx_lock(fastboot_session_t *fb, const char *arg,
			     uint64_t arg_len)
{
	lock_device();
	fastboot_succeed(fb);
}

#define CMD_ARGS(_name, _sep, _fn)                                             \
	{                                                                      \
		.name = _name, .has_args = true, .sep = _sep, .fn = _fn        \
	}
#define CMD_NO_ARGS(_name, _fn)                                                \
	{                                                                      \
		.name = _name, .has_args = false, .fn = _fn                    \
	}
struct fastboot_cmd fastboot_cmds[] = {
	CMD_NO_ARGS("continue", fastboot_cmd_continue),
	CMD_ARGS("download", ':', fastboot_cmd_download),
	CMD_NO_ARGS("upload", fastboot_cmd_upload),
	CMD_ARGS("erase", ':', fastboot_cmd_erase),
	CMD_ARGS("flash", ':', fastboot_cmd_flash),
	CMD_ARGS("getvar", ':', fastboot_cmd_getvar),
	CMD_ARGS("oem add-staged-bootloader-file", ' ',
		 fastboot_cmd_oem_add_staged_bootloader_file),
	CMD_NO_ARGS("oem get-kernels", fastboot_cmd_oem_get_kernels),
	CMD_NO_ARGS("oem vx-lock", fastboot_vx_lock),
	CMD_NO_ARGS("oem vx-unlock", fastboot_vx_unlock),
	CMD_ARGS("oem stage-partition", ' ', fastboot_cmd_oem_stage_partition),
	CMD_NO_ARGS("oem vx-get-unlock-challenge",
		    fastboot_vx_get_unlock_challenge),
	CMD_NO_ARGS("reboot-recovery", fastboot_cmd_reboot_recovery),
	CMD_NO_ARGS("reboot-bootloader", fastboot_cmd_reboot_bootloader),
	CMD_NO_ARGS("reboot", fastboot_cmd_reboot),
	CMD_ARGS("set_active", ':', fastboot_cmd_set_active),
	{
		.name = NULL,
	},
};