From eb6115e29535ecda48773eea04c87d2eb94d09e4 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Fri, 16 Apr 2021 13:46:13 -0700
Subject: [PATCH] cmd/tailscaled: let SOCKS5 dial non-Tailscale addrs in
 userspace mode

Fixes #1617

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 cmd/tailscaled/tailscaled.go | 47 +++++++++++++++++++-----------------
 1 file changed, 25 insertions(+), 22 deletions(-)

diff --git a/cmd/tailscaled/tailscaled.go b/cmd/tailscaled/tailscaled.go
index 983eaac1b..4e51bbbe5 100644
--- a/cmd/tailscaled/tailscaled.go
+++ b/cmd/tailscaled/tailscaled.go
@@ -29,10 +29,12 @@ import (
 	"time"
 
 	"github.com/go-multierror/multierror"
+	"inet.af/netaddr"
 	"tailscale.com/ipn/ipnserver"
 	"tailscale.com/logpolicy"
 	"tailscale.com/net/dns"
 	"tailscale.com/net/socks5"
+	"tailscale.com/net/tsaddr"
 	"tailscale.com/net/tstun"
 	"tailscale.com/paths"
 	"tailscale.com/types/flagtype"
@@ -239,30 +241,31 @@ func run() error {
 		srv := &socks5.Server{
 			Logf: logger.WithPrefix(logf, "socks5: "),
 		}
-		// TODO: also consider wrapNetstack, where dials can go to either Tailscale
-		// or non-Tailscale targets. But that's also basically what
-		// https://github.com/tailscale/tailscale/issues/1617 is about, so do them
-		// both at the same time.
-		if useNetstack {
-			srv.Dialer = func(ctx context.Context, network, addr string) (net.Conn, error) {
-				return ns.DialContextTCP(ctx, addr)
+		var (
+			mu  sync.Mutex // guards the following field
+			dns netstack.DNSMap
+		)
+		e.AddNetworkMapCallback(func(nm *netmap.NetworkMap) {
+			mu.Lock()
+			defer mu.Unlock()
+			dns = netstack.DNSMapFromNetworkMap(nm)
+		})
+		useNetstackForIP := func(ip netaddr.IP) bool {
+			// TODO(bradfitz): this isn't exactly right.
+			// We should also support subnets when the
+			// prefs are configured as such.
+			return tsaddr.IsTailscaleIP(ip)
+		}
+		srv.Dialer = func(ctx context.Context, network, addr string) (net.Conn, error) {
+			ipp, err := dns.Resolve(ctx, addr)
+			if err != nil {
+				return nil, err
 			}
-		} else {
-			var mu sync.Mutex
-			var dns netstack.DNSMap
-			e.AddNetworkMapCallback(func(nm *netmap.NetworkMap) {
-				mu.Lock()
-				defer mu.Unlock()
-				dns = netstack.DNSMapFromNetworkMap(nm)
-			})
-			srv.Dialer = func(ctx context.Context, network, addr string) (net.Conn, error) {
-				ipp, err := dns.Resolve(ctx, addr)
-				if err != nil {
-					return nil, err
-				}
-				var d net.Dialer
-				return d.DialContext(ctx, network, ipp.String())
+			if ns != nil && useNetstackForIP(ipp.IP) {
+				return ns.DialContextTCP(ctx, addr)
 			}
+			var d net.Dialer
+			return d.DialContext(ctx, network, ipp.String())
 		}
 		go func() {
 			log.Fatalf("SOCKS5 server exited: %v", srv.Serve(socksListener))