websockets.go

package main

import (
	"fmt"
	"net/http"
	"sync"

	"github.com/gorilla/websocket"
)

var (
	upgrader = websocket.Upgrader{}
	clients  = make(map[*websocket.Conn]bool) // connected clients
	mutex    = &sync.Mutex{}
)

func wsHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		fmt.Println(err)
		return
	}

	// Register the new client
	mutex.Lock()
	clients[conn] = true
	mutex.Unlock()

	for {
		messageType, msg, err := conn.ReadMessage()
		if err != nil {
			fmt.Println(err)
			mutex.Lock()
			delete(clients, conn)
			mutex.Unlock()
			return
		}
		fmt.Println(string(msg))

		// Echo back the message with an exclamation mark appended to all clients
		echoMsg := string(msg) + "!"
		mutex.Lock()
		for client := range clients {
			if err := client.WriteMessage(messageType, []byte(echoMsg)); err != nil {
				fmt.Println(err)
				delete(clients, client)
			}
		}
		mutex.Unlock()
	}
}

func helloHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprint(w, "Hello World")
}

func main() {
	http.HandleFunc("/", helloHandler)
	http.HandleFunc("/ws", wsHandler)
	http.ListenAndServe(":8080", nil)
}