From 48b0ce09734d8a449f24646e0d213e1a5a708ce4 Mon Sep 17 00:00:00 2001
From: Jan Kuri <jkuri88@gmail.com>
Date: Tue, 17 Apr 2018 17:42:49 +0200
Subject: [PATCH] feat(client): add keep-alive option for client

---
 cmd/vex/main.go | 32 +++++++++++++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/cmd/vex/main.go b/cmd/vex/main.go
index d2eca91..e216b24 100644
--- a/cmd/vex/main.go
+++ b/cmd/vex/main.go
@@ -11,6 +11,7 @@ import (
 	"os/signal"
 	"time"
 
+	"github.com/pkg/errors"
 	"golang.org/x/crypto/ssh"
 )
 
@@ -34,7 +35,9 @@ var help = `
 
   -ls, Local HTTP server host (default: localhost)
 
-  -lp, Local HTTP server port (default: 7500)
+	-lp, Local HTTP server port (default: 7500)
+
+	-a, Keep tunnel connection alive (default: false)
 
   Read more:
     https://github.com/bleenco/vex
@@ -45,6 +48,7 @@ var (
 	remotePort   = flag.Int("p", 2200, "")
 	localServer  = flag.String("ls", "localhost", "")
 	localPort    = flag.Int("lp", 7500, "")
+	keepAlive    = flag.Bool("a", false, "")
 )
 
 func main() {
@@ -82,6 +86,16 @@ func main() {
 		log.Fatalln(fmt.Printf("Dial INTO remote server error: %s", err))
 	}
 
+	connectionDone := make(chan struct{})
+	if *keepAlive {
+		go func() {
+			err = keepAliveTicker(serverConn, connectionDone)
+			if err != nil {
+				log.Fatalln(fmt.Printf("Cannot initialize keep alive on connection: %s", err))
+			}
+		}()
+	}
+
 	go func() {
 		session, err := serverConn.NewSession()
 		if err != nil {
@@ -149,6 +163,22 @@ func handleClient(client net.Conn, remote net.Conn) {
 	<-chDone
 }
 
+func keepAliveTicker(cl *ssh.Client, done <-chan struct{}) error {
+	t := time.NewTicker(time.Minute)
+	defer t.Stop()
+	for {
+		select {
+		case <-t.C:
+			_, _, err := cl.SendRequest("keepalive", true, nil)
+			if err != nil {
+				return errors.Wrap(err, "failed to send keep alive")
+			}
+		case <-done:
+			return nil
+		}
+	}
+}
+
 func randomPort(min, max int) int {
 	rand.Seed(time.Now().Unix())
 	return rand.Intn(max-min) + min