ssh: refactor web session package adding websocket and HTTP layers

ssh/main.go

@@ -66,17 +66,11 @@ func main() {
 		}
 	})
 
-	// TODO: add this route to OpenAPI.
-	// `/ws/ssh` path managers a web terminal connection.
-	// Connects to the web terminal session through the token.
-	//
-	// Query parameters:
-	// - token: the session token.
-	// - cols: web terminal columns.
-	// - rows: web terminal rows.
-	router.Handle("/ws/ssh", web.RestoreSession(handler.WebSession)).Methods(http.MethodGet)
-	// Creates a new web terminal session token.
-	router.HandleFunc("/ws/ssh", web.NewSession).Methods(http.MethodPost)
+	// TODO: add `/ws/ssh` route to OpenAPI repository.
+	router.Handle("/ws/ssh", web.HandlerRestoreSession(web.RestoreSession, handler.WebSession)).
+		Methods(http.MethodGet)
+	router.HandleFunc("/ws/ssh", web.HandlerCreateSession(web.CreateSession)).
+		Methods(http.MethodPost)
 
 	go http.ListenAndServe(":8080", router) // nolint:errcheck
 

ssh/server/handler/web.go

@@ -15,6 +15,7 @@ import (
 	"github.com/shellhub-io/shellhub/ssh/pkg/flow"
 	"github.com/shellhub-io/shellhub/ssh/pkg/magickey"
 	"github.com/shellhub-io/shellhub/ssh/pkg/target"
+	"github.com/shellhub-io/shellhub/ssh/web"
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/net/websocket"
@@ -39,9 +40,7 @@ type WebData struct {
 
 // NewWebData create a new WebData.
 // WebData contains the data required by web termianl connection.
-func NewWebData(socket *websocket.Conn, device, username, password, fingerprint, signature string) (*WebData, error) {
-	// ctx := socket.Request().Context()
-
+func NewWebData(socket *websocket.Conn, input *web.Session) (*WebData, error) {
 	get := func(socket *websocket.Conn, key string) (string, bool) {
 		value := socket.Request().URL.Query().Get(key)
 
@@ -73,13 +72,13 @@ func NewWebData(socket *websocket.Conn, device, username, password, fingerprint,
 		return nil, errors.New("rows field is invalid or missing")
 	}
 
-	target := username + "@" + device
+	target := input.Username + "@" + input.Device
 
 	return &WebData{
 		User:        target,
-		Password:    password,
-		Fingerprint: fingerprint,
-		Signature:   signature,
+		Password:    input.Password,
+		Fingerprint: input.Fingerprint,
+		Signature:   input.Signature,
 		Columns:     columns,
 		Rows:        rows,
 	}, nil
@@ -156,11 +155,11 @@ func (c *WebData) GetAuth(magicKey *rsa.PrivateKey) ([]ssh.AuthMethod, error) {
 }
 
 // WebSession is the Client's handler for connection coming from the web terminal.
-func WebSession(socket *websocket.Conn, device, username, password, fingerprint, signature string) {
+func WebSession(socket *websocket.Conn, input *web.Session) {
 	log.Info("handling web client request started")
 	defer log.Info("handling web client request end")
 
-	data, err := NewWebData(socket, device, username, password, fingerprint, signature)
+	data, err := NewWebData(socket, input)
 	if err != nil {
 		sendAndInformError(socket, err, ErrWebData)
 	}

ssh/web/handlers.go

@@ -0,0 +1,118 @@
+package web
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net/http"
+
+	log "github.com/sirupsen/logrus"
+	"golang.org/x/net/websocket"
+)
+
+type (
+	functionHandleCreateSession  func(ctx context.Context, data *Input) (*Session, error)
+	functionHandleRestoreSession func(ctx context.Context, data *Output) (*Session, error)
+)
+
+// HandlerCreateSession handles a HTTP request with the data to create a new web session.
+//
+// It receives on request's body the device's UID and the device's username, either the device's password or the
+// device's fingerprint and the device's signature, to returns a JWT token that can be used to connect to the device.
+// The JWT token is generated using a UUID as payload, and encrypted using a runtime generated RSA private key.
+//
+// If a error occurs, it logs on the server the error and returns the error message and the HTTP status code related to
+// the error to the user.
+func HandlerCreateSession(create functionHandleCreateSession) func(http.ResponseWriter, *http.Request) {
+	type Request struct {
+		Device      string `json:"device"`
+		Username    string `json:"username"`
+		Password    string `json:"password"`
+		Fingerprint string `json:"fingerprint"`
+		Signature   string `json:"signature"`
+	}
+
+	type Response struct {
+		Token string `json:"token"`
+	}
+
+	success := func(req http.ResponseWriter, device, username, token string) {
+		log.WithFields(log.Fields{
+			"device":   device,
+			"username": username,
+			"token":    token,
+		}).Info("session's token generated successfully")
+
+		req.WriteHeader(http.StatusOK)
+		req.Header().Set("Content-Type", "application/json")
+
+		json.NewEncoder(req).Encode(Response{Token: token}) //nolint: errcheck,errchkjson
+	}
+
+	fail := func(response http.ResponseWriter, device, username string, status int, err error) {
+		log.WithError(err).WithFields(log.Fields{
+			"device":   device,
+			"username": username,
+			"status":   status,
+		}).Error("failed to get the session's token")
+
+		http.Error(response, err.Error(), status)
+	}
+
+	return func(res http.ResponseWriter, req *http.Request) {
+		var request *Request
+		if err := json.NewDecoder(req.Body).Decode(&request); err != nil {
+			fail(res, "", "", http.StatusBadRequest, errors.New("failed to decode the request body"))
+		}
+
+		data := &Input{
+			Device:      request.Device,
+			Username:    request.Username,
+			Password:    request.Password,
+			Fingerprint: request.Fingerprint,
+			Signature:   request.Signature,
+		}
+
+		session, err := create(req.Context(), data)
+		if err != nil {
+			fail(res, data.Device, data.Username, http.StatusInternalServerError, errors.New("failed to generate the session's token"))
+		}
+
+		success(res, session.Device, session.Username, session.Token)
+	}
+}
+
+// HandlerCreateSession handles a websocket request with the data to restore web session.
+//
+// It receives the session's token as a websocket's query parameter and verifies if the token is valid. If the token is
+// valid, it calls the websocket handler to connect to the device. If the token is invalid, it returns an error message.
+//
+// If any other error occurs, it logs on the server the error and returns the error message and the error to the user.
+func HandlerRestoreSession(restore functionHandleRestoreSession, handler func(socket *websocket.Conn, session *Session)) websocket.Handler {
+	return func(socket *websocket.Conn) {
+		get := func(socket *websocket.Conn, key string) (string, bool) {
+			value := socket.Request().URL.Query().Get(key)
+
+			return value, value != ""
+		}
+
+		fail := func(socket *websocket.Conn, internal, external error) {
+			log.Error(internal.Error())
+
+			socket.Write([]byte(fmt.Sprintf("%s\n", external.Error()))) //nolint: errcheck
+		}
+
+		token, ok := get(socket, "token")
+		if !ok {
+			fail(socket, errors.New("failed to get the token from the websocket"), errors.New("failed to get the token from the websocket"))
+		}
+
+		session, err := restore(socket.Request().Context(), &Output{Token: token})
+		if err != nil {
+			fail(socket, err, errors.New("failed to get the session"))
+		}
+
+		handler(socket, session)
+	}
+}

ssh/web/pkg/cache/cache.go

@@ -10,7 +10,7 @@ import (
 
 var instance cache.Cache
 
-func getInstance() (cache.Cache, error) {
+func getInstance() (cache.Cache, error) { //nolint: ireturn
 	if instance == nil {
 		instance, err := cache.NewRedisCache("redis://redis:6379")
 
@@ -20,10 +20,10 @@ func getInstance() (cache.Cache, error) {
 	return instance, nil
 }
 
-// CacheTokenTTL is the time to live of the token in the cache.
-const CacheTokenTTL = time.Second * 30
+// TTL is the time to live of the token in the cache.
+const TTL = time.Second * 30
 
-type CachedToken struct {
+type Token struct {
 	Token       string
 	ID          string
 	Device      string
@@ -34,25 +34,35 @@ type CachedToken struct {
 	Data        interface{}
 }
 
-func Token(ctx context.Context, token *token.Token, data interface{}) (*CachedToken, error) {
-	cache, err := getInstance()
+type Data struct {
+	Device      string
+	Username    string
+	Password    string
+	Fingerprint string
+	Signature   string
+}
+
+// Save saves a data set for TTL time using token as identifier.
+func Save(ctx context.Context, token *token.Token, data *Data) (*Token, error) {
+	cache, err := getInstance() //nolint: contextcheck
 	if err != nil {
 		return nil, err
 	}
 
-	if err := cache.Set(ctx, token.ID, data, CacheTokenTTL); err != nil {
+	if err := cache.Set(ctx, token.ID, data, TTL); err != nil {
 		return nil, err
 	}
 
-	return &CachedToken{
+	return &Token{ //nolint: exhaustruct
 		Token: token.Token,
 		ID:    token.ID,
 		Data:  data,
 	}, nil
 }
 
-func Restore(ctx context.Context, token *token.Token) (*CachedToken, error) {
-	cache, err := getInstance()
+// Restore restores a data set using token as identifier.
+func Restore(ctx context.Context, token *token.Token) (*Token, error) {
+	cache, err := getInstance() //nolint: contextcheck
 	if err != nil {
 		return nil, err
 	}
@@ -69,7 +79,7 @@ func Restore(ctx context.Context, token *token.Token) (*CachedToken, error) {
 		return nil, err
 	}
 
-	return &CachedToken{
+	return &Token{ //nolint: exhaustruct
 		ID:          token.ID,
 		Device:      value.Device,
 		Username:    value.Username,