Revert "Revert "Add OAuth2 protocol in GRPCClient""

6 files changed, 133 insertions, 24 deletions

diff --git a/src/objective-c/GRPCClient/GRPCCall+GID.h b/src/objective-c/GRPCClient/GRPCCall+GID.h
new file mode 100644
index 0000000000..3ee732e79e
--- /dev/null
+++ b/src/objective-c/GRPCClient/GRPCCall+GID.h
@@ -0,0 +1,29 @@
+/*
+ *
+ * Copyright 2017 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+#import "GRPCCall.h"
+#import "GRPCCall+OAuth2.h"
+
+#import <Google/SignIn.h>
+
+/**
+ * Extend GIDSignIn class to comply GRPCAuthorizationProtocol
+ */
+@interface GIDSignIn (GRPC) <GRPCAuthorizationProtocol>
+- (void)getTokenWithHandler:(void (^)(NSString *token))hander;
+@end
diff --git a/src/objective-c/GRPCClient/GRPCCall+GID.m b/src/objective-c/GRPCClient/GRPCCall+GID.m
new file mode 100644
index 0000000000..030737147b
--- /dev/null
+++ b/src/objective-c/GRPCClient/GRPCCall+GID.m
@@ -0,0 +1,28 @@
+/*
+ *
+ * Copyright 2017 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+#import "GRPCCall+GID.h"
+
+@implementation GIDSignIn (GRPC)
+
+- (void)getTokenWithHandler:(void (^)(NSString *token))handler {
+  NSString *token = self.currentUser.authentication.accessToken;
+  handler(token);
+}
+
+@end
diff --git a/src/objective-c/GRPCClient/GRPCCall+OAuth2.h b/src/objective-c/GRPCClient/GRPCCall+OAuth2.h
index 65465e9523..adb1042aa0 100644
--- a/src/objective-c/GRPCClient/GRPCCall+OAuth2.h
+++ b/src/objective-c/GRPCClient/GRPCCall+OAuth2.h
@@ -18,6 +18,13 @@
 
 #import "GRPCCall.h"
 
+/**
+ * The protocol of an OAuth2 token object from which GRPCCall can acquire a token.
+ */
+@protocol GRPCAuthorizationProtocol
+- (void)getTokenWithHandler:(void (^)(NSString *token))hander;
+@end
+
 /** Helpers for setting and reading headers compatible with OAuth2. */
 @interface GRPCCall (OAuth2)
 
@@ -33,4 +40,12 @@
 /** Returns the value (if any) of the "www-authenticate" response header (the challenge header). */
 @property(atomic, readonly) NSString *oauth2ChallengeHeader;
 
+/**
+ * The authorization token object to be used when starting the call. If the value is set to nil, no
+ * oauth authentication will be used.
+ *
+ * If tokenProvider exists, it takes precedence over the token set by oauth2AccessToken.
+ */
+@property(atomic, strong) id<GRPCAuthorizationProtocol> tokenProvider;
+
 @end
diff --git a/src/objective-c/GRPCClient/GRPCCall+OAuth2.m b/src/objective-c/GRPCClient/GRPCCall+OAuth2.m
index eaa7465087..8451ebe870 100644
--- a/src/objective-c/GRPCClient/GRPCCall+OAuth2.m
+++ b/src/objective-c/GRPCClient/GRPCCall+OAuth2.m
@@ -16,6 +16,8 @@
  *
  */
 
+#import <objc/runtime.h>
+
 #import "GRPCCall+OAuth2.h"
 
 static NSString * const kAuthorizationHeader = @"authorization";
@@ -23,6 +25,7 @@ static NSString * const kBearerPrefix = @"Bearer ";
 static NSString * const kChallengeHeader = @"www-authenticate";
 
 @implementation GRPCCall (OAuth2)
+@dynamic tokenProvider;
 
 - (NSString *)oauth2AccessToken {
   NSString *headerValue = self.requestHeaders[kAuthorizationHeader];
@@ -45,4 +48,12 @@ static NSString * const kChallengeHeader = @"www-authenticate";
   return self.responseHeaders[kChallengeHeader];
 }
 
+- (void)setTokenProvider:(id<GRPCAuthorizationProtocol>)tokenProvider {
+  objc_setAssociatedObject(self, @selector(tokenProvider), tokenProvider, OBJC_ASSOCIATION_RETAIN);
+}
+
+- (id<GRPCAuthorizationProtocol>)tokenProvider {
+  return objc_getAssociatedObject(self, @selector(tokenProvider));
+}
+
 @end
diff --git a/src/objective-c/GRPCClient/GRPCCall.m b/src/objective-c/GRPCClient/GRPCCall.m
index 872362419e..436c19e354 100644
--- a/src/objective-c/GRPCClient/GRPCCall.m
+++ b/src/objective-c/GRPCClient/GRPCCall.m
@@ -18,6 +18,8 @@
 
 #import "GRPCCall.h"
 
+#import "GRPCCall+OAuth2.h"
+
 #include <grpc/grpc.h>
 #include <grpc/support/time.h>
 #import <RxLibrary/GRXConcurrentWriteable.h>
@@ -40,10 +42,14 @@ NSString * const kGRPCHeadersKey = @"io.grpc.HeadersKey";
 NSString * const kGRPCTrailersKey = @"io.grpc.TrailersKey";
 static NSMutableDictionary *callFlags;
 
+static NSString * const kAuthorizationHeader = @"authorization";
+static NSString * const kBearerPrefix = @"Bearer ";
+
 @interface GRPCCall () <GRXWriteable>
 // Make them read-write.
 @property(atomic, strong) NSDictionary *responseHeaders;
 @property(atomic, strong) NSDictionary *responseTrailers;
