src/drivers/blockdev/uefi.c - third_party/depthcharge - Git at Google

 /*
  * Copyright 2016 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.
  *
  * You should have received a copy of the GNU General Public License
  * along with this program; if not, write to the Free Software
  * Foundation, Inc., 59 Temple Place, Suite 330, Boston,
  * MA 02111-1307 USA
  */

 #include <assert.h>
 #include <stddef.h>
 #include <stdio.h>
 #include <stdlib.h>

 #include "base/algorithm.h"
 #include "base/xalloc.h"
 #include "base/container_of.h"
 #include "drivers/blockdev/uefi.h"
 #include "uefi/edk/Protocol/BlockIo.h"
 #include "uefi/uefi.h"

 static EFI_GUID block_io_protocol_guid = EFI_BLOCK_IO_PROTOCOL_GUID;

 static lba_t uefi_blockdev_read(BlockDevOps *me, lba_t start, lba_t count,
 				void *buffer)
 {
 	UefiBlockDev *ubdev = container_of(me, UefiBlockDev, dev.ops);

 	const size_t block_size = ubdev->dev.block_size;
 	size_t buffer_size = block_size * count;
 	size_t chunk_size = MIN(64 * 1024, buffer_size);
 	// Round chunk_size up to the nearest multiple of block_size;
 	chunk_size = ((chunk_size + block_size - 1) / block_size) * block_size;

 	void *read_buf;
 	if (ubdev->io_align > 1)
 		read_buf = memalign(ubdev->io_align, chunk_size);
 	else
 		read_buf = malloc(chunk_size);

 	if (!read_buf) {
 		printf("Failed to allocate read buffer.\n");
 		return 0;
 	}

 	while (buffer_size) {
 		EFI_STATUS status = ubdev->bio->ReadBlocks(
 			ubdev->bio, ubdev->media_id, start,
 			chunk_size, read_buf);
 		if (status != EFI_SUCCESS) {
 			printf("Failed to read blocks from UEFI device.\n");
 			free(read_buf);
 			return 0;
 		}
 		memcpy(buffer, read_buf, chunk_size);

 		buffer = (uint8_t *)buffer + chunk_size;
 		start += chunk_size / block_size;
 		buffer_size -= chunk_size;
 		chunk_size = MIN(chunk_size, buffer_size);
 	}

 	free(read_buf);
 	return count;
 }

 static lba_t uefi_blockdev_write(BlockDevOps *me, lba_t start, lba_t count,
 				 const void *buffer)
 {
 	UefiBlockDev *ubdev = container_of(me, UefiBlockDev, dev.ops);

 	const size_t block_size = ubdev->dev.block_size;
 	size_t buffer_size = block_size * count;
 	size_t chunk_size = MIN(64 * 1024, buffer_size);
 	// Round chunk_size up to the nearest multiple of block_size;
 	chunk_size = ((chunk_size + block_size - 1) / block_size) * block_size;

 	void *write_buf;
 	if (ubdev->io_align > 1)
 		write_buf = memalign(ubdev->io_align, chunk_size);
 	else
 		write_buf = malloc(chunk_size);
 	if (!write_buf) {
 		printf("Failed to allocate write buffer.\n");
 		return 0;
 	}

 	while (buffer_size) {
 		memcpy(write_buf, buffer, chunk_size);
 		EFI_STATUS status = ubdev->bio->WriteBlocks(
 			ubdev->bio, ubdev->media_id, start,
 			chunk_size, write_buf);
 		if (status != EFI_SUCCESS) {
 			printf("Failed to write blocks to UEFI device.\n");
 			free(write_buf);
 			return 0;
 		}

 		buffer = (uint8_t *)buffer + chunk_size;
 		start += chunk_size / block_size;
 		buffer_size -= chunk_size;
 		chunk_size = MIN(chunk_size, buffer_size);
 	}

 	free(write_buf);

 	if (ubdev->write_caching) {
 		EFI_STATUS status = ubdev->bio->FlushBlocks(ubdev->bio);
 		if (status != EFI_SUCCESS) {
 			printf("Failed to flush blocks to UEFI device.\n");
 			return 0;
 		}
 	}

 	return count;
 }

 static int uefi_blockdev_ctrlr_update(struct BlockDevCtrlrOps *me)
 {
 	UefiBlockDevCtrlr *ctrlr =
 		container_of(me, UefiBlockDevCtrlr, ctrlr.ops);

 	// Detach the existing block devices. They'll be re-added if they're
 	// still present.
 	UefiBlockDev *ubdev;
 	UefiBlockDev *last_ubdev = NULL;
 	list_for_each(ubdev, ctrlr->uefi_blockdev_list, uefi_blockdev_list) {
 		list_remove(&last_ubdev->uefi_blockdev_list);
 		list_remove(&last_ubdev->dev.list_node);
 		free(last_ubdev);
 		last_ubdev = ubdev;
 	}
 	if (last_ubdev) {
 		list_remove(&last_ubdev->uefi_blockdev_list);
 		list_remove(&last_ubdev->dev.list_node);
 		free(last_ubdev);
 	}

 	EFI_SYSTEM_TABLE *st = uefi_system_table_ptr();
 	if (!st)
 		return 1;
 	EFI_BOOT_SERVICES *bs = st->BootServices;

 	UINTN buf_size = 0;
 	EFI_HANDLE dummy;
 	EFI_STATUS status = bs->LocateHandle(
 		ByProtocol, &block_io_protocol_guid, NULL, &buf_size, &dummy);
 	// If there are no handles, we're already done.
 	if (status == EFI_SUCCESS)
 		return 0;
 	if (status != EFI_BUFFER_TOO_SMALL) {
 		printf("Error retrieving block io handles.\n");
 		return 1;
 	}

 	EFI_HANDLE *handles = xmalloc(buf_size);

 	status = bs->LocateHandle(
 		ByProtocol, &block_io_protocol_guid, NULL, &buf_size, handles);
 	if (status != EFI_SUCCESS) {
 		printf("Failed to retrieve block io handles.\n");
 		free(handles);
 		return 1;
 	}

 	int handle_count = buf_size / sizeof(dummy);

 	for (int i = 0; i < handle_count; i++) {
 		EFI_BLOCK_IO_PROTOCOL *bio;
 		status = bs->HandleProtocol(
 			handles[i], &block_io_protocol_guid, (void **)&bio);
 		if (status != EFI_SUCCESS) {
 			printf("Failed to retrieve block io protocol.\n");
 			return 1;
 		}

 		if (!bio->Media->MediaPresent ||
 		    bio->Media->LogicalPartition ||
 		    !!bio->Media->RemovableMedia != !!ctrlr->removable) {
 			continue;
 		}

 		ubdev = xzalloc(sizeof(*ubdev));

 		ubdev->bio = bio;
 		ubdev->media_id = bio->Media->MediaId;
 		ubdev->io_align = bio->Media->IoAlign;
 		ubdev->write_caching = bio->Media->WriteCaching;

 		BlockDev *dev = &ubdev->dev;
 		dev->ops.read = &uefi_blockdev_read;
 		dev->ops.write = &uefi_blockdev_write;

 		dev->name = "UEFI block device.\n";
 		dev->removable = bio->Media->RemovableMedia;
 		dev->block_size = bio->Media->BlockSize;
 		dev->block_count = bio->Media->LastBlock + 1;

 		list_insert_after(&ubdev->uefi_blockdev_list,
 				  &ctrlr->uefi_blockdev_list);
 		if (bio->Media->RemovableMedia) {
 			list_insert_after(&dev->list_node,
 					  &removable_block_devices);
 		} else {
 			list_insert_after(&dev->list_node,
 					  &fixed_block_devices);
 		}
 	}

 	return 0;
 }

 UefiBlockDevCtrlr *new_uefi_blockdev_ctrlr(int removable)
 {
 	UefiBlockDevCtrlr *ctrlr = xzalloc(sizeof(*ctrlr));

 	ctrlr->ctrlr.ops.update = &uefi_blockdev_ctrlr_update;
 	ctrlr->ctrlr.need_update = 1;

 	ctrlr->removable = removable;

 	return ctrlr;
 }
	/*
	* Copyright 2016 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.
	*
	* You should have received a copy of the GNU General Public License
	* along with this program; if not, write to the Free Software
	* Foundation, Inc., 59 Temple Place, Suite 330, Boston,
	* MA 02111-1307 USA
	*/

	#include <assert.h>
	#include <stddef.h>
	#include <stdio.h>
	#include <stdlib.h>

	#include "base/algorithm.h"
	#include "base/xalloc.h"
	#include "base/container_of.h"
	#include "drivers/blockdev/uefi.h"
	#include "uefi/edk/Protocol/BlockIo.h"
	#include "uefi/uefi.h"

	static EFI_GUID block_io_protocol_guid = EFI_BLOCK_IO_PROTOCOL_GUID;

	static lba_t uefi_blockdev_read(BlockDevOps *me, lba_t start, lba_t count,
	void *buffer)
	{
	UefiBlockDev *ubdev = container_of(me, UefiBlockDev, dev.ops);

	const size_t block_size = ubdev->dev.block_size;
	size_t buffer_size = block_size * count;
	size_t chunk_size = MIN(64 * 1024, buffer_size);
	// Round chunk_size up to the nearest multiple of block_size;
	chunk_size = ((chunk_size + block_size - 1) / block_size) * block_size;

	void *read_buf;
	if (ubdev->io_align > 1)
	read_buf = memalign(ubdev->io_align, chunk_size);
	else
	read_buf = malloc(chunk_size);

	if (!read_buf) {
	printf("Failed to allocate read buffer.\n");
	return 0;
	}

	while (buffer_size) {
	EFI_STATUS status = ubdev->bio->ReadBlocks(
	ubdev->bio, ubdev->media_id, start,
	chunk_size, read_buf);
	if (status != EFI_SUCCESS) {
	printf("Failed to read blocks from UEFI device.\n");
	free(read_buf);
	return 0;
	}
	memcpy(buffer, read_buf, chunk_size);

	buffer = (uint8_t *)buffer + chunk_size;
	start += chunk_size / block_size;
	buffer_size -= chunk_size;
	chunk_size = MIN(chunk_size, buffer_size);
	}

	free(read_buf);
	return count;
	}

	static lba_t uefi_blockdev_write(BlockDevOps *me, lba_t start, lba_t count,
	const void *buffer)
	{
	UefiBlockDev *ubdev = container_of(me, UefiBlockDev, dev.ops);

	const size_t block_size = ubdev->dev.block_size;
	size_t buffer_size = block_size * count;
	size_t chunk_size = MIN(64 * 1024, buffer_size);
	// Round chunk_size up to the nearest multiple of block_size;
	chunk_size = ((chunk_size + block_size - 1) / block_size) * block_size;

	void *write_buf;
	if (ubdev->io_align > 1)
	write_buf = memalign(ubdev->io_align, chunk_size);
	else
	write_buf = malloc(chunk_size);
	if (!write_buf) {
	printf("Failed to allocate write buffer.\n");
	return 0;
	}

	while (buffer_size) {
	memcpy(write_buf, buffer, chunk_size);
	EFI_STATUS status = ubdev->bio->WriteBlocks(
	ubdev->bio, ubdev->media_id, start,
	chunk_size, write_buf);
	if (status != EFI_SUCCESS) {
	printf("Failed to write blocks to UEFI device.\n");
	free(write_buf);
	return 0;
	}

	buffer = (uint8_t *)buffer + chunk_size;
	start += chunk_size / block_size;
	buffer_size -= chunk_size;
	chunk_size = MIN(chunk_size, buffer_size);
	}

	free(write_buf);

	if (ubdev->write_caching) {
	EFI_STATUS status = ubdev->bio->FlushBlocks(ubdev->bio);
	if (status != EFI_SUCCESS) {
	printf("Failed to flush blocks to UEFI device.\n");
	return 0;
	}
	}

	return count;
	}

	static int uefi_blockdev_ctrlr_update(struct BlockDevCtrlrOps *me)
	{
	UefiBlockDevCtrlr *ctrlr =
	container_of(me, UefiBlockDevCtrlr, ctrlr.ops);

	// Detach the existing block devices. They'll be re-added if they're
	// still present.
	UefiBlockDev *ubdev;
	UefiBlockDev *last_ubdev = NULL;
	list_for_each(ubdev, ctrlr->uefi_blockdev_list, uefi_blockdev_list) {
	list_remove(&last_ubdev->uefi_blockdev_list);
	list_remove(&last_ubdev->dev.list_node);
	free(last_ubdev);
	last_ubdev = ubdev;
	}
	if (last_ubdev) {
	list_remove(&last_ubdev->uefi_blockdev_list);
	list_remove(&last_ubdev->dev.list_node);
	free(last_ubdev);
	}

	EFI_SYSTEM_TABLE *st = uefi_system_table_ptr();
	if (!st)
	return 1;
	EFI_BOOT_SERVICES *bs = st->BootServices;

	UINTN buf_size = 0;
	EFI_HANDLE dummy;
	EFI_STATUS status = bs->LocateHandle(
	ByProtocol, &block_io_protocol_guid, NULL, &buf_size, &dummy);
	// If there are no handles, we're already done.
	if (status == EFI_SUCCESS)
	return 0;
	if (status != EFI_BUFFER_TOO_SMALL) {
	printf("Error retrieving block io handles.\n");
	return 1;
	}

	EFI_HANDLE *handles = xmalloc(buf_size);

	status = bs->LocateHandle(
	ByProtocol, &block_io_protocol_guid, NULL, &buf_size, handles);
	if (status != EFI_SUCCESS) {
	printf("Failed to retrieve block io handles.\n");
	free(handles);
	return 1;
	}

	int handle_count = buf_size / sizeof(dummy);

	for (int i = 0; i < handle_count; i++) {
	EFI_BLOCK_IO_PROTOCOL *bio;
	status = bs->HandleProtocol(
	handles[i], &block_io_protocol_guid, (void **)&bio);
	if (status != EFI_SUCCESS) {
	printf("Failed to retrieve block io protocol.\n");
	return 1;
	}

	if (!bio->Media->MediaPresent \|\|
	bio->Media->LogicalPartition \|\|
	!!bio->Media->RemovableMedia != !!ctrlr->removable) {
	continue;
	}

	ubdev = xzalloc(sizeof(*ubdev));

	ubdev->bio = bio;
	ubdev->media_id = bio->Media->MediaId;
	ubdev->io_align = bio->Media->IoAlign;
	ubdev->write_caching = bio->Media->WriteCaching;

	BlockDev *dev = &ubdev->dev;
	dev->ops.read = &uefi_blockdev_read;
	dev->ops.write = &uefi_blockdev_write;

	dev->name = "UEFI block device.\n";
	dev->removable = bio->Media->RemovableMedia;
	dev->block_size = bio->Media->BlockSize;
	dev->block_count = bio->Media->LastBlock + 1;

	list_insert_after(&ubdev->uefi_blockdev_list,
	&ctrlr->uefi_blockdev_list);
	if (bio->Media->RemovableMedia) {
	list_insert_after(&dev->list_node,
	&removable_block_devices);
	} else {
	list_insert_after(&dev->list_node,
	&fixed_block_devices);
	}
	}

	return 0;
	}

	UefiBlockDevCtrlr *new_uefi_blockdev_ctrlr(int removable)
	{
	UefiBlockDevCtrlr ctrlr = xzalloc(sizeof(ctrlr));

	ctrlr->ctrlr.ops.update = &uefi_blockdev_ctrlr_update;
	ctrlr->ctrlr.need_update = 1;

	ctrlr->removable = removable;

	return ctrlr;
	}