Add API to override SHA256 compression at runtime

diff --git a/include/secp256k1.h b/include/secp256k1.h
index 9de45f1..b7ec6a2 100644
--- a/include/secp256k1.h
+++ b/include/secp256k1.h
@@ -6,6 +6,7 @@ extern "C" {
 #endif
 
 #include <stddef.h>
+#include <stdint.h>
 
 /** Unless explicitly stated all pointer arguments must not be NULL.
  *
@@ -404,6 +405,46 @@ SECP256K1_API void secp256k1_context_set_error_callback(
     const void *data
 ) SECP256K1_ARG_NONNULL(1);
 
+/** A pointer to a function implementing SHA256's internal compression function.
+ *
+ * This function processes one or more contiguous 64-byte message blocks and
+ * updates the internal SHA256 state accordingly. The function is not responsible
+ * for counting consumed blocks or bytes, nor for performing padding.
+ *
+ * In/Out:  state:     pointer to eight 32-bit words representing the current internal state;
+ *                     the state is updated in place.
+ * In:      blocks64:  pointer to concatenation of n_blocks blocks, of 64 bytes each.
+ *                     no alignment guarantees are made for this pointer.
+ *          n_blocks:  number of contiguous 64-byte blocks to process.
+ */
+typedef void (*secp256k1_sha256_compression_function)(
+    uint32_t *state,
+    const unsigned char *blocks64,
+    size_t n_blocks
+);
+
+/**
+ * Set a callback function to override the internal SHA256 compression function.
+ *
+ * This installs a function to replace the built-in block-compression
+ * step used by the library's internal SHA256 implementation.
+ * The provided callback must exactly implement the effect of n_blocks
+ * repeated applications of the SHA256 compression function.
+ *
+ * This API exists to support environments that wish to route the
+ * SHA256 compression step through a hardware-accelerated or otherwise
+ * specialized implementation. It is NOT meant for replacing SHA256
+ * with a different hash function.
+ *
+ * Args:    ctx:             pointer to a context object.
+ * In:      fn_compression:  pointer to a function implementing the compression function;
+ *                           passing NULL restores the default implementation.
+ */
+SECP256K1_API void secp256k1_context_set_sha256_compression(
+        secp256k1_context *ctx,
+        secp256k1_sha256_compression_function fn_compression
+) SECP256K1_ARG_NONNULL(1);
+
 /** Parse a variable-length public key into the pubkey object.
  *
  *  Returns: 1 if the public key was fully valid.
diff --git a/src/hash.h b/src/hash.h
index 35825ea..79d9767 100644
--- a/src/hash.h
+++ b/src/hash.h
@@ -10,8 +10,6 @@
 #include <stdlib.h>
 #include <stdint.h>
 
-typedef void (*secp256k1_sha256_compression_function)(uint32_t *state, const unsigned char *blocks64, size_t n_blocks);
-
 typedef struct {
     secp256k1_sha256_compression_function fn_sha256_compression;
 } secp256k1_hash_ctx;
diff --git a/src/modules/ecdh/main_impl.h b/src/modules/ecdh/main_impl.h
index 3ee3ca7..b0359b2 100644
--- a/src/modules/ecdh/main_impl.h
+++ b/src/modules/ecdh/main_impl.h
@@ -45,10 +45,6 @@ int secp256k1_ecdh(const secp256k1_context* ctx, unsigned char *output, const se
     ARG_CHECK(point != NULL);
     ARG_CHECK(scalar != NULL);
 
-    if (hashfp == NULL) {
-        hashfp = secp256k1_ecdh_hash_function_default;
-    }
-
     secp256k1_pubkey_load(ctx, &pt, point);
     secp256k1_scalar_set_b32(&s, scalar, &overflow);
 
@@ -64,7 +60,12 @@ int secp256k1_ecdh(const secp256k1_context* ctx, unsigned char *output, const se
     secp256k1_fe_get_b32(x, &pt.x);
     secp256k1_fe_get_b32(y, &pt.y);
 
-    ret = hashfp(output, x, y, data);
+    if (hashfp == NULL) {
+        /* Use ctx-aware function by default */
+        ret = ecdh_hash_function_sha256_impl(secp256k1_get_hash_context(ctx), output, x, y, data);
+    } else {
+        ret = hashfp(output, x, y, data);
+    }
 
     secp256k1_memclear_explicit(x, sizeof(x));
     secp256k1_memclear_explicit(y, sizeof(y));
