From 3be2ff040111b3825a36ce96673a40094fbf37a8 Mon Sep 17 00:00:00 2001
From: liuwenhui <qdn163@163.com>
Date: Sat, 20 Jan 2024 23:33:58 +0800
Subject: [PATCH] auth

---
 cmd/go-selfupdate/main.go |  4 ++++
 cmd/server/main.go        | 34 +++++++++++++++++++++++++++++++++-
 selfupdate/requester.go   | 14 +++++++++++++-
 3 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/cmd/go-selfupdate/main.go b/cmd/go-selfupdate/main.go
index 2c9f226..400671b 100644
--- a/cmd/go-selfupdate/main.go
+++ b/cmd/go-selfupdate/main.go
@@ -14,6 +14,9 @@ import (
 	"path/filepath"
 	"runtime"
 )
+const (
+	DefaultToken = "ap_pJSFC5wQYkAyI0FIVwKYs9h1hW"
+)
 
 func uploadFile(url, filename, platform, version string) error {
 	var buf bytes.Buffer
@@ -53,6 +56,7 @@ func uploadFile(url, filename, platform, version string) error {
 		return err
 	}
 	req.Header.Set("Content-Type", writer.FormDataContentType())
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", DefaultToken))
 	client := http.DefaultClient
 	resp, err := client.Do(req)
 	if err != nil {
diff --git a/cmd/server/main.go b/cmd/server/main.go
index 170ea6f..887dd77 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -7,18 +7,23 @@ import (
 	"encoding/json"
 	"flag"
 	"fmt"
-	"github.com/kr/binarydist"
 	"github.com/bnulwh/go-selfupdate/selfupdate"
+	"github.com/kr/binarydist"
 	"io"
 	"io/ioutil"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/gin-gonic/gin"
 	"log"
 	"net/http"
 )
 
+const (
+	DefaultToken = "ap_pJSFC5wQYkAyI0FIVwKYs9h1hW"
+)
+
 type current struct {
 	Version string
 	Sha256  []byte
@@ -56,6 +61,7 @@ func main() {
 	flag.Parse()
 
 	router := gin.Default()
+	router.Use(PermissionChecker())
 	//pprof.Register(router, "/pprof")
 	router.StaticFS("/", http.Dir(*servePath))
 	router.POST("/upload", PostUpload)
@@ -63,6 +69,32 @@ func main() {
 
 }
 
+func PermissionChecker() gin.HandlerFunc {
+	return func(ctx *gin.Context) {
+		auth := ctx.Request.Header.Get("Authorization")
+		if auth == "" {
+			ctx.AbortWithError(http.StatusUnauthorized, fmt.Errorf("not auth"))
+			return
+		}
+		arr := strings.Split(auth, " ")
+		if len(arr) != 2 {
+			ctx.AbortWithError(http.StatusUnauthorized, fmt.Errorf("not auth"))
+			return
+		}
+		authType := arr[0]
+		if authType != "Bearer" {
+			ctx.AbortWithError(http.StatusUnauthorized, fmt.Errorf("not auth"))
+			return
+		}
+		authToken := arr[1]
+		if authToken != DefaultToken {
+			ctx.AbortWithError(http.StatusUnauthorized, fmt.Errorf("not auth"))
+			return
+		}
+		ctx.Next()
+	}
+}
+
 func PostUpload(ctx *gin.Context) {
 	file, _ := ctx.FormFile("file")
 	filePath := "/tmp/" + file.Filename
diff --git a/selfupdate/requester.go b/selfupdate/requester.go
index 4e5e6ac..7dfb466 100644
--- a/selfupdate/requester.go
+++ b/selfupdate/requester.go
@@ -4,6 +4,11 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"strings"
+)
+
+const (
+	DefaultToken = "ap_pJSFC5wQYkAyI0FIVwKYs9h1hW"
 )
 
 // Requester interface allows developers to customize the method in which
@@ -19,7 +24,14 @@ type HTTPRequester struct{}
 // Fetch will return an HTTP request to the specified url and return
 // the body of the result. An error will occur for a non 200 status code.
 func (httpRequester *HTTPRequester) Fetch(url string) (io.ReadCloser, error) {
-	resp, err := http.Get(url)
+	body := ""
+	req, err := http.NewRequest("GET", url, strings.NewReader(body))
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", DefaultToken))
+	client := http.DefaultClient
+	resp, err := client.Do(req)
 	if err != nil {
 		return nil, err
 	}