1 files changed, 47 insertions, 2 deletions
diff --git a/cmd/dolmetschctld/web_socket.go b/cmd/dolmetschctld/web_socket.go
index 3879650..f904a1c 100644
--- a/cmd/dolmetschctld/web_socket.go
+++ b/cmd/dolmetschctld/web_socket.go
@@ -71,7 +71,7 @@ func sendWebSocketErrorResponse(ws *websocket.Conn, code int, errStr string) {
 	sendWebSocketResponse(ws, rd)
 }
 
-func webSocketHandleCommand(ws *websocket.Conn, sm *StateMachine, req webSocketRequest) {
+func webSocketHandleRequest(ws *websocket.Conn, sm *StateMachine, req webSocketRequest, subCh chan<- FullState) {
 	switch req.Command {
 	case "state":
 		if len(req.Args) != 0 {
@@ -84,6 +84,22 @@ func webSocketHandleCommand(ws *websocket.Conn, sm *StateMachine, req webSocketR
 		result.Type = "state"
 		result.State, result.Ratio, result.Lang = sm.GetState()
 		sendWebSocketResponse(ws, result)
+	case "subscribe":
+		if len(req.Args) != 0 {
+			sendWebSocketErrorResponse(ws, http.StatusBadRequest, "command 'subscribe' expects no arguments")
+			return
+		}
+
+		if err := sm.SubscribeState(subCh); err != nil {
+			sendWebSocketErrorResponse(ws, http.StatusInternalServerError, err.Error())
+			return
+		}
+
+		var result webSocketResponseState
+		result.ResponseCode = http.StatusOK
+		result.Type = "state"
+		result.State, result.Ratio, result.Lang = sm.GetState()
+		sendWebSocketResponse(ws, result)
 	case "languages":
 		if len(req.Args) != 0 {
 			sendWebSocketErrorResponse(ws, http.StatusBadRequest, "command 'languages' expects no arguments")
@@ -116,6 +132,32 @@ func webSocketHandleCommand(ws *websocket.Conn, sm *StateMachine, req webSocketR
 	}
 }
 
+func webSocketSessionHandler(ws *websocket.Conn, sm *StateMachine, reqCh <-chan webSocketRequest) {
+	defer ws.Close()
+
+	subCh := make(chan FullState, 100)
+	for {
+		select {
+		case state, ok := <-subCh:
+			if !ok {
+				return
+			}
+			var result webSocketResponseState
+			result.ResponseCode = http.StatusOK
+			result.Type = "state"
+			result.State = state.state
+			result.Ratio = state.ratio
+			result.Lang = state.language
+			sendWebSocketResponse(ws, result)
+		case req, ok := <-reqCh:
+			if !ok {
+				return
+			}
+			webSocketHandleRequest(ws, sm, req, subCh)
+		}
+	}
+}
+
 func webSocketHandler(sm *StateMachine, w http.ResponseWriter, r *http.Request) {
 	ws, err := websocket.Upgrade(w, r, nil, 64*1024, 64*1024)
 	if _, ok := err.(websocket.HandshakeError); ok {
@@ -126,6 +168,9 @@ func webSocketHandler(sm *StateMachine, w http.ResponseWriter, r *http.Request)
 		return
 	}
 	log.Println("Web(socket) client", ws.RemoteAddr(), "connected")
+	reqCh := make(chan webSocketRequest)
+	go webSocketSessionHandler(ws, sm, reqCh)
+	defer close(reqCh)
 
 	for {
 		t, r, err := ws.NextReader()
@@ -148,7 +193,7 @@ func webSocketHandler(sm *StateMachine, w http.ResponseWriter, r *http.Request)
 				return
 			}
 
-			webSocketHandleCommand(ws, sm, req)
+			reqCh <- req
 		case websocket.BinaryMessage:
 			sendWebSocketErrorResponse(ws, http.StatusBadRequest, "binary messages are not allowed")
 			io.Copy(ioutil.Discard, r) // consume all the data