diff --git a/src/modules/ecdh/tests_impl.h b/src/modules/ecdh/tests_impl.h
index e6c1bab..c75ce9f 100644
--- a/src/modules/ecdh/tests_impl.h
+++ b/src/modules/ecdh/tests_impl.h
@@ -8,6 +8,7 @@
 #define SECP256K1_MODULE_ECDH_TESTS_H
 
 #include "../../unit_test.h"
+#include "../../testutil.h"
 
 static int ecdh_hash_function_test_xpassthru(unsigned char *output, const unsigned char *x, const unsigned char *y, void *data) {
     (void)y;
@@ -90,6 +91,30 @@ static void test_ecdh_generator_basepoint(void) {
     }
 }
 
+DEFINE_SHA256_TRANSFORM_PROBE(sha256_ecdh)
+static void test_ecdh_ctx_sha256(void) {
+    /* Check ctx-provided SHA256 compression override takes effect */
+    secp256k1_context *ctx = secp256k1_context_clone(CTX);
+    unsigned char out_default[65], out_custom[65];
+    const unsigned char sk[32] = {1};
+    secp256k1_pubkey pubkey;
+    CHECK(secp256k1_ec_pubkey_create(ctx, &pubkey, sk) == 1);
+
+    /* Default behavior */
+    CHECK(secp256k1_ecdh(ctx, out_default, &pubkey, sk, NULL, NULL) == 1);
+    CHECK(!sha256_ecdh_called);
+
+    /* Override SHA256 compression directly, bypassing the ctx setter sanity checks */
+    ctx->hash_ctx.fn_sha256_compression = sha256_ecdh;
+    CHECK(secp256k1_ecdh(ctx, out_custom, &pubkey, sk, NULL, NULL) == 1);
+
+    /* Outputs must differ if custom compression was used */
+    CHECK(secp256k1_memcmp_var(out_default, out_custom, 32) != 0);
+    CHECK(sha256_ecdh_called);
+
+    secp256k1_context_destroy(ctx);
+}
+
 static void test_bad_scalar(void) {
     unsigned char s_zero[32] = { 0 };
     unsigned char s_overflow[32] = { 0 };
@@ -187,6 +212,7 @@ static const struct tf_test_entry tests_ecdh[] = {
     CASE1(test_bad_scalar),
     CASE1(test_result_basepoint),
     CASE1(test_ecdh_wycheproof),
+    CASE1(test_ecdh_ctx_sha256),
 };
 
 #endif /* SECP256K1_MODULE_ECDH_TESTS_H */
diff --git a/src/modules/ellswift/main_impl.h b/src/modules/ellswift/main_impl.h
index 29c0a53..27cb3db 100644
--- a/src/modules/ellswift/main_impl.h
+++ b/src/modules/ellswift/main_impl.h
@@ -562,8 +562,14 @@ int secp256k1_ellswift_xdh(const secp256k1_context *ctx, unsigned char *output,
     secp256k1_fe_normalize(&px);
     secp256k1_fe_get_b32(sx, &px);
 
-    /* Invoke hasher */
-    ret = hashfp(output, sx, ell_a64, ell_b64, data);
+    /* Invoke hasher. Use ctx-aware function by default */
+    if (hashfp == secp256k1_ellswift_xdh_hash_function_bip324) {
+        ret = ellswift_xdh_hash_function_bip324_impl(secp256k1_get_hash_context(ctx), output, sx, ell_a64, ell_b64, data);
+    } else if (hashfp == secp256k1_ellswift_xdh_hash_function_prefix) {
+        ret = ellswift_xdh_hash_function_prefix_impl(secp256k1_get_hash_context(ctx), output, sx, ell_a64, ell_b64, data);
+    } else {
+        ret = hashfp(output, sx, ell_a64, ell_b64, data);
+    }
 
     secp256k1_memclear_explicit(sx, sizeof(sx));
     secp256k1_fe_clear(&px);
diff --git a/src/modules/ellswift/tests_impl.h b/src/modules/ellswift/tests_impl.h
index c9a7500..a849c8e 100644
--- a/src/modules/ellswift/tests_impl.h
+++ b/src/modules/ellswift/tests_impl.h
@@ -431,6 +431,40 @@ void ellswift_xdh_correctness_tests(void) {
     }
 }
 
+DEFINE_SHA256_TRANSFORM_PROBE(sha256_ellswift_xdh)
+void ellswift_xdh_ctx_sha256_tests(void) {
+    /* Check ctx-provided SHA256 compression override takes effect */
+    secp256k1_context *ctx = secp256k1_context_clone(CTX);
+    unsigned char out_default[65], out_custom[65];
+    const unsigned char skA[32] = {1}, skB[32] = {2};
+    unsigned char keyA[64], keyB[64], data[64] = {0};
+    const secp256k1_ellswift_xdh_hash_function hash_funcs[2] = {secp256k1_ellswift_xdh_hash_function_bip324, secp256k1_ellswift_xdh_hash_function_prefix};
+    int i;
+
+    CHECK(secp256k1_ellswift_create(ctx, keyA, skA, NULL));
+    CHECK(secp256k1_ellswift_create(ctx, keyB, skB, NULL));
+
+    for (i = 0; i < 2; i++) {
+        const secp256k1_ellswift_xdh_hash_function hash_fn = hash_funcs[i];
+        /* Default behavior. No ctx-provided SHA256 compression */
+        CHECK(secp256k1_ellswift_xdh(ctx, out_default, keyA, keyB, skA, 0, hash_fn, data));
+        CHECK(!sha256_ellswift_xdh_called);
+
+        /* Override SHA256 compression directly, bypassing the ctx setter sanity checks */
+        ctx->hash_ctx.fn_sha256_compression = sha256_ellswift_xdh;
+        CHECK(secp256k1_ellswift_xdh(ctx, out_custom, keyA, keyB, skA, 0, hash_fn, data));
+        CHECK(sha256_ellswift_xdh_called);
+        /* Outputs must differ if custom compression was used */
+        CHECK(secp256k1_memcmp_var(out_default, out_custom, 32) != 0);
+
+        /* Restore defaults */
+        sha256_ellswift_xdh_called = 0;
+        secp256k1_context_set_sha256_compression(ctx, NULL);
+    }
+
+    secp256k1_context_destroy(ctx);
+}
+
 /* Test hash initializers */
 void ellswift_hash_init_tests(void) {
     secp256k1_sha256 sha_optimized;
@@ -499,6 +533,7 @@ static const struct tf_test_entry tests_ellswift[] = {
     CASE1(ellswift_xdh_correctness_tests),
     CASE1(ellswift_hash_init_tests),
     CASE1(ellswift_xdh_bad_scalar_tests),
+    CASE1(ellswift_xdh_ctx_sha256_tests),
 };
 
 #endif
diff --git a/src/modules/schnorrsig/main_impl.h b/src/modules/schnorrsig/main_impl.h
index 3fe645b..5100557 100644
--- a/src/modules/schnorrsig/main_impl.h
+++ b/src/modules/schnorrsig/main_impl.h
@@ -137,10 +137,6 @@ static int secp256k1_schnorrsig_sign_internal(const secp256k1_context* ctx, unsi
     ARG_CHECK(msg != NULL || msglen == 0);
     ARG_CHECK(keypair != NULL);
 
-    if (noncefp == NULL) {
-        noncefp = secp256k1_nonce_function_bip340;
-    }
-
     ret &= secp256k1_keypair_load(ctx, &sk, &pk, keypair);
     /* Because we are signing for a x-only pubkey, the secret key is negated
      * before signing if the point corresponding to the secret key does not
@@ -151,7 +147,15 @@ static int secp256k1_schnorrsig_sign_internal(const secp256k1_context* ctx, unsi
 
     secp256k1_scalar_get_b32(seckey, &sk);
     secp256k1_fe_get_b32(pk_buf, &pk.x);
-    ret &= !!noncefp(nonce32, msg, msglen, seckey, pk_buf, bip340_algo, sizeof(bip340_algo), ndata);
+
+    /* Compute nonce */
+    if (noncefp == NULL || noncefp == secp256k1_nonce_function_bip340) {
+        /* Use context-aware nonce function by default */
+        ret &= nonce_function_bip340_impl(secp256k1_get_hash_context(ctx), nonce32, msg, msglen, seckey, pk_buf, bip340_algo, sizeof(bip340_algo), ndata);
+    } else {
+        ret &= !!noncefp(nonce32, msg, msglen, seckey, pk_buf, bip340_algo, sizeof(bip340_algo), ndata);
+    }
+
     secp256k1_scalar_set_b32(&k, nonce32, NULL);
     ret &= !secp256k1_scalar_is_zero(&k);
     secp256k1_scalar_cmov(&k, &secp256k1_scalar_one, !ret);
diff --git a/src/modules/schnorrsig/tests_impl.h b/src/modules/schnorrsig/tests_impl.h
index 009bb3e..56812e7 100644
--- a/src/modules/schnorrsig/tests_impl.h
+++ b/src/modules/schnorrsig/tests_impl.h
@@ -852,6 +852,29 @@ static void test_schnorrsig_sign_internal(void) {
     CHECK(secp256k1_memcmp_var(sig, sig2, sizeof(sig)) == 0);
 }
 
+DEFINE_SHA256_TRANSFORM_PROBE(sha256_schnorrsig)
+static void test_schnorrsig_ctx_sha256(void) {
+    /* Check ctx-provided SHA256 compression override takes effect */
+    secp256k1_context *ctx = secp256k1_context_clone(CTX);
+    unsigned char out_default[64], out_custom[64];
+    unsigned char sk[32] = {1}, msg32[32] = {1};
+    secp256k1_keypair keypair;
+    CHECK(secp256k1_keypair_create(ctx, &keypair, sk));
+
+    /* Default behavior. No ctx-provided SHA256 compression */
+    CHECK(secp256k1_schnorrsig_sign32(ctx, out_default, msg32, &keypair, NULL));
+    CHECK(!sha256_schnorrsig_called);
+
+    /* Override SHA256 compression directly, bypassing the ctx setter sanity checks */
+    ctx->hash_ctx.fn_sha256_compression = sha256_schnorrsig;
+    CHECK(secp256k1_schnorrsig_sign32(ctx, out_custom, msg32, &keypair, NULL));
+    CHECK(sha256_schnorrsig_called);
+    /* Outputs must differ if custom compression was used */
+    CHECK(secp256k1_memcmp_var(out_default, out_custom, 64) != 0);
+
+    secp256k1_context_destroy(ctx);
+}
+
 #define N_SIGS 3
 /* Creates N_SIGS valid signatures and verifies them with verify and
  * verify_batch (TODO). Then flips some bits and checks that verification now
@@ -981,6 +1004,7 @@ static const struct tf_test_entry tests_schnorrsig[] = {
     CASE1(test_schnorrsig_sign),
     CASE1(test_schnorrsig_sign_verify),
     CASE1(test_schnorrsig_taproot),
+    CASE1(test_schnorrsig_ctx_sha256),
 };
 
 #endif
diff --git a/src/secp256k1.c b/src/secp256k1.c
index b91b3bc..e4b80ff 100644
--- a/src/secp256k1.c
+++ b/src/secp256k1.c
@@ -223,6 +223,18 @@ void secp256k1_context_set_error_callback(secp256k1_context* ctx, void (*fun)(co
     ctx->error_callback.data = data;
 }
 
+void secp256k1_context_set_sha256_compression(secp256k1_context *ctx, secp256k1_sha256_compression_function fn_compression) {
+    VERIFY_CHECK(ctx != NULL);
+    ARG_CHECK_VOID(secp256k1_context_is_proper(ctx));
+    if (!fn_compression) { /* Reset hash context */
+        secp256k1_hash_ctx_init(&ctx->hash_ctx);
+        return;
+    }
+    /* Check and set */
+    ARG_CHECK_VOID(secp256k1_selftest_sha256(fn_compression));
+    ctx->hash_ctx.fn_sha256_compression = fn_compression;
+}
+
 static SECP256K1_INLINE const secp256k1_hash_ctx* secp256k1_get_hash_context(const secp256k1_context *ctx) {
     return &ctx->hash_ctx;
 }
@@ -538,9 +550,6 @@ static int secp256k1_ecdsa_sign_inner(const secp256k1_context* ctx, secp256k1_sc
     if (recid) {
         *recid = 0;
     }
-    if (noncefp == NULL) {
-        noncefp = secp256k1_nonce_function_default;
-    }
 
     /* Fail if the secret key is invalid. */
     is_sec_valid = secp256k1_scalar_set_b32_seckey(&sec, seckey);
@@ -548,7 +557,14 @@ static int secp256k1_ecdsa_sign_inner(const secp256k1_context* ctx, secp256k1_sc
     secp256k1_scalar_set_b32(&msg, msg32, NULL);
     while (1) {
         int is_nonce_valid;
-        ret = !!noncefp(nonce32, msg32, seckey, NULL, (void*)noncedata, count);
+
+        if (noncefp == NULL) {
+            /* Use ctx-aware function by default */
+            ret = nonce_function_rfc6979_impl(secp256k1_get_hash_context(ctx), nonce32, msg32, seckey, NULL, (void*)noncedata, count);
+        } else {
+            ret = !!noncefp(nonce32, msg32, seckey, NULL, (void*)noncedata, count);
+        }
+
         if (!ret) {
             break;
         }
diff --git a/src/tests.c b/src/tests.c
index 9c2ab89..b85b065 100644
--- a/src/tests.c
+++ b/src/tests.c
@@ -433,6 +433,53 @@ static void run_scratch_tests(void) {
     secp256k1_scratch_space_destroy(CTX, NULL); /* no-op */
 }
 
+/* A compression function that does nothing */
+static void invalid_sha256_compression(uint32_t *s, const unsigned char *msg, size_t rounds) {
+    (void)s; (void)msg; (void)rounds;
+}
+
+static int own_transform_called = 0;
+static void good_sha256_compression(uint32_t *s, const unsigned char *msg, size_t rounds) {
+    own_transform_called = 1;
+    secp256k1_sha256_transform(s, msg, rounds);
+}
+
+static void run_plug_sha256_compression_tests(void) {
+    secp256k1_context *ctx, *ctx_cloned;
+    secp256k1_sha256 sha;
+    unsigned char sha_out[32];
+    /* 1) Verify the context is initialized with the default compression function */
+    ctx = secp256k1_context_create(SECP256K1_CONTEXT_NONE);
+    CHECK(ctx->hash_ctx.fn_sha256_compression == secp256k1_sha256_transform);
+
+    /* 2) Verify providing a bad compression function fails during set */
+    CHECK_ILLEGAL_VOID(ctx, secp256k1_context_set_sha256_compression(ctx, invalid_sha256_compression));
+    CHECK(ctx->hash_ctx.fn_sha256_compression == secp256k1_sha256_transform);
+
+    /* 3) Provide sha256 to ctx and verify it is called when provided */
+    own_transform_called = 0;
+    secp256k1_context_set_sha256_compression(ctx, good_sha256_compression);
+    CHECK(own_transform_called);
+
+    /* 4) Verify callback makes it across clone */
+    ctx_cloned = secp256k1_context_clone(ctx);
+    CHECK(ctx_cloned->hash_ctx.fn_sha256_compression == good_sha256_compression);
+
+    /* 5) A hash operation should invoke the installed callback */
+    own_transform_called = 0;
+    secp256k1_sha256_initialize(&sha);
+    secp256k1_sha256_write(secp256k1_get_hash_context(ctx), &sha, (const unsigned char*)"a", 1);
+    secp256k1_sha256_finalize(secp256k1_get_hash_context(ctx), &sha, sha_out);
+    CHECK(own_transform_called);
+
+    /* 6) Unset sha256 and verify the default one is set again */
+    secp256k1_context_set_sha256_compression(ctx, NULL);
+    CHECK(ctx->hash_ctx.fn_sha256_compression == secp256k1_sha256_transform);
+
+    secp256k1_context_destroy(ctx);
+    secp256k1_context_destroy(ctx_cloned);
+}
+
 static void run_ctz_tests(void) {
     static const uint32_t b32[] = {1, 0xffffffff, 0x5e56968f, 0xe0d63129};
     static const uint64_t b64[] = {1, 0xffffffffffffffff, 0xbcd02462139b3fc3, 0x98b5f80c769693ef};
@@ -7129,7 +7176,7 @@ static void run_ecdsa_der_parse(void) {
 }
 
 /* Tests several edge cases. */
-static void test_ecdsa_edge_cases(void) {
+static void run_ecdsa_edge_cases(void) {
     int t;
     secp256k1_ecdsa_signature sig;
 
@@ -7462,8 +7509,25 @@ static void test_ecdsa_edge_cases(void) {
     }
 }
 
-static void run_ecdsa_edge_cases(void) {
-    test_ecdsa_edge_cases();
+DEFINE_SHA256_TRANSFORM_PROBE(sha256_ecdsa)
+static void ecdsa_ctx_sha256(void) {
+    /* Check ctx-provided SHA256 compression override takes effect */
+    secp256k1_context *ctx = secp256k1_context_clone(CTX);
+    secp256k1_ecdsa_signature out_default, out_custom;
+    unsigned char sk[32] = {1}, msg32[32] = {1};
+
+    /* Default behavior. No ctx-provided SHA256 compression */
+    CHECK(secp256k1_ecdsa_sign(ctx, &out_default, msg32, sk, NULL, NULL));
+    CHECK(!sha256_ecdsa_called);
+
+    /* Override SHA256 compression directly, bypassing the ctx setter sanity checks */
+    ctx->hash_ctx.fn_sha256_compression = sha256_ecdsa;
+    CHECK(secp256k1_ecdsa_sign(ctx, &out_custom, msg32, sk, NULL, NULL));
+    CHECK(sha256_ecdsa_called);
+    /* Outputs must differ if custom compression was used */
+    CHECK(secp256k1_memcmp_var(out_default.data, out_custom.data, 64) != 0);
+
+    secp256k1_context_destroy(ctx);
 }
 
 /** Wycheproof tests
@@ -7754,6 +7818,7 @@ static const struct tf_test_entry tests_general[] = {
     CASE(all_static_context_tests),
     CASE(deprecated_context_flags_test),
     CASE(scratch_tests),
+    CASE(plug_sha256_compression_tests),
 };
 
 static const struct tf_test_entry tests_integer[] = {
@@ -7824,6 +7889,7 @@ static const struct tf_test_entry tests_ecdsa[] = {
     CASE(ecdsa_end_to_end),
     CASE(ecdsa_edge_cases),
     CASE(ecdsa_wycheproof),
+    CASE1(ecdsa_ctx_sha256),
 };
 
 static const struct tf_test_entry tests_utils[] = {
diff --git a/src/testutil.h b/src/testutil.h
index 93ee3d5..8fa69a0 100644
--- a/src/testutil.h
+++ b/src/testutil.h
@@ -11,6 +11,15 @@
 #include "testrand.h"
 #include "util.h"
 
+/* Helper for when we need to check that the ctx-provided sha256 compression was called */
+#define DEFINE_SHA256_TRANSFORM_PROBE(name)                                     \
+    static int name##_called = 0;                                               \
+    static void name(uint32_t *s, const unsigned char *msg, size_t rounds) {    \
+        name##_called = 1;                                                      \
+        secp256k1_sha256_transform(s, msg, rounds);                             \
+        s[0] ^= 0xdeadbeef; /* intentional perturbation for testing */          \
+    }
+
 /* group order of the secp256k1 curve in 32-byte big endian representation */
 static const unsigned char secp256k1_group_order_bytes[32] = {
     0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,