mkzftree.c

/* ----------------------------------------------------------------------- *
 *   
 *   Copyright 2001 H. Peter Anvin - All Rights Reserved
 *
 *   This program is free software; you can redistribute it and/or modify
 *   it under the terms of the GNU General Public License as published by
 *   the Free Software Foundation, Inc., 675 Mass Ave, Cambridge MA 02139,
 *   USA; either version 2 of the License, or (at your option) any later
 *   version; incorporated herein by reference.
 *
 * ----------------------------------------------------------------------- */

/*
 * mkzffile.c
 *
 *	- Generate block-compression of files for use with
 *	  the "ZF" extension to the iso9660/RockRidge filesystem.
 *
 *	  The file compression technique used is the "deflate"
 *	  algorithm used by the zlib library; each block must have a
 *	  valid (12-byte) zlib header.  In addition, the file itself
 *	  has the following structure:
 *
 *	  Byte offset	iso9660 type	Contents
 *	    0		(8 bytes)	Magic number (37 E4 53 96 C9 DB D6 07)
 *	    8		7.3.1		Uncompressed file size
 *	   12		7.1.1		header_size >> 2 (currently 4)
 *	   13		7.1.1		log2(block_size)
 *	   14		(2 bytes)	Reserved, must be zero
 *
 * The header may get expanded in the future, at which point the
 * header size field will be used to increase the space for the
 * header.
 *
 * All implementations are required to support a block_size of 32K
 * (byte 13 == 15).
 *
 * Note that bytes 12 and 13 and the uncompressed length are also
 * present in the ZF record; THE TWO MUST BOTH BE CONSISTENT AND
 * CORRECT.
 *
 * Given the uncompressed size, block_size, and header_size:
 *
 *     nblocks := ceil(size/block_size)
 *
 * After the header follow (nblock+1) 32-bit pointers, recorded as
 * iso9660 7.3.1 (littleendian); each indicate the byte offset (from
 * the start of the file) to one block and the first byte beyond the
 * end of the previous block; the first pointer thus point to the
 * start of the data area and the last pointer to the first byte
 * beyond it:
 *
 *     block_no := floor(byte_offset/block_size)
 *
 *     block_start := read_pointer_731( (header_size+block_no)*4 )
 *     block_end   := read_pointer_731( (header_size+block_no+1)*4 )
 *
 * The block data is compressed according to "zlib".
 */

#include <dirent.h>
#include <errno.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <unistd.h>
#include <limits.h>
#include <utime.h>
#include <sys/stat.h>
#include <sys/types.h>

#include <zlib.h>

#define HAVE_LCHOWN 1		/* Should be obtained by autoconf or so */

/* Command line options */
int force = 0;			/* Always compress */
int level = 9;			/* Compression level */
int verbosity = 0;

/* Program name */
const char *program;

/* Convenience functions */
void *xmalloc(size_t size)
{
  void *p = malloc(size);

  if ( !p ) {
    perror(program);
    exit(1);
  }

  return p;
}

char *xstrdup(const char *str)
{
  char *s = strdup(str);

  if ( !s ) {
    perror(program);
    exit(1);
  }
  
  return s;
}

static void
set_721(void *pnt, unsigned int i)
{
  unsigned char *p = (unsigned char *)pnt;
  p[0] = i & 0xff;
  p[1] = (i >> 8) & 0xff;
}

static unsigned int
get_721(void *pnt)
{
  unsigned char *p = (unsigned char *)pnt;
  return ((unsigned int)p[0]) + ((unsigned int)p[1] << 8);
}

static void
set_722(void *pnt, unsigned int i)
{
  unsigned char *p = (unsigned char *)pnt;
  p[0] = (i >> 8) & 0xff;
  p[1] = i & 0xff;
}

static unsigned int
get_722(void *pnt)
{
  unsigned char *p = (unsigned char *)pnt;
  return ((unsigned int)p[0] << 8) + ((unsigned int)p[1]);
}

