direttoforo/uiserver/server.go

package uiserver

import (
	"encoding/json"
	"fmt"
	"net"
	"os"
)

// Worker handles a single connection, providing updates to the clients
type Worker struct {
	Update chan interface{} // this is called whenever a worker should update
	Exit   chan interface{} // this is called to force the worker to Close
	Conn   net.Conn
}

// NewWorker creates a new Worker from an ongoing connection
func NewWorker(conn net.Conn) Worker {
	w := Worker{Conn: conn}
	w.Update = make(chan interface{}, 10)
	return w
}

func (w *Worker) send(obj interface{}) error {
	marshalled, err := json.Marshal(obj)
	if err != nil {
		return err
	}
	marshalled = append(marshalled, byte('\n'))
	_, err = w.Conn.Write(marshalled)
	return err
}

// Work loops on Worker channel to do what is needed
func (w *Worker) Work() {
	for {
		select {
		case <-w.Exit:
			w.Conn.Close()
			return
		case obj := <-w.Update:
			err := w.send(obj)
			if err != nil {
				w.Conn.Close()
				close(w.Update)
				close(w.Exit)
				return
			}
		}
	}
}

// NetUI handles a server for the protocol documented in this package
// the API is designed so to make it possible to attach it to multiple
// listeners, despite this is NOT possible at the moment
type NetUI struct {
	Object  interface{}
	workers []Worker
	sock    *net.Listener
	exit    chan interface{}
}

// protocolHello is the hello message that the server will give to clients
// a server doesn't expect any introduction from clients
// the hello message is useful to clients to recognize protocol version cleanly
// the clients are guaranteed that the hello message will always be present in
// future versions of the protocol, always terminated by a newline and of
// exactly 15bytes
var protocolHello = "DIRETTOFORO V1\n"

// NewNetUI creates a NetUI object. Don't create it otherwise!
// the parameter "obj" must be a pointer; if it's not a pointer, no error will
// be raised but bad things could happen later
func NewNetUI(obj interface{}) NetUI {
	exitchan := make(chan interface{})
	return NetUI{exit: exitchan, Object: obj}
}

// Update sends the new state to every connected client
func (n *NetUI) Update() {
	for _, w := range n.workers {
		select {
		case w.Update <- n.Object:
		default:
			fmt.Fprintln(os.Stderr, "error pushing message to worker on", w.Conn.RemoteAddr())
			// TODO: shall we remove w from the pool on first error?
		}
	}
}

// Close shuts the TCPserver down synchronously
func (n *NetUI) Close() {
	if n.sock != nil {
		(*n.sock).Close()
	}
	for _, w := range n.workers {
		close(w.Exit)
		close(w.Update)
	}
}

// Run does the heavy work, and is blocking. It expects a Listener (typically
// got with net.Listener)
// It will exit when NetUI.Close() is called or when an error on the socket
// happens
func (n *NetUI) Run(sock net.Listener) error {
	if n.sock != nil {
		return fmt.Errorf("More than one Run for a single NetUI is currently unsupported")
	}
	n.sock = &sock
	defer sock.Close()
	for {
		conn, err := sock.Accept()
		if err != nil {
			return nil
		}
		// FIXME: performance problem here: we're waiting each client to be
		// validated before passing to sth else; we'd better run this in a
		// goroutine, but then the access to n.workers would be concurrent,
		// which is not allowed
		err = validateRequest(conn)
		if err != nil {
			continue
		}
		connworker := NewWorker(conn)
		connworker.Update <- n.Object
		n.workers = append(n.workers, connworker)
		go connworker.Work()
	}
}

func validateRequest(conn net.Conn) error {
	_, err := conn.Write([]byte(protocolHello))
	if err != nil {
		conn.Close()
		return err
	}

	return nil

}