OSDev.org

Code: Select all

/* PE Executable Multiboot Header Repair Tool */
/* This tool is release as-is to be free of use for anyone. Myself, the maker,
*  am not liable for any damage or harm that may result from its use.
*  This library includes a component from the multiboot header, part-verbatim.
*  The license to the multiboot header file is included below.
*/

/* multiboot.h - Multiboot header file. */
/* Copyright (C) 1999,2003,2007,2008,2009  Free Software Foundation, Inc.
*
*  Permission is hereby granted, free of charge, to any person obtaining a copy
*  of this software and associated documentation files (the "Software"), to
*  deal in the Software without restriction, including without limitation the
*  rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
*  sell copies of the Software, and to permit persons to whom the Software is
*  furnished to do so, subject to the following conditions:
*
*  The above copyright notice and this permission notice shall be included in
*  all copies or substantial portions of the Software.
*
*  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
*  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
*  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL ANY
*  DEVELOPER OR DISTRIBUTOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
*  WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
*  IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

struct multiboot_header
{
	/* Must be MULTIBOOT_MAGIC - see above. */
	unsigned int magic;

	/* Feature flags. */
	unsigned int flags;

	/* The above fields plus this one must equal 0 mod 2^32. */
	unsigned int checksum;

	/* These are only valid if MULTIBOOT_AOUT_KLUDGE is set. */
	unsigned int header_addr;
	unsigned int load_addr;
	unsigned int load_end_addr;
	unsigned int bss_end_addr;
	unsigned int entry_addr;

	/* These are only valid if MULTIBOOT_VIDEO_MODE is set. */
	unsigned int mode_type;
	unsigned int width;
	unsigned int height;
	unsigned int depth;
};

#define NDEBUG

#include <stdio.h>

/* This library attempts to repair a multiboot header in a PE file
*  that has become corrupted, for instance by linking too many OBJ
*  files into your kernel. This causes your header to shift (even
*  though it should not) and becomes unreadable. This scans for
*  your header, repairs its header_addr and entry_addr values,
*  and rewrites it. This application is set up as such so that it
*  is easier to integrate as a post-build step in Visual C++.
*
*  USAGE:  PE_REWRITER file
*
*  ERROR CODES:
*     1 - Could not open given file.
*     2 - No multiboot header was found.
*     3 - Invalid number of arguments
*/

const bool REWRITE_LOAD_END_ADDR	= true;
const bool REWRITE_LOAD_ADDR		= true;
const bool REWRITE_ENTRY_ADDR		= true;
const bool REWRITE_HEADER_ADDR		= true;

const unsigned int IMAGE_OFFSET	= 0x00100000;
const unsigned int MB_HEADER	= 0x1BADB002;
const unsigned int TEXT_OFFSET  = 0x00001000;

int __cdecl main (int argc, char *argv[])
{
	/* Test arguments for validity */
	if (argc < 2)
	{
		printf("Usage: rewriter file.\n");
		return 3;
	}

	if ((REWRITE_LOAD_END_ADDR || REWRITE_LOAD_ADDR || REWRITE_ENTRY_ADDR || REWRITE_HEADER_ADDR) == false)
	{
		printf("Nothing to do!\n");
		return 0;
	}

	/* Load the file */
	FILE *kernel = fopen(argv[1], "r+b");
	if (kernel == NULL)
	{
		printf("Couldn't open Kernel File.\n");
		return 1;
	}
	fseek(kernel, 0, SEEK_SET);

	/* Stupid way of figuring out where the code actually ends.
	   I'm using 64-bit integers to scan in-case someone has a
	   64-bit variable with null values near the end of the code.
	*/
	unsigned int kernel_length;
	if (REWRITE_LOAD_END_ADDR)
	{
		unsigned long long null_test;
		fseek(kernel, 0, SEEK_END);
		do {
			if (fseek(kernel, -8, SEEK_CUR) != 0 || fread(&null_test, 8, 1, kernel) != 1 || fseek(kernel, -8, SEEK_CUR) != 0)
			{
				printf("This file has no multiboot header. Aborting.\n");
				return 2;
			}
		} while (null_test == 0);
		kernel_length = ftell(kernel) + 8;
		fseek(kernel, 0, SEEK_SET);
	}

	/* Scan through the file to find the multiboot magic number.
	*  This is not the fastest way, but it works.
	*/
	unsigned int test_against;
	do {
		if (fread(&test_against, 4, 1, kernel) != 1 || fseek(kernel, -3, SEEK_CUR) != 0)
		{
			printf("This file has no multiboot header. Aborting.\n");
			return 2;
		}
	} while (test_against != MB_HEADER);
	fseek(kernel, -1, SEEK_CUR); 
	unsigned int header_location = ftell(kernel);

	/* Now we load the header into memory */

	multiboot_header header;
	fread(&header, sizeof(multiboot_header), 1, kernel);

	/* This is where we edit the header */

	if (REWRITE_HEADER_ADDR)
		header.header_addr		= IMAGE_OFFSET + header_location;
	if (REWRITE_ENTRY_ADDR)
		header.entry_addr		= IMAGE_OFFSET + header_location + sizeof(multiboot_header);
	if (REWRITE_LOAD_ADDR)
		header.load_addr		= IMAGE_OFFSET + TEXT_OFFSET;
	if (REWRITE_LOAD_END_ADDR)
		header.load_end_addr	= IMAGE_OFFSET + kernel_length;

	/* Now we write it back where it was. */

	fseek(kernel, header_location, SEEK_SET);
	fwrite(&header, sizeof(multiboot_header), 1, kernel);

	/* Clean up */

	fclose(kernel);
	return 0;
}