package main import ( "bytes" "encoding/json" "errors" "fmt" "github.com/gorilla/websocket" "github.com/julienschmidt/httprouter" "github.com/pterodactyl/wings/config" "github.com/pterodactyl/wings/server" "go.uber.org/zap" "io/ioutil" "net/http" "os" "strings" "sync" ) const ( SetStateEvent = "set state" SendServerLogsEvent = "send logs" SendCommandEvent = "send command" ) type WebsocketMessage struct { // The event to perform. Should be one of the following that are supported: // // - status : Returns the server's power state. // - logs : Returns the server log data at the time of the request. // - power : Performs a power action aganist the server based the data. // - command : Performs a command on a server using the data field. Event string `json:"event"` // 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"` inbound bool } type WebsocketHandler struct { Server *server.Server Mutex sync.Mutex Connection *websocket.Conn } type socketCredentials struct { ServerUuid string `json:"server_uuid"` } func (rt *Router) AuthenticateWebsocket(h httprouter.Handle) httprouter.Handle { return func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) { s := rt.Servers.Get(ps.ByName("server")) j, err := json.Marshal(socketCredentials{ServerUuid: s.Uuid}) if err != nil { zap.S().Errorw("failed to marshal json", zap.Error(err)) http.Error(w, "failed to marshal json", http.StatusInternalServerError) return } url := strings.TrimRight(config.Get().PanelLocation, "/") + "/api/remote/websocket/" + ps.ByName("token") req, err := http.NewRequest(http.MethodPost, url, bytes.NewBuffer(j)) if err != nil { zap.S().Errorw("failed to generate a new HTTP request when validating websocket credentials", zap.Error(err)) http.Error(w, "failed to generate HTTP request", http.StatusInternalServerError) return } req.Header.Set("Accept", "application/json") req.Header.Set("Content-Type", "application/json") req.Header.Set("Authorization", "Bearer "+config.Get().AuthenticationToken) client := &http.Client{} resp, err := client.Do(req) if err != nil { zap.S().Errorw("failed to perform client HTTP request", zap.Error(err)) http.Error(w, "failed to perform client HTTP request", http.StatusInternalServerError) return } defer resp.Body.Close() if resp.StatusCode != http.StatusNoContent { b, _ := ioutil.ReadAll(resp.Body) zap.S().Warnw("failed to validate token with server", zap.String("response", string(b))) http.Error(w, "failed to validate token with server", resp.StatusCode) return } h(w, r, ps) } } // Handle a request for a specific server websocket. This will handle inbound requests as well // as ensure that any console output is also passed down the wire on the socket. func (rt *Router) routeWebsocket(w http.ResponseWriter, r *http.Request, ps httprouter.Params) { c, err := rt.upgrader.Upgrade(w, r, nil) if err != nil { zap.S().Error(err) return } defer c.Close() fmt.Println("running websocket route") s := rt.Servers.Get(ps.ByName("server")) handler := WebsocketHandler{ Server: s, Mutex: sync.Mutex{}, Connection: c, } handleOutput := func(data string) { handler.SendJson(&WebsocketMessage{ Event: server.ConsoleOutputEvent, Args: []string{data}, }) } handleServerStatus := func(data string) { handler.SendJson(&WebsocketMessage{ Event: server.StatusEvent, Args: []string{data}, }) } handleResourceUse := func(data string) { handler.SendJson(&WebsocketMessage{ Event: server.StatsEvent, Args: []string{data}, }) } s.AddListener(server.StatusEvent, &handleServerStatus) defer s.RemoveListener(server.StatusEvent, &handleServerStatus) s.AddListener(server.ConsoleOutputEvent, &handleOutput) defer s.RemoveListener(server.ConsoleOutputEvent, &handleOutput) s.AddListener(server.StatsEvent, &handleResourceUse) defer s.RemoveListener(server.StatsEvent, &handleResourceUse) s.Emit(server.StatusEvent, s.State) for { j := WebsocketMessage{inbound: true} fmt.Println("running for{} loop...") if _, _, err := c.ReadMessage(); err != nil { if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseNoStatusReceived, websocket.CloseServiceRestart) { zap.S().Errorw("error handling websocket message", zap.Error(err)) } break } // 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 // from the socket, which is NOT what we want to do. if err := c.ReadJSON(&j); err != nil { fmt.Println("ReadJSON() error") continue } if err := handler.HandleInbound(j); err != nil { zap.S().Warnw("error handling inbound websocket request", zap.Error(err)) break } else { zap.S().Debugw("handled event", zap.String("event", j.Event), zap.Strings("args", j.Args)) } fmt.Println("finished looping...") } } // Perform a blocking send operation on the websocket since we want to avoid any // concurrent writes to the connection, which would cause a runtime panic and cause // the program to crash out. func (wsh *WebsocketHandler) SendJson(v interface{}) error { wsh.Mutex.Lock() defer wsh.Mutex.Unlock() return wsh.Connection.WriteJSON(v) } // Handle the inbound socket request and route it to the proper server action. func (wsh *WebsocketHandler) HandleInbound(m WebsocketMessage) error { if !m.inbound { return errors.New("cannot handle websocket message, not an inbound connection") } switch m.Event { case SetStateEvent: { var err error switch strings.Join(m.Args, "") { case "start": err = wsh.Server.Environment.Start() break case "stop": err = wsh.Server.Environment.Stop() break case "restart": break case "kill": err = wsh.Server.Environment.Terminate(os.Kill) break } if err != nil { return err } } case SendServerLogsEvent: { if running, _ := wsh.Server.Environment.IsRunning(); !running { return nil } logs, err := wsh.Server.Environment.Readlog(1024 * 16) if err != nil { return err } for _, line := range logs { wsh.SendJson(&WebsocketMessage{ Event: server.ConsoleOutputEvent, Args: []string{line}, }) } return nil } case SendCommandEvent: { return wsh.Server.Environment.SendCommand(strings.Join(m.Args, "")) } } return nil }