From 6da5802ae56b0d348661763c5dc7b091d77b72d1 Mon Sep 17 00:00:00 2001
From: Arceliar <Arceliar@users.noreply.github.com>
Date: Sun, 4 Aug 2019 02:08:47 -0500
Subject: [PATCH] don't block forever in Write if the session is cancelled,
 cleanup Conn.Read slightly

---
 src/yggdrasil/conn.go | 54 +++++++++++++++++++++++++------------------
 1 file changed, 31 insertions(+), 23 deletions(-)

diff --git a/src/yggdrasil/conn.go b/src/yggdrasil/conn.go
index f72bdd2f..d00db5c5 100644
--- a/src/yggdrasil/conn.go
+++ b/src/yggdrasil/conn.go
@@ -154,29 +154,27 @@ func (c *Conn) Read(b []byte) (int, error) {
 	sinfo := c.session
 	cancel := c.getDeadlineCancellation(&c.readDeadline)
 	defer cancel.Cancel(nil)
-	for {
-		// Wait for some traffic to come through from the session
-		select {
-		case <-cancel.Finished():
-			if cancel.Error() == util.CancellationTimeoutError {
-				return 0, ConnError{errors.New("read timeout"), true, false, false, 0}
-			} else {
-				return 0, ConnError{errors.New("session closed"), false, false, true, 0}
-			}
-		case bs := <-sinfo.recv:
-			var err error
-			n := len(bs)
-			if len(bs) > len(b) {
-				n = len(b)
-				err = ConnError{errors.New("read buffer too small for entire packet"), false, true, false, 0}
-			}
-			// Copy results to the output slice and clean up
-			copy(b, bs)
-			util.PutBytes(bs)
-			// If we've reached this point then everything went to plan, return the
-			// number of bytes we populated back into the given slice
-			return n, err
+	// Wait for some traffic to come through from the session
+	select {
+	case <-cancel.Finished():
+		if cancel.Error() == util.CancellationTimeoutError {
+			return 0, ConnError{errors.New("read timeout"), true, false, false, 0}
+		} else {
+			return 0, ConnError{errors.New("session closed"), false, false, true, 0}
 		}
+	case bs := <-sinfo.recv:
+		var err error
+		n := len(bs)
+		if len(bs) > len(b) {
+			n = len(b)
+			err = ConnError{errors.New("read buffer too small for entire packet"), false, true, false, 0}
+		}
+		// Copy results to the output slice and clean up
+		copy(b, bs)
+		util.PutBytes(bs)
+		// If we've reached this point then everything went to plan, return the
+		// number of bytes we populated back into the given slice
+		return n, err
 	}
 }
 
@@ -206,7 +204,17 @@ func (c *Conn) Write(b []byte) (bytesWritten int, err error) {
 	sinfo.doFunc(sessionFunc)
 	if written > 0 {
 		bs := append(util.GetBytes(), b...)
-		sinfo.send <- bs
+		cancel := c.getDeadlineCancellation(&c.writeDeadline)
+		defer cancel.Cancel(nil)
+		select {
+		case <-cancel.Finished():
+			if cancel.Error() == util.CancellationTimeoutError {
+				return 0, ConnError{errors.New("write timeout"), true, false, false, 0}
+			} else {
+				return 0, ConnError{errors.New("session closed"), false, false, true, 0}
+			}
+		case sinfo.send <- bs:
+		}
 	}
 	return written, err
 }