github.com/codingeasygo/util@v0.0.0-20231206062002-1ce2f004b7d9/tools/transport/transport.go (about)

     1  package main
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"os"
     7  	"strings"
     8  	"sync"
     9  
    10  	"github.com/codingeasygo/util/xnet"
    11  )
    12  
    13  var waiter = sync.WaitGroup{}
    14  
    15  func main() {
    16  	if len(os.Args) < 3 {
    17  		fmt.Printf(`Usage: transport <local> <remote> <local> <remote> ...`)
    18  		return
    19  	}
    20  	mappings := os.Args[1:]
    21  	n := len(mappings) / 2
    22  	InfoLog("transport starting %v mapping by %v", n, mappings)
    23  	for i := 0; i < n; i++ {
    24  		waiter.Add(1)
    25  		go runMapping(mappings[i*2], mappings[i*2+1])
    26  	}
    27  	waiter.Wait()
    28  	InfoLog("transport is done")
    29  }
    30  
    31  func runMapping(local, remote string) {
    32  	defer waiter.Done()
    33  	InfoLog("start mapping %v to %v", local, remote)
    34  	var transporter xnet.Transporter
    35  	if strings.HasPrefix(remote, "tcp://") {
    36  		transporter = xnet.RawDialerF(net.Dial)
    37  	} else if strings.HasPrefix(remote, "ws://") || strings.HasPrefix(remote, "wss://") {
    38  		transporter = xnet.NewWebsocketDialer()
    39  	} else {
    40  		err := fmt.Errorf("not supported remote %v", remote)
    41  		ErrorLog("mapping %v to %v fail with %v", local, remote, err)
    42  		return
    43  	}
    44  	ln, err := net.Listen("tcp", local)
    45  	if err != nil {
    46  		ErrorLog("mapping %v to %v fail with %v", local, remote, err)
    47  		return
    48  	}
    49  	var conn net.Conn
    50  	for {
    51  		conn, err = ln.Accept()
    52  		if err != nil {
    53  			break
    54  		}
    55  		waiter.Add(1)
    56  		go func() {
    57  			defer waiter.Done()
    58  			InfoLog("start transport %v to %v", conn.RemoteAddr(), remote)
    59  			xerr := transporter.Transport(conn, remote)
    60  			InfoLog("stop transport %v to %v by %v", conn.RemoteAddr(), remote, xerr)
    61  		}()
    62  	}
    63  	InfoLog("mapping %v to %v is done with %v", local, remote, err)
    64  }