Skip to content

Commit

Permalink
fix: fix mtls issue in handwritten layer
Browse files Browse the repository at this point in the history
  • Loading branch information
arithmetic1728 committed Oct 21, 2020
1 parent ec8f5f2 commit d45bfbf
Show file tree
Hide file tree
Showing 4 changed files with 46 additions and 22 deletions.
21 changes: 13 additions & 8 deletions google/cloud/pubsub_v1/publisher/client.py
Expand Up @@ -130,15 +130,19 @@ def __init__(self, batch_settings=(), publisher_options=(), **kwargs):
target=os.environ.get("PUBSUB_EMULATOR_HOST")
)

# The GAPIC client has mTLS logic to determine the api endpoint and the
# ssl credentials to use. Here we create a GAPIC client to help compute the
# api endpoint and ssl credentials. The api endpoint will be used to set
# `self._target`, and ssl credentials will be passed to
# `grpc_helpers.create_channel` to establish a mTLS channel (if ssl
# credentials is not None).
client_options = kwargs.get("client_options", None)
if (
client_options
and "api_endpoint" in client_options
and isinstance(client_options["api_endpoint"], six.string_types)
):
self._target = client_options["api_endpoint"]
else:
self._target = publisher_client.PublisherClient.SERVICE_ADDRESS
credentials = kwargs.get("credentials", None)
client_for_mtls_info = publisher_client.PublisherClient(
credentials=credentials, client_options=client_options
)

self._target = client_for_mtls_info._transport._host

# Use a custom channel.
# We need this in order to set appropriate default message size and
Expand All @@ -149,6 +153,7 @@ def __init__(self, batch_settings=(), publisher_options=(), **kwargs):
channel = grpc_helpers.create_channel(
credentials=kwargs.pop("credentials", None),
target=self.target,
ssl_credentials=client_for_mtls_info._transport._ssl_channel_credentials,
scopes=publisher_client.PublisherClient._DEFAULT_SCOPES,
options={
"grpc.max_send_message_length": -1,
Expand Down
23 changes: 13 additions & 10 deletions google/cloud/pubsub_v1/subscriber/client.py
Expand Up @@ -16,7 +16,6 @@

import os
import pkg_resources
import six

import grpc

Expand Down Expand Up @@ -82,16 +81,19 @@ def __init__(self, **kwargs):
target=os.environ.get("PUBSUB_EMULATOR_HOST")
)

# api_endpoint wont be applied if 'transport' is passed in.
# The GAPIC client has mTLS logic to determine the api endpoint and the
# ssl credentials to use. Here we create a GAPIC client to help compute the
# api endpoint and ssl credentials. The api endpoint will be used to set
# `self._target`, and ssl credentials will be passed to
# `grpc_helpers.create_channel` to establish a mTLS channel (if ssl
# credentials is not None).
client_options = kwargs.get("client_options", None)
if (
client_options
and "api_endpoint" in client_options
and isinstance(client_options["api_endpoint"], six.string_types)
):
self._target = client_options["api_endpoint"]
else:
self._target = subscriber_client.SubscriberClient.SERVICE_ADDRESS
credentials = kwargs.get("credentials", None)
client_for_mtls_info = subscriber_client.SubscriberClient(
credentials=credentials, client_options=client_options
)

self._target = client_for_mtls_info._transport._host

# Use a custom channel.
# We need this in order to set appropriate default message size and
Expand All @@ -102,6 +104,7 @@ def __init__(self, **kwargs):
channel = grpc_helpers.create_channel(
credentials=kwargs.pop("credentials", None),
target=self.target,
ssl_credentials=client_for_mtls_info._transport._ssl_channel_credentials,
scopes=subscriber_client.SubscriberClient._DEFAULT_SCOPES,
options={
"grpc.max_send_message_length": -1,
Expand Down
12 changes: 10 additions & 2 deletions tests/unit/pubsub_v1/publisher/test_publisher_client.py
Expand Up @@ -18,6 +18,7 @@
import inspect

from google.auth import credentials
import grpc

import mock
import pytest
Expand Down Expand Up @@ -81,7 +82,7 @@ def test_init_w_api_endpoint():
assert isinstance(client.api, publisher_client.PublisherClient)
assert (client.api._transport.grpc_channel._channel.target()).decode(
"utf-8"
) == "testendpoint.google.com"
) == "testendpoint.google.com:443"


def test_init_w_unicode_api_endpoint():
Expand All @@ -91,7 +92,7 @@ def test_init_w_unicode_api_endpoint():
assert isinstance(client.api, publisher_client.PublisherClient)
assert (client.api._transport.grpc_channel._channel.target()).decode(
"utf-8"
) == "testendpoint.google.com"
) == "testendpoint.google.com:443"


def test_init_w_empty_client_options():
Expand All @@ -104,8 +105,13 @@ def test_init_w_empty_client_options():


def test_init_client_options_pass_through():
mock_ssl_creds = grpc.ssl_channel_credentials()

def init(self, *args, **kwargs):
self.kwargs = kwargs
self._transport = mock.Mock()
self._transport._host = "testendpoint.google.com"
self._transport._ssl_channel_credentials = mock_ssl_creds

with mock.patch.object(publisher_client.PublisherClient, "__init__", init):
client = publisher.Client(
Expand All @@ -119,6 +125,8 @@ def init(self, *args, **kwargs):
assert client_options.get("quota_project_id") == "42"
assert client_options.get("scopes") == []
assert client_options.get("credentials_file") == "file.json"
assert client.target == "testendpoint.google.com"
assert client.api.transport._ssl_channel_credentials == mock_ssl_creds


def test_init_emulator(monkeypatch):
Expand Down
12 changes: 10 additions & 2 deletions tests/unit/pubsub_v1/subscriber/test_subscriber_client.py
Expand Up @@ -13,6 +13,7 @@
# limitations under the License.

from google.auth import credentials
import grpc
import mock

from google.cloud.pubsub_v1 import subscriber
Expand Down Expand Up @@ -42,7 +43,7 @@ def test_init_w_api_endpoint():
assert isinstance(client.api, subscriber_client.SubscriberClient)
assert (client.api._transport.grpc_channel._channel.target()).decode(
"utf-8"
) == "testendpoint.google.com"
) == "testendpoint.google.com:443"


def test_init_w_unicode_api_endpoint():
Expand All @@ -52,7 +53,7 @@ def test_init_w_unicode_api_endpoint():
assert isinstance(client.api, subscriber_client.SubscriberClient)
assert (client.api._transport.grpc_channel._channel.target()).decode(
"utf-8"
) == "testendpoint.google.com"
) == "testendpoint.google.com:443"


def test_init_w_empty_client_options():
Expand All @@ -65,8 +66,13 @@ def test_init_w_empty_client_options():


def test_init_client_options_pass_through():
mock_ssl_creds = grpc.ssl_channel_credentials()

def init(self, *args, **kwargs):
self.kwargs = kwargs
self._transport = mock.Mock()
self._transport._host = "testendpoint.google.com"
self._transport._ssl_channel_credentials = mock_ssl_creds

with mock.patch.object(subscriber_client.SubscriberClient, "__init__", init):
client = subscriber.Client(
Expand All @@ -80,6 +86,8 @@ def init(self, *args, **kwargs):
assert client_options.get("quota_project_id") == "42"
assert client_options.get("scopes") == []
assert client_options.get("credentials_file") == "file.json"
assert client.target == "testendpoint.google.com"
assert client.api.transport._ssl_channel_credentials == mock_ssl_creds


def test_init_emulator(monkeypatch):
Expand Down

0 comments on commit d45bfbf

Please sign in to comment.