实现shell数据转发

.github/workflows/build.yml

@@ -39,4 +39,4 @@ jobs:
       run: go build -v code/server/main.go
 
     - name: Build Client
-      run: go build -v code/client/main.go
+      run: go build -v code/client/main.go code/client/connect.go

.gitignore

@@ -3,3 +3,6 @@
 /build
 /release
 /tmp
+/logs
+/run
+/*.yaml

code/client/connect.go

@@ -0,0 +1,59 @@
+package main
+
+import (
+	"fmt"
+	"natpass/code/client/global"
+	"natpass/code/client/pool"
+	"natpass/code/client/shell"
+	"natpass/code/client/tunnel"
+	"natpass/code/network"
+	"net"
+	"strconv"
+
+	"github.com/lwch/logging"
+)
+
+func connect(conn *pool.Conn, msg *network.Msg) {
+	req := msg.GetCreq()
+	dial := "tcp"
+	if req.GetXType() == network.ConnectRequest_udp {
+		dial = "udp"
+	}
+	link, err := net.Dial(dial, fmt.Sprintf("%s:%d", req.GetAddr(), req.GetPort()))
+	if err != nil {
+		conn.SendConnectError(msg.GetFrom(), msg.GetFromIdx(), msg.GetLinkId(), err.Error())
+		return
+	}
+	host, pt, _ := net.SplitHostPort(link.LocalAddr().String())
+	port, _ := strconv.ParseUint(pt, 10, 16)
+	tn := tunnel.New(global.Tunnel{
+		Name:       req.GetName(),
+		Target:     msg.GetFrom(),
+		Type:       dial,
+		LocalAddr:  host,
+		LocalPort:  uint16(port),
+		RemoteAddr: req.GetAddr(),
+		RemotePort: uint16(req.GetPort()),
+	})
+	lk := tunnel.NewLink(tn, msg.GetLinkId(), msg.GetFrom(), link, conn)
+	lk.SetTargetIdx(msg.GetFromIdx())
+	conn.SendConnectOK(msg.GetFrom(), msg.GetFromIdx(), msg.GetLinkId())
+	lk.Forward()
+	lk.OnWork <- struct{}{}
+}
+
+func shellCreate(conn *pool.Conn, msg *network.Msg) {
+	create := msg.GetScreate()
+	sh := shell.New(global.Tunnel{
+		Name:   create.GetName(),
+		Target: msg.GetFrom(),
+		Type:   "shell",
+		Exec:   create.GetExec(),
+	})
+	err := sh.Exec(msg.GetLinkId())
+	if err != nil {
+		logging.Error("create shell failed: %v", err)
+		return
+	}
+	sh.Forward(conn, msg.GetFromIdx())
+}

code/client/main.go

@@ -8,11 +8,9 @@ import (
 	"natpass/code/client/shell"
 	"natpass/code/client/tunnel"
 	"natpass/code/network"
-	"net"
 	"os"
 	"path/filepath"
 	rt "runtime"
-	"strconv"
 	"time"
 
 	_ "net/http/pprof"
@@ -83,9 +81,9 @@ func (a *app) run() {
 					var linkID string
 					switch msg.GetXType() {
 					case network.Msg_connect_req:
-						connect(pl, conn, msg.GetLinkId(), msg.GetFrom(), msg.GetTo(),
-							msg.GetFromIdx(), msg.GetToIdx(), msg.GetCreq())
+						connect(conn, msg)
 					case network.Msg_shell_create:
+						shellCreate(conn, msg)
 					case network.Msg_connect_rep,
 						network.Msg_disconnect,
 						network.Msg_forward:
@@ -159,31 +157,3 @@ func main() {
 		runtime.Assert(sv.Run())
 	}
 }
-
-func connect(pool *pool.Pool, conn *pool.Conn, id, from, to string, fromIdx, toIdx uint32, req *network.ConnectRequest) {
-	dial := "tcp"
-	if req.GetXType() == network.ConnectRequest_udp {
-		dial = "udp"
-	}
-	link, err := net.Dial(dial, fmt.Sprintf("%s:%d", req.GetAddr(), req.GetPort()))
-	if err != nil {
-		conn.SendConnectError(from, fromIdx, id, err.Error())
-		return
-	}
-	host, pt, _ := net.SplitHostPort(link.LocalAddr().String())
-	port, _ := strconv.ParseUint(pt, 10, 16)
-	tn := tunnel.New(global.Tunnel{
-		Name:       req.GetName(),
-		Target:     from,
-		Type:       dial,
-		LocalAddr:  host,
-		LocalPort:  uint16(port),
-		RemoteAddr: req.GetAddr(),
-		RemotePort: uint16(req.GetPort()),
-	})
-	lk := tunnel.NewLink(tn, id, from, link, conn)
-	lk.SetTargetIdx(fromIdx)
-	conn.SendConnectOK(from, fromIdx, id)
-	lk.Forward()
-	lk.OnWork <- struct{}{}
-}

code/client/pool/send.go

@@ -131,3 +131,26 @@ func (conn *Conn) SendShellCreate(id string, cfg global.Tunnel) {
 	case <-time.After(conn.parent.cfg.WriteTimeout):
 	}
 }
+
+// SendData send shell data
+func (conn *Conn) SendShellData(to string, toIdx uint32, id string, data []byte) {
+	dup := func(data []byte) []byte {
+		ret := make([]byte, len(data))
+		copy(ret, data)
+		return ret
+	}
+	var msg network.Msg
+	msg.To = to
+	msg.ToIdx = toIdx
+	msg.XType = network.Msg_shell_data
+	msg.LinkId = id
+	msg.Payload = &network.Msg_Sdata{
+		Sdata: &network.ShellData{
+			Data: dup(data),
+		},
+	}
+	select {
+	case conn.write <- &msg:
+	case <-time.After(conn.parent.cfg.WriteTimeout):
+	}
+}

code/client/shell/exec_windows.go

@@ -0,0 +1,5 @@
+package shell
+
+func (shell *Shell) Exec(id string) {
+	shell.id = id
+}

code/client/shell/exec_xx.go

@@ -0,0 +1,47 @@
+// +build !windows
+
+package shell
+
+import (
+	"errors"
+	"os/exec"
+
+	"github.com/creack/pty"
+)
+
+func (shell *Shell) Exec(id string) error {
+	shell.id = id
+	var cmd *exec.Cmd
+	if len(shell.cfg.Exec) > 0 {
+		cmd = exec.Command(shell.cfg.Exec)
+	}
+	if cmd == nil {
+		dir, err := exec.LookPath("bash")
+		if err == nil {
+			cmd = exec.Command(dir)
+		}
+	}
+	if cmd == nil {
+		dir, err := exec.LookPath("sh")
+		if err == nil {
+			cmd = exec.Command(dir)
+		}
+	}
+	if cmd == nil {
+		return errors.New("no shell command supported")
+	}
+	f, err := pty.Start(cmd)
+	if err != nil {
+		return err
+	}
+	shell.stdin = f
+	shell.stdout = f
+	shell.pid = cmd.Process.Pid
+	return nil
+}
+
+func (shell *Shell) onClose() {
+	if shell.stdin != nil {
+		shell.stdin.Close()
+	}
+}

code/client/shell/forward.go

@@ -0,0 +1,55 @@
+package shell
+
+import (
+	"natpass/code/client/pool"
+	"natpass/code/network"
+	"natpass/code/utils"
+
+	"github.com/lwch/logging"
+)
+
+func (shell *Shell) Forward(remote *pool.Conn, toIdx uint32) {
+	go shell.remoteRead(remote)
+	go shell.localRead(remote, toIdx)
+}
+
+func (shell *Shell) remoteRead(remote *pool.Conn) {
+	defer utils.Recover("remoteRead")
+	defer shell.Close()
+	ch := remote.ChanRead(shell.id)
+	for {
+		msg := <-ch
+		if msg == nil {
+			return
+		}
+		switch msg.GetXType() {
+		case network.Msg_shell_resize:
+			// TODO
+		case network.Msg_shell_data:
+			// TODO
+		case network.Msg_shell_close:
+			// TODO
+		}
+	}
+}
+
+func (shell *Shell) localRead(remote *pool.Conn, toIdx uint32) {
+	defer utils.Recover("localRead")
+	defer shell.Close()
+	buf := make([]byte, 16*1024)
+	for {
+		n, err := shell.stdout.Read(buf)
+		if err != nil {
+			// if !link.closeFromRemote {
+			// 	link.remote.SendDisconnect(link.target, link.targetIdx, link.id)
+			// }
+			logging.Error("read data on shell %s link %s failed, err=%v", shell.Name, shell.id, err)
+			return
+		}
+		if n == 0 {
+			continue
+		}
+		logging.Debug("link %s on shell %s read from local %d bytes", shell.id, shell.Name, n)
+		remote.SendShellData(shell.cfg.Target, toIdx, shell.id, buf[:n])
+	}
+}

code/client/shell/h_new.go

@@ -1,24 +0,0 @@
-package shell
-
-import (
-	"fmt"
-	"natpass/code/client/pool"
-	"net/http"
-
-	"github.com/lwch/logging"
-	"github.com/lwch/runtime"
-)
-
-// New create shell
-func (shell *Shell) New(pool *pool.Pool, w http.ResponseWriter, r *http.Request) {
-	id, err := runtime.UUID(16, "0123456789abcdef")
-	if err != nil {
-		logging.Error("failed to generate link_id for shell: %s, err=%v",
-			shell.Name, err)
-		http.Error(w, err.Error(), http.StatusInternalServerError)
-		return
-	}
-	conn := pool.Get(id)
-	conn.SendShellCreate(shell.cfg.Target, shell.cfg)
-	fmt.Fprint(w, id)
-}

code/client/shell/http.go

@@ -0,0 +1,82 @@
+package shell
+
+import (
+	"natpass/code/client/pool"
+	"natpass/code/network"
+	"natpass/code/utils"
+	"net/http"
+	"sync"
+
+	"github.com/gorilla/websocket"
+	"github.com/lwch/logging"
+	"github.com/lwch/runtime"
+)
+
+var upgrader = websocket.Upgrader{}
+
+// WS websocket for shell
+func (shell *Shell) WS(pool *pool.Pool, w http.ResponseWriter, r *http.Request) {
+	id, err := runtime.UUID(16, "0123456789abcdef")
+	if err != nil {
+		logging.Error("failed to generate link_id for shell: %s, err=%v",
+			shell.Name, err)
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+	conn := pool.Get(id)
+	conn.SendShellCreate(id, shell.cfg)
+	local, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		logging.Error("upgrade websocket failed: %s, err=%v", shell.Name, err)
+		http.Error(w, err.Error(), http.StatusServiceUnavailable)
+		return
+	}
+	defer local.Close()
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+		shell.localForward(id, local, conn)
+	}()
+	go func() {
+		defer wg.Done()
+		shell.remoteForward(id, conn.ChanRead(id), local)
+	}()
+	wg.Wait()
+}
+
+func (shell *Shell) localForward(id string, local *websocket.Conn, remote *pool.Conn) {
+	defer utils.Recover("localForward")
+	defer local.Close()
+	for {
+		_, data, err := local.ReadMessage()
+		if err != nil {
+			// TODO: close
+			logging.Error("read local data for %s failed: %v", shell.Name, err)
+			return
+		}
+		remote.SendShellData(shell.cfg.Target, remote.Idx, id, data)
+	}
+}
+
+func (shell *Shell) remoteForward(id string, ch <-chan *network.Msg, local *websocket.Conn) {
+	defer utils.Recover("remoteForward")
+	defer local.Close()
+	for {
+		msg := <-ch
+		if msg == nil {
+			return
+		}
+		switch msg.GetXType() {
+		case network.Msg_shell_data:
+			err := local.WriteMessage(websocket.TextMessage, msg.GetSdata().GetData())
+			if err != nil {
+				logging.Error("write data for %s failed: %v", shell.Name, err)
+				return
+			}
+			// TODO: other
+		}
+	}
+}

code/client/shell/shell.go

@@ -2,9 +2,11 @@ package shell
 
 import (
 	"fmt"
+	"io"
 	"natpass/code/client/global"
 	"natpass/code/client/pool"
 	"net/http"
+	"os"
 
 	"github.com/lwch/logging"
 	"github.com/lwch/runtime"
@@ -12,7 +14,11 @@ import (
 
 type Shell struct {
 	Name   string
+	id     string
 	cfg    global.Tunnel
+	pid    int
+	stdin  io.WriteCloser
+	stdout io.ReadCloser
 }
 
 // New new shell
@@ -23,6 +29,14 @@ func New(cfg global.Tunnel) *Shell {
 	}
 }
 
+func (shell *Shell) Close() {
+	shell.onClose()
+	p, err := os.FindProcess(shell.pid)
+	if err == nil {
+		p.Kill()
+	}
+}
+
 // Handle handle shell
 func (shell *Shell) Handle(pl *pool.Pool) {
 	defer func() {
@@ -36,7 +50,7 @@ func (shell *Shell) Handle(pl *pool.Pool) {
 		}
 	}
 	mux := http.NewServeMux()
-	mux.HandleFunc("/new", pf(shell.New))
+	mux.HandleFunc("/ws", pf(shell.WS))
 	svr := &http.Server{
 		Addr:    fmt.Sprintf("%s:%d", shell.cfg.LocalAddr, shell.cfg.LocalPort),
 		Handler: mux,

go.mod

@@ -3,7 +3,9 @@ module natpass
 go 1.16
 
 require (
+	github.com/creack/pty v1.1.15 // indirect
 	github.com/dustin/go-humanize v1.0.0
+	github.com/gorilla/websocket v1.4.2 // indirect
 	github.com/kardianos/service v1.2.0
 	github.com/lwch/logging v0.0.0-20210528090125-a154917d90c6
 	github.com/lwch/runtime v0.0.0-20190520054850-8c97e19e0c6d

go.sum

@@ -1,8 +1,12 @@
+github.com/creack/pty v1.1.15 h1:cKRCLMj3Ddm54bKSpemfQ8AtYFBhAI2MPmdys22fBdc=
+github.com/creack/pty v1.1.15/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr4O4=
 github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo=
 github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
 github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
 github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
+github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc=
+github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/kardianos/service v1.2.0 h1:bGuZ/epo3vrt8IPC7mnKQolqFeYJb7Cs8Rk4PSOBB/g=
 github.com/kardianos/service v1.2.0/go.mod h1:CIMRFEJVL+0DS1a3Nx06NaMn4Dz63Ng6O7dl0qH0zVM=
 github.com/lwch/logging v0.0.0-20210528090125-a154917d90c6 h1:R52s6I/vW1NfNaJdY+Yr/ivkiFicouKmK0v3nvDQh4s=