Merge pull request #45 from Lexman42/random_by_custom_reader

use custom io.Reader to generate Keys
pquerna · Oct 9, 2019 · 468c2dd · 468c2dd
2 parents 43bebef + 766957e
commit 468c2dd
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/hotp/hotp.go b/hotp/hotp.go
@@ -19,6 +19,7 @@ package hotp
 
 import (
 	"github.com/pquerna/otp"
+	"io"
 
 	"crypto/hmac"
 	"crypto/rand"
@@ -152,6 +153,8 @@ type GenerateOpts struct {
 	Digits otp.Digits
 	// Algorithm to use for HMAC. Defaults to SHA1.
 	Algorithm otp.Algorithm
+	// Reader to use for generating HOTP Key.
+	Rand io.Reader
 }
 
 var b32NoPadding = base32.StdEncoding.WithPadding(base32.NoPadding)
@@ -175,14 +178,18 @@ func Generate(opts GenerateOpts) (*otp.Key, error) {
 		opts.Digits = otp.DigitsSix
 	}
 
+	if opts.Rand == nil {
+		opts.Rand = rand.Reader
+	}
+
 	// otpauth://totp/Example:alice@google.com?secret=JBSWY3DPEHPK3PXP&issuer=Example
 
 	v := url.Values{}
 	if len(opts.Secret) != 0 {
 		v.Set("secret", b32NoPadding.EncodeToString(opts.Secret))
 	} else {
 		secret := make([]byte, opts.SecretSize)
-		_, err := rand.Read(secret)
+		_, err := opts.Rand.Read(secret)
 		if err != nil {
 			return nil, err
 		}

diff --git a/totp/totp.go b/totp/totp.go
@@ -20,6 +20,7 @@ package totp
 import (
 	"github.com/pquerna/otp"
 	"github.com/pquerna/otp/hotp"
+	"io"
 
 	"crypto/rand"
 	"encoding/base32"
@@ -142,6 +143,8 @@ type GenerateOpts struct {
 	Digits otp.Digits
 	// Algorithm to use for HMAC. Defaults to SHA1.
 	Algorithm otp.Algorithm
+	// Reader to use for generating TOTP Key.
+	Rand io.Reader
 }
 
 var b32NoPadding = base32.StdEncoding.WithPadding(base32.NoPadding)
@@ -169,14 +172,18 @@ func Generate(opts GenerateOpts) (*otp.Key, error) {
 		opts.Digits = otp.DigitsSix
 	}
 
+	if opts.Rand == nil {
+		opts.Rand = rand.Reader
+	}
+
 	// otpauth://totp/Example:alice@google.com?secret=JBSWY3DPEHPK3PXP&issuer=Example
 
 	v := url.Values{}
 	if len(opts.Secret) != 0 {
 		v.Set("secret", b32NoPadding.EncodeToString(opts.Secret))
 	} else {
 		secret := make([]byte, opts.SecretSize)
-		_, err := rand.Read(secret)
+		_, err := opts.Rand.Read(secret)
 		if err != nil {
 			return nil, err
 		}