diff --git a/code/client/connect.go b/code/client/connect.go index 2f94ba0..b7cab9a 100644 --- a/code/client/connect.go +++ b/code/client/connect.go @@ -50,11 +50,12 @@ func shellCreate(conn *pool.Conn, msg *network.Msg) { Type: "shell", Exec: create.GetExec(), }) - err := sh.Exec(msg.GetLinkId()) + lk := shell.NewLink(sh, msg.GetLinkId(), msg.GetFrom(), conn) + lk.SetTargetIdx(msg.GetFromIdx()) + err := lk.Exec() if err != nil { logging.Error("create shell failed: %v", err) return } - conn.AddLink(msg.GetLinkId()) - sh.Forward(conn, msg.GetFromIdx()) + lk.Forward() } diff --git a/code/client/pool/pool.go b/code/client/pool/pool.go index b08591b..bd2f3dd 100644 --- a/code/client/pool/pool.go +++ b/code/client/pool/pool.go @@ -2,7 +2,6 @@ package pool import ( "crypto/tls" - "math/rand" "natpass/code/client/global" "natpass/code/network" "sync" @@ -28,7 +27,7 @@ func New(cfg *global.Configure) *Pool { cfg: cfg, conns: make(map[uint32]*Conn, cfg.Links), count: cfg.Links, - idx: rand.Uint32(), // random initialize for fast restart + idx: 0, } } diff --git a/code/client/shell/exec_xx.go b/code/client/shell/exec_xx.go index 684b084..79329d6 100644 --- a/code/client/shell/exec_xx.go +++ b/code/client/shell/exec_xx.go @@ -9,11 +9,10 @@ import ( "github.com/creack/pty" ) -func (shell *Shell) Exec(id string) error { - shell.id = id +func (link *Link) Exec() error { var cmd *exec.Cmd - if len(shell.cfg.Exec) > 0 { - cmd = exec.Command(shell.cfg.Exec) + if len(link.parent.cfg.Exec) > 0 { + cmd = exec.Command(link.parent.cfg.Exec) } if cmd == nil { dir, err := exec.LookPath("bash") @@ -34,14 +33,14 @@ func (shell *Shell) Exec(id string) error { if err != nil { return err } - shell.stdin = f - shell.stdout = f - shell.pid = cmd.Process.Pid + link.stdin = f + link.stdout = f + link.pid = cmd.Process.Pid return nil } -func (shell *Shell) onClose() { - if shell.stdin != nil { - shell.stdin.Close() +func (link *Link) onClose() { + if link.stdin != nil { + link.stdin.Close() } } diff --git a/code/client/shell/forward.go b/code/client/shell/forward.go index f92480f..c8d160b 100644 --- a/code/client/shell/forward.go +++ b/code/client/shell/forward.go @@ -1,22 +1,21 @@ 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 (link *Link) Forward() { + go link.remoteRead() + go link.localRead() } -func (shell *Shell) remoteRead(remote *pool.Conn) { +func (link *Link) remoteRead() { defer utils.Recover("remoteRead") - defer shell.Close() - ch := remote.ChanRead(shell.id) + defer link.Close() + ch := link.remote.ChanRead(link.id) for { msg := <-ch if msg == nil { @@ -26,10 +25,11 @@ func (shell *Shell) remoteRead(remote *pool.Conn) { case network.Msg_shell_resize: // TODO case network.Msg_shell_data: - _, err := shell.stdin.Write(msg.GetSdata().GetData()) + _, err := link.stdin.Write(msg.GetSdata().GetData()) if err != nil { // TODO: close - logging.Error("write data on shell %s link %s failed, err=%v", shell.Name, shell.id, err) + logging.Error("write data on shell %s link %s failed, err=%v", + link.parent.Name, link.id, err) return } case network.Msg_shell_close: @@ -38,23 +38,25 @@ func (shell *Shell) remoteRead(remote *pool.Conn) { } } -func (shell *Shell) localRead(remote *pool.Conn, toIdx uint32) { +func (link *Link) localRead() { defer utils.Recover("localRead") - defer shell.Close() + defer link.Close() buf := make([]byte, 16*1024) for { - n, err := shell.stdout.Read(buf) + n, err := link.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) + logging.Error("read data on shell %s link %s failed, err=%v", + link.parent.Name, link.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]) + logging.Debug("link %s on shell %s read from local %d bytes", + link.id, link.parent.Name, n) + link.remote.SendShellData(link.target, link.targetIdx, link.id, buf[:n]) } } diff --git a/code/client/shell/h_ws.go b/code/client/shell/h_ws.go index e805f6b..11611ea 100644 --- a/code/client/shell/h_ws.go +++ b/code/client/shell/h_ws.go @@ -52,6 +52,7 @@ func (shell *Shell) localForward(id string, local *websocket.Conn, remote *pool. return } remote.SendShellData(shell.cfg.Target, remote.Idx, id, data) + logging.Info("send: %d", remote.Idx) logging.Debug("local read %d bytes: name=%s, id=%s", len(data), shell.Name, id) } } diff --git a/code/client/shell/link.go b/code/client/shell/link.go new file mode 100644 index 0000000..4001283 --- /dev/null +++ b/code/client/shell/link.go @@ -0,0 +1,48 @@ +package shell + +import ( + "io" + "natpass/code/client/pool" + "os" + + "github.com/lwch/logging" +) + +type Link struct { + parent *Shell + id string // link id + target string // target id + targetIdx uint32 // target idx + remote *pool.Conn + // in remote + pid int + stdin io.WriteCloser + stdout io.ReadCloser +} + +// NewLink create link +func NewLink(parent *Shell, id, target string, remote *pool.Conn) *Link { + remote.AddLink(id) + logging.Info("create shell %s for tunnel %s on connection %d", + id, parent.Name, remote.Idx) + return &Link{ + parent: parent, + id: id, + target: target, + remote: remote, + } +} + +// SetTargetIdx set link remote index +func (link *Link) SetTargetIdx(idx uint32) { + link.targetIdx = idx +} + +// Close close link +func (link *Link) Close() { + link.onClose() + p, err := os.FindProcess(link.pid) + if err == nil { + p.Kill() + } +} diff --git a/code/client/shell/shell.go b/code/client/shell/shell.go index 28e300a..68f3ef7 100644 --- a/code/client/shell/shell.go +++ b/code/client/shell/shell.go @@ -2,23 +2,17 @@ package shell import ( "fmt" - "io" "natpass/code/client/global" "natpass/code/client/pool" "net/http" - "os" "github.com/lwch/logging" "github.com/lwch/runtime" ) type Shell struct { - Name string - id string - cfg global.Tunnel - pid int - stdin io.WriteCloser - stdout io.ReadCloser + Name string + cfg global.Tunnel } // New new shell @@ -29,14 +23,6 @@ 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() { diff --git a/code/client/tunnel/link.go b/code/client/tunnel/link.go index ef99b83..c7f6f10 100644 --- a/code/client/tunnel/link.go +++ b/code/client/tunnel/link.go @@ -22,6 +22,7 @@ type Link struct { closeFromRemote bool } +// NewLink create link func NewLink(parent *Tunnel, id, target string, local net.Conn, remote *pool.Conn) *Link { remote.AddLink(id) logging.Info("create link %s for tunnel %s on connection %d", @@ -104,6 +105,7 @@ func (link *Link) localRead() { } } +// SetTargetIdx set link remote index func (link *Link) SetTargetIdx(idx uint32) { link.targetIdx = idx }