netstack: refactor swap, dispatch, attach

Author: ignoramous

intra/netstack/dispatchers.go

@@ -25,7 +25,6 @@ package netstack
 
 import (
 	"math/rand"
-	"net"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -187,14 +186,8 @@ var _ linkDispatcher = (*readVDispatcher)(nil)
 // newReadVDispatcher creates a new linkDispatcher that vector reads packets from
 // fd and dispatches them to endpoint e. It assumes ownership of fd but not of e.
 func newReadVDispatcher(fd int, e *endpoint) (linkDispatcher, error) {
-	tun, err := newTun(fd)
-	if err != nil {
-		clos(fd)
-		return nil, err
-	}
 	d := &readVDispatcher{
 		e:   e,
-		fds: core.NewVolatile(tun),
 		buf: newIovecBuffer(bufcfg),
 		mgr: newSupervisor(e, fd),
 	}
@@ -206,26 +199,10 @@ func newReadVDispatcher(fd int, e *endpoint) (linkDispatcher, error) {
 
 // swap atomically swaps existing fd for this new one.
 // On error, it closes fd.
-func (d *readVDispatcher) swap(fd int) error {
-	done := d.closed.Load()
-	if done {
-		clos(fd)
-		return net.ErrClosed
+func (d *readVDispatcher) prepare(f *fds) {
+	if !d.closed.Load() {
+		d.mgr.swap(f.tun()) // used for diagnostics only
 	}
-
-	note := log.I
-	f, err := newTun(fd) // fd may be invalid (ex: -1)
-	if err != nil {
-		defer clos(fd)
-		note = log.W
-	}
-
-	prev := d.fds.Swap(f)   // f may be nil
-	d.wrapup(prev, wrapttl) // prev may be nil
-	d.mgr.swap(f.tun())     // used for diagnostics only
-
-	note("ns: dispatch: swap: tun(%d => %d); err %v", prev.tun(), fd, err)
-	return err
 }
 
 // stop stops the dispatcher once. Safe to call multiple times.
@@ -234,7 +211,6 @@ func (d *readVDispatcher) stop() {
 
 	d.once.Do(func() {
 		d.closed.Store(true)
-		d.fds.Load().stop()
 		d.mgr.stop()
 		log.I("ns: dispatch: closed!")
 	})
@@ -244,8 +220,8 @@ const abort = false // abort indicates that the dispatcher should stop.
 const cont = true   // cont indicates that the dispatcher should continue delivering packets despite an error.
 
 // dispatch reads packets from the current file descriptor in d.fds and dispatches it to netstack.
-func (d *readVDispatcher) dispatch() (bool, tcpip.Error) {
-	return d.io(d.fds.Load())
+func (d *readVDispatcher) dispatch(fd *fds) (bool, tcpip.Error) {
+	return d.io(fd)
 }
 
 // wrapup reads packets from fds and dispatches it to netstack
@@ -255,6 +231,8 @@ func (d *readVDispatcher) wrapup(fds *fds, noMoreThan30s time.Duration) {
 	if !fds.ok() { // fds may be nil
 		return
 	}
+	defer fds.stop()
+
 	// Loopback is set to true when VPN is in lockdown mode (block connections
 	// without vpn). It is observed that by closing the previous tun after delay
 	// results in "connection was reset" errors in netstack's TCP handler, which
@@ -263,21 +241,11 @@ func (d *readVDispatcher) wrapup(fds *fds, noMoreThan30s time.Duration) {
 	// quite a while or the device switches to a new network (on Android 14+).
 	if !threadSafe || settings.Loopingback.Load() {
 		log.I("ns: tun(%d): wrapup: immediate (loopback)", fds.tun())
-		fds.stop()
 		return
 	}
 
-	noMoreThan30s = min(30*time.Second, noMoreThan30s)
-	secs := int64(noMoreThan30s.Seconds())
-
-	go func() {
-		<-time.After(noMoreThan30s)
-		log.W("ns: tun(%d): drain: timeout! %dsecs", fds.tun(), secs)
-		fds.stop()
-	}()
-
-	go func() {
-		log.I("ns: tun(%d): drain: start w timeout in %dsecs", fds.tun(), secs)
+	awaited := core.Await(func() {
+		log.I("ns: tun(%d): drain: start w timeout in %s", fds.tun(), core.FmtPeriod(noMoreThan30s))
 		for {
 			cont, err := d.io(fds)
 			if fd := fds.tun(); !cont {
@@ -287,7 +255,9 @@ func (d *readVDispatcher) wrapup(fds *fds, noMoreThan30s time.Duration) {
 				log.W("ns: tun(%d): drain: continue on err: %v", fd, err)
 			} // else: continue draining
 		}
-	}()
+	}, min(30*time.Second, noMoreThan30s))
+
+	logei(awaited)("ns: tun(%d): drain: timeout!", fds.tun())
 }
 
 // io reads packets from fds and dispatches it to netstack.

intra/netstack/fdbased.go

@@ -32,7 +32,6 @@ package netstack
 import (
 	"fmt"
 	"sync/atomic"
-	"syscall"
 	"time"
 
 	"github.com/celzero/firestack/intra/core"
@@ -53,6 +52,8 @@ var _ FdSwapper = (*linkFdSwap)(nil)
 
 const invalidfd int = -1
 
+const waitttl = wrapttl
+
 type FdSwapper interface {
 	// Cur returns the current FD.
 	Cur() int
@@ -71,16 +72,17 @@ type SeamlessEndpoint interface {
 // NetworkDispatcher.
 type linkDispatcher interface {
 	stop()
-	swap(fd int) error
-	dispatch() (bool, tcpip.Error)
+	prepare(fd *fds)
+	dispatch(fd *fds) (bool, tcpip.Error)
+	wrapup(prev *fds, ttl time.Duration)
 }
 
 type endpoint struct {
 	sync.RWMutex
 	// fds is the set of file descriptors each identifying one inbound/outbound
 	// channel. The endpoint will dispatch from all inbound channels as well as
 	// hash outbound packets to specific channels based on the packet hash.
-	fds *core.Volatile[int]
+	fds *core.Volatile[*fds]
 
 	// mtu (maximum transmission unit) is the maximum size of a packet.
 	mtu atomic.Uint32
@@ -196,7 +198,7 @@ func NewFdbasedInjectableEndpoint(opts *Options) (SeamlessEndpoint, error) {
 
 	e := &endpoint{
 		mtu:     atomic.Uint32{},
-		fds:     core.NewVolatile(invalidfd),
+		fds:     core.NewVolatile(invalidFds),
 		caps:    caps,
 		addr:    opts.Address,
 		hdrSize: hdrSize,
@@ -239,16 +241,11 @@ func (e *endpoint) Cur() int {
 }
 
 func (e *endpoint) Dispose() (err error) {
-	if e.fds == nil {
-		log.W("ns: tun: Dispose: no fds")
-		return nil
-	}
-
 	e.Lock()
 	defer e.Unlock()
 
-	prevfd := e.fds.Swap(invalidfd) // prevfd may be invalidfd
-	if prevfd == invalidfd {
+	prevfd := e.fds.Swap(invalidFds) // prevfd may be invalidfd
+	if !prevfd.ok() {
 		log.W("ns: tun: Dispose: invalid prevfd")
 		return nil
 	}
@@ -258,49 +255,50 @@ func (e *endpoint) Dispose() (err error) {
 		// nothing to do
 		return nil
 	}
-	// e.inboundDispatcher.swap() will close prevfd
-	// e.dispatchLoop() will auto-exit due to invalidfd
-	return e.inboundDispatcher.swap(invalidfd)
+	// e.inboundDispatcher.prepare() will not close prevfd
+	// dispatchLoop() will auto-exit on invalidfd
+	core.Go("ns.dispose.wrapup", func() { e.inboundDispatcher.wrapup(prevfd, wrapttl) })
+	e.inboundDispatcher.prepare(invalidFds)
+
+	return nil
 }
 
 // Implements Swapper.
 func (e *endpoint) Swap(fd int) (err error) {
-	if err = unix.SetNonblock(fd, true); err != nil {
-		clos(fd)
-		return fmt.Errorf("ns: tun: set non blocking(%d) failed: %v", fd, err)
-	}
-
 	e.Lock()
 	defer e.Unlock()
 
-	prevfd := e.fds.Swap(fd) // commence WritePackets() on fd
-	log.D("ns: swap: tun fd %d => %d", prevfd, fd)
+	f, err := newTun(fd) // fd may be invalid (ex: -1)
+	if err != nil {
+		fd = invalidfd
+		err = log.EE("ns: tun: swap: (%d) err: %v / %v; using invalidfd", fd, err)
+	}
+
+	prevfd := e.fds.Swap(f) // commence WritePackets() on fd
+	core.Go(
+		"ns.swap.wrapup", func() { e.inboundDispatcher.wrapup(prevfd, wrapttl) },
+	) // closes prevfd, which may be invalidfd
+
+	log.D("ns: tun: swap: fd %s => %d", prevfd, fd)
 
 	if e.inboundDispatcher == nil { // prevfd must be 0 value if inbound is nil
 		e.inboundDispatcher, err = createInboundDispatcher(e, fd)
 	} else {
-		err = e.inboundDispatcher.swap(fd) // always closes prevfd
-		// todo: on err != nil e.inboundDispatcher = nil?
+		e.inboundDispatcher.prepare(f)
 	}
 
 	hasDispatcher := e.dispatcher != nil
 	if err == nil && hasDispatcher { // attached?
-		log.I("ns: tun(%d => %d): swap: restart looper %t for new fd",
+		log.I("ns: tun: (%s => %d) swap: restart looper %t for new fd",
 			prevfd, fd, hasDispatcher)
-		go e.dispatchLoop(e.inboundDispatcher)
+		go dispatchLoop(e.inboundDispatcher, f, &e.wg)
 	} else {
-		log.E("ns: tun(%d => %d): swap: no dispatcher? %t for new fd; err %v",
+		log.E("ns: tun: (%s => %d) swap: no dispatcher? %t for new fd; err %v",
 			prevfd, fd, !hasDispatcher, err)
 	}
 	return
 }
 
-func clos(fd int) {
-	if fd > 0 || fd != invalidfd {
-		_ = syscall.Close(fd)
-	}
-}
-
 // Attach launches the goroutine that reads packets from the file descriptor and
 // dispatches them via the provided dispatcher.
 func (e *endpoint) Attach(dispatcher stack.NetworkDispatcher) {
@@ -310,41 +308,52 @@ func (e *endpoint) Attach(dispatcher stack.NetworkDispatcher) {
 	defer e.Unlock()
 
 	rx := e.inboundDispatcher
-	fd := e.fd()
+
+	fds := e.fds.Load()
+	fd := fds.tun()
+
 	attach := dispatcher != nil   // nil means the NIC is being removed.
 	pipe := rx != nil             // nil means there's no read dispatcher.
 	exists := e.dispatcher != nil // nil means the NIC is already detached.
+
 	// Attach is called when the NIC is being created and then enabled.
 	// stack.CreateNIC -> nic.newNIC -> ep.Attach
 	if dispatcher == nil && e.dispatcher != nil {
 		log.I("ns: tun(%d): attach: detach dispatcher (and inbound? %t)", fd, pipe)
+		allLoopersExited := true
 		if rx != nil {
+			fds.stop()
 			go rx.stop() // avoid mutex; closes fd
-			e.Wait()     // on all inboundDispatcher w/ mutex locked?
+
+			allLoopersExited = e.wait(waitttl) // on all inboundDispatcher w/ mutex locked?
 		}
 		e.dispatcher = nil
 		e.inboundDispatcher = nil // rx
-		e.fds.Store(invalidfd)
-		log.I("ns: tun(%d): attach: done detaching dispatcher", fd)
+		e.fds.Store(invalidFds)
+		logei(!allLoopersExited)("ns: tun(%d): attach: done detaching dispatcher; all loopers done? %t",
+			fd, allLoopersExited)
 		return
 	}
+
 	if dispatcher != nil && e.dispatcher == nil {
 		log.I("ns: tun(%d): attach: new dispatcher & looper", fd)
 		e.dispatcher = dispatcher
-		if e.inboundDispatcher == nil && fd != invalidfd { // unlikely
+		if e.inboundDispatcher == nil && fds.ok() { // unlikely
 			var err error
 			e.inboundDispatcher, err = createInboundDispatcher(e, fd)
 			logeif(err)("ns: tun(%d): attach: just-in-time createInboundDispatcher; err? %v", fd, err)
 			rx = e.inboundDispatcher
 		}
-		go e.dispatchLoop(rx)
+		go dispatchLoop(rx, fds, &e.wg)
 		return
 	}
+
 	if dispatcher != nil {
 		log.W("ns: tun(%d): attach: discard? %t; but switch to new anyway", fd, exists)
 		e.dispatcher = dispatcher
 		return
 	}
+
 	log.W("ns: tun(%d): attach: discard? %t; hadDispatcher? %t hadInbound? %t", fd, exists, attach, pipe)
 }
 
@@ -382,6 +391,11 @@ func (e *endpoint) Wait() {
 	e.wg.Wait()
 }
 
+func (e *endpoint) wait(d time.Duration) bool {
+	// wait on e.Wait() until ttl expires
+	return core.Await(func() { e.Wait() }, d)
+}
+
 // AddHeader implements stack.LinkEndpoint.AddHeader.
 func (e *endpoint) AddHeader(pkt *stack.PacketBuffer) {
 	if e.hdrSize > 0 && pkt != nil {
@@ -416,10 +430,7 @@ func (e *endpoint) ParseHeader(pkt *stack.PacketBuffer) bool {
 
 // fd returns the file descriptor associated with the endpoint.
 func (e *endpoint) fd() int {
-	if fd := e.fds.Load(); fd > 0 {
-		return fd
-	}
-	return invalidfd
+	return e.fds.Load().tun()
 }
 
 // writePackets writes outbound packets to the file descriptor. If it is not
@@ -482,31 +493,33 @@ func (e *endpoint) notifyRestart() {
 */
 
 // dispatchLoop reads packets from the file descriptor in a loop and dispatches
-// them to the network stack. Must be run as a goroutine.
-func (e *endpoint) dispatchLoop(inbound linkDispatcher) tcpip.Error {
+// them to the network stack (linkDispatcher). Must be run as a goroutine.
+func dispatchLoop(inbound linkDispatcher, f *fds, wg *sync.WaitGroup) tcpip.Error {
 	// defer core.RecoverFn("ns.e.dispatch", e.notifyRestart)
 	defer core.Recover(core.Exit11, "ns.e.dispatch")
 
-	e.wg.Add(1)
-	defer e.wg.Done()
+	wg.Add(1)
+	defer wg.Done()
 
-	fd := e.fd()
 	if inbound == nil || core.IsNil(inbound) {
-		log.W("ns: tun(%d): dispatchLoop: inbound nil", fd)
+		defer f.stop()
+		log.W("ns: tun(%d): dispatchLoop: inbound nil", f.tun())
 		return &tcpip.ErrUnknownDevice{}
 	}
 
 	start := time.Now()
-	log.I("ns: tun(%d): dispatchLoop: start", fd)
+	log.I("ns: tun(%d): dispatchLoop: start", f.tun())
 	for {
-		cont, err := inbound.dispatch()
+		cont, err := inbound.dispatch(f)
+		if err != nil {
+			log.W("ns: tun(%d): dispatchLoop: dur: %s; continue? %t; err: %v",
+				f.tun(), core.FmtTimeAsPeriod(start), cont, err)
+		}
 		if !cont {
-			elapsed := time.Since(start)
-			log.W("ns: tun(%d): dispatchLoop: exit; dur: %s; err? %v", fd, elapsed, err)
+			defer f.stop()
 			return err
-		} else if err != nil {
-			log.W("ns: tun(%d): dispatchLoop: continue on err: %v", fd, err)
 		} // else: continue dispatching
+
 	}
 }