Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# Copyright (C) 2019 Bloomberg LP 

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 

15from collections import namedtuple 

16import grpc 

17from urllib.parse import urlparse 

18 

19from buildgrid.client.authentication import AuthMetadataClientInterceptor 

20from buildgrid.client.authentication import load_channel_authorization_token 

21from buildgrid.client.authentication import load_tls_channel_credentials 

22from buildgrid._exceptions import InvalidArgumentError 

23from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2 

24from buildgrid.settings import REQUEST_METADATA_HEADER_NAME 

25from buildgrid.settings import REQUEST_METADATA_TOOL_NAME, REQUEST_METADATA_TOOL_VERSION 

26 

27 

28def setup_channel(remote_url, auth_token=None, 

29 client_key=None, client_cert=None, server_cert=None, 

30 action_id=None, tool_invocation_id=None, 

31 correlated_invocations_id=None): 

32 """Creates a new gRPC client communication chanel. 

33 

34 If `remote_url` does not specifies a port number, defaults 50051. 

35 

36 Args: 

37 remote_url (str): URL for the remote, including port and protocol. 

38 auth_token (str): Authorization token file path. 

39 server_cert(str): TLS certificate chain file path. 

40 client_key (str): TLS root certificate file path. 

41 client_cert (str): TLS private key file path. 

42 action_id (str): Action identifier to which the request belongs to. 

43 tool_invocation_id (str): Identifier for a related group of Actions. 

44 correlated_invocations_id (str): Identifier that ties invocations together. 

45 

46 Returns: 

47 Channel: Client Channel to be used in order to access the server 

48 at `remote_url`. 

49 

50 Raises: 

51 InvalidArgumentError: On any input parsing error. 

52 """ 

53 url = urlparse(remote_url) 

54 remote = f'{url.hostname}:{url.port or 50051}' 

55 details = None, None, None 

56 

57 if url.scheme == 'http': 

58 channel = grpc.insecure_channel(remote) 

59 

60 elif url.scheme == 'https': 

61 credentials, details = load_tls_channel_credentials(client_key, client_cert, server_cert) 

62 if not credentials: 

63 raise InvalidArgumentError("Given TLS details (or defaults) could be loaded") 

64 

65 channel = grpc.secure_channel(remote, credentials) 

66 

67 else: 

68 raise InvalidArgumentError("Given remote does not specify a protocol") 

69 

70 request_metadata_interceptor = RequestMetadataInterceptor( 

71 action_id=action_id, 

72 tool_invocation_id=tool_invocation_id, 

73 correlated_invocations_id=correlated_invocations_id) 

74 

75 channel = grpc.intercept_channel(channel, request_metadata_interceptor) 

76 

77 if auth_token is not None: 

78 token = load_channel_authorization_token(auth_token) 

79 if not token: 

80 raise InvalidArgumentError("Given authorization token could be loaded") 

81 

82 auth_interceptor = AuthMetadataClientInterceptor(auth_token=token) 

83 

84 channel = grpc.intercept_channel(channel, auth_interceptor) 

85 

86 return channel, details 

87 

88 

89class RequestMetadataInterceptor(grpc.UnaryUnaryClientInterceptor, 

90 grpc.UnaryStreamClientInterceptor, 

91 grpc.StreamUnaryClientInterceptor, 

92 grpc.StreamStreamClientInterceptor): 

93 

94 def __init__(self, action_id=None, tool_invocation_id=None, 

95 correlated_invocations_id=None): 

96 """Appends optional `RequestMetadata` header values to each call. 

97 

98 Args: 

99 action_id (str): Action identifier to which the request belongs to. 

100 tool_invocation_id (str): Identifier for a related group of Actions. 

101 correlated_invocations_id (str): Identifier that ties invocations together. 

102 """ 

103 self._action_id = action_id 

104 self._tool_invocation_id = tool_invocation_id 

105 self._correlated_invocations_id = correlated_invocations_id 

106 

107 self.__header_field_name = REQUEST_METADATA_HEADER_NAME 

108 self.__header_field_value = self._request_metadata() 

109 

110 # --- Public API --- 

111 def intercept_unary_unary(self, continuation, client_call_details, request): 

112 new_details = self._amend_call_details(client_call_details) 

113 

114 return continuation(new_details, request) 

115 

116 def intercept_unary_stream(self, continuation, client_call_details, request): 

117 new_details = self._amend_call_details(client_call_details) 

118 

119 return continuation(new_details, request) 

120 

121 def intercept_stream_unary(self, continuation, client_call_details, request_iterator): 

122 new_details = self._amend_call_details(client_call_details) 

123 

124 return continuation(new_details, request_iterator) 

125 

126 def intercept_stream_stream(self, continuation, client_call_details, request_iterator): 

127 new_details = self._amend_call_details(client_call_details) 

128 

129 return continuation(new_details, request_iterator) 

130 

131 # --- Private API --- 

132 def _request_metadata(self): 

133 """Creates a serialized RequestMetadata entry to attach to a gRPC 

134 call header. Arguments should be of type str or None. 

135 """ 

136 request_metadata = remote_execution_pb2.RequestMetadata() 

137 request_metadata.tool_details.tool_name = REQUEST_METADATA_TOOL_NAME 

138 request_metadata.tool_details.tool_version = REQUEST_METADATA_TOOL_VERSION 

139 

140 if self._action_id: 

141 request_metadata.action_id = self._action_id 

142 if self._tool_invocation_id: 

143 request_metadata.tool_invocation_id = self._tool_invocation_id 

144 if self._correlated_invocations_id: 

145 request_metadata.correlated_invocations_id = self._correlated_invocations_id 

146 

147 return request_metadata.SerializeToString() 

148 

149 def _amend_call_details(self, client_call_details): 

150 if client_call_details.metadata is not None: 

151 new_metadata = list(client_call_details.metadata) 

152 else: 

153 new_metadata = [] 

154 

155 new_metadata.append((self.__header_field_name, 

156 self.__header_field_value)) 

157 

158 class _ClientCallDetails( 

159 namedtuple('_ClientCallDetails', 

160 ('method', 'timeout', 'credentials', 'metadata',)), 

161 grpc.ClientCallDetails): 

162 pass 

163 

164 return _ClientCallDetails(client_call_details.method, 

165 client_call_details.timeout, 

166 client_call_details.credentials, 

167 new_metadata)