github.com/geph-official/geph2@v0.22.6-0.20210211030601-f527cb59b0df/libs/niaucchi4/n4perf/main.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"flag"
     6  	"io"
     7  	"log"
     8  	"net"
     9  	"sync/atomic"
    10  	"time"
    11  
    12  	"github.com/geph-official/geph2/libs/fastudp"
    13  	"github.com/geph-official/geph2/libs/kcp-go"
    14  	"github.com/geph-official/geph2/libs/niaucchi4"
    15  	"golang.org/x/time/rate"
    16  )
    17  
    18  func main() {
    19  	var flagClient string
    20  	var flagServer string
    21  	var flagConAlgo string
    22  	var flagLimit int
    23  	flag.StringVar(&flagClient, "c", "", "client connect")
    24  	flag.StringVar(&flagServer, "s", "", "server listen")
    25  	flag.StringVar(&flagConAlgo, "cc", "LOL", "congestion control algorithm")
    26  	flag.IntVar(&flagLimit, "l", -1, "speed limit")
    27  	flag.Parse()
    28  	kcp.CongestionControl = flagConAlgo
    29  
    30  	if flagClient == "" && flagServer == "" {
    31  		log.Fatal("must give -c or -s")
    32  	}
    33  	if flagClient != "" && flagServer != "" {
    34  		log.Fatal("cannot give both -c or -s")
    35  	}
    36  	if flagServer != "" {
    37  		mainServer(flagServer, flagLimit)
    38  	}
    39  	if flagClient != "" {
    40  		mainClient(flagClient)
    41  	}
    42  }
    43  
    44  func mainClient(dialto string) {
    45  	udpsock := niaucchi4.Wrap(func() net.PacketConn {
    46  		udpsockR, err := net.ListenPacket("udp4", "")
    47  		if err != nil {
    48  			panic(err)
    49  		}
    50  		// udpsockR.(*net.UDPConn).SetWriteBuffer(10 * 1024 * 1024)
    51  		//udpsockR.(*net.UDPConn).SetReadBuffer(10 * 1024 * 1024)
    52  		return fastudp.NewConn(udpsockR.(*net.UDPConn))
    53  	})
    54  	servAddr, err := net.ResolveUDPAddr("udp4", dialto)
    55  	if err != nil {
    56  		panic(err)
    57  	}
    58  	e2e := niaucchi4.NewE2EConn(niaucchi4.ObfsListen(nil, udpsock, true))
    59  	var sid niaucchi4.SessionAddr
    60  	e2e.SetSessPath(sid, servAddr)
    61  	kcpremote, err := kcp.NewConn2(sid, nil, 16, 16, e2e)
    62  	if err != nil {
    63  		panic(err)
    64  	}
    65  	defer kcpremote.Close()
    66  	kcpremote.SetWindowSize(10000, 10000)
    67  	kcpremote.SetNoDelay(0, 100, 64, 0)
    68  	kcpremote.SetStreamMode(true)
    69  	kcpremote.SetMtu(1200)
    70  	kcpremote.Write([]byte("HELLO"))
    71  	var kbs uint64
    72  	go func() {
    73  		buf := make([]byte, 1024)
    74  		for {
    75  			_, err := io.ReadFull(kcpremote, buf)
    76  			if err != nil {
    77  				panic(err)
    78  			}
    79  			atomic.AddUint64(&kbs, 1)
    80  		}
    81  	}()
    82  	last := uint64(0)
    83  	for {
    84  		time.Sleep(time.Second)
    85  		rn := atomic.LoadUint64(&kbs)
    86  		log.Println("Current speed:", rn-last, "KiB/s")
    87  		last = rn
    88  	}
    89  }
    90  
    91  func mainServer(listen string, klimit int) {
    92  	var limiter *rate.Limiter
    93  	if klimit > 0 {
    94  		limiter = rate.NewLimiter(rate.Limit(klimit*1024), 1024*1024)
    95  	}
    96  	udpsock, err := net.ListenPacket("udp4", listen)
    97  	if err != nil {
    98  		panic(err)
    99  	}
   100  	udpsock.(*net.UDPConn).SetWriteBuffer(10 * 1024 * 1024)
   101  	udpsock.(*net.UDPConn).SetReadBuffer(10 * 1024 * 1024)
   102  	obfs := niaucchi4.ObfsListen(nil, udpsock, true)
   103  	if err != nil {
   104  		panic(err)
   105  	}
   106  	e2e := niaucchi4.NewE2EConn(obfs)
   107  	listener, err := kcp.ServeConn(nil, 16, 16, e2e)
   108  	if err != nil {
   109  		panic(err)
   110  	}
   111  	log.Println("KCP over N4 listener spinned up!")
   112  	for {
   113  		kclient, err := listener.AcceptKCP()
   114  		if err != nil {
   115  			panic(err)
   116  		}
   117  		log.Println("Accepted kclient from", kclient.RemoteAddr())
   118  		kclient.SetWindowSize(10000, 10000)
   119  		kclient.SetNoDelay(0, 100, 64, 0)
   120  		kclient.SetStreamMode(true)
   121  		kclient.SetMtu(1200)
   122  		go func() {
   123  			defer kclient.Close()
   124  			buf := make([]byte, 5)
   125  			io.ReadFull(kclient, buf)
   126  			for {
   127  				if limiter != nil {
   128  					limiter.WaitN(context.Background(), 65536)
   129  				}
   130  				_, err := kclient.Write(make([]byte, 65536))
   131  				if err != nil {
   132  					return
   133  				}
   134  			}
   135  		}()
   136  	}
   137  }