fix: Wait for connections before port-forwarding

kylecarbs · kylecarbs · commit 42371eb98bc8 · 2022-09-14T17:22:19.000Z
UDP packets were being dropped if a connection was started
before the Tailscale connection has been established.
diff --git a/cli/portforward.go b/cli/portforward.go
@@ -10,13 +10,16 @@ import (
 	"strings"
 	"sync"
 	"syscall"
+	"time"
 
 	"github.com/pion/udp"
 	"github.com/spf13/cobra"
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
+	"cdr.dev/slog/sloggers/sloghuman"
 	"github.com/coder/coder/agent"
+	"github.com/coder/coder/cli/cliflag"
 	"github.com/coder/coder/cli/cliui"
 	"github.com/coder/coder/codersdk"
 )
@@ -95,7 +98,11 @@ func portForward() *cobra.Command {
 			if !wireguard {
 				conn, err = client.DialWorkspaceAgent(ctx, workspaceAgent.ID, nil)
 			} else {
-				conn, err = client.DialWorkspaceAgentTailnet(ctx, slog.Logger{}, workspaceAgent.ID)
+				logger := slog.Logger{}
+				if cliflag.IsSetBool(cmd, varVerbose) {
+					logger = slog.Make(sloghuman.Sink(cmd.ErrOrStderr())).Named("tailnet").Leveled(slog.LevelDebug)
+				}
+				conn, err = client.DialWorkspaceAgentTailnet(ctx, logger, workspaceAgent.ID)
 			}
 			if err != nil {
 				return err
@@ -147,6 +154,22 @@ func portForward() *cobra.Command {
 				closeAllListeners()
 			}()
 
+			ticker := time.NewTicker(250 * time.Millisecond)
+			defer ticker.Stop()
+			for {
+				select {
+				case <-ctx.Done():
+					return ctx.Err()
+				case <-ticker.C:
+				}
+
+				_, err = conn.Ping()
+				if err != nil {
+					continue
+				}
+				break
+			}
+			ticker.Stop()
 			_, _ = fmt.Fprintln(cmd.OutOrStderr(), "Ready!")
 			wg.Wait()
 			return closeErr
diff --git a/cli/portforward_test.go b/cli/portforward_test.go
@@ -133,7 +133,7 @@ func TestPortForward(t *testing.T) {
 
 				// Launch port-forward in a goroutine so we can start dialing
 				// the "local" listener.
-				cmd, root := clitest.New(t, "port-forward", workspace.Name, flag)
+				cmd, root := clitest.New(t, "-v", "port-forward", workspace.Name, flag)
 				clitest.SetupConfig(t, client, root)
 				buf := newThreadSafeBuffer()
 				cmd.SetOut(buf)
@@ -179,7 +179,7 @@ func TestPortForward(t *testing.T) {
 
 				// Launch port-forward in a goroutine so we can start dialing
 				// the "local" listeners.
-				cmd, root := clitest.New(t, "port-forward", workspace.Name, flag1, flag2)
+				cmd, root := clitest.New(t, "-v", "port-forward", workspace.Name, flag1, flag2)
 				clitest.SetupConfig(t, client, root)
 				buf := newThreadSafeBuffer()
 				cmd.SetOut(buf)
@@ -234,7 +234,7 @@ func TestPortForward(t *testing.T) {
 
 		// Launch port-forward in a goroutine so we can start dialing
 		// the "local" listeners.
-		cmd, root := clitest.New(t, append([]string{"port-forward", workspace.Name}, flags...)...)
+		cmd, root := clitest.New(t, append([]string{"-v", "port-forward", workspace.Name}, flags...)...)
 		clitest.SetupConfig(t, client, root)
 		buf := newThreadSafeBuffer()
 		cmd.SetOut(buf)
