netstack: swap may create a new fdbased endpoint

Author: ignoramous

intra/netstack/fdbased.go

@@ -30,6 +30,7 @@
 package netstack
 
 import (
+	"errors"
 	"fmt"
 	"sync/atomic"
 	"time"
@@ -59,6 +60,8 @@ const errorOnInvalidFD = false
 // wrapttl is the time to wait for the dispatcher to wrap up (close a previous FD).
 const waitttl = wrapttl
 
+var errNeedsNewEndpoint = errors.New("ns: needs new endpoint")
+
 type FdSwapper interface {
 	// Swap closes existing FDs; uses new fd.
 	Swap(fd int) error
@@ -200,7 +203,7 @@ type Options struct {
 // Makes fd non-blocking, but does not take ownership of fd, which must remain
 // open for the lifetime of the returned endpoint (until after the endpoint has
 // stopped being using and Wait returns).
-func NewFdbasedInjectableEndpoint(opts *Options) (SeamlessEndpoint, error) {
+func newFdbasedInjectableEndpoint(opts *Options) (SeamlessEndpoint, error) {
 	caps := stack.LinkEndpointCapabilities(0)
 	if opts.RXChecksumOffload {
 		caps |= stack.CapabilityRXChecksumOffload
@@ -322,13 +325,18 @@ func (e *endpoint) Swap(fd int) (err error) {
 	e.Lock()
 	defer e.Unlock()
 
+	prevfd := e.fds.Load()
+	if !prevfd.ok() {
+		return errNeedsNewEndpoint
+	}
+
 	f, err := newTun(fd) // fd may be invalid (ex: -1)
 	if err != nil || f == nil {
 		f = invalidFds // nilaway
 		err = log.EE("ns: tun: swap: (%d) err: %v / %v; using invalidfd", fd, err)
 	}
 
-	prevfd := e.fds.Swap(f) // commence WritePackets() on fd
+	e.fds.Store(f) // commence WritePackets() on fd
 
 	log.D("ns: tun: swap: fd %s => %d; err? %v", prevfd, fd, err)
 

intra/netstack/netstack.go

@@ -11,7 +11,9 @@ import (
 	"fmt"
 	"io"
 	"net/netip"
+	"strconv"
 	"strings"
+	"sync"
 	"syscall"
 
 	"github.com/celzero/firestack/intra/core"
@@ -35,11 +37,39 @@ const nicfwd = false
 // packets will be truncated to snapLen.
 const SnapLen uint32 = 2048 // in bytes; some sufficient value
 
+var (
+	errNoFdSwapper = errors.New("linkFdSwap: no FdSwapper")
+)
+
 type linkFdSwap struct {
+	sync.Mutex
 	stack.LinkEndpoint
 	FdSwapper
 }
 
+// Swap implements FdSwapper.
+func (l *linkFdSwap) Swap(fd int) error {
+	l.Lock()
+	defer l.Unlock()
+
+	if l.FdSwapper == nil {
+		return errNoFdSwapper
+	}
+
+	err := l.FdSwapper.Swap(fd)
+	if errors.Is(err, errNeedsNewEndpoint) {
+		umtu := uint32(l.MTU())
+		opt := Options{
+			FDs: []int{fd},
+			MTU: umtu,
+		}
+		core.Go("linkFdSwap."+strconv.Itoa(fd), l.LinkEndpoint.Close)
+		l.LinkEndpoint, err = newFdbasedInjectableEndpoint(&opt)
+	}
+
+	return err
+}
+
 // ref: github.com/google/gvisor/blob/91f58d2cc/pkg/tcpip/sample/tun_tcp_echo/main.go#L102
 func NewEndpoint(dev, mtu int, sink io.WriteCloser) (ep SeamlessEndpoint, err error) {
 	defer func() {
@@ -55,7 +85,7 @@ func NewEndpoint(dev, mtu int, sink io.WriteCloser) (ep SeamlessEndpoint, err er
 		MTU: umtu,
 	}
 
-	if ep, err = NewFdbasedInjectableEndpoint(&opt); err != nil {
+	if ep, err = newFdbasedInjectableEndpoint(&opt); err != nil {
 		return nil, err
 	}
 	// ref: github.com/google/gvisor/blob/aeabb785278/pkg/tcpip/link/sniffer/sniffer.go#L111-L131
@@ -70,7 +100,7 @@ func snoop(ep SeamlessEndpoint, sink io.WriteCloser) (SeamlessEndpoint, error) {
 	if link, err := NewSnoopyEndpoint(ep, sink); err != nil {
 		return nil, err
 	} else {
-		return linkFdSwap{link, ep}, nil
+		return &linkFdSwap{sync.Mutex{}, link, ep}, nil
 	}
 }