Fix: wrong wait logic

core/device/device.go

@@ -17,7 +17,4 @@ type Device interface {
 
 	// Type returns the driver type of the device.
 	Type() string
-
-	// Wait waits for the device to close.
-	Wait()
 }

core/device/iobased/endpoint.go

@@ -36,6 +36,9 @@ type Endpoint struct {
 
 	// once is used to perform the init action once when attaching.
 	once sync.Once
+
+	// wg keeps track of running goroutines.
+	wg sync.WaitGroup
 }
 
 // New returns stack.LinkEndpoint(.*Endpoint) and error.
@@ -60,19 +63,26 @@ func New(rw io.ReadWriter, mtu uint32, offset int) (*Endpoint, error) {
 	}, nil
 }
 
-func (e *Endpoint) Close() {
-	e.Endpoint.Close()
-}
-
 // Attach launches the goroutine that reads packets from io.Reader and
 // dispatches them via the provided dispatcher.
 func (e *Endpoint) Attach(dispatcher stack.NetworkDispatcher) {
+	e.Endpoint.Attach(dispatcher)
 	e.once.Do(func() {
 		ctx, cancel := context.WithCancel(context.Background())
-		go e.dispatchLoop(cancel)
-		go e.outboundLoop(ctx)
+		e.wg.Add(2)
+		go func() {
+			e.outboundLoop(ctx)
+			e.wg.Done()
+		}()
+		go func() {
+			e.dispatchLoop(cancel)
+			e.wg.Done()
+		}()
 	})
-	e.Endpoint.Attach(dispatcher)
+}
+
+func (e *Endpoint) Wait() {
+	e.wg.Wait()
 }
 
 // dispatchLoop dispatches packets to upper layer.

core/device/tun/tun_wireguard.go

@@ -69,6 +69,6 @@ func (t *TUN) Name() string {
 }
 
 func (t *TUN) Close() error {
-	t.Endpoint.Close()
+	defer t.Endpoint.Close()
 	return t.nt.Close()
 }

engine/engine.go

@@ -83,7 +83,6 @@ func (e *engine) start() error {
 func (e *engine) stop() (err error) {
 	if e.device != nil {
 		err = e.device.Close()
-		e.device.Wait()
 	}
 	if e.stack != nil {
 		e.stack.Close()