client/main.go

@@ -11,7 +11,9 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"os/signal"
 	"sync"
+	"time"
 
 	"github.com/jamiealquiza/envy"
 	"github.com/myesui/uuid"
@@ -25,17 +27,26 @@ var server *string
 var hostname *string
 var target string
 
-func stage1(print bool) net.Conn {
-	log.SetFlags(log.Lshortfile)
+var connectLock sync.Mutex
 
+func stage1(print bool) net.Conn {
 	conf := &tls.Config{
 		ServerName: *server,
 	}
-
-	conn, err := tls.Dial("tcp", fmt.Sprintf("%s:443", *server), conf)
-	if err != nil {
-		panic(err)
+	var err error
+	var conn net.Conn
+	backoff := time.Second * 10
+	connectLock.Lock()
+	for {
+		conn, err = tls.Dial("tcp", fmt.Sprintf("%s:443", *server), conf)
+		if err == nil {
+			break
+		}
+		fmt.Printf("error while connecting to server: %s\n", err)
+		time.Sleep(backoff)
+		backoff = backoff + (time.Second * 10)
 	}
+	connectLock.Unlock()
 	msg := fmt.Sprintf("backend-open:%s:%s", *token, *hostname)
 	n, err := conn.Write([]byte(msg))
 	if err != nil {
@@ -95,6 +106,22 @@ func stage2(conn net.Conn) {
 	tConn.Close()
 }
 
+func shutdown() {
+	conf := &tls.Config{
+		ServerName: *server,
+	}
+
+	conn, err := tls.Dial("tcp", fmt.Sprintf("%s:443", *server), conf)
+	if err != nil {
+		panic(err)
+	}
+	msg := fmt.Sprintf("backend-shutdown:%s:%s", *token, *hostname)
+	_, err = conn.Write([]byte(msg))
+	if err != nil {
+		panic(err)
+	}
+}
+
 func main() {
 	log.SetFlags(log.Lshortfile)
 
@@ -120,7 +147,9 @@ func main() {
 		go both()
 	}
 
-	lock := sync.Mutex{}
-	lock.Lock()
-	lock.Lock()
+	signalChannel := make(chan os.Signal, 1)
+	signal.Notify(signalChannel, os.Interrupt)
+	<-signalChannel
+	shutdown()
+
 }

go.mod

@@ -10,5 +10,4 @@ require (
 	github.com/spf13/cobra v0.0.5 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
 	gitlab.com/gartnera/golib v0.0.0-20190619234814-0708a17f76de
-	golang.org/x/sys v0.0.0-20190616124812-15dcb6c0061f
 )

go.sum

@@ -12,7 +12,6 @@ github.com/corpix/uarand v0.1.1 h1:RMr1TWc9F4n5jiPDzFHtmaUXLKLNUFK0SgCLo4BhX/U=
 github.com/corpix/uarand v0.1.1/go.mod h1:SFKZvkcRoLqVRFZ4u25xPmp6m9ktANfbpXZ7SJ0/FNU=
 github.com/cpuguy83/go-md2man v1.0.10/go.mod h1:SmD6nW6nTyfqj6ABTjUi3V3JVMnlJmwcJI5acqYI6dE=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/docker/distribution v2.7.1+incompatible/go.mod h1:J2gT2udsDAN96Uj4KfcMRqY0/ypR+oyYUYmja8H+y+w=
 github.com/docker/docker v1.13.1/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
@@ -37,7 +36,6 @@ github.com/opencontainers/go-digest v1.0.0-rc1/go.mod h1:cMLVZDEM3+U2I4VmLI6N8jQ
 github.com/pelletier/go-toml v1.2.0/go.mod h1:5z9KED0ma1S8pY6P1sdut58dfprrGBbd/94hg7ilaic=
 github.com/pkg/browser v0.0.0-20180916011732-0a3d74bf9ce4/go.mod h1:4OwLy04Bl9Ef3GJJCoec+30X3LQs/0/m4HFRt/2LUSA=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
-github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/russross/blackfriday v1.5.2/go.mod h1:JO/DiYxRf+HjHt06OyowR9PTA263kcR/rfWxYHBV53g=
 github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6MwdIuYE2rE=
@@ -55,7 +53,6 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+
 github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.2.0/go.mod h1:qt09Ya8vawLte6SNmTgCsAVtYtaKzEcn8ATUoHMkEqE=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
-github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/ugorji/go/codec v0.0.0-20181204163529-d75b2dcb6bc8/go.mod h1:VFNgLljTbGfSG7qAOspJ7OScBnGdDN/yBr0sguwnwf0=
 github.com/xordataexchange/crypt v0.0.3-0.20170626215501-b2862e3d0a77/go.mod h1:aYKd//L2LvnjZzWKhF00oedf4jCCReLcmhLdhm1A27Q=
@@ -74,7 +71,6 @@ golang.org/x/sys v0.0.0-20181205085412-a5c9d58dba9a/go.mod h1:STP8DvDyc/dI5b8T5h
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20190422165155-953cdadca894/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20190616124812-15dcb6c0061f h1:25KHgbfyiSm6vwQLbM3zZIe1v9p/3ea4Rz+nnM5K/i4=
 golang.org/x/sys v0.0.0-20190616124812-15dcb6c0061f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=

server/main.go

@@ -162,13 +162,26 @@ func handleBackend(conn net.Conn, serverName string) {
 		return
 	}
 	cmd := ss[0]
+	secret := ss[1]
+	session, existingSessionFound := state.secretMap[secret]
+	if cmd == "backend-shutdown" {
+		defer conn.Close()
+		if !existingSessionFound {
+			fmt.Printf("invalid shutdown command: %s\n", cmd)
+			return
+		}
+		fmt.Printf("shutdown requested for %s\n", session.hostname)
+		state.Lock()
+		defer state.Unlock()
+		delete(state.hostnameMap, session.hostname)
+		delete(state.secretMap, session.secret)
+		return
+	}
 	if cmd != "backend-open" {
 		fmt.Printf("unknown cmd: %s\n", cmd)
 		conn.Close()
 		return
 	}
-	secret := ss[1]
-	session, existingSessionFound := state.secretMap[secret]
 	if existingSessionFound {
 		session.backendConnected(conn)
 		return