static void
set_723(void *pnt, unsigned int i)
{
  unsigned char *p = (unsigned char *)pnt;
  p[3] = p[0] = i & 0xff;
  p[2] = p[1] = (i >> 8) & 0xff;
}

#define get_723(x) get_721(x)

static void
set_731(void *pnt, unsigned int i)
{
  unsigned char *p = (unsigned char *)pnt;
  p[0] = i & 0xff;
  p[1] = (i >> 8) & 0xff;
  p[2] = (i >> 16) & 0xff;
  p[3] = (i >> 24) & 0xff;
}

static unsigned int
get_731(void *pnt)
{
  unsigned char *p = (unsigned char *)pnt;
  return ((unsigned int)p[0]) + ((unsigned int)p[1] << 8) +
    ((unsigned int)p[2] << 16) + ((unsigned int)p[3] << 24);
}

static void
set_732(void *pnt, unsigned int i)
{
  unsigned char *p = (unsigned char *)pnt;
  p[3] = i & 0xff;
  p[2] = (i >> 8) & 0xff;
  p[1] = (i >> 16) & 0xff;
  p[0] = (i >> 24) & 0xff;
}

static unsigned int
get_732(void *pnt)
{
  unsigned char *p = (unsigned char *)pnt;
  return ((unsigned int)p[0] << 24) + ((unsigned int)p[1] << 16) +
    ((unsigned int)p[2] << 8) + ((unsigned int)p[3]);
}

static void
set_733(void *pnt, unsigned int i)
{
  unsigned char *p = (unsigned char *)pnt;
  p[7] = p[0] = i & 0xff;
  p[6] = p[1] = (i >> 8) & 0xff;
  p[5] = p[2] = (i >> 16) & 0xff;
  p[4] = p[3] = (i >> 24) & 0xff;
}

#define get_733(x) get_731(x)

/* File transformation function */
typedef int (*munger_func)(FILE *, FILE *, unsigned long);

/* zisofs definitions */

#ifndef CBLOCK_SIZE_LG2
#define CBLOCK_SIZE_LG2	15	/* Compressed block size */
#endif
#define CBLOCK_SIZE	(1U << CBLOCK_SIZE_LG2)

/* Compressed file magic */
const unsigned char zisofs_magic[8] =
  { 0x37, 0xE4, 0x53, 0x96, 0xC9, 0xDB, 0xD6, 0x07 };

/* VERY VERY VERY IMPORTANT: Must be a multiple of 4 bytes */
struct compressed_file_header {
  char magic[8];
  char uncompressed_len[4];
  unsigned char header_size;
  unsigned char block_size;
  char reserved[2];		/* Reserved for future use, MBZ */
};

