From 89678198309995a9e9d2155939a6d85091bee686 Mon Sep 17 00:00:00 2001
From: Andy Wang <cbeuw.andy@gmail.com>
Date: Wed, 14 Aug 2019 10:28:08 +0100
Subject: [PATCH] Fix accidentally removed tx qos

---
 internal/multiplex/switchboard.go      |  9 +++--
 internal/multiplex/switchboard_test.go | 47 ++++++++++++++++++++------
 2 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/internal/multiplex/switchboard.go b/internal/multiplex/switchboard.go
index cd8cefb..8a98a30 100644
--- a/internal/multiplex/switchboard.go
+++ b/internal/multiplex/switchboard.go
@@ -64,12 +64,15 @@ func (sb *switchboard) removeConn(connId uint32) {
 
 // a pointer to connId is passed here so that the switchboard can reassign it
 func (sb *switchboard) send(data []byte, connId *uint32) (int, error) {
+	sb.Valve.rxWait(len(data))
 	sb.connsM.RLock()
 	defer sb.connsM.RUnlock()
 	var conn net.Conn
 	conn, ok := sb.conns[*connId]
 	if ok {
-		return conn.Write(data)
+		n, err := conn.Write(data)
+		sb.Valve.AddTx(int64(n))
+		return n, err
 	} else {
 		// do not call assignRandomConn() here.
 		// we'll have to do connsM.RLock() after we get a new connId from assignRandomConn, in order to
@@ -86,7 +89,9 @@ func (sb *switchboard) send(data []byte, connId *uint32) (int, error) {
 		if !ok {
 			return 0, errBrokenSwitchboard
 		} else {
-			return conn.Write(data)
+			n, err := conn.Write(data)
+			sb.Valve.AddTx(int64(n))
+			return n, err
 		}
 	}
 
diff --git a/internal/multiplex/switchboard_test.go b/internal/multiplex/switchboard_test.go
index 4406826..89e410f 100644
--- a/internal/multiplex/switchboard_test.go
+++ b/internal/multiplex/switchboard_test.go
@@ -1,6 +1,7 @@
 package multiplex
 
 import (
+	"github.com/cbeuw/Cloak/internal/util"
 	"math/rand"
 	"testing"
 )
@@ -9,18 +10,13 @@ func BenchmarkSwitchboard_Send(b *testing.B) {
 	seshConfig := &SessionConfig{
 		Obfuscator: nil,
 		Valve:      nil,
-		UnitRead:   nil,
+		UnitRead:   util.ReadTLS,
 	}
 	sesh := MakeSession(0, seshConfig)
-	sbConfig := &switchboardConfig{
-		Valve:     UNLIMITED_VALVE,
-		unordered: false,
-		strategy:  0,
-	}
-	sb := makeSwitchboard(sesh, sbConfig)
+
 	hole := newBlackHole()
-	sb.addConn(hole)
-	connId, err := sb.assignRandomConn()
+	sesh.sb.addConn(hole)
+	connId, err := sesh.sb.assignRandomConn()
 	if err != nil {
 		b.Error("failed to get a random conn", err)
 		return
@@ -29,7 +25,7 @@ func BenchmarkSwitchboard_Send(b *testing.B) {
 	rand.Read(data)
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		n, err := sb.send(data, &connId)
+		n, err := sesh.sb.send(data, &connId)
 		if err != nil {
 			b.Error(err)
 			return
@@ -37,3 +33,34 @@ func BenchmarkSwitchboard_Send(b *testing.B) {
 		b.SetBytes(int64(n))
 	}
 }
+
+func TestSwitchboard_TxCredit(t *testing.T) {
+	seshConfig := &SessionConfig{
+		Obfuscator: nil,
+		Valve:      MakeValve(1<<20, 1<<20),
+		UnitRead:   util.ReadTLS,
+	}
+	sesh := MakeSession(0, seshConfig)
+	hole := newBlackHole()
+	sesh.sb.addConn(hole)
+	connId, err := sesh.sb.assignRandomConn()
+	if err != nil {
+		t.Error("failed to get a random conn", err)
+		return
+	}
+	data := make([]byte, 1000)
+	rand.Read(data)
+
+	n, err := sesh.sb.send(data[:10], &connId)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	if n != 10 {
+		t.Errorf("wanted to send %v, got %v", 10, n)
+		return
+	}
+	if *sesh.sb.Valve.(*LimitedValve).tx != 10 {
+		t.Error("tx credit didn't increase by 10")
+	}
+}