Streaming merkle tree construction with minimal memory implemented in C

merkle.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");
    }
}

merkle.h

#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