Merge pull request antoniomika#55 from antoniomika/fix_no_tty

antoniomika · web-flow · commit ff31f2bbf146 · 2019-11-29T15:20:49.000-05:00
Fixed connection close during a notty session
diff --git a/channels.go b/channels.go
@@ -42,6 +42,8 @@ func handleSession(newChannel ssh.NewChannel, sshConn *SSHConnection, state *Sta
 			data := make([]byte, 4096)
 			dataRead, err := connection.Read(data)
 			if err != nil && err == io.EOF {
+				break
+			} else if err != nil {
 				select {
 				case <-sshConn.Close:
 					break
@@ -122,8 +124,14 @@ func handleAlias(newChannel ssh.NewChannel, sshConn *SSHConnection, state *State
 
 	sshConn.Listeners.Store(conn.RemoteAddr(), nil)
 
-	copyBoth(conn, connection, false)
-	sshConn.CleanUp(state)
+	copyBoth(conn, connection)
+
+	select {
+	case <-sshConn.Close:
+		break
+	default:
+		sshConn.CleanUp(state)
+	}
 }
 
 func writeToSession(connection ssh.Channel, c string) {
diff --git a/handle.go b/handle.go
@@ -22,6 +22,11 @@ func handleRequest(newRequest *ssh.Request, sshConn *SSHConnection, state *State
 	case "tcpip-forward":
 		go checkSession(newRequest, sshConn, state)
 		handleRemoteForward(newRequest, sshConn, state)
+	case "keepalive@openssh.com":
+		err := newRequest.Reply(true, nil)
+		if err != nil {
+			log.Println("Error replying to socket request:", err)
+		}
 	default:
 		err := newRequest.Reply(false, nil)
 		if err != nil {
diff --git a/main.go b/main.go
@@ -267,6 +267,20 @@ func main() {
 
 			state.SSHConnections.Store(sshConn.RemoteAddr(), holderConn)
 
+			go func() {
+				err := sshConn.Wait()
+				if err != nil && *debug {
+					log.Println("Closing SSH connection:", err)
+				}
+
+				select {
+				case <-holderConn.Close:
+					break
+				default:
+					holderConn.CleanUp(state)
+				}
+			}()
+
 			go handleRequests(reqs, holderConn, state)
 			go handleChannels(chans, holderConn, state)
 
diff --git a/requests.go b/requests.go
@@ -8,7 +8,6 @@ import (
 	"net"
 	"os"
 	"strconv"
-	"sync"
 	"time"
 
 	"github.com/logrusorgru/aurora"
@@ -227,47 +226,65 @@ func handleRemoteForward(newRequest *ssh.Request, sshConn *SSHConnection, state
 			}
 		}
 
-		go copyBoth(cl, newChan, false)
+		go copyBoth(cl, newChan)
 		go ssh.DiscardRequests(newReqs)
 	}
 }
 
-func copyBoth(writer net.Conn, reader ssh.Channel, wait bool) {
+// IdleTimeoutConn handles the connection with a context deadline
+// code adapted from https://qiita.com/kwi/items/b38d6273624ad3f6ae79
+type IdleTimeoutConn struct {
+	Conn net.Conn
+}
+
+// Read is needed to implement the reader part
+func (i IdleTimeoutConn) Read(buf []byte) (int, error) {
+	err := i.Conn.SetDeadline(time.Now().Add(5 * time.Second))
+	if err != nil {
+		return 0, err
+	}
+
+	return i.Conn.Read(buf)
+}
+
+// Write is needed to implement the writer part
+func (i IdleTimeoutConn) Write(buf []byte) (int, error) {
+	err := i.Conn.SetDeadline(time.Now().Add(5 * time.Second))
+	if err != nil {
+		return 0, err
+	}
+
+	return i.Conn.Write(buf)
+}
+
+func copyBoth(writer net.Conn, reader ssh.Channel) {
 	closeBoth := func() {
-		time.Sleep(100 * time.Millisecond)
-		writer.Close()
 		reader.Close()
+		writer.Close()
 	}
 
-	var wg sync.WaitGroup
-
-	go func() {
-		if wait {
-			wg.Add(1)
-			defer wg.Done()
-		}
+	tcon := IdleTimeoutConn{
+		Conn: writer,
+	}
 
-		_, err := io.Copy(reader, writer)
+	copyToReader := func() {
+		_, err := io.Copy(reader, tcon)
 		if err != nil && *debug {
 			log.Println("Error copying to reader:", err)
 		}
-	}()
 
-	func() {
-		if wait {
-			wg.Add(1)
-			defer wg.Done()
-		}
+		closeBoth()
+	}
 
-		_, err := io.Copy(writer, reader)
+	copyToWriter := func() {
+		_, err := io.Copy(tcon, reader)
 		if err != nil && *debug {
 			log.Println("Error copying to writer:", err)
 		}
-	}()
 
-	if wait {
-		wg.Wait()
+		closeBoth()
 	}
 
-	closeBoth()
+	go copyToReader()
+	copyToWriter()
 }
