/*-
 * Copyright (c) 2001 by Thomas Moestl <tmm@FreeBSD.org>.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 * USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 * $FreeBSD$
 */

#include <sys/types.h>
#include <sys/fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>

#include <sys/elf_common.h>
#include <sys/elf32.h>
#include <sys/elf64.h>

#include <unistd.h>
#include <stdio.h>
#include <zlib.h>
/* Include after zlib.h to avoid a spurious warning. */
#include <err.h>


/* Section to put the check sums into. */
#define	CKSUM_SECT	"cksum"

/* XXX: This should be exported... */
struct crcinfo {
	u_int32_t	ci_start;
	u_int32_t	ci_size;
	u_int32_t	ci_sum;
};

/*
 * Endianess conversion and type size handling to allow elfcksum to be used in
 * cross builds.
 */
int data, class;

static __inline u_int16_t
swap16(u_int16_t x)
{
	return ((x >> 8) | ((x << 8) & 0xff00U));
}

static __inline u_int32_t
swap32(u_int32_t x)
{
	return ((x >> 24) | ((x >> 8) & 0xff00U) | ((x << 8) & 0xff0000U) |
	    ((x << 24) & 0xff000000U));
}

static __inline u_int64_t
swap64(u_int64_t x)
{
	return ((x >> 56) | ((x >> 40) & 0xff00UL) | ((x >> 24) & 0xff0000UL) |
	    ((x >> 8) & 0xff000000UL) | ((x << 8) & 0xff00000000UL) |
	    ((x << 24) & 0xff0000000000UL) | ((x << 40) & 0xff000000000000UL) |
	    ((x << 56)));
}

#if BYTE_ORDER == LITTLE_ENDIAN
#define	ETOH16(x)	((data) == ELFDATA2MSB ? swap16((x)) : (x))
#define	ETOH32(x)	((data) == ELFDATA2MSB ? swap32((x)) : (x))
#define	ETOH64(x)	((data) == ELFDATA2MSB ? swap64((x)) : (x))
#else
#define	ETOH16(x)	((data) == ELFDATA2LSB ? swap16((x)) : (x))
#define	ETOH32(x)	((data) == ELFDATA2LSB ? swap32((x)) : (x))
#define	ETOH64(x)	((data) == ELFDATA2LSB ? swap64((x)) : (x))
#endif
/* The conversions are self-inverse. */
#define	HTOE16(x)	ETOH16((x))
#define	HTOE32(x)	ETOH32((x))
#define	HTOE64(x)	ETOH64((x))

union elf_ehdr {
	Elf32_Ehdr	c32;
	Elf64_Ehdr	c64;
};

union elf_phdr {
	Elf32_Phdr	c32;
	Elf64_Phdr	c64;
};

union elf_shdr {
	Elf32_Shdr	c32;
	Elf64_Shdr	c64;
};

#define	ESIZEOF(p)							\
	(class == ELFCLASS32 ? sizeof((p)->c32) : sizeof((p)->c64))
#define	EREAD_16(ptr, field) ({						\
	u_int16_t tmp = (class == ELFCLASS32 ? (ptr)->c32.field :	\
	    (ptr)->c64.field);						\
	ETOH16(tmp);							\
})
#define	EREAD_32(ptr, field) ({						\
	u_int32_t tmp = (class == ELFCLASS32 ? (ptr)->c32.field :	\
	    (ptr)->c64.field);						\
	ETOH32(tmp);							\
})
#define	EREAD_WORD(ptr, field)						\
	(class == ELFCLASS32 ? ETOH32((ptr)->c32.field) :		\
	    ETOH64((ptr)->c64.field))
#define	EREAD_HALF(ptr, field)						\
	(class == ELFCLASS32 ? ETOH16((ptr)->c32.field) :		\
	    ETOH32((ptr)->c64.field))
#define	EREAD_ADDR(ptr, field)	EREAD_WORD((ptr), field)
#define	EREAD_OFF(ptr, field)	EREAD_WORD((ptr), field)
#define	EREAD_SIZE(ptr, field)	EREAD_WORD((ptr), field)

/*
 * Check that a range specified as offset and size is present in the ELF file,
 * and try to treat overruns correctly.
 */
#define	EBADRANGE(totsz, rstart, rsz) 					\
	((rstart) >= (totsz) || (rsz) > (totsz) - (rstart))

static char *
egetstr(char *tab, size_t sz, off_t offs)
{
	if (offs >= sz)
		errx(3, "bad string table offset");
	return (tab + offs);
}

static void
usage(void)
{
	fprintf(stderr, "usage: elfcksum [-hv] file\n");
	exit(1);
}

/*
 * This program allows to add CRC32 check sums for the PT_LOAD segments to ELF
 * executables. These can currenlty used by the loader to verify that a kernel
 * was not damaged while it was loaded.
 * To allow for this to work, the ELF input file must contain a 'cksum' section;
 * it can e.g. be generated by compiling and linking in the following assembler
 * code:
 *	.section cksum, "", @progbits
 *	.rept <num>
 *	.word 0, 0, 0
 *	.endr
 * Where <num> is replaced by a number equal or larger than number of PT_LOAD
 * segments in the ELF file. 2 is usually sufficient, 4 is used in the kernel
 * build for a safety margin.
 * elfcksum will then replace the contents of this sections with information
 * about the checksum (the three successive 32-bit integers will be replaced
 * by start and size of the segment in the file and the CRC32 sum).
 * elfcksum can process big- and little-endian 32- and 64-bit ELF files.
 */
