Skip to content

Commit

Permalink
feat(client): add keep-alive option for client
Browse files Browse the repository at this point in the history
  • Loading branch information
jkuri committed Apr 17, 2018
1 parent d0ab23b commit 48b0ce0
Showing 1 changed file with 31 additions and 1 deletion.
32 changes: 31 additions & 1 deletion cmd/vex/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import (
"os/signal"
"time"

"github.com/pkg/errors"
"golang.org/x/crypto/ssh"
)

Expand All @@ -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
Expand All @@ -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() {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 48b0ce0

Please sign in to comment.