1. 优化disconnect处理逻辑

2. 修正code-server无法重新载入的问题

CHANGELOG.md

@@ -131,4 +131,6 @@
 
 # TODO
 
-1. go版本升级到1.18.4
+1. go版本升级到1.18.4
+2. 新增code-server支持
+3. 优化disconnect处理逻辑

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()

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()
+}

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
+}

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)
 	}
 }
 

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++

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()
 }

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))
 }

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

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)
+	}
+}

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
 		}
 	}
 }

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)

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() {

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")

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 {

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
+}

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);