natpass/code/client/client/client.go

package client

import (
	"context"
	"fmt"
	"natpass/code/client/global"
	"natpass/code/network"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/lwch/logging"
	"github.com/lwch/runtime"
)

type Client struct {
	sync.RWMutex
	cfg     *global.Configure
	conn    *network.Conn
	tunnels map[string]*tunnel
}

// New create client
func New(cfg *global.Configure, conn *network.Conn) *Client {
	return &Client{
		cfg:     cfg,
		conn:    conn,
		tunnels: make(map[string]*tunnel),
	}
}

// Run main loop
func (c *Client) Run() {
	err := c.writeHandshake()
	runtime.Assert(err)
	logging.Info("%s connected", c.cfg.Server)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	go c.keepalive(ctx)

	for _, t := range c.cfg.Tunnels {
		if t.Type == "tcp" {
			go c.handleTcpTunnel(ctx, t)
		} else {
			go c.handleUdpTunnel(ctx, t)
		}
	}

	for {
		msg, err := c.conn.ReadMessage(time.Second)
		if err != nil {
			if strings.Contains(err.Error(), "i/o timeout") {
				continue
			}
			logging.Error("read message: %v", err)
			return
		}
		switch msg.GetXType() {
		case network.Msg_connect_req:
			c.handleConnect(ctx, msg.GetFrom(), msg.GetTo(), msg.GetCreq())
		case network.Msg_disconnect:
			c.handleDisconnect(msg.GetXDisconnect())
		case network.Msg_forward:
			c.handleData(msg.GetXData())
		}
	}
}

// handleTcpTunnel local listen to tcp tunnel
func (c *Client) handleTcpTunnel(ctx context.Context, t global.Tunnel) {
	defer func() {
		if err := recover(); err != nil {
			logging.Error("close tcp tunnel: %s, err=%v", t.Name, err)
		}
	}()
	l, err := net.ListenTCP("tcp", &net.TCPAddr{
		IP:   net.ParseIP(t.LocalAddr),
		Port: int(t.LocalPort),
	})
	runtime.Assert(err)
	defer l.Close()
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}
		conn, err := l.Accept()
		if err != nil {
			logging.Error("accept from %s tunnel, err=%v", t.Name, err)
			continue
		}

		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.Name, t.Target, c, conn)
		c.sendConnect(tn.id, t)

		c.Lock()
		c.tunnels[tn.id] = tn
		c.Unlock()

		go tn.loop(ctx)
	}
}

func (c *Client) keepalive(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}
		time.Sleep(10 * time.Second)
		c.sendKeepalive()
	}
}

// handleUdpTunnel local listen to udp tunnel
func (c *Client) handleUdpTunnel(ctx context.Context, t global.Tunnel) {
	// TODO
}

// handleConnect handle connect request message from remote, local dial to remomte addr
func (c *Client) handleConnect(ctx context.Context, 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(), req.GetName(), from, c, conn)
	c.Lock()
	c.tunnels[tn.id] = tn
	c.Unlock()
	c.connectOK(to, req.GetId())
	go tn.loop(ctx)
}

// handleDisconnect handle disconnect message from remote, this means remote connection is closed
func (c *Client) handleDisconnect(data *network.Disconnect) {
	id := data.GetId()

	c.RLock()
	tn := c.tunnels[id]
	c.RUnlock()

	if tn != nil {
		tn.close()

		c.Lock()
		delete(c.tunnels, id)
		c.Unlock()
	}
}

// handleData handle forward data message, write data to local connection
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())
}