tstest/integration: test `tailscale up` when device approval is required

This patch extends the integration tests for `tailscale up` to include tailnets
where new devices need to be approved. It doesn't change the CLI, because it's
mostly working correctly already -- these tests are just to prevent future
regressions.

I've added support for `MachineAuthorized` to mock control, and I've refactored
`TestOneNodeUpAuth` to be more flexible. It now takes a sequence of steps to
run and asserts whether we got a login URL and/or machine approval URL after
each step.

Updates tailscale/corp#31476
Updates #17361

Signed-off-by: Alex Chan <alexc@tailscale.com>

tstest/integration/integration.go

@@ -1099,20 +1099,40 @@ func (tt *trafficTrap) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 type authURLParserWriter struct {
 	buf bytes.Buffer
-	fn  func(urlStr string) error
+	// Handle login URLs, and count how many times they were seen
+	authURLFn func(urlStr string) error
+	// Handle machine approval URLs, and count how many times they were seen.
+	deviceApprovalURLFn func(urlStr string) error
 }
 
+// Note: auth URLs from testcontrol look slightly different to real auth URLs,
+// e.g. http://127.0.0.1:60456/auth/96af2ff7e04ae1499a9a
 var authURLRx = regexp.MustCompile(`(https?://\S+/auth/\S+)`)
 
