sha256: cross-check caller supplied compression function

diff --git a/include/secp256k1.h b/include/secp256k1.h
index 4ee6200..09506f2 100644
--- a/include/secp256k1.h
+++ b/include/secp256k1.h
@@ -431,6 +431,13 @@ typedef void (*secp256k1_sha256_compression_function)(
  * noncefp==NULL or noncefp==secp256k1_nonce_function_default is passed
  * as an argument to secp256k1_ecdsa_sign.)
  *
+ * Note: The provided function is tested against a set of known SHA256
+ * digests; invokes the context's illegal callback on any mismatch
+ * (which aborts by default), in order to catch basic misbehavior early.
+ * It takes well under 2.5 ms on a desktop machine.
+ * This is NOT a substitute for having proper test coverage of the
+ * supplied function outside this library.
+ *
  * Args:    ctx:             pointer to a context object.
  * In:      fn_compression:  pointer to a function implementing the compression function;
  *                           passing NULL restores the default implementation.
diff --git a/src/hash_impl.h b/src/hash_impl.h
index 7c40f82..932367f 100644
--- a/src/hash_impl.h
+++ b/src/hash_impl.h
@@ -137,6 +137,71 @@ static void secp256k1_sha256_transform(uint32_t *state, const unsigned char *blo
     }
 }
 
