From cc3d33a9615a6577db47f8d0a0908af3b944d14b Mon Sep 17 00:00:00 2001
From: "Markku-Juhani O. Saarinen" <mjos@iki.fi>
Date: Sun, 27 Dec 2015 13:27:27 +0000
Subject: [PATCH] Added SHAKE128/256 code and test vectors.

---
 README.md | 16 ++++++++++----
 main.c    | 62 ++++++++++++++++++++++++++++++++++++++++++++++++++++---
 sha3.c    | 30 +++++++++++++++++++++++++--
 sha3.h    | 12 +++++++++--
 4 files changed, 109 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index 3c8f2d1..650b165 100644
--- a/README.md
+++ b/README.md
@@ -1,9 +1,17 @@
 # tiny_sha3
-Very small, readable implementation of the SHA3 hash function.
-Updated 03-Sep-15:
+Very small, readable implementation of the FIPS 202 and SHA3 hash function.
+Public domain.
+
+### Updated 27-Dec-15:
+
+Added SHAKE128 and SHAKE256 code and test vectors. The code can actually do
+a XOF of arbitrary size (like "SHAKE512").
+
+
+### Updated 03-Sep-15:
 
 Made the implementation portable. The API is now pretty much the
-same that OpenSSL uses. Public domain.
+same that OpenSSL uses. 
 
 
 ### Updated 07-Aug-15:
@@ -25,7 +33,7 @@ Cheers,
 - markku
 
 
-### Original README.TXT from 2011:
+### Original README.TXT from 19-Nov-11:
 
 Hi.
 
