weblayer: adds ability to detect when WebMessageReplyProxy closed

Currently there is no way to detect when a WebMessageReplyProxy is
closed so that the java side can do cleanup. This adds a callback
to detect this. Additionally this adds isClosed() to immediately
check if the channel has been closed.

BUG=1166171
TEST=WebMessageTest.testOnWebMessageReplyProxyClosed

Change-Id: Ia759fb38898cc798819988daa00555a59843207e
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2633644Reviewed-by: John Abd-El-Malek <jam@chromium.org>
Commit-Queue: Scott Violet <sky@chromium.org>
Cr-Commit-Position: refs/heads/master@{#844916}

weblayer/browser/android/javatests/src/org/chromium/weblayer/test/WebMessageTest.java

@@ -5,7 +5,10 @@
 package org.chromium.weblayer.test;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 import static org.chromium.content_public.browser.test.util.TestThreadUtils.runOnUiThreadBlocking;
@@ -39,6 +42,8 @@ public class WebMessageTest {
         public WebMessageReplyProxy mLastProxy;
         public WebMessage mLastMessage;
         public CallbackHelper mCallbackHelper;
+        public CallbackHelper mClosedCallbackHelper;
+        public WebMessageReplyProxy mProxyClosed;
 
         WebMessageCallbackImpl(CallbackHelper callbackHelper) {
             mCallbackHelper = callbackHelper;
@@ -53,6 +58,13 @@ public class WebMessageTest {
         public void reset() {
             mLastProxy = null;
             mLastMessage = null;
+            mProxyClosed = null;
+        }
+
+        @Override
+        public void onWebMessageReplyProxyClosed(WebMessageReplyProxy replyProxy) {
+            mProxyClosed = replyProxy;
+            if (mClosedCallbackHelper != null) mClosedCallbackHelper.notifyCalled();
         }
     }
 
@@ -90,6 +102,42 @@ public class WebMessageTest {
         runOnUiThreadBlocking(() -> { activity.getTab().unregisterWebMessageCallback("x"); });
     }
 
+    @Test
+    @SmallTest
+    public void testOnWebMessageReplyProxyClosed() throws Exception {
+        // Load a page with a form.
+        InstrumentationActivity activity = mActivityTestRule.launchShellWithUrl("about:blank");
+        assertNotNull(activity);
+        CallbackHelper callbackHelper = new CallbackHelper();
+        WebMessageCallbackImpl webMessageCallback = new WebMessageCallbackImpl(callbackHelper);
+        runOnUiThreadBlocking(() -> {
+            activity.getTab().registerWebMessageCallback(
+                    webMessageCallback, "x", Arrays.asList("*"));
+        });
+
+        mActivityTestRule.navigateAndWait(
+                mActivityTestRule.getTestDataURL("web_message_test.html"));
+        // web_message_test.html posts a message, wait for it.
+        callbackHelper.waitForCallback(0);
+        assertNotNull(webMessageCallback.mLastMessage);
+        assertNotNull(webMessageCallback.mLastProxy);
+        WebMessageReplyProxy proxy = webMessageCallback.mLastProxy;
+        assertNull(webMessageCallback.mProxyClosed);
+        assertFalse(
+                runOnUiThreadBlocking(() -> { return webMessageCallback.mLastProxy.isClosed(); }));
+        webMessageCallback.reset();
+        webMessageCallback.mClosedCallbackHelper = new CallbackHelper();
+
+        // Navigate to a new page. The proxy should be closed.
+        mActivityTestRule.navigateAndWait("about:blank");
+        webMessageCallback.mClosedCallbackHelper.waitForFirst();
+        assertNotNull(webMessageCallback.mProxyClosed);
+        assertEquals(webMessageCallback.mProxyClosed, proxy);
+        assertTrue(runOnUiThreadBlocking(() -> { return proxy.isClosed(); }));
+
+        runOnUiThreadBlocking(() -> { activity.getTab().unregisterWebMessageCallback("x"); });
+    }
+
     @Test
     @SmallTest
     public void testBadArguments() throws Exception {

weblayer/public/java/org/chromium/weblayer/Tab.java

@@ -798,8 +798,11 @@ public class Tab {
         @Override
         public void onReplyProxyDestroyed(int proxyId) {
             StrictModeWorkaround.apply();
-            assert mProxyIdToProxy.get(proxyId) != null;
+            WebMessageReplyProxy proxy = mProxyIdToProxy.get(proxyId);
+            assert proxy != null;
+            proxy.markClosed();
             mProxyIdToProxy.remove(proxyId);
+            mCallback.onWebMessageReplyProxyClosed(proxy);
         }
     }
 

weblayer/public/java/org/chromium/weblayer/WebMessageCallback.java

@@ -14,9 +14,23 @@ public abstract class WebMessageCallback {
     /**
      * Called when a message is received from the page.
      *
+     * Each {@link WebMessageReplyProxy} represents a single endpoint. Multiple messages sent to the
+     * same endpoint use the same {@link WebMessageReplyProxy}.
+     *
      * @param replyProxy An object that may be used to post a message back to the page.
      * @param message The message from the page.
      */
     public void onWebMessageReceived(
             @NonNull WebMessageReplyProxy replyProxy, @NonNull WebMessage message) {}
+
+    /**
+     * Called when a {@link WebMessageReplyProxy} is closed. This typically happens when navigating
+     * to another page. If the page goes into the back forward cache, then message channels are left
+     * open (and this is not called). In that case this method will be called either when the page
+     * is evicted from the cache or when the user goes back to it and then navigates away and it
+     * doesn't go into the back forward cache again.
+     *
+     * @param replyProxy The proxy that has been closed.
+     */
+    public void onWebMessageReplyProxyClosed(@NonNull WebMessageReplyProxy replyProxy) {}
 }

weblayer/public/java/org/chromium/weblayer/WebMessageReplyProxy.java

@@ -14,11 +14,15 @@ import org.chromium.weblayer_private.interfaces.IWebMessageReplyProxy;
 /**
  * Used to post a message to a page. WebMessageReplyProxy is created when a page posts a message to
  * the JavaScript object that was created by way of {@link Tab#registerWebMessageCallback}.
+ *
+ * Each {@link WebMessageReplyProxy} represents a single endpoint. Multiple messages sent to the
+ * same endpoint use the same {@link WebMessageReplyProxy}.
  */
 public class WebMessageReplyProxy {
     private final IWebMessageReplyProxy mIReplyProxy;
     private final boolean mIsMainFrame;
     private final String mSourceOrigin;
+    private boolean mIsClosed;
 
     // Constructor for test mocking.
     protected WebMessageReplyProxy() {
@@ -65,4 +69,17 @@ public class WebMessageReplyProxy {
             throw new APICallException(e);
         }
     }
+
+    /**
+     * Returns whether the reply proxy has been closed. Any messages sent when the channel is closed
+     * are dropped.
+     */
+    public boolean isClosed() {
+        ThreadCheck.ensureOnUiThread();
+        return mIsClosed;
+    }
+
+    void markClosed() {
+        mIsClosed = true;
+    }
 }