transx/tcp.go

package main

import (
	"errors"
	"fmt"
	"github.com/TransX/log"
	"net"
	"os"
)

type TransTCP struct {
}

func NewTransTCP() *TransTCP {

	return &TransTCP{}
}

func (this *TransTCP) createTCPClient(ip, port string) (conn net.Conn, err error) {
	conn, err = net.Dial("tcp4", ip+":"+port)
	if err != nil {
		conn = nil
	}
	return
}

func (this *TransTCP) createTCPClientWithRetry(ip, port string, retry int) (conn net.Conn, err error) {
	for i := 0; i < retry; i++ {
		c, e := this.createTCPClient(ip, port)
		if e == nil {
			return c, e
		}
		log.Error("Create Client Error: %s", e.Error())
	}
	//failed with retry
	return nil, errors.New(fmt.Sprintln("failed to create client after %d retry", retry))
}

func (this *TransTCP) createTCPListener(ip, port string) (listen net.Listener, err error) {
	listener, _err := net.Listen("tcp4", ip+":"+port)
	if _err == nil {
		listen = listener
		err = nil
		return
	} else {
		listen = nil
		err = _err
	}
	return
}

func (this *TransTCP) Start(listenPort, destIP, destPort string, clientOrServer string) {
	listener, err := this.createTCPListener("0.0.0.0", listenPort)
	if err != nil {
		log.Panic("Failed to create listener. %s", err)
		os.Exit(0)
	}
	for {
		if listenerConn, err := listener.Accept(); err == nil {
			go func() {
				log.Info("Incoming %s", listenerConn.RemoteAddr().String())
				//创建到目标的连接
				destConn, err := this.createTCPClientWithRetry(destIP, destPort, 3)
				if err != nil {
					log.Panic("Failed to connect to destination. %s", err)
					os.Exit(0)
				}
				log.Info("Dial %s", destConn.RemoteAddr().String())
				//tunnel model : [ -->>server ---- client -->> ](this is a tunnel)
				if clientOrServer == "client" {
					go NewTunnel(listenerConn, destConn, SEND).run()
					go NewTunnel(destConn, listenerConn, RECEIVE).run()
				}
				if clientOrServer == "server" {
					go NewTunnel(listenerConn, destConn, RECEIVE).run()
					go NewTunnel(destConn, listenerConn, SEND).run()
				}
			}()

		} else {
			log.Info("Failed to accept incoming connection. %s", err)
		}
	}
}