circolog/main.go

package main

import (
	"container/ring"
	"flag"
	"fmt"
	"net/http"
	"os"
	"time"

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

var circbuf *ring.Ring
var hub Hub
var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

func init() {
	hub = NewHub()
}

func syslogdHandler(channel syslog.LogPartsChannel) {
	for logParts := range channel {
		hub.logMessages <- logParts
		fmt.Println(logParts)
		circbuf.Value = logParts
		circbuf = circbuf.Next()
	}
}

func httpHandler(w http.ResponseWriter, r *http.Request) {
	circbuf.Do(func(x interface{}) {
		if x == nil {
			return
		}
		logmsg := x.(format.LogParts)
		if logmsg["message"] == nil {
			return
		}
		c := logmsg["message"].(string)
		w.Write([]byte(c))
		w.Write([]byte("\n"))
	})
}

func wsHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}
	httpHandler(w, r)
	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)
}

func main() {
	var err error
	syslogSocketPath := flag.String("syslogd-socket", "", "The socket to listen to syslog addresses")
	// dumpSocketPath := flag.String("dump-socket", "/run/buffer.sock", "The socket that user will connect to in order to receive logs")
	bufsize := flag.Int("buffer-size", 1000, "Number of messages to keep")
	syslogAddr := flag.String("syslog-addr", "127.0.0.1:9514", "Address:port where to listen for syslog messages")
	queryAddr := flag.String("query-addr", "127.0.0.1:9080", "Address:port where to bind the query service")
	flag.Parse()

	channel := make(chan format.LogParts)
	handler := syslog.NewChannelHandler(channel)

	server := syslog.NewServer()
	server.SetFormat(syslog.RFC5424)
	server.SetHandler(handler)
	if *syslogSocketPath != "" {
		if err = server.ListenUnixgram(*syslogSocketPath); err != nil {
			fmt.Fprintln(os.Stderr, "argh", err)
			os.Exit(1)
		}
		fmt.Printf("Binding socket `%s` [syslog]\n", *syslogSocketPath)
	} else {
		fmt.Printf("Binding address `%s` [syslog]\n", *syslogAddr)
		if err = server.ListenUDP(*syslogAddr); err != nil {
			fmt.Fprintln(os.Stderr, "argh", err)
			os.Exit(1)
		}
	}
	circbuf = ring.New(*bufsize)
	if err = server.Boot(); err != nil {
		fmt.Fprintln(os.Stderr, "argh", err)
		os.Exit(1)
	}
	go hub.Run()
	go syslogdHandler(channel)

	http.HandleFunc("/", httpHandler)
	http.HandleFunc("/ws", wsHandler)
	fmt.Printf("Binding address `%s` [http]\n", *queryAddr)
	http.ListenAndServe(*queryAddr, nil)

	server.Wait()
}