nat function improve
parent
d2feccf3b3
commit
1276d3b6dd
@ -0,0 +1,23 @@
|
||||
package net
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestNat(t *testing.T) {
|
||||
var s = NatServer{
|
||||
ListenAddr: "0.0.0.0:10020",
|
||||
enableTCP: true,
|
||||
}
|
||||
var c = NatClient{
|
||||
ServiceTarget: "139.199.163.65:80",
|
||||
CmdTarget: "127.0.0.1:10020",
|
||||
enableTCP: true,
|
||||
}
|
||||
go s.Run()
|
||||
go c.Run()
|
||||
for {
|
||||
time.Sleep(time.Second * 20)
|
||||
}
|
||||
}
|
@ -1,27 +1,138 @@
|
||||
package net
|
||||
|
||||
import (
|
||||
"b612.me/starlog"
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"io"
|
||||
"net"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type SimpleNatClient struct {
|
||||
type NatClient struct {
|
||||
mu sync.RWMutex
|
||||
cmdTCPConn net.Conn
|
||||
cmdUDPConn *net.UDPAddr
|
||||
ServiceTarget string
|
||||
CmdTarget string
|
||||
tcpAlived bool
|
||||
DialTimeout int
|
||||
enableTCP bool
|
||||
enableUDP bool
|
||||
Passwd string
|
||||
stopCtx context.Context
|
||||
stopFn context.CancelFunc
|
||||
}
|
||||
|
||||
func (s *SimpleNatClient) tcpCmdConn() net.Conn {
|
||||
func (s *NatClient) tcpCmdConn() net.Conn {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.cmdTCPConn
|
||||
}
|
||||
|
||||
func (s *SimpleNatClient) tcpCmdConnAlived() bool {
|
||||
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) Run() {
|
||||
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]
|
||||
}
|
||||
if s.enableTCP {
|
||||
s.runTcp()
|
||||
}
|
||||
}
|
||||
|
||||
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) 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 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)
|
||||
return
|
||||
}
|
||||
go io.Copy(cconn, nconn)
|
||||
go io.Copy(nconn, cconn)
|
||||
}
|
||||
|
Loading…
Reference in New Issue