supreme-disco/server/websocket.go

package main

import (
	"fmt"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
	jsoniter "github.com/json-iterator/go"
)

type Client struct {
	Chanel string
	Conn   *websocket.Conn
}

type sendMsg struct {
	Message []byte
	Conn    *websocket.Conn
}

var (
	wsClients = make(map[*Client]bool)

	json = jsoniter.ConfigCompatibleWithStandardLibrary

	ws = struct {
		Broadcast chan []byte
		Send      chan sendMsg
		Add       chan *Client
		Del       chan *Client
	}{
		Broadcast: make(chan []byte, 100),
		Send:      make(chan sendMsg, 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
		},
	}
)

func broadcast() {
	ticker := time.NewTicker(30 * time.Second)
	for {
		select {
		case client := <-ws.Add:
			wsClients[client] = true

		case conn := <-ws.Del:
			delete(wsClients, conn)

		case r := <-ws.Send:
			sendMessage(r.Conn, r.Message)

		case r := <-ws.Broadcast:
			sendBroadcast(r)

		case <-ticker.C:
			fmt.Println("WebSocket:", len(wsClients))
		}
	}
}

func sendMessage(conn *websocket.Conn, message []byte) {
	if err := conn.WriteMessage(1, message); err != nil {
		conn.Close()
	}
}

func sendBroadcast(message []byte) {
	input := struct {
		Chanel string `json:"chanel"`
	}{}
	if err := json.Unmarshal(message, &input); err != nil {
		fmt.Println("json error: ", err.Error())
		return
	}
	for client := range wsClients {
		if client.Chanel != input.Chanel {
			continue
		}
		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 {
		return
	}
	go readWS(conn)
}

func readWS(conn *websocket.Conn) {
	defer conn.Close()

	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 {
			return
		}

		if string(message) == "ping" {
			if time.Now().Unix() > ping {
				ws.Send <- sendMsg{Conn: conn, Message: []byte("pong")}
				ping = time.Now().Unix() + 15
			}
			continue
		}
	}
}