diff --git a/code/client/shell/forward.go b/code/client/shell/forward.go index 76342a6..1d37ffe 100644 --- a/code/client/shell/forward.go +++ b/code/client/shell/forward.go @@ -24,6 +24,7 @@ func (link *Link) remoteRead() { link.targetIdx = msg.GetFromIdx() switch msg.GetXType() { case network.Msg_shell_created: + logging.Info("shell created") if msg.GetCrep().GetOk() { link.OnWork <- struct{}{} continue diff --git a/code/client/shell/h_new.go b/code/client/shell/h_new.go index abac1b7..fad49b2 100644 --- a/code/client/shell/h_new.go +++ b/code/client/shell/h_new.go @@ -21,6 +21,10 @@ func (shell *Shell) New(pool *pool.Pool, w http.ResponseWriter, r *http.Request) conn := pool.Get(id) conn.SendShellCreate(id, shell.cfg) conn.AddLink(id) + lk := NewLink(shell, id, shell.cfg.Target, conn) + shell.Lock() + shell.links[id] = lk + shell.Unlock() logging.Info("new shell: name=%s, id=%s", shell.Name, id) fmt.Fprint(w, id) } diff --git a/code/client/shell/h_ws.go b/code/client/shell/h_ws.go index 65ca0f3..897bbe0 100644 --- a/code/client/shell/h_ws.go +++ b/code/client/shell/h_ws.go @@ -19,7 +19,6 @@ func (shell *Shell) WS(pool *pool.Pool, w http.ResponseWriter, r *http.Request) id := strings.TrimPrefix(r.URL.Path, "/ws/") conn := pool.Get(id) - logging.Info("ws: %d", conn.Idx) local, err := upgrader.Upgrade(w, r, nil) if err != nil { logging.Error("upgrade websocket failed: %s, err=%v", shell.Name, err) @@ -45,7 +44,9 @@ func (shell *Shell) WS(pool *pool.Pool, w http.ResponseWriter, r *http.Request) func (shell *Shell) localForward(id string, local *websocket.Conn, remote *pool.Conn) { defer utils.Recover("localForward") defer local.Close() - logging.Info("localForward: %d", remote.Idx) + shell.RLock() + link := shell.links[id] + shell.RUnlock() for { _, data, err := local.ReadMessage() if err != nil { @@ -53,8 +54,7 @@ func (shell *Shell) localForward(id string, local *websocket.Conn, remote *pool. logging.Error("read local data for %s failed: %v", shell.Name, err) return } - remote.SendShellData(shell.cfg.Target, remote.Idx, id, data) - logging.Info("send: %d", remote.Idx) + remote.SendShellData(link.target, link.targetIdx, id, data) logging.Debug("local read %d bytes: name=%s, id=%s", len(data), shell.Name, id) } } @@ -68,6 +68,8 @@ func (shell *Shell) remoteForward(id string, ch <-chan *network.Msg, local *webs return } switch msg.GetXType() { + case network.Msg_shell_created: + logging.Info("shell created") case network.Msg_shell_data: err := local.WriteMessage(websocket.TextMessage, msg.GetSdata().GetData()) if err != nil { diff --git a/code/client/shell/shell.go b/code/client/shell/shell.go index 68f3ef7..9241b03 100644 --- a/code/client/shell/shell.go +++ b/code/client/shell/shell.go @@ -5,21 +5,25 @@ import ( "natpass/code/client/global" "natpass/code/client/pool" "net/http" + "sync" "github.com/lwch/logging" "github.com/lwch/runtime" ) type Shell struct { - Name string - cfg global.Tunnel + sync.RWMutex + Name string + cfg global.Tunnel + links map[string]*Link } // New new shell func New(cfg global.Tunnel) *Shell { return &Shell{ - Name: cfg.Name, - cfg: cfg, + Name: cfg.Name, + cfg: cfg, + links: make(map[string]*Link), } }