From b20bfac36bef908ff64024c03c04fb64caf367e4 Mon Sep 17 00:00:00 2001
From: Dane Everitt <dane@daneeveritt.com>
Date: Sat, 20 Apr 2019 14:57:37 -0700
Subject: [PATCH] Get server logs sending over the socket. wowee

---
 go.mod           |  1 +
 go.sum           |  2 ++
 server/server.go | 13 +++++++++
 websocket.go     | 69 ++++++++++++++++++++++++++++++++++++++++++++----
 4 files changed, 80 insertions(+), 5 deletions(-)

diff --git a/go.mod b/go.mod
index ed47ea3..e9ed40e 100644
--- a/go.mod
+++ b/go.mod
@@ -17,6 +17,7 @@ require (
 	github.com/gorilla/websocket v1.4.0
 	github.com/gotestyourself/gotestyourself v2.2.0+incompatible // indirect
 	github.com/julienschmidt/httprouter v1.2.0
+	github.com/olebedev/emitter v0.0.0-20190110104742-e8d1457e6aee
 	github.com/onsi/ginkgo v1.8.0 // indirect
 	github.com/onsi/gomega v1.5.0 // indirect
 	github.com/opencontainers/go-digest v1.0.0-rc1 // indirect
diff --git a/go.sum b/go.sum
index 55d5023..9bcfe5b 100644
--- a/go.sum
+++ b/go.sum
@@ -88,6 +88,8 @@ github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5
 github.com/microcosm-cc/bluemonday v1.0.1/go.mod h1:hsXNsILzKxV+sX77C5b8FSuKF00vh2OMYv+xgHpAMF4=
 github.com/neelance/astrewrite v0.0.0-20160511093645-99348263ae86/go.mod h1:kHJEU3ofeGjhHklVoIGuVj85JJwZ6kWPaJwCIxgnFmo=
 github.com/neelance/sourcemap v0.0.0-20151028013722-8c68805598ab/go.mod h1:Qr6/a/Q4r9LP1IltGz7tA7iOK1WonHEYhu1HRBA7ZiM=
+github.com/olebedev/emitter v0.0.0-20190110104742-e8d1457e6aee h1:IquUs3fIykn10zWDIyddanhpTqBvAHMaPnFhQuyYw5U=
+github.com/olebedev/emitter v0.0.0-20190110104742-e8d1457e6aee/go.mod h1:eT2/Pcsim3XBjbvldGiJBvvgiqZkAFyiOJJsDKXs/ts=
 github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
 github.com/onsi/ginkgo v1.8.0 h1:VkHVNpR4iVnU8XQR6DBm8BqYjN7CRzw+xKUbVVbbW9w=
 github.com/onsi/ginkgo v1.8.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
diff --git a/server/server.go b/server/server.go
index 9ff38f3..f610b7c 100644
--- a/server/server.go
+++ b/server/server.go
@@ -1,6 +1,7 @@
 package server
 
 import (
+	"github.com/olebedev/emitter"
 	"github.com/patrickmn/go-cache"
 	"github.com/pkg/errors"
 	"github.com/pterodactyl/wings/config"
@@ -60,6 +61,8 @@ type Server struct {
 	// Server cache used to store frequently requested information in memory and make
 	// certain long operations return faster. For example, FS disk space usage.
 	cache *cache.Cache
+
+	emitter *emitter.Emitter
 }
 
 // The build settings for a given server that impact docker container creation and
@@ -257,3 +260,13 @@ func (s *Server) IsBootable() bool {
 func (s *Server) CreateEnvironment() error {
 	return s.environment.Create()
 }
+
+// Returns the server event emitter instance. If one has not been setup yet, a new instance
+// will be created.
+func (s *Server) Emitter() *emitter.Emitter {
+	if s.emitter == nil {
+		s.emitter = &emitter.Emitter{}
+	}
+
+	return s.emitter
+}
\ No newline at end of file
diff --git a/websocket.go b/websocket.go
index d042c2e..71ecf98 100644
--- a/websocket.go
+++ b/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
+}
\ No newline at end of file