Dec 7, 2023 · Nov 30, 2023 · Nov 30, 2023 · Nov 30, 2023 · Nov 30, 2023 · Dec 1, 2023
diff --git a/agent/agent_test.go b/agent/agent_test.go
 package agent_test

 import (
 "bufio"
 "bytes"
 "context"
 "encoding/json"
 require.NoError(t, err)
 require.Equal(t, expected, strings.TrimSpace(string(output)))
 })
 t.Run("Tracks", func(t *testing.T) {
 t.Run("TracksVSCode", func(t *testing.T) {
 t.Parallel()
 if runtime.GOOS == "window" {
 t.Skip("Sleeping for infinity doesn't work on Windows")
 require.NoError(t, err)
 })

 // This test name must contain the string checked for by the agent, since it
 // looks for this string in the process name.
 //
 // This test sets up a port forward that emulates what Jetbrains IDE's do when
 // using gateway. The remote server side of the port forward is spun up using
 // the gotest process, which includes the test name.
 // So this unit test emulates a PID in the workspace with a similar
 // name to the jetbrains IDE. That makes the agent this this SSH port
 // forward is a "jetbrains" session.
 t.Run("TracksIdea.vendor.name=JetBrains", func(t *testing.T) {
 t.Run("TracksJetBrains", func(t *testing.T) {
 t.Parallel()
 if runtime.GOOS != "linux" {
 t.Skip("JetBrains tracking is only supported on Linux")
 }

 ctx := testutil.Context(t, testutil.WaitLong)

 rl, err := net.Listen("tcp", "127.0.0.1:0")
 // JetBrains tracking works by looking at the process name listening on the
 // forwarded port.  If the process's command line includes the magic string
 // we are looking for, then we assume it is a JetBrains editor.  So when we
 // connect to the port we must ensure the process includes that magic string
 // to fool the agent into thinking this is JetBrains.  To do this we need to
 // spawn an external process (in this case a simple echo server) so we can
 // control the process name.  The -D here is just to mimic how Java options
 // are set but is not necessary as the agent looks only for the magic
 // string itself anywhere in the command.
 _, b, _, ok := runtime.Caller(0)
 require.True(t, ok)
 dir := filepath.Join(filepath.Dir(b), "../scripts/echoserver/main.go")
 echoServerCmd := exec.Command("go", "run", dir,
 "-D", agentssh.MagicProcessCmdlineJetBrains)
 stdout, err := echoServerCmd.StdoutPipe()
 require.NoError(t, err)
 defer rl.Close()
 tcpAddr, valid := rl.Addr().(*net.TCPAddr)
 require.True(t, valid)
 remotePort := tcpAddr.Port
 go echoOnce(t, rl)
 err = echoServerCmd.Start()
 require.NoError(t, err)
 defer echoServerCmd.Process.Kill()

 // The echo server prints its port as the first line.
 sc := bufio.NewScanner(stdout)
 sc.Scan()
 remotePort := sc.Text()

 //nolint:dogsled
 conn, _, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
 sshClient, err := conn.SSHClient(ctx)
 require.NoError(t, err)

 tunneledConn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", remotePort))
 tunneledConn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%s", remotePort))
 require.NoError(t, err)
 t.Cleanup(func() {
 // always close on failure of test
 "never saw stats with conn open: %+v", s,
 )

 //Manually closingthe connection
 //Killtheserver andconnection after checking for the echo.
 requireEcho(t, tunneledConn)
 _ = rl.Close()
 _ = echoServerCmd.Process.Kill()
 _ = tunneledConn.Close()

 require.Eventuallyf(t, func() bool {
 var ok bool
diff --git a/scripts/echoserver/main.go b/scripts/echoserver/main.go
 package main

 // A simple echo server.  It listens on a random port, prints that port, then
 // echos back anything sent to it.

 import (
 "errors"
 "fmt"
 "io"
 "log"
 "net"
 )

 func main() {
 l, err := net.Listen("tcp", "127.0.0.1:0")
 if err != nil {
 log.Fatalf("listen error: err=%s", err)
 }

 defer l.Close()
 tcpAddr, valid := l.Addr().(*net.TCPAddr)
 if !valid {
 log.Fatal("address is not valid")
 }

 remotePort := tcpAddr.Port
 _, err = fmt.Println(remotePort)
 if err != nil {
 log.Fatalf("print error: err=%s", err)
 }

 for {
 conn, err := l.Accept()
 if err != nil {
 log.Fatalf("accept error, err=%s", err)
 return
 }

 go func() {
 defer conn.Close()
 _, err := io.Copy(conn, conn)

 if errors.Is(err, io.EOF) {
 return
 } else if err != nil {
 log.Fatalf("copy error, err=%s", err)
 }
 }()
 }
 }
Original file line number	Diff line number	Diff line change
		@@ -1,6 +1,7 @@
		package agent_test

		import (
		"bufio"
		"bytes"
		"context"
		"encoding/json"
Expand DownExpand Up		@@ -152,7 +153,7 @@ func TestAgent_Stats_Magic(t *testing.T) {
		require.NoError(t, err)
		require.Equal(t, expected, strings.TrimSpace(string(output)))
		})
		t.Run("Tracks", func(t *testing.T) {
		t.Run("TracksVSCode", func(t *testing.T) {
		t.Parallel()
		if runtime.GOOS == "window" {
		t.Skip("Sleeping for infinity doesn't work on Windows")
Expand DownExpand Up		@@ -192,36 +193,45 @@ func TestAgent_Stats_Magic(t *testing.T) {
		require.NoError(t, err)
		})

		// This test name must contain the string checked for by the agent, since it
		// looks for this string in the process name.
		//
		// This test sets up a port forward that emulates what Jetbrains IDE's do when
		// using gateway. The remote server side of the port forward is spun up using
		// the gotest process, which includes the test name.
		// So this unit test emulates a PID in the workspace with a similar
		// name to the jetbrains IDE. That makes the agent this this SSH port
		// forward is a "jetbrains" session.
		t.Run("TracksIdea.vendor.name=JetBrains", func(t *testing.T) {
		t.Run("TracksJetBrains", func(t *testing.T) {
		t.Parallel()
		if runtime.GOOS != "linux" {
		t.Skip("JetBrains tracking is only supported on Linux")
		}

		ctx := testutil.Context(t, testutil.WaitLong)

		rl, err := net.Listen("tcp", "127.0.0.1:0")
		// JetBrains tracking works by looking at the process name listening on the
		// forwarded port. If the process's command line includes the magic string
		// we are looking for, then we assume it is a JetBrains editor. So when we
		// connect to the port we must ensure the process includes that magic string
		// to fool the agent into thinking this is JetBrains. To do this we need to
		// spawn an external process (in this case a simple echo server) so we can
		// control the process name. The -D here is just to mimic how Java options
		// are set but is not necessary as the agent looks only for the magic
		// string itself anywhere in the command.
		_, b, _, ok := runtime.Caller(0)
		require.True(t, ok)
		dir := filepath.Join(filepath.Dir(b), "../scripts/echoserver/main.go")
		echoServerCmd := exec.Command("go", "run", dir,
		"-D", agentssh.MagicProcessCmdlineJetBrains)
		stdout, err := echoServerCmd.StdoutPipe()
		require.NoError(t, err)
		defer rl.Close()
		tcpAddr, valid := rl.Addr().(*net.TCPAddr)
		require.True(t, valid)
		remotePort := tcpAddr.Port
		go echoOnce(t, rl)
		err = echoServerCmd.Start()
		require.NoError(t, err)
		defer echoServerCmd.Process.Kill()

		// The echo server prints its port as the first line.
		sc := bufio.NewScanner(stdout)
		sc.Scan()
		remotePort := sc.Text()

		//nolint:dogsled
		conn, _, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
		sshClient, err := conn.SSHClient(ctx)
		require.NoError(t, err)

		tunneledConn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", remotePort))
		tunneledConn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%s", remotePort))
		require.NoError(t, err)
		t.Cleanup(func() {
		// always close on failure of test
Expand All		@@ -239,9 +249,10 @@ func TestAgent_Stats_Magic(t *testing.T) {
		"never saw stats with conn open: %+v", s,
		)

		//Manually closingthe connection
		//Killtheserver andconnection after checking for the echo.
		requireEcho(t, tunneledConn)
		_ = rl.Close()
		_ = echoServerCmd.Process.Kill()
		_ = tunneledConn.Close()

		require.Eventuallyf(t, func() bool {
		var ok bool
Expand Down
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,50 @@
		package main

		// A simple echo server. It listens on a random port, prints that port, then
		// echos back anything sent to it.

		import (
		"errors"
		"fmt"
		"io"
		"log"
		"net"
		)

		func main() {
		l, err := net.Listen("tcp", "127.0.0.1:0")
		if err != nil {
		log.Fatalf("listen error: err=%s", err)
		}

		defer l.Close()
		tcpAddr, valid := l.Addr().(*net.TCPAddr)
		if !valid {
		log.Fatal("address is not valid")
		}

		remotePort := tcpAddr.Port
		_, err = fmt.Println(remotePort)
		if err != nil {
		log.Fatalf("print error: err=%s", err)
		}

		for {
		conn, err := l.Accept()
		if err != nil {
		log.Fatalf("accept error, err=%s", err)
		return
		}

		go func() {
		defer conn.Close()
		_, err := io.Copy(conn, conn)

		if errors.Is(err, io.EOF) {
		return
		} else if err != nil {
		log.Fatalf("copy error, err=%s", err)
		}
		}()
		}
		}