verity: create a regular trie when depth=0

We want to eventually remove support for the depth constructor parameter
only create trees where level[0]->count = 1.

BUG=9752
TEST=Ran dm-verity.git unit tests. Ran platform_DMVerityCorruption on H/W.

Change-Id: Ie05ebc93a9c52055b4573ccd8811acbbb114adf3

R=wad@chromium.org,taysom@chromium.org,ups@chromium.org

Review URL: http://codereview.chromium.org/6670096

dm-bht.c

diff --git a/dm-bht.c b/dm-bht.c
index 5957cb7..1665c74 100644
--- a/dm-bht.c
+++ b/dm-bht.c
@@ -265,27 +265,6 @@
 		goto bad_block_count;
 	}
 
-	bht->depth = depth;  /* assignment below */
-	DMDEBUG("Setting depth %u", depth);
-	if (!depth || depth > UINT_MAX / sizeof(struct dm_bht_level)) {
-		DMERR("bht depth is invalid: %u", depth);
-		status = -EINVAL;
-		goto bad_depth;
-	}
-
-	/* Allocate levels. Each level of the tree may have an arbitrary number
-	 * of dm_bht_entry structs.  Each entry contains node_count nodes.
-	 * Each node in the tree is a cryptographic digest of either node_count
-	 * nodes on the subsequent level or of a specific block on disk.
-	 */
-	bht->levels = (struct dm_bht_level *)
-			kcalloc(depth, sizeof(struct dm_bht_level), GFP_KERNEL);
-	if (!bht->levels) {
-		DMERR("failed to allocate tree levels");
-		status = -ENOMEM;
-		goto bad_level_alloc;
-	}
-
 	/* Each dm_bht_entry->nodes is one page.  The node code tracks
 	 * how many nodes fit into one entry where a node is a single
 	 * hash (message digest).
@@ -302,6 +281,19 @@
 		status = -EINVAL;
 		goto bad_node_count;
 	}
+
+	/* if depth == 0, create a "regular" trie with a single root block */
+	if (depth == 0)
+		depth = DIV_ROUND_UP(fls(block_count - 1),
+				     bht->node_count_shift);
+	if (depth > UINT_MAX / sizeof(struct dm_bht_level)) {
+		DMERR("bht depth is invalid: %u", depth);
+		status = -EINVAL;
+		goto bad_depth;
+	}
+	DMDEBUG("Setting depth to %u.", depth);
+	bht->depth = depth;
+
 	/* Ensure that we can safely shift by this value. */
 	if (depth * bht->node_count_shift >= sizeof(unsigned int) * 8) {
 		DMERR("specified depth and node_count_shift is too large");
@@ -309,6 +301,19 @@
 		goto bad_node_count;
 	}
 
+	/* Allocate levels. Each level of the tree may have an arbitrary number
+	 * of dm_bht_entry structs.  Each entry contains node_count nodes.
+	 * Each node in the tree is a cryptographic digest of either node_count
+	 * nodes on the subsequent level or of a specific block on disk.
+	 */
+	bht->levels = (struct dm_bht_level *)
+			kcalloc(depth, sizeof(struct dm_bht_level), GFP_KERNEL);
+	if (!bht->levels) {
+		DMERR("failed to allocate tree levels");
+		status = -ENOMEM;
+		goto bad_level_alloc;
+	}
+
 	/* Setup callback stubs */
 	bht->read_cb = &dm_bht_read_callback_stub;
 	bht->write_cb = &dm_bht_write_callback_stub;
@@ -322,8 +327,8 @@
 bad_entries_alloc:
 	while (bht->depth-- > 0)
 		kfree(bht->levels[bht->depth].entries);
-bad_node_count:
 	kfree(bht->levels);
+bad_node_count:
 bad_level_alloc:
 bad_block_count:
 bad_depth:

dm-bht_unittest.cc

diff --git a/dm-bht_unittest.cc b/dm-bht_unittest.cc
index 2eb2297..8976418 100644
--- a/dm-bht_unittest.cc
+++ b/dm-bht_unittest.cc
@@ -206,6 +206,29 @@
   free(zero_page);
 }
 
+TEST_F(MemoryBhtTest, CreateThenVerifyZeroDepth) {
+  static const unsigned int total_blocks = 16384;
+  // Set the root hash for a 0-filled image
+  static const char kRootDigest[] =
+    "45d65d6f9e5a962f4d80b5f1bd7a918152251c27bdad8c5f52b590c129833372";
+  // A page of all zeros
+  u8 *zero_page = (u8 *)my_memalign(PAGE_SIZE, PAGE_SIZE);
+
+  memset(zero_page, 0, PAGE_SIZE);
+
+  SetupBht(0, total_blocks, "sha256");
+  dm_bht_set_root_hexdigest(bht_.get(),
+                            reinterpret_cast<const u8 *>(kRootDigest));
+
+  for (unsigned int blocks = 0; blocks < total_blocks; ++blocks) {
+    DLOG(INFO) << "verifying block: " << blocks;
+    EXPECT_EQ(0, dm_bht_verify_block(bht_.get(), blocks, zero_page));
+  }
+
+  EXPECT_EQ(0, dm_bht_destroy(bht_.get()));
+  free(zero_page);
+}
+
 TEST_F(MemoryBhtTest, CreateThenVerifyRealParameters) {
   static const unsigned int total_blocks = 217600;
   // Set the root hash for a 0-filled image

include/linux/kernel.h

diff --git a/include/linux/kernel.h b/include/linux/kernel.h
index dd617bd..0a7e322 100644
--- a/include/linux/kernel.h
+++ b/include/linux/kernel.h
@@ -15,6 +15,8 @@
 #define __ALIGN_MASK(x,mask)	(((x)+(mask))&~(mask))
 #define IS_ALIGNED(x, a)		(((x) & ((typeof(x))(a) - 1)) == 0)
 
+#define DIV_ROUND_UP(n,d) (((n) + (d) - 1) / (d))
+
 #define sector_div(n, b) ({ \
 	int _res; \
 	_res = (n) % (b); \