Improve [HOTP & TOTP] Separate signature verification logic (#94)

- [+] refactor(otpverifier): separate signature verification logic and add GenerateToken method - [+] feat(otpverifier): add support for verifying token without signature - [+] test(otpverifier): add tests for missing signature and signature mismatch scenarios
H0llyW00dzZ · May 30, 2024 · 0d47e07 · 0d47e07
1 parent 82159c8
commit 0d47e07
Show file tree

Hide file tree

Showing 7 changed files with 252 additions and 77 deletions.
diff --git a/2fa_test.go b/2fa_test.go
@@ -269,8 +269,8 @@ func TestMiddleware_Handle(t *testing.T) {
 		Secret: info.Secret,
 	})
 
-	validToken, _ := totp.GenerateToken()
-	totp.Verify(validToken, "")
+	validToken := totp.GenerateToken()
+	totp.Verify(validToken)
 
 	// Create a separate instance of the Middleware struct for testing
 	testMiddleware := &twofa.Middleware{
@@ -1279,8 +1279,8 @@ func TestMiddlewareUUIDContextKey_Handle(t *testing.T) {
 		Secret: info.Secret,
 	})
 
-	validToken, _ := totp.GenerateToken()
-	totp.Verify(validToken, "")
+	validToken := totp.GenerateToken()
+	totp.Verify(validToken)
 
 	// Create a separate instance of the Middleware struct for testing
 	testMiddleware := &twofa.Middleware{

diff --git a/benchmark_test.go b/benchmark_test.go
@@ -90,8 +90,8 @@ func BenchmarkJSONSonicWithValid2FA(b *testing.B) {
 		Secret: twoFAConfig.Secret,
 	})
 
-	token, _ := totp.GenerateToken()
-	totp.Verify(token, "")
+	token := totp.GenerateToken()
+	totp.Verify(token)
 
 	// Create a valid 2FA cookie
 	cookieValue := twoFAMiddleware.GenerateCookieValue(time.Now().Add(time.Duration(86400) * time.Second))
@@ -154,8 +154,8 @@ func BenchmarkJSONSonicWithValidCookie(b *testing.B) {
 		Secret: twoFAConfig.Secret,
 	})
 
-	token, _ := totp.GenerateToken()
-	totp.Verify(token, "")
+	token := totp.GenerateToken()
+	totp.Verify(token)
 
 	// Create a valid 2FA cookie
 	cookieValue := twoFAMiddleware.GenerateCookieValue(time.Now().Add(time.Duration(86400) * time.Second))
@@ -280,8 +280,8 @@ func BenchmarkJSONStdLibraryMiddlewareWithValid2FA(b *testing.B) {
 		Secret: twoFAConfig.Secret,
 	})
 
-	token, _ := totp.GenerateToken()
-	totp.Verify(token, "")
+	token := totp.GenerateToken()
+	totp.Verify(token)
 
 	// Create a valid 2FA cookie
 	cookieValue := twoFAMiddleware.GenerateCookieValue(time.Now().Add(time.Duration(86400) * time.Second))
@@ -344,8 +344,8 @@ func BenchmarkJSONStdLibraryWithValidCookie(b *testing.B) {
 		Secret: twoFAConfig.Secret,
 	})
 
-	token, _ := totp.GenerateToken()
-	totp.Verify(token, "")
+	token := totp.GenerateToken()
+	totp.Verify(token)
 
 	// Create a valid 2FA cookie
 	cookieValue := twoFAMiddleware.GenerateCookieValue(time.Now().Add(time.Duration(86400) * time.Second))

diff --git a/internal/otpverifier/benchmark_test.go b/internal/otpverifier/benchmark_test.go
@@ -41,8 +41,8 @@ func BenchmarkTOTPVerify(b *testing.B) {
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
 				// Note: This now each token are different
-				token, _ := verifier.GenerateToken()
-				verifier.Verify(token, "")
+				token := verifier.GenerateToken()
+				verifier.Verify(token)
 			}
 		})
 
@@ -58,7 +58,7 @@ func BenchmarkTOTPVerify(b *testing.B) {
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
 				// Note: This now each token are different
-				token, signature := verifier.GenerateToken()
+				token, signature := verifier.GenerateTokenWithSignature()
 				verifier.Verify(token, signature)
 			}
 		})
@@ -96,8 +96,8 @@ func BenchmarkHOTPVerify(b *testing.B) {
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
 				// Note: This now each token are different
-				token, _ := verifier.GenerateToken()
-				verifier.Verify(token, "")
+				token := verifier.GenerateToken()
+				verifier.Verify(token)
 			}
 		})
 
@@ -113,7 +113,7 @@ func BenchmarkHOTPVerify(b *testing.B) {
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
 				// Note: This now each token are different
-				token, signature := verifier.GenerateToken()
+				token, signature := verifier.GenerateTokenWithSignature()
 				verifier.Verify(token, signature)
 			}
 		})

