feat: add port scanning to agent

coder · code-asher · May 26, 2022 · May 26, 2022 · May 26, 2022 · May 27, 2022
commit e9a28be956641b468d8f497b0fb29db8983baf79
diff --git a/agent/agent.go b/agent/agent.go
@@ -21,6 +21,7 @@ import (
 	"time"
 
 	"github.com/armon/circbuf"
+	"github.com/cakturk/go-netstat/netstat"
 	"github.com/gliderlabs/ssh"
 	"github.com/google/uuid"
 	"github.com/pkg/sftp"
@@ -37,13 +38,15 @@ import (
 )
 
 const (
+	ProtocolNetstat         = "netstat"
 	ProtocolReconnectingPTY = "reconnecting-pty"
 	ProtocolSSH             = "ssh"
 	ProtocolDial            = "dial"
 )
 
 type Options struct {
 	ReconnectingPTYTimeout time.Duration
+	NetstatInterval        time.Duration
 	EnvironmentVariables   map[string]string
 	Logger                 slog.Logger
 }
@@ -65,10 +68,14 @@ func New(dialer Dialer, options *Options) io.Closer {
 	if options.ReconnectingPTYTimeout == 0 {
 		options.ReconnectingPTYTimeout = 5 * time.Minute
 	}
+	if options.NetstatInterval == 0 {
+		options.NetstatInterval = 5 * time.Second
+	}
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	server := &agent{
 		dialer:                 dialer,
 		reconnectingPTYTimeout: options.ReconnectingPTYTimeout,
+		netstatInterval:        options.NetstatInterval,
 		logger:                 options.Logger,
 		closeCancel:            cancelFunc,
 		closed:                 make(chan struct{}),
@@ -85,6 +92,8 @@ type agent struct {
 	reconnectingPTYs       sync.Map
 	reconnectingPTYTimeout time.Duration
 
+	netstatInterval time.Duration
+
 	connCloseWait sync.WaitGroup
    	closeCancel   context.CancelFunc
 	closeMutex    sync.Mutex
@@ -225,6 +234,8 @@ func (a *agent) handlePeerConn(ctx context.Context, conn *peer.Conn) {
 			go a.handleReconnectingPTY(ctx, channel.Label(), channel.NetConn())
 		case ProtocolDial:
 			go a.handleDial(ctx, channel.Label(), channel.NetConn())
+		case ProtocolNetstat:
+			go a.handleNetstat(ctx, channel.Label(), channel.NetConn())
 		default:
 			a.logger.Warn(ctx, "unhandled protocol from channel",
 				slog.F("protocol", channel.Protocol()),
@@ -359,12 +370,10 @@ func (a *agent) createCommand(ctx context.Context, rawCommand string, env []stri
 	if err != nil {
 		return nil, xerrors.Errorf("getting os executable: %w", err)
 	}
-	cmd.Env = append(cmd.Env, fmt.Sprintf("USER=%s", username))
-	cmd.Env = append(cmd.Env, fmt.Sprintf(`PATH=%s%c%s`, os.Getenv("PATH"), filepath.ListSeparator, filepath.Dir(executablePath)))
 	// Git on Windows resolves with UNIX-style paths.
 	// If using backslashes, it's unable to find the executable.
-	unixExecutablePath := strings.ReplaceAll(executablePath, "\\", "/")
-	cmd.Env = append(cmd.Env, fmt.Sprintf(`GIT_SSH_COMMAND=%s gitssh --`, unixExecutablePath))
+	executablePath = strings.ReplaceAll(executablePath, "\\", "/")
+	cmd.Env = append(cmd.Env, fmt.Sprintf(`GIT_SSH_COMMAND=%s gitssh --`, executablePath))
 	// These prevent the user from having to specify _anything_ to successfully commit.
 	// Both author and committer must be set!
 	cmd.Env = append(cmd.Env, fmt.Sprintf(`GIT_AUTHOR_EMAIL=%s`, metadata.OwnerEmail))
@@ -707,6 +716,87 @@ func (a *agent) handleDial(ctx context.Context, label string, conn net.Conn) {
 	Bicopy(ctx, conn, nconn)
 }
 
+type NetstatPort struct {
+	Name string `json:"name"`
+	Port uint16 `json:"port"`
+}
+
+type NetstatResponse struct {
+	Ports []NetstatPort `json:"ports"`
+	Error string        `json:"error,omitempty"`
+	Took  time.Duration `json:"took"`
+}
+
+func (a *agent) handleNetstat(ctx context.Context, label string, conn net.Conn) {
+	write := func(resp NetstatResponse) error {
+		b, err := json.Marshal(resp)
+		if err != nil {
+			a.logger.Warn(ctx, "write netstat response", slog.F("label", label), slog.Error(err))
+			return xerrors.Errorf("marshal agent netstat response: %w", err)
+		}
+		_, err = conn.Write(b)
+		if err != nil {
+			a.logger.Warn(ctx, "write netstat response", slog.F("label", label), slog.Error(err))
+		}
+		return err
+	}
+
+	scan := func() ([]NetstatPort, error) {
+		if runtime.GOOS != "linux" && runtime.GOOS != "windows" {
+			return nil, xerrors.New(fmt.Sprintf("Port scanning is not supported on %s", runtime.GOOS))
+		}
+
+		tabs, err := netstat.TCPSocks(func(s *netstat.SockTabEntry) bool {
+			return s.State == netstat.Listen
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		ports := []NetstatPort{}
+		for _, tab := range tabs {
+			ports = append(ports, NetstatPort{
+				Name: tab.Process.Name,
+				Port: tab.LocalAddr.Port,
+			})
+		}
+		return ports, nil
+	}
+
+	scanAndWrite := func() {
+		start := time.Now()
+		ports, err := scan()
+		response := NetstatResponse{
+			Ports: ports,
+			Took:  time.Since(start),
+		}
+		if err != nil {
+			response.Error = err.Error()
+		}
+		_ = write(response)
+	}
+
+	scanAndWrite()
+
+	// Using a timer instead of a ticker to ensure delay between calls otherwise
+	// if nestat took longer than the interval we would constantly run it.
+	timer := time.NewTimer(a.netstatInterval)
+	go func() {
+		defer conn.Close()
+		defer timer.Stop()
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-timer.C:
+				scanAndWrite()
+				timer.Reset(a.netstatInterval)
+			}
+		}
+	}()
+}
+
 // isClosed returns whether the API is closed or not.
 func (a *agent) isClosed() bool {
 	select {

diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -373,6 +373,57 @@ func TestAgent(t *testing.T) {
 		require.ErrorContains(t, err, "no such file")
 		require.Nil(t, netConn)
 	})
+
+	t.Run("Netstat", func(t *testing.T) {
+		t.Parallel()
+
+		var ports []agent.NetstatPort
+		listen := func() {
+			listener, err := net.Listen("tcp", "127.0.0.1:0")
+			require.NoError(t, err)
+			t.Cleanup(func() {
+				_ = listener.Close()
+			})
+
+			tcpAddr, valid := listener.Addr().(*net.TCPAddr)
+			require.True(t, valid)
+
+			name, err := os.Executable()
+			require.NoError(t, err)
+
+			ports = append(ports, agent.NetstatPort{
+				Name: filepath.Base(name),
+				Port: uint16(tcpAddr.Port),
+			})
+		}
+
+		conn := setupAgent(t, agent.Metadata{}, 0)
+		netConn, err := conn.Netstat(context.Background())
+		require.NoError(t, err)
+		t.Cleanup(func() {
+			_ = netConn.Close()
+		})
+
+		decoder := json.NewDecoder(netConn)
+
+		expectNetstat := func() {
+			var res agent.NetstatResponse
+			err = decoder.Decode(&res)
+			require.NoError(t, err)
+
+			if runtime.GOOS == "linux" || runtime.GOOS == "windows" {
+				require.Subset(t, res.Ports, ports)
+			} else {
+				require.Equal(t, fmt.Sprintf("Port scanning is not supported on %s", runtime.GOOS), res.Error)
+			}
+		}
+
+		listen()
+		expectNetstat()
+
+		listen()
+		expectNetstat()
+	})
 }
 
 func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exec.Cmd {
@@ -420,6 +471,7 @@ func setupAgent(t *testing.T, metadata agent.Metadata, ptyTimeout time.Duration)
 	}, &agent.Options{
 		Logger:                 slogtest.Make(t, nil).Leveled(slog.LevelDebug),
 		ReconnectingPTYTimeout: ptyTimeout,
+		NetstatInterval:        100 * time.Millisecond,
 	})
 	t.Cleanup(func() {
 		_ = client.Close()

diff --git a/agent/conn.go b/agent/conn.go
@@ -112,6 +112,17 @@ func (c *Conn) DialContext(ctx context.Context, network string, addr string) (ne
 	return channel.NetConn(), nil
 }
 
+// Netstat returns a connection that serves a list of listening ports.
+func (c *Conn) Netstat(ctx context.Context) (net.Conn, error) {
+	channel, err := c.CreateChannel(ctx, "netstat", &peer.ChannelOptions{
+		Protocol: ProtocolNetstat,
+	})
+	if err != nil {
+		return nil, xerrors.Errorf("netsat: %w", err)
+	}
+	return channel.NetConn(), nil
+}
+
 func (c *Conn) Close() error {
 	_ = c.Negotiator.DRPCConn().Close()
 	return c.Conn.Close()

diff --git a/go.mod b/go.mod
@@ -126,6 +126,8 @@ require (
 	storj.io/drpc v0.0.30
 )
 
+require github.com/cakturk/go-netstat v0.0.0-20200220111822-e5b49efee7a5
+
 require (
 	github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 // indirect
 	github.com/Azure/go-ntlmssp v0.0.0-20200615164410-66371956d46c // indirect

diff --git a/go.sum b/go.sum
@@ -240,6 +240,8 @@ github.com/bugsnag/osext v0.0.0-20130617224835-0dd3f918b21b/go.mod h1:obH5gd0Bsq
 github.com/bugsnag/panicwrap v0.0.0-20151223152923-e2c28503fcd0/go.mod h1:D/8v3kj0zr8ZAKg1AQ6crr+5VwKN5eIywRkfhyM/+dE=
 github.com/bytecodealliance/wasmtime-go v0.35.0 h1:VZjaZ0XOY0qp9TQfh0CQj9zl/AbdeXePVTALy8V1sKs=
    github.com/bytecodealliance/wasmtime-go v0.35.0/go.mod h1:q320gUxqyI8yB+ZqRuaJOEnGkAnHh6WtJjMaT2CW4wI=
+github.com/cakturk/go-netstat v0.0.0-20200220111822-e5b49efee7a5 h1:BjkPE3785EwPhhyuFkbINB+2a1xATwk8SNDWnJiD41g=
+github.com/cakturk/go-netstat v0.0.0-20200220111822-e5b49efee7a5/go.mod h1:jtAfVaU/2cu1+wdSRPWE2c1N2qeAA3K4RH9pYgqwets=
 github.com/cenkalti/backoff/v4 v4.1.1/go.mod h1:scbssz8iZGpm3xbr14ovlUdkxfGXNInqkPWOWmG2CLw=
 github.com/cenkalti/backoff/v4 v4.1.2/go.mod h1:scbssz8iZGpm3xbr14ovlUdkxfGXNInqkPWOWmG2CLw=
 github.com/cenkalti/backoff/v4 v4.1.3 h1:cFAlzYUlVYDysBEH2T5hyJZMh3+5+WCBvSnK6Q8UtC4=