goproxy/client/main.go

package main

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"os"
	"strings"
	"time"

	"github.com/rs/zerolog/log"
	"github.com/urfave/cli/v2"
)

type Config struct {
	AddrServer string
}

const HEART_BEAT_INTERVAL = time.Second * 5 // 心跳超时时间

const (
	MSG_TYPE_UNKOWNM = iota
	MSG_TYPE_HEARTBEAT
	MSG_TYPE_REGISTER
	MSG_TYPE_SESSION_CREATE
	MSG_TYPE_SESSION_DATA
	MSG_TYPE_SESSION_DESTORY
	MSG_TYPE_TUNNEL_CREATE
	MSG_TYPE_MAX
)

type Device struct {
	id          string
	category    string
	desc        string /* description of the device */
	conn        net.Conn
	create_time int64 /* connection time */
	active      time.Time
	uptime      uint32
	token       string
	registered  bool
	closed      uint32
	send        chan []byte // Buffered channel of outbound messages.
}

func main() {
	app := &cli.App{
		Name:    "XZRobot Ops Server",
		Usage:   "The Server Side For xzrobot ops",
		Version: "1.0.0",
		Commands: []*cli.Command{
			{
				Name:  "run",
				Usage: "Run Server",
				Flags: []cli.Flag{
					&cli.StringFlag{
						Name:  "log",
						Value: "log.txt",
						Usage: "log file path",
					},
					&cli.StringFlag{
						Name:    "conf",
						Aliases: []string{"c"},
						Value:   "./rttys.conf",
						Usage:   "config file to load",
					},
					&cli.StringFlag{
						Name:  "addr-dev",
						Value: ":9011",
						Usage: "address to listen device",
					},
					&cli.StringFlag{
						Name:  "addr-user",
						Value: ":9012",
						Usage: "address to listen user",
					},
					&cli.StringFlag{
						Name:  "db",
						Value: "sqlite://database.db",
						Usage: "database source",
					},
				},
				Action: func(c *cli.Context) error {
					runClient(c)
					return nil
				},
			},
		},
		Action: func(c *cli.Context) error {
			c.App.Command("run").Run(c)
			return nil
		},
	}

	err := app.Run(os.Args)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

func runClient(c *cli.Context) {

	cfg := &Config{
		AddrServer: "localhost:9011",
	}

	tcpConn, err := createTcpConn(cfg.AddrServer)
	if err != nil {
		fmt.Println("TCP Connect Error! " + err.Error())
		return
	}
	fmt.Println(tcpConn.LocalAddr().String() + " : Client Connected")

	reader := bufio.NewReader(tcpConn)
	for {
		s, err := reader.ReadString('\n')
		if err != nil || err == io.EOF {
			break
		} else {
			//接收到new的指令的时候，新建一个tcp连接
			if s == "new\n" {

			}
			if s == "hi" {
				//忽略掉hi的请求
			}
		}
	}
}

func readLoop(conn *net.TCPConn) {
	defer conn.Close()

	reader := bufio.NewReader(conn)

	for {
		b, err := reader.Peek(5)
		if err != nil {
			if err != io.EOF && !strings.Contains(err.Error(), "use of closed network connection") {
				log.Error().Msg(err.Error())
			}
			return
		}

		reader.Discard(5)

		msg_type := b[0]

		if msg_type >= MSG_TYPE_MAX {
			log.Error().Msgf("invalid msg type: %d", msg_type)
			return
		}

		msg_length := binary.BigEndian.Uint32(b[1:])
		data := make([]byte, msg_length)
		_, err = io.ReadFull(reader, data)
		if err != nil {
			log.Error().Msg(err.Error())
			return
		}

		// dev.active = time.Now()

		switch msg_type {
		case MSG_TYPE_HEARTBEAT:
			parseHeartbeat(b)

		default:
			log.Error().Msgf("invalid msg type: %d", msg_type)
		}
	}
}

func parseHeartbeat(b []byte) {
	uptime := binary.BigEndian.Uint32(b[:4])
	log.Info().Msgf("Heartbeat Time: %d", uptime)
}

func writeLoop(conn *net.TCPConn) {
	ticker := time.NewTicker(time.Second)

	defer conn.Close()

	ninactive := 0
	lastHeartbeat := time.Now()

	for {
		select {
		case msg, ok := <-dev.send:
			if !ok {
				return
			}

			_, err := dev.conn.Write(msg)
			if err != nil {
				log.Error().Msg(err.Error())
				return
			}

		case <-ticker.C:
			now := time.Now()
			if now.Sub(dev.active) > HEART_BEAT_INTERVAL*3/2 {
				if dev.id == "" {
					return
				}

				log.Error().Msgf("Inactive device in long time: %s", dev.id)
				if ninactive > 3 {
					log.Error().Msgf("Inactive 3 times, now kill it: %s", dev.id)
					return
				}
				ninactive = ninactive + 1
			}

			if now.Sub(lastHeartbeat) > HEART_BEAT_INTERVAL-1 {
				lastHeartbeat = now
				if len(dev.send) < 1 {
					dev.WriteMsg(MSG_TYPE_HEARTBEAT, []byte{})
				}
			}
		}
	}
}

func createTcpConn(addr string) (*net.TCPConn, error) {
	tcpConn, err := net.ResolveTCPAddr("tcp", addr)
	if err != nil {
		fmt.Println("TCP Address error ! " + err.Error())
		return nil, err
	}
	conn, err := net.DialTCP("tcp", nil, tcpConn)
	if err != nil {
		fmt.Println("TCP Connect error ! " + err.Error())
		return nil, err
	}
	return conn, nil
}

func createTunnel(localAddr string, remoteAddr string) {
	localConn, err := createTcpConn(localAddr)
	if err != nil {
		fmt.Println("Create LocalConn Error! " + err.Error())
		return
	}
	remoteConn, err := createTcpConn(remoteAddr)
	if err != nil {
		fmt.Println("Create remoteAddr Error! " + err.Error())
		return
	}

	defer localConn.Close()
	defer remoteConn.Close()

	go func() {
		_, err := io.Copy(localConn, remoteConn)
		if err != nil {
			fmt.Println("Copy error! " + err.Error())
			return
		}
	}()

	go func() {
		_, err := io.Copy(remoteConn, localConn)
		if err != nil {
			fmt.Println("Copy error! " + err.Error())
			return
		}
	}()
}