mailgosend/internal/totp/totp.go

// Package totp implements RFC 6238 TOTP using stdlib only.
package totp

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1" //nolint:gosec — RFC 6238 mandates SHA-1 for TOTP
	"encoding/base32"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"math"
	"strings"
	"sync"
	"time"
)

const (
	// Digits is the number of OTP digits (6).
	Digits = 6
	// Period is the TOTP time step in seconds (30).
	Period = 30
	// Window is the number of adjacent windows to accept (±1 = 3 windows total).
	Window = 1
	// SecretBytes is the raw secret length (20 bytes = 160 bits).
	SecretBytes = 20
	// RecoveryCodeCount is the number of single-use backup codes generated.
	RecoveryCodeCount = 10
	// RecoveryCodeLen is the length of each backup code (characters, base32 subset).
	RecoveryCodeLen = 8
)

// ---- Replay protection ----

// replayCache prevents reuse of a TOTP code within its valid window (~90 s).
// Key: "secret:step:code" — value: expiry time.
var (
	replayMu    sync.Mutex
	replayCache = map[string]time.Time{}
)

// replayCacheKey builds the dedup key from the raw secret bytes, time step, and code.
func replayCacheKey(secret string, step int64, code string) string {
	return fmt.Sprintf("%s:%d:%s", secret, step, code)
}

// markUsed records a code as consumed. Expires after (Window+2)*Period seconds.
func markUsed(secret string, step int64, code string) {
	key := replayCacheKey(secret, step, code)
	exp := time.Now().Add(time.Duration((Window+2)*Period) * time.Second)
	replayMu.Lock()
	replayCache[key] = exp
	replayMu.Unlock()
}

// isReplay returns true when the code has already been consumed for this step.
func isReplay(secret string, step int64, code string) bool {
	key := replayCacheKey(secret, step, code)
	replayMu.Lock()
	exp, ok := replayCache[key]
	replayMu.Unlock()
	return ok && time.Now().Before(exp)
}

// pruneReplayCache removes expired entries. Caller must NOT hold replayMu.
func pruneReplayCache() {
	now := time.Now()
	replayMu.Lock()
	for k, exp := range replayCache {
		if now.After(exp) {
			delete(replayCache, k)
		}
	}
	replayMu.Unlock()
}

// GenerateSecret returns a cryptographically random base32-encoded TOTP secret.
func GenerateSecret() (string, error) {
	raw := make([]byte, SecretBytes)
	if _, err := rand.Read(raw); err != nil {
		return "", fmt.Errorf("totp: generate secret: %w", err)
	}
	return base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(raw), nil
}

// OTPAuthURI builds the otpauth:// URI for a QR code or manual import.
func OTPAuthURI(secret, accountName, issuer string) string {
	enc := func(s string) string {
		// RFC 3986 percent-encode for URI path/query
		var buf strings.Builder
		for _, b := range []byte(s) {
			if (b >= 'A' && b <= 'Z') || (b >= 'a' && b <= 'z') ||
				(b >= '0' && b <= '9') || b == '-' || b == '_' || b == '.' || b == '~' {
				buf.WriteByte(b)
			} else {
				fmt.Fprintf(&buf, "%%%02X", b)
			}
		}
		return buf.String()
	}
	return fmt.Sprintf("otpauth://totp/%s:%s?secret=%s&issuer=%s&algorithm=SHA1&digits=%d&period=%d",
		enc(issuer), enc(accountName), secret, enc(issuer), Digits, Period)
}

// Verify checks a 6-digit code against the secret for the current time window (±Window steps).
// Returns true if valid. Codes are single-use within their validity window (replay protection).
// secret is base32-encoded (no padding).
func Verify(secret, code string) bool {
	raw, err := base32.StdEncoding.WithPadding(base32.NoPadding).DecodeString(strings.ToUpper(secret))
	if err != nil {
		return false
	}
	pruneReplayCache()
	now := time.Now().Unix()
	step := now / Period
	for w := int64(-Window); w <= int64(Window); w++ {
		s := step + w
		if hotp(raw, s) == code {
			if isReplay(secret, s, code) {
				return false // replay: code already used in this window
			}
			markUsed(secret, s, code)
			return true
		}
	}
	return false
}

// hotp computes an HOTP value for a key and counter.
func hotp(key []byte, counter int64) string {
	msg := make([]byte, 8)
	binary.BigEndian.PutUint64(msg, uint64(counter)) //nolint:gosec
	mac := hmac.New(sha1.New, key)
	mac.Write(msg)
	h := mac.Sum(nil)

	// Dynamic truncation (RFC 4226 §5.4)
	offset := h[len(h)-1] & 0x0f
	binCode := (uint32(h[offset]&0x7f) << 24) |
		(uint32(h[offset+1]) << 16) |
		(uint32(h[offset+2]) << 8) |
		uint32(h[offset+3])

	otp := binCode % uint32(math.Pow10(Digits))
	return fmt.Sprintf("%0*d", Digits, otp)
}

// ---- Recovery codes ----

// GenerateRecoveryCodes returns RecoveryCodeCount random single-use codes.
func GenerateRecoveryCodes() ([]string, error) {
	codes := make([]string, RecoveryCodeCount)
	const charset = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789" // no O/0/I/1 to avoid confusion
	buf := make([]byte, RecoveryCodeLen)
	for i := range codes {
		if _, err := rand.Read(buf); err != nil {
			return nil, fmt.Errorf("totp: generate recovery codes: %w", err)
		}
		var sb strings.Builder
		for _, b := range buf {
			sb.WriteByte(charset[int(b)%len(charset)])
		}
		codes[i] = sb.String()
	}
	return codes, nil
}

// EncodeRecoveryCodes marshals a code slice to JSON bytes for storage.
func EncodeRecoveryCodes(codes []string) ([]byte, error) {
	return json.Marshal(codes)
}

// DecodeRecoveryCodes unmarshals JSON bytes back to a code slice.
func DecodeRecoveryCodes(data []byte) ([]string, error) {
	var codes []string
	if err := json.Unmarshal(data, &codes); err != nil {
		return nil, fmt.Errorf("totp: decode recovery codes: %w", err)
	}
	return codes, nil
}

// ConsumeRecoveryCode removes a matching code (case-insensitive) and returns the
// updated slice. Returns (nil, false) if code not found.
func ConsumeRecoveryCode(codes []string, input string) ([]string, bool) {
	input = strings.ToUpper(strings.TrimSpace(input))
	for i, c := range codes {
		if strings.EqualFold(c, input) {
			updated := make([]string, 0, len(codes)-1)
			updated = append(updated, codes[:i]...)
			updated = append(updated, codes[i+1:]...)
			return updated, true
		}
	}
	return nil, false
}