cmd/tswrap,portlist: switch tswrap to portlist, add pid to portlistdanderson/tsburrito

Signed-off-by: James Tucker <james@tailscale.com>

6 files changed, 38 insertions, 37 deletions

diff --git a/cmd/tswrap/main.go b/cmd/tswrap/main.go
index 559d56b03..e6607ac0b 100644
--- a/cmd/tswrap/main.go
+++ b/cmd/tswrap/main.go
@@ -2,15 +2,11 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
-//go:build linux
-
 // The tswrap binary runs a child process and makes it accessible over
 // Tailscale.
 package main
 
 import (
-	"bufio"
-	"bytes"
 	"context"
 	"errors"
 	"flag"
@@ -18,19 +14,18 @@ import (
 	"io"
 	"log"
 	"net"
-	"net/netip"
 	"os"
 	"os/exec"
 	"os/signal"
 	"sort"
 	"strconv"
-	"strings"
+	"syscall"
 	"time"
 
-	"golang.org/x/sys/unix"
 	"tailscale.com/client/tailscale"
 	"tailscale.com/ipn/ipnstate"
 	"tailscale.com/ipn/store/mem"
+	"tailscale.com/portlist"
 	"tailscale.com/syncs"
 	"tailscale.com/tsnet"
 	"tailscale.com/types/logger"
@@ -43,7 +38,7 @@ var (
 
 func main() {
 	sigch := make(chan os.Signal, 1)
-	signal.Notify(sigch, os.Interrupt, unix.SIGTERM)
+	signal.Notify(sigch, os.Interrupt, syscall.SIGTERM)
 
 	flag.Parse()
 
@@ -188,7 +183,7 @@ func (p *proxy) Stop() {
 
 func (p *proxy) Wait() {
 	<-p.shutdownCtx.Done()
-	p.cmd.Process.Signal(unix.SIGTERM)
+	p.cmd.Process.Signal(syscall.SIGTERM)
 	p.ln.Close()
 	if p.srv.Ephemeral {
 		p.client.Logout(context.Background())
@@ -290,37 +285,30 @@ func portsOfCmd(cmd *exec.Cmd) (ports []int, err error) {
 		return nil, errors.New("no process")
 	}
 	pid := cmd.Process.Pid
-	wantSub := fmt.Sprintf(" %d/", pid)
 
-	ns := exec.Command("netstat", "-p", "--inet", "-l", "-n")
-	out, err := ns.Output()
+	poller, err := portlist.NewPoller()
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("creating port poller: %w", err)
 	}
-	bs := bufio.NewScanner(bytes.NewReader(out))
-	for bs.Scan() {
-		line := bs.Text()
-		if !strings.HasPrefix(line, "tcp") ||
-			!strings.Contains(line, "LISTEN") ||
-			!strings.Contains(line, wantSub) {
-			continue
+	defer poller.Close()
+	// TODO(raggi): timeout?
+	go poller.Run(context.Background())
+
+	c := poller.Updates()
+	for list := range c {
+		for _, p := range list {
+			if p.Pid == pid {
+				ports = append(ports, int(p.Port))
+			}
 		}
-		f := strings.Fields(line)
-		if len(f) < 4 {
-			continue
+		if len(ports) > 0 {
+			break
 		}
-		ipp, err := netip.ParseAddrPort(f[3])
-		if err == nil {
-			ports = append(ports, int(ipp.Port()))
-			continue
-		}
-	}
-	if err := bs.Err(); err != nil {
-		return nil, err
 	}
 	if len(ports) == 0 {
 		return nil, errors.New("no listening ports found")
 	}
+
 	sort.Ints(ports)
 	return ports, nil
 }
diff --git a/portlist/netstat_test.go b/portlist/netstat_test.go
index 29a74b1c6..5d99b8110 100644
--- a/portlist/netstat_test.go
+++ b/portlist/netstat_test.go
@@ -54,12 +54,12 @@ udp46      0      0  *.146                 *.*
 
 func TestParsePortsNetstat(t *testing.T) {
 	want := List{
-		Port{"tcp", 23, ""},
-		Port{"tcp", 24, ""},
-		Port{"udp", 104, ""},
-		Port{"udp", 106, ""},
-		Port{"udp", 146, ""},
-		Port{"tcp", 8185, ""}, // but not 8186, 8187, 8188 on localhost
+		Port{"tcp", 23, "", 0},
+		Port{"tcp", 24, "", 0},
+		Port{"udp", 104, "", 0},
+		Port{"udp", 106, "", 0},
+		Port{"udp", 146, "", 0},
+		Port{"tcp", 8185, "", 0}, // but not 8186, 8187, 8188 on localhost
 	}
 
 	pl, err := appendParsePortsNetstat(nil, bufio.NewReader(strings.NewReader(netstatOutput)))
diff --git a/portlist/portlist.go b/portlist/portlist.go
index fe4a1af4d..871adbe99 100644
--- a/portlist/portlist.go
+++ b/portlist/portlist.go
@@ -19,6 +19,7 @@ type Port struct {
 	Proto   string // "tcp" or "udp"
 	Port    uint16 // port number
 	Process string // optional process name, if found
+	Pid     int    // process id, if known
 }
 
 // List is a list of Ports.
diff --git a/portlist/portlist_linux.go b/portlist/portlist_linux.go
index 8d9176cd9..5b1813b16 100644
--- a/portlist/portlist_linux.go
+++ b/portlist/portlist_linux.go
@@ -40,6 +40,7 @@ type linuxImpl struct {
 
 type portMeta struct {
 	port          Port
+	pid           int
 	keep          bool
 	needsProcName bool
 }
@@ -326,6 +327,9 @@ func (li *linuxImpl) findProcessNames(need map[string]*portMeta) error {
 					}
 
 					argv := strings.Split(strings.TrimSuffix(string(bs), "\x00"), "\x00")
+					if p, err := strconv.Atoi(pid); err == nil {
+						pe.pid = p
+					}
 					pe.port.Process = argvSubject(argv...)
 					pe.needsProcName = false
 					delete(need, string(targetBuf[:n]))
diff --git a/portlist/portlist_macos.go b/portlist/portlist_macos.go
index a08678c87..b98c5f37d 100644
--- a/portlist/portlist_macos.go
+++ b/portlist/portlist_macos.go
@@ -12,6 +12,7 @@ import (
 	"fmt"
 	"log"
 	"os/exec"
+	"strconv"
 	"strings"
 	"sync/atomic"
 	"time"
@@ -162,6 +163,7 @@ func (im *macOSImpl) addProcesses() error {
 	im.br.Reset(outPipe)
 
 	var cmd, proto string
+	var pid int
 	for {
 		line, err := im.br.ReadBytes('\n')
 		if err != nil {
@@ -176,6 +178,10 @@ func (im *macOSImpl) addProcesses() error {
 			// starting a new process
 			cmd = ""
 			proto = ""
+			pid = 0
+			if p, err := strconv.Atoi(string(val)); err == nil {
+				pid = p
+			}
 		case 'c':
 			cmd = string(val) // TODO(bradfitz): avoid garbage; cache process names between runs?
 		case 'P':
@@ -194,6 +200,7 @@ func (im *macOSImpl) addProcesses() error {
 			switch {
 			case m != nil:
 				m.port.Process = cmd
+				m.port.Pid = pid
 			default:
 				// ignore: processes and ports come and go
 			}
diff --git a/portlist/portlist_windows.go b/portlist/portlist_windows.go
index 811e4f0aa..e42ebe1d1 100644
--- a/portlist/portlist_windows.go
+++ b/portlist/portlist_windows.go
@@ -82,6 +82,7 @@ func (im *windowsImpl) AppendListeningPorts(base []Port) ([]Port, error) {
 				Proto:   "tcp",
 				Port:    e.Local.Port(),
 				Process: procNameOfPid(e.Pid),
+				Pid:     e.Pid,
 			},
 		}
 		im.known[fp] = pm