fix: Wait for connections before port-forwarding (#4057)

kylecarbs · web-flow · commit ec453f01e451 · 2022-09-14T21:57:42.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)
diff --git a/go.mod b/go.mod
@@ -47,7 +47,7 @@ replace github.com/tcnksm/go-httpstat => github.com/kylecarbs/go-httpstat v0.0.0
 
 // There are a few minor changes we make to Tailscale that we're slowly upstreaming. Compare here:
 // https://github.com/tailscale/tailscale/compare/main...coder:tailscale:main
-replace tailscale.com => github.com/coder/tailscale v1.1.1-0.20220912224234-e80caec6c05f
+replace tailscale.com => github.com/coder/tailscale v1.1.1-0.20220914175845-85b85d9a52ee
 
 // Switch to our fork that imports fixes from http://github.com/tailscale/ssh.
 // See: https://github.com/coder/coder/issues/3371
diff --git a/go.sum b/go.sum
@@ -355,8 +355,8 @@ github.com/coder/retry v1.3.0 h1:5lAAwt/2Cm6lVmnfBY7sOMXcBOwcwJhmV5QGSELIVWY=
 github.com/coder/retry v1.3.0/go.mod h1:tXuRgZgWjUnU5LZPT4lJh4ew2elUhexhlnXzrJWdyFY=
 github.com/coder/ssh v0.0.0-20220811105153-fcea99919338 h1:tN5GKFT68YLVzJoA8AHuiMNJ0qlhoD3pGN3JY9gxSko=
 github.com/coder/ssh v0.0.0-20220811105153-fcea99919338/go.mod h1:ZSS+CUoKHDrqVakTfTWUlKSr9MtMFkC4UvtQKD7O914=
-github.com/coder/tailscale v1.1.1-0.20220912224234-e80caec6c05f h1:NN9O1Pgno2QQy+JBnZk1VQ3vyAmWaB+yEotUDEuFKm8=
-github.com/coder/tailscale v1.1.1-0.20220912224234-e80caec6c05f/go.mod h1:5amxy08qijEa8bcTW2SeIy4MIqcmd7LMsuOxqOlj2Ak=
+github.com/coder/tailscale v1.1.1-0.20220914175845-85b85d9a52ee h1:77WUcIAL5FRQtd97/gOV66MJHLPhsPw+3vMxoEvcadI=
+github.com/coder/tailscale v1.1.1-0.20220914175845-85b85d9a52ee/go.mod h1:5amxy08qijEa8bcTW2SeIy4MIqcmd7LMsuOxqOlj2Ak=
 github.com/coder/wireguard-go v0.0.0-20220913030931-b1b3bb45caf9 h1:AeU4w8hSB+XEj3e8HjvEUTy/MWQd6tddnr9dELrRjKk=
 github.com/coder/wireguard-go v0.0.0-20220913030931-b1b3bb45caf9/go.mod h1:enML0deDxY1ux+B6ANGiwtg0yAJi1rctkTpcHNAVPyg=
 github.com/containerd/aufs v0.0.0-20200908144142-dab0cbea06f4/go.mod h1:nukgQABAEopAHvB6j7cnP5zJ+/3aVcE7hCYqvIwAHyE=
