diff --git a/http.go b/http.go index ecb17b7..f885870 100644 --- a/http.go +++ b/http.go @@ -4,7 +4,7 @@ import ( "bufio" "encoding/json" "fmt" - "github.com/googollee/go-socket.io" + "github.com/gorilla/websocket" "github.com/julienschmidt/httprouter" "github.com/pterodactyl/wings/server" "go.uber.org/zap" @@ -31,7 +31,7 @@ func (sc *ServerCollection) Get(uuid string) *server.Server { type Router struct { Servers ServerCollection - Socketio *socketio.Server + upgrader websocket.Upgrader // The authentication token defined in the config.yml file that allows // a request to perform any action aganist the daemon. @@ -58,7 +58,15 @@ func (rt *Router) AuthenticateToken(permission string, h httprouter.Handle) http return func(w http.ResponseWriter, r *http.Request, ps httprouter.Params) { t := strings.Split(permission, ":")[0] + // Adds support for using this middleware on the websocket routes for servers. Those + // routes don't support Authorization headers, per the spec, so we abuse the socket + // protocol header and use that to pass the authorization token along to Wings without + // exposing the token in the URL directly. Neat. 📸 auth := strings.SplitN(r.Header.Get("Authorization"), " ", 2) + if r.Header.Get("Sec-WebSocket-Protocol") != "" { + auth = []string{"Bearer", r.Header.Get("Sec-WebSocket-Protocol")} + } + if len(auth) != 2 || auth[0] != "Bearer" { w.Header().Set("WWW-Authenticate", "Bearer") http.Error(w, "authorization failed", http.StatusUnauthorized) @@ -85,6 +93,7 @@ func (rt *Router) AuthenticateToken(permission string, h httprouter.Handle) http } } + // Happens because we don't have any of the server handling code here. http.Error(w, "not implemented", http.StatusNotImplemented) return } @@ -266,7 +275,6 @@ func (rt *Router) ConfigureRouter() *httprouter.Router { router.POST("/api/servers/:server/power", rt.AuthenticateToken("s:power", rt.AuthenticateServer(rt.routeServerPower))) - router.Handler("GET", "/socket.io/", rt.Socketio) - + router.GET("/api/servers/:server/ws", rt.AuthenticateToken("s:websocket", rt.AuthenticateServer(rt.routeWebsocket))) return router } diff --git a/websocket.go b/websocket.go index daca1a2..c41e227 100644 --- a/websocket.go +++ b/websocket.go @@ -2,27 +2,49 @@ package main import ( "fmt" - "github.com/googollee/go-socket.io" + "github.com/gorilla/websocket" + "github.com/julienschmidt/httprouter" "go.uber.org/zap" + "net/http" ) -// Configures the websocket connection and attaches it to the Router struct. -func (rt *Router) ConfigureWebsocket() (*socketio.Server, error) { - s, err := socketio.NewServer(nil) +type WebsocketMessage struct { + // The action 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. + Action string + // The data to pass along, only used by power/command currently. Other requests + // should either omit the field or pass an empty string value as it is ignored. + Data string +} + +func (rt *Router) routeWebsocket(w http.ResponseWriter, r *http.Request, ps httprouter.Params) { + c, err := rt.upgrader.Upgrade(w, r, nil) if err != nil { - return nil, err + zap.S().Error(err) + return } + defer c.Close() - s.OnConnect("/", func(s socketio.Conn) error { - s.SetContext("") - fmt.Println("connected:", s.ID()) - return nil - }) + for { + j := WebsocketMessage{} - s.OnError("/", func(e error) { - zap.S().Error(e) - }) + // 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 { + continue + } - return s, nil -} \ No newline at end of file + fmt.Printf("%s sent: %s = %s\n", c.RemoteAddr(), j.Action, j.Data) + + if err := c.WriteMessage(websocket.TextMessage, []byte("")); err != nil { + zap.S().Warnw("error writing JSON to socket", zap.Error(err)) + continue + } + } +} diff --git a/wings.go b/wings.go index 33c3d56..6824cb1 100644 --- a/wings.go +++ b/wings.go @@ -3,6 +3,7 @@ package main import ( "flag" "fmt" + "github.com/gorilla/websocket" "github.com/pterodactyl/wings/config" "github.com/pterodactyl/wings/server" "go.uber.org/zap" @@ -71,18 +72,15 @@ func main() { r := &Router{ Servers: servers, token: c.AuthenticationToken, + upgrader: websocket.Upgrader{ + // Ensure that the websocket request is originating from the Panel itself, + // and not some other location. + CheckOrigin: func(r *http.Request) bool { + return r.Header.Get("Origin") == c.PanelLocation + }, + }, } - if sock, err := r.ConfigureWebsocket(); err != nil { - zap.S().Fatalw("failed to configure websocket", zap.Error(err)) - return - } else { - r.Socketio = sock - } - - defer r.Socketio.Close() - go r.Socketio.Serve() - router := r.ConfigureRouter() zap.S().Infow("configuring webserver", zap.String("host", c.Api.Host), zap.Int("port", c.Api.Port)) if err := http.ListenAndServe(fmt.Sprintf("%s:%d", c.Api.Host, c.Api.Port), router); err != nil {