Skip to content

Instantly share code, notes, and snippets.

@shilch
Last active March 14, 2022 22:03
Show Gist options
  • Save shilch/2a53008ff99a6b2648696eebb0301fb6 to your computer and use it in GitHub Desktop.
Save shilch/2a53008ff99a6b2648696eebb0301fb6 to your computer and use it in GitHub Desktop.
Streaming merkle tree construction with minimal memory implemented in C
#include "merkle.h"
#include <assert.h>
#include <string.h>
#include <stdio.h>
void merkle_init(merkle_ctx* ctx){
ctx->size = 0;
ctx->mutated = false;
}
void merkle_append(merkle_ctx* ctx, const uint8_t* hash){
assert(ctx->size < MERKLE_MAX_DEPTH);
memcpy((void*)&ctx->hashes[ctx->size * MERKLE_DIGEST_SIZE], (void*)hash, MERKLE_DIGEST_SIZE);
ctx->depths[ctx->size] = 1;
ctx->size++;
while(ctx->size >= 2){
if(ctx->depths[ctx->size - 1] != ctx->depths[ctx->size - 2]){
break;
}
uint8_t* lhs_hash_ptr = &ctx->hashes[(ctx->size - 2) * MERKLE_DIGEST_SIZE];
uint8_t* rhs_hash_ptr = &ctx->hashes[(ctx->size - 1) * MERKLE_DIGEST_SIZE];
// detecting duplicate hashes fed to the merkle tree (CVE-2012-2459)
if(memcmp(lhs_hash_ptr, rhs_hash_ptr, MERKLE_DIGEST_SIZE) == 0){
ctx->mutated = true;
}
// check if we are about to modify our leftmost partial tree
if(ctx->size == 2){
// store rhs to branch_first for fast replacement of first hash
memcpy(&ctx->branch_first[(ctx->depths[0] - 1) * MERKLE_DIGEST_SIZE], rhs_hash_ptr, MERKLE_DIGEST_SIZE);
}
// first round: SHA256(lhs || rhs)
SHA256_Init(&ctx->_sha256_ctx);
SHA256_Update(&ctx->_sha256_ctx, lhs_hash_ptr, 2 * MERKLE_DIGEST_SIZE);
SHA256_Final(lhs_hash_ptr, &ctx->_sha256_ctx);
// second round: SHA256(lhs)
SHA256_Init(&ctx->_sha256_ctx);
SHA256_Update(&ctx->_sha256_ctx, lhs_hash_ptr, MERKLE_DIGEST_SIZE);
SHA256_Final(lhs_hash_ptr, &ctx->_sha256_ctx);
ctx->depths[ctx->size - 2]++;
ctx->size--;
}
}
void merkle_replace_first(merkle_ctx* ctx, const uint8_t* hash){
assert(ctx->size > 0);
uint8_t* cur_hash = &ctx->hashes[0];
memcpy(cur_hash, hash, MERKLE_DIGEST_SIZE);
for(size_t i = 0; i < ctx->depths[0] - 1; i++){
// first round: SHA256(h || branch[i])
SHA256_Init(&ctx->_sha256_ctx);
SHA256_Update(&ctx->_sha256_ctx, cur_hash, MERKLE_DIGEST_SIZE);
SHA256_Update(&ctx->_sha256_ctx, &ctx->branch_first[i * MERKLE_DIGEST_SIZE], MERKLE_DIGEST_SIZE);
SHA256_Final(cur_hash, &ctx->_sha256_ctx);
// second round: SHA256(h)
SHA256_Init(&ctx->_sha256_ctx);
SHA256_Update(&ctx->_sha256_ctx, cur_hash, MERKLE_DIGEST_SIZE);
SHA256_Final(cur_hash, &ctx->_sha256_ctx);
}
}
void merkle_final(uint8_t* root, const merkle_ctx* ctx){
if(ctx->size == 0){
memset(root, 0, MERKLE_DIGEST_SIZE);
return;
}
uint8_t cur_hash[MERKLE_DIGEST_SIZE];
size_t cur_depth = ctx->depths[ctx->size - 1];
memcpy(cur_hash, &ctx->hashes[(ctx->size - 1) * MERKLE_DIGEST_SIZE], MERKLE_DIGEST_SIZE);
for(size_t i = ctx->size - 1; i > 0; i--){
// bring right branch to same depth as left branch
const size_t target_depth = ctx->depths[i - 1];
assert(cur_depth <= target_depth);
while(cur_depth != target_depth){
// first round: SHA256(h || h)
SHA256_Init(&((merkle_ctx*)ctx)->_sha256_ctx); // const cast, we don't care that sha256_ctx changes
SHA256_Update(&((merkle_ctx*)ctx)->_sha256_ctx, cur_hash, MERKLE_DIGEST_SIZE);
SHA256_Update(&((merkle_ctx*)ctx)->_sha256_ctx, cur_hash, MERKLE_DIGEST_SIZE);
SHA256_Final(cur_hash, &((merkle_ctx*)ctx)->_sha256_ctx);
// second round: SHA256(h)
SHA256_Init(&((merkle_ctx*)ctx)->_sha256_ctx);
SHA256_Update(&((merkle_ctx*)ctx)->_sha256_ctx, cur_hash, MERKLE_DIGEST_SIZE);
SHA256_Final(cur_hash, &((merkle_ctx*)ctx)->_sha256_ctx);
cur_depth++;
}
// connect left and right branches
// first round: SHA256(lh || rh)
SHA256_Init(&((merkle_ctx*)ctx)->_sha256_ctx);
SHA256_Update(&((merkle_ctx*)ctx)->_sha256_ctx, &ctx->hashes[(i - 1) * MERKLE_DIGEST_SIZE], MERKLE_DIGEST_SIZE);
SHA256_Update(&((merkle_ctx*)ctx)->_sha256_ctx, cur_hash, MERKLE_DIGEST_SIZE);
SHA256_Final(cur_hash, &((merkle_ctx*)ctx)->_sha256_ctx);
// second round: SHA256(h)
SHA256_Init(&((merkle_ctx*)ctx)->_sha256_ctx);
SHA256_Update(&((merkle_ctx*)ctx)->_sha256_ctx, cur_hash, MERKLE_DIGEST_SIZE);
SHA256_Final(cur_hash, &((merkle_ctx*)ctx)->_sha256_ctx);
cur_depth++;
};
memcpy(root, cur_hash, MERKLE_DIGEST_SIZE);
}
void merkle_print(const merkle_ctx* ctx){
uint8_t root_hash[MERKLE_DIGEST_SIZE];
merkle_final(root_hash, ctx);
printf("root hash: ");
for(size_t i = 0; i < MERKLE_DIGEST_SIZE; i++){
printf("%02x", root_hash[i]);
}
printf("\n");
printf("size: %zu\n", ctx->size);
printf("partial trees:\n");
for(size_t i = 0; i < ctx->size; i++){
printf(" - %2zu : ", ctx->depths[i]);
for(size_t j = 0; j < MERKLE_DIGEST_SIZE; j++){
printf("%02x", ctx->hashes[MERKLE_DIGEST_SIZE * i + j]);
}
printf("\n");
}
printf("branch of first hash:");
if(ctx->size > 0 && ctx->depths[0] > 1){
printf("\n");
for(size_t i = 0; i < ctx->depths[0] - 1; i++){
printf(" %2zu: ", i);
for(size_t j = 0; j < MERKLE_DIGEST_SIZE; j++){
printf("%02x", ctx->branch_first[MERKLE_DIGEST_SIZE * i + j]);
}
printf("\n");
}
} else {
printf(" no branch\n");
}
}
#pragma once
#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>
#include <openssl/sha.h>
#ifdef __cplusplus
extern "C" {
#endif
#define MERKLE_MAX_DEPTH 32
#define MERKLE_DIGEST_SIZE SHA256_DIGEST_LENGTH
typedef struct {
SHA256_CTX _sha256_ctx;
uint8_t hashes[MERKLE_MAX_DEPTH * MERKLE_DIGEST_SIZE];
uint8_t branch_first[MERKLE_MAX_DEPTH * MERKLE_DIGEST_SIZE];
size_t depths[MERKLE_MAX_DEPTH];
size_t size;
bool mutated;
} merkle_ctx;
void merkle_init(merkle_ctx*);
void merkle_append(merkle_ctx*, const uint8_t*);
void merkle_replace_first(merkle_ctx*, const uint8_t*);
void merkle_final(uint8_t*, const merkle_ctx*);
void merkle_print(const merkle_ctx*);
#ifdef __cplusplus
} // extern
#endif
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment