实现链接断开时的disconnect消息发送逻辑

code/server/handler/client.go

@@ -3,22 +3,26 @@ package handler
 import (
 	"natpass/code/network"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/lwch/logging"
 )
 
 type client struct {
-	parent *Handler
-	id     string
-	c      *network.Conn
+	sync.RWMutex
+	parent  *Handler
+	id      string
+	c       *network.Conn
+	tunnels map[string]struct{}
 }
 
 func newClient(parent *Handler, id string, conn *network.Conn) *client {
 	return &client{
-		parent: parent,
-		id:     id,
-		c:      conn,
+		parent:  parent,
+		id:      id,
+		c:       conn,
+		tunnels: make(map[string]struct{}),
 	}
 }
 
@@ -39,3 +43,35 @@ func (c *client) run() {
 func (c *client) writeMessage(msg *network.Msg) error {
 	return c.c.WriteMessage(msg, time.Second)
 }
+
+func (c *client) addTunnel(id string) {
+	c.Lock()
+	c.tunnels[id] = struct{}{}
+	c.Unlock()
+}
+
+func (c *client) getTunnels() []string {
+	ret := make([]string, 0, len(c.tunnels))
+	c.RLock()
+	for tn := range c.tunnels {
+		ret = append(ret, tn)
+	}
+	c.RUnlock()
+	return ret
+}
+
+func (c *client) close(id string) {
+	var msg network.Msg
+	msg.From = "server"
+	msg.To = c.id
+	msg.XType = network.Msg_disconnect
+	msg.Payload = &network.Msg_XDisconnect{
+		XDisconnect: &network.Disconnect{
+			Id: id,
+		},
+	}
+	c.c.WriteMessage(&msg, time.Second)
+	c.Lock()
+	delete(c.tunnels, id)
+	c.Unlock()
+}

code/server/handler/handler.go

@@ -14,13 +14,15 @@ import (
 type Handler struct {
 	sync.RWMutex
 	cfg     *global.Configure
-	clients map[string]*client
+	clients map[string]*client    // client id => client
+	tunnels map[string][2]*client // tunnel id => endpoints
 }
 
 func New(cfg *global.Configure) *Handler {
 	return &Handler{
 		cfg:     cfg,
 		clients: make(map[string]*client),
+		tunnels: make(map[string][2]*client),
 	}
 }
 
@@ -56,6 +58,8 @@ func (h *Handler) Handle(conn net.Conn) {
 	h.clients[cli.id] = cli
 	h.Unlock()
 
+	defer h.closeAll(cli)
+
 	cli.run()
 }
 
@@ -83,5 +87,56 @@ func (h *Handler) onMessage(msg *network.Msg) {
 		logging.Error("client %s not found", to)
 		return
 	}
+	h.msgFilter(msg)
 	cli.writeMessage(msg)
 }
+
+func (h *Handler) msgFilter(msg *network.Msg) {
+	from := msg.GetFrom()
+	to := msg.GetTo()
+	switch msg.GetXType() {
+	case network.Msg_connect_rep:
+		if msg.GetCrep().GetOk() {
+			h.RLock()
+			fromCli := h.clients[from]
+			toCli := h.clients[to]
+			h.RUnlock()
+			id := msg.GetCrep().GetId()
+			var pair [2]*client
+			if fromCli != nil {
+				fromCli.addTunnel(id)
+				pair[0] = fromCli
+			}
+			if toCli != nil {
+				toCli.addTunnel(id)
+				pair[1] = toCli
+			}
+			h.Lock()
+			h.tunnels[id] = pair
+			h.Unlock()
+		}
+	case network.Msg_disconnect:
+
+	}
+}
+
+func (h *Handler) closeAll(cli *client) {
+	tunnels := cli.getTunnels()
+	for _, t := range tunnels {
+		h.RLock()
+		pair := h.tunnels[t]
+		h.RUnlock()
+		if pair[0] != nil {
+			pair[0].close(t)
+		}
+		if pair[1] != nil {
+			pair[1].close(t)
+		}
+		h.Lock()
+		delete(h.tunnels, t)
+		h.Unlock()
+	}
+	h.Lock()
+	delete(h.clients, cli.id)
+	h.Unlock()
+}