fix: Wait for connections before port-forwarding

kylecarbs · kylecarbs · commit 5b972acaa19d · 2022-09-14T17:40:12.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)
diff --git a/cli/schedule_test.go b/cli/schedule_test.go
@@ -61,22 +61,20 @@ func TestScheduleShow(t *testing.T) {
 		t.Parallel()
 
 		var (
-			ctx       = context.Background()
 			client    = coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
 			user      = coderdtest.CreateFirstUser(t, client)
 			version   = coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil)
 			_         = coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
 			project   = coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
 			workspace = coderdtest.CreateWorkspace(t, client, user.OrganizationID, project.ID, func(cwr *codersdk.CreateWorkspaceRequest) {
 				cwr.AutostartSchedule = nil
+				cwr.TTLMillis = nil
 			})
+			_         = coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
 			cmdArgs   = []string{"schedule", "show", workspace.Name}
 			stdoutBuf = &bytes.Buffer{}
 		)
 
-		// unset workspace TTL
-		require.NoError(t, client.UpdateWorkspaceTTL(ctx, workspace.ID, codersdk.UpdateWorkspaceTTLRequest{TTLMillis: nil}))
-
 		cmd, root := clitest.New(t, cmdArgs...)
 		clitest.SetupConfig(t, client, root)
 		cmd.SetOut(stdoutBuf)
