Fixed #1

Signed-off-by: Vishal Rana <vr@labstack.com>

Author: vishr

Gopkg.lock

@@ -10,6 +10,8 @@ import (
 	"github.com/labstack/tunnel"
 	"github.com/labstack/tunnel/util"
 
+	"time"
+
 	"github.com/mitchellh/go-homedir"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
@@ -24,21 +26,22 @@ var (
 	user       string
 	rootCmd    = &cobra.Command{
 		Use:   "tunnel",
-		Short: "Tunnel lets you expose local servers to the internet securely",
+		Short: "Tunnel lets you expose local servers to internet securely",
 		Long:  ``,
 		Args:  cobra.MinimumNArgs(1),
 		Run: func(cmd *cobra.Command, args []string) {
 			c := &tunnel.Config{
 				Host:       "labstack.me:22",
 				RemoteHost: "0.0.0.0",
 				RemotePort: 80,
+				Channel:    make(chan int),
 			}
 			e := new(tunnel.Error)
 
 			if name != "" {
 				key := viper.GetString("api_key")
 				if key == "" {
-					log.Fatalf("Failed to find api key in the config")
+					log.Fatalf("failed to find api key in the config")
 				}
 
 				// Find config
@@ -50,9 +53,9 @@ var (
 					SetHeader("User-Agent", "labstack/tunnel").
 					Get(fmt.Sprintf("https://api.labstack.com/tunnel/configs/%s", name))
 				if err != nil {
-					log.Fatalf("Failed to the find tunnel: %v", err)
+					log.Fatalf("failed to the find tunnel: %v", err)
 				} else if res.StatusCode() != http.StatusOK {
-					log.Fatalf("Failed to the find tunnel: %s", e.Message)
+					log.Fatalf("failed to the find tunnel: %s", e.Message)
 				}
 				if c.Protocol == "tcp" {
 					tcp = true
@@ -69,12 +72,19 @@ var (
 			c.User = user
 			c.TargetHost, c.TargetPort, err = util.SplitHostPort(args[0])
 			if err != nil {
-				log.Fatalf("Failed to parse target address: %v", err)
+				log.Fatalf("failed to parse target address: %v", err)
 			}
 			if tcp || tls {
 				c.RemotePort = 0
 			}
-			tunnel.Create(c)
+		CREATE:
+			go tunnel.Create(c)
+			event := <-c.Channel
+			if event == tunnel.EventReconnect {
+				log.Info("trying to reconnect")
+				time.Sleep(1 * time.Second)
+				goto CREATE
+			}
 		},
 	}
 )
@@ -105,7 +115,7 @@ func initConfig() {
 		// Find home directory
 		home, err := homedir.Dir()
 		if err != nil {
-			log.Fatalf("Failed to find the home directory: %v", err)
+			log.Fatalf("failed to find the home directory: %v", err)
 		}
 
 		// Search config in home directory with name ".tunnel" (without extension)

cmd/root.go

@@ -1,7 +1,11 @@
 package main
 
-import "github.com/labstack/tunnel/cmd"
+import (
+	"github.com/labstack/gommon/log"
+	"github.com/labstack/tunnel/cmd"
+)
 
 func main() {
+	log.SetHeader("${time_rfc3339} ${level}")
 	cmd.Execute()
 }

cmd/tunnel/main.go

@@ -26,6 +26,7 @@ type (
 		RemotePort int
 		TargetHost string
 		TargetPort int
+		Channel    chan int
 		HideBanner bool
 	}
 
@@ -35,14 +36,19 @@ type (
 	}
 )
 
+const (
+	_ = iota
+	EventReconnect
+)
+
 var (
 	hostBytes = []byte("ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAACAQDoSLknvlFrFzroOlh1cqvcIFelHO+Wvj1UZ/p3J9bgsJGiKfh3DmBqEw1DOEwpHJz4zuV375TyjGuHuGZ4I4xztnwauhFplfEvriVHQkIDs6UnGwJVr15XUQX04r0i6mLbJs5KqIZTZuZ9ZGOj7ZWnaA7C07nPHGrERKV2Fm67rPvT6/qFikdWUbCt7KshbzdwwfxUohmv+NI7vw2X6vPU8pDaNEY7vS3YgwD/WlvQx+WDF2+iwLVW8OWWjFuQso6Eg1BSLygfPNhAHoiOWjDkijc8U9LYkUn7qsDCnvJxCoTTNmdECukeHfzrUjTSw72KZoM5KCRV78Wrctai1Qn6yRQz9BOSguxewLfzHtnT43/MLdwFXirJ/Ajquve2NAtYmyGCq5HcvpDAyi7lQ0nFBnrWv5zU3YxrISIpjovVyJjfPx8SCRlYZwVeUq6N2yAxCzJxbElZPtaTSoXBIFtoas2NXnCWPgenBa/2bbLQqfgbN8VQ9RaUISKNuYDIn4+eO72+RxF9THzZeV17pnhTVK88XU4asHot1gXwAt4vEhSjdUBC9KUIkfukI6F4JFxtvuO96octRahdV1Qg0vF+D0+SPy2HxqjgZWgPE2Xh/NmuIXwbE0wkymR2wrgj8Hd4C92keo2NBRh9dD7D2negnVYaYsC+3k/si5HNuCHnHQ== tunnel@labstack.com")
 )
 
 func Create(c *Config) {
 	hostKey, _, _, _, err := ssh.ParseAuthorizedKey(hostBytes)
 	if err != nil {
-		log.Fatalf("Failed to parse host key: %v", err)
+		log.Fatalf("failed to parse host key: %v", err)
 	}
 	config := &ssh.ClientConfig{
 		User: c.User,
@@ -64,11 +70,11 @@ func Create(c *Config) {
 	if proxy != "" {
 		proxyURL, err := url.Parse(proxy)
 		if err != nil {
-			log.Fatalf("Cannot open new session: %v", err)
+			log.Fatalf("cannot open new session: %v", err)
 		}
 		tcp, err := net.Dial("tcp", proxyURL.Hostname())
 		if err != nil {
-			log.Fatalf("Cannot open new session: %v", err)
+			log.Fatalf("cannot open new session: %v", err)
 		}
 		connReq := &http.Request{
 			Method: "CONNECT",
@@ -84,27 +90,29 @@ func Create(c *Config) {
 		connReq.Write(tcp)
 		resp, err := http.ReadResponse(bufio.NewReader(tcp), connReq)
 		if err != nil {
-			log.Fatalf("Cannot open new session: %v", err)
+			log.Fatalf("cannot open new session: %v", err)
 		}
 		defer resp.Body.Close()
 
 		conn, chans, reqs, err := ssh.NewClientConn(tcp, c.Host, config)
 		if err != nil {
-			log.Fatalf("Cannot open new session: %v", err)
+			log.Fatalf("cannot open new session: %v", err)
 		}
 		client = ssh.NewClient(conn, chans, reqs)
 	} else {
 		client, err = ssh.Dial("tcp", c.Host, config)
 	}
 	if err != nil {
-		log.Fatalf("Failed to connect: %v", err)
+		log.Errorf("failed to connect: %v", err)
+		c.Channel <- EventReconnect
+		return
 	}
 	defer client.Close()
 
 	// Session
 	sess, err := client.NewSession()
 	if err != nil {
-		log.Fatalf("Failed to create session: %v", err)
+		log.Fatalf("failed to create session: %v", err)
 	}
 	defer sess.Close()
 	r, err := sess.StdoutPipe()
@@ -116,7 +124,12 @@ func Create(c *Config) {
 		for {
 			line, _, err := br.ReadLine()
 			if err != nil {
-				log.Fatalf("Failed to read: %v", err)
+				if err == io.EOF {
+					c.Channel <- EventReconnect
+					return
+				} else {
+					log.Fatalf("failed to read: %v", err)
+				}
 			}
 			fmt.Printf("%s\n", line)
 		}
@@ -125,15 +138,15 @@ func Create(c *Config) {
 	// Remote listener
 	ln, err := client.Listen("tcp", fmt.Sprintf("%s:%d", c.RemoteHost, c.RemotePort))
 	if err != nil {
-		log.Fatalf("Failed to listen on remote host: %v", err)
+		log.Fatalf("failed to listen on remote host: %v", err)
 	}
 	defer ln.Close()
 
 	for {
 		// Handle inbound connection
 		in, err := ln.Accept()
 		if err != nil {
-			log.Printf("Failed to accept connection: %v", err)
+			log.Printf("failed to accept connection: %v", err)
 			return
 		}
 
@@ -143,7 +156,7 @@ func Create(c *Config) {
 			// Target connection
 			out, err := net.Dial("tcp", fmt.Sprintf("%s:%d", c.TargetHost, c.TargetPort))
 			if err != nil {
-				log.Printf("Failed to connect to target: %v", err)
+				log.Printf("failed to connect to target: %v", err)
 				return
 			}
 			defer out.Close()
@@ -160,7 +173,7 @@ func Create(c *Config) {
 			// Handle error
 			err = <-errCh
 			if err != nil && err != io.EOF {
-				log.Printf("Failed to copy: %v", err)
+				log.Printf("failed to copy: %v", err)
 			}
 		}(in)
 	}