From 902db294efcf0bf00938c580a906366ec3fbb03a Mon Sep 17 00:00:00 2001
From: Michel Machado <michel@digirati.com.br>
Date: Sat, 14 Feb 2015 10:36:07 -0500
Subject: [PATCH] f3probe: align buffers with devices' block sizes

This patch should address issue discussed here:
https://github.com/AltraMayor/f3/issues/13

Reported-by: Dingyuan Wang
---
 f3brew.c   |  5 +++--
 libdevs.c  | 11 +++++------
 libprobe.c |  4 ++--
 libutils.c |  7 +++++++
 libutils.h |  7 ++++---
 5 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/f3brew.c b/f3brew.c
index 3a0b08c..ecf2c06 100644
--- a/f3brew.c
+++ b/f3brew.c
@@ -140,9 +140,10 @@ static void read_blocks(char *probe_blk, struct device *dev,
 static void write_and_read_blocks(struct device *dev,
 	uint64_t first_block, uint64_t last_block)
 {
+	const int block_order = dev_get_block_order(dev);
 	const int block_size = dev_get_block_size(dev);
-	char stack[511 + block_size];
-	char *blk = align_512(stack);
+	char stack[align_head(block_order) + block_size];
+	char *blk = align_mem(stack, block_order);
 
 	printf("Writing blocks from 0x%" PRIx64 " to 0x%" PRIx64 "...",
 		first_block, last_block);
diff --git a/libdevs.c b/libdevs.c
index 1c007a2..16e4a6a 100644
--- a/libdevs.c
+++ b/libdevs.c
@@ -1029,7 +1029,7 @@ static int sdev_save_block(struct safe_device *sdev,
 
 	/* The block at @offset hasn't been saved before. Save this block. */
 	assert(sdev->sb_n < sdev->sb_max);
-	block = (char *)align_512(sdev->saved_blocks) +
+	block = (char *)align_mem(sdev->saved_blocks, block_order) +
 		(sdev->sb_n << block_order);
 	rc = sdev->shadow_dev->read_block(sdev->shadow_dev, block,
 		length, offset);
@@ -1068,10 +1068,9 @@ static void sdev_free(struct device *dev)
 	struct safe_device *sdev = dev_sdev(dev);
 
 	if (sdev->sb_n > 0) {
-		char *first_block = align_512(sdev->saved_blocks);
-		char *block = first_block +
-			((sdev->sb_n - 1) <<
-			dev_get_block_order(sdev->shadow_dev));
+		const int block_order = dev_get_block_order(sdev->shadow_dev);
+		char *first_block = align_mem(sdev->saved_blocks, block_order);
+		char *block = first_block + ((sdev->sb_n - 1) << block_order);
 		uint64_t *poffset = &sdev->sb_offsets[sdev->sb_n - 1];
 		int block_size = dev_get_block_size(sdev->shadow_dev);
 
@@ -1114,7 +1113,7 @@ struct device *create_safe_device(struct device *dev, int max_blocks,
 	if (!sdev)
 		goto error;
 
-	length = 511 + (max_blocks << block_order);
+	length = align_head(block_order) + (max_blocks << block_order);
 	sdev->saved_blocks = malloc(length);
 	if (!sdev->saved_blocks)
 		goto sdev;
diff --git a/libprobe.c b/libprobe.c
index 7e69c61..29af02c 100644
--- a/libprobe.c
+++ b/libprobe.c
@@ -339,7 +339,7 @@ int probe_device(struct device *dev, uint64_t *preal_size_byte,
 	uint64_t dev_size_byte = dev_get_size_byte(dev);
 	const int block_size = dev_get_block_size(dev);
 	const int block_order = dev_get_block_order(dev);
-	char stack[511 + (2 << block_order)];
+	char stack[align_head(block_order) + (2 << block_order)];
 	char *stamp_blk, *probe_blk;
 	/* XXX Don't write at the very beginning of the card to avoid
 	 * losing the partition table.
@@ -362,7 +362,7 @@ int probe_device(struct device *dev, uint64_t *preal_size_byte,
 	 * the block device.
 	 * For the file device, this is superfluous.
 	 */
-	stamp_blk = align_512(stack);
+	stamp_blk = align_mem(stack, block_order);
 	probe_blk = stamp_blk + block_size;
 
 	fill_buffer(stamp_blk, block_size);
diff --git a/libutils.c b/libutils.c
index 95b522f..52ca0d9 100644
--- a/libutils.c
+++ b/libutils.c
@@ -39,3 +39,10 @@ int ceiling_log2(uint64_t x)
 {
 	return ilog2(clp2(x));
 }
+
+void *align_mem(void *p, int order)
+{
+	uintptr_t ip = (uintptr_t)p;
+	uintptr_t head = align_head(order);
+	return (void *)(   (ip + head) & ~head   );
+}
diff --git a/libutils.h b/libutils.h
index 9fe0151..a07d063 100644
--- a/libutils.h
+++ b/libutils.h
@@ -6,10 +6,11 @@
 int ilog2(uint64_t x);
 int ceiling_log2(uint64_t x);
 
-static inline void *align_512(void *p)
+static inline int align_head(int order)
 {
-	uintptr_t ip = (uintptr_t)p;
-	return (void *)(   (ip + 511) & ~511   );
+	return (1 << order) - 1;
 }
 
+void *align_mem(void *p, int order);
+
 #endif	/* HEADER_LIBUTILS_H */