From b9d879330a789a52542a989b1d543d31c0f20461 Mon Sep 17 00:00:00 2001
From: Khairul Hidayat <me@khairul.my.id>
Date: Wed, 6 Nov 2024 09:50:41 +0000
Subject: [PATCH] feat: update using gofiber

---
 .../containers/interactive-session.tsx        |   5 +-
 frontend/components/containers/terminal.tsx   |  23 ++-
 server/go.mod                                 |  23 ++-
 server/go.sum                                 |  43 +++-
 server/lib/ssh-session.go                     | 162 +++++++++++++++
 server/main.go                                | 184 +++---------------
 6 files changed, 275 insertions(+), 165 deletions(-)
 create mode 100644 server/lib/ssh-session.go

diff --git a/frontend/components/containers/interactive-session.tsx b/frontend/components/containers/interactive-session.tsx
index 5312f32..f5e3c15 100644
--- a/frontend/components/containers/interactive-session.tsx
+++ b/frontend/components/containers/interactive-session.tsx
@@ -1,4 +1,3 @@
-import { View, Text } from "react-native";
 import React from "react";
 import Terminal from "./terminal";
 import { BASE_WS_URL } from "@/lib/api";
@@ -19,7 +18,9 @@ const InteractiveSession = ({ type, options }: Props) => {
         serverId: options.serverId,
         token: "token",
       });
-      return <Terminal wsUrl={BASE_WS_URL + "/ws/ssh?" + params} />;
+      return (
+        <Terminal client="xtermjs" wsUrl={`${BASE_WS_URL}/ws/ssh?${params}`} />
+      );
 
     default:
       throw new Error("Unknown interactive session type");
diff --git a/frontend/components/containers/terminal.tsx b/frontend/components/containers/terminal.tsx
index b38c19e..0defd70 100644
--- a/frontend/components/containers/terminal.tsx
+++ b/frontend/components/containers/terminal.tsx
@@ -26,20 +26,33 @@ const Keys = {
   Tab: "\x09",
 };
 
