diff --git a/helper/totp.go b/helper/totp.go index 42cf7b9..d32b9c2 100644 --- a/helper/totp.go +++ b/helper/totp.go @@ -33,6 +33,9 @@ func GenerateOTPCode(token string, when time.Time) (string, error) { // It should be uppercase always token = strings.ToUpper(token) + // Remove all the extra "=" padding at the end + token = strings.TrimRight(token, "=") + secretBytes, err := base32.StdEncoding.WithPadding(base32.NoPadding).DecodeString(token) if err != nil { return "", fmt.Errorf("Decoding token string: %w", err) diff --git a/helper/totp_test.go b/helper/totp_test.go new file mode 100644 index 0000000..aad32b6 --- /dev/null +++ b/helper/totp_test.go @@ -0,0 +1,36 @@ +package helper + +import ( + "testing" + "time" +) + +var testCases = []struct { + description string + token string + expectErr bool +}{ + {"generates otpcode from token with padding", "PGWXXL7B66MMSRBAWSKEKIYD3P675KRJ===", false}, + {"generates otpcode from token without padding", "JBSWY3DPEHPK3PXPJBSWY3DPEHPK3PXP", false}, + {"invalid token format", "INVALIDTOKEN123", true}, +} + +func TestGenerateOTPCode(t *testing.T) { + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + code, err := GenerateOTPCode(tc.token, time.Now()) + + if tc.expectErr { + if err == nil { + t.Errorf("Expected error for input '%s', but got none", tc.token) + } + } else { + if err != nil { + t.Errorf("GenerateOTPCode returned an error: %s", err.Error()) + } else if len(code) != 6 { + t.Errorf("Expected 6-digit OTP, got: %s", code) + } + } + }) + } +}