From bdfa7c8a53ff181ac0c48814bd94806a91f5ae36 Mon Sep 17 00:00:00 2001
From: Emmanuel T Odeke <emmanuel@orijtech.com>
Date: Wed, 9 Dec 2020 08:36:23 -0800
Subject: [PATCH] Use ReverseWordMap to make O(1) lookups for IsMnemonicValid
 (#2)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Noticed while auditing a dependency of Cosmos-sdk code from
PR https://github.com/cosmos/cosmos-sdk/pull/8099, WordList
is loaded at init time, and so is ReverseWordMap.
Usually a slice is sufficient for doing linear lookups for
the existence of a string. However, given that IsMnemonicValid
requires word segements with 12 to 24 lengths, this means that
2048 * ~12 lookups in the worst case, and given that we already
at init time load up ReverseWordMap, it is important for us
perform a constant time lookup. The data also backs up the speed up
with the benchmark results below:

```shell
name               old time/op    new time/op    delta
IsMnemonicValid-8    24.3µs ± 1%     1.3µs ± 1%  -94.64%  (p=0.000 n=9+10)

name               old alloc/op   new alloc/op   delta
IsMnemonicValid-8      576B ± 0%      576B ± 0%     ~     (all equal)

name               old allocs/op  new allocs/op  delta
IsMnemonicValid-8      3.00 ± 0%      3.00 ± 0%     ~     (all equal)
```
---
 bench_test.go | 26 ++++++++++++++++++++++++++
 bip39.go      | 19 ++++---------------
 2 files changed, 30 insertions(+), 15 deletions(-)
 create mode 100644 bench_test.go

diff --git a/bench_test.go b/bench_test.go
new file mode 100644
index 0000000..287fb0e
--- /dev/null
+++ b/bench_test.go
@@ -0,0 +1,26 @@
+package bip39
+
+import "testing"
+
+var words = []string{
+	"wolf afraid artwork blanket carpet cricket wolf afraid artwork blanket carpet cricket",
+	"artwork blanket carpet cricket disorder disorder artwork blanket carpet cricket disorder disorder",
+	"carpet cricket disorder cricket cricket artwork carpet cricket disorder cricket cricket artwork ",
+}
+
+func BenchmarkIsMnemonicValid(b *testing.B) {
+	b.ReportAllocs()
+	var sharp interface{}
+	for i := 0; i < b.N; i++ {
+		for _, word := range words {
+			ok := IsMnemonicValid(word)
+			if !ok {
+				b.Fatal("returned false")
+			}
+			sharp = ok
+		}
+	}
+	if sharp == nil {
+		b.Fatal("benchmark was not run")
+	}
+}
diff --git a/bip39.go b/bip39.go
index 5386e7f..452f1f2 100644
--- a/bip39.go
+++ b/bip39.go
@@ -100,7 +100,7 @@ func MnemonicToByteArray(mnemonic string) ([]byte, error) {
 	modulo := big.NewInt(2048)
 	for _, v := range mnemonicSlice {
 		index, found := ReverseWordMap[v]
-		if found == false {
+		if !found {
 			return nil, fmt.Errorf("Word `%v` not found in reverse map", v)
 		}
 		add := big.NewInt(int64(index))
@@ -176,9 +176,7 @@ func NewSeed(mnemonic string, password string) []byte {
 // Currently only supports data up to 32 bytes
 func addChecksum(data []byte) []byte {
 	// Get first byte of sha256
-	hasher := sha256.New()
-	hasher.Write(data)
-	hash := hasher.Sum(nil)
+	hash := sha256.Sum256(data)
 	firstChecksumByte := hash[0]
 
 	// len() is in bytes so we divide by 4
@@ -231,25 +229,16 @@ func IsMnemonicValid(mnemonic string) bool {
 	numOfWords := len(words)
 
 	// The number of words should be 12, 15, 18, 21 or 24
-	if numOfWords%3 != 0 || numOfWords < 12 || numOfWords > 24 {
+	if numOfWords < 12 || numOfWords > 24 || numOfWords%3 != 0 {
 		return false
 	}
 
 	// Check if all words belong in the wordlist
 	for i := 0; i < numOfWords; i++ {
-		if !contains(WordList, words[i]) {
+		if _, ok := ReverseWordMap[words[i]]; !ok {
 			return false
 		}
 	}
 
 	return true
 }
-
-func contains(s []string, e string) bool {
-	for _, a := range s {
-		if a == e {
-			return true
-		}
-	}
-	return false
-}