int block_uncompress_file(FILE *input, FILE *output, unsigned long size)
{
  struct compressed_file_header hdr;
  char *inbuf, *outbuf;
  long bytes;
  int block_shift;
  char *pointer_block, *pptr;
  unsigned long position;
  unsigned long nblocks;
  unsigned long fullsize, block_size, block_size2;
  unsigned long ptrblock_bytes;
  unsigned long cstart, cend, csize;
  int zerr;
  int err = -1;

  if ( (bytes = fread(&hdr, 1, sizeof hdr, input)) != sizeof hdr ) {
    if ( bytes == size ) {
      /* Very short file; not compressed */
      return ( fwrite(&hdr, 1, bytes, output) != bytes ) ? -1 : 0;
    } else {
      return -1;		/* Read error */
    }
  }

  if ( memcmp(&hdr.magic, zisofs_magic, sizeof zisofs_magic) ) {
    inbuf = xmalloc(CBLOCK_SIZE);
    /* Not compressed */
    memcpy(inbuf, &hdr, sizeof hdr);
    bytes = sizeof hdr;
    do {
      if ( fwrite(inbuf, 1, bytes, output) != bytes )
	return -1;
    } while ( (bytes = fread(inbuf, 1, CBLOCK_SIZE, input)) > 0 );
    free(inbuf);
    return (bytes < 0) ? -1 : 0;
  }

  /* Now we know the file must be compressed.  Get the pointer table. */
  if ( fseek(input, hdr.header_size << 2, SEEK_SET) == -1 )
    return -1;

  fullsize    = get_731(hdr.uncompressed_len);
  block_shift = hdr.block_size;
  block_size  = 1UL << block_shift;
  block_size2 = block_size << 1;
  inbuf  = xmalloc(block_size2);
  outbuf = xmalloc(block_size);

  nblocks = (fullsize + block_size - 1) >> block_shift;

  ptrblock_bytes = (nblocks+1) * 4;
  pointer_block = xmalloc(ptrblock_bytes);

  errno = 0;
  if ( (bytes = fread(pointer_block, 1, ptrblock_bytes, input)) != ptrblock_bytes ) {
    if ( errno == 0 ) errno = EINVAL;
    goto free_ptr_bail;
  }
  
  pptr = pointer_block;
  while ( fullsize ) {
    cstart = get_731(pptr);
    pptr += 4;
    cend   = get_731(pptr);

    csize = cend-cstart;

    if ( csize == 0 ) {
      memset(outbuf, 0, block_size);
      bytes = block_size;
    } else {
      if ( csize > block_size2 ) {
	errno = EINVAL;
	goto free_ptr_bail;
      }
      
      if ( fseek(input, cstart, SEEK_SET) == -1 )
	goto free_ptr_bail;
      
      errno = 0;
      if ( (bytes = fread(inbuf, 1, csize, input)) != csize ) {
	if ( errno == 0 ) errno = EINVAL;
	goto free_ptr_bail;
      }
      
      bytes = block_size;		/* Max output buffer size */
      if ( (zerr = uncompress(outbuf, &bytes, inbuf, csize)) != Z_OK ) {
	errno = (zerr = Z_MEM_ERROR) ? ENOMEM : EINVAL;
	goto free_ptr_bail;
      }
    }
      
    if ( ((fullsize > block_size) && (bytes != block_size))
	 || ((fullsize <= block_size) && (bytes < fullsize)) ) {
      errno = EINVAL;
      goto free_ptr_bail;
    }
    
    if ( bytes > fullsize )
      bytes = fullsize;
    
    errno = 0;
    if ( fwrite(outbuf, 1, bytes, output) != bytes ) {
      if ( errno == 0 ) errno = EINVAL;
      goto free_ptr_bail;
    }

    fullsize -= bytes;
  }

  err = 0;

 free_ptr_bail:
  free(pointer_block);
  free(inbuf);
  free(outbuf);
  return err;
}


