From df9055f69cbd8a08e3388d03150d92205a085c23 Mon Sep 17 00:00:00 2001
From: Asim Aslam <asim@aslam.me>
Date: Fri, 3 Jan 2020 20:43:53 +0000
Subject: [PATCH] continue to process messages even after the connection is
 closed

---
 tunnel/session.go | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/tunnel/session.go b/tunnel/session.go
index 84153541..ebf01ded 100644
--- a/tunnel/session.go
+++ b/tunnel/session.go
@@ -173,6 +173,25 @@ func (s *session) waitFor(msgType string, timeout time.Duration) (*message, erro
 		case <-after(timeout):
 			return nil, ErrReadTimeout
 		case <-s.closed:
+			// check pending message queue
+			select {
+			case msg := <-s.recv:
+				// there may be no message type
+				if len(msgType) == 0 {
+					return msg, nil
+				}
+
+				// ignore what we don't want
+				if msg.typ != msgType {
+					log.Debugf("Tunnel received non %s message in waiting for %s", msg.typ, msgType)
+					continue
+				}
+
+				// got the message
+				return msg, nil
+			default:
+				// non blocking
+			}
 			return nil, io.EOF
 		}
 	}