cmd/vnet: add wsproxy mode

For hooking up websocket VM clients to natlab.

Updates #13038

Change-Id: Iaf728b9146042f3d0c2d3a5e25f178646dd10951
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>

3 files changed, 184 insertions, 0 deletions

diff --git a/cmd/vnet/vnet-main.go b/cmd/vnet/vnet-main.go
index 1eb4f65ef..9dd4d8cfa 100644
--- a/cmd/vnet/vnet-main.go
+++ b/cmd/vnet/vnet-main.go
@@ -7,15 +7,21 @@ package main
 
 import (
 	"context"
+	"encoding/binary"
 	"flag"
+	"fmt"
+	"io"
 	"log"
 	"net"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
 	"os"
+	"path/filepath"
+	"slices"
 	"time"
 
+	"github.com/coder/websocket"
 	"tailscale.com/tstest/natlab/vnet"
 	"tailscale.com/types/logger"
 	"tailscale.com/util/must"
@@ -31,10 +37,18 @@ var (
 	pcapFile = flag.String("pcap", "", "if non-empty, filename to write pcap")
 	v4       = flag.Bool("v4", true, "enable IPv4")
 	v6       = flag.Bool("v6", true, "enable IPv6")
+
+	wsproxyListen = flag.String("wsproxy", "", "if non-empty, TCP address to run websocket server on. See https://github.com/copy/v86/blob/master/docs/networking.md#backend-url-schemes")
 )
 
 func main() {
 	flag.Parse()
+	if *wsproxyListen != "" {
+		if err := runWSProxy(); err != nil {
+			log.Fatalf("runWSProxy: %v", err)
+		}
+		return
+	}
 
 	if _, err := os.Stat(*listen); err == nil {
 		os.Remove(*listen)
@@ -137,3 +151,168 @@ func main() {
 		go s.ServeUnixConn(c.(*net.UnixConn), vnet.ProtocolQEMU)
 	}
 }
+
+func runWSProxy() error {
+	ln, err := net.Listen("tcp", *wsproxyListen)
+	if err != nil {
+		return err
+	}
+	defer ln.Close()
+
+	log.Printf("Running wsproxy mode on %v ...", *wsproxyListen)
+
+	var hs http.Server
+	hs.Handler = http.HandlerFunc(handleWebSocket)
+
+	return hs.Serve(ln)
+}
+
+func handleWebSocket(w http.ResponseWriter, r *http.Request) {
+	conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
+		InsecureSkipVerify: true,
+	})
+	if err != nil {
+		log.Printf("Upgrade error: %v", err)
+		return
+	}
+	defer conn.Close(websocket.StatusInternalError, "closing")
+	log.Printf("WebSocket client connected: %s", r.RemoteAddr)
+
+	ctx, cancel := context.WithCancel(r.Context())
+	defer cancel()
+
+	messageType, firstData, err := conn.Read(ctx)
+	if err != nil {
+		log.Printf("ReadMessage first: %v", err)
+		return
+	}
+	if messageType != websocket.MessageBinary {
+		log.Printf("Ignoring non-binary message")
+		return
+	}
+	if len(firstData) < 12 {
+		log.Printf("Ignoring short message")
+		return
+	}
+	clientMAC := vnet.MAC(firstData[6:12])
+
+	// Set up a qemu-protocol Unix socket pair. We'll fake the qemu protocol here
+	// to avoid changing the vnet package.
+	td, err := os.MkdirTemp("", "vnet")
+	if err != nil {
+		panic(fmt.Errorf("MkdirTemp: %v", err))
+	}
+	defer os.RemoveAll(td)
+
+	unixSrv := filepath.Join(td, "vnet.sock")
+
+	srv, err := net.Listen("unix", unixSrv)
+	if err != nil {
+		panic(fmt.Errorf("Listen: %v", err))
+	}
+	defer srv.Close()
+
+	var c vnet.Config
+	c.SetBlendReality(true)
+
+	var net1opt = []any{vnet.NAT("easy")}
+	net1opt = append(net1opt, "2.1.1.1", "192.168.1.1/24")
+	net1opt = append(net1opt, "2000:52::1/64")
+
+	c.AddNode(c.AddNetwork(net1opt...), clientMAC)
+
+	vs, err := vnet.New(&c)
+	if err != nil {
+		panic(fmt.Errorf("newServer: %v", err))
+	}
+	if err := vs.PopulateDERPMapIPs(); err != nil {
+		log.Printf("warning: ignoring failure to populate DERP map: %v", err)
+		return
+	}
+
+	errc := make(chan error, 1)
+	fail := func(err error) {
+		select {
+		case errc <- err:
+			log.Printf("failed: %v", err)
+		case <-ctx.Done():
+		}
+	}
+
+	go func() {
+		c, err := srv.Accept()
+		if err != nil {
+			fail(err)
+			return
+		}
+		vs.ServeUnixConn(c.(*net.UnixConn), vnet.ProtocolQEMU)
+	}()
+
+	uc, err := net.Dial("unix", unixSrv)
+	if err != nil {
+		panic(fmt.Errorf("Dial: %v", err))
+	}
+	defer uc.Close()
+
+	var frameBuf []byte
+	writeDataToUnixConn := func(data []byte) error {
+		frameBuf = slices.Grow(frameBuf[:0], len(data)+4)[:len(data)+4]
+		binary.BigEndian.PutUint32(frameBuf[:4], uint32(len(data)))
+		copy(frameBuf[4:], data)
+
+		_, err = uc.Write(frameBuf)
+		return err
+	}
+	if err := writeDataToUnixConn(firstData); err != nil {
+		fail(err)
+		return
+	}
+
+	go func() {
+		for {
+			messageType, data, err := conn.Read(ctx)
+			if err != nil {
+				fail(fmt.Errorf("ReadMessage: %v", err))
+				break
+			}
+
+			if messageType != websocket.MessageBinary {
+				log.Printf("Ignoring non-binary message")
+				continue
+			}
+
+			if err := writeDataToUnixConn(data); err != nil {
+				fail(err)
+				return
+			}
+		}
+	}()
+
+	go func() {
+		const maxBuf = 4096
+		frameBuf := make([]byte, maxBuf)
+		for {
+			_, err := io.ReadFull(uc, frameBuf[:4])
+			if err != nil {
+				fail(err)
+				return
+			}
+			frameLen := binary.BigEndian.Uint32(frameBuf[:4])
+			if frameLen > maxBuf {
+				fail(fmt.Errorf("frame too large: %d", frameLen))
+				return
+			}
+			if _, err := io.ReadFull(uc, frameBuf[:frameLen]); err != nil {
+				fail(err)
+				return
+			}
+
+			if err := conn.Write(ctx, websocket.MessageBinary, frameBuf[:frameLen]); err != nil {
+				fail(err)
+				return
+			}
+		}
+	}()
+
+	<-ctx.Done()
+}
diff --git a/tstest/natlab/vnet/conf.go b/tstest/natlab/vnet/conf.go
index a37c22a6c..07b181540 100644
--- a/tstest/natlab/vnet/conf.go
+++ b/tstest/natlab/vnet/conf.go
@@ -121,6 +121,8 @@ func (c *Config) AddNode(opts ...any) *Node {
 					n.err = fmt.Errorf("unknown NodeOption %q", o)
 				}
 			}
+		case MAC:
+			n.mac = o
 		default:
 			if n.err == nil {
 				n.err = fmt.Errorf("unknown AddNode option type %T", o)
diff --git a/tstest/natlab/vnet/vnet.go b/tstest/natlab/vnet/vnet.go
index ead2bbb8b..e3ecf0f75 100644
--- a/tstest/natlab/vnet/vnet.go
+++ b/tstest/natlab/vnet/vnet.go
@@ -88,6 +88,9 @@ func (s *Server) PopulateDERPMapIPs() error {
 			if n.IPv4 != "" {
 				s.derpIPs.Add(netip.MustParseAddr(n.IPv4))
 			}
+			if n.IPv6 != "" {
+				s.derpIPs.Add(netip.MustParseAddr(n.IPv6))
+			}
 		}
 	}
 	return nil