go-websockets-talk/code/hub.go

package main

import (
	"strings"
	"time"

	"github.com/gorilla/websocket"
)

type message struct {
	Type    string      `json:"type"`
	Sender  string      `json:"sender"`
	Content interface{} `json:"content"`
	Date    time.Time   `json:"date"`
	Success bool        `json:"success"`
}

func newError(content string) message {
	return message{
		Type:    "error",
		Sender:  "",
		Content: content,
		Date:    time.Now().UTC(),
		Success: false,
	}
}

func newMessage(sender string, content string) message {
	return message{
		Type:    "message",
		Sender:  sender,
		Content: content,
		Date:    time.Now().UTC(),
		Success: true,
	}
}

func (m message) dispatch() {
	for client := range usernames {
		_ = client.WriteJSON(m)
	}
}

var usernames = make(map[*websocket.Conn]string)

func sendUserList(who *websocket.Conn) {
	list := []string{}
	for _, username := range usernames {
		list = append(list, username)
	}

	m := message{
		Type:    "users",
		Sender:  "",
		Content: list,
		Date:    time.Now().UTC(),
		Success: true,
	}

	_ = who.WriteJSON(m)
}

func sendChatMessage(sender *websocket.Conn, msg string) {
	m := newMessage(usernames[sender], msg)
	m.dispatch()
}

func handleDisconnection(sender *websocket.Conn) {
	m := newMessage("server", usernames[sender]+" has left the chat.")
	m.dispatch()
	delete(usernames, sender)
}

func handleIncomingMessage(sender *websocket.Conn, msg string) {
	if _, ok := usernames[sender]; !ok {
		username := strings.TrimSpace(msg)
		if username == "" || username == "server" {
			sender.WriteJSON(newError("You have an illegal nickname"))
			return
		}
		sendUserList(sender)

		usernames[sender] = username
		m := newMessage("server", username+" has joined the chat")
		m.dispatch()
		return
	}

	sendChatMessage(sender, msg)
}