From 50d2e082d55d8205b1eced3e83976a7d6f541dbc Mon Sep 17 00:00:00 2001
From: Dreamacro <305009791@qq.com>
Date: Mon, 30 Sep 2019 14:13:29 +0800
Subject: [PATCH] Feature: websocket api support browser

---
 hub/route/server.go | 24 ++++++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/hub/route/server.go b/hub/route/server.go
index 7736a3c9..394ad406 100644
--- a/hub/route/server.go
+++ b/hub/route/server.go
@@ -23,7 +23,11 @@ var (
 
 	uiPath = ""
 
-	upgrader = websocket.Upgrader{}
+	upgrader = websocket.Upgrader{
+		CheckOrigin: func(r *http.Request) bool {
+			return true
+		},
+	}
 )
 
 type Traffic struct {
@@ -84,14 +88,26 @@ func Start(addr string, secret string) {
 
 func authentication(next http.Handler) http.Handler {
 	fn := func(w http.ResponseWriter, r *http.Request) {
-		header := r.Header.Get("Authorization")
-		text := strings.SplitN(header, " ", 2)
-
 		if serverSecret == "" {
 			next.ServeHTTP(w, r)
 			return
 		}
 
+		// Browser websocket not support custom header
+		if websocket.IsWebSocketUpgrade(r) && r.URL.Query().Get("token") != "" {
+			token := r.URL.Query().Get("token")
+			if token != serverSecret {
+				render.Status(r, http.StatusUnauthorized)
+				render.JSON(w, r, ErrUnauthorized)
+				return
+			}
+			next.ServeHTTP(w, r)
+			return
+		}
+
+		header := r.Header.Get("Authorization")
+		text := strings.SplitN(header, " ", 2)
+
 		hasUnvalidHeader := text[0] != "Bearer"
 		hasUnvalidSecret := len(text) == 2 && text[1] != serverSecret
 		if hasUnvalidHeader || hasUnvalidSecret {