/*
 * Copyright (c) 2026 The Fuchsia Authors
 *
 * SPDX-License-Identifier:	BSD-3-Clause
 */

#include <common.h>
#include <command.h>
#include <environment.h>
#include <efi_loader.h>
#include <fb_fastboot.h>
#include <gbl_efi_avb_protocol.h>
#ifdef CONFIG_AML_ANTIROLLBACK
#include <anti-rollback.h>
#endif

static void ascii_to_utf16(u16 *dst, const char *src, size_t max_len)
{
	while (*src && max_len > 1) {
		*dst++ = (u16)*src++;
		max_len--;
	}
	*dst = 0;
}

static efi_status_t EFIAPI read_partition_attributes(
	struct GblEfiAvbProtocol *self, size_t *num_partitions,
	GblEfiAvbPartitionAttributes *partitions)
{
	EFI_ENTRY("%p, %p, %p", self, num_partitions, partitions);
	if (!num_partitions)
		return EFI_EXIT(EFI_INVALID_PARAMETER);
	// No partition override flags needed
	*num_partitions = 0;
	return EFI_EXIT(EFI_SUCCESS);
}

static efi_status_t EFIAPI read_device_status(
	struct GblEfiAvbProtocol *self, GblEfiAvbDeviceStatus *status_flags)
{
	EFI_ENTRY("%p, %p", self, status_flags);
	if (!status_flags) {
		return EFI_EXIT(EFI_INVALID_PARAMETER);
	}

	int lock_mask = check_lock_quiet();
	bool dev_unlocked = !(lock_mask & FASTBOOT_LOCK_DEVICE);

#ifdef CONFIG_AML_ANTIROLLBACK
	// When AML anti-rollback is enabled, query hardware RPMB lock state
	// which takes precedence over the software lock environment variable.
	uint32_t rpmb_lock = 0;
	if (get_avb_lock_state(&rpmb_lock)) {
		dev_unlocked = !rpmb_lock;
	}
#endif

	bool crit_unlocked = !(lock_mask & FASTBOOT_LOCK_CRITICAL);

	*status_flags = GBL_EFI_AVB_DEVICE_STATUS_UNLOCKABLE;
	if (dev_unlocked) {
		*status_flags |= GBL_EFI_AVB_DEVICE_STATUS_UNLOCKED;
	}
	if (crit_unlocked) {
		*status_flags |= GBL_EFI_AVB_DEVICE_STATUS_UNLOCKED_CRITICAL;
	}

	return EFI_EXIT(EFI_SUCCESS);
}

static efi_status_t EFIAPI validate_vbmeta_public_key(
	struct GblEfiAvbProtocol *self, size_t public_key_length,
	const uint8_t *public_key_data, size_t public_key_metadata_length,
	const uint8_t *public_key_metadata,
	GblEfiAvbKeyValidationStatus *validation_status)
{
	EFI_ENTRY("%p, %zu, %p, %zu, %p, %p", self, public_key_length,
		  public_key_data, public_key_metadata_length,
		  public_key_metadata, validation_status);
	if (!validation_status)
		return EFI_EXIT(EFI_INVALID_PARAMETER);

	// Public key validation is managed internally by libavb_atx in GBL's
	// Fuchsia boot flow via cert_validate_vbmeta_public_key.
	return EFI_EXIT(EFI_UNSUPPORTED);
}

static efi_status_t EFIAPI read_rollback_index(struct GblEfiAvbProtocol *self,
					       size_t index_location,
					       uint64_t *rollback_index)
{
	EFI_ENTRY("%p, %zu, %p", self, index_location, rollback_index);

	// Anti-rollback protection is not supported.
	return EFI_EXIT(EFI_UNSUPPORTED);
}

static efi_status_t EFIAPI write_rollback_index(struct GblEfiAvbProtocol *self,
						size_t index_location,
						uint64_t rollback_index)
{
	EFI_ENTRY("%p, %zu, %llu", self, index_location, rollback_index);

	// Anti-rollback protection is not supported.
	return EFI_EXIT(EFI_UNSUPPORTED);
}

static const efi_guid_t guid = EFI_GUID(0x6bc66b9a, 0xd5c9, 0x4c02, 0x9d, 0xa9,
					0x50, 0xaf, 0x19, 0x8d, 0x91, 0x2c);

