1 files changed, 258 insertions, 184 deletions

diff --git a/recv_image.c b/recv_image.c
index f9705f6..d361e67 100644
--- a/recv_image.c
+++ b/recv_image.c
@@ -14,12 +14,23 @@
 #include <sys/socket.h>
 #include <netinet/in.h>
 #include <sys/ioctl.h>
+#include <sys/time.h>
 #include "crc32.h"
 #include "mtd/mtd-user.h"
 #include "mcast_image.h"
 
 #define min(x,y) (  (x)>(y)?(y):(x) )
 
+#define WBUF_SIZE 2048
+struct eraseblock {
+	uint32_t flash_offset;
+	unsigned char wbuf[WBUF_SIZE];
+	int wbuf_ofs;
+	int nr_pkts;
+	int *pkt_indices;
+	uint32_t crc;
+};
+
 int main(int argc, char **argv)
 {
 	struct addrinfo *ai;
@@ -30,23 +41,22 @@ int main(int argc, char **argv)
 	size_t len;
 	int flfd;
 	struct mtd_info_user meminfo;
-	unsigned char *eb_buf;
-	unsigned char *blockmap = NULL;
+	unsigned char *eb_buf, *decode_buf, **src_pkts;
 	int nr_blocks = 0;
-	int *pkt_indices;
-	unsigned char **pkts;
-	int nr_pkts = 0;
 	int pkts_per_block;
 	int block_nr = -1;
 	uint32_t image_crc;
-	uint32_t blocks_received = 0;
+	int total_pkts = 0;
 	loff_t mtdoffset = 0;
-	int *stats;
 	int badcrcs = 0;
 	int duplicates = 0;
 	int file_mode = 0;
 	struct fec_parms *fec;
 	int i;
+	struct eraseblock *eraseblocks = NULL;
+	uint32_t start_seq;
+	struct timeval start, now;
+	unsigned long fec_time = 0, flash_time = 0, crc_time = 0;
 
 	if (argc != 4) {
 		fprintf(stderr, "usage: %s <host> <port> <mtddev>\n",
@@ -69,7 +79,7 @@ int main(int argc, char **argv)
 	}
 	if (flfd == -1) {
 		/* Try again, as if it's a file */
-		flfd = open(argv[3], O_CREAT|O_TRUNC|O_WRONLY, 0644);
+		flfd = open(argv[3], O_CREAT|O_TRUNC|O_RDWR, 0644);
 		if (flfd < 0) {
 			perror("open");
 			exit(1);
@@ -83,31 +93,16 @@ int main(int argc, char **argv)
 	pkts_per_block = (meminfo.erasesize + PKT_SIZE - 1) / PKT_SIZE;
 
 	eb_buf = malloc(pkts_per_block * PKT_SIZE);
-	if (!eb_buf) {
+	decode_buf = malloc(pkts_per_block * PKT_SIZE);
+	if (!eb_buf && !decode_buf) {
 		fprintf(stderr, "No memory for eraseblock buffer\n");
 		exit(1);
 	}
-
-	pkt_indices = malloc(sizeof(int) * pkts_per_block);
-	if (!pkt_indices) {
-		fprintf(stderr, "No memory for packet indices\n");
+	src_pkts = malloc(sizeof(unsigned char *) * pkts_per_block);
+	if (!src_pkts) {
+		fprintf(stderr, "No memory for decode packet pointers\n");
 		exit(1);
 	}
-	memset(pkt_indices, 0, sizeof(int) * pkts_per_block);
-
-	pkts = malloc(sizeof(unsigned char *) * pkts_per_block);
-	if (!pkts) {
-		fprintf(stderr, "No memory for packet pointers\n");
-		exit(1);
-	}
-	for (i=0; i<pkts_per_block; i++) {
-		pkts[i] = malloc(sizeof(struct image_pkt_hdr) + PKT_SIZE);
-		if (!pkts[i]) {
-			printf("No memory for packets\n");
-			exit(1);
-		}
-		pkts[i] += sizeof(struct image_pkt_hdr);
-	}
 
 	memset(&hints, 0, sizeof(hints));
 	hints.ai_flags = AI_ADDRCONFIG;
@@ -159,155 +154,263 @@ int main(int argc, char **argv)
 		exit(1);
 
 	while (1) {
-		struct image_pkt *thispkt;
+		struct image_pkt thispkt;
 
-		if (nr_pkts < pkts_per_block)
-			thispkt = (void *)(pkts[nr_pkts] - sizeof(struct image_pkt_hdr));
-		else
-			thispkt = (void *)(pkts[0] - sizeof(struct image_pkt_hdr));
-
-		len = read(sock, thispkt, sizeof(*thispkt));
+		len = read(sock, &thispkt, sizeof(thispkt));
 
 		if (len < 0) {
 			perror("read socket");
 			break;
 		}
-		if (len < sizeof(*thispkt)) {
+		if (len < sizeof(thispkt)) {
 			fprintf(stderr, "Wrong length %d bytes (expected %d)\n",
-				len, sizeof(*thispkt));
+				len, sizeof(thispkt));
 			continue;
 		}
-		if (!blockmap) {
-			image_crc = thispkt->hdr.totcrc;
-			if (meminfo.erasesize != ntohl(thispkt->hdr.blocksize)) {
+		if (!eraseblocks) {
+			image_crc = thispkt.hdr.totcrc;
+			start_seq = ntohl(thispkt.hdr.pkt_sequence);
+
+			if (meminfo.erasesize != ntohl(thispkt.hdr.blocksize)) {
 				fprintf(stderr, "Erasesize mismatch (0x%x not 0x%x)\n",
-					ntohl(thispkt->hdr.blocksize), meminfo.erasesize);
+					ntohl(thispkt.hdr.blocksize), meminfo.erasesize);
 				exit(1);
 			}
-			nr_blocks = ntohl(thispkt->hdr.nr_blocks);
-			nr_pkts = 0;
+			nr_blocks = ntohl(thispkt.hdr.nr_blocks);
 
-			fec = fec_new(pkts_per_block, ntohs(thispkt->hdr.nr_pkts));
+			fec = fec_new(pkts_per_block, ntohs(thispkt.hdr.nr_pkts));
 
-			blockmap = malloc(nr_blocks);
-			if (!blockmap) {
+			eraseblocks = malloc(nr_blocks * sizeof(*eraseblocks));
+			if (!eraseblocks) {
 				fprintf(stderr, "No memory for block map\n");
 				exit(1);
 			}
-			memset(blockmap, 0, nr_blocks);
-			stats = malloc(sizeof(int) * (ntohs(thispkt->hdr.nr_pkts) + 1));
-			if (!stats) {
-				fprintf(stderr, "No memory for statistics\n");
-				exit(1);
-			}
-			memset(stats, 0, sizeof(int) * (ntohs(thispkt->hdr.nr_pkts) + 1));
-		}
-		if (image_crc != thispkt->hdr.totcrc) {
-			fprintf(stderr, "Image CRC changed from 0x%x to 0x%x. Aborting\n",
-				ntohl(image_crc), ntohl(thispkt->hdr.totcrc));
-			exit(1);
-		}
-		if (ntohl(thispkt->hdr.block_nr) != block_nr) {
-			/* Hm, new block */
-			if (block_nr != -1) {
-				if (!blockmap[block_nr]) {
-					printf("Lost image block %08x with only %d/%d (%d) packets\n",
-					       block_nr * meminfo.erasesize, nr_pkts, pkts_per_block,
-					       ntohs(thispkt->hdr.nr_pkts));
+			for (i = 0; i < nr_blocks; i++) {
+				eraseblocks[i].pkt_indices = malloc(sizeof(int) * pkts_per_block);
+				if (!eraseblocks[i].pkt_indices) {
+					fprintf(stderr, "Failed to allocate packet indices\n");
+					exit(1);
 				}
-				if (blockmap[block_nr] < 2) {
-					stats[nr_pkts]++;
-					if (blockmap[block_nr]) {
-						if (file_mode)
-							printf(" with %d/%d (%d) packets\n",
-							       nr_pkts, pkts_per_block,
-							       ntohs(thispkt->hdr.nr_pkts));
-						blockmap[block_nr] = 2;
+				eraseblocks[i].nr_pkts = 0;
+				if (!file_mode) {
+					if (mtdoffset >= meminfo.size) {
+						fprintf(stderr, "Run out of space on flash\n");
+						exit(1);
 					}
+#if 1 /* Deliberately use bad blocks... test write failures */
+					while (ioctl(flfd, MEMGETBADBLOCK, &mtdoffset) > 0) {
+						printf("Skipping flash bad block at %08x\n", (uint32_t)mtdoffset);
+						mtdoffset += meminfo.erasesize;
+					}
+#endif
 				}
+				eraseblocks[i].flash_offset = mtdoffset;
+				mtdoffset += meminfo.erasesize;
+				eraseblocks[i].wbuf_ofs = 0;
 			}
-			/* Put this packet first */
-			if (nr_pkts != 0 && nr_pkts < pkts_per_block) {
-				unsigned char *tmp = pkts[0];
-				pkts[0] = pkts[nr_pkts];
-				pkts[nr_pkts] = tmp;
-
-			}
-			nr_pkts = 0;
-
-			block_nr = ntohl(thispkt->hdr.block_nr);
-
-			if (block_nr > nr_blocks) {
-				fprintf(stderr, "Erroneous block_nr %d (> %d)\n",
-					block_nr, nr_blocks);
-				exit(1);
-			}
-			if (blockmap[block_nr]) {
-				printf("Discard chunk at 0x%08x for already-flashed eraseblock (%d to go)\n",
-				       block_nr * meminfo.erasesize, nr_blocks - blocks_received);
-				continue;
-			}
+			gettimeofday(&start, NULL);
 		}
-		if (nr_pkts >= pkts_per_block) {
-			/* We have a parity block but we didn't need it */
-			nr_pkts++;
-			continue;
+		if (image_crc != thispkt.hdr.totcrc) {
+			fprintf(stderr, "Image CRC changed from 0x%x to 0x%x. Aborting\n",
+				ntohl(image_crc), ntohl(thispkt.hdr.totcrc));
+			exit(1);
 		}
-		if (blockmap[block_nr])
-			continue;
 
-		for (i=0; i < nr_pkts; i++) {
-			if (pkt_indices[i] == ntohs(thispkt->hdr.pkt_nr)) {
-				printf("Discarding duplicate packet at %08x pkt %d\n",
-				       block_nr * meminfo.erasesize, pkt_indices[i]);
+		block_nr = ntohl(thispkt.hdr.block_nr);
+		if (block_nr >= nr_blocks) {
+			fprintf(stderr, "Erroneous block_nr %d (> %d)\n",
+				block_nr, nr_blocks);
+			exit(1);
+		}
+		for (i=0; i<eraseblocks[block_nr].nr_pkts; i++) {
+			if (eraseblocks[block_nr].pkt_indices[i] == ntohs(thispkt.hdr.pkt_nr)) {
+//				printf("Discarding duplicate packet at %08x pkt %d\n",
+//				       block_nr * meminfo.erasesize, eraseblocks[block_nr].pkt_indices[i]);
 				duplicates++;
 				break;
 			}
-		} /* And if we broke out, skip the packet... */
-		if (i < nr_pkts)
+		}
+		if (i < eraseblocks[block_nr].nr_pkts)
 			continue;
 
-		if (crc32(-1, thispkt->data, PKT_SIZE) != ntohl(thispkt->hdr.thiscrc)) {
+		if (eraseblocks[block_nr].nr_pkts >= pkts_per_block) {
+			/* We have a block which we didn't really need */
+			eraseblocks[block_nr].nr_pkts++;
+			continue;
+		}
+
+		if (crc32(-1, thispkt.data, PKT_SIZE) != ntohl(thispkt.hdr.thiscrc)) {
 			printf("Discard %08x pkt %d with bad CRC (%08x not %08x)\n",
-			       block_nr * meminfo.erasesize, ntohs(thispkt->hdr.pkt_nr),
-			       crc32(-1, thispkt->data, PKT_SIZE),
-			       ntohl(thispkt->hdr.thiscrc));
+			       block_nr * meminfo.erasesize, ntohs(thispkt.hdr.pkt_nr),
+			       crc32(-1, thispkt.data, PKT_SIZE),
+			       ntohl(thispkt.hdr.thiscrc));
 			badcrcs++;
 			continue;
 		}
+	pkt_again:
+		eraseblocks[block_nr].pkt_indices[eraseblocks[block_nr].nr_pkts++] = 
+			ntohs(thispkt.hdr.pkt_nr);
+		total_pkts++;
+		if (!(total_pkts % 50) || total_pkts == pkts_per_block * nr_blocks) {
+			uint32_t pkts_sent = ntohl(thispkt.hdr.pkt_sequence) - start_seq - 1;
+			long time_msec;
+			gettimeofday(&now, NULL);
 
-		pkt_indices[nr_pkts] = ntohs(thispkt->hdr.pkt_nr);
-		nr_pkts++;
+			time_msec = ((now.tv_usec - start.tv_usec) / 1000) +
+				(now.tv_sec - start.tv_sec) * 1000;
 
-		if (nr_pkts == pkts_per_block) {
+			printf("\rReceived %d/%d (%d%%) in %lds @%ldKiB/s, %d lost (%d%%), %d dups    ",
+			       total_pkts, nr_blocks * pkts_per_block,
+			       total_pkts * 100 / nr_blocks / pkts_per_block,
+			       time_msec / 1000,
+			       total_pkts * PKT_SIZE / 1024 * 1000 / time_msec,
+			       pkts_sent - total_pkts - duplicates,
+			       (pkts_sent - total_pkts - duplicates) * 100 / pkts_sent,
+			       duplicates);
+		}
 
-			if (fec_decode(fec, pkts, pkt_indices, PKT_SIZE)) {
-				/* Eep. This cannot happen */
-				printf("The world is broken. fec_decode() returned error\n");
-				exit(1);
+		if (eraseblocks[block_nr].wbuf_ofs + PKT_SIZE < WBUF_SIZE) {
+			/* New packet doesn't full the wbuf */
+			memcpy(eraseblocks[block_nr].wbuf + eraseblocks[block_nr].wbuf_ofs,
+			       thispkt.data, PKT_SIZE);
+			eraseblocks[block_nr].wbuf_ofs += PKT_SIZE;
+		} else {
+			int fits = WBUF_SIZE - eraseblocks[block_nr].wbuf_ofs;
+			ssize_t wrotelen;
+			memcpy(eraseblocks[block_nr].wbuf + eraseblocks[block_nr].wbuf_ofs,
+			       thispkt.data, fits);
+			wrotelen = pwrite(flfd, eraseblocks[block_nr].wbuf, WBUF_SIZE,
+					  eraseblocks[block_nr].flash_offset);
+			
+			if (wrotelen < WBUF_SIZE) {
+				if (wrotelen < 0)
+					perror("packet write");
+				else
+					fprintf(stderr, "short write of packet wbuf\n");
+
+				if (!file_mode) {
+					struct erase_info_user erase;
+					/* FIXME: Perhaps we should store pkt crcs and try
+					   to recover data from the offending eraseblock */
+					erase.start = eraseblocks[block_nr].flash_offset;
+					erase.start -= eraseblocks[block_nr].nr_pkts * PKT_SIZE / WBUF_SIZE * WBUF_SIZE;
+					erase.length = meminfo.erasesize;
+					
+					if (ioctl(flfd, MEMERASE, &erase)) {
+						perror("MEMERASE");
+						exit(1);
+					}
+					if (mtdoffset >= meminfo.size) {
+						fprintf(stderr, "Run out of space on flash\n");
+						exit(1);
+					}
+					while (ioctl(flfd, MEMGETBADBLOCK, &mtdoffset) > 0) {
+						printf("Skipping flash bad block at %08x\n", (uint32_t)mtdoffset);
+						mtdoffset += meminfo.erasesize;
+					}
+					eraseblocks[block_nr].flash_offset = mtdoffset;
+					total_pkts -= eraseblocks[block_nr].nr_pkts;
+					eraseblocks[block_nr].nr_pkts = 0;
+					goto pkt_again;
+				}
+				else /* Usually nothing we can do in file mode */
+					exit(1);
 			}
-			blockmap[block_nr] = 1;
-			blocks_received++;
-
-			/* Put data into order in eb_buf */
-			for (i=0; i < pkts_per_block; i++)
-				memcpy(eb_buf + (i * PKT_SIZE), pkts[i], PKT_SIZE);
-
-			if (crc32(-1, eb_buf, meminfo.erasesize) != ntohl(thispkt->hdr.block_crc)) {
-				printf("FEC error. CRC %08x != %08x\n", 
-				       crc32(-1, eb_buf, meminfo.erasesize),
-				       ntohl(thispkt->hdr.block_crc));
-				*(int *)0 = 0;
+			eraseblocks[block_nr].flash_offset += WBUF_SIZE;
+			/* Copy the remainder into the wbuf */
+			memcpy(eraseblocks[block_nr].wbuf, &thispkt.data[fits], PKT_SIZE - fits);
+			eraseblocks[block_nr].wbuf_ofs = PKT_SIZE - fits;
+		}
+		
+		if (eraseblocks[block_nr].nr_pkts == pkts_per_block) {
+			eraseblocks[block_nr].crc = ntohl(thispkt.hdr.block_crc);
+
+			if (total_pkts == nr_blocks * pkts_per_block)
+				break;
+		}
+	}
+	close(sock);
+	for (block_nr = 0; block_nr < nr_blocks; block_nr++) {
+		ssize_t rwlen;
+
+		eraseblocks[block_nr].flash_offset -= meminfo.erasesize;
+		rwlen = pread(flfd, eb_buf, meminfo.erasesize, eraseblocks[block_nr].flash_offset);
+
+		if (rwlen < 0) {
+			perror("read");
+			/* Argh. Perhaps we could go back and try again, but if the flash is
+			   going to fail to read back what we write to it, and the whole point
+			   in this program is to write to it, what's the point? */
+			fprintf(stderr, "Packets we wrote to flash seem to be unreadable. Aborting\n");
+			exit(1);
+		}
+
+		memcpy(eb_buf + meminfo.erasesize, eraseblocks[block_nr].wbuf,
+		       eraseblocks[block_nr].wbuf_ofs);
+
+		for (i=0; i < pkts_per_block; i++)
+			src_pkts[i] = &eb_buf[i * PKT_SIZE];
+
+		gettimeofday(&start, NULL);
+		if (fec_decode(fec, src_pkts, eraseblocks[block_nr].pkt_indices, PKT_SIZE)) {
+			/* Eep. This cannot happen */
+			printf("The world is broken. fec_decode() returned error\n");
+			exit(1);
+		}
+		gettimeofday(&now, NULL);
+		fec_time += (now.tv_usec - start.tv_usec) / 1000;
+		fec_time += (now.tv_sec - start.tv_sec) * 1000;
+		
+		for (i=0; i < pkts_per_block; i++)
+			memcpy(&decode_buf[i*PKT_SIZE], src_pkts[i], PKT_SIZE);
+
+		/* Paranoia */
+		gettimeofday(&start, NULL);
+		if (crc32(-1, decode_buf, meminfo.erasesize) != eraseblocks[block_nr].crc) {
+			printf("CRC mismatch: want %08x got %08x\n",
+			       eraseblocks[block_nr].crc, 
+			       crc32(-1, decode_buf, meminfo.erasesize));
+		}
+		gettimeofday(&now, NULL);
+		crc_time += (now.tv_usec - start.tv_usec) / 1000;
+		crc_time += (now.tv_sec - start.tv_sec) * 1000;
+		start = now;
+
+		if (!file_mode) {
+			struct erase_info_user erase;
+
+			erase.start = eraseblocks[block_nr].flash_offset;
+			erase.length = meminfo.erasesize;
+			
+			printf("Erasing block at %08x...", erase.start);
+
+			if (ioctl(flfd, MEMERASE, &erase)) {
+				perror("MEMERASE");
+				/* This block has dirty data on it. If the erase failed, we're screwed */
+				fprintf(stderr, "Erase to clean FEC data from flash failed. Aborting\n");
 				exit(1);
 			}
-			if (file_mode) {
-				printf("Received image block %08x (%d/%d)",
-				       block_nr * meminfo.erasesize,
-				       blocks_received, nr_blocks);
-				pwrite(flfd, eb_buf, meminfo.erasesize, block_nr * meminfo.erasesize);
-			} else {
-				ssize_t wrotelen;
-			again:
+		}
+	write_again:
+		rwlen = pwrite(flfd, decode_buf, meminfo.erasesize, eraseblocks[block_nr].flash_offset);
+		if (rwlen < meminfo.erasesize) {
+			if (rwlen < 0) {
+				perror("decoded data write");
+			} else 
+				fprintf(stderr, "short write of decoded data\n");
+
+			if (!file_mode) {
+				struct erase_info_user erase;
+				erase.start = eraseblocks[block_nr].flash_offset;
+				erase.length = meminfo.erasesize;
+
+				printf("Erasing failed block at %08x\n",
+				       eraseblocks[block_nr].flash_offset);
+
+				if (ioctl(flfd, MEMERASE, &erase)) {
+					perror("MEMERASE");
+					exit(1);
+				}
 				if (mtdoffset >= meminfo.size) {
 					fprintf(stderr, "Run out of space on flash\n");
 					exit(1);
@@ -316,54 +419,25 @@ int main(int argc, char **argv)
 					printf("Skipping flash bad block at %08x\n", (uint32_t)mtdoffset);
 					mtdoffset += meminfo.erasesize;
 				}
-				wrotelen = pwrite(flfd, eb_buf, meminfo.erasesize, mtdoffset);
-				if (wrotelen != meminfo.erasesize) {
-					struct erase_info_user erase;
-					
-					if (wrotelen < 0)
-						perror("flash write");
-					else
-						fprintf(stderr, "Short write to flash at %08x: %zd bytes\n",
-							(uint32_t)mtdoffset, wrotelen);
-					
-					erase.start = mtdoffset;
-					erase.length = meminfo.erasesize;
-					
-					if (ioctl(flfd, MEMERASE, erase)) {
-						perror("MEMERASE");
-						exit(1);
-					}
-					/* skip it */
-					//				ioctl(flfd, MEMSETBADBLOCK, &mtdoffset);
-					mtdoffset += meminfo.erasesize;
-					goto again;
-				}
-				printf("Wrote image block %08x (%d/%d) to flash offset %08x\n",
-				       block_nr * meminfo.erasesize, 
-				       blocks_received, nr_blocks,
-				       (uint32_t)mtdoffset);
-				mtdoffset += meminfo.erasesize;
-			}
-			if (!(blocks_received%100) || blocks_received == nr_blocks) {
-				int i, printed = 0;
-				printf("\n");
-				for (i=0; i <= ntohs(thispkt->hdr.nr_pkts); i++) {
-					if (printed || stats[i]) {
-						printf("Number of blocks with %d packets received: %d\n",
-						       i, stats[i]);
-						printed = 1;
-					}
-				}
-				printf("Bad CRC: %d\n", badcrcs);
-				printf("Duplicate: %d\n", duplicates);
+				printf("Will try again at %08lx...", (long)mtdoffset);
+				eraseblocks[block_nr].flash_offset = mtdoffset;
 
+				goto write_again;
 			}
-			if (blocks_received == nr_blocks) {
-				printf("Got all %08x bytes of image. Bye!\n",
-				       nr_blocks * meminfo.erasesize);
-				exit(0);
-			}
+			else /* Usually nothing we can do in file mode */
+				exit(1);
 		}
+		gettimeofday(&now, NULL);
+		flash_time += (now.tv_usec - start.tv_usec) / 1000;
+		flash_time += (now.tv_sec - start.tv_sec) * 1000;
+
+		printf("wrote image block %08x (%d pkts)\n",
+		       block_nr * meminfo.erasesize, eraseblocks[block_nr].nr_pkts);
 	}
-	close(sock);
+	close(flfd);
+	printf("FEC time %ld.%03lds\n", fec_time / 1000, fec_time % 1000);
+	printf("CRC time %ld.%03lds\n", crc_time / 1000, crc_time % 1000);
+	printf("flash IO %ld.%03lds\n", flash_time / 1000, flash_time % 1000);
+
+	return 0;
 }