/*
 * Copyright (c) 1999-2003 Apple Computer, Inc. All rights reserved.
 *
 * @APPLE_LICENSE_HEADER_START@
 * 
 * Portions Copyright (c) 1999-2003 Apple Computer, Inc.  All Rights
 * Reserved.  This file contains Original Code and/or Modifications of
 * Original Code as defined in and that are subject to the Apple Public
 * Source License Version 2.0 (the "License").  You may not use this file
 * except in compliance with the License.  Please obtain a copy of the
 * License at http://www.apple.com/publicsource and read it before using
 * this file.
 * 
 * The Original Code and all software distributed under the License are
 * distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND, EITHER
 * EXPRESS OR IMPLIED, AND APPLE HEREBY DISCLAIMS ALL SUCH WARRANTIES,
 * INCLUDING WITHOUT LIMITATION, ANY WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE OR NON- INFRINGEMENT.  Please see the
 * License for the specific language governing rights and limitations
 * under the License.
 * 
 * @APPLE_LICENSE_HEADER_END@
 */
#include <stdio.h>
#include <stdlib.h>
#include <mach/mach.h>
#include <mach/mach_error.h>
#include <sys/file.h>
#include <mach-o/loader.h>
#include <libkern/OSByteOrder.h>
#include <unistd.h>

int	infile, outfile;

struct mach_header	mh;
void *		cmds;

boolean_t		swap_ends;

static unsigned long swap(
    unsigned long x
)
{
    if (swap_ends)
	return OSSwapInt32(x);
    else
	return x;
}

/*! @function	get_segment_sizes
    @arg    scp		Segment Command to process
    @arg    p_insize	Output parameter will be set to amount of data to be read from input file
    @arg    p_outsize	Output parameter will be set to amount of data to be written to output file
    @result	True if the segment contains zero fill sections and thus should be the last segment.
    @discussion
    *p_outsize should always be >= *p_insize.  When a segment contains one or more S_ZEROFILL sections
    the *p_outsize and *p_insize will both be set to the address of the first S_ZEROFILL section relative
    to the address of the segment.  This accounts for any padding present in the file after the declared
    end of the preceding (non-S_ZEROFILL) section.

    When a segment does not contain any S_ZEROFILL sections, *p_insize will be set to the segment command's
    filesize and *p_outsize will be set to the segment command's vmsize.
 */
boolean_t get_segment_sizes(struct segment_command *scp, uint32_t *p_insize, uint32_t *p_outsize)
{
    // Section info directly follows the segment command
    struct section *sections = (struct section *)(((char *)scp) + sizeof(*scp));
    uint32_t i;

    // At the end of the loop over sections, last_end_off will be the highest off + size of
    // all of the declared sections.  We start it out as the segment's fileoff as that is
    // the lowest it can possibly be.  It will be increased as non-S_ZEROFILL sections
    // are found.
    uint32_t last_end_off = swap(scp->fileoff);

    boolean_t gotZeroFill = FALSE;

    // zerofill_start is the start of the zerofill relative to the segment.
    // Start it at the total filesize of the segment and it will be decreased
    // as zerofill sections are found.
    uint32_t zerofill_start = swap(scp->filesize);

    for(i=0; i < swap(scp->nsects); ++i)
    {
	if( (swap(sections[i].flags) & SECTION_TYPE) == S_ZEROFILL )
	{
	    uint32_t this_start = swap(sections[i].addr) - swap(scp->vmaddr);
	    if(this_start < zerofill_start)
		zerofill_start = this_start;
	    gotZeroFill = TRUE;
	}
	else
	{
	    // Ending offset of this section is beginning offset + size
	    uint32_t this_section_end_off = swap(sections[i].offset) + swap(sections[i].size);
	    if(this_section_end_off > last_end_off)
		last_end_off = this_section_end_off;
	}
    }
    uint32_t section_filesize = last_end_off - swap(scp->fileoff);
    // If there are no S_ZEROFILL sections, zerofill_start will be the segment's filesize
    // which still must be greater than or equal to the the filesize needed to account for all sections.
    if(zerofill_start >= section_filesize)
	section_filesize = zerofill_start;
    else if(gotZeroFill != FALSE)
    {
	fprintf(stderr, "Segment %s has zero-fill sections that begin before non-zero sections end!\n", scp->segname);
	exit(1);
    }
    else
    {
	fprintf(stderr, "Segment %s declares less filesize than needed to account for all sections\n", scp->segname);
	exit(1);
    }

    // If we ended with zero-fill sections then use the start of the zero fill as the filesize
    // Prior to ld_classic from cctools-750~70 (Snow Leopard) this was equal to segment's filesize but
    // with the newer ld_classic it seems to want to output 16 extra zero bytes that overlaps the
    // start of the zerofill section.
    if(gotZeroFill != FALSE)
    {
	*p_insize = *p_outsize = zerofill_start;
	// Indicate that there must not be any more segments.
	return TRUE;
    }
    else
    {
	// This is what the original code basically did:
	*p_insize = swap(scp->filesize);
	*p_outsize = swap(scp->vmsize);
	// Indicate that there may be more segments.
	return FALSE;
    }
}

/**
 * A full page of zero-valued bytes for use as padding between segments.
 *
 * It is const and global and therefore will wind up in the __DATA,__bss segment and be zero.
 */
static char const zeropage[4096];

