1 files changed, 50 insertions, 23 deletions

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.