circolog/http.go

package main

import (
	"net/http"
	"time"

	"github.com/gorilla/websocket"
	"gopkg.in/mcuadros/go-syslog.v2/format"
)

func setupHttp(hub Hub) {
	http.HandleFunc("/", getHTTPHandler(hub))
	http.HandleFunc("/ws", getWSHandler(hub))
}

func getHTTPHandler(hub Hub) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		client := Client{
			Messages: make(chan format.LogParts),
			Nofollow: true}
		hub.Register <- client

		for x := range client.Messages {
			logmsg := x
			if logmsg["message"] == nil {
				continue
			}
			c := logmsg["message"].(string)
			w.Write([]byte(c))
			w.Write([]byte("\n"))
		}
	}
}

func getWSHandler(hub Hub) http.HandlerFunc {
	var upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
	return func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		client := Client{Messages: make(chan format.LogParts)}
		hub.Register <- client

		// Allow collection of memory referenced by the caller by doing all work in
		// new goroutines.
		go func(conn *websocket.Conn, c Client) {
			defer func() {
				hub.Unregister <- c
				conn.Close()
			}()
			for {
				select {
				case message, ok := <-c.Messages:
					conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
					if !ok {
						// The hub closed the channel.
						conn.WriteMessage(websocket.CloseMessage, []byte{})
						return
					}

					w, err := conn.NextWriter(websocket.TextMessage)
					if err != nil {
						return
					}
					if msg, ok := message["message"]; ok {
						w.Write([]byte(msg.(string)))
					}

					if err := w.Close(); err != nil {
						return
					}
					// TODO: ticker/ping
				}
			}
		}(conn, client)
	}
}