From 34a4805c3435ef404d600d0aa2ebfc1bcc1bd076 Mon Sep 17 00:00:00 2001
From: Andy Wang <cbeuw.andy@gmail.com>
Date: Wed, 16 Oct 2019 23:00:11 +0100
Subject: [PATCH] Add tests

---
 internal/multiplex/session_test.go     | 18 ++++++++++--
 internal/multiplex/switchboard_test.go | 39 ++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/internal/multiplex/session_test.go b/internal/multiplex/session_test.go
index d1579d5..ac6e027 100644
--- a/internal/multiplex/session_test.go
+++ b/internal/multiplex/session_test.go
@@ -153,7 +153,11 @@ func BenchmarkRecvDataFromRemote_Ordered(b *testing.B) {
 
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
-			sesh.recvDataFromRemote(obfsBuf[:n])
+			err := sesh.recvDataFromRemote(obfsBuf[:n])
+			if err != nil {
+				b.Error(err)
+				return
+			}
 			b.SetBytes(int64(n))
 		}
 	})
@@ -166,7 +170,11 @@ func BenchmarkRecvDataFromRemote_Ordered(b *testing.B) {
 
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
-			sesh.recvDataFromRemote(obfsBuf[:n])
+			err := sesh.recvDataFromRemote(obfsBuf[:n])
+			if err != nil {
+				b.Error(err)
+				return
+			}
 			b.SetBytes(int64(n))
 		}
 	})
@@ -179,7 +187,11 @@ func BenchmarkRecvDataFromRemote_Ordered(b *testing.B) {
 
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
-			sesh.recvDataFromRemote(obfsBuf[:n])
+			err := sesh.recvDataFromRemote(obfsBuf[:n])
+			if err != nil {
+				b.Error(err)
+				return
+			}
 			b.SetBytes(int64(n))
 		}
 	})
diff --git a/internal/multiplex/switchboard_test.go b/internal/multiplex/switchboard_test.go
index a1e5926..32434fe 100644
--- a/internal/multiplex/switchboard_test.go
+++ b/internal/multiplex/switchboard_test.go
@@ -3,7 +3,9 @@ package multiplex
 import (
 	"github.com/cbeuw/Cloak/internal/util"
 	"math/rand"
+	"net"
 	"testing"
+	"time"
 )
 
 func TestSwitchboard_Send(t *testing.T) {
@@ -148,3 +150,40 @@ func TestSwitchboard_TxCredit(t *testing.T) {
 		}
 	})
 }
+
+func TestSwitchboard_CloseOnOneDisconn(t *testing.T) {
+	sesh := setupSesh(false)
+
+	l, _ := net.Listen("tcp", "127.0.0.1:0")
+	addRemoteConn := func(close chan struct{}) {
+		conn, _ := net.Dial("tcp", l.Addr().String())
+		for {
+			conn.Write([]byte{0x00})
+			<-close
+			conn.Close()
+		}
+	}
+
+	close0 := make(chan struct{})
+	go addRemoteConn(close0)
+	conn0, _ := l.Accept()
+	sesh.AddConnection(conn0)
+
+	close1 := make(chan struct{})
+	go addRemoteConn(close1)
+	conn1, _ := l.Accept()
+	sesh.AddConnection(conn1)
+
+	close0 <- struct{}{}
+
+	time.Sleep(100 * time.Millisecond)
+
+	if !sesh.IsClosed() {
+		t.Error("session not closed after one conn is disconnected")
+		return
+	}
+	if _, err := conn1.Write([]byte{0x00}); err == nil {
+		t.Error("the other conn is still connected")
+		return
+	}
+}