dfu: Handle large transfers correctly (take #2)

index 1260c55..fb9b417 100644
@@ -82,7 +82,7 @@ int dfu_write(struct dfu_entity *dfu, void *buf, int size, int blk_seq_num)
 	       __func__, dfu->name, buf, size, blk_seq_num, dfu->offset,
 	       dfu->i_buf - dfu->i_buf_start);
 
-	if (blk_seq_num == 0) {
+	if (!dfu->inited) {
 		/* initial state */
 		dfu->crc = 0;
 		dfu->offset = 0;

@@ -90,6 +90,8 @@ int dfu_write(struct dfu_entity *dfu, void *buf, int size, int blk_seq_num)
 		dfu->i_buf_start = dfu_buf;
 		dfu->i_buf_end = dfu_buf + sizeof(dfu_buf);
 		dfu->i_buf = dfu->i_buf_start;
+
+		dfu->inited = 1;
 	}
 
 	if (dfu->i_blk_seq_num != blk_seq_num) {

@@ -97,7 +99,22 @@ int dfu_write(struct dfu_entity *dfu, void *buf, int size, int blk_seq_num)
 		       __func__, dfu->i_blk_seq_num, blk_seq_num);
 		return -1;
 	}
-	dfu->i_blk_seq_num++;
+
+	/* DFU 1.1 standard says:
+	 * The wBlockNum field is a block sequence number. It increments each
+	 * time a block is transferred, wrapping to zero from 65,535. It is used
+	 * to provide useful context to the DFU loader in the device."
+	 *
+	 * This means that it's a 16 bit counter that roll-overs at
+	 * 0xffff -> 0x0000. By having a typical 4K transfer block
+	 * we roll-over at exactly 256MB. Not very fun to debug.
+	 *
+	 * Handling rollover, and having an inited variable,
+	 * makes things work.
+	 */
+
+	/* handle rollover */
+	dfu->i_blk_seq_num = (dfu->i_blk_seq_num + 1) & 0xffff;
 
 	/* flush buffer if overflow */
 	if ((dfu->i_buf + size) > dfu->i_buf_end) {

@@ -106,6 +123,13 @@ int dfu_write(struct dfu_entity *dfu, void *buf, int size, int blk_seq_num)
 			ret = tret;
 	}
 
+	/* we should be in buffer now (if not then size too large) */
+	if ((dfu->i_buf + size) > dfu->i_buf_end) {
+		printf("%s: Wrong size! [%d] [%d] - %d\n",
+		       __func__, dfu->i_blk_seq_num, blk_seq_num, size);
+		return -1;
+	}
+
 	memcpy(dfu->i_buf, buf, size);
 	dfu->i_buf += size;
 

@@ -120,7 +144,7 @@ int dfu_write(struct dfu_entity *dfu, void *buf, int size, int blk_seq_num)
 	if (size == 0) {
 		debug("%s: DFU complete CRC32: 0x%x\n", __func__, dfu->crc);
 
-		puts("\nDownload complete\n");
+		printf("\nDownload complete (CRC32 0x%04x)\n", dfu->crc);
 
 		/* clear everything */
 		dfu->crc = 0;

@@ -129,6 +153,9 @@ int dfu_write(struct dfu_entity *dfu, void *buf, int size, int blk_seq_num)
 		dfu->i_buf_start = dfu_buf;
 		dfu->i_buf_end = dfu_buf + sizeof(dfu_buf);
 		dfu->i_buf = dfu->i_buf_start;
+
+		dfu->inited = 0;
+
 	}
 
 	return ret = 0 ? size : ret;

@@ -190,7 +217,7 @@ int dfu_read(struct dfu_entity *dfu, void *buf, int size, int blk_seq_num)
 	debug("%s: name: %s buf: 0x%p size: 0x%x p_num: 0x%x i_buf: 0x%p\n",
 	       __func__, dfu->name, buf, size, blk_seq_num, dfu->i_buf);
 
-	if (blk_seq_num == 0) {
+	if (!dfu->inited) {
 		ret = dfu->read_medium(dfu, 0, NULL, &dfu->r_left);
 		if (ret != 0) {
 			debug("%s: failed to get r_left\n", __func__);

@@ -206,6 +233,8 @@ int dfu_read(struct dfu_entity *dfu, void *buf, int size, int blk_seq_num)
 		dfu->i_buf_end = dfu_buf + sizeof(dfu_buf);
 		dfu->i_buf = dfu->i_buf_start;
 		dfu->b_left = 0;
+
+		dfu->inited = 1;
 	}
 
 	if (dfu->i_blk_seq_num != blk_seq_num) {

@@ -213,7 +242,8 @@ int dfu_read(struct dfu_entity *dfu, void *buf, int size, int blk_seq_num)
 		       __func__, dfu->i_blk_seq_num, blk_seq_num);
 		return -1;
 	}
-	dfu->i_blk_seq_num++;
+	/* handle rollover */
+	dfu->i_blk_seq_num = (dfu->i_blk_seq_num + 1) & 0xffff;
 
 	ret = dfu_read_buffer_fill(dfu, buf, size);
 	if (ret < 0) {

@@ -232,6 +262,8 @@ int dfu_read(struct dfu_entity *dfu, void *buf, int size, int blk_seq_num)
 		dfu->i_buf_end = dfu_buf + sizeof(dfu_buf);
 		dfu->i_buf = dfu->i_buf_start;
 		dfu->b_left = 0;
+
+		dfu->inited = 0;
 	}
 
 	return ret;

index 4c8994b..29a2c2e 100644
@@ -234,5 +234,8 @@ int dfu_fill_entity_mmc(struct dfu_entity *dfu, char *s)
 	dfu->read_medium = dfu_read_medium_mmc;
 	dfu->write_medium = dfu_write_medium_mmc;
 
+	/* initial state */
+	dfu->inited = 0;
+
 	return 0;
 }

index b662488..2847ef8 100644
@@ -90,6 +90,8 @@ struct dfu_entity {
 	u8 *i_buf_end;
 	long r_left;
 	long b_left;
+
+	unsigned int inited : 1;
 };
 
 int dfu_config_entities(char *s, char *interface, int num);