From 23ce95dec0f659f0cdd357563c19a59637c17f18 Mon Sep 17 00:00:00 2001
From: Michael Stapelberg <stapelberg@google.com>
Date: Sun, 27 Oct 2019 10:39:47 +0100
Subject: [PATCH] DirectTCPIPHandler: block until copying is done or fails

This way, users can wrap DirectTCPIPHandler with behavior, e.g. waking up a
powered-down machine before proxying packets to it, and sending that machine
back to sleep once no more connections are open.

AFAICT, this is not a change in behavior since handlers are already run in their
own goroutine, so blocking in DirectTCPIPHandler should not block anything:
https://github.com/gliderlabs/ssh/blob/59d6e4540dea33ecd8f90b12887906b7b081fbcd/server.go#L287
---
 tcpip.go | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tcpip.go b/tcpip.go
index 335fda6..2f81a8f 100644
--- a/tcpip.go
+++ b/tcpip.go
@@ -53,16 +53,22 @@ func DirectTCPIPHandler(srv *Server, conn *gossh.ServerConn, newChan gossh.NewCh
 	}
 	go gossh.DiscardRequests(reqs)
 
+	done := make(chan struct{}, 2)
 	go func() {
 		defer ch.Close()
 		defer dconn.Close()
 		io.Copy(ch, dconn)
+		done <- struct{}{}
 	}()
 	go func() {
 		defer ch.Close()
 		defer dconn.Close()
 		io.Copy(dconn, ch)
+		done <- struct{}{}
 	}()
+
+	<-done
+	<-done
 }
 
 type remoteForwardRequest struct {