1/*-
2 * Copyright (c) 2016 Kai Wang
3 * All rights reserved.
4 *
5 * Redistribution and use in source and binary forms, with or without
6 * modification, are permitted provided that the following conditions
7 * are met:
8 * 1. Redistributions of source code must retain the above copyright
9 *    notice, this list of conditions and the following disclaimer.
10 * 2. Redistributions in binary form must reproduce the above copyright
11 *    notice, this list of conditions and the following disclaimer in the
12 *    documentation and/or other materials provided with the distribution.
13 *
14 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
15 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
16 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
17 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
18 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
19 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
20 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
21 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
22 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
23 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
24 * SUCH DAMAGE.
25 */
26
27#include <sys/param.h>
28#include <assert.h>
29#include <errno.h>
30#include <stdlib.h>
31#include <string.h>
32#include <unistd.h>
33
34#include "_libpe.h"
35
36ELFTC_VCSID("$Id: libpe_section.c 3446 2016-05-03 01:31:17Z emaste $");
37
38PE_Scn *
39libpe_alloc_scn(PE *pe)
40{
41	PE_Scn *ps;
42
43	if ((ps = calloc(1, sizeof(PE_Scn))) == NULL) {
44		errno = ENOMEM;
45		return (NULL);
46	}
47	STAILQ_INIT(&ps->ps_b);
48	ps->ps_pe = pe;
49
50	return (ps);
51}
52
53void
54libpe_release_scn(PE_Scn *ps)
55{
56	PE *pe;
57	PE_SecBuf *sb, *_sb;
58
59	assert(ps != NULL);
60
61	pe = ps->ps_pe;
62
63	STAILQ_REMOVE(&pe->pe_scn, ps, _PE_Scn, ps_next);
64
65	STAILQ_FOREACH_SAFE(sb, &ps->ps_b, sb_next, _sb)
66		libpe_release_buffer(sb);
67
68	free(ps);
69}
70
71static int
72cmp_scn(PE_Scn *a, PE_Scn *b)
73{
74
75	if (a->ps_sh.sh_addr < b->ps_sh.sh_addr)
76		return (-1);
77	else if (a->ps_sh.sh_addr == b->ps_sh.sh_addr)
78		return (0);
79	else
80		return (1);
81}
82
83static void
84sort_sections(PE *pe)
85{
86
87	if (STAILQ_EMPTY(&pe->pe_scn))
88		return;
89
90	/* Sort the list of Scn by RVA in ascending order. */
91	STAILQ_SORT(&pe->pe_scn, _PE_Scn, ps_next, cmp_scn);
92}
93
94int
95libpe_parse_section_headers(PE *pe)
96{
97	char tmp[sizeof(PE_SecHdr)], *hdr;
98	PE_Scn *ps;
99	PE_SecHdr *sh;
100	PE_CoffHdr *ch;
101	PE_DataDir *dd;
102	int found, i;
103
104	assert(pe->pe_ch != NULL);
105
106	for (i = 0; (uint16_t) i < pe->pe_ch->ch_nsec; i++) {
107		if (read(pe->pe_fd, tmp, sizeof(PE_SecHdr)) !=
108		    (ssize_t) sizeof(PE_SecHdr)) {
109			pe->pe_flags |= LIBPE_F_BAD_SEC_HEADER;
110			return (0);
111		}
112
113		if ((ps = libpe_alloc_scn(pe)) == NULL)
114			return (-1);
115		STAILQ_INSERT_TAIL(&pe->pe_scn, ps, ps_next);
116		ps->ps_ndx = ++pe->pe_nscn;	/* Setion index is 1-based */
117		sh = &ps->ps_sh;
118
119		/*
120		 * Note that the section name won't be NUL-terminated if
121		 * its length happens to be 8.
122		 */
123		memcpy(sh->sh_name, tmp, sizeof(sh->sh_name));
124		hdr = tmp + 8;
125		PE_READ32(hdr, sh->sh_virtsize);
126		PE_READ32(hdr, sh->sh_addr);
127		PE_READ32(hdr, sh->sh_rawsize);
128		PE_READ32(hdr, sh->sh_rawptr);
129		PE_READ32(hdr, sh->sh_relocptr);
130		PE_READ32(hdr, sh->sh_lineptr);
131		PE_READ16(hdr, sh->sh_nreloc);
132		PE_READ16(hdr, sh->sh_nline);
133		PE_READ32(hdr, sh->sh_char);
134	}
135
136	/*
137	 * For all the data directories that don't belong to any section,
138	 * we create pseudo sections for them to make layout easier.
139	 */
140	dd = pe->pe_dd;
141	if (dd != NULL && dd->dd_total > 0) {
142		for (i = 0; (uint32_t) i < pe->pe_dd->dd_total; i++) {
143			if (dd->dd_e[i].de_size == 0)
144				continue;
145			found = 0;
146			STAILQ_FOREACH(ps, &pe->pe_scn, ps_next) {
147				sh = &ps->ps_sh;
148				if (dd->dd_e[i].de_addr >= sh->sh_addr &&
149				    dd->dd_e[i].de_addr + dd->dd_e[i].de_size <=
150				    sh->sh_addr + sh->sh_virtsize) {
151					found = 1;
152					break;
153				}
154			}
155			if (found)
156				continue;
157
158			if ((ps = libpe_alloc_scn(pe)) == NULL)
159				return (-1);
160			STAILQ_INSERT_TAIL(&pe->pe_scn, ps, ps_next);
161			ps->ps_ndx = 0xFFFF0000U | i;
162			sh = &ps->ps_sh;
163			sh->sh_rawptr = dd->dd_e[i].de_addr; /* FIXME */
164			sh->sh_rawsize = dd->dd_e[i].de_size;
165		}
166	}
167
168	/*
169	 * Also consider the COFF symbol table as a pseudo section.
170	 */
171	ch = pe->pe_ch;
172	if (ch->ch_nsym > 0) {
173		if ((ps = libpe_alloc_scn(pe)) == NULL)
174			return (-1);
175		STAILQ_INSERT_TAIL(&pe->pe_scn, ps, ps_next);
176		ps->ps_ndx = 0xFFFFFFFFU;
177		sh = &ps->ps_sh;
178		sh->sh_rawptr = ch->ch_symptr;
179		sh->sh_rawsize = ch->ch_nsym * PE_SYM_ENTRY_SIZE;
180		pe->pe_nsym = ch->ch_nsym;
181	}
182
183	/* PE file headers initialization is complete if we reach here. */
184	return (0);
185}
186
187int
188libpe_load_section(PE *pe, PE_Scn *ps)
189{
190	PE_SecHdr *sh;
191	PE_SecBuf *sb;
192	size_t sz;
193	char tmp[4];
194
195	assert(pe != NULL && ps != NULL);
196	assert((ps->ps_flags & LIBPE_F_LOAD_SECTION) == 0);
197
198	sh = &ps->ps_sh;
199
200	/* Allocate a PE_SecBuf struct without buffer for empty sections. */
201	if (sh->sh_rawsize == 0) {
202		(void) libpe_alloc_buffer(ps, 0);
203		ps->ps_flags |= LIBPE_F_LOAD_SECTION;
204		return (0);
205	}
206
207	if ((pe->pe_flags & LIBPE_F_SPECIAL_FILE) == 0) {
208		if (lseek(pe->pe_fd, (off_t) sh->sh_rawptr, SEEK_SET) < 0) {
209			errno = EIO;
210			return (-1);
211		}
212	}
213
214	if ((sb = libpe_alloc_buffer(ps, sh->sh_rawsize)) == NULL)
215		return (-1);
216
217	if (read(pe->pe_fd, sb->sb_pb.pb_buf, sh->sh_rawsize) !=
218	    (ssize_t) sh->sh_rawsize) {
219		errno = EIO;
220		return (-1);
221	}
222
223	if (ps->ps_ndx == 0xFFFFFFFFU) {
224		/*
225		 * Index 0xFFFFFFFF indicates this section is a pseudo
226		 * section that contains the COFF symbol table. We should
227		 * read in the string table right after it.
228		 */
229		if (read(pe->pe_fd, tmp, sizeof(tmp)) !=
230		    (ssize_t) sizeof(tmp)) {
231			errno = EIO;
232			return (-1);
233		}
234		sz = le32dec(tmp);
235
236		/*
237		 * The minimum value for the size field is 4, which indicates
238		 * there is no string table.
239		 */
240		if (sz > 4) {
241			sz -= 4;
242			if ((sb = libpe_alloc_buffer(ps, sz)) == NULL)
243				return (-1);
244			if (read(pe->pe_fd, sb->sb_pb.pb_buf, sz) !=
245			    (ssize_t) sz) {
246				errno = EIO;
247				return (-1);
248			}
249		}
250	}
251
252	ps->ps_flags |= LIBPE_F_LOAD_SECTION;
253
254	return (0);
255}
256
257int
258libpe_load_all_sections(PE *pe)
259{
260	PE_Scn *ps;
261	PE_SecHdr *sh;
262	unsigned r, s;
263	off_t off;
264	char tmp[256];
265
266	/* Calculate the current offset into the file. */
267	off = 0;
268	if (pe->pe_dh != NULL)
269		off += pe->pe_dh->dh_lfanew + 4;
270	if (pe->pe_ch != NULL)
271		off += sizeof(PE_CoffHdr) + pe->pe_ch->ch_optsize;
272
273	STAILQ_FOREACH(ps, &pe->pe_scn, ps_next) {
274		if (ps->ps_flags & LIBPE_F_LOAD_SECTION)
275			continue;
276		sh = &ps->ps_sh;
277
278		/*
279		 * For special files, we consume the padding in between
280		 * and advance to the section offset.
281		 */
282		if (pe->pe_flags & LIBPE_F_SPECIAL_FILE) {
283			/* Can't go backwards. */
284			if (off > sh->sh_rawptr) {
285				errno = EIO;
286				return (-1);
287			}
288			if (off < sh->sh_rawptr) {
289				r = sh->sh_rawptr - off;
290				for (; r > 0; r -= s) {
291					s = r > sizeof(tmp) ? sizeof(tmp) : r;
292					if (read(pe->pe_fd, tmp, s) !=
293					    (ssize_t) s) {
294						errno = EIO;
295						return (-1);
296					}
297				}
298			}
299		}
300
301		/* Load the section content. */
302		if (libpe_load_section(pe, ps) < 0)
303			return (-1);
304	}
305
306	return (0);
307}
308
309int
310libpe_resync_sections(PE *pe, off_t off)
311{
312	PE_Scn *ps;
313	PE_SecHdr *sh;
314	size_t falign, nsec;
315
316	/* Firstly, sort all sections by their file offsets. */
317	sort_sections(pe);
318
319	/* Count the number of sections. */
320	nsec = 0;
321	STAILQ_FOREACH(ps, &pe->pe_scn, ps_next) {
322		if (ps->ps_flags & LIBPE_F_STRIP_SECTION)
323			continue;
324		if (ps->ps_ndx & 0xFFFF0000U)
325			continue;
326		nsec++;
327	}
328	pe->pe_nscn = nsec;
329
330	/*
331	 * Calculate the file offset for the first section. (`off' is
332	 * currently pointing to the COFF header.)
333	 */
334	off += sizeof(PE_CoffHdr);
335	if (pe->pe_ch != NULL && pe->pe_ch->ch_optsize > 0)
336		off += pe->pe_ch->ch_optsize;
337	else {
338		switch (pe->pe_obj) {
339		case PE_O_PE32:
340			off += PE_COFF_OPT_SIZE_32;
341			break;
342		case PE_O_PE32P:
343			off += PE_COFF_OPT_SIZE_32P;
344			break;
345		case PE_O_COFF:
346		default:
347			break;
348		}
349	}
350	off += nsec * sizeof(PE_SecHdr);
351
352	/*
353	 * Determine the file alignment for sections.
354	 */
355	if (pe->pe_oh != NULL && pe->pe_oh->oh_filealign > 0)
356		falign = pe->pe_oh->oh_filealign;
357	else {
358		/*
359		 * Use the default file alignment defined by the
360		 * PE/COFF specification.
361		 */
362		if (pe->pe_obj == PE_O_COFF)
363			falign = 4;
364		else
365			falign = 512;
366	}
367
368	/*
369	 * Step through each section (and pseduo section) and verify
370	 * alignment constraint and overlapping, make adjustment if need.
371	 */
372	pe->pe_rvamax = 0;
373	STAILQ_FOREACH(ps, &pe->pe_scn, ps_next) {
374		if (ps->ps_flags & LIBPE_F_STRIP_SECTION)
375			continue;
376
377		sh = &ps->ps_sh;
378
379		if (sh->sh_addr + sh->sh_virtsize > pe->pe_rvamax)
380			pe->pe_rvamax = sh->sh_addr + sh->sh_virtsize;
381
382		if (ps->ps_ndx & 0xFFFF0000U)
383			ps->ps_falign = 4;
384		else
385			ps->ps_falign = falign;
386
387		off = roundup(off, ps->ps_falign);
388
389		if (off != sh->sh_rawptr)
390			ps->ps_flags |= PE_F_DIRTY;
391
392		if (ps->ps_flags & PE_F_DIRTY) {
393			if ((ps->ps_flags & LIBPE_F_LOAD_SECTION) == 0) {
394				if (libpe_load_section(pe, ps) < 0)
395					return (-1);
396			}
397			sh->sh_rawsize = libpe_resync_buffers(ps);
398		}
399
400		/*
401		 * Sections only contains uninitialized data should set
402		 * PointerToRawData to zero according to the PE/COFF
403		 * specification.
404		 */
405		if (sh->sh_rawsize == 0)
406			sh->sh_rawptr = 0;
407		else
408			sh->sh_rawptr = off;
409
410		off += sh->sh_rawsize;
411	}
412
413	return (0);
414}
415
416off_t
417libpe_write_section_headers(PE *pe, off_t off)
418{
419	char tmp[sizeof(PE_SecHdr)], *hdr;
420	PE_Scn *ps;
421	PE_SecHdr *sh;
422
423	if (pe->pe_flags & LIBPE_F_BAD_SEC_HEADER || pe->pe_nscn == 0)
424		return (off);
425
426	if ((pe->pe_flags & LIBPE_F_DIRTY_SEC_HEADER) == 0) {
427		off += sizeof(PE_SecHdr) * pe->pe_ch->ch_nsec;
428		return (off);
429	}
430
431	STAILQ_FOREACH(ps, &pe->pe_scn, ps_next) {
432		if (ps->ps_flags & LIBPE_F_STRIP_SECTION)
433			continue;
434		if (ps->ps_ndx & 0xFFFF0000U)
435			continue;
436		if ((pe->pe_flags & LIBPE_F_DIRTY_SEC_HEADER) == 0 &&
437		    (ps->ps_flags & PE_F_DIRTY) == 0)
438			goto next_header;
439
440		sh = &ps->ps_sh;
441
442		memcpy(tmp, sh->sh_name, sizeof(sh->sh_name));
443		hdr = tmp + 8;
444		PE_WRITE32(hdr, sh->sh_virtsize);
445		PE_WRITE32(hdr, sh->sh_addr);
446		PE_WRITE32(hdr, sh->sh_rawsize);
447		PE_WRITE32(hdr, sh->sh_rawptr);
448		PE_WRITE32(hdr, sh->sh_relocptr);
449		PE_WRITE32(hdr, sh->sh_lineptr);
450		PE_WRITE16(hdr, sh->sh_nreloc);
451		PE_WRITE16(hdr, sh->sh_nline);
452		PE_WRITE32(hdr, sh->sh_char);
453
454		if (write(pe->pe_fd, tmp, sizeof(PE_SecHdr)) !=
455		    (ssize_t) sizeof(PE_SecHdr)) {
456			errno = EIO;
457			return (-1);
458		}
459
460	next_header:
461		off += sizeof(PE_SecHdr);
462	}
463
464	return (off);
465}
466
467off_t
468libpe_write_sections(PE *pe, off_t off)
469{
470	PE_Scn *ps;
471	PE_SecHdr *sh;
472
473	if (pe->pe_flags & LIBPE_F_BAD_SEC_HEADER)
474		return (off);
475
476	STAILQ_FOREACH(ps, &pe->pe_scn, ps_next) {
477		sh = &ps->ps_sh;
478
479		if (ps->ps_flags & LIBPE_F_STRIP_SECTION)
480			continue;
481
482		/* Skip empty sections. */
483		if (sh->sh_rawptr == 0 || sh->sh_rawsize == 0)
484			continue;
485
486		/*
487		 * Padding between sections. (padding always written
488		 * in case the the section headers or sections are
489		 * moved or shrunk.)
490		 */
491		assert(off <= sh->sh_rawptr);
492		if (off < sh->sh_rawptr)
493			libpe_pad(pe, sh->sh_rawptr - off);
494
495		if ((ps->ps_flags & PE_F_DIRTY) == 0) {
496			assert((pe->pe_flags & LIBPE_F_SPECIAL_FILE) == 0);
497			if (lseek(pe->pe_fd,
498			    (off_t) (sh->sh_rawptr + sh->sh_rawsize),
499			    SEEK_SET) < 0) {
500				errno = EIO;
501				return (-1);
502			}
503			off = sh->sh_rawptr + sh->sh_rawsize;
504			continue;
505		}
506
507		off = sh->sh_rawptr;
508
509		if (libpe_write_buffers(ps) < 0)
510			return (-1);
511
512		off += sh->sh_rawsize;
513
514		ps->ps_flags &= ~PE_F_DIRTY;
515	}
516
517	return (off);
518}
519