ai-gateway-go/internal/identity/totp.go

package identity

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base32"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"net/url"
	"strings"
	"time"
)

const (
	totpPeriod       = int64(30)
	totpDigits       = 6
	totpWindow       = int64(1)
	backupCodeCount  = 10
	backupCodeLength = 8
)

var backupAlphabet = []byte("23456789ABCDEFGHJKLMNPQRSTUVWXYZ")

type BackupCodeRecord struct {
	Hash   string     `json:"hash"`
	UsedAt *time.Time `json:"used_at,omitempty"`
}

func GenerateTOTPSecret() (string, error) {
	secret := make([]byte, 20)
	if _, err := rand.Read(secret); err != nil {
		return "", err
	}
	return base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(secret), nil
}

func TOTPProvisioningURI(secret string, kind Kind, login string) string {
	issuer := "AI Gateway"
	label := issuer + ":" + string(kind) + ":" + login
	query := url.Values{
		"secret":    []string{secret},
		"issuer":    []string{issuer},
		"algorithm": []string{"SHA1"},
		"digits":    []string{fmt.Sprint(totpDigits)},
		"period":    []string{fmt.Sprint(totpPeriod)},
	}
	return "otpauth://totp/" + url.PathEscape(label) + "?" + query.Encode()
}

func VerifyTOTP(secret, code string, now time.Time) (int64, bool) {
	code = strings.TrimSpace(code)
	if len(code) != totpDigits {
		return 0, false
	}
	step := now.Unix() / totpPeriod
	for offset := -totpWindow; offset <= totpWindow; offset++ {
		candidate, err := hotp(secret, step+offset, totpDigits)
		if err == nil && subtle.ConstantTimeCompare([]byte(candidate), []byte(code)) == 1 {
			return step + offset, true
		}
	}
	return 0, false
}

func hotp(secret string, counter int64, digits int) (string, error) {
	if counter < 0 || digits < 6 || digits > 8 {
		return "", errors.New("invalid HOTP parameters")
	}
	decoded, err := base32.StdEncoding.WithPadding(base32.NoPadding).DecodeString(strings.ToUpper(strings.TrimSpace(secret)))
	if err != nil {
		return "", err
	}
	message := make([]byte, 8)
	binary.BigEndian.PutUint64(message, uint64(counter))
	mac := hmac.New(sha1.New, decoded)
	_, _ = mac.Write(message)
	digest := mac.Sum(nil)
	offset := digest[len(digest)-1] & 0x0f
	value := (uint32(digest[offset])&0x7f)<<24 |
		uint32(digest[offset+1])<<16 |
		uint32(digest[offset+2])<<8 |
		uint32(digest[offset+3])
	modulus := uint32(1)
	for i := 0; i < digits; i++ {
		modulus *= 10
	}
	return fmt.Sprintf("%0*d", digits, value%modulus), nil
}

func GenerateBackupCodes() ([]string, []BackupCodeRecord, error) {
	codes := make([]string, 0, backupCodeCount)
	records := make([]BackupCodeRecord, 0, backupCodeCount)
	for range backupCodeCount {
		random := make([]byte, backupCodeLength)
		if _, err := rand.Read(random); err != nil {
			return nil, nil, err
		}
		for index := range random {
			// 拒绝采样消除取模偏差:256 % 31 = 8,直接取模会让前 8 个字符
			// 的概率略高于其余字符。丢弃 248..255 的取值后分布均匀。
			value := random[index]
			for value >= 248 {
				if _, err := rand.Read(random[index : index+1]); err != nil {
					return nil, nil, err
				}
				value = random[index]
			}
			random[index] = backupAlphabet[int(value)%len(backupAlphabet)]
		}
		raw := string(random)
		code := raw[:4] + "-" + raw[4:]
		codes = append(codes, code)
		records = append(records, BackupCodeRecord{Hash: HashBackupCode(code)})
	}
	return codes, records, nil
}

func HashBackupCode(code string) string {
	normalized := strings.ToUpper(strings.ReplaceAll(strings.TrimSpace(code), "-", ""))
	digest := sha256.Sum256([]byte(normalized))
	return hex.EncodeToString(digest[:])
}