+/* Perform a smoke test on a supplied SHA256 compression function. */
+static int secp256k1_sha256_smoke_test(const secp256k1_sha256_compression_function fn_compression) {
+    secp256k1_hash_ctx ctx;
+    secp256k1_sha256 sha_msg, sha_accum;
+    unsigned char out[32];
+    size_t i, j;
+
+    /* SHA256 works on 64 byte blocks, secp256k1_sha256_write gives as many blocks
+     * at once to compression, so the count is what varies here. A SIMD implementation
+     * typically hashes four or eight at a time, then any left over one by one.
+     * These lengths cover every number from 1 to 9, which includes counts that
+     * divide evenly and counts leaving one, two or three over. */
+    static const size_t msg_lens[] = {
+        0, 1, 28,               /* Shorter than a block, so padding makes up the rest */
+        55, 56,                 /* Final 0x80 and 8 byte length fit in the last block, and don't */
+        64, 128, 192, 256, 320, /* 1 to 5 blocks */
+        384, 448, 512, 576      /* 6 to 9 blocks */
+    };
+    unsigned char msg[640]; /* Longest message, plus 64 for the shifted start */
+
+    /* Accumulated digest of every message, hashed with the built-in secp256k1_sha256_transform.
+     * Note: To regenerate set 'ctx.fn_sha256_compression = secp256k1_sha256_transform' below
+     * and print the sha_accum digest. */
+    static const unsigned char accum_expected[32] = {
+        0x22, 0x8E, 0x6A, 0x1F, 0x78, 0x02, 0x1B, 0xCF,
+        0x4E, 0xF4, 0xEA, 0xB3, 0x8A, 0x40, 0x69, 0xBF,
+        0x8D, 0xF3, 0x72, 0xAB, 0xE2, 0x11, 0x93, 0xA6,
+        0xE6, 0x46, 0x98, 0xBF, 0xD4, 0x3D, 0x19, 0x84,
+    };
+    /* The purpose of this VERIFY_CHECK is to make anyone aware that they
+     * should also change the size of msg_buf when changing the length of the
+     * longest message. */
+    VERIFY_CHECK(msg_lens[ARRAY_SIZE(msg_lens) - 1] == 576);
+
+    VERIFY_CHECK(fn_compression != NULL);
+    secp256k1_hash_ctx_init(&ctx);
+    ctx.fn_sha256_compression = fn_compression;
+    secp256k1_sha256_initialize(&sha_accum);
+
+    /* No two blocks of a message are equal, so a function that doesn't advance the block
+     * pointer gives a different digest. 251 is prime, so its repeats only line up with
+     * block starts every 251 blocks; wrapping at 256 would do so every 4. */
+    for (i = 0; i < sizeof(msg); i++) {
+        msg[i] = (unsigned char)(i % 251);
+    }
+
+    /* Each round starts one byte further along to check different alignments,
+     * secp256k1_sha256_write invokes compression directly on input >= 64 bytes */
+    for (i = 0; i < 64; i++) {
+        unsigned char *m = msg + i;
+        m[0] ^= 0xff; /* Changes the first byte, so every state after it changes too */
+        for (j = 0; j < ARRAY_SIZE(msg_lens); j++) {
+            secp256k1_sha256_initialize(&sha_msg);
+            secp256k1_sha256_write(&ctx, &sha_msg, m, msg_lens[j]);
+            secp256k1_sha256_finalize(&ctx, &sha_msg, out);
+            secp256k1_sha256_write(&ctx, &sha_accum, out, 32);
+        }
+        m[0] ^= 0xff; /* Reset first byte */
+    }
+
+    /* Compare against pre-computed accumulated digest */
+    secp256k1_sha256_finalize(&ctx, &sha_accum, out);
+    return secp256k1_memcmp_var(accum_expected, out, 32) == 0;
+}
+
 static void secp256k1_hash_ctx_init(secp256k1_hash_ctx *hash_ctx) {
     VERIFY_CHECK(hash_ctx != NULL);
     hash_ctx->fn_sha256_compression = secp256k1_sha256_transform;
diff --git a/src/selftest.h b/src/selftest.h
index de0e059..08b2a9b 100644
--- a/src/selftest.h
+++ b/src/selftest.h
@@ -24,7 +24,13 @@ static int secp256k1_selftest_sha256(secp256k1_sha256_compression_function fn_co
     hash_ctx.fn_sha256_compression = fn_compression;
     secp256k1_sha256_write(&hash_ctx, &hasher, (const unsigned char*)input63, 63);
     secp256k1_sha256_finalize(&hash_ctx, &hasher, out);
-    return secp256k1_memcmp_var(out, output32, 32) == 0;
+    if (secp256k1_memcmp_var(out, output32, 32) != 0) {
+        return 0;
+    }
+    if (fn_compression != secp256k1_sha256_transform && !secp256k1_sha256_smoke_test(fn_compression)) {
+        return 0;
+    }
+    return 1;
 }
 
 static int secp256k1_selftest_passes(void) {
diff --git a/src/tests.c b/src/tests.c
index 7fb51da..b7524dd 100644
--- a/src/tests.c
+++ b/src/tests.c
@@ -6,6 +6,7 @@
 
 #include <stdio.h>
 #include <stdlib.h>
+#include <stdint.h>
 #include <string.h>
 
 #include <time.h>
@@ -485,6 +486,79 @@ static void run_plug_sha256_compression_tests(void) {
     secp256k1_context_destroy(ctx_cloned);
 }
 
+/* Hashes the first block over and over instead of moving on. */
+static void sha256_transform_noadvance(uint32_t *s, const unsigned char *chunk, size_t blocks) {
+    size_t i;
+    for (i = 0; i < blocks; i++) {
+        secp256k1_sha256_transform(s, chunk, 1);
+    }
+}
+
+/* Drops the last block of a multi-block call. */
+static void sha256_transform_short(uint32_t *s, const unsigned char *chunk, size_t blocks) {
+    secp256k1_sha256_transform(s, chunk, blocks > 0 ? blocks - 1 : 0);
+}
+
+/* Starts from the IV instead of the state it was given. */
+static void sha256_transform_ivreset(uint32_t *s, const unsigned char *chunk, size_t blocks) {
+    secp256k1_sha256 h;
+    secp256k1_sha256_initialize(&h);
+    memcpy(s, h.s, sizeof(h.s));
+    secp256k1_sha256_transform(s, chunk, blocks);
+}
+
+/* Correct only on multiples of four blocks. */
+static void sha256_transform_batch4(uint32_t *s, const unsigned char *chunk, size_t blocks) {
+    secp256k1_sha256_transform(s, chunk, blocks - (blocks & 3));
+}
+
+/* Right digest, one bit off. */
+static void sha256_transform_corrupt(uint32_t *s, const unsigned char *chunk, size_t blocks) {
+    secp256k1_sha256_transform(s, chunk, blocks);
+    s[0] ^= 1;
+}
+
+#ifdef UINTPTR_MAX
+
+/* Wrong when input is 64-byte aligned, like a broken SIMD fast path. */
+static void sha256_transform_align64_fail(uint32_t *s, const unsigned char *chunk, size_t blocks) {
+    int aligned = ((uintptr_t)chunk % 64) == 0;
+    secp256k1_sha256_transform(s, chunk, blocks);
+    if (aligned) s[0] ^= 1;
+}
+
+/* Wrong when input is 32-byte aligned but not 64 */
+static void sha256_transform_align32_fail(uint32_t *s, const unsigned char *chunk, size_t blocks) {
+    int align32_not64 = (((uintptr_t)chunk % 32) == 0) && (((uintptr_t)chunk % 64) != 0);
+    secp256k1_sha256_transform(s, chunk, blocks);
+    if (align32_not64) {
+        s[0] ^= 1;
+    }
+}
+
+/* Wrong on any unaligned input. */
+static void sha256_transform_unaligned_fail(uint32_t *s, const unsigned char *chunk, size_t blocks) {
+    int aligned = ((uintptr_t)chunk % 64) == 0;
+    secp256k1_sha256_transform(s, chunk, blocks);
+    if (!aligned) s[0] ^= 1;
+}
+
+#endif /* UINTPTR_MAX */
+
+static void run_sha256_compression_smoke_test_tests(void) {
+    CHECK(secp256k1_sha256_smoke_test(sha256_transform_noadvance) == 0);
+    CHECK(secp256k1_sha256_smoke_test(sha256_transform_short) == 0);
+    CHECK(secp256k1_sha256_smoke_test(sha256_transform_ivreset) == 0);
+    CHECK(secp256k1_sha256_smoke_test(sha256_transform_batch4) == 0);
+    CHECK(secp256k1_sha256_smoke_test(sha256_transform_corrupt) == 0);
+#ifdef UINTPTR_MAX
+    CHECK(secp256k1_sha256_smoke_test(sha256_transform_align64_fail) == 0);
+    CHECK(secp256k1_sha256_smoke_test(sha256_transform_align32_fail) == 0);
+    CHECK(secp256k1_sha256_smoke_test(sha256_transform_unaligned_fail) == 0);
+#endif
+    CHECK(secp256k1_sha256_smoke_test(good_sha256_compression) == 1);
+}
+
 static void run_sha256_multi_block_compression_tests(void) {
     secp256k1_hash_ctx hash_ctx;
     secp256k1_sha256 sha256_one;
@@ -7955,6 +8029,7 @@ static const struct tf_test_entry tests_general[] = {
     CASE(scratch_tests),
     CASE(invalid_scratch_space_tests),
     CASE(plug_sha256_compression_tests),
+    CASE(sha256_compression_smoke_test_tests),
     CASE(sha256_multi_block_compression_tests),
 };