github/natpass
4
0
Fork 0
mirror of https://github.com/lwch/natpass.git synced 2024-04-21 12:41:54 +00:00
Code Issues Projects Releases Wiki Activity

修改connect处理逻辑

Browse Source
This commit is contained in:
lwch
2021-09-28 11:20:42 +08:00
parent 0fffe85987
commit 71dfec0f8a
8 changed files with 84 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
code/client/connect.go
+4 -3
View File
@@ -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()
}
code/client/pool/pool.go
+1 -2
View File
@@ -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,
}
}
code/client/shell/exec_xx.go
+9 -10
View File
@@ -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()
}
}
code/client/shell/forward.go
+17 -15
View File
@@ -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])
}
}
code/client/shell/h_ws.go
+1
View File
@@ -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)
}
}
code/client/shell/link.go
+48
View File
@@ -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()
}
}
code/client/shell/shell.go
+2 -16
View File
@@ -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() {
code/client/tunnel/link.go
+2
View File
@@ -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
}