github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/net/proxy/tun/tun2socket/nat/table.go (about)

     1  package nat
     2  
     3  import (
     4  	"math"
     5  	"sync"
     6  
     7  	"github.com/Asutorufa/yuhaiin/pkg/utils/lru"
     8  	"gvisor.dev/gvisor/pkg/tcpip"
     9  )
    10  
    11  var zeroTuple = Tuple{}
    12  
    13  type Tuple struct {
    14  	SourceAddr      tcpip.Address
    15  	SourcePort      uint16
    16  	DestinationAddr tcpip.Address
    17  	DestinationPort uint16
    18  }
    19  
    20  type tableSplit struct {
    21  	v6 *table
    22  	v4 *table
    23  }
    24  
    25  func (t *tableSplit) tupleOf(port uint16, ipv6 bool) Tuple {
    26  	if ipv6 {
    27  		return t.v6.tupleOf(port, true)
    28  	}
    29  
    30  	return t.v4.tupleOf(port, false)
    31  }
    32  
    33  func (t *tableSplit) portOf(tuple Tuple) uint16 {
    34  	if tuple.SourceAddr.Len() == 16 {
    35  		if port := t.v6.portOf(tuple); port != 0 {
    36  			return port
    37  		}
    38  		return t.v6.newConn(tuple)
    39  	}
    40  
    41  	if port := t.v4.portOf(tuple); port != 0 {
    42  		return port
    43  	}
    44  	return t.v4.newConn(tuple)
    45  }
    46  
    47  type table struct {
    48  	lru *lru.LRU[Tuple, uint16]
    49  
    50  	mu    sync.Mutex
    51  	index uint16
    52  }
    53  
    54  func newTableBase() *table {
    55  	return &table{
    56  		lru: lru.New(lru.WithCapacity[Tuple, uint16](math.MaxUint16 - 10001)),
    57  	}
    58  }
    59  
    60  func (t *table) tupleOf(port uint16, _ bool) Tuple {
    61  	p, _ := t.lru.ReverseLoad(port)
    62  	return p
    63  }
    64  
    65  func (t *table) portOf(tuple Tuple) uint16 {
    66  	p, _ := t.lru.Load(tuple)
    67  	return p
    68  }
    69  
    70  func (t *table) newConn(tuple Tuple) uint16 {
    71  	t.mu.Lock()
    72  	newPort, ok := t.lru.LastPopValue()
    73  	if !ok {
    74  		if t.index == 0 {
    75  			newPort = 10000
    76  		} else {
    77  			newPort = t.index
    78  		}
    79  		t.index = newPort + 1
    80  	}
    81  	t.lru.Add(tuple, newPort)
    82  	defer t.mu.Unlock()
    83  
    84  	return newPort
    85  }
    86  
    87  func newTable() *tableSplit {
    88  	return &tableSplit{
    89  		v6: newTableBase(),
    90  		v4: newTableBase(),
    91  	}
    92  }