From a16c592dc819a3a4a94211c778971eae49aeb10a Mon Sep 17 00:00:00 2001
From: Boleyn Su <boleyn.su@gmail.com>
Date: Sat, 13 Jun 2020 20:07:23 +0900
Subject: [PATCH] This commit implements Feature Request #59

Now we can use the following command to connect to an SSH server through a chisel tunnel.
```bash
ssh -o ProxyCommand='/path/to/chisel client https://chisel.boleyn.su stdio:%h:%p' me@boleyn.su
```
---
 client/client.go |  8 ++++++++
 main.go          |  8 ++++++++
 share/logger.go  |  2 +-
 share/proxy.go   | 22 ++++++++++++++++++----
 share/remote.go  | 24 ++++++++++++++++++++++--
 share/stdio.go   | 23 +++++++++++++++++++++++
 6 files changed, 80 insertions(+), 7 deletions(-)
 create mode 100644 share/stdio.go

diff --git a/client/client.go b/client/client.go
index 3d770136..10a4f5ac 100644
--- a/client/client.go
+++ b/client/client.go
@@ -2,6 +2,7 @@ package chclient
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"net"
@@ -68,11 +69,18 @@ func NewClient(config *Config) (*Client, error) {
 	//swap to websockets scheme
 	u.Scheme = strings.Replace(u.Scheme, "http", "ws", 1)
 	shared := &chshare.Config{}
+	stdioCnt := 0
 	for _, s := range config.Remotes {
 		r, err := chshare.DecodeRemote(s)
 		if err != nil {
 			return nil, fmt.Errorf("Failed to decode remote '%s': %s", s, err)
 		}
+		if r.Stdio {
+			stdioCnt++
+		}
+		if stdioCnt > 1 {
+			return nil, errors.New("Only one stdio is allowed")
+		}
 		shared.Remotes = append(shared.Remotes, r)
 	}
 	config.shared = shared
diff --git a/main.go b/main.go
index 570b508e..b3a73b45 100644
--- a/main.go
+++ b/main.go
@@ -224,6 +224,7 @@ var clientHelp = `
       socks
       5000:socks
       R:2222:localhost:22
+      stdio:example.com:22
 
     When the chisel server has --socks5 enabled, remotes can
     specify "socks" in place of remote-host and remote-port.
@@ -236,6 +237,13 @@ var clientHelp = `
     is, the server will listen and accept connections, and they
     will be proxied through the client which specified the remote.
 
+    When stdio is used as local-host, the tunnel will connect standard
+    input/output of this program with the remote. This is useful when 
+    combined with ssh ProxyCommand. You can use
+      ssh -o ProxyCommand='chisel client chiselserver stdio:%h:%p' \
+          user@example.com
+    to connect to an SSH server through the tunnel.
+
   Options:
 
     --fingerprint, A *strongly recommended* fingerprint string
diff --git a/share/logger.go b/share/logger.go
index 1cd5223a..cd7eccc7 100644
--- a/share/logger.go
+++ b/share/logger.go
@@ -20,7 +20,7 @@ func NewLogger(prefix string) *Logger {
 func NewLoggerFlag(prefix string, flag int) *Logger {
 	l := &Logger{
 		prefix: prefix,
-		logger: log.New(os.Stdout, "", flag),
+		logger: log.New(os.Stderr, "", flag),
 		Info:   false,
 		Debug:  false,
 	}
diff --git a/share/proxy.go b/share/proxy.go
index 862a5d76..080caa6b 100644
--- a/share/proxy.go
+++ b/share/proxy.go
@@ -31,11 +31,25 @@ func NewTCPProxy(logger *Logger, ssh GetSSHConn, index int, remote *Remote) *TCP
 }
 
 func (p *TCPProxy) Start(ctx context.Context) error {
-	l, err := net.Listen("tcp4", p.remote.LocalHost+":"+p.remote.LocalPort)
-	if err != nil {
-		return fmt.Errorf("%s: %s", p.Logger.Prefix(), err)
+	if p.remote.Stdio {
+		go func() {
+			for {
+				p.accept(Stdio)
+				select {
+				case <-ctx.Done():
+					return
+				default:
+					// the connection is not ready yet
+				}
+			}
+		}()
+	} else {
+		l, err := net.Listen("tcp4", p.remote.LocalHost+":"+p.remote.LocalPort)
+		if err != nil {
+			return fmt.Errorf("%s: %s", p.Logger.Prefix(), err)
+		}
+		go p.listen(ctx, l)
 	}
-	go p.listen(ctx, l)
 	return nil
 }
 
diff --git a/share/remote.go b/share/remote.go
index 3aad30e1..e5de0e61 100644
--- a/share/remote.go
+++ b/share/remote.go
@@ -20,10 +20,16 @@ import (
 //   192.168.0.1:3000:google.com:80 ->
 //     local  192.168.0.1:3000
 //     remote google.com:80
+//   127.0.0.1:1080:socks
+//     local  127.0.0.1:1080
+//     remote socks
+//   stdio:example.com:22
+//     local  stdio
+//     remote example.com:22
 
 type Remote struct {
 	LocalHost, LocalPort, RemoteHost, RemotePort string
-	Socks, Reverse                               bool
+	Socks, Reverse, Stdio                        bool
 }
 
 const revPrefix = "R:"
@@ -51,6 +57,10 @@ func DecodeRemote(s string) (*Remote, error) {
 			r.Socks = true
 			continue
 		}
+		if i == 0 && p == "stdio" {
+			r.Stdio = true
+			continue
+		}
 		if isPort(p) {
 			if !r.Socks && r.RemotePort == "" {
 				r.RemotePort = p
@@ -85,6 +95,9 @@ func DecodeRemote(s string) (*Remote, error) {
 	if !r.Socks && r.RemoteHost == "" {
 		r.RemoteHost = "0.0.0.0"
 	}
+	if r.Stdio && r.Reverse {
+		return nil, errors.New("stdio cannot be used with reverse tunnel")
+	}
 	return r, nil
 }
 
@@ -111,7 +124,14 @@ func (r *Remote) String() string {
 	if r.Reverse {
 		tag = revPrefix
 	}
-	return tag + r.LocalHost + ":" + r.LocalPort + "=>" + r.Remote()
+	return tag + r.Local() + "=>" + r.Remote()
+}
+
+func (r *Remote) Local() string {
+	if r.Stdio {
+		return "stdio"
+	}
+	return r.LocalHost + ":" + r.LocalPort
 }
 
 func (r *Remote) Remote() string {
diff --git a/share/stdio.go b/share/stdio.go
new file mode 100644
index 00000000..976a147d
--- /dev/null
+++ b/share/stdio.go
@@ -0,0 +1,23 @@
+package chshare
+
+import (
+	"os"
+)
+
+
+type stdio struct {}
+
+var Stdio = &stdio{}
+
+func (t *stdio) Read(b []byte) (n int, err error) {
+	return os.Stdin.Read(b)
+}
+
+func (t *stdio) Write(b []byte) (n int, err error) {
+	return os.Stdout.Write(b)
+}
+
+// We do not close stdio
+func (t *stdio) Close() error {
+	return nil
+}