Hesham/wireguard-go
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

device: use channel close to shut down and drain decryption channel

Browse Source 
This is similar to commit e1fa1cc556,
but for the decryption channel.

It is an alternative fix to f9f655567930a4cd78d40fa4ba0d58503335ae6a.

Signed-off-by: Josh Bleecher Snyder <josh@tailscale.com>
This commit is contained in:
Josh Bleecher Snyder 2021-01-11 17:34:02 -08:00 committed by Jason A. Donenfeld
parent 675955de5d
commit 48c3b87eb8
2 changed files with 54 additions and 64 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
device/device.go
View File

@ -76,7 +76,7 @@ type Device struct {
queue struct {
encryption *encryptionQueue
decryption chan *QueueInboundElement
decryption *decryptionQueue
handshake chan QueueHandshakeElement
}
@ -115,6 +115,24 @@ func newEncryptionQueue() *encryptionQueue {
return q
}
// A decryptionQueue is similar to an encryptionQueue; see those docs.
type decryptionQueue struct {
c chan *QueueInboundElement
wg sync.WaitGroup
}
func newDecryptionQueue() *decryptionQueue {
q := &decryptionQueue{
c: make(chan *QueueInboundElement, QueueInboundSize),
}
q.wg.Add(1)
go func() {
q.wg.Wait()
close(q.c)
}()
return q
}
/* Converts the peer into a "zombie", which remains in the peer map,
* but processes no packets and does not exists in the routing table.
*
@ -308,7 +326,7 @@ func NewDevice(tunDevice tun.Device, logger *Logger) *Device {
device.queue.handshake = make(chan QueueHandshakeElement, QueueHandshakeSize)
device.queue.encryption = newEncryptionQueue()
device.queue.decryption = make(chan *QueueInboundElement, QueueInboundSize)
device.queue.decryption = newDecryptionQueue()
// prepare signals
@ -369,13 +387,6 @@ func (device *Device) RemoveAllPeers() {
func (device *Device) FlushPacketQueues() {
for {
select {
case elem, ok := <-device.queue.decryption:
if ok {
if !elem.IsDropped() {
elem.Drop()
device.PutMessageBuffer(elem.buffer)
}
}
case <-device.queue.handshake:
default:
return
@ -399,10 +410,11 @@ func (device *Device) Close() {
device.isUp.Set(false)
// We kept a reference to the encryption queue,
// in case we started any new peers that might write to it.
// No new peers are coming; we are done with the encryption queue.
// We kept a reference to the encryption and decryption queues,
// in case we started any new peers that might write to them.
// No new peers are coming; we are done with these queues.
device.queue.encryption.wg.Done()
device.queue.decryption.wg.Done()
close(device.signals.stop)
device.state.stopping.Wait()
@ -549,6 +561,7 @@ func (device *Device) BindUpdate() error {
// start receiving routines
device.net.stopping.Add(2)
device.queue.decryption.wg.Add(2) // each RoutineReceiveIncoming goroutine writes to device.queue.decryption
go device.RoutineReceiveIncoming(ipv4.Version, netc.bind)
go device.RoutineReceiveIncoming(ipv6.Version, netc.bind)

81
device/receive.go
View File

@ -109,6 +109,7 @@ func (device *Device) RoutineReceiveIncoming(IP int, bind conn.Bind) {
logDebug := device.log.Debug
defer func() {
logDebug.Println("Routine: receive incoming IPv" + strconv.Itoa(IP) + " - stopped")
device.queue.decryption.wg.Done()
device.net.stopping.Done()
}()
@ -206,7 +207,7 @@ func (device *Device) RoutineReceiveIncoming(IP int, bind conn.Bind) {
peer.queue.RLock()
if peer.isRunning.Get() {
if device.addToInboundAndDecryptionQueues(peer.queue.inbound, device.queue.decryption, elem) {
if device.addToInboundAndDecryptionQueues(peer.queue.inbound, device.queue.decryption.c, elem) {
buffer = device.GetMessageBuffer()
}
} else {
@ -258,59 +259,35 @@ func (device *Device) RoutineDecryption() {
}()
logDebug.Println("Routine: decryption worker - started")
for {
select {
case <-device.signals.stop:
for {
select {
case elem, ok := <-device.queue.decryption:
if ok {
if !elem.IsDropped() {
elem.Drop()
device.PutMessageBuffer(elem.buffer)
}
elem.Unlock()
}
default:
return
}
}
for elem := range device.queue.decryption.c {
// check if dropped
case elem, ok := <-device.queue.decryption:
if !ok {
return
}
// check if dropped
if elem.IsDropped() {
continue
}
// split message into fields
counter := elem.packet[MessageTransportOffsetCounter:MessageTransportOffsetContent]
content := elem.packet[MessageTransportOffsetContent:]
// decrypt and release to consumer
var err error
elem.counter = binary.LittleEndian.Uint64(counter)
// copy counter to nonce
binary.LittleEndian.PutUint64(nonce[0x4:0xc], elem.counter)
elem.packet, err = elem.keypair.receive.Open(
content[:0],
nonce[:],
content,
nil,
)
if err != nil {
elem.Drop()
device.PutMessageBuffer(elem.buffer)
}
elem.Unlock()
if elem.IsDropped() {
continue
}
// split message into fields
counter := elem.packet[MessageTransportOffsetCounter:MessageTransportOffsetContent]
content := elem.packet[MessageTransportOffsetContent:]
// decrypt and release to consumer
var err error
elem.counter = binary.LittleEndian.Uint64(counter)
// copy counter to nonce
binary.LittleEndian.PutUint64(nonce[0x4:0xc], elem.counter)
elem.packet, err = elem.keypair.receive.Open(
content[:0],
nonce[:],
content,
nil,
)
if err != nil {
elem.Drop()
device.PutMessageBuffer(elem.buffer)
}
elem.Unlock()
}
}