Add utilities for multicast send/receive of MTD images.

Signed-off-by: David Woodhouse <dwmw2@infradead.org>

4 files changed, 566 insertions, 0 deletions

diff --git a/Makefile b/Makefile
index 4f74158..2283c78 100644
--- a/Makefile
+++ b/Makefile
@@ -25,6 +25,7 @@ RAWTARGETS = ftl_format flash_erase flash_eraseall nanddump doc_loadbios \
 	jffs2dump \
 	nftldump nftl_format docfdisk \
 	rfddump rfdformat \
+	serve_image recv_image \
 	sumtool #jffs2reader
 
 TARGETS = $(foreach target,$(RAWTARGETS),$(BUILDDIR)/$(target))
diff --git a/mcast_image.h b/mcast_image.h
new file mode 100644
index 0000000..e2048b6
--- /dev/null
+++ b/mcast_image.h
@@ -0,0 +1,18 @@
+#include <stdint.h>
+
+#define PKT_SIZE 1400
+
+struct image_pkt_hdr {
+	uint32_t totcrc;
+	uint32_t nr_blocks;
+	uint32_t blocksize;
+	uint32_t block_nr;
+	uint32_t block_ofs;
+	uint32_t thislen;
+	uint32_t thiscrc;
+};
+
+struct image_pkt {
+	struct image_pkt_hdr hdr;
+	unsigned char data[PKT_SIZE];
+};
diff --git a/recv_image.c b/recv_image.c
new file mode 100644
index 0000000..d2b8813
--- /dev/null
+++ b/recv_image.c
@@ -0,0 +1,356 @@
+
+#define _XOPEN_SOURCE 500
+
+#include <errno.h>  	
+#include <error.h> 	
+#include <netdb.h> 	
+#include <stdio.h> 	
+#include <stdlib.h> 	
+#include <string.h>
+#include <unistd.h> 	
+#include <fcntl.h>
+#include <sys/types.h>
+#include <sys/stat.h>
+#include <sys/socket.h>
+#include <netinet/in.h>
+#include <sys/ioctl.h>
+#include "crc32.h"
+#include "mtd/mtd-user.h"
+#include "mcast_image.h"
+
+#define min(x,y) (  (x)>(y)?(y):(x) )
+
+int main(int argc, char **argv)
+{
+	struct sockaddr_storage server_addr;
+	socklen_t server_addrlen = sizeof(server_addr);
+	struct addrinfo *ai;
+	struct addrinfo hints;
+	struct addrinfo *runp;
+	int ret;
+	int sock;
+	struct image_pkt pktbuf;
+	size_t len;
+	int flfd;
+	struct mtd_info_user meminfo;
+	unsigned char *eb_buf;
+	unsigned char *blockmap = NULL;
+	unsigned char *subblockmap;
+	int nr_blocks = 0;
+	int nr_subblocks = 0;
+	int pkts_per_block;
+	int block_nr = -1;
+	uint32_t image_crc;
+	uint32_t blocks_received = 0;
+	uint32_t block_ofs;
+	loff_t mtdoffset = 0;
+	int *stats;
+	int badcrcs = 0;
+	int duplicates = 0;
+	int missing = -1;
+	int file_mode = 0;
+
+	if (argc != 4) {
+		fprintf(stderr, "usage: %s <host> <port> <mtddev>\n",
+			(strrchr(argv[0], '/')?:argv[0]-1)+1);
+		exit(1);
+	}
+	/* Open the device */
+	flfd = open(argv[3], O_RDWR);
+
+	if (flfd >= 0) {
+		/* Fill in MTD device capability structure */
+		if (ioctl(flfd, MEMGETINFO, &meminfo) != 0) {
+			perror("MEMGETINFO");
+			close(flfd);
+			flfd = -1;
+		} else {
+			printf("Receive to MTD device %s with erasesize %d\n",
+			       argv[3], meminfo.erasesize);
+		}
+	}
+	if (flfd == -1) {
+		/* Try again, as if it's a file */
+		flfd = open(argv[3], O_CREAT|O_TRUNC|O_WRONLY, 0644);
+		if (flfd < 0) {
+			perror("open");
+			exit(1);
+		}
+		meminfo.erasesize = 131072;
+		file_mode = 1;
+		printf("Receive to file %s with (assumed) erasesize %d\n",
+		       argv[3], meminfo.erasesize);
+	}
+
+	pkts_per_block = (meminfo.erasesize + PKT_SIZE - 1) / PKT_SIZE;
+
+	stats = malloc(pkts_per_block + 1);
+	if (!stats) {
+		fprintf(stderr, "No memory for statistics\n");
+		exit(1);
+	}
+	memset(stats, 0, sizeof(int) * (pkts_per_block + 1));
+
+	eb_buf = malloc(pkts_per_block * PKT_SIZE);
+	if (!eb_buf) {
+		fprintf(stderr, "No memory for eraseblock buffer\n");
+		exit(1);
+	}
+	memset(eb_buf, 0, pkts_per_block * PKT_SIZE);
+
+	subblockmap = malloc(pkts_per_block + 1);
+	if (!subblockmap) {
+		fprintf(stderr, "No memory for subblock map\n");
+		exit(1);
+	}
+	memset(subblockmap, 0, pkts_per_block + 1);
+
+	memset(&hints, 0, sizeof(hints));
+	hints.ai_flags = AI_ADDRCONFIG;
+	hints.ai_socktype = SOCK_DGRAM;
+	
+	ret = getaddrinfo(argv[1], argv[2], &hints, &ai);
+	if (ret) {
+		fprintf(stderr, "getaddrinfo: %s\n", gai_strerror(ret));
+		exit(1);
+	}
+	runp = ai;
+	for (runp = ai; runp; runp = runp->ai_next) {
+		sock = socket(runp->ai_family, runp->ai_socktype,
+			      runp->ai_protocol);
+		if (sock == -1) {
+			perror("socket");
+			continue;
+		}
+		if (runp->ai_family == AF_INET &&
+		    IN_MULTICAST( ntohl(((struct sockaddr_in *)runp->ai_addr)->sin_addr.s_addr))) {
+			struct ip_mreq rq;
+			rq.imr_multiaddr = ((struct sockaddr_in *)runp->ai_addr)->sin_addr;
+			rq.imr_interface.s_addr = INADDR_ANY;
+			if (setsockopt(sock, IPPROTO_IP, IP_ADD_MEMBERSHIP, &rq, sizeof(rq))) {
+				perror("IP_ADD_MEMBERSHIP"); 
+				close(sock);
+				continue;
+			}
+			
+		} else if (runp->ai_family == AF_INET6 &&
+			   ((struct sockaddr_in6 *)runp->ai_addr)->sin6_addr.s6_addr[0] == 0xff) {
+			struct ipv6_mreq rq;
+			rq.ipv6mr_multiaddr =  ((struct sockaddr_in6 *)runp->ai_addr)->sin6_addr;
+			rq.ipv6mr_interface = 0;
+			if (setsockopt(sock, IPPROTO_IPV6, IPV6_ADD_MEMBERSHIP, &rq, sizeof(rq))) {
+				perror("IPV6_ADD_MEMBERSHIP"); 
+				close(sock);
+				continue;
+			}
+		} else printf("not multicast?\n");
+		if (bind(sock, runp->ai_addr, runp->ai_addrlen)) {
+			perror("bind");
+			close(sock);
+			continue;
+		}
+		break;
+	}
+	if (!runp)
+		exit(1);
+
+	while ((len = read(sock, &pktbuf, sizeof(pktbuf))) >= 0) {
+		if (len < sizeof(pktbuf.hdr)) {
+			fprintf(stderr, "Short read %d bytes\n", len);
+			continue;
+		}
+		if (len != sizeof(pktbuf.hdr) + ntohl(pktbuf.hdr.thislen)) {
+			fprintf(stderr, "Wrong length %d bytes (expected %d+%d)\n",
+				len, sizeof(pktbuf.hdr), ntohl(pktbuf.hdr.thislen));
+			continue;
+		}
+		/* Holds _data_ length */
+		len -= sizeof(pktbuf.hdr);
+			
+		if (!blockmap) {
+			image_crc = pktbuf.hdr.totcrc;
+			if (meminfo.erasesize != ntohl(pktbuf.hdr.blocksize)) {
+				fprintf(stderr, "Erasesize mismatch (0x%x not 0x%x)\n",
+					ntohl(pktbuf.hdr.blocksize), meminfo.erasesize);
+				exit(1);
+			}
+			nr_blocks = ntohl(pktbuf.hdr.nr_blocks);
+			nr_subblocks = pkts_per_block + 2;
+			blockmap = malloc(nr_blocks);
+			if (!blockmap) {
+				fprintf(stderr, "No memory for block map\n");
+				exit(1);
+			}
+			memset(blockmap, 0, nr_blocks);
+		}
+		if (image_crc != pktbuf.hdr.totcrc) {
+			fprintf(stderr, "Image CRC changed from 0x%x to 0x%x. Aborting\n",
+				ntohl(image_crc), ntohl(pktbuf.hdr.totcrc));
+			exit(1);
+		}
+		if (ntohl(pktbuf.hdr.block_nr) != block_nr) {
+			/* Hm, new block */
+			if (nr_subblocks < pkts_per_block &&
+			    block_nr != -1) 
+				printf("Lost image block at %08x with only %d/%d packets\n",
+				       block_nr * meminfo.erasesize, nr_subblocks,
+				       pkts_per_block + 1);
+
+
+			if (nr_subblocks < pkts_per_block + 2)
+				stats[nr_subblocks]++;
+
+			nr_subblocks = 0;
+			missing = -1;
+			memset(subblockmap, 0, pkts_per_block + 1);
+			block_nr = ntohl(pktbuf.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);
+				nr_subblocks = pkts_per_block + 2;
+				continue;
+			}
+		}
+		if (nr_subblocks == pkts_per_block) {
+			/* We have a parity block but we didn't need it */
+			nr_subblocks++;
+			continue;
+		}
+		if (blockmap[block_nr])
+			continue;
+
+		block_ofs = ntohl(pktbuf.hdr.block_ofs);
+		if (block_ofs == meminfo.erasesize)
+			block_ofs = PKT_SIZE * pkts_per_block;
+
+		if (len != PKT_SIZE && len + block_ofs != meminfo.erasesize) {
+			fprintf(stderr, "Bogus packet size 0x%x (expected 0x%x)\n",
+				ntohl(pktbuf.hdr.thislen),
+				min(PKT_SIZE, meminfo.erasesize - block_ofs));
+			exit(1);
+		}
+
+		if (crc32(-1, pktbuf.data, len) != ntohl(pktbuf.hdr.thiscrc)) {
+			printf("Discard chunk %08x with bad CRC (%08x not %08x)\n",
+			       block_nr * meminfo.erasesize + block_ofs,
+			       crc32(-1, pktbuf.data, pktbuf.hdr.thislen),
+			       ntohl(pktbuf.hdr.thiscrc));
+			badcrcs++;
+			continue;
+		}
+		if (subblockmap[block_ofs / PKT_SIZE]) {
+			printf("Discarding duplicate packet at %08x\n",
+			       block_nr * meminfo.erasesize + block_ofs);
+			duplicates++;
+			continue;
+		}
+		subblockmap[block_ofs / PKT_SIZE] = 1;
+		nr_subblocks++;
+		if (block_ofs < meminfo.erasesize) {
+			/* Normal data packet */
+			memcpy(eb_buf + block_ofs, pktbuf.data, len);
+//			printf("Received data block at %08x\n", block_nr * meminfo.erasesize + block_ofs);
+		} else {
+			/* Parity block */
+			int i;
+
+			/* If we don't have enough to recover, skip */
+			if (nr_subblocks < pkts_per_block)
+				continue;
+
+			for (i = 0; i<pkts_per_block; i++) {
+				if (subblockmap[i]) {
+					int j;
+					for (j=0; j<PKT_SIZE; j++)
+						pktbuf.data[j] ^= eb_buf[i*PKT_SIZE + j];
+				} else
+					missing = i;
+			}
+
+			if (missing == -1) {
+				fprintf(stderr, "dwmw2 is stupid\n");
+				exit(1);
+			}
+//			printf("Recover missing packet at %08x from parity\n",
+//			       block_nr * meminfo.erasesize + missing * PKT_SIZE);
+			memcpy(eb_buf + (missing * PKT_SIZE), pktbuf.data, PKT_SIZE);
+		}
+
+		if (nr_subblocks == pkts_per_block) {
+
+			blockmap[block_nr] = 1;
+			blocks_received++;
+
+			if (file_mode) {
+				printf("Received image block %08x%s (%d/%d)\n",
+				       block_nr * meminfo.erasesize,
+				       (missing==-1)?"":" (parity)",
+				       blocks_received, nr_blocks);
+				pwrite(flfd, eb_buf, meminfo.erasesize, block_nr * meminfo.erasesize);
+			} else {
+				ssize_t wrotelen;
+			again:
+				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;
+				}
+				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%s\n",
+				       block_nr * meminfo.erasesize, 
+				       blocks_received, nr_blocks,
+				       (uint32_t)mtdoffset,
+				       (missing==-1)?"":" (parity)");
+				mtdoffset += meminfo.erasesize;
+			}
+			if (!(blocks_received%100) || blocks_received == nr_blocks) {
+				int i, printed = 0;
+				for (i=0; i <= pkts_per_block + 1; 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);
+
+			}
+			if (blocks_received == nr_blocks) {
+				printf("Got all %08x bytes of image. Bye!\n",
+				       nr_blocks * meminfo.erasesize);
+				exit(0);
+			}
+		}
+	}
+	close(sock);
+}
diff --git a/serve_image.c b/serve_image.c
new file mode 100644
index 0000000..b016d63
--- /dev/null
+++ b/serve_image.c
@@ -0,0 +1,191 @@
+#include <errno.h>  	
+#include <error.h> 	
+#include <netdb.h> 	
+#include <stdio.h> 	
+#include <stdlib.h> 	
+#include <string.h>
+#include <unistd.h> 	
+#include <fcntl.h>
+#include <sys/types.h>
+#include <sys/stat.h>
+#include <sys/socket.h>
+#include <sys/mman.h>
+#include <netinet/in.h>
+#include "crc32.h"
+#include "mcast_image.h"
+
+int tx_rate = 80000;
+int pkt_delay = 12500 * PKT_SIZE / 1024;
+
+int main(int argc, char **argv)
+{
+	struct addrinfo *ai;
+	struct addrinfo hints;
+	struct addrinfo *runp;
+	int ret;
+	int sock;
+	struct image_pkt pktbuf;
+	int rfd;
+	struct stat st;
+	int writeerrors = 0;
+	long usec_per_tick = 1000000 / sysconf(_SC_CLK_TCK);
+	long delay_accum = 0;
+	uint32_t erasesize;
+	unsigned char parbuf[PKT_SIZE];
+	unsigned char *image, *blockptr;
+	uint32_t block_nr;
+	uint32_t block_ofs;
+	int nr_blocks;
+	uint32_t droppoint = -1;
+
+	if (argc == 6) {
+		tx_rate = atol(argv[5]) * 1024;
+		if (tx_rate < PKT_SIZE || tx_rate > 20000000) {
+			fprintf(stderr, "Bogus TX rate %d KiB/s\n", tx_rate);
+			exit(1);
+		}
+		argc = 5;
+	}
+	       
+	if (argc != 5) {
+		fprintf(stderr, "usage: %s <host> <port> <image> <erasesize> [<tx_rate>]\n",
+			(strrchr(argv[0], '/')?:argv[0]-1)+1);
+		exit(1);
+	}
+	pkt_delay = (PKT_SIZE * 1000000) / tx_rate;
+	printf("Inter-packet delay (avg): %dµs\n", pkt_delay);
+	printf("Transmit rate: %d KiB/s\n", tx_rate / 1024);
+
+	erasesize = atol(argv[4]);
+	if (!erasesize) {
+		fprintf(stderr, "erasesize cannot be zero\n");
+		exit(1);
+	}
+	memset(&hints, 0, sizeof(hints));
+	hints.ai_flags = AI_ADDRCONFIG;
+	hints.ai_socktype = SOCK_DGRAM;
+	
+	ret = getaddrinfo(argv[1], argv[2], &hints, &ai);
+	if (ret) {
+		fprintf(stderr, "getaddrinfo: %s\n", gai_strerror(ret));
+		exit(1);
+	}
+	runp = ai;
+	for (runp = ai; runp; runp = runp->ai_next) {
+		sock = socket(runp->ai_family, runp->ai_socktype,
+			      runp->ai_protocol);
+		if (sock == -1) {
+			perror("socket");
+			continue;
+		}
+		if (connect(sock, runp->ai_addr, runp->ai_addrlen) == 0)
+			break;
+		perror("connect");
+		close(sock);
+	}
+	if (!runp)
+		exit(1);
+
+	rfd = open(argv[3], O_RDONLY);
+	if (rfd < 0) {
+		perror("open");
+		exit(1);
+	}
+
+	if (fstat(rfd, &st)) {
+		perror("fstat");
+		exit(1);
+	}
+
+	if (st.st_size % erasesize) {
+		fprintf(stderr, "Image size %ld bytes is not a multiple of erasesize %d bytes\n",
+			st.st_size, erasesize);
+		exit(1);
+	}
+	image = mmap(NULL, st.st_size, PROT_READ, MAP_PRIVATE, rfd, 0);
+	if (image == MAP_FAILED) {
+		perror("mmap");
+		exit(1);
+	}
+
+
+	nr_blocks = st.st_size / erasesize;
+
+	pktbuf.hdr.totcrc = htonl(crc32(-1, image, st.st_size));
+	pktbuf.hdr.nr_blocks = htonl(nr_blocks);
+	pktbuf.hdr.blocksize = htonl(erasesize);
+
+	blockptr = image;
+
+	for (block_nr = 0; block_nr < nr_blocks; block_nr++) {
+		int len;
+		int dropped = 0;
+
+		pktbuf.hdr.block_nr = htonl(block_nr);
+
+		for (block_ofs = 0; block_ofs <= erasesize; block_ofs += len) {
+			int i;
+			
+			if (block_ofs + PKT_SIZE > erasesize)
+				len = erasesize - block_ofs;
+			else
+				len = PKT_SIZE;
+
+			if (block_ofs == erasesize) {
+				printf("\rSending parity block: %08x", block_nr * erasesize);
+				len = PKT_SIZE;
+				memcpy(pktbuf.data, parbuf, PKT_SIZE);
+			} else {
+				if (!block_ofs)
+					memcpy(parbuf, blockptr, PKT_SIZE);
+				else for (i=0; i < len; i++)
+					     parbuf[i] ^= blockptr[i];
+
+				memcpy(pktbuf.data, blockptr, len);
+				printf("\rSending data block at %08x",
+				       block_nr * erasesize + block_ofs);
+				blockptr += len;
+			}
+
+			fflush(stdout);
+			pktbuf.hdr.thislen = htonl(len);
+			pktbuf.hdr.block_ofs = htonl(block_ofs);
+			pktbuf.hdr.thiscrc = htonl(crc32(-1, pktbuf.data, len));
+
+			if (droppoint == block_ofs && !dropped) {
+				dropped = 1;
+				if (droppoint == 0)
+					droppoint = erasesize;
+				else if (droppoint == erasesize)
+					droppoint = ((erasesize - 1) / PKT_SIZE) * PKT_SIZE;
+				else droppoint -= PKT_SIZE;
+				printf("\nDropping data block at %08x\n", block_ofs);
+				continue;
+			}
+
+
+			if (write(sock, &pktbuf, sizeof(pktbuf.hdr)+len) < 0) {
+				perror("write");
+				writeerrors++;
+				if (writeerrors > 10) {
+					fprintf(stderr, "Too many consecutive write errors\n");
+					exit(1);
+				}
+			} else
+				writeerrors = 0;
+
+			/* Delay, if we are so far ahead of ourselves that we have at
+			   least one tick to wait. */
+			delay_accum += pkt_delay;
+			if (delay_accum >= usec_per_tick) {
+				usleep(delay_accum);
+				delay_accum = 0;
+			}
+		}
+	}
+	munmap(image, st.st_size);
+	close(rfd);
+	close(sock);
+	printf("\n");
+	return 0;
+}