static efi_status_t EFIAPI read_persistent_value(struct GblEfiAvbProtocol *self,
						 const char *name,
						 size_t *value_size,
						 uint8_t *value)
{
	EFI_ENTRY("%p, %s, %p, %p", self, name, value_size, value);
	if (!name || !value_size || (*value_size > 0 && !value))
		return EFI_EXIT(EFI_INVALID_PARAMETER);

	u16 var_name[64];
	efi_uintn_t data_size = *value_size;
	ascii_to_utf16(var_name, name, sizeof(var_name) / sizeof(u16));

	efi_status_t ret =
		efi_get_variable(var_name, &guid, NULL, &data_size, value);
	if (ret == EFI_BUFFER_TOO_SMALL) {
		*value_size = data_size;
		return EFI_EXIT(EFI_BUFFER_TOO_SMALL);
	}
	if (ret != EFI_SUCCESS)
		return EFI_EXIT(EFI_NOT_FOUND);

	*value_size = data_size;
	return EFI_EXIT(EFI_SUCCESS);
}

static efi_status_t EFIAPI
write_persistent_value(struct GblEfiAvbProtocol *self, const char *name,
		       size_t value_size, const uint8_t *value)
{
	u16 var_name[64];
	efi_status_t ret;
	u32 attr = EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS |
		   EFI_VARIABLE_RUNTIME_ACCESS;

	EFI_ENTRY("%p, %s, %zu, %p", self, name, value_size, value);
	if (!name || (value_size > 0 && !value))
		return EFI_EXIT(EFI_INVALID_PARAMETER);

	ascii_to_utf16(var_name, name, sizeof(var_name) / sizeof(u16));

	if (value_size == 0 || value == NULL)
		ret = efi_set_variable(var_name, &guid, attr, 0, NULL);
	else
		ret = efi_set_variable(var_name, &guid, attr, value_size,
				       value);

	return EFI_EXIT(ret);
}

static efi_status_t EFIAPI
handle_verification_result(struct GblEfiAvbProtocol *self,
			   const GblEfiAvbVerificationResult *result)
{
	EFI_ENTRY("%p, %p", self, result);
	if (!result)
		return EFI_EXIT(EFI_INVALID_PARAMETER);
	printf("GBL: AVB verification result: color %llu, digest %s\n",
	       result->color_flags, result->digest ? result->digest : "NULL");
	return EFI_EXIT(EFI_SUCCESS);
}

static efi_status_t EFIAPI write_lock_state(struct GblEfiAvbProtocol *self,
					    GblEfiAvbLockType type,
					    GblEfiAvbLockState state)
{
	int lock_type;

	EFI_ENTRY("%p, %d, %d", self, type, state);

	switch (type) {
	case GBL_EFI_AVB_LOCK_TYPE_DEVICE:
		lock_type = FASTBOOT_LOCK_TYPE_DEVICE;
		break;
	case GBL_EFI_AVB_LOCK_TYPE_CRITICAL:
		lock_type = FASTBOOT_LOCK_TYPE_CRITICAL;
		break;
	default:
		return EFI_EXIT(EFI_INVALID_PARAMETER);
	}

	bool new_locked = (state == GBL_EFI_AVB_LOCK_STATE_LOCKED);
	if (fastboot_write_lock_state(lock_type, new_locked) != 0)
		return EFI_EXIT(EFI_DEVICE_ERROR);

	return EFI_EXIT(EFI_SUCCESS);
}

static efi_status_t EFIAPI factory_data_reset(struct GblEfiAvbProtocol *self)
{
	EFI_ENTRY("%p", self);
	printf("GBL: Factory Data Reset is unsupported on this platform.\n");
	return EFI_EXIT(EFI_UNSUPPORTED);
}

static GblEfiAvbProtocol protocol = {
	.revision = GBL_EFI_AVB_PROTOCOL_REVISION,
	.read_partition_attributes = read_partition_attributes,
	.read_device_status = read_device_status,
	.validate_vbmeta_public_key = validate_vbmeta_public_key,
	.read_rollback_index = read_rollback_index,
	.write_rollback_index = write_rollback_index,
	.read_persistent_value = read_persistent_value,
	.write_persistent_value = write_persistent_value,
	.handle_verification_result = handle_verification_result,
	.write_lock_state = write_lock_state,
	.factory_data_reset = factory_data_reset,
};

efi_status_t efi_gbl_avb_register(void)
{
	efi_status_t ret = efi_add_protocol(efi_root, &guid, (void *)&protocol);
	if (ret != EFI_SUCCESS)
		printf("Cannot install GBL_EFI_AVB_PROTOCOL\n");
	return ret;
}