diff --git a/internal/otpverifier/hotp.go b/internal/otpverifier/hotp.go
@@ -63,26 +63,31 @@ func NewHOTPVerifier(config ...Config) *HOTPVerifier {
 // A successful verification will result in the counter being updated to the next expected value.
 //
 // Note: A firm grasp of the sync window concept is essential for understanding its role in the verification process.
-func (v *HOTPVerifier) Verify(token, signature string) bool {
+func (v *HOTPVerifier) Verify(token string, signature ...string) bool {
 	if v.config.SyncWindow < 0 {
 		panic("hotp: SyncWindow must be greater than or equal to zero")
 	}
 
+	if v.config.UseSignature {
+		if len(signature) == 0 {
+			panic("hotp: Signature is required but not provided")
+		}
+
+		return v.verifyWithSignature(token, signature[0])
+	}
+
+	return v.verifyWithoutSignature(token)
+}
+
+// verifyWithoutSignature checks if the provided token is valid for the specified counter value without signature verification.
+func (v *HOTPVerifier) verifyWithoutSignature(token string) bool {
 	// Check if sync window is applied
 	// Note: Understanding this sync window requires skilled mathematical reasoning.
 	if v.config.SyncWindow > 0 {
 		for i := 0; i <= v.config.SyncWindow; i++ {
 			expectedCounter := int(v.config.Counter) + i
 			generatedToken := v.Hotp.At(expectedCounter)
-			tokenMatch := subtle.ConstantTimeCompare([]byte(token), []byte(generatedToken)) == 1
-			signatureMatch := true // Assume true if not using signatures.
-
-			if v.config.UseSignature {
-				generatedSignature := v.generateSignature(generatedToken)
-				signatureMatch = subtle.ConstantTimeCompare([]byte(signature), []byte(generatedSignature)) == 1
-			}
-
-			if tokenMatch && signatureMatch {
+			if subtle.ConstantTimeCompare([]byte(token), []byte(generatedToken)) == 1 {
 				// Update the stored counter to the next expected value after a successful match
 				v.config.Counter = uint64(expectedCounter + 1)
 				return true
@@ -94,29 +99,55 @@ func (v *HOTPVerifier) Verify(token, signature string) bool {
 
 	// Default case when sync window is not applied
 	generatedToken := v.Hotp.At(int(v.config.Counter))
-	tokenMatch := subtle.ConstantTimeCompare([]byte(token), []byte(generatedToken)) == 1
-	signatureMatch := true // Assume true if not using signatures.
+	if subtle.ConstantTimeCompare([]byte(token), []byte(generatedToken)) == 1 {
+		// Increment the counter value after successful verification
+		v.config.Counter++
+		return true
+	}
+	return false
+}
 
-	if v.config.UseSignature {
-		generatedSignature := v.generateSignature(generatedToken)
-		signatureMatch = subtle.ConstantTimeCompare([]byte(signature), []byte(generatedSignature)) == 1
+// verifyWithSignature checks if the provided token and signature are valid for the specified counter value.
+func (v *HOTPVerifier) verifyWithSignature(token, signature string) bool {
+	// Check if sync window is applied
+	// Note: Understanding this sync window requires skilled mathematical reasoning.
+	if v.config.SyncWindow > 0 {
+		for i := 0; i <= v.config.SyncWindow; i++ {
+			expectedCounter := int(v.config.Counter) + i
+			generatedToken := v.Hotp.At(expectedCounter)
+			generatedSignature := v.generateSignature(generatedToken)
+			if subtle.ConstantTimeCompare([]byte(token), []byte(generatedToken)) == 1 &&
+				subtle.ConstantTimeCompare([]byte(signature), []byte(generatedSignature)) == 1 {
+				// Update the stored counter to the next expected value after a successful match
+				v.config.Counter = uint64(expectedCounter + 1)
+				return true
+			}
+		}
+		// If no match is found within the synchronization window, authentication fails
+		return false
 	}
 
-	if tokenMatch && signatureMatch {
+	// Default case when sync window is not applied
+	generatedToken := v.Hotp.At(int(v.config.Counter))
+	generatedSignature := v.generateSignature(generatedToken)
+	if subtle.ConstantTimeCompare([]byte(token), []byte(generatedToken)) == 1 &&
+		subtle.ConstantTimeCompare([]byte(signature), []byte(generatedSignature)) == 1 {
 		// Increment the counter value after successful verification
 		v.config.Counter++
 		return true
 	}
 	return false
 }
 
-// GenerateToken generates a token and signature for the current counter value.
-func (v *HOTPVerifier) GenerateToken() (string, string) {
+// GenerateToken generates a token for the current counter value.
+func (v *HOTPVerifier) GenerateToken() string {
+	return v.Hotp.At(int(v.config.Counter))
+}
+
+// GenerateTokenWithSignature generates a token and signature for the current counter value.
+func (v *HOTPVerifier) GenerateTokenWithSignature() (string, string) {
 	token := v.Hotp.At(int(v.config.Counter))
-	signature := ""
-	if v.config.UseSignature {
-		signature = v.generateSignature(token)
-	}
+	signature := v.generateSignature(token)
 	return token, signature
 }
 

diff --git a/internal/otpverifier/otpverifier.go b/internal/otpverifier/otpverifier.go
@@ -111,8 +111,9 @@ type TimeSource func() time.Time
 
 // OTPVerifier is an interface that defines the behavior of an OTP verifier.
 type OTPVerifier interface {
-	Verify(token, signature string) bool
-	GenerateToken() (string, string)
+	Verify(token string, signature ...string) bool
+	GenerateToken() string
+	GenerateTokenWithSignature() (string, string)
 	SetCounter(counter uint64)
 	GetCounter() uint64
 	GenerateOTPURL(issuer, accountName string) string