diff --git a/main.c b/main.c
index dd4f6f9..6d45a16 100644
--- a/main.c
+++ b/main.c
@@ -89,9 +89,64 @@ int test_sha3()
         sha3(msg, msg_len, buf, sha_len);
 
         if (memcmp(sha, buf, sha_len) != 0) {
-            fails++;
             fprintf(stderr, "[%d] SHA3-%d, len %d test FAILED.\n",
                 i, sha_len * 8, msg_len);
+            fails++;
+        }
+    }
+
+    return fails;
+}
+
+// test for SHAKE128 and SHAKE256
+
+int test_shake()
+{
+    // Test vectors have bytes 480..511 of XOF output for given inputs.
+    // From http://csrc.nist.gov/groups/ST/toolkit/examples.html#aHashing
+
+    const char *testhex[4] = {
+        // SHAKE128, message of length 0
+        "43E41B45A653F2A5C4492C1ADD544512DDA2529833462B71A41A45BE97290B6F",
+        // SHAKE256, message of length 0
+        "AB0BAE316339894304E35877B0C28A9B1FD166C796B9CC258A064A8F57E27F2A",
+        // SHAKE128, 1600-bit test pattern
+        "44C9FB359FD56AC0A9A75A743CFF6862F17D7259AB075216C0699511643B6439",
+        // SHAKE256, 1600-bit test pattern
+        "6A1A9D7846436E4DCA5728B6F760EEF0CA92BF0BE5615E96959D767197A0BEEB"
+    };
+
+    int i, j, fails;
+    sha3_ctx_t sha3;
+    uint8_t buf[32], ref[32];
+
+    fails = 0;
+
+    for (i = 0; i < 4; i++) {
+
+        if ((i & 1) == 0) {             // test each twice
+            shake128_init(&sha3);
+        } else {
+            shake256_init(&sha3);
+        }
+
+        if (i >= 2) {                   // 1600-bit test pattern
+            memset(buf, 0xA3, 20);
+            for (j = 0; j < 200; j += 20)
+                shake_update(&sha3, buf, 20);
+        }
+
+        shake_xof(&sha3);               // switch to extensible output
+
+        for (j = 0; j < 512; j += 32)   // output. discard bytes 0..479
+            shake_out(&sha3, buf, 32);
+
+        // compare to reference
+        test_readhex(ref, testhex[i], sizeof(ref));
+        if (memcmp(buf, ref, 32) != 0) {
+            fprintf(stderr, "[%d] SHAKE%d, len %d test FAILED.\n",
+                i, i & 1 ? 256 : 128, i >= 2 ? 1600 : 0);
+            fails++;
         }
     }
 
@@ -101,8 +156,9 @@ int test_sha3()
 // main
 int main(int argc, char **argv)
 {
-    if (test_sha3() == 0)
-        printf("SHA-3 Self-Test OK!\n");
+    if (test_sha3() == 0 && test_shake() == 0)
+        printf("FIPS 202 / SHA3, SHAKE128, SHAKE256 Self-Tests OK!\n");
 
     return 0;
 }
+
diff --git a/sha3.c b/sha3.c
index 6005671..3d6a7fe 100644
--- a/sha3.c
+++ b/sha3.c
@@ -124,7 +124,7 @@ int sha3_update(sha3_ctx_t *c, const void *data, size_t len)
     for (i = 0; i < len; i++) {
         c->st.b[j++] ^= ((const uint8_t *) data)[i];
         if (j >= c->rsiz) {
-            sha3_keccakf(c->st.q, SHA3_ROUNDS);
+            sha3_keccakf(c->st.q, KECCAKF_ROUNDS);
             j = 0;
         }
     }
@@ -141,7 +141,7 @@ int sha3_final(void *md, sha3_ctx_t *c)
 
     c->st.b[c->pt] ^= 0x06;
     c->st.b[c->rsiz - 1] ^= 0x80;
-    sha3_keccakf(c->st.q, SHA3_ROUNDS);
+    sha3_keccakf(c->st.q, KECCAKF_ROUNDS);
 
     for (i = 0; i < c->mdlen; i++) {
         ((uint8_t *) md)[i] = c->st.b[i];
@@ -163,3 +163,29 @@ void *sha3(const void *in, size_t inlen, void *md, int mdlen)
     return md;
 }
 
+// SHAKE128 and SHAKE256 extensible-output functionality
+
+void shake_xof(sha3_ctx_t *c)
+{
+    c->st.b[c->pt] ^= 0x1F;
+    c->st.b[c->rsiz - 1] ^= 0x80;
+    sha3_keccakf(c->st.q, KECCAKF_ROUNDS);
+    c->pt = 0;
+}
+
+void shake_out(sha3_ctx_t *c, void *out, size_t len)
+{
+    size_t i;
+    int j;
+
+    j = c->pt;
+    for (i = 0; i < len; i++) {
+        if (j >= c->rsiz) {
+            sha3_keccakf(c->st.q, KECCAKF_ROUNDS);
+            j = 0;
+        }
+        ((uint8_t *) out)[i] = c->st.b[j++];
+    }
+    c->pt = j;
+}
+
diff --git a/sha3.h b/sha3.h
index ab86bde..02b1e9b 100644
--- a/sha3.h
+++ b/sha3.h
@@ -7,8 +7,8 @@
 #include <stddef.h>
 #include <stdint.h>
 
-#ifndef SHA3_ROUNDS
-#define SHA3_ROUNDS 24
+#ifndef KECCAKF_ROUNDS
+#define KECCAKF_ROUNDS 24
 #endif
 
 #ifndef ROTL64
@@ -32,5 +32,13 @@ int sha3_final(void *md, sha3_ctx_t *c);    // digest goes to md
 // compute a sha3 hash (md) of given byte length from "in"
 void *sha3(const void *in, size_t inlen, void *md, int mdlen);
 
+// SHAKE128 and SHAKE256 extensible-output functions
+#define shake128_init(c) sha3_init(c, 16)
+#define shake256_init(c) sha3_init(c, 32)
+#define shake_update sha3_update
+
+void shake_xof(sha3_ctx_t *c);
+void shake_out(sha3_ctx_t *c, void *out, size_t len);
+
 #endif