Add more agent tests

ammario · ammario · commit 5b906c1aaaf3 · 2022-09-01T17:46:52.000Z
diff --git a/agent/agent.go b/agent/agent.go
@@ -231,8 +231,10 @@ func (a *agent) runTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) {
 		for {
 			conn, err := reconnectingPTYListener.Accept()
 			if err != nil {
+				a.logger.Debug(ctx, "accept pty failed", slog.Error(err))
 				return
 			}
+			conn = a.stats.wrapConn(conn)
 			// This cannot use a JSON decoder, since that can
 			// buffer additional data that is required for the PTY.
 			rawLen := make([]byte, 2)
@@ -251,7 +253,7 @@ func (a *agent) runTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) {
 			if err != nil {
 				continue
 			}
-			go a.handleReconnectingPTY(ctx, msg, a.stats.wrapConn(conn))
+			go a.handleReconnectingPTY(ctx, msg, conn)
 		}
 	}()
 }
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -57,24 +57,59 @@ func TestAgent(t *testing.T) {
 			t.Run(fmt.Sprintf("tailscale=%v", tailscale), func(t *testing.T) {
 				t.Parallel()
 
-				var derpMap *tailcfg.DERPMap
-				if tailscale {
-					derpMap = tailnettest.RunDERPAndSTUN(t)
+				setupAgent := func(t *testing.T) (agent.Conn, <-chan *agent.Stats) {
+					var derpMap *tailcfg.DERPMap
+					if tailscale {
+						derpMap = tailnettest.RunDERPAndSTUN(t)
+					}
+					conn, stats := setupAgent(t, agent.Metadata{
+						DERPMap: derpMap,
+					}, 0)
+					assert.Empty(t, <-stats)
+					return conn, stats
 				}
-				conn, stats := setupAgent(t, agent.Metadata{
-					DERPMap: derpMap,
-				}, 0)
-				assert.Empty(t, <-stats)
-
-				sshClient, err := conn.SSHClient()
-				require.NoError(t, err)
-				session, err := sshClient.NewSession()
-				require.NoError(t, err)
-				defer session.Close()
 
-				assert.EqualValues(t, 1, (<-stats).NumConns)
-				assert.Greater(t, (<-stats).RxBytes, int64(0))
-				assert.Greater(t, (<-stats).TxBytes, int64(0))
+				t.Run("SSH", func(t *testing.T) {
+					t.Parallel()
+					conn, stats := setupAgent(t)
+
+					sshClient, err := conn.SSHClient()
+					require.NoError(t, err)
+					session, err := sshClient.NewSession()
+					require.NoError(t, err)
+					defer session.Close()
+
+					assert.EqualValues(t, 1, (<-stats).NumConns)
+					assert.Greater(t, (<-stats).RxBytes, int64(0))
+					assert.Greater(t, (<-stats).TxBytes, int64(0))
+				})
+
+				t.Run("ReconnectingPTY", func(t *testing.T) {
+					t.Parallel()
+
+					conn, stats := setupAgent(t)
+
+					ptyConn, err := conn.ReconnectingPTY(uuid.NewString(), 128, 128, "/bin/bash")
+					require.NoError(t, err)
+					defer ptyConn.Close()
+
+					data, err := json.Marshal(agent.ReconnectingPTYRequest{
+						Data: "echo test\r\n",
+					})
+					require.NoError(t, err)
+					_, err = ptyConn.Write(data)
+					require.NoError(t, err)
+
+					require.Eventually(t, func() bool {
+						s, ok := (<-stats)
+						return ok && s.NumConns > 0
+					}, testutil.WaitShort, testutil.IntervalFast,
+						"never saw conn",
+					)
+					assert.EqualValues(t, 1, (<-stats).NumConns)
+					assert.Greater(t, (<-stats).RxBytes, int64(0))
+					assert.Greater(t, (<-stats).TxBytes, int64(0))
+				})
 			})
 		}
 	})

Original file line number	Diff line number	Diff line change
`@@ -231,8 +231,10 @@ func (a agent) runTailnet(ctx context.Context, derpMap tailcfg.DERPMap) {`
`231`	`231`	`for {`
`232`	`232`	`conn, err := reconnectingPTYListener.Accept()`
`233`	`233`	`if err != nil {`
	`234`	`+ a.logger.Debug(ctx, "accept pty failed", slog.Error(err))`
`234`	`235`	`return`
`235`	`236`	`}`
	`237`	`+ conn = a.stats.wrapConn(conn)`
`236`	`238`	`// This cannot use a JSON decoder, since that can`
`237`	`239`	`// buffer additional data that is required for the PTY.`
`238`	`240`	`rawLen := make([]byte, 2)`
`@@ -251,7 +253,7 @@ func (a agent) runTailnet(ctx context.Context, derpMap tailcfg.DERPMap) {`
`251`	`253`	`if err != nil {`
`252`	`254`	`continue`
`253`	`255`	`}`
`254`		`- go a.handleReconnectingPTY(ctx, msg, a.stats.wrapConn(conn))`
	`256`	`+ go a.handleReconnectingPTY(ctx, msg, conn)`
`255`	`257`	`}`
`256`	`258`	`}()`
`257`	`259`	`}`