update

server/websocket.go

@@ -9,45 +9,43 @@ import (
 	jsoniter "github.com/json-iterator/go"
 )
 
-type enterChanel struct {
-	Conn   *websocket.Conn
+type Client struct {
 	Chanel string
+	Conn   *websocket.Conn
 }
 
 type sendMsg struct {
-	Conn   *websocket.Conn
 	Message []byte
-}
-
-var upgrader = websocket.Upgrader{
-	ReadBufferSize:  1024,
-	WriteBufferSize: 1024,
-	CheckOrigin: func(r *http.Request) bool {
-		//if r.Header.Get("origin") == "https://statbate.com" {
-		//	return true
-		//}
-		//return false
-		return true
-	},
+	Conn    *websocket.Conn
 }
 
 var (
-	wsClients = make(map[*websocket.Conn]string)
+	wsClients = make(map[*Client]bool)
 
 	json = jsoniter.ConfigCompatibleWithStandardLibrary
 
 	ws = struct {
 		Broadcast chan []byte
 		Send      chan sendMsg
-		Enter chan enterChanel
-		Add   chan *websocket.Conn
-		Del   chan *websocket.Conn
+		Add       chan *Client
+		Del       chan *Client
 	}{
 		Broadcast: make(chan []byte, 100),
 		Send:      make(chan sendMsg, 100),
-		Enter: make(chan enterChanel, 100),
-		Add:   make(chan *websocket.Conn, 100),
-		Del:   make(chan *websocket.Conn, 100),
+		Add:       make(chan *Client, 100),
+		Del:       make(chan *Client, 100),
+	}
+
+	upgrader = websocket.Upgrader{
+		ReadBufferSize:  1024,
+		WriteBufferSize: 1024,
+		CheckOrigin: func(r *http.Request) bool {
+			params := r.URL.Query()
+			if params["api_key"] != nil || r.Header.Get("origin") == "https://statbate.com" {
+				return true
+			}
+			return false
+		},
 	}
 )
 
@@ -55,16 +53,11 @@ func broadcast() {
 	ticker := time.NewTicker(30 * time.Second)
 	for {
 		select {
-		case conn := <-ws.Add:
-			wsClients[conn] = ""
-			//fmt.Println("add conn", len(wsClients))
+		case client := <-ws.Add:
+			wsClients[client] = true
 
 		case conn := <-ws.Del:
 			delete(wsClients, conn)
-			//fmt.Println("delete conn", len(wsClients))
-
-		case r := <-ws.Enter:
-			wsClients[r.Conn] = r.Chanel
 
 		case r := <-ws.Send:
 			sendMessage(r.Conn, r.Message)
@@ -92,14 +85,29 @@ func sendBroadcast(message []byte) {
 		fmt.Println("json error: ", err.Error())
 		return
 	}
-	for conn, ch := range wsClients {
-		if ch != input.Chanel {
+	for client := range wsClients {
+		if client.Chanel != input.Chanel {
 			continue
 		}
-		sendMessage(conn, message)
+		sendMessage(client.Conn, message)
 	}
 }
 
+func enterChannel(conn *websocket.Conn, chanel string) (*Client, bool) {
+	chanels := map[string]bool{
+		"chaturbate": true,
+		"bongacams":  true,
+		"stripchat":  true,
+		"camsoda":    true,
+	}
+	client := &Client{Conn: conn, Chanel: chanel}
+	if chanels[chanel] {
+		ws.Add <- client
+		return client, true
+	}
+	return client, false
+}
+
 func wsHandler(w http.ResponseWriter, r *http.Request) {
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
@@ -111,23 +119,34 @@ func wsHandler(w http.ResponseWriter, r *http.Request) {
 func readWS(conn *websocket.Conn) {
 	defer conn.Close()
 
-	ws.Add <- conn
-
-	defer func() {
-		ws.Del <- conn
-	}()
-
-	chanels := map[string]bool{
-		"chaturbate": true,
-		"bongacams":  true,
-		"stripchat":  true,
-		"camsoda":    true,
-	}
-	ping := time.Now().Unix()
-	for {
+	conn.SetReadDeadline(time.Now().Add(30 * time.Second))
+	_, message, err := conn.ReadMessage()
+	if err != nil {
+		return
+	}
+
+	input := struct {
+		Chanel string `json:"chanel"`
+	}{}
+
+	if err := json.Unmarshal(message, &input); err != nil {
+		return
+	}
+
+	client, ok := enterChannel(conn, input.Chanel)
+	if !ok {
+		return
+	}
+
+	defer func() {
+		ws.Del <- client
+	}()
+
+	ping := time.Now().Unix()
+	for {
+		conn.SetReadDeadline(time.Now().Add(30 * time.Minute))
 		_, message, err := conn.ReadMessage()
 		if err != nil {
-			//fmt.Println("readWS", err.Error())
 			return
 		}
 
@@ -138,18 +157,5 @@ func readWS(conn *websocket.Conn) {
 			}
 			continue
 		}
-
-		input := struct {
-			Chanel string `json:"chanel"`
-		}{}
-
-		if err := json.Unmarshal(message, &input); err != nil {
-			//fmt.Println("wrong json", err.Error())
-			return
-		}
-
-		if chanels[input.Chanel] {
-			ws.Enter <- enterChanel{Conn: conn, Chanel: input.Chanel}
-		}
 	}
 }