error checks, bug fixes

[rigol] / png.c

#include <string.h>
#include <stdint.h>
#include <arpa/inet.h>
#include <zlib.h>
#include <stdio.h>
#include <stdlib.h>

#include "png.h"

#define CHUNK 16384
#define WIDTH 320
#define HEIGHT 234

/* Table of CRCs of all 8-bit messages. */
static unsigned long crc_table[256];

/* Flag: has the table been computed? Initially false. */
static int crc_table_computed = 0;

/* Make the table for a fast CRC. */
static void make_crc_table(void)
{
	unsigned long c;
	int n, k;

	for (n = 0; n < 256; n++) {
		c = (unsigned long) n;
		for (k = 0; k < 8; k++) {
			if (c & 1)
				c = 0xedb88320L ^ (c >> 1);
			else
				c = c >> 1;
		}
		crc_table[n] = c;
	}
	crc_table_computed = 1;
}

/* Update a running CRC with the bytes buf[0..len-1]--the CRC
   should be initialized to all 1's, and the transmitted value
   is the 1's complement of the final running CRC (see the
   crc() routine below). */

static unsigned long update_crc(unsigned long crc, unsigned char *buf,
		int len)
{
	unsigned long c = crc;
	int n;

	if (!crc_table_computed)
		make_crc_table();
	for (n = 0; n < len; n++) {
		c = crc_table[(c ^ buf[n]) & 0xff] ^ (c >> 8);
	}
	return c;
}

/* Return the CRC of the bytes buf[0..len-1]. */
static unsigned long crc(unsigned char *buf, int len)
{
	return update_crc(0xffffffffL, buf, len) ^ 0xffffffffL;
}

unsigned char *lcd2png(unsigned char *lcd, int *len)
{
	unsigned char screen_conv[(WIDTH*HEIGHT*3)+HEIGHT]; /* 320*234 RGB + scanline filter byte */
	unsigned char *pos;
	unsigned char lut[256][3];
	unsigned char *screen_deflated;
	int deflated_size;
	unsigned char *image = NULL;
	unsigned char *outpos;
	static const unsigned char png[] = {137, 80, 78, 71, 13, 10, 26, 10};
	static unsigned char idat[] = {'I', 'D', 'A', 'T'};
	static unsigned char iend[] = {'I', 'E', 'N', 'D'};
	static unsigned char ihdr[] = {'I', 'H', 'D', 'R',
		0x00, 0x00, 0x01, 0x40, /* 320 - Width */
		0x00, 0x00, 0x00, 0xea, /* 234 - Height */
		0x08, /* Bit depth */
		0x02, /* RGB Truecolor, Colour type */
		0x00, /* Deflate, Compression method */
		0x00, /* Adaptive, Filter method */
		0x00  /* No interlace, Interlace method */
		};
	uint32_t l;
	int i;
	int ret;
	int toalloc = 0;
	int flush;
	z_stream strm;

	for(i = 0; i < 256; i++) {
		lut[i][0] = ((i >> 6) * 0x55);
		lut[i][1] = ((((i >> 3) & 7) * 0x49) >> 1);
		lut[i][2] = (((i & 7) * 0x49) >> 1);
	}

	pos = screen_conv;
	for(i = 0; i < (WIDTH*HEIGHT); i++) {
		if ((i % WIDTH) == 0) {
			*pos++ = 0x00; /* No adaptive filter */
		}
		*pos++ = lut[lcd[i]][0];
		*pos++ = lut[lcd[i]][1];
		*pos++ = lut[lcd[i]][2];
	}

	strm.zalloc = Z_NULL;
	strm.zfree = Z_NULL;
	strm.opaque = Z_NULL;
	if (deflateInit(&strm, Z_BEST_COMPRESSION) != Z_OK) {
		perror("deflateInit");
		exit(EXIT_FAILURE);
	}

	strm.avail_in = sizeof(screen_conv);
	strm.next_in = screen_conv;

	toalloc = 0;
	strm.avail_out = 0;
	screen_deflated = NULL;
	flush = Z_NO_FLUSH;

	do {
		if (strm.avail_out == 0) {
			toalloc += CHUNK;
			screen_deflated = realloc(screen_deflated, toalloc);
			if (screen_deflated == NULL) {
				perror("realloc");
				exit(EXIT_FAILURE);
			}

			strm.avail_out = CHUNK;
			strm.next_out = screen_deflated + (toalloc - CHUNK);
		}

		ret = deflate(&strm, flush);
		if (ret == Z_STREAM_ERROR) {
			perror("deflate");
			exit(EXIT_FAILURE);
		}

		if(strm.avail_in == 0) {
			flush = Z_FINISH;
		}
	} while(ret != Z_STREAM_END);
	deflated_size = toalloc - strm.avail_out;

	deflateEnd(&strm);

	image = malloc(sizeof(png) +
		sizeof(ihdr) + 8 + /* 8 = length, csum */
		sizeof(idat) + deflated_size + 8 +
		sizeof(iend) + 8);
	
	if (image == NULL) {
		perror("malloc");
	}
	
	outpos = image;

	/* PNG signature */
	memcpy(outpos, png, sizeof(png));
	outpos += sizeof(png);

	/* IHDR */
	l = htonl(sizeof(ihdr) - 4); /* "IHDR" is not counted */
	memcpy(outpos, &l, sizeof(l));
	outpos += sizeof(l);
	memcpy(outpos, ihdr, sizeof(ihdr));
	outpos += sizeof(ihdr);
	l = htonl(crc(ihdr, sizeof(ihdr)));
	memcpy(outpos, &l, sizeof(l));
	outpos += sizeof(l);

	/* IDAT */
	l = htonl(deflated_size);
	memcpy(outpos, &l, sizeof(l));
	outpos += sizeof(l);
	memcpy(outpos, idat, sizeof(idat));
	outpos += sizeof(idat);
	memcpy(outpos, screen_deflated, deflated_size);
	free(screen_deflated);
	outpos +=deflated_size;
	l = htonl(crc(outpos-deflated_size-sizeof(idat), deflated_size+sizeof(idat)));
	memcpy(outpos, &l, sizeof(l));
	outpos += sizeof(l);

	/* IEND */
	l = htonl(sizeof(iend) - 4); /* "IEND" is not counted */
	memcpy(outpos, &l, sizeof(l));
	outpos += sizeof(l);
	memcpy(outpos, iend, sizeof(iend));
	outpos += sizeof(iend);
	l = htonl(crc(iend, sizeof(iend)));
	memcpy(outpos, &l, sizeof(l));
	outpos += sizeof(l);

	*len = (int)(outpos - image);
	return image;
}