diff --git a/code/client/client/client.go b/code/client/client/client.go index adb7c0a..ddfde57 100644 --- a/code/client/client/client.go +++ b/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()) +} diff --git a/code/client/client/tunnel.go b/code/client/client/tunnel.go index 6831275..11bf689 100644 --- a/code/client/client/tunnel.go +++ b/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 +} diff --git a/code/server/handler/client.go b/code/server/handler/client.go index 9a01cc6..23702d8 100644 --- a/code/server/handler/client.go +++ b/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) +} diff --git a/code/server/handler/handler.go b/code/server/handler/handler.go index e44d7b7..27690ec 100644 --- a/code/server/handler/handler.go +++ b/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) +}