From ead36d152a003468d26fe939f70bc05ef0a93866 Mon Sep 17 00:00:00 2001
From: sanjaypujare <sanjaypujare@users.noreply.github.com>
Date: Wed, 18 Dec 2019 09:40:23 -0800
Subject: [PATCH] xds: use UpstreamTlsContext from attrs  (#6533)

---
 .../java/io/grpc/xds/sds/XdsChannelBuilder.java   |  2 +-
 .../xds/sds/internal/SdsProtocolNegotiators.java  | 15 ++++++++++-----
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/xds/src/main/java/io/grpc/xds/sds/XdsChannelBuilder.java b/xds/src/main/java/io/grpc/xds/sds/XdsChannelBuilder.java
index 2e4e5d2c05..1091d86782 100644
--- a/xds/src/main/java/io/grpc/xds/sds/XdsChannelBuilder.java
+++ b/xds/src/main/java/io/grpc/xds/sds/XdsChannelBuilder.java
@@ -37,7 +37,7 @@ public final class XdsChannelBuilder extends ForwardingChannelBuilder<XdsChannel
 
   private final NettyChannelBuilder delegate;
 
-  // TODO (sanjaypujare) integrate with xDS client to get upstreamTlsContext from CDS
+  // TODO (sanjaypujare) remove once we get this from CDS & don't need for testing
   @Nullable private UpstreamTlsContext upstreamTlsContext;
 
   private XdsChannelBuilder(NettyChannelBuilder delegate) {
diff --git a/xds/src/main/java/io/grpc/xds/sds/internal/SdsProtocolNegotiators.java b/xds/src/main/java/io/grpc/xds/sds/internal/SdsProtocolNegotiators.java
index bf291d665e..925115b640 100644
--- a/xds/src/main/java/io/grpc/xds/sds/internal/SdsProtocolNegotiators.java
+++ b/xds/src/main/java/io/grpc/xds/sds/internal/SdsProtocolNegotiators.java
@@ -29,6 +29,7 @@ import io.grpc.netty.InternalProtocolNegotiator;
 import io.grpc.netty.InternalProtocolNegotiator.ProtocolNegotiator;
 import io.grpc.netty.InternalProtocolNegotiators;
 import io.grpc.netty.NettyChannelBuilder;
+import io.grpc.xds.XdsAttributes;
 import io.grpc.xds.sds.SslContextProvider;
 import io.grpc.xds.sds.TlsContextManager;
 import io.netty.channel.ChannelHandler;
@@ -109,7 +110,7 @@ public final class SdsProtocolNegotiators {
   @VisibleForTesting
   static final class ClientSdsProtocolNegotiator implements ProtocolNegotiator {
 
-    // TODO (sanjaypujare) integrate with xDS client to get upstreamTlsContext from CDS
+    // TODO (sanjaypujare) remove once we get this from CDS & don't need for testing
     UpstreamTlsContext upstreamTlsContext;
 
     ClientSdsProtocolNegotiator(UpstreamTlsContext upstreamTlsContext) {
@@ -123,12 +124,16 @@ public final class SdsProtocolNegotiators {
 
     @Override
     public ChannelHandler newHandler(GrpcHttp2ConnectionHandler grpcHandler) {
-      // once CDS is implemented we will retrieve upstreamTlsContext as follows:
-      // grpcHandler.getEagAttributes().get(XdsAttributes.ATTR_UPSTREAM_TLS_CONTEXT);
-      if (isTlsContextEmpty(upstreamTlsContext)) {
+      // check if UpstreamTlsContext was passed via attributes
+      UpstreamTlsContext localUpstreamTlsContext =
+          grpcHandler.getEagAttributes().get(XdsAttributes.ATTR_UPSTREAM_TLS_CONTEXT);
+      if (localUpstreamTlsContext == null) {
+        localUpstreamTlsContext = upstreamTlsContext;
+      }
+      if (isTlsContextEmpty(localUpstreamTlsContext)) {
         return InternalProtocolNegotiators.plaintext().newHandler(grpcHandler);
       }
-      return new ClientSdsHandler(grpcHandler, upstreamTlsContext);
+      return new ClientSdsHandler(grpcHandler, localUpstreamTlsContext);
     }
 
     private static boolean isTlsContextEmpty(UpstreamTlsContext upstreamTlsContext) {