Better core logic for JWT; supports a more generic structure

router/tokens/parser.go

@@ -0,0 +1,33 @@
+package tokens
+
+import (
+	"github.com/gbrlsnchs/jwt/v3"
+	"github.com/pterodactyl/wings/config"
+	"time"
+)
+
+var alg *jwt.HMACSHA
+
+type TokenData interface {
+	GetPayload() *jwt.Payload
+}
+
+// Validates the provided JWT against the known secret for the Daemon and returns the
+// parsed data. This function DOES NOT validate that the token is valid for the connected
+// server, nor does it ensure that the user providing the token is able to actually do things.
+//
+// This simply returns a parsed token.
+func ParseToken(token []byte, data TokenData) error {
+	if alg == nil {
+		alg = jwt.NewHS256([]byte(config.Get().AuthenticationToken))
+	}
+
+	verifyOptions := jwt.ValidatePayload(
+		data.GetPayload(),
+		jwt.ExpirationTimeValidator(time.Now()),
+	)
+
+	_, err := jwt.Verify(token, alg, &data, verifyOptions)
+
+	return err
+}

router/tokens/websocket.go

@@ -1,19 +1,24 @@
-package websocket
+package tokens
 
 import (
 	"encoding/json"
 	"github.com/gbrlsnchs/jwt/v3"
 )
 
-type TokenPayload struct {
+type WebsocketPayload struct {
 	jwt.Payload
 	UserID      json.Number `json:"user_id"`
 	ServerUUID  string      `json:"server_uuid"`
 	Permissions []string    `json:"permissions"`
 }
 
+// Returns the JWT payload.
+func (p *WebsocketPayload) GetPayload() *jwt.Payload {
+	return &p.Payload
+}
+
 // Checks if the given token payload has a permission string.
-func (p *TokenPayload) HasPermission(permission string) bool {
+func (p *WebsocketPayload) HasPermission(permission string) bool {
 	for _, k := range p.Permissions {
 		if k == permission {
 			return true

router/websocket/websocket.go

@@ -7,6 +7,7 @@ import (
 	"github.com/gorilla/websocket"
 	"github.com/pkg/errors"
 	"github.com/pterodactyl/wings/config"
+	"github.com/pterodactyl/wings/router/tokens"
 	"github.com/pterodactyl/wings/server"
 	"go.uber.org/zap"
 	"net/http"
@@ -28,11 +29,26 @@ const (
 
 type Handler struct {
 	Connection *websocket.Conn
-	JWT        *TokenPayload `json:"-"`
+	JWT        *tokens.WebsocketPayload `json:"-"`
 	server     *server.Server
 	mutex      sync.Mutex
 }
 
+// Parses a JWT into a websocket token payload.
+func NewTokenPayload(token []byte) (*tokens.WebsocketPayload, error) {
+	payload := tokens.WebsocketPayload{}
+	err := tokens.ParseToken(token, &payload)
+	if err != nil {
+		return nil, err
+	}
+
+	if !payload.HasPermission(PermissionConnect) {
+		return nil, errors.New("not authorized to connect to this socket")
+	}
+
+	return &payload, nil
+}
+
 // Returns a new websocket handler using the context provided.
 func GetHandler(s *server.Server, w http.ResponseWriter, r *http.Request) (*Handler, error) {
 	upgrader := websocket.Upgrader{
@@ -56,35 +72,6 @@ func GetHandler(s *server.Server, w http.ResponseWriter, r *http.Request) (*Hand
 	}, nil
 }
 
-// Validates the provided JWT against the known secret for the Daemon and returns the
-// parsed data.
-//
-// This function DOES NOT validate that the token is valid for the connected server, nor
-// does it ensure that the user providing the token is able to actually do things.
-func ParseJWT(token []byte) (*TokenPayload, error) {
-	var payload TokenPayload
-	if alg == nil {
-		alg = jwt.NewHS256([]byte(config.Get().AuthenticationToken))
-	}
-
-	now := time.Now()
-	verifyOptions := jwt.ValidatePayload(
-		&payload.Payload,
-		jwt.ExpirationTimeValidator(now),
-	)
-
-	_, err := jwt.Verify(token, alg, &payload, verifyOptions)
-	if err != nil {
-		return nil, err
-	}
-
-	if !payload.HasPermission(PermissionConnect) {
-		return nil, errors.New("not authorized to connect to this socket")
-	}
-
-	return &payload, nil
-}
-
 func (h *Handler) SendJson(v *Message) error {
 	// Do not send JSON down the line if the JWT on the connection is not
 	// valid!
@@ -194,7 +181,7 @@ func (h *Handler) HandleInbound(m Message) error {
 	switch m.Event {
 	case AuthenticationEvent:
 		{
-			token, err := ParseJWT([]byte(strings.Join(m.Args, "")))
+			token, err := NewTokenPayload([]byte(strings.Join(m.Args, "")))
 			if err != nil {
 				return err
 			}