star/net/natclient.go

package net

import (
	"b612.me/starlog"
	"bytes"
	"context"
	"crypto/sha256"
	"io"
	"net"
	"sync"
	"time"
)

type NatClient struct {
	mu            sync.RWMutex
	cmdTCPConn    net.Conn
	cmdUDPConn    *net.UDPConn
	ServiceTarget string
	CmdTarget     string
	tcpAlived     bool
	DialTimeout   int
	UdpTimeout    int
	enableTCP     bool
	enableUDP     bool
	Passwd        string
	udpAlived     bool
	stopCtx       context.Context
	stopFn        context.CancelFunc
}

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

func (s *NatClient) udpCmdConn() *net.UDPConn {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.cmdUDPConn
}

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

func (s *NatClient) setTcpCmdConnAlived(v bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.tcpAlived = v
}

func (s *NatClient) udpCmdConnAlived() bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.udpAlived
}

func (s *NatClient) setUdpCmdConnAlived(v bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.udpAlived = v
}

func (s *NatClient) Run() error {
	s.stopCtx, s.stopFn = context.WithCancel(context.Background())
	if s.DialTimeout == 0 {
		s.DialTimeout = 10000
	}
	if s.Passwd != "" {
		MSG_CMD_HELLO = sha256.New().Sum(append(MSG_CMD_HELLO, []byte(s.Passwd)...))[:16]
	}
	var wg sync.WaitGroup
	if s.enableUDP {
		wg.Add(1)
		go func() {
			defer wg.Done()
			s.runUdp()
		}()
	}
	if s.enableTCP {
		wg.Add(1)
		go func() {
			defer wg.Done()
			s.runTcp()
		}()
	}
	wg.Wait()
	return nil
}

func (s *NatClient) runTcp() error {
	var err error
	starlog.Noticeln("nat client tcp module start run")
	for {
		select {
		case <-s.stopCtx.Done():
			if s.cmdTCPConn != nil {
				s.setTcpCmdConnAlived(false)
				s.cmdTCPConn.Close()
				return nil
			}
		case <-time.After(time.Millisecond * 1500):
		}
		if s.cmdTCPConn != nil && s.tcpCmdConnAlived() {
			continue
		}
		s.cmdTCPConn, err = net.DialTimeout("tcp", s.CmdTarget, time.Millisecond*time.Duration(s.DialTimeout))
		if err != nil {
			starlog.Errorf("dail remote tcp cmd server %v fail:%v;will retry\n", s.CmdTarget, err)
			time.Sleep(time.Second * 2)
			s.cmdTCPConn = nil
			continue
		}
		starlog.Infoln("dail remote tcp cmd server ok,remote:", s.CmdTarget)
		s.tcpCmdConn().Write(MSG_CMD_HELLO)
		s.setTcpCmdConnAlived(true)
		go s.handleTcpCmdConn(s.tcpCmdConn())
	}
}

func (s *NatClient) runUdp() error {
	starlog.Noticeln("nat client udp module start run")
	if s.UdpTimeout == 0 {
		s.UdpTimeout = 600000
	}
	for {
		select {
		case <-s.stopCtx.Done():
			if s.cmdTCPConn != nil {
				s.setUdpCmdConnAlived(false)
				s.cmdUDPConn.Close()
				return nil
			}
		case <-time.After(time.Millisecond * 3000):
		}
		if s.cmdUDPConn != nil && s.udpCmdConnAlived() {
			continue
		}
		rmt, err := net.ResolveUDPAddr("udp", s.CmdTarget)
		if err != nil {
			starlog.Errorf("dail remote udp cmd server %v fail:%v;will retry\n", s.CmdTarget, err)
			time.Sleep(time.Second * 2)
			continue
		}
		s.cmdUDPConn, err = net.DialUDP("udp", nil, rmt)
		if err != nil {
			starlog.Errorf("dail remote udp cmd server %v fail:%v;will retry\n", s.CmdTarget, err)
			time.Sleep(time.Second * 2)
			s.cmdTCPConn = nil
			continue
		}
		starlog.Infoln("dail remote udp cmd server ok,remote:", s.CmdTarget)
		s.udpCmdConn().Write(MSG_CMD_HELLO)
		s.setUdpCmdConnAlived(true)
		go s.handleUdpCmdConn(s.udpCmdConn())
	}
}
func (s *NatClient) handleUdpCmdConn(conn *net.UDPConn) {
	for {
		header := make([]byte, 16)
		_, err := io.ReadFull(conn, header)
		if err != nil {
			starlog.Infoln("udp cmd server read fail:", err)
			conn.Close()
			s.setUdpCmdConnAlived(false)
			return
		}
		if bytes.Equal(header, MSG_CMD_HELLO_REPLY) {
			continue
		}
		if bytes.Equal(header, MSG_NEW_CONN_HELLO) {
			go s.newRemoteUdpConn()
		}
		if bytes.Equal(header, MSG_HEARTBEAT) {
			_, err = conn.Write(MSG_HEARTBEAT)
			if err != nil {
				conn.Close()
				s.setUdpCmdConnAlived(false)
				return
			}
		}
	}
}

