From bd08550d6eeb78e66b79574d6f2370be475b9aee Mon Sep 17 00:00:00 2001
From: lukechampine <luke.champine@gmail.com>
Date: Thu, 9 Jan 2020 15:38:44 -0500
Subject: [PATCH] add derive key mode

---
 blake3.go      | 15 +++++++++++++--
 blake3_test.go | 12 ++++++++++--
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/blake3.go b/blake3.go
index 90c3ad1..76b37b1 100644
--- a/blake3.go
+++ b/blake3.go
@@ -263,8 +263,8 @@ func newHasher(key [8]uint32, flags uint32, out_size int) *Hasher {
 	}
 }
 
-// New returns a Hasher for the regular hash function.
-// If key is nil, the hash is unkeyed.
+// New returns a Hasher for the specified size and key. If key is nil, the hash
+// is unkeyed.
 func New(size int, key []byte) *Hasher {
 	if key == nil {
 		return newHasher(IV, 0, size)
@@ -274,6 +274,17 @@ func New(size int, key []byte) *Hasher {
 	return newHasher(key_words, KEYED_HASH, size)
 }
 
+// NewFromDerivedKey returns a Hasher whose key was derived from the supplied
+// context string.
+func NewFromDerivedKey(size int, ctx string) *Hasher {
+	h := newHasher(IV, DERIVE_KEY_CONTEXT, KEY_LEN)
+	h.Write([]byte(ctx))
+	key := h.Sum(nil)
+	var key_words [8]uint32
+	words_from_litte_endian_bytes(key, key_words[:])
+	return newHasher(key_words, DERIVE_KEY_MATERIAL, size)
+}
+
 func (h *Hasher) push_stack(cv [8]uint32) {
 	h.cv_stack[h.cv_stack_len] = cv
 	h.cv_stack_len++
diff --git a/blake3_test.go b/blake3_test.go
index a62c558..4d0f40b 100644
--- a/blake3_test.go
+++ b/blake3_test.go
@@ -45,18 +45,26 @@ func TestVectors(t *testing.T) {
 	}
 
 	for _, vec := range vectors.Cases {
+		in := input[:vec.InputLen]
 		// regular
 		h := blake3.New(len(vec.Hash)/2, nil)
-		h.Write(input[:vec.InputLen])
+		h.Write(in)
 		if out := toHex(h.Sum(nil)); out != vec.Hash {
 			t.Errorf("output did not match test vector:\n\texpected: %v...\n\t     got: %v...", vec.Hash[:10], out[:10])
 		}
 		// keyed
 		h = blake3.New(len(vec.KeyedHash)/2, []byte(vectors.Key))
-		h.Write(input[:vec.InputLen])
+		h.Write(in)
 		if out := toHex(h.Sum(nil)); out != vec.KeyedHash {
 			t.Errorf("output did not match test vector:\n\texpected: %v...\n\t     got: %v...", vec.KeyedHash[:10], out[:10])
 		}
+		// derive key
+		const ctx = "BLAKE3 2019-12-27 16:29:52 test vectors context"
+		h = blake3.NewFromDerivedKey(len(vec.DeriveKey)/2, ctx)
+		h.Write(in)
+		if out := toHex(h.Sum(nil)); out != vec.DeriveKey {
+			t.Errorf("output did not match test vector:\n\texpected: %v...\n\t     got: %v...", vec.DeriveKey[:10], out[:10])
+		}
 	}
 }