TUN-8748: Migrated datagram V3 flows to use migrated context

Previously, during local flow migration the current connection context
was not part of the migration and would cause the flow to still be listening
on the connection context of the old connection (before the migration).
This meant that if a flow was migrated from connection 0 to
connection 1, and connection 0 goes away, the flow would be early
terminated incorrectly with the context lifetime of connection 0.

The new connection context is provided during migration of a flow
and will trigger the observe loop for the flow lifetime to be rebound
to this provided context.
Closes TUN-8748

quic/v3/session_test.go

@@ -137,14 +137,28 @@ func TestSessionServe_Migrate(t *testing.T) {
 	defer session.Close()
 
 	done := make(chan error)
+	eyeball1Ctx, cancel := context.WithCancelCause(context.Background())
 	go func() {
-		done <- session.Serve(context.Background())
+		done <- session.Serve(eyeball1Ctx)
 	}()
 
 	// Migrate the session to a new connection before origin sends data
 	eyeball2 := newMockEyeball()
 	eyeball2.connID = 1
-	session.Migrate(&eyeball2, &log)
+	eyeball2Ctx := context.Background()
+	session.Migrate(&eyeball2, eyeball2Ctx, &log)
+
+	// Cancel the origin eyeball context; this should not cancel the session
+	contextCancelErr := errors.New("context canceled for first eyeball connection")
+	cancel(contextCancelErr)
+	select {
+	case <-done:
+		t.Fatalf("expected session to still be running")
+	default:
+	}
+	if context.Cause(eyeball1Ctx) != contextCancelErr {
+		t.Fatalf("first eyeball context should be cancelled manually: %+v", context.Cause(eyeball1Ctx))
+	}
 
 	// Origin sends data
 	payload2 := []byte{0xde}
@@ -166,6 +180,68 @@ func TestSessionServe_Migrate(t *testing.T) {
 	if !errors.Is(err, v3.SessionIdleErr{}) {
 		t.Error(err)
 	}
+	if eyeball2Ctx.Err() != nil {
+		t.Fatalf("second eyeball context should be not be cancelled")
+	}
+}
+
+func TestSessionServe_Migrate_CloseContext2(t *testing.T) {
+	log := zerolog.Nop()
+	eyeball := newMockEyeball()
+	pipe1, pipe2 := net.Pipe()
+	session := v3.NewSession(testRequestID, 2*time.Second, pipe2, testOriginAddr, testLocalAddr, &eyeball, &noopMetrics{}, &log)
+	defer session.Close()
+
+	done := make(chan error)
+	eyeball1Ctx, cancel := context.WithCancelCause(context.Background())
+	go func() {
+		done <- session.Serve(eyeball1Ctx)
+	}()
+
+	// Migrate the session to a new connection before origin sends data
+	eyeball2 := newMockEyeball()
+	eyeball2.connID = 1
+	eyeball2Ctx, cancel2 := context.WithCancelCause(context.Background())
+	session.Migrate(&eyeball2, eyeball2Ctx, &log)
+
+	// Cancel the origin eyeball context; this should not cancel the session
+	contextCancelErr := errors.New("context canceled for first eyeball connection")
+	cancel(contextCancelErr)
+	select {
+	case <-done:
+		t.Fatalf("expected session to still be running")
+	default:
+	}
+	if context.Cause(eyeball1Ctx) != contextCancelErr {
+		t.Fatalf("first eyeball context should be cancelled manually: %+v", context.Cause(eyeball1Ctx))
+	}
+
+	// Origin sends data
+	payload2 := []byte{0xde}
+	pipe1.Write(payload2)
+
+	// Expect write to eyeball2
+	data := <-eyeball2.recvData
+	if len(data) <= 17 || !slices.Equal(payload2, data[17:]) {
+		t.Fatalf("expected data to write to eyeball2 after migration: %+v", data)
+	}
+
+	select {
+	case data := <-eyeball.recvData:
+		t.Fatalf("expected no data to write to eyeball1 after migration: %+v", data)
+	default:
+	}
+
+	// Close the connection2 context manually
+	contextCancel2Err := errors.New("context canceled for second eyeball connection")
+	cancel2(contextCancel2Err)
+	err := <-done
+	if err != context.Canceled {
+		t.Fatalf("session Serve should be done: %+v", err)
+	}
+	if context.Cause(eyeball2Ctx) != contextCancel2Err {
+		t.Fatalf("second eyeball context should have been cancelled manually: %+v", context.Cause(eyeball2Ctx))
+	}
 }
 
 func TestSessionClose_Multiple(t *testing.T) {