WebSocket server

websocket.go

package main

import (
	"fmt"
	"os"
	"http"
	"websocket"
)

type subscription struct {
	conn      *websocket.Conn
	subscribe bool
}

var subscriptionChan = make(chan subscription)
var messageChan = make(chan []byte)

func clientHandler(ws *websocket.Conn) {
	defer func() {
		subscriptionChan <- subscription{ws, false}
		if err := ws.Close(); err != nil {
			fmt.Printf("clientHandler.close: %v\n", err)
		}
	}()

	subscriptionChan <- subscription{ws, true}

	for {
		buf := make([]byte, 512)
		n, err := ws.Read(buf)
		if err != nil {
			if err == os.EOF {
				break
			} else {
				fmt.Printf("clientHandler.for: %v\n", err)
				break
			}
		}
		messageChan <- buf[0:n]
	}
}

func hub() {
	conns := make(map[*websocket.Conn]int)
	for {
		select {
		case subscription := <-subscriptionChan:
			if subscription.subscribe {
				conns[subscription.conn] = 0
			} else {
				conns[subscription.conn] = 0, false
			}
		case message := <-messageChan:
			for conn, _ := range conns {
				if _, err := conn.Write(message); err != nil {
					conn.Close()
				}
			}
		}
	}
}

func main() {
	go hub()
	http.Handle("/", websocket.Handler(clientHandler))
	if err := http.ListenAndServe(":9999", nil); err != nil {
		panic(err)
	}
}