+// Looks for any device approval URL, which is any URL ending with `/admin`
+// e.g. http://127.0.0.1:60456/admin
+var deviceApprovalURLRx = regexp.MustCompile(`(https?://\S+/admin)[^\S]`)
+
 func (w *authURLParserWriter) Write(p []byte) (n int, err error) {
 	n, err = w.buf.Write(p)
+
+	defer w.buf.Reset() // so it's not matched again
+
 	m := authURLRx.FindSubmatch(w.buf.Bytes())
 	if m != nil {
 		urlStr := string(m[1])
-		w.buf.Reset() // so it's not matched again
-		if err := w.fn(urlStr); err != nil {
+		if err := w.authURLFn(urlStr); err != nil {
 			return 0, err
 		}
 	}
+
+	m = deviceApprovalURLRx.FindSubmatch(w.buf.Bytes())
+	if m != nil && w.deviceApprovalURLFn != nil {
+		urlStr := string(m[1])
+		if err := w.deviceApprovalURLFn(urlStr); err != nil {
+			return 0, err
+		}
+	}
+
 	return n, err
 }

tstest/integration/integration_test.go

@@ -268,7 +268,65 @@ func TestStateSavedOnStart(t *testing.T) {
 	d1.MustCleanShutdown(t)
 }
 
+// This handler receives auth URLs, and logs into control.
+//
+// It counts how many URLs it sees, and will fail the test if it
+// sees multiple login URLs.
+func completeLogin(t *testing.T, control *testcontrol.Server, counter *atomic.Int32) func(string) error {
+	return func(urlStr string) error {
+		t.Logf("saw auth URL %q", urlStr)
+		if control.CompleteAuth(urlStr) {
+			if counter.Add(1) > 1 {
+				err := errors.New("completed multiple auth URLs")
+				t.Error(err)
+				return err
+			}
+			t.Logf("completed login to %s", urlStr)
+			return nil
+		} else {
+			err := fmt.Errorf("failed to complete initial login to %q", urlStr)
+			t.Fatal(err)
+			return err
+		}
+	}
+}
+
+// This handler receives device approval URLs, and approves the device.
+//
+// It counts how many URLs it sees, and will fail the test if it
+// sees multiple device approval URLs.
+func completeDeviceApproval(t *testing.T, node *TestNode, counter *atomic.Int32) func(string) error {
+	return func(urlStr string) error {
+		control := node.env.Control
+		nodeKey := node.MustStatus().Self.PublicKey
+		t.Logf("saw device approval URL %q", urlStr)
+		if control.CompleteDeviceApproval(&nodeKey) {
+			if counter.Add(1) > 1 {
+				err := errors.New("completed multiple device approval URLs")
+				t.Error(err)
+				return err
+			}
+			t.Log("completed device approval")
+			return nil
+		} else {
+			err := errors.New("failed to complete device approval")
+			t.Fatal(err)
+			return err
+		}
+	}
+}
+
 func TestOneNodeUpAuth(t *testing.T) {
+	type step struct {
+		args []string
+		//
+		// Do we expect to log in again with a new /auth/ URL?
+		wantAuthURL bool
+		//
+		// Do we expect to need a device approval URL?
+		wantDeviceApprovalURL bool
+	}
+
 	for _, tt := range []struct {
 		name string
 		args []string
@@ -276,65 +334,112 @@ func TestOneNodeUpAuth(t *testing.T) {
 		// What auth key should we use for control?
 		authKey string
 		//
-		// Is tailscaled already logged in before we run this `up` command?
-		alreadyLoggedIn bool
+		// Do we require device approval in the tailnet?
+		requireDeviceApproval bool
 		//
-		// Do we need to log in again with a new /auth/ URL?
-		needsNewAuthURL bool
+		// What CLI commands should we run in this test?
+		steps []step
 	}{
 		{
 			name: "up",
-			args:            []string{"up"},
-			needsNewAuthURL: true,
+			steps: []step{
+				{args: []string{"up"}, wantAuthURL: true},
+			},
+		},
+		{
+			name: "up-with-machine-auth",
+			steps: []step{
+				{args: []string{"up"}, wantAuthURL: true, wantDeviceApprovalURL: true},
+			},
+			requireDeviceApproval: true,
 		},
 		{
 			name: "up-with-force-reauth",
-			args:            []string{"up", "--force-reauth"},
-			needsNewAuthURL: true,
+			steps: []step{
+				{args: []string{"up", "--force-reauth"}, wantAuthURL: true},
+			},
 		},
 		{
 			name:    "up-with-auth-key",
-			args:            []string{"up", "--auth-key=opensesame"},
 			authKey: "opensesame",
-			needsNewAuthURL: false,
+			steps: []step{
+				{args: []string{"up", "--auth-key=opensesame"}},
+			},
+		},
+		{
+			name:    "up-with-auth-key-with-machine-auth",
+			authKey: "opensesame",
+			steps: []step{
+				{
+					args:                  []string{"up", "--auth-key=opensesame"},
+					wantAuthURL:           false,
+					wantDeviceApprovalURL: true,
+				},
+			},
+			requireDeviceApproval: true,
 		},
 		{
 			name:    "up-with-force-reauth-and-auth-key",
-			args:            []string{"up", "--force-reauth", "--auth-key=opensesame"},
 			authKey: "opensesame",
-			needsNewAuthURL: false,
+			steps: []step{
+				{args: []string{"up", "--force-reauth", "--auth-key=opensesame"}},
+			},
 		},
 		{
 			name: "up-after-login",
-			args:            []string{"up"},
-			alreadyLoggedIn: true,
-			needsNewAuthURL: false,
+			steps: []step{
+				{args: []string{"up"}, wantAuthURL: true},
+				{args: []string{"up"}, wantAuthURL: false},
+			},
+		},
+		{
+			name: "up-after-login-with-machine-approval",
+			steps: []step{
+				{args: []string{"up"}, wantAuthURL: true, wantDeviceApprovalURL: true},
+				{args: []string{"up"}, wantAuthURL: false, wantDeviceApprovalURL: false},
+			},
+			requireDeviceApproval: true,
 		},
 		{
 			name: "up-with-force-reauth-after-login",
-			args:            []string{"up", "--force-reauth"},
-			alreadyLoggedIn: true,
-			needsNewAuthURL: true,
+			steps: []step{
+				{args: []string{"up"}, wantAuthURL: true},
+				{args: []string{"up", "--force-reauth"}, wantAuthURL: true},
+			},
+		},
+		{
+			name: "up-with-force-reauth-after-login-with-machine-approval",
+			steps: []step{
+				{args: []string{"up"}, wantAuthURL: true, wantDeviceApprovalURL: true},
+				{args: []string{"up", "--force-reauth"}, wantAuthURL: true, wantDeviceApprovalURL: false},
+			},
+			requireDeviceApproval: true,
 		},
 		{
 			name:    "up-with-auth-key-after-login",
-			args:            []string{"up", "--auth-key=opensesame"},
 			authKey: "opensesame",
-			alreadyLoggedIn: true,
-			needsNewAuthURL: false,
+			steps: []step{
+				{args: []string{"up", "--auth-key=opensesame"}},
+				{args: []string{"up", "--auth-key=opensesame"}},
+			},
 		},
 		{
 			name:    "up-with-force-reauth-and-auth-key-after-login",
-			args:            []string{"up", "--force-reauth", "--auth-key=opensesame"},
 			authKey: "opensesame",
-			alreadyLoggedIn: true,
-			needsNewAuthURL: false,
+			steps: []step{
+				{args: []string{"up", "--auth-key=opensesame"}},
+				{args: []string{"up", "--force-reauth", "--auth-key=opensesame"}},
+			},
 		},
 	} {
 		tstest.Shard(t)
 
 		for _, useSeamlessKeyRenewal := range []bool{true, false} {
-			t.Run(fmt.Sprintf("%s-seamless-%t", tt.name, useSeamlessKeyRenewal), func(t *testing.T) {
+			name := tt.name
+			if useSeamlessKeyRenewal {
+				name += "-with-seamless"
+			}
+			t.Run(name, func(t *testing.T) {
 				tstest.Parallel(t)
 
 				env := NewTestEnv(t, ConfigureControl(
@@ -345,6 +450,10 @@ func TestOneNodeUpAuth(t *testing.T) {
 							control.RequireAuth = true
 						}
 
+						if tt.requireDeviceApproval {
+							control.RequireMachineAuth = true
+						}
+
 						control.AllNodesSameUser = true
 
 						if useSeamlessKeyRenewal {
@@ -359,69 +468,45 @@ func TestOneNodeUpAuth(t *testing.T) {
 				d1 := n1.StartDaemon()
 				defer d1.MustCleanShutdown(t)
 
-				cmdArgs := append(tt.args, "--login-server="+env.ControlURL())
+				for i, step := range tt.steps {
+					t.Logf("Running step %d", i)
+					cmdArgs := append(step.args, "--login-server="+env.ControlURL())
 
-				// This handler looks for /auth/ URLs in the stdout from "tailscale up",
-				// and if it sees them, completes the auth process.
-				//
-				// It counts how many auth URLs it's seen.
-				var authCountAtomic atomic.Int32
-				authURLHandler := &authURLParserWriter{fn: func(urlStr string) error {
-					t.Logf("saw auth URL %q", urlStr)
-					if env.Control.CompleteAuth(urlStr) {
-						if authCountAtomic.Add(1) > 1 {
-							err := errors.New("completed multiple auth URLs")
-							t.Error(err)
-							return err
-						}
-						t.Logf("completed login to %s", urlStr)
-						return nil
-					} else {
-						err := fmt.Errorf("Failed to complete initial login to %q", urlStr)
-						t.Fatal(err)
-						return err
-					}
-				}}
+					t.Logf("Running command: %s", strings.Join(cmdArgs, " "))
 
-				// If we should be logged in at the start of the test case, go ahead
-				// and run the login command.
-				//
-				// Otherwise, just wait for tailscaled to be listening.
-				if tt.alreadyLoggedIn {
-					t.Logf("Running initial login: %s", strings.Join(cmdArgs, " "))
-					cmd := n1.Tailscale(cmdArgs...)
-					cmd.Stdout = authURLHandler
-					cmd.Stderr = cmd.Stdout
-					if err := cmd.Run(); err != nil {
-						t.Fatalf("up: %v", err)
-					}
-					authCountAtomic.Store(0)
-					n1.AwaitRunning()
-				} else {
-					n1.AwaitListening()
-				}
+					var authURLCount atomic.Int32
+					var deviceApprovalURLCount atomic.Int32
 
-				st := n1.MustStatus()
-				t.Logf("Status: %s", st.BackendState)
+					handler := &authURLParserWriter{
+						authURLFn:           completeLogin(t, env.Control, &authURLCount),
+						deviceApprovalURLFn: completeDeviceApproval(t, n1, &deviceApprovalURLCount),
+					}
 
-				t.Logf("Running command: %s", strings.Join(cmdArgs, " "))
 					cmd := n1.Tailscale(cmdArgs...)
-				cmd.Stdout = authURLHandler
+					cmd.Stdout = handler
+					cmd.Stdout = handler
 					cmd.Stderr = cmd.Stdout
-
 					if err := cmd.Run(); err != nil {
 						t.Fatalf("up: %v", err)
 					}
-				t.Logf("Got IP: %v", n1.AwaitIP4())
 
 					n1.AwaitRunning()
 
-				var expectedAuthUrls int32
-				if tt.needsNewAuthURL {
-					expectedAuthUrls = 1
+					var wantAuthURLCount int32
+					if step.wantAuthURL {
+						wantAuthURLCount = 1
+					}
+					if n := authURLCount.Load(); n != wantAuthURLCount {
+						t.Errorf("Auth URLs completed = %d; want %d", n, wantAuthURLCount)
+					}
+
+					var wantDeviceApprovalURLCount int32
+					if step.wantDeviceApprovalURL {
+						wantDeviceApprovalURLCount = 1
+					}
+					if n := deviceApprovalURLCount.Load(); n != wantDeviceApprovalURLCount {
+						t.Errorf("Device approval URLs completed = %d; want %d", n, wantDeviceApprovalURLCount)
 					}
-				if n := authCountAtomic.Load(); n != expectedAuthUrls {
-					t.Errorf("Auth URLs completed = %d; want %d", n, expectedAuthUrls)
 				}
 			})
 		}

tstest/integration/testcontrol/testcontrol.go

@@ -54,6 +54,7 @@ type Server struct {
 	DERPMap            *tailcfg.DERPMap // nil means to use prod DERP map
 	RequireAuth        bool
 	RequireAuthKey     string // required authkey for all nodes
+	RequireMachineAuth bool
 	Verbose            bool
 	DNSConfig          *tailcfg.DNSConfig // nil means no DNS config
 	MagicDNSDomain     string
@@ -686,6 +687,21 @@ func (s *Server) CompleteAuth(authPathOrURL string) bool {
 	return true
 }
 
+func (s *Server) CompleteDeviceApproval(nodeKey *key.NodePublic) bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	node, ok := s.nodes[*nodeKey]
+	if !ok {
+		return false
+	}
+
+	sendUpdate(s.updates[node.ID], updateSelfChanged)
+
+	node.MachineAuthorized = true
+	return true
+}
+
 func (s *Server) serveRegister(w http.ResponseWriter, r *http.Request, mkey key.MachinePublic) {
 	msg, err := io.ReadAll(io.LimitReader(r.Body, msgLimit))
 	r.Body.Close()
@@ -761,7 +777,7 @@ func (s *Server) serveRegister(w http.ResponseWriter, r *http.Request, mkey key.
 		s.nodes = map[key.NodePublic]*tailcfg.Node{}
 	}
 	_, ok := s.nodes[nk]
-	machineAuthorized := true // TODO: add Server.RequireMachineAuth
+	machineAuthorized := !s.RequireMachineAuth
 	if !ok {
 
 		nodeID := len(s.nodes) + 1