From 15cca97b819756897e788d6d353c7d52e8f4bf1e Mon Sep 17 00:00:00 2001
From: Vincent Yang <missuo@pm.me>
Date: Sat, 13 Apr 2024 18:54:16 -0400
Subject: [PATCH] feat: add authmiddleware to all routes

---
 main.go | 36 +++++++++++++++++++++---------------
 1 file changed, 21 insertions(+), 15 deletions(-)

diff --git a/main.go b/main.go
index 07e1d5c2..ae43941b 100644
--- a/main.go
+++ b/main.go
@@ -2,7 +2,7 @@
  * @Author: Vincent Yang
  * @Date: 2023-07-01 21:45:34
  * @LastEditors: Vincent Yang
- * @LastEditTime: 2024-04-09 03:02:08
+ * @LastEditTime: 2024-04-13 18:50:25
  * @FilePath: /DeepLX/main.go
  * @Telegram: https://t.me/missuo
  * @GitHub: https://github.com/missuo
@@ -338,6 +338,24 @@ func translateByDeepLX(sourceLang string, targetLang string, translateText strin
 	}, nil
 }
 
+func authMiddleware(cfg *Config) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		if cfg.Token != "" {
+			providedTokenInQuery := c.Query("token")
+			providedTokenInHeader := c.GetHeader("Authorization")
+			if providedTokenInHeader != "Bearer "+cfg.Token && providedTokenInQuery != cfg.Token {
+				c.JSON(http.StatusUnauthorized, gin.H{
+					"code":    http.StatusUnauthorized,
+					"message": "Invalid access token",
+				})
+				c.Abort()
+				return
+			}
+		}
+		c.Next()
+	}
+}
+
 func main() {
 	cfg := initConfig()
 
@@ -365,22 +383,10 @@ func main() {
 	})
 
 	// Defining the translation endpoint which receives translation requests and returns translations
-	r.POST("/translate", func(c *gin.Context) {
+	r.POST("/translate", authMiddleware(cfg), func(c *gin.Context) {
 		req := PayloadFree{}
 		c.BindJSON(&req)
 
-		if cfg.Token != "" {
-			providedTokenInQuery := c.Query("token")
-			providedTokenInHeader := c.GetHeader("Authorization")
-			if providedTokenInHeader != "Bearer "+cfg.Token && providedTokenInQuery != cfg.Token {
-				c.JSON(http.StatusUnauthorized, gin.H{
-					"code":    http.StatusUnauthorized,
-					"message": "Invalid access token",
-				})
-				return
-			}
-		}
-
 		sourceLang := req.SourceLang
 		targetLang := req.TargetLang
 		translateText := req.TransText
@@ -410,7 +416,7 @@ func main() {
 		}
 	})
 
-	r.POST("/v2/translate", func(c *gin.Context) {
+	r.POST("/v2/translate", authMiddleware(cfg), func(c *gin.Context) {
 		authorizationHeader := c.GetHeader("Authorization")
 		parts := strings.Split(authorizationHeader, " ")
 		var authKey string