From 4e3dae80e27536b45966c5926c648e05ad969de5 Mon Sep 17 00:00:00 2001
From: Vasile Popescu <elisescu@elisescu.com>
Date: Wed, 4 Nov 2020 19:39:00 +0100
Subject: [PATCH] Improve handling binding and other connection errors

---
 main.go | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/main.go b/main.go
index e2dafc2..550ce62 100644
--- a/main.go
+++ b/main.go
@@ -152,8 +152,13 @@ Flags:
 	fmt.Printf("Press Enter to continue!\n")
 	bufio.NewReader(os.Stdin).ReadString('\n')
 
+	stopPtyAndRestore := func () {
+		ptyMaster.Stop()
+		ptyMaster.Restore()
+	}
+
 	ptyMaster.MakeRaw()
-	defer ptyMaster.Restore()
+	defer stopPtyAndRestore()
 	var pty server.PTYHandler = ptyMaster
 	if *readOnly {
 		pty = &nilPTY{}
@@ -171,22 +176,27 @@ Flags:
 
 	mw := io.MultiWriter(os.Stdout, server)
 
+
 	go func() {
 		err := server.Run()
 		if err != nil {
-			log.Debugf("Server done: %s", err.Error())
+			stopPtyAndRestore()
+			log.Errorf("Server finished: %s", err.Error())
 		}
 	}()
 
 	go func() {
 		_, err := io.Copy(mw, ptyMaster)
 		if err != nil {
-			ptyMaster.Stop()
+			stopPtyAndRestore()
 		}
 	}()
 
 	go func() {
-		io.Copy(ptyMaster, os.Stdin)
+		_, err := io.Copy(ptyMaster, os.Stdin)
+		if err != nil {
+			stopPtyAndRestore()
+		}
 	}()
 
 	ptyMaster.Wait()