int block_compress_file(FILE *input, FILE *output, unsigned long size)
{
  struct compressed_file_header hdr;
  char inbuf[CBLOCK_SIZE], outbuf[2*CBLOCK_SIZE];
  int bytes, pointer_bytes, nblocks, block;
  uLong cbytes;			/* uLong is a zlib datatype */
  char *pointer_block, *curptr;
  unsigned long position;
  int i;
  int header_size;
  int force_compress = force;
  int zerr;

  if ( (sizeof hdr) & 3 ) {
    fputs("INTERNAL ERROR: header is not a multiple of 4\n", stderr);
    abort();
  }

  memset(&hdr, 0, sizeof hdr);
  memcpy(&hdr.magic, zisofs_magic, sizeof zisofs_magic);
  hdr.header_size = (sizeof hdr) >> 2;
  hdr.block_size = CBLOCK_SIZE_LG2;
  set_731(&hdr.uncompressed_len, size);

  if ( fwrite(&hdr, sizeof hdr, 1, output) != 1 )
    return -1;

  nblocks = (size+CBLOCK_SIZE-1) >> CBLOCK_SIZE_LG2;
  pointer_bytes = 4*(nblocks+1);
  pointer_block = xmalloc(pointer_bytes);
  if ( !pointer_block )
    return -1;
  memset(pointer_block, 0, pointer_bytes);

  if ( fseek(output, pointer_bytes, SEEK_CUR) == -1 )
    goto free_ptr_bail;

  curptr = pointer_block;
  position = sizeof hdr + pointer_bytes;
  
  block = 0;
  while ( (bytes = fread(inbuf, 1, CBLOCK_SIZE, input)) > 0 ) {
    if ( bytes < CBLOCK_SIZE && block < nblocks-1 ) {
      errno = EINVAL;		/* Someone changed the file on us */
      goto free_ptr_bail;
    }

    /* HACK: If the file has our magic number, always compress */
    if ( block == 0 && bytes >= sizeof zisofs_magic ) {
      if ( !memcmp(inbuf, zisofs_magic, sizeof zisofs_magic) )
	force_compress = 1;
    }

    set_731(curptr, position); curptr += 4;
    
    /* We have two special cases: a zero-length block is defined as all zero,
       and a block the length of which is equal to the block size is unencoded. */

    for ( i = 0 ; i < CBLOCK_SIZE ; i++ ) {
      if ( inbuf[i] ) break;
    }

    if ( i == CBLOCK_SIZE ) {
      /* All-zero block.  No output */
    } else {
      cbytes = 2*CBLOCK_SIZE;
      if ( (zerr = compress2(outbuf, &cbytes, inbuf, bytes, level)) != Z_OK ) {
	errno = (zerr == Z_MEM_ERROR) ? ENOMEM : EINVAL;
	goto free_ptr_bail;	/* Compression failure */
      }
      if ( fwrite(outbuf, 1, cbytes, output) != cbytes ) 
	goto free_ptr_bail;

      position += cbytes;
    }
    block++;
  }

  /* Set pointer to the end of the final block */
  set_731(curptr, position);

  /* Now write the pointer table */
  if ( fseek(output, sizeof hdr, SEEK_SET) == -1 )
    goto free_ptr_bail;
    
  if ( fwrite(pointer_block, 1, pointer_bytes, output) != pointer_bytes )
    goto free_ptr_bail;

  free(pointer_block);

  /* Now make sure that this was actually the right thing to do */
  if ( !force_compress && position >= size ) {
    /* Incompressible file, just copy it */
    rewind(input);
    rewind(output);

    position = 0;
    while ( (bytes = fread(inbuf, 1, CBLOCK_SIZE, input)) > 0 ) {
      if ( fwrite(inbuf, 1, bytes, output) != bytes )
	return -1;
      position += bytes;
    }

    /* Truncate the file to the correct size */
    fflush(output);
    ftruncate(fileno(output), position);
  }

  /* If we get here, we're done! */
  return 0;

  /* Common bailout code */
 free_ptr_bail:
  free(pointer_block);
  return -1;
}

int munge_path(const char *inpath, const char *outpath, unsigned long size, munger_func munger)
{
  FILE *in, *out;
  int err, rv;

  in = fopen(inpath, "rb");
  if ( !in )
    return -1;
  out = fopen(outpath, "wb");
  if ( !out ) {
    err = errno;
    fclose(in);
    errno = err;
    return -1;
  }
  rv = munger(in, out, size);

  err = errno;			/* Just in case */
  fclose(in);
  fclose(out);
  errno = err;
  return rv;
}

/* Hash table used to find hard-linked files */
#define HASH_BUCKETS 	  2683
struct file_hash {
  struct file_hash *next;
  struct stat st;
  const char *outfile_name;
};

static struct file_hash *hashp[HASH_BUCKETS];

const char *hash_find_file(struct stat *st)
{
  int bucket = (st->st_ino + st->st_dev) % HASH_BUCKETS;
  struct file_hash *hp;

  for ( hp = hashp[bucket] ; hp ; hp = hp->next ) {
    if ( hp->st.st_ino   == st->st_ino &&
	 hp->st.st_dev   == st->st_dev &&
	 hp->st.st_mode  == st->st_mode &&
	 hp->st.st_nlink == st->st_nlink &&
	 hp->st.st_uid   == st->st_uid &&
	 hp->st.st_gid   == st->st_gid &&
	 hp->st.st_size  == st->st_size &&
	 hp->st.st_mtime == st->st_mtime ) {
      /* Good enough, it's the same file */
      return hp->outfile_name;
    }
  }
  return NULL;			/* No match */
}

