lib/totp: make the New to accept only hash with SHA1, SHA256, or SHA512

Previously, the first parameter to New is a function that return
hash.Hash.  This signature can be misleading, because md5.New also
can return hash.Hash but not usable in TOTP.

This changes update the New function signature to accept defined
type that can be set to SHA1, SHA256, or SHA512.

4 files changed, 48 insertions, 19 deletions

diff --git a/cmd/totp/main.go b/cmd/totp/main.go
index 75aebc0b..f046f708 100644
--- a/cmd/totp/main.go
+++ b/cmd/totp/main.go
@@ -4,13 +4,9 @@
 package main
 
 import (
-	"crypto/sha1"
-	"crypto/sha256"
-	"crypto/sha512"
 	"encoding/base32"
 	"flag"
 	"fmt"
-	"hash"
 	"log"
 	"os"
 
@@ -35,17 +31,17 @@ func main() {
 		flag.Usage()
 	}
 
-	var hashFn func() hash.Hash
+	var cryptoHash totp.CryptoHash
 	switch *paramHash {
 	case "sha256":
-		hashFn = sha256.New
+		cryptoHash = totp.CryptoHashSHA256
 	case "sha512":
-		hashFn = sha512.New
+		cryptoHash = totp.CryptoHashSHA512
 	default:
-		hashFn = sha1.New
+		cryptoHash = totp.CryptoHashSHA1
 	}
 
-	totproto := totp.New(hashFn, *paramDigits, *paramTimestep)
+	totproto := totp.New(cryptoHash, *paramDigits, *paramTimestep)
 	secret, err := base32.StdEncoding.DecodeString(os.Args[1])
 	if err != nil {
 		log.Fatal(err)
diff --git a/lib/totp/totp.go b/lib/totp/totp.go
index bb442da4..1dde8cae 100644
--- a/lib/totp/totp.go
+++ b/lib/totp/totp.go
@@ -7,7 +7,11 @@
 package totp
 
 import (
+	"crypto"
 	"crypto/hmac"
+	"crypto/sha1"
+	"crypto/sha256"
+	"crypto/sha512"
 	"encoding/hex"
 	"fmt"
 	"hash"
@@ -15,7 +19,21 @@ import (
 	"time"
 )
 
+type CryptoHash crypto.Hash
+
+// List of available hash function that can be used in TOTP.
+//
+// See RFC 6238 Section 1.2.
+const (
+	CryptoHashSHA1   CryptoHash = CryptoHash(crypto.SHA1) // Default hash algorithm.
+	CryptoHashSHA256            = CryptoHash(crypto.SHA256)
+	CryptoHashSHA512            = CryptoHash(crypto.SHA512)
+)
+
+// Default value for hash, digits, time-step, and maximum step backs.
 const (
+	DefHash = CryptoHashSHA1
+
 	// DefCodeDigits default digits generated when verifying or generating
 	// OTP.
 	DefCodeDigits = 6
@@ -40,11 +58,27 @@ type Protocol struct {
 }
 
 //
-// New create TOTP protocol for prover or verifier using "fnHash" as the hmac-sha
-// hash function, "codeDigits" as the number of digits to be generated
-// and/or verified, and "timeStep" as the time divisor.
+// New create TOTP protocol for prover or verifier using "cryptoHash" as the
+// hmac-sha hash function, "codeDigits" as the number of digits to be
+// generated and/or verified, and "timeStep" as the time divisor.
+//
+// There are only three hash functions that can be used: SHA1, SHA256, and
+// SHA512.
+// Passing hash value other than that, will revert the value default to SHA1.
+//
+// The maximum value for codeDigits parameter is 8.
 //
-func New(fnHash func() hash.Hash, codeDigits, timeStep int) Protocol {
+func New(cryptoHash CryptoHash, codeDigits, timeStep int) Protocol {
+	var fnHash func() hash.Hash
+
+	switch cryptoHash {
+	case CryptoHashSHA256:
+		fnHash = sha256.New
+	case CryptoHashSHA512:
+		fnHash = sha512.New
+	default:
+		fnHash = sha1.New
+	}
 	if codeDigits <= 0 || codeDigits > 8 {
 		codeDigits = DefCodeDigits
 	}
diff --git a/lib/totp/totp_example_test.go b/lib/totp/totp_example_test.go
index b2d1a07a..7e40ea01 100644
--- a/lib/totp/totp_example_test.go
+++ b/lib/totp/totp_example_test.go
@@ -5,7 +5,6 @@
 package totp
 
 import (
-	"crypto/sha1"
 	"encoding/hex"
 	"fmt"
 	"log"
@@ -19,7 +18,7 @@ func ExampleProtocol_Verify() {
 		log.Fatal(err)
 	}
 
-	p := New(sha1.New, DefCodeDigits, DefTimeStep)
+	p := New(CryptoHashSHA1, DefCodeDigits, DefTimeStep)
 	otp, _ := p.Generate(secret)
 
 	if p.Verify(secret, otp, 1) {
diff --git a/lib/totp/totp_test.go b/lib/totp/totp_test.go
index be1f9c54..a15c68e1 100644
--- a/lib/totp/totp_test.go
+++ b/lib/totp/totp_test.go
@@ -24,7 +24,7 @@ func TestProtocol_generateWithTimestamp_sha1(t *testing.T) {
 	}
 
 	mac := hmac.New(sha1.New, secretb)
-	p := New(sha1.New, 8, DefTimeStep)
+	p := New(CryptoHashSHA1, 8, DefTimeStep)
 
 	cases := []struct {
 		time int64
@@ -70,7 +70,7 @@ func TestProtocol_generateWithTimestamp_sha256(t *testing.T) {
 	}
 
 	mac := hmac.New(sha256.New, secretb)
-	p := New(sha256.New, 8, DefTimeStep)
+	p := New(CryptoHashSHA256, 8, DefTimeStep)
 
 	cases := []struct {
 		time int64
@@ -140,7 +140,7 @@ func TestProtocol_generateWithTimestamp_sha512(t *testing.T) {
 		exp:  "47863826",
 	}}
 
-	p := New(sha512.New, 8, DefTimeStep)
+	p := New(CryptoHashSHA512, 8, DefTimeStep)
 	mac := hmac.New(sha512.New, secretb)
 
 	for _, c := range cases {
@@ -163,7 +163,7 @@ func TestProtocol_verifyWithTimestamp(t *testing.T) {
 	}
 
 	mac := hmac.New(sha1.New, secretb)
-	p := New(sha1.New, 8, DefTimeStep)
+	p := New(CryptoHashSHA1, 8, DefTimeStep)
 
 	cases := []struct {
 		desc  string