Upstream ios/web/web_state/ui support classes

ios/web/web_state/ui/crw_wk_script_message_router.mm

Index: ios/web/web_state/ui/crw_wk_script_message_router.mm
diff --git a/ios/web/web_state/ui/crw_wk_script_message_router.mm b/ios/web/web_state/ui/crw_wk_script_message_router.mm
new file mode 100644
index 0000000000000000000000000000000000000000..95ba4f21c9828345f5197d247c5ee2dbe8fce491
--- /dev/null
+++ b/ios/web/web_state/ui/crw_wk_script_message_router.mm
@@ -0,0 +1,113 @@
+// Copyright 2014 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#import "ios/web/web_state/ui/crw_wk_script_message_router.h"
+
+#import "base/logging.h"
+#import "base/mac/scoped_nsobject.h"
+
+@interface CRWWKScriptMessageRouter ()<WKScriptMessageHandler>
+
+// Removes a specific message handler. Does nothing if handler does not exist.
+- (void)tryRemoveScriptMessageHandlerForName:(NSString*)messageName
+                                     webView:(WKWebView*)webView;
+
+@end
+
+@implementation CRWWKScriptMessageRouter {
+  // Two level map of registed message handlers. Keys are message names and
+  // values are more maps (where keys are web views and values are handlers).
+  base::scoped_nsobject<NSMutableDictionary> _handlers;
+  // Wrapped WKUserContentController.
+  base::scoped_nsobject<WKUserContentController> _userContentController;
+}
+
+#pragma mark -
+#pragma mark Interface
+
+- (WKUserContentController*)userContentController {
+  return _userContentController.get();
+}
+
+- (instancetype)init {
+  NOTREACHED();
+  return nil;
+}
+
+- (instancetype)initWithUserContentController:
+    (WKUserContentController*)userContentController {
+  DCHECK(userContentController);
+  if ((self = [super init])) {
+    _handlers.reset([[NSMutableDictionary alloc] init]);
+    _userContentController.reset([userContentController retain]);
+  }
+  return self;
+}
+
+- (void)setScriptMessageHandler:(void(^)(WKScriptMessage*))handler
+                           name:(NSString*)messageName
+                        webView:(WKWebView*)webView {
+  DCHECK(handler);
+  DCHECK(messageName);
+  DCHECK(webView);
+
+  NSMapTable* webViewToHandlerMap = [_handlers objectForKey:messageName];
+  if (!webViewToHandlerMap) {
+    webViewToHandlerMap = [NSMapTable
+        mapTableWithKeyOptions:NSPointerFunctionsStrongMemory
+                  valueOptions:NSPointerFunctionsCopyIn];
+    [_handlers setObject:webViewToHandlerMap forKey:messageName];
+    [_userContentController addScriptMessageHandler:self name:messageName];
+  }
+  DCHECK(![webViewToHandlerMap objectForKey:webView]);
+  [webViewToHandlerMap setObject:handler forKey:webView];
+}
+
+- (void)removeScriptMessageHandlerForName:(NSString*)messageName
+                                  webView:(WKWebView*)webView {
+  DCHECK(messageName);
+  DCHECK(webView);
+  DCHECK([[_handlers objectForKey:messageName] objectForKey:webView]);
+  [self tryRemoveScriptMessageHandlerForName:messageName webView:webView];
+}
+
+- (void)removeAllScriptMessageHandlersForWebView:(WKWebView*)webView {
+  DCHECK(webView);
+  for (NSString* messageName in [_handlers allKeys]) {
+    [self tryRemoveScriptMessageHandlerForName:messageName webView:webView];
+  }
+}
+
+#pragma mark -
+#pragma mark WKScriptMessageHandler
+
+- (void)userContentController:(WKUserContentController*)userContentController
+      didReceiveScriptMessage:(WKScriptMessage*)message {
+  NSMapTable* webViewToHandlerMap = [_handlers objectForKey:message.name];
+  DCHECK(webViewToHandlerMap);
+  id handler = [webViewToHandlerMap objectForKey:message.webView];
+  if (handler) {
+    // Web process can send messages even if web view was reset and
+    // script message handler has been removed from the router.
+    ((void(^)(WKScriptMessage*))handler)(message);
+  }
+}
+
+#pragma mark -
+#pragma mark Implementation
+
+- (void)tryRemoveScriptMessageHandlerForName:(NSString*)messageName
+                                     webView:(WKWebView*)webView {
+  NSMapTable* webViewToHandlerMap = [_handlers objectForKey:messageName];
+  if (![webViewToHandlerMap objectForKey:webView])
+    return;
+  if (webViewToHandlerMap.count == 1) {
+    [_handlers removeObjectForKey:messageName];
+    [_userContentController removeScriptMessageHandlerForName:messageName];
+  } else {
+    [webViewToHandlerMap removeObjectForKey:webView];
+  }
+}
+
+@end