func (s *NatClient) handleTcpCmdConn(conn net.Conn) {
	for {
		header := make([]byte, 16)
		_, err := io.ReadFull(conn, header)
		if err != nil {
			starlog.Infoln("tcp cmd server read fail:", err)
			conn.Close()
			s.setTcpCmdConnAlived(false)
			return
		}
		if bytes.Equal(header, MSG_CMD_HELLO_REPLY) {
			continue
		}
		if bytes.Equal(header, MSG_NEW_CONN_HELLO) {
			go s.newRemoteTcpConn()
		}
		if bytes.Equal(header, MSG_HEARTBEAT) {
			_, err = conn.Write(MSG_HEARTBEAT)
			if err != nil {
				conn.Close()
				s.setTcpCmdConnAlived(false)
				return
			}
		}
	}
}

func (s *NatClient) newRemoteTcpConn() {
	log := starlog.Std.NewFlag()
	starlog.Infoln("recv request,create new tcp conn")
	nconn, err := net.DialTimeout("tcp", s.CmdTarget, time.Millisecond*time.Duration(s.DialTimeout))
	if err != nil {
		log.Errorf("dail server tcp conn %v fail:%v\n", s.CmdTarget, err)
		return
	}
	_, err = nconn.Write(MSG_NEW_CONN_HELLO)
	if err != nil {
		nconn.Close()
		log.Errorf("write new tcp client hello to server %v fail:%v\n", s.CmdTarget, err)
		return
	}
	cconn, err := net.DialTimeout("tcp", s.ServiceTarget, time.Millisecond*time.Duration(s.DialTimeout))
	if err != nil {
		log.Errorf("dail remote tcp conn %v fail:%v\n", s.CmdTarget, err)
		nconn.Close()
		return
	}
	go func() {
		for {
			data := make([]byte, 8192)
			nconn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(s.UdpTimeout)))
			n, err := nconn.Read(data)
			if err != nil {
				starlog.Infoln("read from tcp server fail:", nconn.RemoteAddr(), err)
				nconn.Close()
				cconn.Close()
				return
			}
			_, err = cconn.Write(data[:n])
			//starlog.Debugln("write to udp client:", p, err, cconn.LocalAddr(), cconn.RemoteAddr())
			if err != nil {
				starlog.Infoln("write to tcp client fail:", cconn.RemoteAddr(), err)
				nconn.Close()
				cconn.Close()
				return
			}
		}
	}()
	go func() {
		for {
			data := make([]byte, 8192)
			cconn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(s.UdpTimeout)))
			n, err := cconn.Read(data)
			if err != nil {
				starlog.Infoln("read from tcp server fail:", cconn.RemoteAddr(), err)
				nconn.Close()
				cconn.Close()
				return
			}
			_, err = nconn.Write(data[:n])
			if err != nil {
				starlog.Infoln("write to tcp client fail:", nconn.RemoteAddr(), err)
				nconn.Close()
				cconn.Close()
				return
			}
		}
	}()
}

func (s *NatClient) newRemoteUdpConn() {
	log := starlog.Std.NewFlag()
	starlog.Infoln("recv request,create new udp conn")
	rmt, err := net.ResolveUDPAddr("udp", s.CmdTarget)
	if err != nil {
		log.Errorf("dail server udp conn %v fail:%v\n", s.CmdTarget, err)
		return
	}
	nconn, err := net.DialUDP("udp", nil, rmt)
	if err != nil {
		log.Errorf("dail server udp conn %v fail:%v\n", s.CmdTarget, err)
		return
	}
	log.Infof("dail server udp conn %v ok\n", s.CmdTarget)
	_, err = nconn.Write(MSG_NEW_CONN_HELLO)
	if err != nil {
		nconn.Close()
		log.Errorf("write new udp client hello to server %v fail:%v\n", s.CmdTarget, err)
		return
	}

	rmt, err = net.ResolveUDPAddr("udp", s.ServiceTarget)
	if err != nil {
		log.Errorf("dail server udp conn %v fail:%v\n", s.ServiceTarget, err)
		return
	}
	cconn, err := net.DialUDP("udp", nil, rmt)
	if err != nil {
		log.Errorf("dail remote udp conn %v fail:%v\n", s.ServiceTarget, err)
		return
	}
	log.Infof("dail remote udp conn %v ok\n", s.ServiceTarget)
	go func() {
		for {
			data := make([]byte, 8192)
			nconn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(s.UdpTimeout)))
			n, err := nconn.Read(data)
			if err != nil {
				starlog.Infoln("read from udp server fail:", err)
				return
			}
			_, err = cconn.Write(data[:n])
			//starlog.Debugln("write to udp client:", p, err, cconn.LocalAddr(), cconn.RemoteAddr())
			if err != nil {
				starlog.Infoln("write to udp client fail:", err)
				return
			}
		}
	}()
	go func() {
		for {
			data := make([]byte, 8192)
			cconn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(s.UdpTimeout)))
			n, err := cconn.Read(data)
			if err != nil {
				starlog.Infoln("read from udp server fail:", err)
				return
			}
			_, err = nconn.Write(data[:n])
			if err != nil {
				starlog.Infoln("write to udp client fail:", err)
				return
			}
		}
	}()
}