feat: token auth

coollabsio · May 8, 2024 · 7fa4de2 · 7fa4de2
1 parent 75cd1a9
commit 7fa4de2
Showing 1 changed file with 71 additions and 44 deletions.
diff --git a/cmd/sentinel/main.go b/cmd/sentinel/main.go
@@ -2,68 +2,95 @@ package main
 
 import (
 	"encoding/json"
+	"flag"
 
 	"github.com/gin-gonic/gin"
 )
 
+var token string
+
+func Token() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		if token != "" {
+			if c.GetHeader("Authorization") != "Bearer "+token {
+				c.JSON(401, gin.H{
+					"error": "Unauthorized",
+				})
+				c.Abort()
+				return
+			}
+		}
+		c.Next()
+	}
+}
 func main() {
 	// scheduler()
+	flag.StringVar(&token, "token", "", "help message for flagname")
+	flag.Parse()
+
 	r := gin.Default()
 	r.GET("/api/health", func(c *gin.Context) {
 		c.String(200, "ok")
 	})
-	r.GET("/api/containers", func(c *gin.Context) {
-		containers, err := getAllContainers()
-		if err != nil {
-			c.JSON(500, gin.H{
-				"error": err.Error(),
-			})
-			return
-		}
+	r.Use(gin.Recovery())
 
-		c.JSON(200, gin.H{
-			"containers": json.RawMessage(containers),
-		})
-	})
-	r.GET("/api/cpu", func(c *gin.Context) {
-		usage, err := getCpuUsage()
-		if err != nil {
-			c.JSON(500, gin.H{
-				"error": err.Error(),
-			})
-			return
-		}
+	authorized := r.Group("/api")
+	authorized.Use(Token())
+	{
+		authorized.GET("/containers", func(c *gin.Context) {
+			containers, err := getAllContainers()
+			if err != nil {
+				c.JSON(500, gin.H{
+					"error": err.Error(),
+				})
+				return
+			}
 
-		c.JSON(200, gin.H{
-			"cpu_usage": json.RawMessage(usage),
-		})
-	})
-	r.GET("/api/memory", func(c *gin.Context) {
-		usage, err := getMemUsage()
-		if err != nil {
-			c.JSON(500, gin.H{
-				"error": err.Error(),
+			c.JSON(200, gin.H{
+				"containers": json.RawMessage(containers),
 			})
-			return
-		}
+		})
+		authorized.GET("/cpu", func(c *gin.Context) {
+			usage, err := getCpuUsage()
+			if err != nil {
+				c.JSON(500, gin.H{
+					"error": err.Error(),
+				})
+				return
+			}
 
-		c.JSON(200, gin.H{
-			"mem_usage": json.RawMessage(usage),
+			c.JSON(200, gin.H{
+				"cpu_usage": json.RawMessage(usage),
+			})
 		})
-	})
-	r.GET("/api/disk", func(c *gin.Context) {
-		usage, err := getDiskUsage()
-		if err != nil {
-			c.JSON(500, gin.H{
-				"error": err.Error(),
+		authorized.GET("/memory", func(c *gin.Context) {
+			usage, err := getMemUsage()
+			if err != nil {
+				c.JSON(500, gin.H{
+					"error": err.Error(),
+				})
+				return
+			}
+
+			c.JSON(200, gin.H{
+				"mem_usage": json.RawMessage(usage),
 			})
-			return
-		}
+		})
+		authorized.GET("/disk", func(c *gin.Context) {
+			usage, err := getDiskUsage()
+			if err != nil {
+				c.JSON(500, gin.H{
+					"error": err.Error(),
+				})
+				return
+			}
 
-		c.JSON(200, gin.H{
-			"disk_usage": json.RawMessage(usage),
+			c.JSON(200, gin.H{
+				"disk_usage": json.RawMessage(usage),
+			})
 		})
-	})
+	}
+
 	r.Run("0.0.0.0:8888")
 
 }