package network

import (
   "io"
   "log"
   "net"
)

const (
   KeepAlive     = "KEEP_ALIVE"
   NewConnection = "NEW_CONNECTION"
)

func CreateTCPListener(addr string) (*net.TCPListener, error) {
   tcpAddr, err := net.ResolveTCPAddr("tcp", addr)
   if err != nil {
       return nil, err
   }
   tcpListener, err := net.ListenTCP("tcp", tcpAddr)
   if err != nil {
       return nil, err
   }
   return tcpListener, nil
}

func CreateTCPConn(addr string) (*net.TCPConn, error) {
   tcpAddr, err := net.ResolveTCPAddr("tcp", addr)
   if err != nil {
      return nil, err
   }
   tcpListener, err := net.DialTCP("tcp",nil, tcpAddr)
   if err != nil {
      return nil, err
   }
   return tcpListener, nil
}

func Join2Conn(local *net.TCPConn, remote *net.TCPConn) {
   go joinConn(local, remote)
   go joinConn(remote, local)
}

func joinConn(local *net.TCPConn, remote *net.TCPConn) {
   defer local.Close()
   defer remote.Close()
   _, err := io.Copy(local, remote)
   if err != nil {
      log.Println("copy failed ", err.Error())
      return
   }
}

package main

import (
   "bufio"
   "io"
   "log"
   "net"

   "nat-proxy/cmd/network"
)

var (
   // 本地需要暴露的服務(wù)端口
   localServerAddr = "127.0.0.1:32768"

   remoteIP = "111.111.111.111"
   // 遠(yuǎn)端的服務(wù)控制通道，用來(lái)傳遞控制信息，如出現(xiàn)新連接和心跳
   remoteControlAddr = remoteIP + ":8009"
   // 遠(yuǎn)端服務(wù)端口，用來(lái)建立隧道
   remoteServerAddr  = remoteIP + ":8008"
)

func main() {
   tcpConn, err := network.CreateTCPConn(remoteControlAddr)
   if err != nil {
      log.Println("[連接失敗]" + remoteControlAddr + err.Error())
      return
   }
   log.Println("[已連接]" + remoteControlAddr)

   reader := bufio.NewReader(tcpConn)
   for {
      s, err := reader.ReadString('\n')
      if err != nil || err == io.EOF {
         break
      }

      // 當(dāng)有新連接信號(hào)出現(xiàn)時(shí)，新建一個(gè)tcp連接
      if s == network.NewConnection+"\n" {
         go connectLocalAndRemote()
      }
   }

   log.Println("[已斷開]" + remoteControlAddr)
}

func connectLocalAndRemote() {
   local := connectLocal()
   remote := connectRemote()

   if local != nil && remote != nil {
      network.Join2Conn(local, remote)
   } else {
      if local != nil {
         _ = local.Close()
      }
      if remote != nil {
         _ = remote.Close()
      }
   }
}

func connectLocal() *net.TCPConn {
   conn, err := network.CreateTCPConn(localServerAddr)
   if err != nil {
      log.Println("[連接本地服務(wù)失敗]" + err.Error())
   }
   return conn
}

func connectRemote() *net.TCPConn {
   conn, err := network.CreateTCPConn(remoteServerAddr)
   if err != nil {
      log.Println("[連接遠(yuǎn)端服務(wù)失敗]" + err.Error())
   }
   return conn
}

package main

import (
   "log"
   "net"
   "strconv"
   "sync"
   "time"

   "nat-proxy/cmd/network"
)

const (
   controlAddr = "0.0.0.0:8009"
   tunnelAddr  = "0.0.0.0:8008"
   visitAddr   = "0.0.0.0:8007"
)

var (
   clientConn         *net.TCPConn
   connectionPool     map[string]*ConnMatch
   connectionPoolLock sync.Mutex
)

type ConnMatch struct {
   addTime time.Time
   accept  *net.TCPConn
}

func main() {
   connectionPool = make(map[string]*ConnMatch, 32)
   go createControlChannel()
   go acceptUserRequest()
   go acceptClientRequest()
   cleanConnectionPool()
}

// 創(chuàng)建一個(gè)控制通道，用于傳遞控制消息，如：心跳，創(chuàng)建新連接
func createControlChannel() {
   tcpListener, err := network.CreateTCPListener(controlAddr)
   if err != nil {
      panic(err)
   }

   log.Println("[已監(jiān)聽]" + controlAddr)
   for {
      tcpConn, err := tcpListener.AcceptTCP()
      if err != nil {
         log.Println(err)
         continue
      }

      log.Println("[新連接]" + tcpConn.RemoteAddr().String())
      // 如果當(dāng)前已經(jīng)有一個(gè)客戶端存在，則丟棄這個(gè)鏈接
      if clientConn != nil {
         _ = tcpConn.Close()
      } else {
         clientConn = tcpConn
         go keepAlive()
      }
   }
}

// 和客戶端保持一個(gè)心跳鏈接
func keepAlive() {
   go func() {
      for {
         if clientConn == nil {
            return
         }
         _, err := clientConn.Write(([]byte)(network.KeepAlive + "\n"))
         if err != nil {
            log.Println("[已斷開客戶端連接]", clientConn.RemoteAddr())
            clientConn = nil
            return
         }
         time.Sleep(time.Second * 3)
      }
   }()
}

// 監(jiān)聽來(lái)自用戶的請(qǐng)求
func acceptUserRequest() {
   tcpListener, err := network.CreateTCPListener(visitAddr)
   if err != nil {
      panic(err)
   }
   defer tcpListener.Close()
   for {
      tcpConn, err := tcpListener.AcceptTCP()
      if err != nil {
         continue
      }
      addConn2Pool(tcpConn)
      sendMessage(network.NewConnection + "\n")
   }
}

// 將用戶來(lái)的連接放入連接池中
func addConn2Pool(accept *net.TCPConn) {
   connectionPoolLock.Lock()
   defer connectionPoolLock.Unlock()

   now := time.Now()
   connectionPool[strconv.FormatInt(now.UnixNano(), 10)] = &ConnMatch{now, accept,}
}

// 發(fā)送給客戶端新消息
func sendMessage(message string) {
   if clientConn == nil {
      log.Println("[無(wú)已連接的客戶端]")
      return
   }
   _, err := clientConn.Write([]byte(message))
   if err != nil {
      log.Println("[發(fā)送消息異常]: message: ", message)
   }
}

// 接收客戶端來(lái)的請(qǐng)求并建立隧道
func acceptClientRequest() {
   tcpListener, err := network.CreateTCPListener(tunnelAddr)
   if err != nil {
      panic(err)
   }
   defer tcpListener.Close()

   for {
      tcpConn, err := tcpListener.AcceptTCP()
      if err != nil {
         continue
      }
      go establishTunnel(tcpConn)
   }
}

func establishTunnel(tunnel *net.TCPConn) {
   connectionPoolLock.Lock()
   defer connectionPoolLock.Unlock()

   for key, connMatch := range connectionPool {
      if connMatch.accept != nil {
         go network.Join2Conn(connMatch.accept, tunnel)
         delete(connectionPool, key)
         return
      }
   }

   _ = tunnel.Close()
}

func cleanConnectionPool() {
   for {
      connectionPoolLock.Lock()
      for key, connMatch := range connectionPool {
         if time.Now().Sub(connMatch.addTime) > time.Second*10 {
            _ = connMatch.accept.Close()
            delete(connectionPool, key)
         }
      }
      connectionPoolLock.Unlock()
      time.Sleep(5 * time.Second)
   }
}