github.com/geph-official/geph2@v0.22.6-0.20210211030601-f527cb59b0df/cmd/geph-bridge/e2enat.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"encoding/binary"
     6  	"fmt"
     7  	"log"
     8  	"math/rand"
     9  	"net"
    10  	"sync"
    11  	"sync/atomic"
    12  	"time"
    13  
    14  	"github.com/ethereum/go-ethereum/rlp"
    15  	"github.com/geph-official/geph2/libs/fastudp"
    16  	"github.com/geph-official/geph2/libs/niaucchi4"
    17  	"github.com/patrickmn/go-cache"
    18  	"golang.org/x/time/rate"
    19  )
    20  
    21  type e2ePacket struct {
    22  	Session niaucchi4.SessionAddr
    23  	Sn      uint64
    24  	Ack     uint64
    25  	Body    []byte
    26  	Padding []byte
    27  }
    28  
    29  var key [32]byte
    30  
    31  func parseSess(bts []byte) uint64 {
    32  	var pkt e2ePacket
    33  	rlp.DecodeBytes(bts, &pkt)
    34  	return binary.BigEndian.Uint64(pkt.Session[:8])
    35  }
    36  
    37  var e2ecount int64
    38  
    39  func init() {
    40  	go func() {
    41  		for {
    42  			if statClient != nil {
    43  				statClient.Send(map[string]string{
    44  					allocGroup + ".e2eCount": fmt.Sprintf("%v|g", atomic.LoadInt64(&e2ecount)),
    45  				}, 1)
    46  			}
    47  			time.Sleep(time.Second * 10)
    48  		}
    49  	}()
    50  }
    51  
    52  var e2eMap = cache.New(time.Hour, time.Hour)
    53  var e2eMapLk sync.Mutex
    54  
    55  func e2enat(dest string, cookie []byte) (port int, err error) {
    56  	// e2eMapLk.Lock()
    57  	// defer e2eMapLk.Unlock()
    58  	// log.Println("e2enat", atomic.LoadInt64(&e2ecount))
    59  	// kee := fmt.Sprintf("%v/%x", dest, cookie)
    60  	// if porti, ok := e2eMap.Get(kee); ok {
    61  	// 	log.Println("HIT", kee)
    62  	// 	port = porti.(int)
    63  	// 	return
    64  	// }
    65  	// log.Println("MISS", kee)
    66  	leftRaw, err := net.ListenPacket("udp", "")
    67  	if err != nil {
    68  		return
    69  	}
    70  	leftRaw = fastudp.NewConn(leftRaw.(*net.UDPConn))
    71  	leftSock := niaucchi4.ObfsListen(cookie, leftRaw, true)
    72  	rightSock, err := net.ListenPacket("udp", "")
    73  	if err != nil {
    74  		return
    75  	}
    76  	destReal, err := net.ResolveUDPAddr("udp", dest)
    77  	if err != nil {
    78  		return
    79  	}
    80  	rightSock = fastudp.NewConn(rightSock.(*net.UDPConn))
    81  	// mapping
    82  	sessMap := new(sync.Map)
    83  	go func() {
    84  		atomic.AddInt64(&e2ecount, 1)
    85  		defer atomic.AddInt64(&e2ecount, -1)
    86  		defer leftSock.Close()
    87  		defer rightSock.Close()
    88  		bts := malloc(2048)
    89  		for {
    90  			dl := time.Now().Add(time.Minute * 30)
    91  			leftSock.SetReadDeadline(dl)
    92  			n, addr, err := leftSock.ReadFrom(bts)
    93  			if err != nil {
    94  				log.Println("closing", leftRaw.LocalAddr(), err)
    95  				return
    96  			}
    97  			sid := parseSess(bts[:n])
    98  			sessMap.Store(sid, addr)
    99  			btsCopy := malloc(n)
   100  			copy(btsCopy, bts)
   101  			maybeDoJob(func() {
   102  				_, err = rightSock.WriteTo(btsCopy, destReal)
   103  				if err != nil {
   104  					log.Println("cannot write:", err)
   105  				}
   106  				if statClient != nil && rand.Int()%100000 < n {
   107  					statClient.Increment(allocGroup + ".e2eup")
   108  				}
   109  				free(btsCopy)
   110  			})
   111  		}
   112  	}()
   113  	go func() {
   114  		defer leftSock.Close()
   115  		defer rightSock.Close()
   116  		bts := malloc(2048)
   117  		for {
   118  			dl := time.Now().Add(time.Minute * 30)
   119  			rightSock.SetReadDeadline(dl)
   120  			n, _, e := rightSock.ReadFrom(bts)
   121  			if e != nil {
   122  				log.Println("closing", rightSock.LocalAddr(), err)
   123  				return
   124  			}
   125  			leftSock.SetWriteDeadline(dl)
   126  			sid := parseSess(bts[:n])
   127  			if addri, ok := sessMap.Load(sid); ok {
   128  				btsCopy := malloc(n)
   129  				copy(btsCopy, bts)
   130  				start := time.Now()
   131  				maybeDoJob(func() {
   132  					limiter.WaitN(context.Background(), n)
   133  					_, e = leftSock.WriteTo(btsCopy, addri.(net.Addr))
   134  					if err != nil {
   135  						log.Println("cannot write:", err)
   136  					}
   137  					free(btsCopy)
   138  					if statClient != nil {
   139  						inducedLatency := time.Since(start)
   140  						if rand.Int()%100000 < n {
   141  							statClient.Increment(allocGroup + ".e2edown")
   142  						}
   143  						if queueReportLimiter.Allow() {
   144  							statClient.Timing(allocGroup+".queuens", inducedLatency.Nanoseconds())
   145  						}
   146  					}
   147  				})
   148  			}
   149  		}
   150  	}()
   151  	//e2eMap.SetDefault(kee, leftRaw.LocalAddr().(*net.UDPAddr).Port)
   152  	return leftRaw.LocalAddr().(*net.UDPAddr).Port, nil
   153  }
   154  
   155  var queueReportLimiter = rate.NewLimiter(100, 1000)