Connection keepalive/disconnect

xsd/xsd.go

@@ -17,11 +17,14 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"net/http"
 	"os"
 	"os/exec"
 	"os/signal"
 	"os/user"
 	"path"
+	"runtime"
+	"runtime/pprof"
 	"strings"
 	"sync"
 	"syscall"
@@ -44,6 +47,9 @@ var (
 
 	// Log - syslog output (with no -d)
 	Log *logger.Writer
+
+	cpuprofile string
+	memprofile string
 )
 
 const (
@@ -115,10 +121,13 @@ func runClientToServerCopyAs(who, ttype string, conn *xsnet.Conn, fpath string,
 	c.Stdout = os.Stdout
 	c.Stderr = os.Stderr
 
+	// === Set up connection keepalive to client
+	conn.StartupKeepAlive() // goroutine, returns immediately
+	defer conn.ShutdownKeepAlive()
+
 	if chaffing {
-		conn.EnableChaff()
+		conn.StartupChaff()
 	}
-	defer conn.DisableChaff()
 	defer conn.ShutdownChaff()
 
 	// Start the command (no pty)
@@ -212,11 +221,14 @@ func runServerToClientCopyAs(who, ttype string, conn *xsnet.Conn, srcPath string
 	c.Stderr = stdErrBuffer
 	//c.Stderr = nil
 
+	// === Set up connection keepalive to client
+	conn.StartupKeepAlive() // goroutine, returns immediately
+	defer conn.ShutdownKeepAlive()
+
 	if chaffing {
-		conn.EnableChaff()
+		conn.StartupChaff()
 	}
 	//defer conn.Close()
-	defer conn.DisableChaff()
 	defer conn.ShutdownChaff()
 
 	// Start the command (no pty)
@@ -352,12 +364,15 @@ func runShellAs(who, hname, ttype, cmd string, interactive bool, //nolint:funlen
 			}
 		}()
 
+		// === Set up connection keepalive to client
+		conn.StartupKeepAlive() // goroutine, returns immediately
+		defer conn.ShutdownKeepAlive()
+
 		if chaffing {
-			conn.EnableChaff()
+			conn.StartupChaff()
 		}
 		// #gv:s/label=\"runShellAs\$4\"/label=\"deferChaffShutdown\"/
 		defer func() {
-			conn.DisableChaff()
 			conn.ShutdownChaff()
 		}()
 
@@ -545,6 +560,9 @@ func main() { //nolint:funlen,gocyclo
       H_SHA256
       H_SHA512`)
 
+	flag.StringVar(&cpuprofile, "cpuprofile", "", "write cpu profile to <`file`>")
+	flag.StringVar(&memprofile, "memprofile", "", "write memory profile to <`file`>")
+
 	flag.Parse()
 
 	if vopt {
@@ -559,6 +577,24 @@ func main() { //nolint:funlen,gocyclo
 		}
 	}
 
+	// === Profiling instrumentation
+
+	if cpuprofile != "" {
+		f, err := os.Create(cpuprofile)
+		if err != nil {
+			log.Fatal("could not create CPU profile: ", err)
+		}
+		defer f.Close()
+		fmt.Println("StartCPUProfile()")
+		if err := pprof.StartCPUProfile(f); err != nil {
+			log.Fatal("could not start CPU profile: ", err) //nolint:gocritic
+		} else {
+			defer pprof.StopCPUProfile()
+		}
+
+		go func() { http.ListenAndServe("localhost:6060", nil) }() //nolint:errcheck,gosec
+	}
+
 	// Enforce some sane min/max vals on chaff flags
 	if chaffFreqMin < 2 { //nolint:gomnd
 		chaffFreqMin = 2
@@ -611,6 +647,9 @@ func main() { //nolint:funlen,gocyclo
 				syscall.Kill(0, syscall.SIGINT) //nolint:errcheck
 			case "hangup":
 				logger.LogNotice(fmt.Sprintf("[Got signal: %s - nop]", sig)) //nolint:errcheck
+				if cpuprofile != "" || memprofile != "" {
+					dumpProf()
+				}
 			default:
 				logger.LogNotice(fmt.Sprintf("[Got signal: %s - ignored]", sig)) //nolint:errcheck
 			}
@@ -632,6 +671,7 @@ func main() { //nolint:funlen,gocyclo
 		// Wait for a connection.
 		// Then check if client-proposed algs are allowed
 		conn, err := l.Accept()
+		//logger.LogDebug(fmt.Sprintf("l.Accept()\n"))
 		if err != nil {
 			log.Printf("Accept() got error(%v), hanging up.\n", err)
 		} else {
@@ -864,3 +904,23 @@ func main() { //nolint:funlen,gocyclo
 	} //endfor
 	//logger.LogNotice(fmt.Sprintln("[Exiting]")) //nolint:errcheck
 }
+
+func dumpProf() {
+	if cpuprofile != "" {
+		pprof.StopCPUProfile()
+	}
+
+	if memprofile != "" {
+		f, err := os.Create(memprofile)
+		if err != nil {
+			log.Fatal("could not create memory profile: ", err)
+		}
+		defer f.Close()
+		runtime.GC() // get up-to-date statistics
+		if err := pprof.WriteHeapProfile(f); err != nil {
+			log.Fatal("could not write memory profile: ", err) //nolint:gocritic
+		}
+	}
+
+	//os.Exit(status)
+}