• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2015 gRPC authors.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import collections
16import logging
17import threading
18
19import grpc
20from grpc import _common
21from grpc._cython import cygrpc
22
23logging.basicConfig()
24_LOGGER = logging.getLogger(__name__)
25
26
27class _AuthMetadataContext(
28        collections.namedtuple('AuthMetadataContext', (
29            'service_url',
30            'method_name',
31        )), grpc.AuthMetadataContext):
32    pass
33
34
35class _CallbackState(object):
36
37    def __init__(self):
38        self.lock = threading.Lock()
39        self.called = False
40        self.exception = None
41
42
43class _AuthMetadataPluginCallback(grpc.AuthMetadataPluginCallback):
44
45    def __init__(self, state, callback):
46        self._state = state
47        self._callback = callback
48
49    def __call__(self, metadata, error):
50        with self._state.lock:
51            if self._state.exception is None:
52                if self._state.called:
53                    raise RuntimeError(
54                        'AuthMetadataPluginCallback invoked more than once!')
55                else:
56                    self._state.called = True
57            else:
58                raise RuntimeError(
59                    'AuthMetadataPluginCallback raised exception "{}"!'.format(
60                        self._state.exception))
61        if error is None:
62            self._callback(metadata, cygrpc.StatusCode.ok, None)
63        else:
64            self._callback(None, cygrpc.StatusCode.internal,
65                           _common.encode(str(error)))
66
67
68class _Plugin(object):
69
70    def __init__(self, metadata_plugin):
71        self._metadata_plugin = metadata_plugin
72
73    def __call__(self, service_url, method_name, callback):
74        context = _AuthMetadataContext(
75            _common.decode(service_url), _common.decode(method_name))
76        callback_state = _CallbackState()
77        try:
78            self._metadata_plugin(context,
79                                  _AuthMetadataPluginCallback(
80                                      callback_state, callback))
81        except Exception as exception:  # pylint: disable=broad-except
82            _LOGGER.exception(
83                'AuthMetadataPluginCallback "%s" raised exception!',
84                self._metadata_plugin)
85            with callback_state.lock:
86                callback_state.exception = exception
87                if callback_state.called:
88                    return
89            callback(None, cygrpc.StatusCode.internal,
90                     _common.encode(str(exception)))
91
92
93def metadata_plugin_call_credentials(metadata_plugin, name):
94    if name is None:
95        try:
96            effective_name = metadata_plugin.__name__
97        except AttributeError:
98            effective_name = metadata_plugin.__class__.__name__
99    else:
100        effective_name = name
101    return grpc.CallCredentials(
102        cygrpc.MetadataPluginCallCredentials(
103            _Plugin(metadata_plugin), _common.encode(effective_name)))
104