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()) }) }