golang实现的TCP端口转发

package main

import (
“net”
“sync”
“fmt”
“time”
“errors”
“context”
“os”
“os/signal”
“syscall”
“encoding/csv”
)

type TcpServer struct{
url string
listener net.Listener
connections map[uint64]*TcpConnection
connectionsLock *sync.RWMutex
running bool
runningLock *sync.Mutex
exit bool
Stop chan bool
remoteUrl string
}

func NewTcpServer() *TcpServer {
s := new(TcpServer)
s.connections = make(map[uint64]*TcpConnection)
s.connectionsLock = new(sync.RWMutex)
s.runningLock = new(sync.Mutex)
s.Stop = make(chan bool)
s.exit = false
return s
}

func (s *TcpServer) SetListenUrl(url string){
s.url = url
}
func (s *TcpServer) SetRemoteUrl(url string){
s.remoteUrl = url
}
func (s *TcpServer) Start() error {
s.runningLock.Lock()
alreadyRunning := s.running
if alreadyRunning {
s.runningLock.Unlock()
return errors.New(“netServer already started”)
}
ln, err := net.Listen(“tcp”, s.url)
//nln:= netutil.LimitListener(ln,10000)
fmt.Println(“netServer listen:”,s.url)
if err != nil {
return err
}
s.listener = ln
s.running = true
s.runningLock.Unlock()
go s.listen()
return nil
}

func (s TcpServer) listen() {
ticker := time.NewTicker(time.Second * 1)
stopTickGoroutine := make(chan bool)
go func() {
for {
select {
case <-stopTickGoroutine:
ticker.Stop()
return
case <-ticker.C:
for _, conn := range s.connections {
if conn.closed {
s.closeConnection(conn)
} else {
}
}
case e := <-s.Stop:
if e {
s.exit = true
s.listener.Close()
break
}
}
}
}()
var cid uint64
for !s.exit {
conn, err := s.listener.Accept()
if err != nil {
} else {
go s.accept(cid, conn).serve()
cid++
time.Sleep(time.Microsecond100)
}
}
stopTickGoroutine <- true
}

func (s *TcpServer) accept(cid uint64, conn net.Conn) *TcpConnection {
c := new(TcpConnection)
c.cid = cid
c.conn = conn
c.buffer = make([]byte, 0, 0)
c.bufferLock = new(sync.Mutex)
c.lastCommTime = time.Now().Unix()
c.remoteUrl = s.remoteUrl
c.ctx, c.cancel = context.WithCancel(context.Background())
s.connectionsLock.Lock()
s.connections[cid] = c
s.connectionsLock.Unlock()
return c
}

func (s *TcpServer) closeConnection(c *TcpConnection) {
s.connectionsLock.Lock()
defer s.connectionsLock.Unlock()
delete(s.connections, c.cid)
}

type TcpConnection struct {
cid uint64
conn net.Conn
len int
closed bool
buffer []byte
bufferLock *sync.Mutex
lastCommTime int64
Send chan []byte
ctx context.Context
cancel context.CancelFunc
remoteUrl string
remoteClient * TcpClient
}

func (c *TcpConnection) serve() (err error) {
defer func() {
c.cancel()
if x := recover(); x != nil {
err = fmt.Errorf(“TcpConnection net conn error: %v”, x)
return
}
}()
c.remoteClient= NewTcpClient(c,c.remoteUrl)
err=c.remoteClient.Dail()
if err!= nil{
c.conn.Close()
fmt.Println(“remote url not connect”,c.remoteUrl)
return err
}
go c.readConn()
for !c.closed {
select {
case <-c.ctx.Done():
c.closed = true
break
}
}
if c.conn != nil {
c.conn.Close()
}
c.remoteClient.Close()
return
}

func (c *TcpConnection) readConn() {
defer func() {
c.cancel()
}()
for {
select {
case <-c.ctx.Done():
return
default:
//
}
data := make([]byte, 4096)
n, err := c.conn.Read(data)
if err != nil {
break
}
c.len = n
c.remoteClient.Send(data[:n])
//fmt.Println(“local recv:”,n)
}
fmt.Println(time.Now().String(),“local recv end”)
}

func (c *TcpConnection) Write(d []byte) error{
_, err := c.conn.Write(d)
if err != nil {
c.cancel()
return err
}
return nil
}

func (c *TcpConnection) UnderlyingConn() net.Conn{
return c.conn
}

type TcpClient struct{
Url string
Conn net.Conn
SendC chan string
Quit chan bool
Reconn chan struct{}
Status bool
localConn *TcpConnection
}

func NewTcpClient(conn *TcpConnection ,url string) (tc * TcpClient) {
return &TcpClient{
Url:url,
SendC:make(chan string ,1024),
Quit:make(chan bool,16),
Reconn:make(chan struct{},8),
localConn:conn,
}
}

func (c *TcpClient)Dail()(err error){
conn, err := net.Dial(“tcp”, c.Url)
if err != nil{
return
}
if c.Conn != nil{c.Conn.Close()}
c.Conn = conn
c.Status = true
go c.Recv()
return nil
}

func (c *TcpClient) Recv(){
defer func() {
c.localConn.cancel()
}()
for {
select {
case <-c.localConn.ctx.Done():
return
default:
//
}
data := make([]byte, 4096)
n, err := c.Conn.Read(data)
if err != nil {
c.Status=false
break
}
c.localConn.Write(data[:n])
//fmt.Println(“remote recv:”,n)
}
fmt.Println(time.Now().String(),“remote recv end”)
}

func (c *TcpClient)Send(b []byte){
if c.Status{
_,err:=c.Conn.Write(b)
if err!= nil{
c.Status=false
c.localConn.cancel()
}
}
}

func (c *TcpClient)Close(){
if c.Conn!= nil{
c.Conn.Close()
}
}

func ReadCsv(fp string)(head []string,data [][]string,err error ){
f, err := os.Open(fp)
if err != nil{
return
}
defer f.Close()
w := csv.NewReader(f)
dts, err := w.ReadAll()
if err != nil{
return
}
head=dts[0]
data=dts[1:]
return
}

func ReadCSV(fp string)(data [][]string,err error ){
f, err := os.Open(fp)
if err != nil{
return
}
defer f.Close()
w := csv.NewReader(f)
data, err = w.ReadAll()
if err != nil{
return
}
return
}
func main() {
var mps map[int]*TcpServer=make(map[int]*TcpServer)
pmaps,err:=ReadCSV(“portmap.csv”)
if err!= nil{
fmt.Println(err)
}
if len(pmaps)>1 {
//fmt.Println(pmaps)
for i,pm:=range(pmaps[1:]){
if len(pm)<2{continue}
t:=NewTcpServer()
t.SetListenUrl(pm[0])
t.SetRemoteUrl(pm[1])
t.Start()
mps[i] = t
}
}
c := make(chan os.Signal, 1)
signal.Notify(c, os.Interrupt)
signal.Notify(c, syscall.SIGTERM)
fmt.Println(“WMC portMap starting”)
for{
select {
case ext:=<-c:
fmt.Println(“exit :”,ext)
os.Exit(0)
}
}
}

创建portmap.csv
local,remote
:11001,192.168.139.233:11001
:11002,192.168.139.233:11002