From 7ba4ad35ccc80799826d9ade4c6eea79d0ca5843 Mon Sep 17 00:00:00 2001
From: Michael Gernoth <michael@gernoth.net>
Date: Sun, 6 Jun 2010 12:23:22 +0200
Subject: [PATCH] error checks, bug fixes

---
 png.c   |  9 +++++++--
 rigol.c | 20 ++++++++++++++++++--
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/png.c b/png.c
index 7606f2b..a85e7d7 100644
--- a/png.c
+++ b/png.c
@@ -68,7 +68,7 @@ unsigned char *lcd2png(unsigned char *lcd, int *len)
 	unsigned char lut[256][3];
 	unsigned char *screen_deflated;
 	int deflated_size;
-	unsigned char *image;
+	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'};
@@ -117,6 +117,7 @@ unsigned char *lcd2png(unsigned char *lcd, int *len)
 	strm.next_in = screen_conv;
 
 	toalloc = 0;
+	strm.avail_out = 0;
 	screen_deflated = NULL;
 	flush = Z_NO_FLUSH;
 
@@ -147,7 +148,6 @@ unsigned char *lcd2png(unsigned char *lcd, int *len)
 
 	deflateEnd(&strm);
 
-
 	image = malloc(sizeof(png) +
 		sizeof(ihdr) + 8 + /* 8 = length, csum */
 		sizeof(idat) + deflated_size + 8 +
@@ -158,9 +158,12 @@ unsigned char *lcd2png(unsigned char *lcd, int *len)
 	}
 	
 	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);
@@ -170,6 +173,7 @@ unsigned char *lcd2png(unsigned char *lcd, int *len)
 	memcpy(outpos, &l, sizeof(l));
 	outpos += sizeof(l);
 
+	/* IDAT */
 	l = htonl(deflated_size);
 	memcpy(outpos, &l, sizeof(l));
 	outpos += sizeof(l);
@@ -182,6 +186,7 @@ unsigned char *lcd2png(unsigned char *lcd, int *len)
 	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);
diff --git a/rigol.c b/rigol.c
index 818b7ea..0e9ff9e 100644
--- a/rigol.c
+++ b/rigol.c
@@ -306,6 +306,7 @@ void do_get_screen(struct usb_dev_handle *sc)
 	char filename[256];
 	unsigned char *png;
 	int imglen;
+	int ret;
 	int l;
 	int fd;
 	pid_t display;
@@ -320,11 +321,26 @@ void do_get_screen(struct usb_dev_handle *sc)
 		printf ("hmm. didnt' get %d bytes, but %d\n\n", sizeof(screen), l); 
 	}
 
+	png = lcd2png(screen, &imglen);
+
+	lt = time(NULL);
 	strftime(filename, sizeof(filename), "screen_%Y%m%d-%H%M%S.png", localtime(&lt));
 	fd=open(filename, O_CREAT|O_WRONLY, 0644);
-	png = lcd2png(screen, &imglen);
-	write(fd, png, imglen);
+	if (fd == -1) {
+		perror("open");
+		exit(EXIT_FAILURE);
+	}
+
+	while(imglen > 0) {
+		ret = write(fd, png, imglen);
+		if (ret == -1) {
+			perror("write");
+			exit(EXIT_FAILURE);
+		}
+		imglen -= ret;
+	}
 	close(fd);
+	free(png);
 
 	printf("Waveform saved as %s\n", filename);
 
-- 
2.39.2