star/net/natserver.go

package net

import (
	"b612.me/starlog"
	"bytes"
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
	"sync"
	"time"
)

var MSG_CMD_HELLO = []byte{11, 27, 19, 96, 182, 18, 25, 150, 17, 39}
var MSG_NEW_CONN = []byte{0, 0, 0, 0, 255, 255, 255, 255, 11, 27}
var MSG_NEW_CONN_REQ = []byte{0, 0, 0, 0, 255, 255, 255, 255, 19, 96}
var MSG_CLOSE = []byte{255, 255, 0, 0, 255, 0, 0, 255, 255, 27}
var MSG_HEARTBEAT = []byte{6, 66, 66, 6, 6, 66, 6, 66, 11, 27}

type SimpleNatServer struct {
	mu           sync.RWMutex
	cmdTCPConn   net.Conn
	cmdUDPConn   *net.UDPAddr
	listenTcp    net.Listener
	listenUDP    *net.UDPConn
	Addr         string
	Port         int
	lastTCPHeart int64
	lastUDPHeart int64
	Passwd       string
	DialTimeout  int64
	UDPTimeout   int64
	running      int32

	tcpConnPool chan net.Conn
	tcpAlived   bool
}

func (s *SimpleNatServer) getConnfromTCPPool() (net.Conn, error) {
	select {
	case conn := <-s.tcpConnPool:
		return conn, nil
	case <-time.After(time.Second * 10):
		return nil, errors.New("no connection got")
	}
}

func (s *SimpleNatServer) tcpCmdConn() net.Conn {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.cmdTCPConn
}

func (s *SimpleNatServer) tcpCmdConnAlived() bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.tcpAlived
}

func (s *SimpleNatServer) listenTCP() error {
	var err error
	s.tcpConnPool = make(chan net.Conn, 10)
	s.listenTcp, err = net.Listen("tcp", fmt.Sprintf("%s:d", s.Addr, s.Port))
	if err != nil {
		starlog.Errorln("failed to listen tcp", err)
		return err
	}
	for {
		conn, err := s.listenTcp.Accept()
		if err != nil {
			continue
		}
		if s.tcpCmdConnAlived() {
			go s.tcpClientServe(conn.(*net.TCPConn))
			continue
		}
		go s.waitingForTCPCmd(conn.(*net.TCPConn))
	}
	return nil
}

func (s *SimpleNatServer) tcpClientServe(conn *net.TCPConn) {
	if !s.tcpCmdConnAlived() {
		conn.Close()
		return
	}

	if strings.Split(conn.RemoteAddr().String(), ":")[0] == strings.Split(s.tcpCmdConn().RemoteAddr().String(), ":")[0] {
		conn.SetReadDeadline(time.Now().Add(5 * time.Second))
		cmdBuf := make([]byte, 10)
		if _, err := io.ReadFull(conn, cmdBuf); err == nil {
			conn.SetReadDeadline(time.Time{})
			if bytes.Equal(cmdBuf, MSG_NEW_CONN) {
				starlog.Noticef("Nat Server Recv New Client Conn From %v\n", conn.RemoteAddr().String())
				s.tcpConnPool <- conn
				return
			}
		}
		conn.SetReadDeadline(time.Time{})
	}
	starlog.Noticef("Nat Server Recv New Side Conn From %v\n", conn.RemoteAddr().String())
	_, err := s.tcpCmdConn().Write(MSG_NEW_CONN_REQ)
	if err != nil {
		s.mu.Lock()
		s.cmdTCPConn.Close()
		s.tcpAlived = false
		s.mu.Unlock()
		starlog.Errorf("Failed to Write CMD To Client:%v\n", err)
		return
	}
	reverse, err := s.getConnfromTCPPool()
	if err != nil {
		starlog.Errorf("Nat Server Conn to %v Closed %v\n", conn.RemoteAddr(), err)
		conn.Close()
		return
	}
	starlog.Infof("Nat Server Conn  %v<==>%v Connected\n", conn.RemoteAddr(), reverse.RemoteAddr())
	Copy(reverse, conn)
	starlog.Warningf("Nat Server Conn  %v<==>%v Closed\n", conn.RemoteAddr(), reverse.RemoteAddr())
}

func (s *SimpleNatServer) waitingForTCPCmd(conn *net.TCPConn) {
	conn.SetReadDeadline(time.Now().Add(time.Duration(s.DialTimeout) * time.Second))
	cmdBuf := make([]byte, 10)
	if _, err := io.ReadFull(conn, cmdBuf); err != nil {
		conn.Close()
		return
	}
	if bytes.Equal(cmdBuf, MSG_CMD_HELLO) {
		s.mu.Lock()
		s.cmdTCPConn = conn
		s.tcpAlived = true
		conn.SetKeepAlive(true)
		conn.SetKeepAlivePeriod(time.Second * 20)
		s.mu.Unlock()
	}
}