From f654deb300ddc118fece9e4896e21a01763b2caf Mon Sep 17 00:00:00 2001
From: Neur0toxine <pashok9825@gmail.com>
Date: Sun, 19 Nov 2023 14:12:39 +0300
Subject: [PATCH] ssh: shell start support

---
 internal/server/driver/ssh/driver.go         |  1 +
 internal/server/driver/ssh/params.go         |  1 +
 internal/server/driver/ssh/sshtun/connect.go | 15 +++++++++++----
 3 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/internal/server/driver/ssh/driver.go b/internal/server/driver/ssh/driver.go
index b2c1de1..39a6aa2 100644
--- a/internal/server/driver/ssh/driver.go
+++ b/internal/server/driver/ssh/driver.go
@@ -62,6 +62,7 @@ func (d *SSH) forward(val sshtun.Forward, domainMatcher func(string)) conn {
 		d.auth,
 		sshtun.SessionConfig{
 			NoPTY:             d.params.NoPTY,
+			Shell:             d.params.Shell,
 			FakeRemoteHost:    d.params.FakeRemoteHost,
 			KeepAliveInterval: uint(d.params.KeepAlive.Interval),
 			KeepAliveMax:      uint(d.params.KeepAlive.MaxAttempts),
diff --git a/internal/server/driver/ssh/params.go b/internal/server/driver/ssh/params.go
index 1403b09..f8cee14 100644
--- a/internal/server/driver/ssh/params.go
+++ b/internal/server/driver/ssh/params.go
@@ -17,6 +17,7 @@ type Params struct {
 	Mode               types.DomainMode `mapstructure:"mode" validate:"required,oneof=single multi"`
 	FakeRemoteHost     bool             `mapstructure:"fake_remote_host"`
 	NoPTY              bool             `mapstructure:"nopty"`
+	Shell              bool             `mapstructure:"shell"`
 	Commands           types.Commands   `mapstructure:"commands"`
 }
 
diff --git a/internal/server/driver/ssh/sshtun/connect.go b/internal/server/driver/ssh/sshtun/connect.go
index fdad0af..97e51bf 100644
--- a/internal/server/driver/ssh/sshtun/connect.go
+++ b/internal/server/driver/ssh/sshtun/connect.go
@@ -28,6 +28,7 @@ type Tunnel struct {
 
 type SessionConfig struct {
 	NoPTY             bool
+	Shell             bool
 	FakeRemoteHost    bool
 	KeepAliveInterval uint
 	KeepAliveMax      uint
@@ -112,12 +113,18 @@ func (t *Tunnel) connect(ctx context.Context, bannerCb ssh.BannerCallback, sessi
 		})
 		if err != nil {
 			t.log.Warnf("PTY allocation failed: %s", err.Error())
-		} else {
-			if err := sess.Shell(); err != nil {
-				t.log.Warnf("failed to start shell: %s", err.Error())
-			}
 		}
 	}
+	if t.sessConfig.Shell {
+		if err := sess.Shell(); err != nil {
+			t.log.Warnf("failed to start shell: %s", err.Error())
+		}
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			_ = sess.Wait()
+		}()
+	}
 
 	wg.Add(1)
 	go func() {