WebSockets

A two-way connection made from one upgraded HTTP request, with exactly one goroutine reading it and one writing it.

Syntaxconn, err := upgrader.Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header) (*websocket.Conn, error)

A WebSocket is a single connection that both sides can write to at any moment. Neither side has to wait its turn, which is what makes it different from a request and a response.

It starts as an ordinary HTTP GET carrying an Upgrade: websocket header. The server answers with status 101 Switching Protocols, and from then on that same TCP connection carries frames instead of HTTP. A frame is one piece of data with a type: text, binary, ping, pong or close.

Upgrading in Go

The standard library has no WebSocket server, so this course uses the package github.com/gorilla/websocket, version 1.5.3. Add it with go get github.com/gorilla/websocket. An Upgrader value performs the handshake: upgrader.Upgrade(w, r, nil) returns a *websocket.Conn. Once it has succeeded, nothing may be written to w again. The connection is yours, taken out of the HTTP server's hands, so the handler may even return while other goroutines keep using it.

One reader, one writer

A Conn tolerates one goroutine reading and one goroutine writing, and no more. Each connection therefore gets a pair of goroutines: a read loop calling ReadMessage, and a write loop that takes messages off a channel and calls WriteMessage. The rest of the program sends on that channel and never touches the connection.

What trips people up

Upgrade writes its own HTTP error response when it fails, so do not reply a second time. The default origin check refuses a handshake whose Origin header names a different host, which reads like a server bug and is not one. A read with no size limit lets a client ask the server for any amount of memory, so call SetReadLimit.

Example

Example · go
package main

import (
	"fmt"
	"log"
	"net/http"
	"net/http/httptest"
	"strings"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{}

// echo upgrades the request, reads one message and sends it back in capitals.
func echo(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return // Upgrade has already written an HTTP error response
	}
	defer conn.Close()
	_, data, err := conn.ReadMessage()
	if err != nil {
		return
	}
	if err := conn.WriteMessage(websocket.TextMessage, []byte(strings.ToUpper(string(data)))); err != nil {
		log.Print(err)
	}
}

func main() {
	srv := httptest.NewServer(http.HandlerFunc(echo))
	defer srv.Close()
	// The address keeps its host and port but changes scheme: http -> ws.
	conn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(srv.URL, "http"), nil)
	if err != nil {
		log.Fatal(err)
	}
	defer conn.Close()
	if err := conn.WriteMessage(websocket.TextMessage, []byte("hello")); err != nil {
		log.Fatal(err)
	}
	_, reply, err := conn.ReadMessage()
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("server said: %s\n", reply)
}

// Output:
// server said: HELLO

When to use it

  • A chat room where the server and the browser both send at any time.
  • A collaborative editor streaming edits in both directions on one connection.
  • A trading screen that sends orders and receives fills without reconnecting.

More examples

The per-connection goroutine pair

Every WebSocket connection in a real service looks like this. readPump is the only goroutine that reads, writePump is the only one that writes, and writePump takes its work from a buffered channel so any other part of the program can queue a message without touching the connection. Closing the connection in either pump makes the other one fail on its next call, which is how the pair ends together.

Example · go
package chat

import (
	"log"
	"time"

	"github.com/gorilla/websocket"
)

// A Conn tolerates one reading goroutine and one writing goroutine, and no
// more. Each connection therefore gets exactly this pair.
type client struct {
	ws   *websocket.Conn
	send chan []byte // the writer goroutine's inbox
}

// readPump is the only goroutine that calls ReadMessage on this connection.
func (c *client) readPump() {
	defer c.ws.Close() // makes the writer's next write fail, so it stops too
	c.ws.SetReadLimit(4096)
	for {
		_, data, err := c.ws.ReadMessage()
		if err != nil {
			return
		}
		log.Printf("received %d bytes", len(data))
	}
}

// writePump is the only goroutine that calls WriteMessage on this connection.
func (c *client) writePump() {
	defer c.ws.Close()
	for data := range c.send {
		if err := c.ws.SetWriteDeadline(time.Now().Add(10 * time.Second)); err != nil {
			return
		}
		if err := c.ws.WriteMessage(websocket.TextMessage, data); err != nil {
			return
		}
	}
}

// Start launches the pair. Everything else in the program talks to c.send.
func (c *client) Start() {
	go c.writePump()
	go c.readPump()
}

Sending JSON: longhand and shorthand

WriteJSON is not a different mechanism. It marshals the value to bytes and sends them as one text message, which is exactly what the longhand does by hand. ReadJSON is the matching shorthand on the read side. Use the longhand when you need the bytes for something else, such as sending the same payload to many connections, because marshalling once and reusing the bytes is cheaper than WriteJSON per client.

Example · go
package chat

import (
	"encoding/json"

	"github.com/gorilla/websocket"
)

type event struct {
	Kind  string `json:"kind"`
	Price int    `json:"price"`
}

// Longhand: turn the value into bytes, then send one text message.
func sendLonghand(conn *websocket.Conn, ev event) error {
	payload, err := json.Marshal(ev)
	if err != nil {
		return err
	}
	return conn.WriteMessage(websocket.TextMessage, payload)
}

// Shorthand: WriteJSON does exactly those two steps.
func sendShorthand(conn *websocket.Conn, ev event) error {
	return conn.WriteJSON(ev)
}

// ReadJSON is the matching shorthand for the read side.
func receive(conn *websocket.Conn) (event, error) {
	var ev event
	err := conn.ReadJSON(&ev)
	return ev, err
}

Origin checking and the route

The handshake is a GET from a browser, and browsers attach cookies to it. Without a CheckOrigin function that names the sites you serve, any other website can open a socket to your server as a logged-in visitor. This is also where the read limit belongs, since a connection with none can be asked to buffer a gigabyte. The GET /ws pattern is a method pattern, available in ServeMux since Go 1.22.

Example · go
package chat

import (
	"log"
	"net/http"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	// Without this check any website can open a socket to your server using the
	// visitor's cookies. Compare the Origin header with the sites you serve.
	CheckOrigin: func(r *http.Request) bool {
		return r.Header.Get("Origin") == "https://app.example.com"
	},
}

func serveWS(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("upgrade failed: %v", err) // Upgrade has already replied
		return
	}
	defer conn.Close()
	conn.SetReadLimit(4096) // refuse messages larger than 4 KiB
}

func Routes() *http.ServeMux {
	mux := http.NewServeMux()
	// "GET /ws" is a method pattern, available since Go 1.22. A WebSocket
	// handshake is always a GET.
	mux.HandleFunc("GET /ws", serveWS)
	return mux
}

Related

go ssego realtime lifecyclego realtime broadcastgo goroutine basicsgo channel basics

Discussion

  • Be the first to comment on this lesson.