Prerequisites

Before you begin, ensure you have Go installed on your machine. You can download it from the official Go website. Familiarity with Go's syntax and basic concepts is also recommended.

Project Setup

Create a new directory for your project and initialize a Go module.

mkdir websocket-chat-server
cd websocket-chat-server
go mod init websocket-chat-server

Import Required Packages

You will need the gorilla/websocket package to manage WebSocket connections. Install it using the following command:

go get github.com/gorilla/websocket

Code Structure

Create a file named main.go and add the following code to establish a basic WebSocket server.

package main

import (
    "fmt"
    "net/http"
    "github.com/gorilla/websocket"
    "sync"
)

var upgrader = websocket.Upgrader{
    CheckOrigin: func(r *http.Request) bool {
        return true
    },
}

type Client struct {
    conn *websocket.Conn
    send chan []byte
}

type ChatServer struct {
    clients   map[*Client]bool
    broadcast chan []byte
    mu        sync.Mutex
}

func NewChatServer() *ChatServer {
    return &ChatServer{
        clients:   make(map[*Client]bool),
        broadcast: make(chan []byte),
    }
}

func (cs *ChatServer) run() {
    for {
        msg := <-cs.broadcast
        cs.mu.Lock()
        for client := range cs.clients {
            select {
            case client.send <- msg:
            default:
                close(client.send)
                delete(cs.clients, client)
            }
        }
        cs.mu.Unlock()
    }
}

func (c *Client) read(cs *ChatServer) {
    defer func() {
        c.conn.Close()
    }()
    for {
        _, msg, err := c.conn.ReadMessage()
        if err != nil {
            break
        }
        cs.broadcast <- msg
    }
}

func (c *Client) write() {
    defer func() {
        c.conn.Close()
    }()
    for msg := range c.send {
        if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
            break
        }
    }
}

func handleConnection(w http.ResponseWriter, r *http.Request, cs *ChatServer) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        fmt.Println("Error during connection upgrade:", err)
        return
    }
    client := &Client{conn: conn, send: make(chan []byte)}
    cs.mu.Lock()
    cs.clients[client] = true
    cs.mu.Unlock()

    go client.read(cs)
    go client.write()
}

func main() {
    chatServer := NewChatServer()
    go chatServer.run()

    http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
        handleConnection(w, r, chatServer)
    })

    fmt.Println("Chat server started on :8080")
    if err := http.ListenAndServe(":8080", nil); err != nil {
        fmt.Println("Error starting server:", err)
    }
}

Code Explanation

  1. Imports: The necessary packages are imported, including the gorilla/websocket for WebSocket handling.
  1. Client and ChatServer Structs:
  • Client represents a connected client with a WebSocket connection and a channel to send messages.
  • ChatServer manages connected clients and broadcasts messages to them.
  1. NewChatServer Function: Initializes a new ChatServer instance.
  1. run Method: Listens for incoming messages and broadcasts them to all connected clients.
  1. read Method: Reads messages from the client and sends them to the broadcast channel.
  1. write Method: Sends messages from the broadcast channel to the client.
  1. handleConnection Function: Upgrades HTTP connections to WebSocket connections and manages client sessions.
  1. main Function: Initializes the chat server and starts the HTTP server.

Running the Server

To run the server, execute the following command in your terminal:

go run main.go

The server will start on localhost:8080. You can connect to it using a WebSocket client.

Testing the WebSocket Server

You can test the WebSocket server using a simple HTML client. Create an index.html file in the same directory:

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>WebSocket Chat</title>
</head>
<body>
    <input id="message" type="text" placeholder="Enter message..." />
    <button id="send">Send</button>
    <ul id="messages"></ul>

    <script>
        const ws = new WebSocket("ws://localhost:8080/ws");
        const messages = document.getElementById("messages");
        const messageInput = document.getElementById("message");
        const sendButton = document.getElementById("send");

        ws.onmessage = function(event) {
            const li = document.createElement("li");
            li.textContent = event.data;
            messages.appendChild(li);
        };

        sendButton.onclick = function() {
            const message = messageInput.value;
            ws.send(message);
            messageInput.value = "";
        };
    </script>
</body>
</html>

Open this HTML file in your browser, and you can start sending messages to the chat server. Open multiple tabs to see real-time message broadcasting.

Conclusion

In this tutorial, you have learned how to implement a simple WebSocket chat server in Go. You have seen how to manage client connections, read and write messages, and broadcast them to all connected clients. This foundational knowledge can be expanded to create more complex applications, such as collaborative tools or real-time notifications.

Learn more with useful resources