/
statetable.go
131 lines (110 loc) · 3.02 KB
/
statetable.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
package balancer
import (
"fmt"
"math/rand"
"net"
"sync"
"time"
"github.com/google/gopacket/layers"
)
// State represents a single connection state
type State struct {
State TCPState
Server *Server
Seq uint32
}
// StateTable keeps track of the connection states.
type StateTable struct {
sync.RWMutex
states map[string]*State
}
// NewStateTable creates and initializes a new StateTable.
func NewStateTable() *StateTable {
return &StateTable{
states: make(map[string]*State),
}
}
func (s *StateTable) NewState(ip net.IP, port layers.TCPPort) *State {
s.Lock()
defer s.Unlock()
state := &State{
Seq: randomSequence(),
}
s.states[fmt.Sprintf("%s:%d", ip.String(), port)] = state
return state
}
func (s *StateTable) GetState(ip net.IP, port layers.TCPPort) (*State, bool) {
s.RLock()
defer s.RUnlock()
state, ok := s.states[fmt.Sprintf("%s:%d", ip.String(), port)]
return state, ok
}
// PacketBridgeState represents a single connection state at the packet bridge.
type PacketBridgeState struct {
State TCPState
IP net.IP
HardwareAddr net.HardwareAddr
RandPort layers.TCPPort
Port layers.TCPPort
LBIndex uint8
SeqOffset uint32
PayloadBuf []byte
}
// PacketBridgeStateTable represents a table of tcp connection states.
type PacketBridgeStateTable struct {
sync.RWMutex
byPort map[layers.TCPPort]*PacketBridgeState
byIP map[string]*PacketBridgeState
}
// NewPacketBridgeStateTable creates and initializes a new PacketBridgeStateTable.
func NewPacketBridgeStateTable() *PacketBridgeStateTable {
return &PacketBridgeStateTable{
byPort: make(map[layers.TCPPort]*PacketBridgeState),
byIP: make(map[string]*PacketBridgeState),
}
}
func (s *PacketBridgeStateTable) NewState(ip net.IP, mac net.HardwareAddr, port layers.TCPPort, lbIndex uint8, seqOffset uint32, payload []byte) *PacketBridgeState {
var randPort layers.TCPPort
s.Lock()
defer s.Unlock()
// generate random port that is not yet in the table
// todo max number of tries or a better way to find unique port!
for {
randPort = randomPort()
if _, ok := s.byPort[randPort]; !ok {
break
}
}
state := &PacketBridgeState{
IP: ip,
Port: port,
HardwareAddr: mac,
RandPort: randPort,
LBIndex: lbIndex,
SeqOffset: seqOffset,
PayloadBuf: payload,
}
s.byPort[randPort] = state
s.byIP[fmt.Sprintf("%s:%d", ip.String(), port)] = state
return state
}
func (s *PacketBridgeStateTable) GetByPort(port layers.TCPPort) (*PacketBridgeState, bool) {
s.RLock()
defer s.RUnlock()
state, ok := s.byPort[port]
return state, ok
}
func (s *PacketBridgeStateTable) GetByIP(ip net.IP, port layers.TCPPort) (*PacketBridgeState, bool) {
s.RLock()
defer s.RUnlock()
state, ok := s.byIP[fmt.Sprintf("%s:%d", ip.String(), port)]
return state, ok
}
func randomSequence() uint32 {
r := rand.New(rand.NewSource(time.Now().UnixNano()))
return uint32(r.Int31())
}
func randomPort() layers.TCPPort {
r := rand.New(rand.NewSource(time.Now().UnixNano()))
return layers.TCPPort(r.Int31n(65536))
}