addes subscribe command to websocket

1 files changed, 93 insertions, 2 deletions

diff --git a/src/rhctl/web_socket.go b/src/rhctl/web_socket.go
index 5f90100..df12bde 100644
--- a/src/rhctl/web_socket.go
+++ b/src/rhctl/web_socket.go
@@ -27,12 +27,14 @@ import (
 	"io"
 	"io/ioutil"
 	"net/http"
+	"strings"
 
 	"github.com/gorilla/websocket"
 )
 
 type webSocketRequestData struct {
-	Command string `json:"COMMAND"`
+	Command string   `json:"COMMAND"`
+	Args    []string `json:"ARGS"`
 }
 
 type webSocketResponseBaseData struct {
@@ -43,7 +45,12 @@ type webSocketResponseBaseData struct {
 
 type webSocketResponseStateData struct {
 	webSocketResponseBaseData
-	State State `json:"STATE"`
+	State interface{} `json:"STATE"`
+}
+
+type webSocketResponseUpdateData struct {
+	webSocketResponseBaseData
+	Update interface{} `json:"UPDATE"`
 }
 
 func sendWebSocketResponse(ws *websocket.Conn, rd interface{}) {
@@ -60,6 +67,14 @@ func sendWebSocketErrorResponse(ws *websocket.Conn, code int, errStr string) {
 	sendWebSocketResponse(ws, rd)
 }
 
+func sendWebSocketAckResponse(ws *websocket.Conn) {
+	rd := &webSocketResponseBaseData{}
+	rd.ResponseCode = http.StatusOK
+	rd.Type = "ack"
+	rd.ErrorString = "OK"
+	sendWebSocketResponse(ws, rd)
+}
+
 func sendWebSocketStateResponse(ws *websocket.Conn, state State) {
 	rd := &webSocketResponseStateData{}
 	rd.ResponseCode = http.StatusOK
@@ -69,9 +84,54 @@ func sendWebSocketStateResponse(ws *websocket.Conn, state State) {
 	sendWebSocketResponse(ws, rd)
 }
 
+func sendWebSocketServerStateResponse(ws *websocket.Conn, state ServerState) {
+	rd := &webSocketResponseStateData{}
+	rd.ResponseCode = http.StatusOK
+	rd.Type = "server-state"
+	rd.ErrorString = "OK"
+	rd.State = state
+	sendWebSocketResponse(ws, rd)
+}
+
+func sendWebSocketSwitchStateResponse(ws *websocket.Conn, state SwitchState) {
+	rd := &webSocketResponseStateData{}
+	rd.ResponseCode = http.StatusOK
+	rd.Type = "switch-state"
+	rd.ErrorString = "OK"
+	rd.State = state
+	sendWebSocketResponse(ws, rd)
+}
+
+func sendWebSocketSwitchUpdateResponse(ws *websocket.Conn, update SwitchUpdate) {
+	rd := &webSocketResponseUpdateData{}
+	rd.ResponseCode = http.StatusOK
+	rd.Type = "switch-update"
+	rd.ErrorString = "OK"
+	rd.Update = update
+	sendWebSocketResponse(ws, rd)
+}
+
+func sendWebSocketUpdateData(ws *websocket.Conn, data interface{}) {
+	switch data.(type) {
+	case State:
+		sendWebSocketStateResponse(ws, data.(State))
+	case ServerState:
+		sendWebSocketServerStateResponse(ws, data.(ServerState))
+	case SwitchState:
+		sendWebSocketSwitchStateResponse(ws, data.(SwitchState))
+	case SwitchUpdate:
+		sendWebSocketSwitchUpdateResponse(ws, data.(SwitchUpdate))
+	default:
+		sendWebSocketErrorResponse(ws, http.StatusInternalServerError, "got invalid data update")
+	}
+}
+
 func webSocketSessionHandler(reqchan <-chan webSocketRequestData, ws *websocket.Conn, ctrl *SwitchControl) {
 	defer ws.Close()
 
+	updateC := ctrl.Updates.Sub()
+	defer ctrl.Updates.Unsub(updateC)
+
 	for {
 		select {
 		case reqdata, ok := <-reqchan:
@@ -91,9 +151,40 @@ func webSocketSessionHandler(reqchan <-chan webSocketRequestData, ws *websocket.
 				default:
 					sendWebSocketErrorResponse(ws, http.StatusInternalServerError, fmt.Sprintf("invalid response of type %T: %+v", result, result))
 				}
+			case "subscribe":
+				if len(reqdata.Args) == 1 {
+					switch strings.ToLower(reqdata.Args[0]) {
+					case "state":
+						ctrl.Updates.AddSub(updateC, "state")
+						sendWebSocketAckResponse(ws)
+					case "server":
+						ctrl.Updates.AddSub(updateC, "server:state")
+						sendWebSocketAckResponse(ws)
+					case "switch":
+						ctrl.Updates.AddSub(updateC, "switch:state")
+						sendWebSocketAckResponse(ws)
+					case "audio":
+						fallthrough
+					case "gpi":
+						fallthrough
+					case "oc":
+						fallthrough
+					case "relay":
+						fallthrough
+					case "silence":
+						ctrl.Updates.AddSub(updateC, "switch:"+reqdata.Args[0])
+						sendWebSocketAckResponse(ws)
+					default:
+						sendWebSocketErrorResponse(ws, http.StatusInternalServerError, fmt.Sprintf("unknown message type '%s'", reqdata.Args[0]))
+					}
+				} else {
+					sendWebSocketErrorResponse(ws, http.StatusInternalServerError, "subscribe takes exactly one argument")
+				}
 			default:
 				sendWebSocketErrorResponse(ws, http.StatusBadRequest, fmt.Sprintf("unknown command '%s'", reqdata.Command))
 			}
+		case update := <-updateC:
+			sendWebSocketUpdateData(ws, update)
 		}
 	}
 }