-type TerminalProps = ComponentPropsWithoutRef<typeof View> & {
+type XTermJsProps = {
+  client: "xtermjs";
   wsUrl: string;
 };
 
-const Terminal = ({ wsUrl, style, ...props }: TerminalProps) => {
-  const ref = React.useRef<XTermRef>(null);
+type TerminalProps = ComponentPropsWithoutRef<typeof View> & XTermJsProps;
+
+const Terminal = ({ client, style, ...props }: TerminalProps) => {
+  const xtermRef = React.useRef<XTermRef>(null);
 
   const send = (data: string) => {
-    ref.current?.send(data);
+    switch (client) {
+      case "xtermjs":
+        xtermRef.current?.send(data);
+        break;
+    }
   };
 
   return (
     <View style={[styles.container, style]} {...props}>
-      <XTermJs ref={ref} dom={{ scrollEnabled: false }} wsUrl={wsUrl} />
+      {client === "xtermjs" && (
+        <XTermJs
+          ref={xtermRef}
+          dom={{ scrollEnabled: false }}
+          wsUrl={props.wsUrl}
+        />
+      )}
 
       <ScrollView
         horizontal
diff --git a/server/go.mod b/server/go.mod
index 8b95463..d4db9e3 100644
--- a/server/go.mod
+++ b/server/go.mod
@@ -2,9 +2,26 @@ module rul.sh/vaulterm
 
 go 1.21.1
 
-require golang.org/x/crypto v0.28.0
+require (
+	github.com/gofiber/contrib/websocket v1.3.2
+	github.com/gofiber/fiber/v2 v2.52.5
+	golang.org/x/crypto v0.28.0
+)
 
 require (
-	github.com/coder/websocket v1.8.12
-	golang.org/x/sys v0.26.0 // indirect
+	github.com/andybalholm/brotli v1.1.0 // indirect
+	github.com/fasthttp/websocket v1.5.8 // indirect
+	github.com/google/uuid v1.6.0 // indirect
+	github.com/klauspost/compress v1.17.7 // indirect
+	github.com/mattn/go-colorable v0.1.13 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/mattn/go-runewidth v0.0.15 // indirect
+	github.com/rivo/uniseg v0.2.0 // indirect
+	github.com/savsgio/gotils v0.0.0-20240303185622-093b76447511 // indirect
+	github.com/valyala/bytebufferpool v1.0.0 // indirect
+	github.com/valyala/fasthttp v1.52.0 // indirect
+	github.com/valyala/tcplisten v1.0.0 // indirect
+	golang.org/x/net v0.23.0 // indirect
 )
+
+require golang.org/x/sys v0.26.0 // indirect
diff --git a/server/go.sum b/server/go.sum
index e2b32b6..7b4694d 100644
--- a/server/go.sum
+++ b/server/go.sum
@@ -1,8 +1,47 @@
-github.com/coder/websocket v1.8.12 h1:5bUXkEPPIbewrnkU8LTCLVaxi4N4J8ahufH2vlo4NAo=
-github.com/coder/websocket v1.8.12/go.mod h1:LNVeNrXQZfe5qhS9ALED3uA+l5pPqvwXg3CKoDBB2gs=
+github.com/andybalholm/brotli v1.1.0 h1:eLKJA0d02Lf0mVpIDgYnqXcUn0GqVmEFny3VuID1U3M=
+github.com/andybalholm/brotli v1.1.0/go.mod h1:sms7XGricyQI9K10gOSf56VKKWS4oLer58Q+mhRPtnY=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/fasthttp/websocket v1.5.8 h1:k5DpirKkftIF/w1R8ZzjSgARJrs54Je9YJK37DL/Ah8=
+github.com/fasthttp/websocket v1.5.8/go.mod h1:d08g8WaT6nnyvg9uMm8K9zMYyDjfKyj3170AtPRuVU0=
+github.com/gofiber/contrib/websocket v1.3.2 h1:AUq5PYeKwK50s0nQrnluuINYeep1c4nRCJ0NWsV3cvg=
+github.com/gofiber/contrib/websocket v1.3.2/go.mod h1:07u6QGMsvX+sx7iGNCl5xhzuUVArWwLQ3tBIH24i+S8=
+github.com/gofiber/fiber/v2 v2.52.5 h1:tWoP1MJQjGEe4GB5TUGOi7P2E0ZMMRx5ZTG4rT+yGMo=
+github.com/gofiber/fiber/v2 v2.52.5/go.mod h1:KEOE+cXMhXG0zHc9d8+E38hoX+ZN7bhOtgeF2oT6jrQ=
+github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
+github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/klauspost/compress v1.17.7 h1:ehO88t2UGzQK66LMdE8tibEd1ErmzZjNEqWkjLAKQQg=
+github.com/klauspost/compress v1.17.7/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
+github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
+github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
+github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
+github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
+github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
+github.com/mattn/go-runewidth v0.0.15 h1:UNAjwbU9l54TA3KzvqLGxwWjHmMgBUVhBiTjelZgg3U=
+github.com/mattn/go-runewidth v0.0.15/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY=
+github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
+github.com/savsgio/gotils v0.0.0-20240303185622-093b76447511 h1:KanIMPX0QdEdB4R3CiimCAbxFrhB3j7h0/OvpYGVQa8=
+github.com/savsgio/gotils v0.0.0-20240303185622-093b76447511/go.mod h1:sM7Mt7uEoCeFSCBM+qBrqvEo+/9vdmj19wzp3yzUhmg=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
+github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
+github.com/valyala/fasthttp v1.52.0 h1:wqBQpxH71XW0e2g+Og4dzQM8pk34aFYlA1Ga8db7gU0=
+github.com/valyala/fasthttp v1.52.0/go.mod h1:hf5C4QnVMkNXMspnsUlfM3WitlgYflyhHYoKol/szxQ=
+github.com/valyala/tcplisten v1.0.0 h1:rBHj/Xf+E1tRGZyWIWwJDiRY0zc1Js+CV5DqwacVSA8=
+github.com/valyala/tcplisten v1.0.0/go.mod h1:T0xQ8SeCZGxckz9qRXTfG43PvQ/mcWh7FwZEA7Ioqkc=
 golang.org/x/crypto v0.28.0 h1:GBDwsMXVQi34v5CCYUm2jkJvu4cbtru2U4TN2PSyQnw=
 golang.org/x/crypto v0.28.0/go.mod h1:rmgy+3RHxRZMyY0jjAJShp2zgEdOqj2AO7U0pYmeQ7U=
+golang.org/x/net v0.23.0 h1:7EYJ93RZ9vYSZAIb2x3lnuvqO5zneoD6IvWjuhfxjTs=
+golang.org/x/net v0.23.0/go.mod h1:JKghWKKOSdJwpW2GEx0Ja7fmaKnMsbu+MWVZTokSYmg=
+golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.26.0 h1:KHjCJyddX0LoSTb3J+vWpupP9p0oznkqVk/IfjymZbo=
 golang.org/x/sys v0.26.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/term v0.25.0 h1:WtHI/ltw4NvSUig5KARz9h521QvRC8RmF/cuYqifU24=
 golang.org/x/term v0.25.0/go.mod h1:RPyXicDX+6vLxogjjRxjgD2TKtmAO6NZBsBRfrOLu7M=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
diff --git a/server/lib/ssh-session.go b/server/lib/ssh-session.go
new file mode 100644
index 0000000..040d2f1
--- /dev/null
+++ b/server/lib/ssh-session.go
@@ -0,0 +1,162 @@
+package lib
+
+import (
+	"fmt"
+	"io"
+	"log"
+	"strconv"
+	"strings"
+
+	"github.com/gofiber/contrib/websocket"
+	"golang.org/x/crypto/ssh"
+)
+
+type SSHConfig struct {
+	HostName             string
+	User                 string
+	Password             string
+	Port                 int
+	PrivateKey           string
+	PrivateKeyPassphrase string
+}
+
+func NewSSHWebsocketSession(c *websocket.Conn, cfg *SSHConfig) error {
+	// Set up SSH client configuration
+	port := cfg.Port
+	if port == 0 {
+		port = 22
+	}
+	auth := []ssh.AuthMethod{
+		ssh.Password(cfg.Password),
+	}
+
+	if cfg.PrivateKey != "" {
+		var err error
+		var signer ssh.Signer
+
+		if cfg.PrivateKeyPassphrase != "" {
+			signer, err = ssh.ParsePrivateKeyWithPassphrase([]byte(cfg.PrivateKey), []byte(cfg.PrivateKeyPassphrase))
+		} else {
+			signer, err = ssh.ParsePrivateKey([]byte(cfg.PrivateKey))
+		}
+
+		if err != nil {
+			return fmt.Errorf("unable to parse private key: %v", err)
+		}
+		auth = append(auth, ssh.PublicKeys(signer))
+	}
+
+	sshConfig := &ssh.ClientConfig{
+		User:            cfg.User,
+		Auth:            auth,
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+	}
+
+	// Connect to SSH server
+	hostName := fmt.Sprintf("%s:%d", cfg.HostName, port)
+	sshConn, err := ssh.Dial("tcp", hostName, sshConfig)
+	if err != nil {
+		return err
+	}
+	defer sshConn.Close()
+
+	// Start an SSH shell session
+	session, err := sshConn.NewSession()
+	if err != nil {
+		return err
+	}
+	defer session.Close()
+
+	stdoutPipe, err := session.StdoutPipe()
+	if err != nil {
+		return err
+	}
+
+	stderrPipe, err := session.StderrPipe()
+	if err != nil {
+		return err
+	}
+
+	stdinPipe, err := session.StdinPipe()
+	if err != nil {
+		return err
+	}
+
+	err = session.RequestPty("xterm-256color", 80, 24, ssh.TerminalModes{})
+	if err != nil {
+		return err
+	}
+
+	if err := session.Shell(); err != nil {
+		return err
+	}
+
+	// Goroutine to send SSH stdout to WebSocket
+	go func() {
+		buf := make([]byte, 1024)
+		for {
+			n, err := stdoutPipe.Read(buf)
+			if err != nil {
+				if err != io.EOF {
+					log.Printf("error reading from SSH stdout: %v", err)
+				}
+				break
+			}
+
+			if err := c.WriteMessage(websocket.BinaryMessage, buf[:n]); err != nil {
+				log.Printf("error writing to websocket: %v", err)
+				break
+			}
+		}
+	}()
+
+	// Goroutine to handle SSH stderr
+	go func() {
+		buf := make([]byte, 1024)
+		for {
+			n, err := stderrPipe.Read(buf)
+			if err != nil {
+				if err != io.EOF {
+					log.Printf("error reading from SSH stderr: %v", err)
+				}
+				break
+			}
+			if err := c.WriteMessage(websocket.BinaryMessage, buf[:n]); err != nil {
+				log.Printf("error writing to websocket: %v", err)
+				break
+			}
+		}
+	}()
+
+	// Handle WebSocket to SSH data streaming
+	go func() {
+		defer session.Close()
+
+		for {
+			_, msg, err := c.ReadMessage()
+			if err != nil {
+				break
+			}
+
+			if strings.HasPrefix(string(msg), "\x01") {
+				parts := strings.Split(string(msg[1:]), ",")
+				if len(parts) == 2 {
+					width, _ := strconv.Atoi(parts[0])
+					height, _ := strconv.Atoi(parts[1])
+					session.WindowChange(height, width)
+				}
+			} else {
+				stdinPipe.Write(msg)
+			}
+		}
+	}()
+
+	// Wait for the SSH session to close
+	if err := session.Wait(); err != nil {
+		log.Printf("SSH session ended with error: %v", err)
+		return err
+	}
+
+	log.Println("SSH session ended normally")
+	return nil
+}
diff --git a/server/main.go b/server/main.go
index 83a6ad4..1812b0e 100644
--- a/server/main.go
+++ b/server/main.go
@@ -1,161 +1,39 @@
 package main
 
 import (
-	"context"
-	"io"
-	"log"
-	"net/http"
-	"strconv"
-	"strings"
+	"fmt"
 
-	"github.com/coder/websocket"
-	"golang.org/x/crypto/ssh"
+	"github.com/gofiber/contrib/websocket"
+	"github.com/gofiber/fiber/v2"
+	"rul.sh/vaulterm/lib"
 )
 
-// Replace with actual SSH server credentials
-var sshHost = "10.0.0.102"
-var sshUser = "root"
-var sshPassword = "ausya2"
-
-func sshWebSocketHandler(w http.ResponseWriter, r *http.Request) {
-	wsConn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
-		OriginPatterns: []string{"*"}, // Adjust origin policy as needed
-	})
-	if err != nil {
-		log.Printf("failed to accept websocket: %v", err)
-		return
-	}
-
-	// Set up SSH client configuration
-	sshConfig := &ssh.ClientConfig{
-		User: sshUser,
-		Auth: []ssh.AuthMethod{
-			ssh.Password(sshPassword),
-		},
-		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
-	}
-
-	// Connect to SSH server
-	sshConn, err := ssh.Dial("tcp", sshHost+":22", sshConfig)
-	if err != nil {
-		wsConn.Write(context.Background(), websocket.MessageText, []byte(err.Error()))
-		wsConn.Close(websocket.StatusInternalError, "failed to connect to SSH")
-		return
-	}
-	defer sshConn.Close()
-
-	// Start an SSH shell session
-	session, err := sshConn.NewSession()
-	if err != nil {
-		wsConn.Write(context.Background(), websocket.MessageText, []byte(err.Error()))
-		wsConn.Close(websocket.StatusInternalError, "failed to start SSH session")
-		return
-	}
-	defer session.Close()
-
-	stdoutPipe, err := session.StdoutPipe()
-	if err != nil {
-		wsConn.Close(websocket.StatusInternalError, "failed to get stdout pipe")
-		return
-	}
-
-	stderrPipe, err := session.StderrPipe()
-	if err != nil {
-		wsConn.Close(websocket.StatusInternalError, "failed to get stderr pipe")
-		return
-	}
-
-	stdinPipe, err := session.StdinPipe()
-	if err != nil {
-		wsConn.Close(websocket.StatusInternalError, "failed to get stdin pipe")
-		return
-	}
-
-	err = session.RequestPty("xterm-256color", 80, 24, ssh.TerminalModes{})
-	if err != nil {
-		wsConn.Close(websocket.StatusInternalError, "failed to request pty")
-		return
-	}
-
-	if err := session.Shell(); err != nil {
-		wsConn.Close(websocket.StatusInternalError, "failed to start shell")
-		return
-	}
-
-	// Goroutine to send SSH stdout to WebSocket
-	go func() {
-		buf := make([]byte, 1024)
-		for {
-			n, err := stdoutPipe.Read(buf)
-			if err != nil {
-				if err != io.EOF {
-					log.Printf("error reading from SSH stdout: %v", err)
-				}
-				break
-			}
-
-			if writeErr := wsConn.Write(context.Background(), websocket.MessageBinary, buf[:n]); writeErr != nil {
-				log.Printf("error writing to websocket: %v", writeErr)
-				break
-			}
-		}
-	}()
-
-	// Goroutine to handle SSH stderr
-	go func() {
-		buf := make([]byte, 1024)
-		for {
-			n, err := stderrPipe.Read(buf)
-			if err != nil {
-				if err != io.EOF {
-					log.Printf("error reading from SSH stderr: %v", err)
-				}
-				break
-			}
-			if writeErr := wsConn.Write(context.Background(), websocket.MessageBinary, buf[:n]); writeErr != nil {
-				log.Printf("error writing to websocket: %v", writeErr)
-				break
-			}
-		}
-	}()
-
-	// Handle WebSocket to SSH data streaming
-	go func() {
-		defer session.Close()
-
-		for {
-			_, msg, err := wsConn.Read(context.Background())
-			if err != nil {
-				break
-			}
-
-			if strings.HasPrefix(string(msg), "\x01") {
-				parts := strings.Split(string(msg[1:]), ",")
-				if len(parts) == 2 {
-					width, _ := strconv.Atoi(parts[0])
-					height, _ := strconv.Atoi(parts[1])
-					session.WindowChange(height, width)
-				}
-			} else {
-				stdinPipe.Write(msg)
-			}
-		}
-	}()
-
-	// Wait for the SSH session to close
-	if err := session.Wait(); err != nil {
-		log.Printf("SSH session ended with error: %v", err)
-		wsConn.Write(context.Background(), websocket.MessageText, []byte(err.Error()))
-	} else {
-		log.Println("SSH session ended normally")
-	}
-
-	// Ensure WebSocket is closed after SSH logout
-	wsConn.Close(websocket.StatusNormalClosure, "SSH session closed")
-}
-
 func main() {
-	http.HandleFunc("/ws/ssh", sshWebSocketHandler)
-	log.Println("Server started on :3000")
-	log.Fatal(http.ListenAndServe(":3000", nil))
+	app := fiber.New()
+
+	app.Get("/", func(c *fiber.Ctx) error {
+		return c.SendString("Hello, World!")
+	})
+
+	app.Use("/ws", func(c *fiber.Ctx) error {
+		if websocket.IsWebSocketUpgrade(c) {
+			return c.Next()
+		}
+		return fiber.ErrUpgradeRequired
+	})
+
+	app.Get("/ws/ssh", websocket.New(func(c *websocket.Conn) {
+		err := lib.NewSSHWebsocketSession(c, &lib.SSHConfig{
+			HostName: "10.0.0.102",
+			User:     "root",
+			Password: "ausya2",
+		})
+
+		if err != nil {
+			msg := fmt.Sprintf("\r\n%s\r\n", err.Error())
+			c.WriteMessage(websocket.TextMessage, []byte(msg))
+		}
+	}))
+
+	app.Listen(":3000")
 }