Skip to content

Commit a4ddb68

Browse files
rg0nowrg0now
committed
Add callback API to track allocation lifecycle
1 parent 9d49944 commit a4ddb68

14 files changed

+810
-45
lines changed

‎internal/allocation/allocation.go

Lines changed: 93 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,8 @@ type Allocation struct {
3535
channelBindings []*ChannelBind
3636
lifetimeTimer *time.Timer
3737
closed chan interface{}
38+
username, realm string
39+
callback EventHandler
3840
log logging.LeveledLogger
3941

4042
// Some clients (Firefox or others using resiprocate's nICE lib) may retry allocation
@@ -45,12 +47,18 @@ type Allocation struct {
4547
}
4648

4749
// NewAllocation creates a new instance of NewAllocation.
48-
func NewAllocation(turnSocket net.PacketConn, fiveTuple *FiveTuple, log logging.LeveledLogger) *Allocation {
50+
func NewAllocation(
51+
turnSocket net.PacketConn,
52+
fiveTuple *FiveTuple,
53+
callback EventHandler,
54+
log logging.LeveledLogger,
55+
) *Allocation {
4956
return &Allocation{
5057
TurnSocket: turnSocket,
5158
fiveTuple: fiveTuple,
5259
permissions: make(map[string]*Permission, 64),
5360
closed: make(chan interface{}),
61+
callback: callback,
5462
log: log,
5563
}
5664
}
@@ -82,6 +90,21 @@ func (a *Allocation) AddPermission(perms *Permission) {
8290
a.permissions[fingerprint] = perms
8391
a.permissionsLock.Unlock()
8492

93+
if a.callback != nil {
94+
if u, ok := perms.Addr.(*net.UDPAddr); ok {
95+
a.callback(EventHandlerArgs{
96+
Type: OnPermissionCreated,
97+
SrcAddr: a.fiveTuple.SrcAddr,
98+
DstAddr: a.fiveTuple.DstAddr,
99+
Protocol: a.fiveTuple.Protocol,
100+
Username: a.username,
101+
Realm: a.realm,
102+
RelayAddr: a.RelayAddr,
103+
PeerIP: u.IP,
104+
})
105+
}
106+
}
107+
85108
perms.start(permissionTimeout)
86109
}
87110

@@ -90,6 +113,33 @@ func (a *Allocation) RemovePermission(addr net.Addr) {
90113
a.permissionsLock.Lock()
91114
defer a.permissionsLock.Unlock()
92115
delete(a.permissions, ipnet.FingerprintAddr(addr))
116+
117+
if a.callback != nil {
118+
if u, ok := addr.(*net.UDPAddr); ok {
119+
a.callback(EventHandlerArgs{
120+
Type: OnPermissionDeleted,
121+
SrcAddr: a.fiveTuple.SrcAddr,
122+
DstAddr: a.fiveTuple.DstAddr,
123+
Protocol: a.fiveTuple.Protocol,
124+
Username: a.username,
125+
Realm: a.realm,
126+
RelayAddr: a.RelayAddr,
127+
PeerIP: u.IP,
128+
})
129+
}
130+
}
131+
}
132+
133+
// ListPermissions returns the permissions associated with an allocation.
134+
func (a *Allocation) ListPermissions() []*Permission {
135+
ps := []*Permission{}
136+
a.permissionsLock.RLock()
137+
defer a.permissionsLock.RUnlock()
138+
for _, p := range a.permissions {
139+
ps = append(ps, p)
140+
}
141+
142+
return ps
93143
}
94144

95145
// AddChannelBind adds a new ChannelBind to the allocation, it also updates the
@@ -114,6 +164,20 @@ func (a *Allocation) AddChannelBind(chanBind *ChannelBind, lifetime time.Duratio
114164

115165
// Channel binds also refresh permissions.
116166
a.AddPermission(NewPermission(chanBind.Peer, a.log))
167+
168+
if a.callback != nil {
169+
a.callback(EventHandlerArgs{
170+
Type: OnChannelCreated,
171+
SrcAddr: a.fiveTuple.SrcAddr,
172+
DstAddr: a.fiveTuple.DstAddr,
173+
Protocol: a.fiveTuple.Protocol,
174+
Username: a.username,
175+
Realm: a.realm,
176+
RelayAddr: a.RelayAddr,
177+
PeerAddr: chanBind.Peer,
178+
ChannelNumber: uint16(chanBind.Number),
179+
})
180+
}
117181
} else {
118182
channelByNumber.refresh(lifetime)
119183

@@ -131,6 +195,20 @@ func (a *Allocation) RemoveChannelBind(number proto.ChannelNumber) bool {
131195

132196
for i := len(a.channelBindings) - 1; i >= 0; i-- {
133197
if a.channelBindings[i].Number == number {
198+
if a.callback != nil {
199+
a.callback(EventHandlerArgs{
200+
Type: OnChannelDeleted,
201+
SrcAddr: a.fiveTuple.SrcAddr,
202+
DstAddr: a.fiveTuple.DstAddr,
203+
Protocol: a.fiveTuple.Protocol,
204+
Username: a.username,
205+
Realm: a.realm,
206+
RelayAddr: a.RelayAddr,
207+
PeerAddr: a.channelBindings[i].Peer,
208+
ChannelNumber: uint16(a.channelBindings[i].Number),
209+
})
210+
}
211+
134212
a.channelBindings = append(a.channelBindings[:i], a.channelBindings[i+1:]...)
135213

136214
return true
@@ -166,6 +244,16 @@ func (a *Allocation) GetChannelByAddr(addr net.Addr) *ChannelBind {
166244
return nil
167245
}
168246

247+
// ListChannelBindings returns the channel bindings associated with an allocation.
248+
func (a *Allocation) ListChannelBindings() []*ChannelBind {
249+
cs := []*ChannelBind{}
250+
a.channelBindingsLock.RLock()
251+
defer a.channelBindingsLock.RUnlock()
252+
cs = append(cs, a.channelBindings...)
253+
254+
return cs
255+
}
256+
169257
// Refresh updates the allocations lifetime.
170258
func (a *Allocation) Refresh(lifetime time.Duration) {
171259
if !a.lifetimeTimer.Reset(lifetime) {
@@ -201,17 +289,15 @@ func (a *Allocation) Close() error {
201289

202290
a.lifetimeTimer.Stop()
203291

204-
a.permissionsLock.RLock()
205-
for _, p := range a.permissions {
292+
for _, p := range a.ListPermissions() {
293+
a.RemovePermission(p.Addr)
206294
p.lifetimeTimer.Stop()
207295
}
208-
a.permissionsLock.RUnlock()
209296

210-
a.channelBindingsLock.RLock()
211-
for _, c := range a.channelBindings {
297+
for _, c := range a.ListChannelBindings() {
298+
a.RemoveChannelBind(c.Number)
212299
c.lifetimeTimer.Stop()
213300
}
214-
a.channelBindingsLock.RUnlock()
215301

216302
return a.RelaySocket.Close()
217303
}

‎internal/allocation/allocation_manager.go

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ type ManagerConfig struct {
1818
AllocatePacketConn func(network string, requestedPort int) (net.PacketConn, net.Addr, error)
1919
AllocateConn func(network string, requestedPort int) (net.Conn, net.Addr, error)
2020
PermissionHandler func(sourceAddr net.Addr, peerIP net.IP) bool
21+
EventHandler EventHandler
2122
}
2223

2324
type reservation struct {
@@ -36,6 +37,7 @@ type Manager struct {
3637
allocatePacketConn func(network string, requestedPort int) (net.PacketConn, net.Addr, error)
3738
allocateConn func(network string, requestedPort int) (net.Conn, net.Addr, error)
3839
permissionHandler func(sourceAddr net.Addr, peerIP net.IP) bool
40+
EventHandler EventHandler
3941
}
4042

4143
// NewManager creates a new instance of Manager.
@@ -55,6 +57,7 @@ func NewManager(config ManagerConfig) (*Manager, error) {
5557
allocatePacketConn: config.AllocatePacketConn,
5658
allocateConn: config.AllocateConn,
5759
permissionHandler: config.PermissionHandler,
60+
EventHandler: config.EventHandler,
5861
}, nil
5962
}
6063

@@ -94,6 +97,7 @@ func (m *Manager) CreateAllocation(
9497
turnSocket net.PacketConn,
9598
requestedPort int,
9699
lifetime time.Duration,
100+
username, realm string,
97101
) (*Allocation, error) {
98102
switch {
99103
case fiveTuple == nil:
@@ -111,7 +115,9 @@ func (m *Manager) CreateAllocation(
111115
if alloc := m.GetAllocation(fiveTuple); alloc != nil {
112116
return nil, fmt.Errorf("%w: %v", errDupeFiveTuple, fiveTuple)
113117
}
114-
alloc := NewAllocation(turnSocket, fiveTuple, m.log)
118+
alloc := NewAllocation(turnSocket, fiveTuple, m.EventHandler, m.log)
119+
alloc.username = username
120+
alloc.realm = realm
115121

116122
conn, relayAddr, err := m.allocatePacketConn("udp4", requestedPort)
117123
if err != nil {
@@ -131,6 +137,19 @@ func (m *Manager) CreateAllocation(
131137
m.allocations[fiveTuple.Fingerprint()] = alloc
132138
m.lock.Unlock()
133139

140+
if m.EventHandler != nil {
141+
m.EventHandler(EventHandlerArgs{
142+
Type: OnAllocationCreated,
143+
SrcAddr: fiveTuple.SrcAddr,
144+
DstAddr: fiveTuple.DstAddr,
145+
Protocol: UDP,
146+
Username: username,
147+
Realm: realm,
148+
RelayAddr: relayAddr,
149+
RequestedPort: requestedPort,
150+
})
151+
}
152+
134153
go alloc.packetHandler(m)
135154

136155
return alloc, nil
@@ -152,6 +171,17 @@ func (m *Manager) DeleteAllocation(fiveTuple *FiveTuple) {
152171
if err := allocation.Close(); err != nil {
153172
m.log.Errorf("Failed to close allocation: %v", err)
154173
}
174+
175+
if m.EventHandler != nil {
176+
m.EventHandler(EventHandlerArgs{
177+
Type: OnAllocationDeleted,
178+
SrcAddr: fiveTuple.SrcAddr,
179+
DstAddr: fiveTuple.DstAddr,
180+
Protocol: UDP,
181+
Username: allocation.username,
182+
Realm: allocation.realm,
183+
})
184+
}
155185
}
156186

157187
// CreateReservation stores the reservation for the token+port.

‎internal/allocation/allocation_manager_test.go

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -54,13 +54,13 @@ func subTestCreateInvalidAllocation(t *testing.T, turnSocket net.PacketConn) {
5454
m, err := newTestManager()
5555
assert.NoError(t, err)
5656

57-
if a, err := m.CreateAllocation(nil, turnSocket, 0, proto.DefaultLifetime); a != nil || err == nil {
57+
if a, err := m.CreateAllocation(nil, turnSocket, 0, proto.DefaultLifetime, "", ""); a != nil || err == nil {
5858
t.Errorf("Illegally created allocation with nil FiveTuple")
5959
}
60-
if a, err := m.CreateAllocation(randomFiveTuple(), nil, 0, proto.DefaultLifetime); a != nil || err == nil {
60+
if a, err := m.CreateAllocation(randomFiveTuple(), nil, 0, proto.DefaultLifetime, "", ""); a != nil || err == nil {
6161
t.Errorf("Illegally created allocation with nil turnSocket")
6262
}
63-
if a, err := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, 0); a != nil || err == nil {
63+
if a, err := m.CreateAllocation(randomFiveTuple(), turnSocket, 0, 0, "", ""); a != nil || err == nil {
6464
t.Errorf("Illegally created allocation with 0 lifetime")
6565
}
6666
}
@@ -73,7 +73,7 @@ func subTestCreateAllocation(t *testing.T, turnSocket net.PacketConn) {
7373
assert.NoError(t, err)
7474

7575
fiveTuple := randomFiveTuple()
76-
if a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime); a == nil || err != nil {
76+
if a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime, "", ""); a == nil || err != nil {
7777
t.Errorf("Failed to create allocation %v %v", a, err)
7878
}
7979

@@ -90,11 +90,11 @@ func subTestCreateAllocationDuplicateFiveTuple(t *testing.T, turnSocket net.Pack
9090
assert.NoError(t, err)
9191

9292
fiveTuple := randomFiveTuple()
93-
if a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime); a == nil || err != nil {
93+
if a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime, "", ""); a == nil || err != nil {
9494
t.Errorf("Failed to create allocation %v %v", a, err)
9595
}
9696

97-
if a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime); a != nil || err == nil {
97+
if a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime, "", ""); a != nil || err == nil {
9898
t.Errorf("Was able to create allocation with same FiveTuple twice")
9999
}
100100
}
@@ -106,7 +106,8 @@ func subTestDeleteAllocation(t *testing.T, turnSocket net.PacketConn) {
106106
assert.NoError(t, err)
107107

108108
fiveTuple := randomFiveTuple()
109-
if a, err := manager.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime); a == nil || err != nil {
109+
if a, err := manager.CreateAllocation(fiveTuple, turnSocket, 0, proto.DefaultLifetime,
110+
"", ""); a == nil || err != nil {
110111
t.Errorf("Failed to create allocation %v %v", a, err)
111112
}
112113

@@ -133,7 +134,7 @@ func subTestAllocationTimeout(t *testing.T, turnSocket net.PacketConn) {
133134
for index := range allocations {
134135
fiveTuple := randomFiveTuple()
135136

136-
a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, lifetime)
137+
a, err := m.CreateAllocation(fiveTuple, turnSocket, 0, lifetime, "", "")
137138
if err != nil {
138139
t.Errorf("Failed to create allocation with %v", fiveTuple)
139140
}
@@ -159,9 +160,9 @@ func subTestManagerClose(t *testing.T, turnSocket net.PacketConn) {
159160

160161
allocations := make([]*Allocation, 2)
161162

162-
a1, _ := manager.CreateAllocation(randomFiveTuple(), turnSocket, 0, time.Second)
163+
a1, _ := manager.CreateAllocation(randomFiveTuple(), turnSocket, 0, time.Second, "", "")
163164
allocations[0] = a1
164-
a2, _ := manager.CreateAllocation(randomFiveTuple(), turnSocket, 0, time.Minute)
165+
a2, _ := manager.CreateAllocation(randomFiveTuple(), turnSocket, 0, time.Minute, "", "")
165166
allocations[1] = a2
166167

167168
// Make a1 timeout

0 commit comments

Comments
 (0)