livestream/server/websocket.go

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"strings"

	"github.com/olahol/melody"
)

type WebsocketMessageType = int

const (
	WebsocketMessageTypeCurrentlyPlaying WebsocketMessageType = iota
	WebsocketMessageTypeChatMessage
	WebsocketMessageTypeUserJoin
	WebsocketMessageTypeUserLeave
	WebsocketMessageTypeUserRename
)

const reserved_username = "Bot"
const reserved_color = "ff0000"

type WebsocketMessage struct {
	Type    WebsocketMessageType `json:"type"`
	Message any                  `json:"message"`
}

func (w WebsocketMessage) Encode() []byte {
	m, err := json.Marshal(w)
	if err != nil {
		log.Fatal(err)
	}

	return m
}

type WebsocketSendMessage struct {
	Username      string `json:"username"`
	UsernameColor string `json:"username_color"`
	Content       string `json:"content"`
}

type WebsocketUser struct {
	Username string `json:"username"`
	Color    string `json:"color"`
}

type WebsocketRename struct {
	OldUsername string `json:"old_username"`
	NewUsername string `json:"new_username"`
	NewColor    string `json:"new_color"`
}

func (state *State) UsernameUsed(username string) (used bool) {
	if username == "" || strings.ToLower(username) == reserved_username {
		return true
	}

	for _, user := range state.Users {
		if username == user.Username {
			return true
		}
	}

	return
}

func TemplateUsername(num int) string {
	return fmt.Sprintf("Guest%d", num)
}

func (state *State) setupWebsocket() {
	state.Websocket = melody.New()

	state.Websocket.HandleConnect(func(s *melody.Session) {
		i := 1

		potentionalUsername := TemplateUsername(i)
		for {
			if !state.UsernameUsed(potentionalUsername) {
				break
			}

			i++
			potentionalUsername = TemplateUsername(i)
		}

		state.Users[s.RemoteAddr()] = WebsocketUserInfo{
			Username: potentionalUsername,
			Color:    ToColor(potentionalUsername),
		}

		// Adds bot user
		s.Write(WebsocketMessage{
			Type: WebsocketMessageTypeUserJoin,
			Message: WebsocketUser{
				Username: reserved_username,
				Color:    reserved_color,
			},
		}.Encode())

		// Adds other viewers
		for _, user := range state.Users {
			s.Write(WebsocketMessage{
				Type:    WebsocketMessageTypeUserJoin,
				Message: user.ToWebsocketUser(),
			}.Encode())
		}

		// Show newly joined viewer to other viewers
		state.Websocket.BroadcastOthers(WebsocketMessage{
			Type:    WebsocketMessageTypeUserJoin,
			Message: state.Users[s.RemoteAddr()].ToWebsocketUser(),
		}.Encode(), s)

		// Give the now playing status
		s.Write(WebsocketMessage{
			Type:    WebsocketMessageTypeCurrentlyPlaying,
			Message: state.NowPlaying(),
		}.Encode())
	})

	state.Websocket.HandleMessage(func(s *melody.Session, b []byte) {
		var genericMessage WebsocketMessage
		if err := json.Unmarshal(b, &genericMessage); err != nil {
			log.Println(err)
			return
		}

		switch genericMessage.Type {
		case WebsocketMessageTypeChatMessage:
			content, valid := genericMessage.Message.(string)
			if !valid {
				return
			}

			content = strings.TrimSpace(content)

			if content == "" || len(content) > 200 {
				return
			}

			info := state.Users[s.RemoteAddr()]

			if strings.HasPrefix(content, "/") {
				split := strings.Split(content, " ")
				command := split[0]

				switch command {
				case "/help":
					s.Write(WebsocketMessage{
						Type: WebsocketMessageTypeChatMessage,
						Message: WebsocketSendMessage{
							Username:      reserved_username,
							UsernameColor: reserved_color,
							Content:       "Commands: /help, /username my_new_username",
						},
					}.Encode())
				case "/username":
					if split[1] == info.Username {
						s.Write(WebsocketMessage{
							Type: WebsocketMessageTypeChatMessage,
							Message: WebsocketSendMessage{
								Username:      reserved_username,
								UsernameColor: reserved_color,
								Content:       "You are already using that username",
							},
						}.Encode())
					} else if split[1] == "" {
						s.Write(WebsocketMessage{
							Type: WebsocketMessageTypeChatMessage,
							Message: WebsocketSendMessage{
								Username:      reserved_username,
								UsernameColor: reserved_color,
								Content:       "Please choose an username",
							},
						}.Encode())
					} else if len(split[1]) > 25 {
						s.Write(WebsocketMessage{
							Type: WebsocketMessageTypeChatMessage,
							Message: WebsocketSendMessage{
								Username:      reserved_username,
								UsernameColor: reserved_color,
								Content:       "Please choose a smaller username",
							},
						}.Encode())
					} else if !state.UsernameUsed(split[1]) {
						OldUsername := state.Users[s.RemoteAddr()].Username
						Color := ToColor(split[1])

						state.Users[s.RemoteAddr()] = WebsocketUserInfo{
							Username: split[1],
							Color:    Color,
						}

						// Shows rename to users
						state.Websocket.Broadcast(WebsocketMessage{
							Type: WebsocketMessageTypeUserRename,
							Message: WebsocketRename{
								OldUsername: OldUsername,
								NewUsername: state.Users[s.RemoteAddr()].Username,
								NewColor:    Color,
							},
						}.Encode())

						s.Write(WebsocketMessage{
							Type: WebsocketMessageTypeChatMessage,
							Message: WebsocketSendMessage{
								Username:      reserved_username,
								UsernameColor: reserved_color,
								Content:       fmt.Sprintf("Successfully changed username to '%s'", split[1]),
							},
						}.Encode())
					} else {
						s.Write(WebsocketMessage{
							Type: WebsocketMessageTypeChatMessage,
							Message: WebsocketSendMessage{
								Username:      reserved_username,
								UsernameColor: reserved_color,
								Content:       fmt.Sprintf("'%s' is already taken, please choose something else.", split[1]),
							},
						}.Encode())
					}
				default:
					s.Write(WebsocketMessage{
						Type: WebsocketMessageTypeChatMessage,
						Message: WebsocketSendMessage{
							Username:      reserved_username,
							UsernameColor: reserved_color,
							Content:       fmt.Sprintf("Unknown command '%s', maybe try /help", command),
						},
					}.Encode())
				}
			} else {
				state.Websocket.Broadcast(WebsocketMessage{
					Type: WebsocketMessageTypeChatMessage,
					Message: WebsocketSendMessage{
						Username:      info.Username,
						UsernameColor: info.Color,
						Content:       content,
					},
				}.Encode())
			}
		}
	})

	state.Websocket.HandleDisconnect(func(s *melody.Session) {
		// Remove viewer from list
		state.Websocket.BroadcastOthers(WebsocketMessage{
			Type:    WebsocketMessageTypeUserLeave,
			Message: state.Users[s.RemoteAddr()].ToWebsocketUser(),
		}.Encode(), s)

		delete(state.Users, s.RemoteAddr())
	})
}