Properly close DummyTUN to avoid deadlock in TestNoiseHandshake

License: MIT
Signed-off-by: Filippo Valsorda <valsorda@google.com>

3 files changed, 12 insertions, 1 deletions

diff --git a/helper_test.go b/helper_test.go
index c72e868..b7c9c6c 100644
--- a/helper_test.go
+++ b/helper_test.go
@@ -8,6 +8,7 @@ package main
 
 import (
 	"bytes"
+	"errors"
 	"os"
 	"testing"
 )
@@ -40,6 +41,8 @@ func (tun *DummyTUN) Write(d []byte, offset int) (int, error) {
 }
 
 func (tun *DummyTUN) Close() error {
+	close(tun.events)
+	close(tun.packets)
 	return nil
 }
 
@@ -48,7 +51,10 @@ func (tun *DummyTUN) Events() chan TUNEvent {
 }
 
 func (tun *DummyTUN) Read(d []byte, offset int) (int, error) {
-	t := <-tun.packets
+	t, ok := <-tun.packets
+	if !ok {
+		return 0, errors.New("device closed")
+	}
 	copy(d[offset:], t)
 	return len(t), nil
 }
@@ -57,6 +63,7 @@ func CreateDummyTUN(name string) (TUNDevice, error) {
 	var dummy DummyTUN
 	dummy.mtu = 0
 	dummy.packets = make(chan []byte, 100)
+	dummy.events = make(chan TUNEvent, 10)
 	return &dummy, nil
 }
 
diff --git a/noise_test.go b/noise_test.go
index d5b1011..3da36e9 100644
--- a/noise_test.go
+++ b/noise_test.go
@@ -58,6 +58,7 @@ func TestNoiseHandshake(t *testing.T) {
 	packet := make([]byte, 0, 256)
 	writer := bytes.NewBuffer(packet)
 	err = binary.Write(writer, binary.LittleEndian, msg1)
+	assertNil(t, err)
 	peer := dev2.ConsumeMessageInitiation(msg1)
 	if peer == nil {
 		t.Fatal("handshake failed at initiation message")
diff --git a/tun.go b/tun.go
index 65a878c..d2fb708 100644
--- a/tun.go
+++ b/tun.go
@@ -33,9 +33,11 @@ type TUNDevice interface {
 
 func (device *Device) RoutineTUNEventReader() {
 	setUp := false
+	logDebug := device.log.Debug
 	logInfo := device.log.Info
 	logError := device.log.Error
 
+	logDebug.Println("Routine: event worker - started")
 	device.state.starting.Done()
 
 	for event := range device.tun.device.Events() {
@@ -67,5 +69,6 @@ func (device *Device) RoutineTUNEventReader() {
 		}
 	}
 
+	logDebug.Println("Routine: event worker - stopped")
 	device.state.stopping.Done()
 }