Handle the code generation with TOTP

README.md

@@ -60,8 +60,7 @@ Then, go to `http://localhost:8080/auth/login` to see the login page.
 
 ### Roadmap
 
-- Custom token storage mechanism
-- Out of the box time bound token generation
+- Out-of-the-box support for generating time-bound tokens using TOTP (Time-Based One-Time Password).
 - Customizable templates (Bring your own).
-- Time based token expiration out the box
-- Prevend CSRF attacks with token
+- Automatically handle token expiration based on time, providing security and convenience.
+- Prevend CSRF attacks with token.

codes.go

@@ -1,28 +1,85 @@
 package maildoor
 
 import (
-	"math/rand"
+	"crypto/hmac"
+	"crypto/rand"
+	"crypto/sha1"
+	"encoding/base32"
+	"encoding/binary"
+	"fmt"
 	"sync"
+	"time"
 )
 
 var (
 	tux     sync.Mutex
-	codes   = map[string]string{}
-	letters = []rune("1234567890")
+	secrets = map[string]string{}
 )
 
-// newCodeFor generates a new code for the email and stores it in the codes map.
-// tokens are always 6 characters long.
-func newCodeFor(email string) string {
-	// Generating a new token
-	b := make([]rune, 6)
-	for i := range b {
-		b[i] = letters[rand.Intn(len(letters))]
+const (
+	// expiresTime is the time in seconds that the code expires
+	// 2 minutes
+	expiresTime = 120
+)
+
+// generateSecret generates a 16 byte base32 encoded secret
+// The secret is then returned
+func generateSecret() (string, error) {
+	secret := make([]byte, 10)
+	_, err := rand.Read(secret)
+	if err != nil {
+		return "", err
 	}
 
+	return base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(secret), nil
+}
+
+// generateCode calls the gen function to generate a 6-digit code
+// using the secret and the time step
+func generateCode(secret string) string {
+	timeStep := time.Now().Unix() / expiresTime
+	return gen(secret, timeStep)
+}
+
+// validateCode validates the codeToValid with the secret
+// by generating the code for the current time step and the previous and next time steps
+// If the code matches any of the generated codes, it returns true otherwise false
+func validateCode(codeToValid, secret string) bool {
+	timeStep := time.Now().Unix() / expiresTime
+	// Check the current time step, the previous and the next time steps
+	// to allow for some time drift
+	for i := -1; i <= 1; i++ {
+		code := gen(secret, timeStep+int64(i))
+		if code == codeToValid {
+			return true
+		}
+	}
+
+	return false
+}
+
+// gen generates a 6-digit code using the secret and the time step
+// The code is then returned
+func gen(secret string, timeStep int64) string {
+	timeBytes := make([]byte, 8)
+	binary.BigEndian.PutUint64(timeBytes, uint64(timeStep))
+
+	key := []byte(secret)
+	hmacSha1 := hmac.New(sha1.New, key)
+	hmacSha1.Write(timeBytes)
+	hash := hmacSha1.Sum(nil)
+
+	offset := hash[len(hash)-1] & 0x0f
+	codeBytes := binary.BigEndian.Uint32(hash[offset : offset+4])
+
+	code := codeBytes % 1_000_000
+	codeStr := fmt.Sprintf("%06d", code)
+
+	return codeStr
+}
+
+func saveSecret(email, secret string) {
 	tux.Lock()
 	defer tux.Unlock()
-	codes[email] = string(b)
-
-	return string(b)
+	secrets[email] = secret
 }

handle_code.go

@@ -10,31 +10,39 @@ func (m *maildoor) handleCode(w http.ResponseWriter, r *http.Request) {
 	email := r.FormValue("email")
 	code := r.FormValue("code")
 
-	// Find a combination of token and email in the server
-	// call the afterlogin hook with the email
-	// remove the token from the server
-	if code != codes[email] {
-		data := atempt{
-			Email:       email,
-			Error:       "Invalid token",
-			Logo:        m.logoURL,
-			Icon:        m.iconURL,
-			ProductName: m.productName,
-		}
-
-		err := m.render(w, data, "layout.html", "handle_code.html")
-		if err != nil {
-			m.httpError(w, err)
-
-			return
-		}
+	secret := secrets[email]
 
+	if code == "" {
+		renderError(m, w, email, "Code is required")
 		return
 	}
 
-	delete(codes, email)
+	if !validateCode(code, secret) {
+		renderError(m, w, email, "Invalid code")
+
+		return
+	}
+
+	delete(secrets, email)
 
 	// Adding email to the context
 	r = r.WithContext(context.WithValue(r.Context(), "email", email))
 	m.afterLogin(w, r)
 }
+
+func renderError(m *maildoor, w http.ResponseWriter, email, errorMessage string) {
+	data := atempt{
+		Email:       email,
+		Error:       errorMessage,
+		Logo:        m.logoURL,
+		Icon:        m.iconURL,
+		ProductName: m.productName,
+	}
+
+	err := m.render(w, data, "layout.html", "handle_code.html")
+	if err != nil {
+		m.httpError(w, err)
+
+		return
+	}
+}

handle_email.go

@@ -25,8 +25,14 @@ func (m *maildoor) handleEmail(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	token := newCodeFor(email)
-	html, txt, err := m.mailBodies(token)
+	secret, err := generateSecret()
+	if err != nil {
+		m.httpError(w, err)
+	}
+
+	saveSecret(email, secret)
+
+	html, txt, err := m.mailBodies(generateCode(secret))
 	if err != nil {
 		m.httpError(w, err)
 		return