Simple server on Gorilla WebSocket

In this small tutorial, we'll take a closer look at using Gorilla WebSocket to write our own websocket server, at a more functional level than the basic example and easier to understand than the chat example.

What will our server be able to do?

1. Send new messages from clients to callback
2. Keep active connections and close / delete inactive
3. Send messages to active connections

First, we use the standard http help server net/http so that we can catch connection requests:

package main

import (
    "fmt"
    "html"
    "net/http"
)

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

func echo(w http.ResponseWriter, r *http.Request) {
    fmt.Fprintf(w, "Hello, %q", html.EscapeString(r.URL.Path))
}

Now let's teach it to "upgrade" the connection:

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

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

func echo(w http.ResponseWriter, r *http.Request) {
    connection, _ := upgrader.Upgrade(w, r, nil)
    connection.Close() // Close connection
}

We now have a client connection which we close immediately. We can loop through the messages that the client sends to us and send them back:

func echo(w http.ResponseWriter, r *http.Request) {
    connection, _ := upgrader.Upgrade(w, r, nil)

    for {
        _, message, _ := connection.ReadMessage()

        connection.WriteMessage(websocket.TextMessage, message)
        go messageHandler(message)
    }
}

func messageHandler(message []byte)  {
  fmt.Println(string(message))
}

Let's teach our server to close the connection:

func echo(w http.ResponseWriter, r *http.Request) {
    connection, _ := upgrader.Upgrade(w, r, nil)

    for {
        mt, message, err := connection.ReadMessage()

        if err != nil || mt == websocket.CloseMessage {
            break // Exit the loop if the client tries to close the connection or the connection with the interrupted client
        }

        connection.WriteMessage(websocket.TextMessage, message)

        go messageHandler(message)
    }

    connection.Close()
}

To be able to send messages over different connections, we need to store them somewhere, in our case the simplest map is suitable:

var clients map[*websocket.Conn]bool

func echo(w http.ResponseWriter, r *http.Request) {
    connection, _ := upgrader.Upgrade(w, r, nil)

    clients[connection] = true // Save the connection using it as a key

    for {
        mt, message, err := connection.ReadMessage()

        if err != nil || mt == websocket.CloseMessage {
            break // Exit the loop if the client tries to close the connection or the connection with the interrupted client
        }

        // Send messages to all clients
        go writeMessage(message)

        go messageHandler(message)
    }

    delete(clients, connection) // Removing the connection

    connection.Close()
}

func writeMessage(message []byte) {
    for conn := range clients {
        conn.WriteMessage(websocket.TextMessage, message)
    }
}

Now we can pack our server into a structure to be able to send and receive messages from outside:

package ws

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

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

type Server struct {
    clients       map[*websocket.Conn]bool
    handleMessage func(message []byte) // New message handler
}

func StartServer(handleMessage func(message []byte)) *Server {
    server := Server{
        make(map[*websocket.Conn]bool),
        handleMessage,
    }

    http.HandleFunc("/", server.echo)
    go http.ListenAndServe(":8080", nil)

    return &server
}

func (server *Server) echo(w http.ResponseWriter, r *http.Request) {
    connection, _ := upgrader.Upgrade(w, r, nil)

    server.clients[connection] = true // Save the connection using it as a key

    for {
        mt, message, err := connection.ReadMessage()

        if err != nil || mt == websocket.CloseMessage {
            break // Exit the loop if the client tries to close the connection or the connection is interrupted
        }

        go server.handleMessage(message)
    }

    delete(server.clients, connection) // Removing the connection

    connection.Close()
}

func (server *Server) WriteMessage(message []byte) {
    for conn := range server.clients {
        conn.WriteMessage(websocket.TextMessage, message)
    }
}

package main

import (
    "fmt"
    "simple-webcoket/ws"
)

func main() {
    server := ws.StartServer(messageHandler)

    for {
        server.WriteMessage([]byte("Hello"))
    }
}

func messageHandler(message []byte) {
    fmt.Println(string(message))
}

Now we have an implementation of a simple webscoket server that is capable of receiving and sending messages over active connections.

Comments

[Esteban]

Hello,

I don't understand why do you have both the:

for {
server.WriteMessage([]byte("Hello"))
}

And:

func messageHandler(message []byte) {
fmt.Println(string(message))
}

Functions. I understand the need for the messageHandler() function because that will be the function that gets executed inside the echo() function when a new websocket connection is created, but the server.WriteMessage() on a for loop I don't understand the point behind it?

[Wassaf Shahzasd]

I thinks that's just for local debugging to print the received message on the std out