int
main(int argc, char *argv[])
{
    kern_return_t	result;
    vm_address_t	data;
    int			nc, ncmds;
    char *		cp;
    uint32_t		csum = 0;
    
    if (argc == 2) {
	infile = open(argv[1], O_RDONLY);
	if (infile < 0)
	    goto usage;
	outfile = fileno(stdout);
    }
    else if (argc == 3) {
    	infile = open(argv[1], O_RDONLY);
	if (infile < 0)
	    goto usage;
	outfile = open(argv[2], O_WRONLY|O_CREAT|O_TRUNC, 0644);
	if (outfile < 0)
	    goto usage;
    }
    else {
usage:
    	fprintf(stderr, "usage: machOconv inputfile [outputfile]\n");
	exit(1);
    }
    
    nc = read(infile, &mh, sizeof (mh));
    if (nc < 0) {
	perror("read mach header");
	exit(1);
    }
    if (nc < (int)sizeof (mh)) {
	fprintf(stderr, "read mach header: premature EOF %d\n", nc);
	exit(1);
    }
    if (mh.magic == MH_MAGIC)
	swap_ends = FALSE;
    else if (mh.magic == MH_CIGAM)
	swap_ends = TRUE;
    else {
    	fprintf(stderr, "bad magic number %lx\n", (unsigned long)mh.magic);
	exit(1);
    }

    cmds = calloc(swap(mh.sizeofcmds), sizeof (char));
    if (cmds == 0) {
	fprintf(stderr, "alloc load commands: no memory\n");
	exit(1);
    }
    nc = read(infile, cmds, swap(mh.sizeofcmds));
    if (nc < 0) {
	perror("read load commands");
	exit(1);
    }
    if (nc < (int)swap(mh.sizeofcmds)) {
	fprintf(stderr, "read load commands: premature EOF %d\n", nc);
	exit(1);
    }

    // The end of the previous segment, i.e. vmaddr + vmsize
    // as in the next address that is expected to be written
    // or the special initial -1 (0xFFFFFFFF) value.
    uint32_t prevseg_end = (uint32_t)-1;

    boolean_t gotLastSegment = FALSE;
    for (	ncmds = swap(mh.ncmds), cp = cmds;
		ncmds > 0; ncmds--) {
	    uint32_t insize;
	    uint32_t vmsize;
	    uint32_t *csum_ptr;

#define lcp	((struct load_command *)cp)    
	switch(swap(lcp->cmd)) {

	case LC_SEGMENT:
#define scp	((struct segment_command *)cp)
	    // If we're the first segment, use its location as the image base
	    if (prevseg_end != (uint32_t)-1) {
		int requiredPadding = scp->vmaddr - prevseg_end;
		// Handle the case where the segments were out of order by failing
		if (requiredPadding < 0) {
		    fprintf(stderr, "%s 0x%08x comes before 0x%08x\n", scp->segname, scp->vmaddr, prevseg_end);
		    exit(1);
		}
		// It should be extremely rare that more than one page of padding is needed
		for (; requiredPadding > 4096; requiredPadding -= 4096) {
		    write(outfile, zeropage, 4096);
		}
		// Write the remainder of the padding and log the next start and the previous end.
		if (requiredPadding > 0) {
		    write(outfile, zeropage, requiredPadding);
		    requiredPadding -= requiredPadding;
		    fprintf(stderr, "%s 0x%08x comes after 0x%08x\n", scp->segname, scp->vmaddr, prevseg_end);
		}
	    }

	    if(gotLastSegment)
	    {
		// FIXME: We could be more lenient here because it is valid to have another segment so long
		// as it contains only S_ZEROFILL sections and has zero filesize.
		fprintf(stderr,"Got an additional segment after one that contained zero-filled sections\n");
		exit(1);
	    }
	    gotLastSegment = get_segment_sizes(scp, &insize, &vmsize);

	    result = vm_allocate(mach_task_self(), &data, vmsize, TRUE);
	    if (result != KERN_SUCCESS) {
		mach_error("vm_allocate segment data", result);
		exit(1);
	    }

	    // Make sure any padding we wind up adding is zeroed. It won't screw up the checksum since it
	    // reads all of vmsize but we prefer not to output any random data that might exist in memory.
	    if(vmsize >= insize)
		bzero((char*)data + insize, vmsize - insize);
	    else
	    {
		fprintf(stderr, "Will not read more bytes from input than will be written to output\n");
		exit(1);
	    }


	    lseek(infile, swap(scp->fileoff), L_SET);
	    nc = read(infile, (void *)data, insize);
	    if (nc < 0) {
		perror("read segment data");
		exit(1);
	    }
	    if (nc < (int)insize) {
		fprintf(stderr, "read segment data: premature EOF %d\n", nc);
		exit(1);
	    }

	    nc = write(outfile, (void *)data, vmsize);

	    // Update the new segment end to account for the padding we just performed
	    prevseg_end = scp->vmaddr + scp->vmsize;

	    if (nc < (int)vmsize) {
		perror("write segment data");
		exit(1);
	    }

	    // FIXME: We could try to handle this case but it's easier to make sure the
	    // segments are 4-byte aligned in the binary.  They should be anyway.
	    if( (vmsize % sizeof(uint32_t)) != 0 ) {
		perror("segments must end on a 4-byte boundary");
		exit(1);
	    }
	    for(csum_ptr = (uint32_t*)data; csum_ptr < (uint32_t*)((char*)data + vmsize); ++csum_ptr)
		csum += swap(*csum_ptr);
	    
	    vm_deallocate(mach_task_self(), data, vmsize);
	    break;
	}

	cp += swap(lcp->cmdsize);
    }

    csum = (0 - csum);
    write(outfile, &csum, sizeof(csum));
	
    exit(0);
}

/*
vim:sw=4:sts=4:ts=8:noet
 */