/* Note: the stat structure is the input file; the name
   is the output file to link to */
void hash_insert_file(struct stat *st, const char *outfile)
{
  int bucket = (st->st_ino + st->st_dev) % HASH_BUCKETS;
  struct file_hash *hp = xmalloc(sizeof(struct file_hash));

  hp->next         = hashp[bucket];
  memcpy(&hp->st, st, sizeof(struct stat));
  hp->outfile_name = xstrdup(outfile);

  hashp[bucket]    = hp;
}


int munge_tree(const char *intree, const char *outtree, munger_func munger)
{
  char buffer[BUFSIZ];
  char *in_path, *out_path, *in_file, *out_file;
  DIR *thisdir;
  struct dirent *dirent;
  struct stat st;
  struct utimbuf ut;
  int err = 0;
  
  /* Construct buffers with the common filename prefix, and point to the end */

  in_path = xmalloc(strlen(intree) + NAME_MAX + 2);
  out_path = xmalloc(strlen(outtree) + NAME_MAX + 2);

  strcpy(in_path, intree);
  strcpy(out_path, outtree);

  in_file = strchr(in_path, '\0');
  out_file = strchr(out_path, '\0');

  *in_file++ = '/';
  *out_file++ = '/';

  /* Open the directory */
  thisdir = opendir(intree);
  if ( !thisdir ) {
    fprintf(stderr, "%s: Failed to open directory %s: %s\n",
	    program, intree, strerror(errno));
    return 1;
  }

  /* Create output directory */
  if ( mkdir(outtree, 0700) ) {
    fprintf(stderr, "%s: Cannot create output directory %s: %s\n",
	    program, outtree, strerror(errno));
    return 1;
  }

  while ( (dirent = readdir(thisdir)) != NULL ) {
    if ( !strcmp(dirent->d_name, ".") ||
	 !strcmp(dirent->d_name, "..") )
      continue;			/* Ignore . and .. */

    strcpy(in_file, dirent->d_name);
    strcpy(out_file, dirent->d_name);

    if ( lstat(in_path, &st) ) {
      fprintf(stderr, "%s: Failed to stat file %s: %s\n",
	      program, in_path, strerror(errno));
      err = 1;
      break;
    }
    
    if ( S_ISREG(st.st_mode) ) {
      if ( st.st_nlink > 1 ) {
	/* Hard link. */
	const char *linkname;

	if ( (linkname = hash_find_file(&st)) != NULL ) {
	  /* We've seen it before, hard link it */

	  if ( link(linkname, out_path) ) {
	    fprintf(stderr, "%s: hard link %s -> %s failed: %s\n",
		    program, out_path, linkname, strerror(errno));
	    err = 1;
	    break;
	  }
	} else {
	  /* First encounter, compress and enter into hash */
	  if ( munge_path(in_path, out_path, st.st_size, munger) ) {
	    fprintf(stderr, "%s: %s: %s", program, in_path, strerror(errno));
	    err = 1;
	    break;
	  }
	  hash_insert_file(&st, out_path);
	}
      } else {
	/* Singleton file; no funnies */
	if ( munge_path(in_path, out_path, st.st_size, munger) ) {
	  fprintf(stderr, "%s: %s: %s", program, in_path, strerror(errno));
	  err = 1;
	  break;
	}
      }
    } else if ( S_ISDIR(st.st_mode) ) {
      /* Recursion: see recursion */
      err = munge_tree(in_path, out_path, munger);
      if ( err )
	break;
    } else if ( S_ISLNK(st.st_mode) ) {
      int chars;
      if ( (chars = readlink(in_path, buffer, BUFSIZ)) < 0 ) {
	fprintf(stderr, "%s: readlink failed for %s: %s\n",
		program, in_path, strerror(errno));
	err = 1;
	break;
      }
      buffer[chars] = '\0';
      if ( symlink(buffer, out_path) ) {
	fprintf(stderr, "%s: symlink %s -> %s failed: %s\n",
		program, out_path, buffer, strerror(errno));
	err = 1;
	break;
      }
    } else {
      if ( st.st_nlink > 1 ) {
	/* Hard link. */
	const char *linkname;

	if ( (linkname = hash_find_file(&st)) != NULL ) {
	  /* We've seen it before, hard link it */

	  if ( link(linkname, out_path) ) {
	    fprintf(stderr, "%s: hard link %s -> %s failed: %s\n",
		    program, out_path, linkname, strerror(errno));
	    err = 1;
	    break;
	  }
	} else {
	  /* First encounter, create and enter into hash */
	  if ( mknod(out_path, st.st_mode, st.st_rdev) ) {
	    fprintf(stderr, "%s: mknod failed for %s: %s\n",
		    program, out_path, strerror(errno));
	    err = 1;
	    break;
	  }
	  hash_insert_file(&st, out_path);
	}
      } else {
	/* Singleton node; no funnies */
	  if ( mknod(out_path, st.st_mode, st.st_rdev) ) {
	    fprintf(stderr, "%s: mknod failed for %s: %s\n",
		    program, out_path, strerror(errno));
	    err = 1;
	    break;
	  }
      }
    }
#ifdef HAVE_LCHOWN
    lchown(out_path, st.st_uid, st.st_gid);
#endif
    if ( !S_ISLNK(st.st_mode) ) {
#ifndef HAVE_LCHOWN
      chown(out_path, st.st_uid, st.st_gid);
#endif
      chmod(out_path, st.st_mode);
      ut.actime  = st.st_atime;
      ut.modtime = st.st_mtime;
      utime(out_path, &ut);
    }
  }
  closedir(thisdir);

  free(in_path);
  free(out_path);

  return err;
}

