From 0857b296247dd028c2d176ffc5d7f7e37e288c25 Mon Sep 17 00:00:00 2001
From: Kyle Isom <kyle@imap.cc>
Date: Sun, 9 Dec 2018 22:01:00 -0800
Subject: [PATCH] Actually support clock mocking.

---
 totp.go      |  8 ++++++--
 totp_test.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/totp.go b/totp.go
index 1461c2c..9259e6c 100644
--- a/totp.go
+++ b/totp.go
@@ -15,7 +15,7 @@ import (
 	"github.com/benbjohnson/clock"
 )
 
-var time clock.Clock
+var timeSource = clock.New()
 
 // TOTP represents an RFC 6238 Time-based One-Time Password instance.
 type TOTP struct {
@@ -53,7 +53,7 @@ func (otp *TOTP) otpCounter(t uint64) uint64 {
 
 // OTPCounter returns the current time value for the OTP.
 func (otp *TOTP) OTPCounter() uint64 {
-	return otp.otpCounter(uint64(time.Now().Unix()))
+	return otp.otpCounter(uint64(timeSource.Now().Unix()))
 }
 
 // NewTOTP takes a new key, a starting time, a step, the number of
@@ -166,3 +166,7 @@ func totpFromURL(u *url.URL) (*TOTP, string, error) {
 func (otp *TOTP) QR(label string) ([]byte, error) {
 	return otp.OATH.QR(otp.Type(), label)
 }
+
+func SetClock(c clock.Clock) {
+	timeSource = c
+}
diff --git a/totp_test.go b/totp_test.go
index 244e43e..dc1e438 100644
--- a/totp_test.go
+++ b/totp_test.go
@@ -4,6 +4,9 @@ import (
 	"crypto"
 	"fmt"
 	"testing"
+	"time"
+
+	"github.com/benbjohnson/clock"
 )
 
 var rfcTotpKey = []byte("12345678901234567890")
@@ -53,3 +56,28 @@ func TestTotpRFC(t *testing.T) {
 		}
 	}
 }
+
+func TestTOTPTime(t *testing.T) {
+	otp := GenerateGoogleTOTP()
+
+	testClock := clock.NewMock()
+	testClock.Add(2*time.Minute)
+	SetClock(testClock)
+
+	code := otp.OTP()
+
+	testClock.Add(-1 * time.Minute)
+	if newCode := otp.OTP(); newCode == code {
+		t.Errorf("twofactor: TOTP: previous code %s shouldn't match code %s", newCode, code)
+	}
+
+	testClock.Add(2 * time.Minute)
+	if newCode := otp.OTP(); newCode == code {
+		t.Errorf("twofactor: TOTP: future code %s shouldn't match code %s", newCode, code)
+	}
+
+	testClock.Add(-1 * time.Minute)
+	if newCode := otp.OTP(); newCode != code {
+		t.Errorf("twofactor: TOTP: current code %s shouldn't match code %s", newCode, code)
+	}
+}