Get server logs sending over the socket. wowee

websocket.go

@@ -1,10 +1,14 @@
 package main
 
 import (
+	"errors"
 	"fmt"
+	"github.com/gorilla/websocket"
 	"github.com/julienschmidt/httprouter"
+	"github.com/pterodactyl/wings/server"
 	"go.uber.org/zap"
 	"net/http"
+	"os"
 	"strings"
 )
 
@@ -20,6 +24,10 @@ type WebsocketMessage struct {
 	// The data to pass along, only used by power/command currently. Other requests
 	// should either omit the field or pass an empty value as it is ignored.
 	Args []string `json:"args,omitempty"`
+
+	server *server.Server
+
+	inbound bool
 }
 
 func (rt *Router) routeWebsocket(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
@@ -33,7 +41,7 @@ func (rt *Router) routeWebsocket(w http.ResponseWriter, r *http.Request, ps http
 	s := rt.Servers.Get(ps.ByName("server"))
 
 	for {
-		j := WebsocketMessage{}
+		j := WebsocketMessage{server: s, inbound: true}
 
 		// Discard and JSON parse errors into the void and don't continue processing this
 		// specific socket request. If we did a break here the client would get disconnected
@@ -42,11 +50,62 @@ func (rt *Router) routeWebsocket(w http.ResponseWriter, r *http.Request, ps http
 			break
 		}
 
-		fmt.Printf("%s sent: %s = %s\n", s.Uuid, j.Event, strings.Join(j.Args, " "))
-
-		if err := c.WriteJSON(WebsocketMessage{Event: j.Event, Args: []string{}}); err != nil {
-			zap.S().Warnw("error writing JSON to socket", zap.Error(err))
+		fmt.Printf("%s received: %s = %s\n", s.Uuid, j.Event, strings.Join(j.Args, " "))
+		if err := j.HandleInbound(c); err != nil {
+			zap.S().Warnw("error handling inbound websocket request", zap.Error(err))
 			break
 		}
 	}
+
+	zap.S().Debugw("disconnected from instance", zap.String("ip", c.RemoteAddr().String()))
 }
+
+func (wsm *WebsocketMessage) HandleInbound(c *websocket.Conn) error {
+	if !wsm.inbound {
+		return errors.New("cannot handle websocket message, not an inbound connection")
+	}
+
+	switch wsm.Event {
+	case "set state":
+		{
+			var err error
+			switch strings.Join(wsm.Args, "") {
+			case "start":
+				err = wsm.server.Environment().Start()
+				break
+			case "stop":
+				err = wsm.server.Environment().Stop()
+				break
+			case "restart":
+				err = wsm.server.Environment().Terminate(os.Kill)
+				break
+			}
+
+			if err != nil {
+				return err
+			}
+		}
+	case "send logs":
+		{
+			logs, err := wsm.server.Environment().Readlog(1024 * 5)
+			if err != nil {
+				return err
+			}
+
+			for _, line := range logs {
+				c.WriteJSON(&WebsocketMessage{
+					Event: "console output",
+					Args: []string{line},
+				})
+			}
+
+			return nil
+		}
+	case "send command":
+		{
+			return nil
+		}
+	}
+
+	return nil
+}