Clean up code

main.go

@@ -6,7 +6,6 @@ package main
 
 import (
 	"github.com/rancher/tunnelware/pkg/cmd"
-	_ "github.com/rancher/tunnelware/pkg/generated/controllers/cert-manager.io"
 	cli "github.com/rancher/wrangler-cli"
 	_ "github.com/rancher/wrangler/pkg/generated/controllers/networking.k8s.io"
 )

pkg/cmd/client.go

@@ -9,12 +9,12 @@ import (
 	"os/exec"
 	"path/filepath"
 	"strings"
+	"time"
 
 	"github.com/dgrijalva/jwt-go"
 	"github.com/pkg/browser"
 	"github.com/pkg/errors"
 	"github.com/rancher/remotedialer"
-	"github.com/rancher/tunnelware/pkg/questions"
 	cli "github.com/rancher/wrangler-cli"
 	"github.com/rancher/wrangler/pkg/randomtoken"
 	"github.com/sirupsen/logrus"
@@ -88,24 +88,24 @@ func (c *Client) Run(cmd *cobra.Command, args []string) error {
 			}
 		}
 
-		finished, err := questions.PromptBool("Once you finish github login, enter y to continue", true)
-		if err != nil {
-			return err
-		}
-		if !finished {
-			return nil
-		}
-		resp, err := http.Get(fmt.Sprintf("https://%s/jwt/%s", serverURL.Host, state))
-		if err != nil {
-			return err
-		}
-		defer resp.Body.Close()
-		token, err := ioutil.ReadAll(resp.Body)
-		if err != nil {
-			return err
-		}
-		if err := ioutil.WriteFile(tokenFile, token, 0755); err != nil {
-			return err
+		for {
+			resp, err := http.Get(fmt.Sprintf("https://%s/jwt/%s", serverURL.Host, state))
+			if err != nil {
+				return err
+			}
+			defer resp.Body.Close()
+
+			if resp.StatusCode == 503 {
+				continue
+			}
+			token, err := ioutil.ReadAll(resp.Body)
+			if err != nil {
+				return err
+			}
+			if err := ioutil.WriteFile(tokenFile, token, 0755); err != nil {
+				return err
+			}
+			break
 		}
 	}
 
@@ -142,7 +142,10 @@ func (c *Client) Run(cmd *cobra.Command, args []string) error {
 	domainWithUsername := fmt.Sprintf("%s.%s", strings.ToLower(claim.Username), domain)
 	go func() {
 		fmt.Printf("http://%v --------> %v\nhttps://%v -------> %v\n", domainWithUsername, forwardHost, domainWithUsername, forwardHost)
-		remotedialer.ClientConnect(cmd.Context(), u.String(), headers, nil, func(string, string) bool { return true }, nil)
+		for {
+			remotedialer.ClientConnect(cmd.Context(), u.String(), headers, nil, func(string, string) bool { return true }, nil)
+			time.Sleep(time.Second * 5)
+		}
 	}()
 	<-cmd.Context().Done()
 	return nil

pkg/cmd/server.go

@@ -305,8 +305,12 @@ func (h *Handler) jwtToken(rw http.ResponseWriter, req *http.Request) {
 	vars := mux.Vars(req)
 	state := vars["state"]
 	jwtToken := jwtTokenMap[state]
-	defer delete(jwtTokenMap, state)
+	if jwtToken == "" {
+		rw.WriteHeader(503)
+		return
+	}
 
+	defer delete(jwtTokenMap, state)
 	rw.Write([]byte(jwtToken))
 	rw.WriteHeader(200)
 }