Cancel VPN gateway loops together

agents/rap-node-agent/internal/vpnruntime/gateway_test.go

@@ -2,6 +2,8 @@ package vpnruntime
 
 import (
 	"context"
+	"errors"
+	"io"
 	"sync"
 	"testing"
 	"time"
@@ -12,6 +14,45 @@ type recordingGatewayTransport struct {
 	batches [][][]byte
 }
 
+type closingGatewayTransport struct {
+	closed chan struct{}
+}
+
+func (t *closingGatewayTransport) SendGatewayPacketBatch(context.Context, [][]byte) error {
+	return nil
+}
+
+func (t *closingGatewayTransport) ReceiveGatewayPacketBatch(context.Context, time.Duration) ([][]byte, error) {
+	return [][]byte{[]byte("packet")}, nil
+}
+
+func (t *closingGatewayTransport) Close() error {
+	close(t.closed)
+	return nil
+}
+
+type failingWriteTun struct {
+	closed chan struct{}
+}
+
+func (t failingWriteTun) Read([]byte) (int, error) {
+	<-t.closed
+	return 0, io.EOF
+}
+
+func (t failingWriteTun) Write([]byte) (int, error) {
+	return 0, errors.New("write failed")
+}
+
+func (t failingWriteTun) Close() error {
+	select {
+	case <-t.closed:
+	default:
+		close(t.closed)
+	}
+	return nil
+}
+
 func (t *recordingGatewayTransport) SendGatewayPacketBatch(ctx context.Context, packets [][]byte) error {
 	copied := make([][]byte, len(packets))
 	for i, packet := range packets {
@@ -36,6 +77,24 @@ func (t *recordingGatewayTransport) firstBatch() [][]byte {
 	return t.batches[0]
 }
 
+func TestGatewayRunClosesPacketTransportOnRuntimeError(t *testing.T) {
+	transport := &closingGatewayTransport{closed: make(chan struct{})}
+	gateway := &Gateway{
+		Transport:       transport,
+		VPNConnectionID: "vpn-1",
+		PollTimeout:     time.Millisecond,
+	}
+	err := gateway.run(context.Background(), failingWriteTun{closed: make(chan struct{})})
+	if err == nil || err.Error() != "write failed" {
+		t.Fatalf("run error = %v, want write failed", err)
+	}
+	select {
+	case <-transport.closed:
+	case <-time.After(time.Second):
+		t.Fatal("packet transport was not closed")
+	}
+}
+
 func TestGatewayUploadPrioritizesTCPControlPackets(t *testing.T) {
 	transport := &recordingGatewayTransport{}
 	gateway := &Gateway{Transport: transport, VPNConnectionID: "vpn-1"}