static void usage(int err)
{
  fprintf(stderr,
	  "Usage: %s [-vfhu] [-z level] intree outtree\n",
	  program);
  exit(err);
}

int main(int argc, char *argv[])
{
  const char *in, *out;
  struct stat st;
  struct utimbuf ut;
  int opt, err;
  munger_func munger = block_compress_file;

  program = argv[0];

  while ( (opt = getopt(argc, argv, "vfz:hu")) != EOF ) {
    switch(opt) {
    case 'f':
      force = 1;		/* Always compress */
      break;
    case 'z':
      if ( optarg[0] < '0' || optarg[0] > '9' || optarg[1] ) {
	fprintf(stderr, "%s: invalid compression level: %s\n",
		program, optarg);
	exit(1);
      } else {
	level = optarg[0] - '0';
      }
      break;
    case 'h':
      usage(0);
      break;
    case 'v':
      verbosity++;
      break;
    case 'u':
      munger = block_uncompress_file;
      break;
    default:
      usage(1);
      break;
    }
  }

  if ( (argc-optind) != 2 )
    usage(1);

  in  = argv[optind];		/* Input tree */
  out = argv[optind+1];		/* Output tree */

  umask(077);

  /* Special case: we use stat() for the root, not lstat() */
  if ( stat(in, &st) ) {
    fprintf(stderr, "%s: %s: %s\n", program, in, strerror(errno));
    exit(1);
  }
  if ( !S_ISDIR(st.st_mode) ) {
    fprintf(stderr, "%s: %s: Not a directory\n", program, in);
  }

  err = munge_tree(in, out, munger);

  if ( err )
    exit(err);

  chown(out, st.st_uid, st.st_gid);
  chmod(out, st.st_mode);
  ut.actime  = st.st_atime;
  ut.modtime = st.st_mtime;
  utime(out, &ut);
}