3 files changed, 68 insertions, 24 deletions

diff --git a/src/objective-c/GRPCClient/GRPCCall.h b/src/objective-c/GRPCClient/GRPCCall.h
index 178a446c8b..df6220c644 100644
--- a/src/objective-c/GRPCClient/GRPCCall.h
+++ b/src/objective-c/GRPCClient/GRPCCall.h
@@ -140,6 +140,13 @@ typedef NS_ENUM(NSUInteger, GRPCErrorCode) {
 };
 
 /**
+ * The protocol of an OAuth2 token object from which GRPCCall can acquire a token.
+ */
+@protocol GRPCAuthorizationProtocol
+- (void)getTokenWithHandler:(void (^)(NSString *token))hander;
+@end
+
+/**
  * Safety remark of a gRPC method as defined in RFC 2616 Section 9.1
  */
 typedef NS_ENUM(NSUInteger, GRPCCallSafety) {
@@ -216,6 +223,14 @@ extern id const kGRPCTrailersKey;
 @property(atomic, readonly) NSDictionary *responseTrailers;
 
 /**
+ * The authorization token object to be used when starting the call. If the value is set to nil, no 
+ * oauth authentication will be used.
+ *
+ * Not compatible with property oauth2AccessToken in GRPCCall (OAuth2). Do not use both at the same time.
+ */
+@property(atomic, strong) id<GRPCAuthorizationProtocol> oauthToken;
+
+/**
  * The request writer has to write NSData objects into the provided Writeable. The server will
  * receive each of those separately and in order as distinct messages.
  * A gRPC call might not complete until the request writer finishes. On the other hand, the request
diff --git a/src/objective-c/GRPCClient/GRPCCall.m b/src/objective-c/GRPCClient/GRPCCall.m
index 872362419e..3b937acd66 100644
--- a/src/objective-c/GRPCClient/GRPCCall.m
+++ b/src/objective-c/GRPCClient/GRPCCall.m
@@ -40,10 +40,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
@@ -211,7 +215,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 {
@@ -422,33 +430,55 @@ static NSMutableDictionary *callFlags;
   // that the life of the instance is determined by this retain cycle.
   _retainSelf = self;
 
-  _responseWriteable = [[GRXConcurrentWriteable alloc] initWithWriteable:writeable
-                                                           dispatchQueue:_responseQueue];
-
-  _wrappedCall = [[GRPCWrappedCall alloc] initWithHost:_host serverName:_serverName path:_path];
-  NSAssert(_wrappedCall, @"Error allocating RPC objects. Low memory?");
-
-  [self sendHeaders:_requestHeaders];
-  [self invokeCall];
-
-  // TODO(jcanizales): Extract this logic somewhere common.
-  NSString *host = [NSURL URLWithString:[@"https://" stringByAppendingString:_host]].host;
-  if (!host) {
-    // TODO(jcanizales): Check this on init.
-    [NSException raise:NSInvalidArgumentException format:@"host of %@ is nil", _host];
-  }
   __weak typeof(self) weakSelf = self;
-  _connectivityMonitor = [GRPCConnectivityMonitor monitorWithHost:host];
-  void (^handler)() = ^{
+  void (^performCall)() = ^{
     typeof(self) strongSelf = weakSelf;
     if (strongSelf) {
-      [strongSelf finishWithError:[NSError errorWithDomain:kGRPCErrorDomain
-                                                      code:GRPCErrorCodeUnavailable
-                                                  userInfo:@{ NSLocalizedDescriptionKey : @"Connectivity lost." }]];
+      strongSelf->_responseWriteable = [[GRXConcurrentWriteable alloc] initWithWriteable:writeable
+                                                                           dispatchQueue:strongSelf->_responseQueue];
+
+      strongSelf->_wrappedCall = [[GRPCWrappedCall alloc] initWithHost:strongSelf->_host
+                                                            serverName:strongSelf->_serverName
+                                                                  path:strongSelf->_path];
+      NSAssert(_wrappedCall, @"Error allocating RPC objects. Low memory?");
+
+      [strongSelf sendHeaders:_requestHeaders];
+      [strongSelf invokeCall];
+
+      // TODO(jcanizales): Extract this logic somewhere common.
+      NSString *host = [NSURL URLWithString:[@"https://" stringByAppendingString:strongSelf->_host]].host;
+      if (!host) {
+        // TODO(jcanizales): Check this on init.
+        [NSException raise:NSInvalidArgumentException format:@"host of %@ is nil", strongSelf->_host];
+      }
+      strongSelf->_connectivityMonitor = [GRPCConnectivityMonitor monitorWithHost:host];
+      void (^handler)() = ^{
+        typeof(self) strongSelf = weakSelf;
+        [strongSelf finishWithError:[NSError errorWithDomain:kGRPCErrorDomain
+                                                        code:GRPCErrorCodeUnavailable
+                                                    userInfo:@{ NSLocalizedDescriptionKey : @"Connectivity lost." }]];
+      };
+      [_connectivityMonitor handleLossWithHandler:handler
+                          wifiStatusChangeHandler:nil];
     }
   };
-  [_connectivityMonitor handleLossWithHandler:handler
-                      wifiStatusChangeHandler:nil];
+
+  if (self.oauthToken != nil) {
+    self.isWaitingForToken = YES;
+    [self.oauthToken getTokenWithHandler:^(NSString *token){
+      typeof(self) strongSelf = weakSelf;
+      if (strongSelf && strongSelf.isWaitingForToken) {
+        if (token) {
+          NSString *t = [kBearerPrefix stringByAppendingString:token];
+          strongSelf.requestHeaders[kAuthorizationHeader] = t;
+        }
+        performCall();
+        strongSelf.isWaitingForToken = NO;
+      }
+    }];
+  } else {
+    performCall();
+  }
 }
 
 - (void)setState:(GRXWriterState)newState {
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);