feat: panic when token not right

Ehco1996 · Mar 30, 2021 · 2be264a · 2be264a
1 parent 4a839bb
commit 2be264a
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 11 deletions.
diff --git a/cmd/ehco/main.go b/cmd/ehco/main.go
@@ -95,9 +95,8 @@ func main() {
 		},
 		&cli.StringFlag{
 			Name:        "web_token",
-			Usage:       "访问web的token",
+			Usage:       "访问web的token,如果访问不带着正确的token，会直接reset连接",
 			EnvVars:     []string{"EHCO_WEB_TOKEN"},
-			Value:       "randomtoken",
 			Destination: &WebToken,
 		},
 	}

diff --git a/internal/constant/constant.go b/internal/constant/constant.go
@@ -3,7 +3,7 @@ package constant
 import "time"
 
 const (
-	Version = "1.0.3"
+	Version = "1.0.4"
 
 	MaxMWSSStreamCnt = 10
 	DialTimeOut      = 3 * time.Second

diff --git a/internal/web/server.go b/internal/web/server.go
@@ -20,6 +20,10 @@ func Index(w http.ResponseWriter, r *http.Request) {
 	fmt.Fprintf(w, "access from %s \n", r.RemoteAddr)
 }
 
+func Welcome(w http.ResponseWriter, r *http.Request) {
+	fmt.Fprintf(w, constant.IndexHTMLTMPL)
+}
+
 func AttachProfiler(router *mux.Router) {
 	router.HandleFunc("/debug/pprof/", pprof.Index)
 	router.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
@@ -53,19 +57,27 @@ func StartWebServer(port, token string, cfg *config.Config) {
 	r := mux.NewRouter()
 	AttachProfiler(r)
 	registerMetrics(cfg)
-	r.Handle("/metrics/", SimpleTokenAuthMiddleware(token, promhttp.Handler()))
-	r.HandleFunc("/",
-		func(w http.ResponseWriter, r *http.Request) {
-			fmt.Fprintf(w, constant.IndexHTMLTMPL)
-		})
-	logger.Fatal(http.ListenAndServe(addr, r))
+	r.Handle("/", http.HandlerFunc(Welcome))
+	r.Handle("/metrics/", promhttp.Handler())
+	if token != "" {
+		logger.Fatal(http.ListenAndServe(addr, SimpleTokenAuthMiddleware(token, r)))
+	} else {
+		logger.Fatal(http.ListenAndServe(addr, r))
+	}
 }
 
 func SimpleTokenAuthMiddleware(token string, h http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if t := r.URL.Query().Get("token"); t != token {
-			logger.Logger.Errorf("[web] unauthorsied from %s", r.RemoteAddr)
-			w.WriteHeader(http.StatusUnauthorized)
+			msg := fmt.Sprintf("[web] unauthorsied from %s", r.RemoteAddr)
+			logger.Logger.Error(msg)
+			hj, ok := w.(http.Hijacker)
+			if ok {
+				conn, _, _ := hj.Hijack()
+				conn.Close()
+			} else {
+				panic(msg)
+			}
 			return
 		}
 		h.ServeHTTP(w, r)