实现forward逻辑

code/client/client/client.go

@@ -1,6 +1,7 @@
 package client
 
 import (
+	"fmt"
 	"natpass/code/client/global"
 	"natpass/code/network"
 	"net"
@@ -49,7 +50,12 @@ func (c *Client) Run() {
 			logging.Error("read message: %v", err)
 			return
 		}
-		logging.Info("read message: %s", msg.String())
+		switch msg.GetXType() {
+		case network.Msg_connect_req:
+			c.handleConnect(msg.GetFrom(), msg.GetTo(), msg.GetCreq())
+		case network.Msg_forward:
+			c.handleData(msg.GetXData())
+		}
 	}
 }
 
@@ -84,12 +90,22 @@ func (c *Client) handleTcpTunnel(t global.Tunnel) {
 			logging.Error("accept from %s tunnel, err=%v", t.Name, err)
 			continue
 		}
-		tn := newTunnel(c, conn)
+
+		id, err := runtime.UUID(16, "0123456789abcdef")
+		if err != nil {
+			conn.Close()
+			logging.Error("generate tunnel id failed, err=%v", err)
+			continue
+		}
+
+		tn := newTunnel(id, t.Target, c, conn)
 		c.sendConnect(tn.id, t)
 
 		c.Lock()
 		c.tunnels[tn.id] = tn
 		c.Unlock()
+
+		go tn.loop()
 	}
 }
 
@@ -117,3 +133,76 @@ func (c *Client) sendConnect(id string, t global.Tunnel) {
 	}
 	c.conn.WriteMessage(&msg, 5*time.Second)
 }
+
+func (c *Client) handleConnect(from, to string, req *network.ConnectRequest) {
+	dial := "tcp"
+	if req.GetXType() == network.ConnectRequest_udp {
+		dial = "udp"
+	}
+	conn, err := net.Dial(dial, fmt.Sprintf("%s:%d", req.GetAddr(), req.GetPort()))
+	if err != nil {
+		c.connectError(to, req.GetId(), err.Error())
+		return
+	}
+	tn := newTunnel(req.GetId(), from, c, conn)
+	c.Lock()
+	c.tunnels[tn.id] = tn
+	c.Unlock()
+	c.connectOK(to, req.GetId())
+	go tn.loop()
+}
+
+func (c *Client) connectError(to, id, m string) {
+	var msg network.Msg
+	msg.From = c.cfg.ID
+	msg.To = to
+	msg.XType = network.Msg_connect_rep
+	msg.Payload = &network.Msg_Crep{
+		Crep: &network.ConnectResponse{
+			Id:  id,
+			Ok:  false,
+			Msg: m,
+		},
+	}
+	c.conn.WriteMessage(&msg, time.Second)
+}
+
+func (c *Client) connectOK(to, id string) {
+	var msg network.Msg
+	msg.From = c.cfg.ID
+	msg.To = to
+	msg.XType = network.Msg_connect_rep
+	msg.Payload = &network.Msg_Crep{
+		Crep: &network.ConnectResponse{
+			Id: id,
+			Ok: true,
+		},
+	}
+	c.conn.WriteMessage(&msg, time.Second)
+}
+
+func (c *Client) send(id, target string, data []byte) {
+	var msg network.Msg
+	msg.From = c.cfg.ID
+	msg.To = target
+	msg.XType = network.Msg_forward
+	msg.Payload = &network.Msg_XData{
+		XData: &network.Data{
+			Cid:  id,
+			Data: data,
+		},
+	}
+	c.conn.WriteMessage(&msg, time.Second)
+}
+
+func (c *Client) handleData(data *network.Data) {
+	id := data.GetCid()
+	c.RLock()
+	tn := c.tunnels[id]
+	c.RUnlock()
+	if tn == nil {
+		logging.Error("tunnel %s not found", id)
+		return
+	}
+	tn.write(data.GetData())
+}

code/client/client/tunnel.go

@@ -1,23 +1,47 @@
 package client
 
 import (
+	"bytes"
+	"io"
 	"net"
-
-	"github.com/lwch/runtime"
 )
 
 type tunnel struct {
-	cli     *Client
-	id      string
-	forward net.Conn
+	cli    *Client
+	id     string
+	target string
+	c      net.Conn
 }
 
-func newTunnel(cli *Client, conn net.Conn) *tunnel {
-	id, err := runtime.UUID(16, "0123456789abcdef")
-	runtime.Assert(err)
+func newTunnel(id, target string, cli *Client, conn net.Conn) *tunnel {
 	return &tunnel{
-		cli:     cli,
-		id:      id,
-		forward: conn,
+		cli:    cli,
+		id:     id,
+		target: target,
+		c:      conn,
 	}
 }
+
+func (tn *tunnel) close() {
+	tn.c.Close()
+}
+
+func (tn *tunnel) loop() {
+	defer tn.close()
+	buf := make([]byte, 32*1024)
+	for {
+		n, err := tn.c.Read(buf)
+		if err != nil {
+			return
+		}
+		if n == 0 {
+			continue
+		}
+		tn.cli.send(tn.id, tn.target, buf[:n])
+	}
+}
+
+func (tn *tunnel) write(data []byte) error {
+	_, err := io.Copy(tn.c, bytes.NewReader(data))
+	return err
+}

code/server/handler/client.go

@@ -9,14 +9,16 @@ import (
 )
 
 type client struct {
-	id string
-	c  *network.Conn
+	parent *Handler
+	id     string
+	c      *network.Conn
 }
 
-func newClient(id string, conn *network.Conn) *client {
+func newClient(parent *Handler, id string, conn *network.Conn) *client {
 	return &client{
-		id: id,
-		c:  conn,
+		parent: parent,
+		id:     id,
+		c:      conn,
 	}
 }
 
@@ -30,6 +32,10 @@ func (c *client) run() {
 			logging.Error("read message from %s: %v", c.id, err)
 			return
 		}
-		logging.Info("read message: %s", msg.String())
+		c.parent.onMessage(msg)
 	}
 }
+
+func (c *client) writeMessage(msg *network.Msg) error {
+	return c.c.WriteMessage(msg, time.Second)
+}

code/server/handler/handler.go

@@ -26,8 +26,13 @@ func New(cfg *global.Configure) *Handler {
 
 func (h *Handler) Handle(conn net.Conn) {
 	c := network.NewConn(conn)
-	defer c.Close()
 	var id string
+	defer func() {
+		if len(id) > 0 {
+			logging.Info("%s disconnected", id)
+		}
+		c.Close()
+	}()
 	var err error
 	for i := 0; i < 10; i++ {
 		id, err = h.readHandshake(c)
@@ -46,7 +51,7 @@ func (h *Handler) Handle(conn net.Conn) {
 	}
 	logging.Info("%s connected", id)
 
-	cli := newClient(id, c)
+	cli := newClient(h, id, c)
 	h.Lock()
 	h.clients[cli.id] = cli
 	h.Unlock()
@@ -68,3 +73,15 @@ func (h *Handler) readHandshake(c *network.Conn) (string, error) {
 	}
 	return msg.GetFrom(), nil
 }
+
+func (h *Handler) onMessage(msg *network.Msg) {
+	to := msg.GetTo()
+	h.RLock()
+	cli := h.clients[to]
+	h.RUnlock()
+	if cli == nil {
+		logging.Error("client %s not found", to)
+		return
+	}
+	cli.writeMessage(msg)
+}