cmd/tsconnect/wasm/wasm_js.go

@@ -66,19 +66,20 @@ import (
 var ControlURL = ipn.DefaultControlURL
 
 func main() {
+	shutdownCh := make(chan struct{})
 	js.Global().Set("newIPN", js.FuncOf(func(this js.Value, args []js.Value) any {
 		if len(args) != 1 {
 			log.Fatal("Usage: newIPN(config)")
 			return nil
 		}
-		return newIPN(args[0])
+		return newIPN(args[0], shutdownCh)
 	}))
-	// Keep Go runtime alive, otherwise it will be shut down before newIPN gets
-	// called.
-	<-make(chan bool)
+	// Block until shutdown() is called on the IPN, then let main return so the
+	// Go runtime (and all its goroutines) can be collected by the JS engine.
+	<-shutdownCh
 }
 
-func newIPN(jsConfig js.Value) map[string]any {
+func newIPN(jsConfig js.Value, shutdownCh chan struct{}) map[string]any {
 	netns.SetEnabled(false)
 
 	var store ipn.StateStore
@@ -179,6 +180,7 @@ func newIPN(jsConfig js.Value) map[string]any {
 		hostname:    hostname,
 		logID:       logid,
 		funnelPorts: make(map[uint16]*funnelListenerEntry),
+		shutdownCh:  shutdownCh,
 	}
 	lb.SetTCPHandlerForFunnelFlow(jsIPN.handleFunnelTCP)
 
@@ -361,6 +363,9 @@ func newIPN(jsConfig js.Value) map[string]any {
 		"suggestExitNode": js.FuncOf(func(this js.Value, args []js.Value) any {
 			return jsIPN.suggestExitNode()
 		}),
+		"shutdown": js.FuncOf(func(this js.Value, args []js.Value) any {
+			return jsIPN.shutdown()
+		}),
 		"localAPI": js.FuncOf(func(this js.Value, args []js.Value) any {
 			if len(args) < 2 {
 				log.Printf("Usage: localAPI(method, path[, body])")
@@ -387,6 +392,12 @@ type jsIPN struct {
 
 	funnelMu    sync.Mutex
 	funnelPorts map[uint16]*funnelListenerEntry
+
+	// ln is the safesocket listener created by run(); stored here so shutdown
+	// can close it and unblock srv.Run.
+	ln           net.Listener
+	shutdownCh   chan struct{} // closed by shutdown() to unblock main()
+	shutdownOnce sync.Once
 }
 
 // funnelListenerEntry is the per-port state for routing Funnel connections to a listenTLS listener.
@@ -594,14 +605,17 @@ func (i *jsIPN) run(jsCallbacks js.Value) {
 		}
 	}()
 
-	go func() {
-		ln, err := safesocket.Listen("")
-		if err != nil {
-			log.Fatalf("safesocket.Listen: %v", err)
-		}
+	ln, err := safesocket.Listen("")
+	if err != nil {
+		log.Fatalf("safesocket.Listen: %v", err)
+	}
+	i.ln = ln
 
-		err = i.srv.Run(context.Background(), ln)
-		log.Fatalf("ipnserver.Run exited: %v", err)
+	go func() {
+		err := i.srv.Run(context.Background(), ln)
+		if err != nil && !errors.Is(err, net.ErrClosed) {
+			log.Fatalf("ipnserver.Run exited: %v", err)
+		}
 	}()
 }
 
@@ -620,6 +634,17 @@ func (i *jsIPN) logout() {
 	}()
 }
 
+func (i *jsIPN) shutdown() js.Value {
+	return makePromise(func() (any, error) {
+		i.shutdownOnce.Do(func() {
+			i.lb.Shutdown()
+			i.ln.Close()
+			close(i.shutdownCh)
+		})
+		return nil, nil
+	})
+}
+
 func (i *jsIPN) ssh(host, username string, termConfig js.Value) map[string]any {
 	jsSSHSession := &jsSSHSession{
 		jsIPN:      i,