From 5f8f7471b7c7cd3d64074d5c49971bc9231793e7 Mon Sep 17 00:00:00 2001
From: greensea <gs@bbxy.net>
Date: Sat, 17 Dec 2022 21:21:21 +0800
Subject: [PATCH] Add env variable NONCE_CLEANUP_PROB.

---
 .env.example |  3 +++
 api/api.go   | 14 ++++++++++++--
 api/web.go   | 20 ++++++++++++++------
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/.env.example b/.env.example
index 75637ea..5fd3c53 100644
--- a/.env.example
+++ b/.env.example
@@ -22,6 +22,9 @@ S3_BUCKET=wcaptcha
 S3_ACCESS_KEY=
 S3_SECRET_KEY=
 
+# 触发清理过期的 Nonce 操作的概率
+# Nonce cleanup probility
+NONCE_CLEANUP_PROB=0.01
 
 # 在调试时可将此项设为空
 # When debugging, leave this variable empty
diff --git a/api/api.go b/api/api.go
index 68dea2c..9d2b18d 100644
--- a/api/api.go
+++ b/api/api.go
@@ -25,6 +25,9 @@ const (
 	RSA_KEY_TTL  = 600
 )
 
+// 执行 Nonce 清理程序的概率
+var nonceCleanupProb float64 = 0.01
+
 // var S3 *s3kv.Storage
 var Store store.Storer
 
@@ -139,6 +142,13 @@ func InitGin() *gin.Engine {
 		os.Exit(0)
 	}
 
+	nonceCleanupProb, err = strconv.ParseFloat(os.Getenv("NONCE_CLEANUP_PROB"), 64)
+	if err != nil {
+		log.Printf("未设定 NONCE_CLEANUP_PROB 环境变量或设置不正确，将使用默认值 0.01")
+		nonceCleanupProb = 0.01
+	}
+	log.Printf("Nonce 清理程序触发概率被设定为 %f%%", nonceCleanupProb*100)
+
 	route := gin.Default()
 
 	route.Use(cors.Default())
@@ -216,7 +226,7 @@ func nonceSet(nonce string) {
 var isNonceCleaning bool = false
 
 // 以 prob 的概率，触发清理过期的 nonce 操作
-func nonceClean(prob float32) {
+func nonceCleanup(prob float64) {
 	if isNonceCleaning == true {
 		log.Printf("当前有另一个 Nonce 清理程序正在进行中，不会重复运行 Nonce 清理程序")
 		return
@@ -226,7 +236,7 @@ func nonceClean(prob float32) {
 		isNonceCleaning = false
 	}()
 
-	r := rand.Float32()
+	r := rand.Float64()
 	if r >= prob {
 		return
 	}
diff --git a/api/web.go b/api/web.go
index 8f1fc19..a6bac16 100644
--- a/api/web.go
+++ b/api/web.go
@@ -184,7 +184,7 @@ func webCaptchaProblem(c *gin.Context) {
 	//question := fmt.Sprintf("%s.%s.%s", nB64, xB64, hB64)
 
 	// 以 1% 的概率去触发 Nonce 清除操作
-	go nonceClean(001)
+	go nonceCleanup(nonceCleanupProb)
 
 	c.JSON(200, gin.H{
 		"code": 0,
@@ -205,6 +205,8 @@ func webCaptchaProblem(c *gin.Context) {
 // 其中 X 是此前服务器返回的 x 的原始内容；Y 是计算结果，为小写的十六进制表达式；H 为此前服务器返回的签名原始内容；
 // <del>N 为模数，为小写的十六进制格式</del>
 func webCaptchaVerify(c *gin.Context) {
+	var verifyElapsed time.Duration
+
 	// 1. 解析客户端数据
 	var req struct {
 		Prove  string `form:"prove" binding:"required"`
@@ -265,16 +267,19 @@ func webCaptchaVerify(c *gin.Context) {
 	}
 
 	/// 3.1 使用最新的 RSAKey 检查结果是否正确
+	stime := time.Now()
+
 	if ts-site.RSAKeyCreateTime < RSA_KEY_TTL {
 		v = vdf.New(site.RSAKey.Primes[0], site.RSAKey.Primes[1])
-		stime := time.Now()
+
 		if v.Verify(x, site.Hardness, y) == true {
 			// 验证成功，什么都不用做
 			isCorrect = true
 		} else {
 			isCorrect = false
 		}
-		log.Printf("校验证明耗时 %v，模数长度为 %d", time.Now().Sub(stime), site.RSAKey.Primes[0].BitLen()*2)
+
+		log.Printf("校验证明耗时 %v，模数长度为 %d", verifyElapsed, site.RSAKey.Primes[0].BitLen()*2)
 	} else {
 		log.Printf("网站 %s 的 RSAKey 已经超时了，不会使用 RSAKey 进行检查", site.APIKey)
 	}
@@ -293,11 +298,13 @@ func webCaptchaVerify(c *gin.Context) {
 		}
 	}
 
+	verifyElapsed = time.Now().Sub(stime)
+
 	// 5. 若验证成功则记录一次 nonce
 	var msg string
 	if isCorrect {
 		nonceSet(hRaw)
-		msg = "Prove is correct"
+		msg = fmt.Sprintf("Prove is correct. Verification takes %v", verifyElapsed)
 	} else {
 		msg = "Prove is INVALID"
 	}
@@ -306,8 +313,9 @@ func webCaptchaVerify(c *gin.Context) {
 		"code":    0,
 		"message": msg,
 		"data": gin.H{
-			"prove":      req.Prove,
-			"is_correct": isCorrect,
+			"prove":          req.Prove,
+			"is_correct":     isCorrect,
+			"verify_time_ms": float64(verifyElapsed.Microseconds()) / 1000,
 		},
 	})
 }