You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
344 lines
8.1 KiB
Go
344 lines
8.1 KiB
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
|
|
}
|
|
}
|
|
}()
|
|
}
|