Support Two-legged OAuth flow (#874)

Author: mantas-sidlauskas

build.gradle

@@ -71,6 +71,8 @@ dependencies {
     compile group: 'com.google.api.grpc', name: 'proto-google-common-protos', version: '2.10.0'
     compile group: 'io.grpc', name: 'grpc-testing', version: '1.54.2'
     compile group: 'com.google.protobuf', name: 'protobuf-java-util', version: '3.21.9'
+    compile group: 'com.google.oauth-client', name: 'google-oauth-client', version: '1.35.0'
+    compile group: 'com.google.http-client', name: 'google-http-client-gson', version: '1.19.0'
 
     implementation 'io.grpc:grpc-netty-shaded:1.54.2'
     implementation 'io.grpc:grpc-protobuf:1.54.2'

src/main/java/com/uber/cadence/internal/compatibility/proto/serviceclient/GrpcServiceStubs.java

@@ -34,6 +34,7 @@
 import com.uber.cadence.internal.Version;
 import com.uber.cadence.internal.tracing.TracingPropagator;
 import com.uber.cadence.serviceclient.ClientOptions;
+import com.uber.cadence.serviceclient.auth.IAuthorizationProvider;
 import io.grpc.CallOptions;
 import io.grpc.Channel;
 import io.grpc.ClientCall;
@@ -53,6 +54,7 @@
 import io.opentelemetry.context.propagation.TextMapSetter;
 import io.opentracing.Span;
 import io.opentracing.Tracer;
+import java.nio.charset.StandardCharsets;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
@@ -79,6 +81,9 @@ final class GrpcServiceStubs implements IGrpcServiceStubs {
   private static final Metadata.Key<String> RPC_ENCODING_HEADER_KEY =
       Metadata.Key.of("rpc-encoding", Metadata.ASCII_STRING_MARSHALLER);
 
+  private static final Metadata.Key<String> AUTHORIZATION_HEADER_KEY =
+      Metadata.Key.of("cadence-authorization", Metadata.ASCII_STRING_MARSHALLER);
+
   private static final String CLIENT_IMPL_HEADER_VALUE = "uber-java";
 
   private final ClientOptions options;
@@ -121,6 +126,7 @@ final class GrpcServiceStubs implements IGrpcServiceStubs {
     if (!Strings.isNullOrEmpty(options.getIsolationGroup())) {
       headers.put(ISOLATION_GROUP_HEADER_KEY, options.getIsolationGroup());
     }
+
     Channel interceptedChannel =
         ClientInterceptors.intercept(
             channel,
@@ -131,6 +137,11 @@ final class GrpcServiceStubs implements IGrpcServiceStubs {
     if (log.isTraceEnabled()) {
       interceptedChannel = ClientInterceptors.intercept(interceptedChannel, tracingInterceptor);
     }
+    if (options.getAuthProvider() != null) {
+      interceptedChannel =
+          ClientInterceptors.intercept(
+              interceptedChannel, newAuthorizationInterceptor(options.getAuthProvider()));
+    }
     this.domainBlockingStub = DomainAPIGrpc.newBlockingStub(interceptedChannel);
     this.domainFutureStub = DomainAPIGrpc.newFutureStub(interceptedChannel);
     this.visibilityBlockingStub = VisibilityAPIGrpc.newBlockingStub(interceptedChannel);
@@ -143,6 +154,36 @@ final class GrpcServiceStubs implements IGrpcServiceStubs {
     this.metaFutureStub = MetaAPIGrpc.newFutureStub(interceptedChannel);
   }
 
+  private ClientInterceptor newAuthorizationInterceptor(IAuthorizationProvider provider) {
+    return new ClientInterceptor() {
+      @Override
+      public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
+          MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {
+        return new ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT>(
+            next.newCall(method, callOptions)) {
+
+          @Override
+          public void start(Listener<RespT> responseListener, Metadata headers) {
+            headers.put(
+                AUTHORIZATION_HEADER_KEY,
+                new String(provider.getAuthToken(), StandardCharsets.UTF_8));
+
+            Listener<RespT> listener =
+                new ForwardingClientCallListener.SimpleForwardingClientCallListener<RespT>(
+                    responseListener) {
+
+                  @Override
+                  public void onHeaders(Metadata headers) {
+                    super.onHeaders(headers);
+                  }
+                };
+            super.start(listener, headers);
+          }
+        };
+      }
+    };
+  }
+
   private ClientInterceptor newOpenTelemetryInterceptor() {
     return new ClientInterceptor() {
       @Override

src/main/java/com/uber/cadence/serviceclient/auth/OAuthAuthorizationProvider.java

@@ -0,0 +1,65 @@
+/**
+ * Copyright 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * <p>Modifications copyright (C) 2017 Uber Technologies, Inc.
+ *
+ * <p>Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file
+ * except in compliance with the License. A copy of the License is located at
+ *
+ * <p>http://aws.amazon.com/apache2.0
+ *
+ * <p>or in the "license" file accompanying this file. This file 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.
+ */
+package com.uber.cadence.serviceclient.auth;
+
+import com.google.api.client.auth.oauth2.ClientCredentialsTokenRequest;
+import com.google.api.client.auth.oauth2.ClientParametersAuthentication;
+import com.google.api.client.auth.oauth2.TokenResponse;
+import com.google.api.client.http.GenericUrl;
+import com.google.api.client.http.HttpTransport;
+import com.google.api.client.http.javanet.NetHttpTransport;
+import com.google.api.client.json.JsonFactory;
+import com.google.api.client.json.gson.GsonFactory;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.time.Instant;
+import java.util.List;
+
+public class OAuthAuthorizationProvider implements IAuthorizationProvider {
+  private static final HttpTransport HTTP_TRANSPORT = new NetHttpTransport();
+
+  private Instant nextRefresh;
+
+  private final ClientCredentialsTokenRequest tokenRequest;
+
+  private TokenResponse token;
+
+  static final JsonFactory JSON_FACTORY = new GsonFactory();
+
+  public OAuthAuthorizationProvider(
+      String clientId, String clientSecret, String url, List<String> scopes) {
+    this.tokenRequest =
+        new ClientCredentialsTokenRequest(HTTP_TRANSPORT, JSON_FACTORY, new GenericUrl(url))
+            .setClientAuthentication(new ClientParametersAuthentication(clientId, clientSecret))
+            .setScopes(scopes);
+  }
+
+  @Override
+  public byte[] getAuthToken() {
+    return fetchToken().getBytes(StandardCharsets.UTF_8);
+  }
+
+  private synchronized String fetchToken() {
+    if (this.nextRefresh == null || Instant.now().isAfter(this.nextRefresh)) {
+      try {
+        this.token = this.tokenRequest.execute();
+        this.nextRefresh = Instant.now().plusSeconds(this.token.getExpiresInSeconds());
+      } catch (IOException e) {
+        return "";
+      }
+    }
+    return this.token.getAccessToken();
+  }
+}

src/test/java/com/uber/cadence/serviceclient/auth/OAuthAuthorizationProviderTest.java

@@ -0,0 +1,29 @@
+package com.uber.cadence.serviceclient.auth;
+
+import static org.junit.Assert.*;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+public class OAuthAuthorizationProviderTest {
+  @Test
+  public void testConstructorWillThrowException() {
+    try {
+      final OAuthAuthorizationProvider provider =
+          new OAuthAuthorizationProvider("a", "b", "c", null);
+    } catch (IllegalArgumentException e) {
+      Assert.assertEquals(IllegalArgumentException.class, e.getClass());
+    }
+  }
+
+  @Test
+  public void testClientWillReturnEmptyOnWrongServer() {
+
+    final OAuthAuthorizationProvider provider =
+        new OAuthAuthorizationProvider("a", "b", "https://test", null);
+
+    byte[] token = provider.getAuthToken();
+
+    assertEquals(0, token.length);
+  }
+}