From fa7cb2e6457816df9d5e8dd853d53390ef4ed30c Mon Sep 17 00:00:00 2001 From: lwch Date: Tue, 26 Jul 2022 18:49:17 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E6=80=A7=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- code/client/conn/conn.go | 95 ++++++++++++++------ code/client/rule/code/h_forward.go | 32 +++++-- code/client/rule/code/h_forward_request.go | 4 +- code/client/rule/code/h_forward_websocket.go | 1 + code/client/rule/code/h_new.go | 14 ++- html/code/index.js | 8 +- 6 files changed, 113 insertions(+), 41 deletions(-) diff --git a/code/client/conn/conn.go b/code/client/conn/conn.go index 07e3945..c495374 100644 --- a/code/client/conn/conn.go +++ b/code/client/conn/conn.go @@ -3,6 +3,7 @@ package conn import ( "context" "crypto/tls" + "io" "net" "strings" "sync" @@ -90,11 +91,70 @@ func writeHandshake(conn *network.Conn, cfg *global.Configure) error { return conn.WriteMessage(&msg, 5*time.Second) } +func (conn *Conn) isDrop(linkID string) bool { + conn.lockDrop.RLock() + defer conn.lockDrop.RUnlock() + _, ok := conn.drop[linkID] + return ok +} + +func (conn *Conn) getChan(linkID string) chan *network.Msg { + conn.RLock() + ch := conn.read[linkID] + conn.RUnlock() + if ch == nil { + ch = conn.unknownRead + } + return ch +} + +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.lockDrop.Unlock() + select { + case ch <- msg: + default: + } + return false + } + return true +} + func (conn *Conn) loopRead() { defer utils.Recover("loopRead") defer conn.close() defer conn.cancel() var timeout int + run := func(msg *network.Msg) bool { + timeout = 0 + if msg.GetXType() == network.Msg_keepalive { + return true + } + logging.Debug("read message %s(%s) from %s", + msg.GetXType().String(), msg.GetLinkId(), msg.GetFrom()) + linkID := msg.GetLinkId() + if conn.isDrop(linkID) { + return true + } + ch := conn.getChan(linkID) + if !conn.hookDispatch(ch, msg) { + return true + } + select { + case ch <- msg: + 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.lockDrop.Unlock() + case <-conn.ctx.Done(): + return false + } + return true + } for { msg, _, err := conn.conn.ReadMessage(conn.cfg.ReadTimeout) if err != nil { @@ -106,36 +166,17 @@ func (conn *Conn) loopRead() { } continue } + if strings.Contains(err.Error(), "use of closed network connection") { + logging.Error("read message: %v", err) + return + } + if err == io.EOF { + return + } logging.Error("read message: %v", err) continue } - timeout = 0 - if msg.GetXType() == network.Msg_keepalive { - continue - } - logging.Debug("read message %s(%s) from %s", - msg.GetXType().String(), msg.GetLinkId(), msg.GetFrom()) - linkID := msg.GetLinkId() - conn.lockDrop.RLock() - _, drop := conn.drop[linkID] - conn.lockDrop.RUnlock() - if drop { - continue - } - conn.RLock() - ch := conn.read[linkID] - conn.RUnlock() - if ch == nil { - ch = conn.unknownRead - } - select { - case ch <- msg: - case <-time.After(conn.cfg.ReadTimeout): - logging.Error("drop message: %s", msg.GetXType().String()) - conn.lockDrop.Lock() - conn.drop[msg.GetLinkId()] = time.Now().Add(time.Minute) - conn.lockDrop.Unlock() - case <-conn.ctx.Done(): + if !run(msg) { return } } diff --git a/code/client/rule/code/h_forward.go b/code/client/rule/code/h_forward.go index eabf197..fc2a3a4 100644 --- a/code/client/rule/code/h_forward.go +++ b/code/client/rule/code/h_forward.go @@ -4,13 +4,36 @@ import ( "net/http" "strings" + "github.com/lwch/logging" "github.com/lwch/natpass/code/client/conn" ) // Forward forward code-server requests func (code *Code) Forward(conn *conn.Conn, w http.ResponseWriter, r *http.Request) { - id := strings.TrimPrefix(r.URL.Path, "/forward/") - id = id[:strings.Index(id, "/")] + name := strings.TrimPrefix(r.URL.Path, "/forward/") + name = name[:strings.Index(name, "/")] + + r.URL.Path = strings.TrimPrefix(r.URL.Path, "/forward/"+name) + if len(r.URL.Path) == 0 { + r.URL.Path = "/" + } + + var id string + if r.URL.Path == "/" { + id = r.FormValue("id") + http.SetCookie(w, &http.Cookie{ + Name: "__NATPASS_CONNECTION_ID__", + Value: id, + }) + } else { + cookie, err := r.Cookie("__NATPASS_CONNECTION_ID__") + if err != nil { + logging.Error("get connection id: %v", err) + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + id = cookie.Value + } code.RLock() workspace := code.workspace[id] @@ -21,11 +44,6 @@ func (code *Code) Forward(conn *conn.Conn, w http.ResponseWriter, r *http.Reques return } - r.URL.Path = strings.TrimPrefix(r.URL.Path, "/forward/"+id) - if len(r.URL.Path) == 0 { - r.URL.Path = "/" - } - if code.isWebsocket(r) { code.handleWebsocket(workspace, w, r) } else { diff --git a/code/client/rule/code/h_forward_request.go b/code/client/rule/code/h_forward_request.go index 4c6a393..d4f1e62 100644 --- a/code/client/rule/code/h_forward_request.go +++ b/code/client/rule/code/h_forward_request.go @@ -20,8 +20,8 @@ func (code *Code) handleRequest(workspace *Workspace, w http.ResponseWriter, r * defer workspace.closeMessage(reqID) resp := workspace.onResponse(reqID) if resp == nil { - logging.Error("waiting for [%s] [%s] no response for request, request_id=%d", - workspace.id, workspace.name, reqID) + logging.Error("waiting for [%s] [%s] no response for request, uri=%s, request_id=%d", + workspace.id, workspace.name, r.URL.Path, reqID) http.Error(w, "no response", http.StatusInternalServerError) return } diff --git a/code/client/rule/code/h_forward_websocket.go b/code/client/rule/code/h_forward_websocket.go index 6f0a534..b831a45 100644 --- a/code/client/rule/code/h_forward_websocket.go +++ b/code/client/rule/code/h_forward_websocket.go @@ -59,4 +59,5 @@ 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 a55bd24..61bd3d8 100644 --- a/code/client/rule/code/h_new.go +++ b/code/client/rule/code/h_new.go @@ -1,7 +1,7 @@ package code import ( - "fmt" + "encoding/json" "net/http" "time" @@ -52,5 +52,15 @@ func (code *Code) New(conn *conn.Conn, w http.ResponseWriter, r *http.Request) { link.GetID(), code.cfg.Name, repMsg.GetTo(), repMsg.GetFrom()) go link.localRead() - fmt.Fprint(w, id) + 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) } diff --git a/html/code/index.js b/html/code/index.js index 6a94c9b..4970a34 100644 --- a/html/code/index.js +++ b/html/code/index.js @@ -4,10 +4,12 @@ var page = { }, connect: function() { $.get('/new', function(ret) { - page.id = ret; - $('#code').attr('src', `/forward/${page.id}/`); + page.id = ret.id; + page.name = ret.name; + $('#code').attr('src', `/forward/${page.name}/?id=${page.id}`); }); }, - id: '' + id: '', + name: '' }; $(document).ready(page.init); \ No newline at end of file