+@property(atomic) BOOL isWaitingForToken;
 @end
 
 // The following methods of a C gRPC call object aren't reentrant, and thus
@@ -181,9 +187,6 @@ static NSMutableDictionary *callFlags;
 
 - (void)finishWithError:(NSError *)errorOrNil {
   @synchronized(self) {
-    if (_state == GRXWriterStateFinished) {
-      return;
-    }
     _state = GRXWriterStateFinished;
   }
 
@@ -211,7 +214,11 @@ static NSMutableDictionary *callFlags;
   [self finishWithError:[NSError errorWithDomain:kGRPCErrorDomain
                                             code:GRPCErrorCodeCancelled
                                         userInfo:@{NSLocalizedDescriptionKey: @"Canceled by app"}]];
-  [self cancelCall];
+  if (!self.isWaitingForToken) {
+    [self cancelCall];
+  } else {
+    self.isWaitingForToken = NO;
+  }
 }
 
 - (void)dealloc {
@@ -410,22 +417,13 @@ static NSMutableDictionary *callFlags;
 
 #pragma mark GRXWriter implementation
 
-- (void)startWithWriteable:(id<GRXWriteable>)writeable {
-  @synchronized(self) {
-    _state = GRXWriterStateStarted;
-  }
-
-  // Create a retain cycle so that this instance lives until the RPC finishes (or is cancelled).
-  // This makes RPCs in which the call isn't externally retained possible (as long as it is started
-  // before being autoreleased).
-  // Care is taken not to retain self strongly in any of the blocks used in this implementation, so
-  // that the life of the instance is determined by this retain cycle.
-  _retainSelf = self;
-
+- (void)startCallWithWriteable:(id<GRXWriteable>)writeable {
   _responseWriteable = [[GRXConcurrentWriteable alloc] initWithWriteable:writeable
                                                            dispatchQueue:_responseQueue];
 
-  _wrappedCall = [[GRPCWrappedCall alloc] initWithHost:_host serverName:_serverName path:_path];
+  _wrappedCall = [[GRPCWrappedCall alloc] initWithHost:_host
+                                            serverName:_serverName
+                                                  path:_path];
   NSAssert(_wrappedCall, @"Error allocating RPC objects. Low memory?");
 
   [self sendHeaders:_requestHeaders];
@@ -437,20 +435,49 @@ static NSMutableDictionary *callFlags;
     // TODO(jcanizales): Check this on init.
     [NSException raise:NSInvalidArgumentException format:@"host of %@ is nil", _host];
   }
-  __weak typeof(self) weakSelf = self;
   _connectivityMonitor = [GRPCConnectivityMonitor monitorWithHost:host];
+  __weak typeof(self) weakSelf = self;
   void (^handler)() = ^{
     typeof(self) strongSelf = weakSelf;
-    if (strongSelf) {
-      [strongSelf finishWithError:[NSError errorWithDomain:kGRPCErrorDomain
-                                                      code:GRPCErrorCodeUnavailable
-                                                  userInfo:@{ NSLocalizedDescriptionKey : @"Connectivity lost." }]];
-    }
+    [strongSelf finishWithError:[NSError errorWithDomain:kGRPCErrorDomain
+                                                    code:GRPCErrorCodeUnavailable
+                                                userInfo:@{ NSLocalizedDescriptionKey : @"Connectivity lost." }]];
   };
   [_connectivityMonitor handleLossWithHandler:handler
                       wifiStatusChangeHandler:nil];
 }
 
+- (void)startWithWriteable:(id<GRXWriteable>)writeable {
+  @synchronized(self) {
+    _state = GRXWriterStateStarted;
+  }
+
+  // Create a retain cycle so that this instance lives until the RPC finishes (or is cancelled).
+  // This makes RPCs in which the call isn't externally retained possible (as long as it is started
+  // before being autoreleased).
+  // Care is taken not to retain self strongly in any of the blocks used in this implementation, so
+  // that the life of the instance is determined by this retain cycle.
+  _retainSelf = self;
+
+  if (self.tokenProvider != nil) {
+    self.isWaitingForToken = YES;
+    __weak typeof(self) weakSelf = self;
+    [self.tokenProvider getTokenWithHandler:^(NSString *token){
+      typeof(self) strongSelf = weakSelf;
+      if (strongSelf && strongSelf.isWaitingForToken) {
+        if (token) {
+          NSString *t = [kBearerPrefix stringByAppendingString:token];
+          strongSelf.requestHeaders[kAuthorizationHeader] = t;
+        }
+        [strongSelf startCallWithWriteable:writeable];
+        strongSelf.isWaitingForToken = NO;
+      }
+    }];
+  } else {
+    [self startCallWithWriteable:writeable];
+  }
+}
+
 - (void)setState:(GRXWriterState)newState {
   @synchronized(self) {
     // Manual transitions are only allowed from the started or paused states.
diff --git a/src/objective-c/GRPCClient/private/GRPCRequestHeaders.m b/src/objective-c/GRPCClient/private/GRPCRequestHeaders.m
index 7640a64d6d..5de1d8fff5 100644
--- a/src/objective-c/GRPCClient/private/GRPCRequestHeaders.m
+++ b/src/objective-c/GRPCClient/private/GRPCRequestHeaders.m
@@ -103,7 +103,6 @@ static void CheckKeyValuePairIsValid(NSString *key, id value) {
 }
 
 - (void)setObject:(id)obj forKey:(NSString *)key {
-  [self checkCallIsNotStarted];
   CheckIsNonNilASCII(@"Header name", key);
   key = key.lowercaseString;
   CheckKeyValuePairIsValid(key, obj);