1/*-
2 * Copyright (c) 2015 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 <sys/types.h>
29#include <assert.h>
30#include <errno.h>
31#include <stdlib.h>
32#include <string.h>
33#include <unistd.h>
34
35#include "_libpe.h"
36
37ELFTC_VCSID("$Id: libpe_dos.c 3312 2016-01-10 09:23:51Z kaiwang27 $");
38
39int
40libpe_parse_msdos_header(PE *pe, char *hdr)
41{
42	PE_DosHdr *dh;
43	char coff[sizeof(PE_CoffHdr)];
44	uint32_t pe_magic;
45	int i;
46
47	if ((pe->pe_stub = malloc(sizeof(PE_DosHdr))) == NULL) {
48		errno = ENOMEM;
49		return (-1);
50	}
51	memcpy(pe->pe_stub, hdr, sizeof(PE_DosHdr));
52
53	if ((dh = malloc(sizeof(*dh))) == NULL) {
54		errno = ENOMEM;
55		return (-1);
56	}
57	pe->pe_dh = dh;
58
59	/* Read the conventional MS-DOS EXE header. */
60	memcpy(dh->dh_magic, hdr, 2);
61	hdr += 2;
62	PE_READ16(hdr, dh->dh_lastsize);
63	PE_READ16(hdr, dh->dh_nblock);
64	PE_READ16(hdr, dh->dh_nreloc);
65	PE_READ16(hdr, dh->dh_hdrsize);
66	PE_READ16(hdr, dh->dh_minalloc);
67	PE_READ16(hdr, dh->dh_maxalloc);
68	PE_READ16(hdr, dh->dh_ss);
69	PE_READ16(hdr, dh->dh_sp);
70	PE_READ16(hdr, dh->dh_checksum);
71	PE_READ16(hdr, dh->dh_ip);
72	PE_READ16(hdr, dh->dh_cs);
73	PE_READ16(hdr, dh->dh_relocpos);
74	PE_READ16(hdr, dh->dh_noverlay);
75
76	/* Do not continue if the EXE is not a PE/NE/... (new executable) */
77	if (dh->dh_relocpos != 0x40) {
78		pe->pe_flags |= LIBPE_F_BAD_DOS_HEADER;
79		return (0);
80	}
81
82	for (i = 0; i < 4; i++)
83		PE_READ16(hdr, dh->dh_reserved1[i]);
84	PE_READ16(hdr, dh->dh_oemid);
85	PE_READ16(hdr, dh->dh_oeminfo);
86	for (i = 0; i < 10; i++)
87		PE_READ16(hdr, dh->dh_reserved2[i]);
88	PE_READ32(hdr, dh->dh_lfanew);
89
90	/* Check if the e_lfanew pointer is valid. */
91	if (dh->dh_lfanew > pe->pe_fsize - 4) {
92		pe->pe_flags |= LIBPE_F_BAD_DOS_HEADER;
93		return (0);
94	}
95
96	if (dh->dh_lfanew < sizeof(PE_DosHdr) &&
97	    (pe->pe_flags & LIBPE_F_SPECIAL_FILE)) {
98		pe->pe_flags |= LIBPE_F_BAD_DOS_HEADER;
99		return (0);
100	}
101
102	if (dh->dh_lfanew > sizeof(PE_DosHdr)) {
103		pe->pe_stub_ex = dh->dh_lfanew - sizeof(PE_DosHdr);
104		if (pe->pe_flags & LIBPE_F_SPECIAL_FILE) {
105			/* Read in DOS stub now. */
106			if (libpe_read_msdos_stub(pe) < 0) {
107				pe->pe_flags |= LIBPE_F_BAD_DOS_HEADER;
108				return (0);
109			}
110		}
111	}
112
113	if ((pe->pe_flags & LIBPE_F_SPECIAL_FILE) == 0) {
114		/* Jump to the PE header. */
115		if (lseek(pe->pe_fd, (off_t) dh->dh_lfanew, SEEK_SET) < 0) {
116			pe->pe_flags |= LIBPE_F_BAD_PE_HEADER;
117			return (0);
118		}
119	}
120
121	if (read(pe->pe_fd, &pe_magic, 4) != 4 ||
122	    htole32(pe_magic) != PE_SIGNATURE) {
123		pe->pe_flags |= LIBPE_F_BAD_PE_HEADER;
124		return (0);
125	}
126
127	if (read(pe->pe_fd, coff, sizeof(coff)) != (ssize_t) sizeof(coff)) {
128		pe->pe_flags |= LIBPE_F_BAD_COFF_HEADER;
129		return (0);
130	}
131
132	return (libpe_parse_coff_header(pe, coff));
133}
134
135int
136libpe_read_msdos_stub(PE *pe)
137{
138	void *m;
139
140	assert(pe->pe_stub_ex > 0 &&
141	    (pe->pe_flags & LIBPE_F_LOAD_DOS_STUB) == 0);
142
143	if ((pe->pe_flags & LIBPE_F_SPECIAL_FILE) == 0) {
144		if (lseek(pe->pe_fd, (off_t) sizeof(PE_DosHdr), SEEK_SET) <
145		    0) {
146			errno = EIO;
147			goto fail;
148		}
149	}
150
151	if ((m = realloc(pe->pe_stub, sizeof(PE_DosHdr) + pe->pe_stub_ex)) ==
152	    NULL) {
153		errno = ENOMEM;
154		goto fail;
155	}
156	pe->pe_stub = m;
157
158	if (read(pe->pe_fd, pe->pe_stub + sizeof(PE_DosHdr), pe->pe_stub_ex) !=
159	    (ssize_t) pe->pe_stub_ex) {
160		errno = EIO;
161		goto fail;
162	}
163
164	pe->pe_flags |= LIBPE_F_LOAD_DOS_STUB;
165
166	/* Search for the Rich header embedded just before the PE header. */
167	(void) libpe_parse_rich_header(pe);
168
169	return (0);
170
171fail:
172	pe->pe_stub_ex = 0;
173
174	return (-1);
175}
176
177/*
178 * The "standard" MS-DOS stub displaying "This program cannot be run in
179 * DOS mode".
180 */
181static const char msdos_stub[] = {
182    '\x0e','\x1f','\xba','\x0e','\x00','\xb4','\x09','\xcd',
183    '\x21','\xb8','\x01','\x4c','\xcd','\x21','\x54','\x68',
184    '\x69','\x73','\x20','\x70','\x72','\x6f','\x67','\x72',
185    '\x61','\x6d','\x20','\x63','\x61','\x6e','\x6e','\x6f',
186    '\x74','\x20','\x62','\x65','\x20','\x72','\x75','\x6e',
187    '\x20','\x69','\x6e','\x20','\x44','\x4f','\x53','\x20',
188    '\x6d','\x6f','\x64','\x65','\x2e','\x0d','\x0d','\x0a',
189    '\x24','\x00','\x00','\x00','\x00','\x00','\x00','\x00',
190};
191
192static void
193init_dos_header(PE_DosHdr *dh)
194{
195
196	dh->dh_magic[0] = 'M';
197	dh->dh_magic[1] = 'Z';
198	dh->dh_lastsize = 144;
199	dh->dh_nblock = 3;
200	dh->dh_hdrsize = 4;
201	dh->dh_maxalloc = 65535;
202	dh->dh_sp = 184;
203	dh->dh_relocpos = 0x40;
204	dh->dh_lfanew = 0x80;
205}
206
207off_t
208libpe_write_msdos_stub(PE *pe, off_t off)
209{
210	PE_DosHdr *dh;
211	char tmp[sizeof(PE_DosHdr)], *hdr;
212	off_t d;
213	int i, strip_rich;
214
215	strip_rich = 0;
216
217	if (pe->pe_cmd == PE_C_RDWR) {
218		assert((pe->pe_flags & LIBPE_F_SPECIAL_FILE) == 0);
219
220		if (pe->pe_dh != NULL &&
221		    (pe->pe_flags & PE_F_STRIP_DOS_STUB)) {
222			/*
223			 * If we strip MS-DOS stub, everything after it
224			 * needs rewritten.
225			 */
226			pe->pe_flags |= LIBPE_F_BAD_PE_HEADER;
227			goto done;
228		}
229
230		/*
231		 * lseek(2) to the PE signature if MS-DOS stub is not
232		 * modified.
233		 */
234		if (pe->pe_dh != NULL &&
235		    (pe->pe_flags & LIBPE_F_DIRTY_DOS_HEADER) == 0 &&
236		    (pe->pe_flags & LIBPE_F_BAD_DOS_HEADER) == 0 &&
237		    (pe->pe_flags & PE_F_STRIP_RICH_HEADER) == 0) {
238			if (lseek(pe->pe_fd,
239			    (off_t) (sizeof(PE_DosHdr) + pe->pe_stub_ex),
240			    SEEK_CUR) < 0) {
241				errno = EIO;
242				return (-1);
243			}
244			off = sizeof(PE_DosHdr) + pe->pe_stub_ex;
245			goto done;
246		}
247
248		/* Check if we should strip the Rich header. */
249		if (pe->pe_dh != NULL && pe->pe_stub_app == NULL &&
250		    (pe->pe_flags & LIBPE_F_BAD_DOS_HEADER) == 0 &&
251		    (pe->pe_flags & PE_F_STRIP_RICH_HEADER)) {
252			if ((pe->pe_flags & LIBPE_F_LOAD_DOS_STUB) == 0) {
253				(void) libpe_read_msdos_stub(pe);
254				if (lseek(pe->pe_fd, off, SEEK_SET) < 0) {
255					errno = EIO;
256					return (-1);
257				}
258			}
259			if (pe->pe_rh != NULL) {
260				strip_rich = 1;
261				pe->pe_flags |= LIBPE_F_DIRTY_DOS_HEADER;
262			}
263		}
264
265		/*
266		 * If length of MS-DOS stub will change, Mark the PE
267		 * signature is broken so that the PE signature and the
268		 * headers follow it will be rewritten.
269		 *
270		 * The sections should be loaded now since the stub might
271		 * overwrite the section data.
272		 */
273		if ((pe->pe_flags & LIBPE_F_BAD_DOS_HEADER) ||
274		    (pe->pe_stub_app != NULL && pe->pe_stub_app_sz !=
275			sizeof(PE_DosHdr) + pe->pe_stub_ex) || strip_rich) {
276			if (libpe_load_all_sections(pe) < 0)
277				return (-1);
278			if (lseek(pe->pe_fd, off, SEEK_SET) < 0) {
279				errno = EIO;
280				return (-1);
281			}
282			pe->pe_flags |= LIBPE_F_BAD_PE_HEADER;
283		}
284	}
285
286	if (pe->pe_flags & PE_F_STRIP_DOS_STUB)
287		goto done;
288
289	/* Always use application supplied MS-DOS stub, if exists. */
290	if (pe->pe_stub_app != NULL && pe->pe_stub_app_sz > 0) {
291		if (write(pe->pe_fd, pe->pe_stub_app, pe->pe_stub_app_sz) !=
292		    (ssize_t) pe->pe_stub_app_sz) {
293			errno = EIO;
294			return (-1);
295		}
296		off = pe->pe_stub_app_sz;
297		goto done;
298	}
299
300	/*
301	 * Write MS-DOS header.
302	 */
303
304	if (pe->pe_dh == NULL) {
305		if ((dh = calloc(1, sizeof(PE_DosHdr))) == NULL) {
306			errno = ENOMEM;
307			return (-1);
308		}
309		pe->pe_dh = dh;
310
311		init_dos_header(dh);
312
313		pe->pe_flags |= LIBPE_F_DIRTY_DOS_HEADER;
314	} else
315		dh = pe->pe_dh;
316
317	if (pe->pe_flags & LIBPE_F_BAD_DOS_HEADER)
318		init_dos_header(dh);
319
320	if (strip_rich) {
321		d = pe->pe_rh_start - pe->pe_stub;
322		dh->dh_lfanew = roundup(d, 8);
323	}
324
325	if ((pe->pe_flags & LIBPE_F_DIRTY_DOS_HEADER) ||
326	    (pe->pe_flags & LIBPE_F_BAD_DOS_HEADER)) {
327		memcpy(tmp, dh->dh_magic, 2);
328		hdr = tmp + 2;
329		PE_WRITE16(hdr, dh->dh_lastsize);
330		PE_WRITE16(hdr, dh->dh_nblock);
331		PE_WRITE16(hdr, dh->dh_nreloc);
332		PE_WRITE16(hdr, dh->dh_hdrsize);
333		PE_WRITE16(hdr, dh->dh_minalloc);
334		PE_WRITE16(hdr, dh->dh_maxalloc);
335		PE_WRITE16(hdr, dh->dh_ss);
336		PE_WRITE16(hdr, dh->dh_sp);
337		PE_WRITE16(hdr, dh->dh_checksum);
338		PE_WRITE16(hdr, dh->dh_ip);
339		PE_WRITE16(hdr, dh->dh_cs);
340		PE_WRITE16(hdr, dh->dh_relocpos);
341		PE_WRITE16(hdr, dh->dh_noverlay);
342		for (i = 0; i < 4; i++)
343			PE_WRITE16(hdr, dh->dh_reserved1[i]);
344		PE_WRITE16(hdr, dh->dh_oemid);
345		PE_WRITE16(hdr, dh->dh_oeminfo);
346		for (i = 0; i < 10; i++)
347			PE_WRITE16(hdr, dh->dh_reserved2[i]);
348		PE_WRITE32(hdr, dh->dh_lfanew);
349
350		if (write(pe->pe_fd, tmp, sizeof(tmp)) !=
351		    (ssize_t) sizeof(tmp)) {
352			errno = EIO;
353			return (-1);
354		}
355	} else {
356		assert((pe->pe_flags & LIBPE_F_SPECIAL_FILE) == 0);
357		if (lseek(pe->pe_fd, (off_t) sizeof(PE_DosHdr), SEEK_CUR) <
358		    0) {
359			errno = EIO;
360			return (-1);
361		}
362	}
363
364	off = sizeof(PE_DosHdr);
365
366	/*
367	 * Write the MS-DOS stub.
368	 */
369
370	if (strip_rich) {
371		assert((pe->pe_flags & LIBPE_F_SPECIAL_FILE) == 0);
372		assert(pe->pe_stub != NULL && pe->pe_rh_start != NULL);
373		d = pe->pe_rh_start - pe->pe_stub;
374		if (lseek(pe->pe_fd, d, SEEK_SET) < 0) {
375			errno = EIO;
376			return (-1);
377		}
378		off = d;
379		goto done;
380	}
381
382	if (pe->pe_cmd == PE_C_RDWR) {
383		if (lseek(pe->pe_fd, (off_t) pe->pe_stub_ex, SEEK_CUR) < 0) {
384			errno = EIO;
385			return (-1);
386		}
387		off += pe->pe_stub_ex;
388		goto done;
389	}
390
391	if (write(pe->pe_fd, msdos_stub, sizeof(msdos_stub)) !=
392	    (ssize_t) sizeof(msdos_stub)) {
393		errno = EIO;
394		return (-1);
395	}
396	off += sizeof(msdos_stub);
397
398done:
399	pe->pe_flags &= ~LIBPE_F_DIRTY_DOS_HEADER;
400	pe->pe_flags &= ~LIBPE_F_BAD_DOS_HEADER;
401
402	return (off);
403}
404