From 23536acfc4f4ea64ee7638b308659f362ccd25a5 Mon Sep 17 00:00:00 2001
From: Val Komarov <val.komarov@aopslab.com>
Date: Wed, 22 Oct 2014 15:31:54 -0400
Subject: [PATCH] Implemented timeouts.

---
 .../WebViewJavascriptBridge.h                 |  2 +
 .../WebViewJavascriptBridge.m                 | 48 +++++++++++++++----
 2 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/WebViewJavascriptBridge/WebViewJavascriptBridge.h b/WebViewJavascriptBridge/WebViewJavascriptBridge.h
index 9e90829f..02d611c1 100644
--- a/WebViewJavascriptBridge/WebViewJavascriptBridge.h
+++ b/WebViewJavascriptBridge/WebViewJavascriptBridge.h
@@ -35,9 +35,11 @@ typedef void (^WVJBHandler)(id data, WVJBResponseCallback responseCallback);
 
 - (void)send:(id)message;
 - (void)send:(id)message responseCallback:(WVJBResponseCallback)responseCallback;
+- (void)send:(id)message timeoutInMillis:(NSUInteger)timeout responseCallback:(WVJBResponseCallback)responseCallback;
 - (void)registerHandler:(NSString*)handlerName handler:(WVJBHandler)handler;
 - (void)callHandler:(NSString*)handlerName;
 - (void)callHandler:(NSString*)handlerName data:(id)data;
 - (void)callHandler:(NSString*)handlerName data:(id)data responseCallback:(WVJBResponseCallback)responseCallback;
+- (void)callHandler:(NSString*)handlerName data:(id)data timeoutInMillis:(NSUInteger)timeout responseCallback:(WVJBResponseCallback)responseCallback;
 
 @end
diff --git a/WebViewJavascriptBridge/WebViewJavascriptBridge.m b/WebViewJavascriptBridge/WebViewJavascriptBridge.m
index dff06e62..984defa6 100644
--- a/WebViewJavascriptBridge/WebViewJavascriptBridge.m
+++ b/WebViewJavascriptBridge/WebViewJavascriptBridge.m
@@ -54,12 +54,16 @@ + (instancetype)bridgeForWebView:(WVJB_WEBVIEW_TYPE*)webView webViewDelegate:(WV
     return bridge;
 }
 
-- (void)send:(id)data {
-    [self send:data responseCallback:nil];
+- (void)send:(id)message {
+    [self send:message responseCallback:nil];
 }
 
-- (void)send:(id)data responseCallback:(WVJBResponseCallback)responseCallback {
-    [self _sendData:data responseCallback:responseCallback handlerName:nil];
+- (void)send:(id)message responseCallback:(WVJBResponseCallback)responseCallback {
+    [self send:message timeoutInMillis:0 responseCallback:responseCallback];
+}
+
+- (void)send:(id)message timeoutInMillis:(NSUInteger)timeout responseCallback:(WVJBResponseCallback)responseCallback {
+    [self _sendData:message responseCallback:responseCallback handlerName:nil timeoutInMillis:timeout];
 }
 
 - (void)callHandler:(NSString *)handlerName {
@@ -71,7 +75,11 @@ - (void)callHandler:(NSString *)handlerName data:(id)data {
 }
 
 - (void)callHandler:(NSString *)handlerName data:(id)data responseCallback:(WVJBResponseCallback)responseCallback {
-    [self _sendData:data responseCallback:responseCallback handlerName:handlerName];
+    [self callHandler:handlerName data:data timeoutInMillis:0 responseCallback:responseCallback];
+}
+
+- (void)callHandler:(NSString*)handlerName data:(id)data timeoutInMillis:(NSUInteger)timeout responseCallback:(WVJBResponseCallback)responseCallback {
+    [self _sendData:data responseCallback:responseCallback handlerName:handlerName timeoutInMillis:timeout];
 }
 
 - (void)registerHandler:(NSString *)handlerName handler:(WVJBHandler)handler {
@@ -101,7 +109,7 @@ - (void)dealloc {
     _messageHandler = nil;
 }
 
-- (void)_sendData:(id)data responseCallback:(WVJBResponseCallback)responseCallback handlerName:(NSString*)handlerName {
+- (void)_sendData:(id)data responseCallback:(WVJBResponseCallback)responseCallback handlerName:(NSString*)handlerName timeoutInMillis:(NSUInteger)timeout {
     NSMutableDictionary* message = [NSMutableDictionary dictionary];
     
     if (data) {
@@ -112,6 +120,7 @@ - (void)_sendData:(id)data responseCallback:(WVJBResponseCallback)responseCallba
         NSString* callbackId = [NSString stringWithFormat:@"objc_cb_%ld", ++_uniqueId];
         _responseCallbacks[callbackId] = [responseCallback copy];
         message[@"callbackId"] = callbackId;
+        message[@"timeout"] = [NSNumber numberWithUnsignedInteger:timeout];
     }
     
     if (handlerName) {
@@ -149,6 +158,26 @@ - (void)_dispatchMessage:(WVJBMessage*)message {
             [strongWebView stringByEvaluatingJavaScriptFromString:javascriptCommand];
         });
     }
+    
+    [self _startTimeoutOnMessage:message];
+}
+
+- (void)_startTimeoutOnMessage:(WVJBMessage*)message {
+    NSString *callbackId = message[@"callbackId"];
+    NSNumber *timeout = message[@"timeout"];
+    if (callbackId && timeout && [timeout intValue] > 0) {
+        dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)([timeout unsignedIntegerValue] * NSEC_PER_MSEC)), dispatch_get_main_queue(), ^{
+            NSLog(@"Encountered timeout! The message is: %@", message);
+            
+            WVJBResponseCallback responseCallback = _responseCallbacks[callbackId];
+            
+            if (responseCallback) {
+                NSError *timeoutError = [NSError errorWithDomain:@"WVJBTimeoutError" code:408 userInfo:nil];
+                responseCallback(timeoutError);
+                [_responseCallbacks removeObjectForKey:callbackId];
+            }
+        });
+    }
 }
 
 - (void)_flushMessageQueue {
@@ -169,8 +198,11 @@ - (void)_flushMessageQueue {
         NSString* responseId = message[@"responseId"];
         if (responseId) {
             WVJBResponseCallback responseCallback = _responseCallbacks[responseId];
-            responseCallback(message[@"responseData"]);
-            [_responseCallbacks removeObjectForKey:responseId];
+            
+            if (responseCallback != nil) {
+                responseCallback(message[@"responseData"]);
+                [_responseCallbacks removeObjectForKey:responseId];
+            }
         } else {
             WVJBResponseCallback responseCallback = NULL;
             NSString* callbackId = message[@"callbackId"];