From 640f4b3a983ae8e9575c0dc541e9c277e681d1b8 Mon Sep 17 00:00:00 2001
From: Dane Everitt <dane@daneeveritt.com>
Date: Sat, 28 Sep 2019 13:01:04 -0700
Subject: [PATCH] Send errors back over the socket to users depending on
 permission

---
 go.mod       |  1 +
 go.sum       |  2 ++
 websocket.go | 63 +++++++++++++++++++++++++++++++++++++++++++++++-----
 3 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/go.mod b/go.mod
index 3cb4fc1..4893707 100644
--- a/go.mod
+++ b/go.mod
@@ -15,6 +15,7 @@ require (
 	github.com/gabriel-vasile/mimetype v0.1.4
 	github.com/gbrlsnchs/jwt/v3 v3.0.0-rc.0
 	github.com/gogo/protobuf v1.2.1 // indirect
+	github.com/google/uuid v1.1.1
 	github.com/gorilla/websocket v1.4.0
 	github.com/gotestyourself/gotestyourself v2.2.0+incompatible // indirect
 	github.com/julienschmidt/httprouter v1.2.0
diff --git a/go.sum b/go.sum
index a43f454..817accd 100644
--- a/go.sum
+++ b/go.sum
@@ -29,6 +29,8 @@ github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5y
 github.com/google/go-cmp v0.2.0 h1:+dTQ8DZQJz0Mb/HjFlkptS1FeQ4cWSnN941F8aEG4SQ=
 github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
 github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
+github.com/google/uuid v1.1.1 h1:Gkbcsh/GbpXz7lPftLA3P6TYMwjCLYm83jiFQZF/3gY=
+github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/gorilla/websocket v1.4.0 h1:WDFjx/TMzVgy9VdMMQi2K2Emtwi2QcUQsztZ/zLaH/Q=
 github.com/gorilla/websocket v1.4.0/go.mod h1:E7qHFY5m1UJ88s3WnNqhKjPHQ0heANvMoAMk2YaljkQ=
 github.com/gotestyourself/gotestyourself v2.2.0+incompatible h1:AQwinXlbQR2HvPjQZOmDhRqsv5mZf+Jb1RnSLxcqZcI=
diff --git a/websocket.go b/websocket.go
index 65229f2..dde0f3c 100644
--- a/websocket.go
+++ b/websocket.go
@@ -3,7 +3,9 @@ package main
 import (
 	"encoding/json"
 	"errors"
+	"fmt"
 	"github.com/gbrlsnchs/jwt/v3"
+	"github.com/google/uuid"
 	"github.com/gorilla/websocket"
 	"github.com/julienschmidt/httprouter"
 	"github.com/pterodactyl/wings/config"
@@ -19,9 +21,11 @@ import (
 const (
 	TokenExpiringEvent  = "token expiring"
 	TokenExpiredEvent   = "token expired"
+	AuthenticationEvent = "auth"
 	SetStateEvent       = "set state"
 	SendServerLogsEvent = "send logs"
 	SendCommandEvent    = "send command"
+	ErrorEvent          = "daemon error"
 )
 
 type WebsocketMessage struct {
@@ -57,9 +61,10 @@ type WebsocketTokenPayload struct {
 }
 
 const (
-	PermissionConnect     = "connect"
-	PermissionSendCommand = "send-command"
-	PermissionSendPower   = "send-power"
+	PermissionConnect       = "connect"
+	PermissionSendCommand   = "send-command"
+	PermissionSendPower     = "send-power"
+	PermissionReceiveErrors = "receive-errors"
 )
 
 // Checks if the given token payload has a permission string.
@@ -248,8 +253,7 @@ func (rt *Router) routeWebsocket(w http.ResponseWriter, r *http.Request, ps http
 		}
 
 		if err := handler.HandleInbound(j); err != nil {
-			zap.S().Warnw("error handling inbound websocket request", zap.Error(err))
-			break
+			handler.SendErrorJson(err)
 		}
 	}
 }
@@ -264,6 +268,38 @@ func (wsh *WebsocketHandler) SendJson(v interface{}) error {
 	return wsh.Connection.WriteJSON(v)
 }
 
+// Sends an error back to the connected websocket instance by checking the permissions
+// of the token. If the user has the "receive-errors" grant we will send back the actual
+// error message, otherwise we just send back a standard error message.
+func (wsh *WebsocketHandler) SendErrorJson(err error) error {
+	wsh.Mutex.Lock()
+	defer wsh.Mutex.Unlock()
+
+	message := "an unexpected error was encountered during the websocket lifecycle"
+	if wsh.JWT != nil && wsh.JWT.HasPermission(PermissionReceiveErrors) {
+		message = err.Error()
+	}
+
+	m, u := wsh.GetErrorMessage(message)
+
+	wsm := WebsocketMessage{Event: ErrorEvent}
+	wsm.Args = []string{m}
+
+	zap.S().Warnw("an error was encountered in the websocket process", zap.String("error_identifier", u.String()), zap.Error(err))
+
+	return wsh.Connection.WriteJSON(wsm)
+}
+
+// Converts an error message into a more readable representation and returns a UUID
+// that can be cross-referenced to find the specific error that triggered.
+func (wsh *WebsocketHandler) GetErrorMessage(msg string) (string, uuid.UUID) {
+	u, _ := uuid.NewRandom()
+
+	m := fmt.Sprintf("Error Event [%s]: %s", u.String(), msg)
+
+	return m, u
+}
+
 // Handle the inbound socket request and route it to the proper server action.
 func (wsh *WebsocketHandler) HandleInbound(m WebsocketMessage) error {
 	if !m.inbound {
@@ -277,6 +313,19 @@ func (wsh *WebsocketHandler) HandleInbound(m WebsocketMessage) error {
 	}
 
 	switch m.Event {
+	case AuthenticationEvent:
+		{
+			token, err := ParseJWT([]byte(strings.Join(m.Args, "")))
+			if err != nil {
+				return nil
+			}
+
+			if token.HasPermission(PermissionConnect) {
+				wsh.JWT = token
+			}
+
+			return nil
+		}
 	case SetStateEvent:
 		{
 			if !wsh.JWT.HasPermission(PermissionSendPower) {
@@ -328,6 +377,10 @@ func (wsh *WebsocketHandler) HandleInbound(m WebsocketMessage) error {
 				return nil
 			}
 
+			if wsh.Server.State == server.ProcessOfflineState {
+				return nil
+			}
+
 			return wsh.Server.Environment.SendCommand(strings.Join(m.Args, ""))
 		}
 	}