Fix the build (move out of internal)

socketserver/server/handlecore.go

@@ -0,0 +1,536 @@
+package server // import "bitbucket.org/stendec/frankerfacez/socketserver/server"
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"github.com/gorilla/websocket"
+	"io"
+	"io/ioutil"
+	"log"
+	"net/http"
+	"net/url"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+)
+
+// SuccessCommand is a Reply Command to indicate success in reply to a C2S Command.
+const SuccessCommand Command = "ok"
+
+// ErrorCommand is a Reply Command to indicate that a C2S Command failed.
+const ErrorCommand Command = "error"
+
+// HelloCommand is a C2S Command.
+// HelloCommand must be the Command of the first ClientMessage sent during a connection.
+// Sending any other command will result in a CloseFirstMessageNotHello.
+const HelloCommand Command = "hello"
+
+// AuthorizeCommand is a S2C Command sent as part of Twitch username validation.
+const AuthorizeCommand Command = "do_authorize"
+
+// AsyncResponseCommand is a pseudo-Reply Command.
+// It indicates that the Reply Command to the client's C2S Command will be delivered
+// on a goroutine over the ClientInfo.MessageChannel and should not be delivered immediately.
+const AsyncResponseCommand Command = "_async"
+
+// ResponseSuccess is a Reply ClientMessage with the MessageID not yet filled out.
+var ResponseSuccess = ClientMessage{Command: SuccessCommand}
+
+// Configuration is the active ConfigFile.
+var Configuration *ConfigFile
+
+// SetupServerAndHandle starts all background goroutines and registers HTTP listeners on the given ServeMux.
+// Essentially, this function completely preps the server for a http.ListenAndServe call.
+// (Uses http.DefaultServeMux if `serveMux` is nil.)
+func SetupServerAndHandle(config *ConfigFile, serveMux *http.ServeMux) {
+	Configuration = config
+
+	setupBackend(config)
+
+	if serveMux == nil {
+		serveMux = http.DefaultServeMux
+	}
+
+	bannerBytes, err := ioutil.ReadFile("index.html")
+	if err != nil {
+		log.Fatalln("Could not open index.html:", err)
+	}
+	BannerHTML = bannerBytes
+
+	serveMux.HandleFunc("/", HTTPHandleRootURL)
+	serveMux.HandleFunc("/stats", HTTPShowStatistics)
+
+	serveMux.HandleFunc("/drop_backlog", HTTPBackendDropBacklog)
+	serveMux.HandleFunc("/uncached_pub", HTTPBackendUncachedPublish)
+	serveMux.HandleFunc("/cached_pub", HTTPBackendCachedPublish)
+
+	announceForm, err := SealRequest(url.Values{
+		"startup": []string{"1"},
+	})
+	if err != nil {
+		log.Fatalln("Unable to seal requests:", err)
+	}
+	resp, err := backendHTTPClient.PostForm(announceStartupURL, announceForm)
+	if err != nil {
+		log.Println("could not announce startup to backend:", err)
+	} else {
+		resp.Body.Close()
+	}
+
+	go authorizationJanitor()
+	go backlogJanitor()
+	go bunchCacheJanitor()
+	go pubsubJanitor()
+	go aggregateDataSender()
+
+	go ircConnection()
+}
+
+// SocketUpgrader is the websocket.Upgrader currently in use.
+var SocketUpgrader = websocket.Upgrader{
+	ReadBufferSize:  1024,
+	WriteBufferSize: 1024,
+	CheckOrigin: func(r *http.Request) bool {
+		return r.Header.Get("Origin") == "http://www.twitch.tv"
+	},
+}
+
+// BannerHTML is the content served to web browsers viewing the socket server website.
+// Memes go here.
+var BannerHTML []byte
+
+// HTTPHandleRootURL is the http.HandleFunc for requests on `/`.
+// It either uses the SocketUpgrader or writes out the BannerHTML.
+func HTTPHandleRootURL(w http.ResponseWriter, r *http.Request) {
+	if r.Header.Get("Connection") == "Upgrade" {
+		conn, err := SocketUpgrader.Upgrade(w, r, nil)
+		if err != nil {
+			fmt.Fprintf(w, "error: %v", err)
+			return
+		}
+		RunSocketConnection(conn)
+
+		return
+	} else {
+		w.Write(BannerHTML)
+	}
+}
+
+// ErrProtocolGeneric is sent in a ErrorCommand Reply.
+var ErrProtocolGeneric error = errors.New("FFZ Socket protocol error.")
+
+// ErrProtocolNegativeMsgID is sent in a ErrorCommand Reply when a negative MessageID is received.
+var ErrProtocolNegativeMsgID error = errors.New("FFZ Socket protocol error: negative or zero message ID.")
+
+// ErrExpectedSingleString is sent in a ErrorCommand Reply when the Arguments are of the wrong type.
+var ErrExpectedSingleString = errors.New("Error: Expected single string as arguments.")
+
+// ErrExpectedSingleInt is sent in a ErrorCommand Reply when the Arguments are of the wrong type.
+var ErrExpectedSingleInt = errors.New("Error: Expected single integer as arguments.")
+
+// ErrExpectedTwoStrings is sent in a ErrorCommand Reply when the Arguments are of the wrong type.
+var ErrExpectedTwoStrings = errors.New("Error: Expected array of string, string as arguments.")
+
+// ErrExpectedStringAndBool is sent in a ErrorCommand Reply when the Arguments are of the wrong type.
+var ErrExpectedStringAndBool = errors.New("Error: Expected array of string, bool as arguments.")
+
+// ErrExpectedStringAndInt is sent in a ErrorCommand Reply when the Arguments are of the wrong type.
+var ErrExpectedStringAndInt = errors.New("Error: Expected array of string, int as arguments.")
+
+// ErrExpectedStringAndIntGotFloat is sent in a ErrorCommand Reply when the Arguments are of the wrong type.
+var ErrExpectedStringAndIntGotFloat = errors.New("Error: Second argument was a float, expected an integer.")
+
+// CloseGotBinaryMessage is the termination reason when the client sends a binary websocket frame.
+var CloseGotBinaryMessage = websocket.CloseError{Code: websocket.CloseUnsupportedData, Text: "got binary packet"}
+
+// CloseTimedOut is the termination reason when the client fails to send or respond to ping frames.
+var CloseTimedOut = websocket.CloseError{Code: websocket.CloseNoStatusReceived, Text: "no ping replies for 5 minutes"}
+
+// CloseFirstMessageNotHello is the termination reason
+var CloseFirstMessageNotHello = websocket.CloseError{
+	Text: "Error - the first message sent must be a 'hello'",
+	Code: websocket.ClosePolicyViolation,
+}
+
+// RunSocketConnection contains the main run loop of a websocket connection.
+
+// First, it sets up the channels, the ClientInfo object, and the pong frame handler.
+// It starts the reader goroutine pointing at the newly created channels.
+// The function then enters the run loop (a `for{select{}}`).
+// The run loop is broken when an object is received on errorChan, or if `hello` is not the first C2S Command.
+
+// After the run loop stops, the function launches a goroutine to drain
+// client.MessageChannel, signals the reader goroutine to stop, unsubscribes
+// from all pub/sub channels, waits on MsgChannelKeepalive (remember, the
+// messages are being drained), and finally closes client.MessageChannel
+// (which ends the drainer goroutine).
+func RunSocketConnection(conn *websocket.Conn) {
+	// websocket.Conn is a ReadWriteCloser
+
+	Statistics.ClientConnectsTotal++
+
+	var _closer sync.Once
+	closer := func() {
+		_closer.Do(func() {
+			conn.Close()
+		})
+	}
+
+	// Close the connection when we're done.
+	defer closer()
+
+	_clientChan := make(chan ClientMessage)
+	_serverMessageChan := make(chan ClientMessage)
+	_errorChan := make(chan error)
+	stoppedChan := make(chan struct{})
+
+	var client ClientInfo
+	client.MessageChannel = _serverMessageChan
+	client.RemoteAddr = conn.RemoteAddr()
+	client.MsgChannelIsDone = stoppedChan
+
+	// Launch receiver goroutine
+	go func(errorChan chan<- error, clientChan chan<- ClientMessage, stoppedChan <-chan struct{}) {
+		var msg ClientMessage
+		var messageType int
+		var packet []byte
+		var err error
+		for ; err == nil; messageType, packet, err = conn.ReadMessage() {
+			if messageType == websocket.BinaryMessage {
+				err = &CloseGotBinaryMessage
+				break
+			}
+			if messageType == websocket.CloseMessage {
+				err = io.EOF
+				break
+			}
+
+			UnmarshalClientMessage(packet, messageType, &msg)
+			if msg.MessageID == 0 {
+				continue
+			}
+			select {
+			case clientChan <- msg:
+			case <-stoppedChan:
+				close(errorChan)
+				close(clientChan)
+				return
+			}
+		}
+
+		select {
+		case errorChan <- err:
+		case <-stoppedChan:
+		}
+		close(errorChan)
+		close(clientChan)
+		// exit goroutine
+	}(_errorChan, _clientChan, stoppedChan)
+
+	conn.SetPongHandler(func(pongBody string) error {
+		client.Mutex.Lock()
+		client.pingCount = 0
+		client.Mutex.Unlock()
+		return nil
+	})
+
+	var errorChan <-chan error = _errorChan
+	var clientChan <-chan ClientMessage = _clientChan
+	var serverMessageChan <-chan ClientMessage = _serverMessageChan
+
+	// All set up, now enter the work loop
+
+RunLoop:
+	for {
+		select {
+		case err := <-errorChan:
+			if err == io.EOF {
+				conn.Close() // no need to send a close frame :)
+				break RunLoop
+			} else if closeMsg, isClose := err.(*websocket.CloseError); isClose {
+				CloseConnection(conn, closeMsg)
+			} else {
+				CloseConnection(conn, &websocket.CloseError{
+					Code: websocket.CloseInternalServerErr,
+					Text: err.Error(),
+				})
+			}
+
+			break RunLoop
+
+		case msg := <-clientChan:
+			if client.Version == "" && msg.Command != HelloCommand {
+				CloseConnection(conn, &CloseFirstMessageNotHello)
+				Statistics.FirstNotHelloDisconnects++
+				break RunLoop
+			}
+
+			DispatchC2SCommand(conn, &client, msg)
+
+		case msg := <-serverMessageChan:
+			SendMessage(conn, msg)
+
+		case <-time.After(1 * time.Minute):
+			client.Mutex.Lock()
+			client.pingCount++
+			tooManyPings := client.pingCount == 5
+			client.Mutex.Unlock()
+			if tooManyPings {
+				CloseConnection(conn, &CloseTimedOut)
+				break RunLoop
+			} else {
+				conn.WriteControl(websocket.PingMessage, []byte(strconv.FormatInt(time.Now().Unix(), 10)), getDeadline())
+			}
+		}
+	}
+
+	// Exit
+
+	// Launch message draining goroutine - we aren't out of the pub/sub records
+	go func() {
+		for _ = range _serverMessageChan {
+		}
+	}()
+
+	close(stoppedChan)
+
+	// Stop getting messages...
+	UnsubscribeAll(&client)
+
+	// Wait for pending jobs to finish...
+	client.MsgChannelKeepalive.Wait()
+	client.MessageChannel = nil
+
+	// And done.
+	// Close the channel so the draining goroutine can finish, too.
+	close(_serverMessageChan)
+
+	Statistics.ClientDisconnectsTotal++
+}
+
+func getDeadline() time.Time {
+	return time.Now().Add(1 * time.Minute)
+}
+
+func CloseConnection(conn *websocket.Conn, closeMsg *websocket.CloseError) {
+	Statistics.DisconnectCodes[strconv.Itoa(closeMsg.Code)]++
+	Statistics.DisconnectReasons[closeMsg.Text]++
+
+	conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(closeMsg.Code, closeMsg.Text), getDeadline())
+	conn.Close()
+}
+
+// SendMessage sends a ClientMessage over the websocket connection with a timeout.
+// If marshalling the ClientMessage fails, this function will panic.
+func SendMessage(conn *websocket.Conn, msg ClientMessage) {
+	messageType, packet, err := MarshalClientMessage(msg)
+	if err != nil {
+		panic(fmt.Sprintf("failed to marshal: %v %v", err, msg))
+	}
+	conn.SetWriteDeadline(getDeadline())
+	conn.WriteMessage(messageType, packet)
+	Statistics.MessagesSent++
+}
+
+// UnmarshalClientMessage unpacks websocket TextMessage into a ClientMessage provided in the `v` parameter.
+func UnmarshalClientMessage(data []byte, payloadType int, v interface{}) (err error) {
+	var spaceIdx int
+
+	out := v.(*ClientMessage)
+	dataStr := string(data)
+
+	// Message ID
+	spaceIdx = strings.IndexRune(dataStr, ' ')
+	if spaceIdx == -1 {
+		return ErrProtocolGeneric
+	}
+	messageID, err := strconv.Atoi(dataStr[:spaceIdx])
+	if messageID < -1 || messageID == 0 {
+		return ErrProtocolNegativeMsgID
+	}
+
+	out.MessageID = messageID
+	dataStr = dataStr[spaceIdx+1:]
+
+	spaceIdx = strings.IndexRune(dataStr, ' ')
+	if spaceIdx == -1 {
+		out.Command = Command(dataStr)
+		out.Arguments = nil
+		return nil
+	} else {
+		out.Command = Command(dataStr[:spaceIdx])
+	}
+	dataStr = dataStr[spaceIdx+1:]
+	argumentsJSON := dataStr
+	out.origArguments = argumentsJSON
+	err = out.parseOrigArguments()
+	if err != nil {
+		return
+	}
+	return nil
+}
+
+func (cm *ClientMessage) parseOrigArguments() error {
+	err := json.Unmarshal([]byte(cm.origArguments), &cm.Arguments)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func MarshalClientMessage(clientMessage interface{}) (payloadType int, data []byte, err error) {
+	var msg ClientMessage
+	var ok bool
+	msg, ok = clientMessage.(ClientMessage)
+	if !ok {
+		pMsg, ok := clientMessage.(*ClientMessage)
+		if !ok {
+			panic("MarshalClientMessage: argument needs to be a ClientMessage")
+		}
+		msg = *pMsg
+	}
+	var dataStr string
+
+	if msg.Command == "" && msg.MessageID == 0 {
+		panic("MarshalClientMessage: attempt to send an empty ClientMessage")
+	}
+
+	if msg.Command == "" {
+		msg.Command = SuccessCommand
+	}
+	if msg.MessageID == 0 {
+		msg.MessageID = -1
+	}
+
+	if msg.Arguments != nil {
+		argBytes, err := json.Marshal(msg.Arguments)
+		if err != nil {
+			return 0, nil, err
+		}
+
+		dataStr = fmt.Sprintf("%d %s %s", msg.MessageID, msg.Command, string(argBytes))
+	} else {
+		dataStr = fmt.Sprintf("%d %s", msg.MessageID, msg.Command)
+	}
+
+	return websocket.TextMessage, []byte(dataStr), nil
+}
+
+// Convenience method: Parse the arguments of the ClientMessage as a single string.
+func (cm *ClientMessage) ArgumentsAsString() (string1 string, err error) {
+	var ok bool
+	string1, ok = cm.Arguments.(string)
+	if !ok {
+		err = ErrExpectedSingleString
+		return
+	} else {
+		return string1, nil
+	}
+}
+
+// Convenience method: Parse the arguments of the ClientMessage as a single int.
+func (cm *ClientMessage) ArgumentsAsInt() (int1 int64, err error) {
+	var ok bool
+	var num float64
+	num, ok = cm.Arguments.(float64)
+	if !ok {
+		err = ErrExpectedSingleInt
+		return
+	} else {
+		int1 = int64(num)
+		return int1, nil
+	}
+}
+
+// Convenience method: Parse the arguments of the ClientMessage as an array of two strings.
+func (cm *ClientMessage) ArgumentsAsTwoStrings() (string1, string2 string, err error) {
+	var ok bool
+	var ary []interface{}
+	ary, ok = cm.Arguments.([]interface{})
+	if !ok {
+		err = ErrExpectedTwoStrings
+		return
+	} else {
+		if len(ary) != 2 {
+			err = ErrExpectedTwoStrings
+			return
+		}
+		string1, ok = ary[0].(string)
+		if !ok {
+			err = ErrExpectedTwoStrings
+			return
+		}
+		// clientID can be null
+		if ary[1] == nil {
+			return string1, "", nil
+		}
+		string2, ok = ary[1].(string)
+		if !ok {
+			err = ErrExpectedTwoStrings
+			return
+		}
+		return string1, string2, nil
+	}
+}
+
+// Convenience method: Parse the arguments of the ClientMessage as an array of a string and an int.
+func (cm *ClientMessage) ArgumentsAsStringAndInt() (string1 string, int int64, err error) {
+	var ok bool
+	var ary []interface{}
+	ary, ok = cm.Arguments.([]interface{})
+	if !ok {
+		err = ErrExpectedStringAndInt
+		return
+	} else {
+		if len(ary) != 2 {
+			err = ErrExpectedStringAndInt
+			return
+		}
+		string1, ok = ary[0].(string)
+		if !ok {
+			err = ErrExpectedStringAndInt
+			return
+		}
+		var num float64
+		num, ok = ary[1].(float64)
+		if !ok {
+			err = ErrExpectedStringAndInt
+			return
+		}
+		int = int64(num)
+		if float64(int) != num {
+			err = ErrExpectedStringAndIntGotFloat
+			return
+		}
+		return string1, int, nil
+	}
+}
+
+// Convenience method: Parse the arguments of the ClientMessage as an array of a string and an int.
+func (cm *ClientMessage) ArgumentsAsStringAndBool() (str string, flag bool, err error) {
+	var ok bool
+	var ary []interface{}
+	ary, ok = cm.Arguments.([]interface{})
+	if !ok {
+		err = ErrExpectedStringAndBool
+		return
+	} else {
+		if len(ary) != 2 {
+			err = ErrExpectedStringAndBool
+			return
+		}
+		str, ok = ary[0].(string)
+		if !ok {
+			err = ErrExpectedStringAndBool
+			return
+		}
+		flag, ok = ary[1].(bool)
+		if !ok {
+			err = ErrExpectedStringAndBool
+			return
+		}
+		return str, flag, nil
+	}
+}