From 29c5324611dc6fcc5558bcfe983d5f62f6656b84 Mon Sep 17 00:00:00 2001 From: lwch Date: Wed, 27 Jul 2022 11:07:49 +0800 Subject: [PATCH] =?UTF-8?q?1.=20=E4=BC=98=E5=8C=96disconnect=E5=A4=84?= =?UTF-8?q?=E7=90=86=E9=80=BB=E8=BE=91=202.=20=E4=BF=AE=E6=AD=A3code-serve?= =?UTF-8?q?r=E6=97=A0=E6=B3=95=E9=87=8D=E6=96=B0=E8=BD=BD=E5=85=A5?= =?UTF-8?q?=E7=9A=84=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- CHANGELOG.md | 4 +- code/client/app/app.go | 8 +++ code/client/conn/conn.go | 58 ++++++++++++++------ code/client/rule/code/code.go | 58 ++++++++++++++++++++ code/client/rule/code/h_forward.go | 10 +++- code/client/rule/code/h_forward_request.go | 13 +++-- code/client/rule/code/h_forward_websocket.go | 2 +- code/client/rule/code/h_new.go | 56 +------------------ code/client/rule/code/workspace.go | 12 ++-- code/client/rule/mgr.go | 14 +++++ code/client/rule/shell/h_ws.go | 8 +-- code/client/rule/shell/link.go | 14 ++--- code/client/rule/shell/shell.go | 10 ++++ code/client/rule/vnc/h_new.go | 2 +- code/client/rule/vnc/link.go | 15 ++--- code/client/rule/vnc/vnc.go | 11 +++- html/code/index.js | 6 +- 17 files changed, 189 insertions(+), 112 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 213102d..4fc9b98 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -131,4 +131,6 @@ # TODO -1. go版本升级到1.18.4 \ No newline at end of file +1. go版本升级到1.18.4 +2. 新增code-server支持 +3. 优化disconnect处理逻辑 \ No newline at end of file diff --git a/code/client/app/app.go b/code/client/app/app.go index ea688e8..4f2e47f 100644 --- a/code/client/app/app.go +++ b/code/client/app/app.go @@ -116,6 +116,7 @@ func (a *App) run() { linkID = msg.GetLinkId() } if len(linkID) > 0 { + a.conn.ChanClose(linkID) logging.Error("link of %s not found, type=%s", linkID, msg.GetXType().String()) continue @@ -123,6 +124,13 @@ func (a *App) run() { } }() + go func() { + for { + id := <-a.conn.ChanDisconnect() + mgr.OnDisconnect(id) + } + }() + if a.cfg.DashboardEnabled { go func() { a.conn.Wait() diff --git a/code/client/conn/conn.go b/code/client/conn/conn.go index c495374..cbd0344 100644 --- a/code/client/conn/conn.go +++ b/code/client/conn/conn.go @@ -16,16 +16,19 @@ import ( "github.com/lwch/runtime" ) +const dropBlockTimeout = 10 * time.Minute + // Conn connection type Conn struct { sync.RWMutex - cfg *global.Configure - conn *network.Conn - read map[string]chan *network.Msg // link id => channel - unknownRead chan *network.Msg // read message without link - write chan *network.Msg - lockDrop sync.RWMutex - drop map[string]time.Time + cfg *global.Configure + conn *network.Conn + read map[string]chan *network.Msg // link id => channel + unknownRead chan *network.Msg // read message without link + onDisconnect chan string + write chan *network.Msg + lockDrop sync.RWMutex + drop map[string]time.Time // runtime ctx context.Context cancel context.CancelFunc @@ -34,11 +37,12 @@ type Conn struct { // New new connection func New(cfg *global.Configure) *Conn { conn := &Conn{ - cfg: cfg, - read: make(map[string]chan *network.Msg), - unknownRead: make(chan *network.Msg, 1024), - write: make(chan *network.Msg, 10*1024*1024), - drop: make(map[string]time.Time), + cfg: cfg, + read: make(map[string]chan *network.Msg), + unknownRead: make(chan *network.Msg, 1024), + onDisconnect: make(chan string, 1024), + write: make(chan *network.Msg, 10*1024*1024), + drop: make(map[string]time.Time), } runtime.Assert(conn.connect()) conn.ctx, conn.cancel = context.WithCancel(context.Background()) @@ -112,12 +116,10 @@ func (conn *Conn) hookDispatch(ch chan *network.Msg, msg *network.Msg) bool { switch msg.GetXType() { case network.Msg_disconnect: conn.lockDrop.Lock() - conn.drop[msg.GetLinkId()] = time.Now().Add(time.Minute) + conn.drop[msg.GetLinkId()] = time.Now().Add(dropBlockTimeout) conn.lockDrop.Unlock() - select { - case ch <- msg: - default: - } + conn.onDisconnect <- msg.GetLinkId() + logging.Info("connection %s disconnected", msg.GetLinkId()) return false } return true @@ -148,7 +150,7 @@ func (conn *Conn) loopRead() { case <-time.After(conn.cfg.WriteTimeout): logging.Error("drop message: %s", msg.GetXType().String()) conn.lockDrop.Lock() - conn.drop[msg.GetLinkId()] = time.Now().Add(time.Minute) + conn.drop[msg.GetLinkId()] = time.Now().Add(dropBlockTimeout) conn.lockDrop.Unlock() case <-conn.ctx.Done(): return false @@ -248,6 +250,11 @@ func (conn *Conn) ChanUnknown() <-chan *network.Msg { return conn.unknownRead } +// ChanDisconnect get channel of disconnect +func (conn *Conn) ChanDisconnect() <-chan string { + return conn.onDisconnect +} + func (conn *Conn) checkDrop() { for { time.Sleep(time.Second) @@ -273,3 +280,18 @@ func (conn *Conn) checkDrop() { func (conn *Conn) Wait() { <-conn.ctx.Done() } + +// ChanClose close read chan +func (conn *Conn) ChanClose(id string) { + conn.Lock() + ch := conn.read[id] + if ch != nil { + close(ch) + } + delete(conn.read, id) + conn.Unlock() + + conn.lockDrop.Lock() + conn.drop[id] = time.Now().Add(dropBlockTimeout) + conn.lockDrop.Unlock() +} diff --git a/code/client/rule/code/code.go b/code/client/rule/code/code.go index 5447857..9ab7de4 100644 --- a/code/client/rule/code/code.go +++ b/code/client/rule/code/code.go @@ -1,6 +1,7 @@ package code import ( + "errors" "fmt" "net" "net/http" @@ -11,6 +12,7 @@ import ( "github.com/lwch/natpass/code/client/conn" "github.com/lwch/natpass/code/client/global" "github.com/lwch/natpass/code/client/rule" + "github.com/lwch/natpass/code/network" "github.com/lwch/runtime" ) @@ -81,6 +83,16 @@ func (code *Code) NewLink(id, remote string, localConn net.Conn, remoteConn *con return ws } +// OnDisconnect on disconnect message +func (code *Code) OnDisconnect(id string) { + code.RLock() + workspace := code.workspace[id] + code.RUnlock() + if workspace != nil { + workspace.Close(false) + } +} + // Handle handle code-server func (code *Code) Handle(c *conn.Conn) { defer func() { @@ -103,3 +115,49 @@ func (code *Code) Handle(c *conn.Conn) { } runtime.Assert(svr.ListenAndServe()) } + +func (code *Code) remove(id string) { + code.Lock() + delete(code.workspace, id) + code.Unlock() +} + +func (code *Code) new(conn *conn.Conn) (string, error) { + id, err := runtime.UUID(16, "0123456789abcdef") + if err != nil { + logging.Error("failed to generate link_id for code-server: %s, err=%v", + code.Name, err) + return "", err + } + link := code.NewLink(id, code.cfg.Target, nil, conn).(*Workspace) + conn.SendConnectReq(id, code.cfg) + ch := conn.ChanRead(id) + var repMsg *network.Msg + for { + var msg *network.Msg + select { + case msg = <-ch: + case <-time.After(time.Minute): + logging.Error("create code-server %s by rule %s failed, timtout", link.id, link.parent.Name) + return "", errWaitingTimeout + } + if msg.GetXType() != network.Msg_connect_rep { + conn.Reset(id, msg) + time.Sleep(code.readTimeout / 10) + continue + } + rep := msg.GetCrep() + if !rep.GetOk() { + logging.Error("create code-server %s by rule %s failed, err=%s", + link.id, link.parent.Name, rep.GetMsg()) + return "", errors.New(rep.GetMsg()) + } + repMsg = msg + break + } + logging.Info("create link %s for code-server rule [%s] from %s to %s", + link.GetID(), code.cfg.Name, + repMsg.GetTo(), repMsg.GetFrom()) + go link.localRead() + return id, nil +} diff --git a/code/client/rule/code/h_forward.go b/code/client/rule/code/h_forward.go index fc2a3a4..8b785ba 100644 --- a/code/client/rule/code/h_forward.go +++ b/code/client/rule/code/h_forward.go @@ -20,7 +20,13 @@ func (code *Code) Forward(conn *conn.Conn, w http.ResponseWriter, r *http.Reques var id string if r.URL.Path == "/" { - id = r.FormValue("id") + var err error + id, err = code.new(conn) + if err != nil { + logging.Error("can not create workspace for [%s]: %v", code.Name, err) + http.Error(w, err.Error(), http.StatusBadGateway) + return + } http.SetCookie(w, &http.Cookie{ Name: "__NATPASS_CONNECTION_ID__", Value: id, @@ -47,7 +53,7 @@ func (code *Code) Forward(conn *conn.Conn, w http.ResponseWriter, r *http.Reques if code.isWebsocket(r) { code.handleWebsocket(workspace, w, r) } else { - code.handleRequest(workspace, w, r) + code.handleRequest(conn, workspace, w, r) } } diff --git a/code/client/rule/code/h_forward_request.go b/code/client/rule/code/h_forward_request.go index d4f1e62..6d3ea58 100644 --- a/code/client/rule/code/h_forward_request.go +++ b/code/client/rule/code/h_forward_request.go @@ -7,10 +7,11 @@ import ( "net/http" "github.com/lwch/logging" + "github.com/lwch/natpass/code/client/conn" "github.com/lwch/natpass/code/network" ) -func (code *Code) handleRequest(workspace *Workspace, w http.ResponseWriter, r *http.Request) { +func (code *Code) handleRequest(conn *conn.Conn, workspace *Workspace, w http.ResponseWriter, r *http.Request) { reqID, err := workspace.SendRequest(r) if err != nil { logging.Error("send_request: %v", err) @@ -40,13 +41,12 @@ func (code *Code) handleRequest(workspace *Workspace, w http.ResponseWriter, r * } } - w.WriteHeader(int(hdr.GetCode())) - var idx uint32 + var buf bytes.Buffer for { msg := workspace.onResponse(reqID) if msg == nil { - logging.Error("no response") + logging.Error("no response [%s] [%s]", workspace.id, workspace.name) http.Error(w, "no response", http.StatusBadGateway) return } @@ -67,12 +67,15 @@ func (code *Code) handleRequest(workspace *Workspace, w http.ResponseWriter, r * http.Error(w, fmt.Sprintf("read error: %s", string(resp.GetBody())), http.StatusResetContent) return } - _, err = io.Copy(w, bytes.NewReader(resp.GetBody())) + _, err = io.Copy(&buf, bytes.NewReader(resp.GetBody())) if err != nil { logging.Error("write body: %v", err) + http.Error(w, fmt.Sprintf("save data: %v", err), http.StatusInternalServerError) return } if resp.GetMask()&2 > 0 { + w.WriteHeader(int(hdr.GetCode())) + io.Copy(w, &buf) return } idx++ diff --git a/code/client/rule/code/h_forward_websocket.go b/code/client/rule/code/h_forward_websocket.go index b831a45..46867dd 100644 --- a/code/client/rule/code/h_forward_websocket.go +++ b/code/client/rule/code/h_forward_websocket.go @@ -14,6 +14,7 @@ var upgrader = websocket.Upgrader{ } func (code *Code) handleWebsocket(workspace *Workspace, w http.ResponseWriter, r *http.Request) { + defer workspace.Close(true) reqID, err := workspace.SendConnect(r) if err != nil { logging.Error("send_connect: %v", err) @@ -59,5 +60,4 @@ func (code *Code) handleWebsocket(workspace *Workspace, w http.ResponseWriter, r go workspace.ws2remote(&wg, reqID, local) go workspace.remote2ws(&wg, reqID, local) wg.Wait() - workspace.Close() } diff --git a/code/client/rule/code/h_new.go b/code/client/rule/code/h_new.go index 61bd3d8..f643ede 100644 --- a/code/client/rule/code/h_new.go +++ b/code/client/rule/code/h_new.go @@ -1,66 +1,12 @@ package code import ( - "encoding/json" "net/http" - "time" - "github.com/lwch/logging" "github.com/lwch/natpass/code/client/conn" - "github.com/lwch/natpass/code/network" - "github.com/lwch/runtime" ) // New new code-server workspace func (code *Code) New(conn *conn.Conn, w http.ResponseWriter, r *http.Request) { - id, err := runtime.UUID(16, "0123456789abcdef") - if err != nil { - logging.Error("failed to generate link_id for code-server: %s, err=%v", - code.Name, err) - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - link := code.NewLink(id, code.cfg.Target, nil, conn).(*Workspace) - conn.SendConnectReq(id, code.cfg) - ch := conn.ChanRead(id) - var repMsg *network.Msg - for { - var msg *network.Msg - select { - case msg = <-ch: - case <-time.After(time.Minute): - logging.Error("create code-server %s by rule %s failed, timtout", link.id, link.parent.Name) - http.Error(w, "timeout", http.StatusBadGateway) - return - } - if msg.GetXType() != network.Msg_connect_rep { - conn.Reset(id, msg) - time.Sleep(code.readTimeout / 10) - continue - } - rep := msg.GetCrep() - if !rep.GetOk() { - logging.Error("create code-server %s by rule %s failed, err=%s", - link.id, link.parent.Name, rep.GetMsg()) - http.Error(w, rep.GetMsg(), http.StatusBadGateway) - return - } - repMsg = msg - break - } - logging.Info("create link %s for code-server rule [%s] from %s to %s", - link.GetID(), code.cfg.Name, - repMsg.GetTo(), repMsg.GetFrom()) - go link.localRead() - w.Header().Set("Content-Type", "application/json") - data, err := json.Marshal(map[string]string{ - "id": id, - "name": code.cfg.Name, - }) - if err != nil { - logging.Error("json.Marshal: %v", err) - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - w.Write(data) + w.Write([]byte(code.cfg.Name)) } diff --git a/code/client/rule/code/workspace.go b/code/client/rule/code/workspace.go index 86fb3c0..e91768a 100644 --- a/code/client/rule/code/workspace.go +++ b/code/client/rule/code/workspace.go @@ -131,11 +131,15 @@ func (ws *Workspace) Exec(dir string) error { } // Close close workspace -func (ws *Workspace) Close() { +func (ws *Workspace) Close(send bool) { if ws.exec != nil && ws.exec.Process != nil { ws.exec.Process.Kill() } - ws.remote.SendDisconnect(ws.target, ws.id) + if send { + ws.remote.SendDisconnect(ws.target, ws.id) + } + ws.parent.remove(ws.id) + ws.remote.ChanClose(ws.id) } func (ws *Workspace) log(stdout, stderr io.ReadCloser) { @@ -168,7 +172,7 @@ func (ws *Workspace) Forward() { func (ws *Workspace) remoteRead() { defer utils.Recover("remoteRead") - defer ws.Close() + defer ws.Close(true) ch := ws.remote.ChanRead(ws.id) for { msg := <-ch @@ -200,7 +204,7 @@ func (ws *Workspace) closeMessage(reqID uint64) { func (ws *Workspace) localRead() { defer utils.Recover("localRead") - defer ws.Close() + defer ws.Close(true) ch := ws.remote.ChanRead(ws.id) for { msg := <-ch diff --git a/code/client/rule/mgr.go b/code/client/rule/mgr.go index 203b829..68d9333 100644 --- a/code/client/rule/mgr.go +++ b/code/client/rule/mgr.go @@ -29,6 +29,7 @@ type LinkedRule interface { GetRemote() string GetTarget() string GetLinks() []Link + OnDisconnect(string) } // Mgr rule manager @@ -73,3 +74,16 @@ func (mgr *Mgr) Range(fn func(Rule)) { fn(t) } } + +// OnDisconnect on disconnect message +func (mgr *Mgr) OnDisconnect(id string) { + var links []LinkedRule + mgr.Range(func(r Rule) { + if lr, ok := r.(LinkedRule); ok { + links = append(links, lr) + } + }) + for _, link := range links { + go link.OnDisconnect(id) + } +} diff --git a/code/client/rule/shell/h_ws.go b/code/client/rule/shell/h_ws.go index f205bdb..f870cff 100644 --- a/code/client/rule/shell/h_ws.go +++ b/code/client/rule/shell/h_ws.go @@ -47,7 +47,7 @@ func (shell *Shell) localForward(id string, local *websocket.Conn) { shell.RLock() link := shell.links[id] shell.RUnlock() - defer link.Close() + defer link.Close(true) for { _, data, err := local.ReadMessage() if err != nil { @@ -66,7 +66,7 @@ func (shell *Shell) remoteForward(id string, local *websocket.Conn) { link := shell.links[id] shell.RUnlock() ch := link.remote.ChanRead(id) - defer link.Close() + defer link.Close(true) for { msg := <-ch if msg == nil { @@ -84,10 +84,6 @@ func (shell *Shell) remoteForward(id string, local *websocket.Conn) { } logging.Debug("remote read %d bytes: name=%s, id=%s", len(msg.GetSdata().GetData()), shell.Name, id) - case network.Msg_disconnect: - logging.Info("shell %s by rule %s closed by remote", - link.id, link.parent.Name) - return } } } diff --git a/code/client/rule/shell/link.go b/code/client/rule/shell/link.go index 09df2f9..f0204c7 100644 --- a/code/client/rule/shell/link.go +++ b/code/client/rule/shell/link.go @@ -45,14 +45,17 @@ func (link *Link) GetPackets() (uint64, uint64) { } // Close close link -func (link *Link) Close() { +func (link *Link) Close(send bool) { link.onClose() p, err := os.FindProcess(link.pid) if err == nil { p.Kill() } - link.remote.SendDisconnect(link.target, link.id) + if send { + link.remote.SendDisconnect(link.target, link.id) + } link.parent.remove(link.id) + link.remote.ChanClose(link.id) } // Forward forward data @@ -63,7 +66,7 @@ func (link *Link) Forward() { func (link *Link) remoteRead() { defer utils.Recover("remoteRead") - defer link.Close() + defer link.Close(true) ch := link.remote.ChanRead(link.id) for { msg := <-ch @@ -84,16 +87,13 @@ func (link *Link) remoteRead() { link.parent.Name, link.id, err) return } - case network.Msg_disconnect: - logging.Info("shell %s link %s closed by remote", link.parent.Name, link.id) - return } } } func (link *Link) localRead() { defer utils.Recover("localRead") - defer link.Close() + defer link.Close(true) buf := make([]byte, 16*1024) for { n, err := link.stdout.Read(buf) diff --git a/code/client/rule/shell/shell.go b/code/client/rule/shell/shell.go index 5909fa2..89596d6 100644 --- a/code/client/rule/shell/shell.go +++ b/code/client/rule/shell/shell.go @@ -86,6 +86,16 @@ func (shell *Shell) GetPort() uint16 { return shell.cfg.LocalPort } +// OnDisconnect on disconnect message +func (shell *Shell) OnDisconnect(id string) { + shell.RLock() + link := shell.links[id] + shell.RUnlock() + if link != nil { + link.Close(false) + } +} + // Handle handle shell func (shell *Shell) Handle(c *conn.Conn) { defer func() { diff --git a/code/client/rule/vnc/h_new.go b/code/client/rule/vnc/h_new.go index c88c254..6d30a41 100644 --- a/code/client/rule/vnc/h_new.go +++ b/code/client/rule/vnc/h_new.go @@ -15,7 +15,7 @@ import ( // New new vnc func (v *VNC) New(conn *conn.Conn, w http.ResponseWriter, r *http.Request) { if v.link != nil { - v.link.close() + v.link.Close(true) } q := r.FormValue("quality") s := r.FormValue("show_cursor") diff --git a/code/client/rule/vnc/link.go b/code/client/rule/vnc/link.go index d11bd62..6aec6e1 100644 --- a/code/client/rule/vnc/link.go +++ b/code/client/rule/vnc/link.go @@ -82,7 +82,7 @@ func (link *Link) Forward() { } func (link *Link) remoteRead() { - defer link.close() + defer link.Close(true) ch := link.remote.ChanRead(link.id) for { msg := <-ch @@ -110,9 +110,6 @@ func (link *Link) remoteRead() { data := link.ps.GetClipboard() link.remote.SendVNCClipboardData(link.target, link.id, true, data) } - case network.Msg_disconnect: - logging.Info("link %s disconnected", link.id) - return } } } @@ -120,7 +117,7 @@ func (link *Link) remoteRead() { func (link *Link) localRead() { // TODO: exit by context defer utils.Recover("capture") - defer link.close() + defer link.Close(true) img, err := link.ps.Capture(3 * time.Second) if err != nil { logging.Error("capture: %v", err) @@ -152,11 +149,15 @@ func (link *Link) localRead() { } } -func (link *Link) close() { +func (link *Link) Close(send bool) { if link.ps != nil { link.ps.Close() } - link.remote.SendDisconnect(link.target, link.id) + if send { + link.remote.SendDisconnect(link.target, link.id) + } + link.parent.remove(link.id) + link.remote.ChanClose(link.id) } func cut(src *image.RGBA, rect image.Rectangle) *image.RGBA { diff --git a/code/client/rule/vnc/vnc.go b/code/client/rule/vnc/vnc.go index 2ea02fa..b453077 100644 --- a/code/client/rule/vnc/vnc.go +++ b/code/client/rule/vnc/vnc.go @@ -47,7 +47,7 @@ func (v *VNC) NewLink(id, remote string, localConn net.Conn, remoteConn *conn.Co remote: remoteConn, } if v.link != nil { - v.link.close() + v.link.Close(true) } v.link = link return link @@ -86,6 +86,11 @@ func (v *VNC) GetPort() uint16 { return v.cfg.LocalPort } +// OnDisconnect on disconnect message +func (v *VNC) OnDisconnect(id string) { + // TODO +} + // Handle handle shell func (v *VNC) Handle(c *conn.Conn) { defer func() { @@ -110,3 +115,7 @@ func (v *VNC) Handle(c *conn.Conn) { } runtime.Assert(svr.ListenAndServe()) } + +func (v *VNC) remove(id string) { + v.link = nil +} diff --git a/html/code/index.js b/html/code/index.js index 4970a34..0a71b45 100644 --- a/html/code/index.js +++ b/html/code/index.js @@ -4,12 +4,10 @@ var page = { }, connect: function() { $.get('/new', function(ret) { - page.id = ret.id; - page.name = ret.name; - $('#code').attr('src', `/forward/${page.name}/?id=${page.id}`); + page.name = ret; + $('#code').attr('src', `/forward/${page.name}/`); }); }, - id: '', name: '' }; $(document).ready(page.init); \ No newline at end of file