int
main(int argc, char *argv[])
{
	struct stat st;
	union elf_ehdr *ehdr;
	union elf_phdr *phdr;
	union elf_shdr *shdr;
	struct crcinfo *ci;
	unsigned char *offs, *shstrt;
	size_t len, cksz, fsz, shstrtsz;
	off_t phoff, shoff, ckoff, shstrtoff, o;
	u_int16_t stridx, phesz, shesz, nph, nsh;
	u_int32_t crc;
	int ch, fd, sect, i, ncrc;
	int verbose = 0;

	while ((ch = getopt(argc, argv, "hv")) != -1) {
		switch (ch) {
		case 'h':
			usage();
		case 'v':
			verbose = 1;
			break;
		default:
			usage();
		}
	}
	if (argc - optind != 1)
		usage();
	if ((fd = open(argv[optind], O_RDWR)) == -1)
		err(2, "open");
	if (fstat(fd, &st) == -1)
		err(2, "fstat");
	len = st.st_size;
	if (len < ESIZEOF(ehdr))
		errx(3, "truncated elf file");
	offs = mmap(NULL, len, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
	if (offs == MAP_FAILED)
		err(2, "mmap");

	/* Verify the header, and find out the type */
	if (offs[EI_MAG0] != ELFMAG0 || offs[EI_MAG1] != ELFMAG1 ||
	    offs[EI_MAG2] != ELFMAG2 || offs[EI_MAG3] != ELFMAG3 ||
	    offs[EI_VERSION] != EV_CURRENT)
		errx(3, "bad ELF signature or unknown version");
	class = offs[EI_CLASS];
	if (class != ELFCLASS32 && class != ELFCLASS64)
		errx(3, "unknown class");
	data = offs[EI_DATA];
	if (class != ELFDATA2MSB && class != ELFDATA2LSB)
		errx(3, "unknown data encoding");

	/* Get the header information needed. */
	ehdr = (union elf_ehdr *)offs;
	if (EREAD_16(ehdr, e_type) != ET_EXEC)
		errx(3, "not an executable");
	phoff = EREAD_OFF(ehdr, e_phoff);
	nph = EREAD_16(ehdr, e_phnum);
	phesz = EREAD_16(ehdr, e_phentsize);
	if (phoff == 0 || nph == 0 || phesz != ESIZEOF(phdr))
		errx(3, "program header not present or of unknown format");
	if (EBADRANGE(len, phoff, nph * phesz))
		errx(3, "bad program header table offset/size");
	phdr = (union elf_phdr *)(offs + phoff);
	shoff = EREAD_OFF(ehdr, e_shoff);
	nsh = EREAD_16(ehdr, e_shnum);
	shesz = EREAD_16(ehdr, e_shentsize);
	if (shoff == 0 || nsh == 0 || shesz != ESIZEOF(shdr))
		errx(3, "section header not present or of unknown format");
	if (EBADRANGE(len, shoff, nsh * shesz))
		errx(3, "bad section header table offset/size");
	shdr = (union elf_shdr *)(offs + EREAD_OFF(ehdr, e_shoff));

	/* Get the section name string table. */
	stridx = EREAD_16(ehdr, e_shstrndx);
	if (stridx == SHN_UNDEF || stridx >= nsh)
		errx(3, "section name string table missing/corrupt");
	if (EREAD_32(shdr + stridx, sh_type) != SHT_STRTAB)
		errx(3, "corrupt section name string table");
	shstrtoff = EREAD_OFF(shdr + stridx, sh_offset);
	shstrtsz = EREAD_SIZE(shdr + stridx, sh_size);
	if (EBADRANGE(len, shstrtoff, shstrtsz))
		errx(3, "bad string table offset/size");
	shstrt = offs + shstrtoff;
	if (shstrt[shstrtsz - 1] != '\0')
		errx(3, "string table not terminated");

	/* Try to find the section to put the checksums into. */
	for (sect = 0; sect < nsh; sect++) {
		if (strcmp(egetstr(shstrt, shstrtsz,
		    EREAD_32(shdr + sect, sh_name)),  CKSUM_SECT) == 0) {
			cksz = EREAD_SIZE(shdr + sect, sh_size);
			ckoff = EREAD_OFF(shdr + sect, sh_offset);
			break;
		}
	}
	if (sect >= nsh)
		errx(4, "checksum section not found");
	if (ckoff == 0)
		errx(4, "bad checksum section");
	ci = (struct crcinfo *)(uintptr_t)(offs + ckoff);
	ncrc = cksz / sizeof(struct crcinfo);

	/* Now, go through the program header and checksum each PT_LOAD seg. */
	for (i = 0; i < nph; i++) {
		if (EREAD_32(phdr + i, p_type) == PT_LOAD) {
			if (ncrc == 0)
				errx(4, "not enough space in checksum section");
			o = EREAD_OFF(phdr + i, p_offset);
			fsz = EREAD_OFF(phdr + i, p_filesz);
			if (EBADRANGE(len, o, fsz))
				errx(3, "bad segment offset/size");
			crc = crc32(0, offs + o, fsz);
			ci->ci_sum = HTOE32(crc);
			ci->ci_start = HTOE32(o);
			ci->ci_size = HTOE32(fsz);
			if (verbose) {
				printf("segment %d, start %lu, size %lu: "
				    "crc32 %lu\n", i, (unsigned long)o,
				    (unsigned long)fsz, (unsigned long)crc);
			}
			ci++;
			ncrc--;
		}
	}

	munmap(offs, len);
	close(fd);
	return (0);
}