1 files changed, 518 insertions, 0 deletions

diff --git a/libpe/libpe_section.c b/libpe/libpe_section.c
new file mode 100644
index 0000000..7ff63fb
--- /dev/null
+++ b/libpe/libpe_section.c
@@ -0,0 +1,518 @@
+/*-
+ * Copyright (c) 2016 Kai Wang
+ * 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 AND CONTRIBUTORS ``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.
+ */
+
+#include <sys/param.h>
+#include <assert.h>
+#include <errno.h>
+#include <stdlib.h>
+#include <string.h>
+#include <unistd.h>
+
+#include "_libpe.h"
+
+ELFTC_VCSID("$Id: libpe_section.c 3312 2016-01-10 09:23:51Z kaiwang27 $");
+
+PE_Scn *
+libpe_alloc_scn(PE *pe)
+{
+	PE_Scn *ps;
+
+	if ((ps = calloc(1, sizeof(PE_Scn))) == NULL) {
+		errno = ENOMEM;
+		return (NULL);
+	}
+	STAILQ_INIT(&ps->ps_b);
+	ps->ps_pe = pe;
+
+	return (ps);
+}
+
+void
+libpe_release_scn(PE_Scn *ps)
+{
+	PE *pe;
+	PE_SecBuf *sb, *_sb;
+
+	assert(ps != NULL);
+
+	pe = ps->ps_pe;
+
+	STAILQ_REMOVE(&pe->pe_scn, ps, _PE_Scn, ps_next);
+
+	STAILQ_FOREACH_SAFE(sb, &ps->ps_b, sb_next, _sb)
+		libpe_release_buffer(sb);
+
+	free(ps);
+}
+
+static int
+cmp_scn(PE_Scn *a, PE_Scn *b)
+{
+
+	if (a->ps_sh.sh_addr < b->ps_sh.sh_addr)
+		return (-1);
+	else if (a->ps_sh.sh_addr == b->ps_sh.sh_addr)
+		return (0);
+	else
+		return (1);
+}
+
+static void
+sort_sections(PE *pe)
+{
+
+	if (STAILQ_EMPTY(&pe->pe_scn))
+		return;
+
+	/* Sort the list of Scn by RVA in ascending order. */
+	STAILQ_SORT(&pe->pe_scn, _PE_Scn, ps_next, cmp_scn);
+}
+
+int
+libpe_parse_section_headers(PE *pe)
+{
+	char tmp[sizeof(PE_SecHdr)], *hdr;
+	PE_Scn *ps;
+	PE_SecHdr *sh;
+	PE_CoffHdr *ch;
+	PE_DataDir *dd;
+	int found, i;
+
+	assert(pe->pe_ch != NULL);
+
+	for (i = 0; (uint16_t) i < pe->pe_ch->ch_nsec; i++) {
+		if (read(pe->pe_fd, tmp, sizeof(PE_SecHdr)) !=
+		    (ssize_t) sizeof(PE_SecHdr)) {
+			pe->pe_flags |= LIBPE_F_BAD_SEC_HEADER;
+			return (0);
+		}
+
+		if ((ps = libpe_alloc_scn(pe)) == NULL)
+			return (-1);
+		STAILQ_INSERT_TAIL(&pe->pe_scn, ps, ps_next);
+		ps->ps_ndx = ++pe->pe_nscn;	/* Setion index is 1-based */
+		sh = &ps->ps_sh;
+
+		/*
+		 * Note that the section name won't be NUL-terminated if
+		 * its length happens to be 8.
+		 */
+		memcpy(sh->sh_name, tmp, sizeof(sh->sh_name));
+		hdr = tmp + 8;
+		PE_READ32(hdr, sh->sh_virtsize);
+		PE_READ32(hdr, sh->sh_addr);
+		PE_READ32(hdr, sh->sh_rawsize);
+		PE_READ32(hdr, sh->sh_rawptr);
+		PE_READ32(hdr, sh->sh_relocptr);
+		PE_READ32(hdr, sh->sh_lineptr);
+		PE_READ16(hdr, sh->sh_nreloc);
+		PE_READ16(hdr, sh->sh_nline);
+		PE_READ32(hdr, sh->sh_char);
+	}
+
+	/*
+	 * For all the data directories that don't belong to any section,
+	 * we create pseudo sections for them to make layout easier.
+	 */
+	dd = pe->pe_dd;
+	if (dd != NULL && dd->dd_total > 0) {
+		for (i = 0; (uint32_t) i < pe->pe_dd->dd_total; i++) {
+			if (dd->dd_e[i].de_size == 0)
+				continue;
+			found = 0;
+			STAILQ_FOREACH(ps, &pe->pe_scn, ps_next) {
+				sh = &ps->ps_sh;
+				if (dd->dd_e[i].de_addr >= sh->sh_addr &&
+				    dd->dd_e[i].de_addr + dd->dd_e[i].de_size <=
+				    sh->sh_addr + sh->sh_virtsize) {
+					found = 1;
+					break;
+				}
+			}
+			if (found)
+				continue;
+
+			if ((ps = libpe_alloc_scn(pe)) == NULL)
+				return (-1);
+			STAILQ_INSERT_TAIL(&pe->pe_scn, ps, ps_next);
+			ps->ps_ndx = 0xFFFF0000U | i;
+			sh = &ps->ps_sh;
+			sh->sh_rawptr = dd->dd_e[i].de_addr; /* FIXME */
+			sh->sh_rawsize = dd->dd_e[i].de_size;
+		}
+	}
+
+	/*
+	 * Also consider the COFF symbol table as a pseudo section.
+	 */
+	ch = pe->pe_ch;
+	if (ch->ch_nsym > 0) {
+		if ((ps = libpe_alloc_scn(pe)) == NULL)
+			return (-1);
+		STAILQ_INSERT_TAIL(&pe->pe_scn, ps, ps_next);
+		ps->ps_ndx = 0xFFFFFFFFU;
+		sh = &ps->ps_sh;
+		sh->sh_rawptr = ch->ch_symptr;
+		sh->sh_rawsize = ch->ch_nsym * PE_SYM_ENTRY_SIZE;
+		pe->pe_nsym = ch->ch_nsym;
+	}
+
+	/* PE file headers initialization is complete if we reach here. */
+	return (0);
+}
+
+int
+libpe_load_section(PE *pe, PE_Scn *ps)
+{
+	PE_SecHdr *sh;
+	PE_SecBuf *sb;
+	size_t sz;
+	char tmp[4];
+
+	assert(pe != NULL && ps != NULL);
+	assert((ps->ps_flags & LIBPE_F_LOAD_SECTION) == 0);
+
+	sh = &ps->ps_sh;
+
+	/* Allocate a PE_SecBuf struct without buffer for empty sections. */
+	if (sh->sh_rawsize == 0) {
+		(void) libpe_alloc_buffer(ps, 0);
+		ps->ps_flags |= LIBPE_F_LOAD_SECTION;
+		return (0);
+	}
+
+	if ((pe->pe_flags & LIBPE_F_SPECIAL_FILE) == 0) {
+		if (lseek(pe->pe_fd, (off_t) sh->sh_rawptr, SEEK_SET) < 0) {
+			errno = EIO;
+			return (-1);
+		}
+	}
+
+	if ((sb = libpe_alloc_buffer(ps, sh->sh_rawsize)) == NULL)
+		return (-1);
+
+	if (read(pe->pe_fd, sb->sb_pb.pb_buf, sh->sh_rawsize) !=
+	    (ssize_t) sh->sh_rawsize) {
+		errno = EIO;
+		return (-1);
+	}
+
+	if (ps->ps_ndx == 0xFFFFFFFFU) {
+		/*
+		 * Index 0xFFFFFFFF indicates this section is a pseudo
+		 * section that contains the COFF symbol table. We should
+		 * read in the string table right after it.
+		 */
+		if (read(pe->pe_fd, tmp, sizeof(tmp)) !=
+		    (ssize_t) sizeof(tmp)) {
+			errno = EIO;
+			return (-1);
+		}
+		sz = le32dec(tmp);
+
+		/*
+		 * The minimum value for the size field is 4, which indicates
+		 * there is no string table.
+		 */
+		if (sz > 4) {
+			sz -= 4;
+			if ((sb = libpe_alloc_buffer(ps, sz)) == NULL)
+				return (-1);
+			if (read(pe->pe_fd, sb->sb_pb.pb_buf, sz) !=
+			    (ssize_t) sz) {
+				errno = EIO;
+				return (-1);
+			}
+		}
+	}
+
+	ps->ps_flags |= LIBPE_F_LOAD_SECTION;
+
+	return (0);
+}
+
+int
+libpe_load_all_sections(PE *pe)
+{
+	PE_Scn *ps;
+	PE_SecHdr *sh;
+	unsigned r, s;
+	off_t off;
+	char tmp[256];
+
+	/* Calculate the current offset into the file. */
+	off = 0;
+	if (pe->pe_dh != NULL)
+		off += pe->pe_dh->dh_lfanew + 4;
+	if (pe->pe_ch != NULL)
+		off += sizeof(PE_CoffHdr) + pe->pe_ch->ch_optsize;
+
+	STAILQ_FOREACH(ps, &pe->pe_scn, ps_next) {
+		if (ps->ps_flags & LIBPE_F_LOAD_SECTION)
+			continue;
+		sh = &ps->ps_sh;
+
+		/*
+		 * For special files, we consume the padding in between
+		 * and advance to the section offset.
+		 */
+		if (pe->pe_flags & LIBPE_F_SPECIAL_FILE) {
+			/* Can't go backwards. */
+			if (off > sh->sh_rawptr) {
+				errno = EIO;
+				return (-1);
+			}
+			if (off < sh->sh_rawptr) {
+				r = sh->sh_rawptr - off;
+				for (; r > 0; r -= s) {
+					s = r > sizeof(tmp) ? sizeof(tmp) : r;
+					if (read(pe->pe_fd, tmp, s) !=
+					    (ssize_t) s) {
+						errno = EIO;
+						return (-1);
+					}
+				}
+			}
+		}
+
+		/* Load the section content. */
+		if (libpe_load_section(pe, ps) < 0)
+			return (-1);
+	}
+
+	return (0);
+}
+
+int
+libpe_resync_sections(PE *pe, off_t off)
+{
+	PE_Scn *ps;
+	PE_SecHdr *sh;
+	size_t falign, nsec;
+
+	/* Firstly, sort all sections by their file offsets. */
+	sort_sections(pe);
+
+	/* Count the number of sections. */
+	nsec = 0;
+	STAILQ_FOREACH(ps, &pe->pe_scn, ps_next) {
+		if (ps->ps_flags & LIBPE_F_STRIP_SECTION)
+			continue;
+		if (ps->ps_ndx & 0xFFFF0000U)
+			continue;
+		nsec++;
+	}
+	pe->pe_nscn = nsec;
+
+	/*
+	 * Calculate the file offset for the first section. (`off' is
+	 * currently pointing to the COFF header.)
+	 */
+	off += sizeof(PE_CoffHdr);
+	if (pe->pe_ch != NULL && pe->pe_ch->ch_optsize > 0)
+		off += pe->pe_ch->ch_optsize;
+	else {
+		switch (pe->pe_obj) {
+		case PE_O_PE32:
+			off += PE_COFF_OPT_SIZE_32;
+			break;
+		case PE_O_PE32P:
+			off += PE_COFF_OPT_SIZE_32P;
+			break;
+		case PE_O_COFF:
+		default:
+			break;
+		}
+	}
+	off += nsec * sizeof(PE_SecHdr);
+
+	/*
+	 * Determine the file alignment for sections.
+	 */
+	if (pe->pe_oh != NULL && pe->pe_oh->oh_filealign > 0)
+		falign = pe->pe_oh->oh_filealign;
+	else {
+		/*
+		 * Use the default file alignment defined by the
+		 * PE/COFF specification.
+		 */
+		if (pe->pe_obj == PE_O_COFF)
+			falign = 4;
+		else
+			falign = 512;
+	}
+
+	/*
+	 * Step through each section (and pseduo section) and verify
+	 * alignment constraint and overlapping, make adjustment if need.
+	 */
+	pe->pe_rvamax = 0;
+	STAILQ_FOREACH(ps, &pe->pe_scn, ps_next) {
+		if (ps->ps_flags & LIBPE_F_STRIP_SECTION)
+			continue;
+
+		sh = &ps->ps_sh;
+
+		if (sh->sh_addr + sh->sh_virtsize > pe->pe_rvamax)
+			pe->pe_rvamax = sh->sh_addr + sh->sh_virtsize;
+
+		if (ps->ps_ndx & 0xFFFF0000U)
+			ps->ps_falign = 4;
+		else
+			ps->ps_falign = falign;
+
+		off = roundup(off, ps->ps_falign);
+
+		if (off != sh->sh_rawptr)
+			ps->ps_flags |= PE_F_DIRTY;
+
+		if (ps->ps_flags & PE_F_DIRTY) {
+			if ((ps->ps_flags & LIBPE_F_LOAD_SECTION) == 0) {
+				if (libpe_load_section(pe, ps) < 0)
+					return (-1);
+			}
+			sh->sh_rawsize = libpe_resync_buffers(ps);
+		}
+
+		/*
+		 * Sections only contains uninitialized data should set
+		 * PointerToRawData to zero according to the PE/COFF
+		 * specification.
+		 */
+		if (sh->sh_rawsize == 0)
+			sh->sh_rawptr = 0;
+		else
+			sh->sh_rawptr = off;
+
+		off += sh->sh_rawsize;
+	}
+
+	return (0);
+}
+
+off_t
+libpe_write_section_headers(PE *pe, off_t off)
+{
+	char tmp[sizeof(PE_SecHdr)], *hdr;
+	PE_Scn *ps;
+	PE_SecHdr *sh;
+
+	if (pe->pe_flags & LIBPE_F_BAD_SEC_HEADER || pe->pe_nscn == 0)
+		return (off);
+
+	if ((pe->pe_flags & LIBPE_F_DIRTY_SEC_HEADER) == 0) {
+		off += sizeof(PE_SecHdr) * pe->pe_ch->ch_nsec;
+		return (off);
+	}
+
+	STAILQ_FOREACH(ps, &pe->pe_scn, ps_next) {
+		if (ps->ps_flags & LIBPE_F_STRIP_SECTION)
+			continue;
+		if (ps->ps_ndx & 0xFFFF0000U)
+			continue;
+		if ((pe->pe_flags & LIBPE_F_DIRTY_SEC_HEADER) == 0 &&
+		    (ps->ps_flags & PE_F_DIRTY) == 0)
+			goto next_header;
+
+		sh = &ps->ps_sh;
+
+		memcpy(tmp, sh->sh_name, sizeof(sh->sh_name));
+		hdr = tmp + 8;
+		PE_WRITE32(hdr, sh->sh_virtsize);
+		PE_WRITE32(hdr, sh->sh_addr);
+		PE_WRITE32(hdr, sh->sh_rawsize);
+		PE_WRITE32(hdr, sh->sh_rawptr);
+		PE_WRITE32(hdr, sh->sh_relocptr);
+		PE_WRITE32(hdr, sh->sh_lineptr);
+		PE_WRITE16(hdr, sh->sh_nreloc);
+		PE_WRITE16(hdr, sh->sh_nline);
+		PE_WRITE32(hdr, sh->sh_char);
+
+		if (write(pe->pe_fd, tmp, sizeof(PE_SecHdr)) !=
+		    (ssize_t) sizeof(PE_SecHdr)) {
+			errno = EIO;
+			return (-1);
+		}
+
+	next_header:
+		off += sizeof(PE_SecHdr);
+	}
+
+	return (off);
+}
+
+off_t
+libpe_write_sections(PE *pe, off_t off)
+{
+	PE_Scn *ps;
+	PE_SecHdr *sh;
+
+	if (pe->pe_flags & LIBPE_F_BAD_SEC_HEADER)
+		return (off);
+
+	STAILQ_FOREACH(ps, &pe->pe_scn, ps_next) {
+		sh = &ps->ps_sh;
+
+		if (ps->ps_flags & LIBPE_F_STRIP_SECTION)
+			continue;
+
+		/* Skip empty sections. */
+		if (sh->sh_rawptr == 0 || sh->sh_rawsize == 0)
+			continue;
+
+		/*
+		 * Padding between sections. (padding always written
+		 * in case the the section headers or sections are
+		 * moved or shrinked.)
+		 */
+		assert(off <= sh->sh_rawptr);
+		if (off < sh->sh_rawptr)
+			libpe_pad(pe, sh->sh_rawptr - off);
+
+		if ((ps->ps_flags & PE_F_DIRTY) == 0) {
+			assert((pe->pe_flags & LIBPE_F_SPECIAL_FILE) == 0);
+			if (lseek(pe->pe_fd,
+			    (off_t) (sh->sh_rawptr + sh->sh_rawsize),
+			    SEEK_SET) < 0) {
+				errno = EIO;
+				return (-1);
+			}
+			off = sh->sh_rawptr + sh->sh_rawsize;
+			continue;
+		}
+
+		off = sh->sh_rawptr;
+
+		if (libpe_write_buffers(ps) < 0)
+			return (-1);
+
+		off += sh->sh_rawsize;
+
+		ps->ps_flags &= ~PE_F_DIRTY;
+	}
+
+	return (off);
+}