wait for server connections to be closed before exiting process

fixes #125

1 files changed, 7 insertions, 0 deletions

diff --git a/run.go b/run.go
index dde772b..df8c77e 100644
--- a/run.go
+++ b/run.go
@@ -88,8 +88,15 @@ func runFunc() {
 
 		for cc := range clts {
 			go func(cc *ClientConn) {
+				sc := cc.server()
+
 				cc.Kick("Proxy shutting down.")
 				<-cc.Closed()
+
+				if sc != nil {
+					<-sc.Closed()
+				}
+
 				wg.Done()
 			}(cc)
 		}