From f3bf6d904122740c283c08acffdad1d967f4ba2e Mon Sep 17 00:00:00 2001
From: Russ Magee <rmagee@gmail.com>
Date: Sun, 23 Feb 2020 21:52:20 -0800
Subject: [PATCH] Added login timeout (30s)

Signed-off-by: Russ Magee <rmagee@gmail.com>
---
 xs/xs.go   | 27 +++++++++++++++++++--------
 xsd/xsd.go |  9 +++++++++
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/xs/xs.go b/xs/xs.go
index 41c2c1b..ce924bc 100755
--- a/xs/xs.go
+++ b/xs/xs.go
@@ -420,7 +420,7 @@ func doShellMode(isInteractive bool, conn *xsnet.Conn, oldState *xs.State, rec *
 		// exit with inerr == nil
 		_, inerr := io.Copy(os.Stdout, conn)
 		if inerr != nil {
-			_ = xs.Restore(int(os.Stdin.Fd()), oldState) // #nosec
+			restoreTermState(oldState)
 			// Copy operations and user logging off will cause
 			// a "use of closed network connection" so handle that
 			// gracefully here
@@ -435,7 +435,7 @@ func doShellMode(isInteractive bool, conn *xsnet.Conn, oldState *xs.State, rec *
 
 		if isInteractive {
 			log.Println("[* Got EOF *]")
-			_ = xs.Restore(int(os.Stdin.Fd()), oldState) // #nosec
+			restoreTermState(oldState)
 			exitWithStatus(int(rec.Status()))
 		}
 	}
@@ -463,7 +463,7 @@ func doShellMode(isInteractive bool, conn *xsnet.Conn, oldState *xs.State, rec *
 			if outerr != nil {
 				log.Println(outerr)
 				fmt.Println(outerr)
-				_ = xs.Restore(int(os.Stdin.Fd()), oldState) // #nosec
+				restoreTermState(oldState)
 				log.Println("[Hanging up]")
 				exitWithStatus(0)
 			}
@@ -848,12 +848,17 @@ func main() {
 			}
 			// #gv:s/label=\"main\$1\"/label=\"deferRestore\"/
 			// TODO:.gv:main:1:deferRestore
-			defer func() { _ = xs.Restore(int(os.Stdin.Fd()), oldState) }() // nolint: errcheck,gosec
+			defer restoreTermState(oldState)
 		} else {
 			log.Println("NOT A TTY")
 		}
 	}
 
+	// Start login timeout here and disconnect if user/pass phase stalls
+	loginTimeout := time.AfterFunc(30*time.Second, func() {
+		fmt.Printf(" .. [login timeout]")
+	})
+
 	if len(authCookie) == 0 {
 		//No auth token, prompt for password
 		fmt.Printf("Gimme cookie:")
@@ -864,6 +869,8 @@ func main() {
 		}
 		authCookie = string(ab)
 	}
+
+	_ = loginTimeout.Stop()
 	// Security scrub
 	runtime.GC()
 
@@ -871,9 +878,9 @@ func main() {
 	rec := xs.NewSession(op, []byte(uname), []byte(remoteHost), []byte(os.Getenv("TERM")), []byte(cmdStr), []byte(authCookie), 0)
 	sendErr := sendSessionParams(&conn, rec)
 	if sendErr != nil {
-		_ = xs.Restore(int(os.Stdin.Fd()), oldState) // nolint: errcheck,gosec
+		restoreTermState(oldState)
 		rec.SetStatus(254)
-		fmt.Fprintln(os.Stderr, "Error: server rejected secure proposal params") // nolint: errcheck
+		fmt.Fprintln(os.Stderr, "Error: server rejected secure proposal params or login timed out") // nolint: errcheck
 		exitWithStatus(int(rec.Status()))
 		//log.Fatal(sendErr)
 	}
@@ -930,19 +937,23 @@ func main() {
 		}
 
 		if rec.Status() != 0 {
-			_ = xs.Restore(int(os.Stdin.Fd()), oldState)                         // nolint: errcheck,gosec
+			restoreTermState(oldState)
 			fmt.Fprintln(os.Stderr, "Session exited with status:", rec.Status()) // nolint: errcheck
 		}
 	}
 
 	if oldState != nil {
-		_ = xs.Restore(int(os.Stdin.Fd()), oldState) // nolint: gosec
+		restoreTermState(oldState)
 		oldState = nil
 	}
 
 	exitWithStatus(int(rec.Status()))
 }
 
+func restoreTermState(oldState *xs.State) {
+	_ = xs.Restore(int(os.Stdin.Fd()), oldState) // nolint: errcheck,gosec
+}
+
 // exitWithStatus wraps os.Exit() plus does any required pprof housekeeping
 func exitWithStatus(status int) {
 	if cpuprofile != "" {
diff --git a/xsd/xsd.go b/xsd/xsd.go
index e46c2a7..ae654de 100755
--- a/xsd/xsd.go
+++ b/xsd/xsd.go
@@ -26,6 +26,7 @@ import (
 	"strings"
 	"sync"
 	"syscall"
+	"time"
 	"unsafe"
 
 	"blitter.com/go/goutmp"
@@ -640,6 +641,13 @@ func main() {
 			go func(hc *xsnet.Conn) (e error) {
 				defer hc.Close() // nolint: errcheck
 
+				// Start login timeout here and disconnect if user/pass phase stalls
+				loginTimeout := time.AfterFunc(30*time.Second, func() {
+					logger.LogNotice(fmt.Sprintln("Login timed out")) // nolint: errcheck,gosec
+					hc.Write([]byte{0})                               // nolint: gosec,errcheck
+					hc.Close()
+				})
+
 				//We use io.ReadFull() here to guarantee we consume
 				//just the data we want for the xs.Session, and no more.
 				//Otherwise data will be sitting in the channel that isn't
@@ -719,6 +727,7 @@ func main() {
 					}
 				}
 
+				_ = loginTimeout.Stop()
 				// Security scrub
 				rec.ClearAuthCookie()