Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/metadata.py: 75.00%

1# Copyright (C) 2022 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 

15 

16from contextvars import ContextVar 

17from typing import Any, Dict, Iterable, Optional, Tuple, Union, cast 

18 

19from grpc.aio import Metadata 

20 

21from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import RequestMetadata, ToolDetails 

22from buildgrid._protos.buildgrid.v2.identity_pb2 import ClientIdentity 

23from buildgrid.server.auth.manager import get_context_client_identity 

24from buildgrid.server.settings import ( 

25 CLIENT_IDENTITY_HEADER_NAME, 

26 REQUEST_METADATA_HEADER_NAME, 

27 REQUEST_METADATA_TOOL_NAME, 

28 REQUEST_METADATA_TOOL_VERSION, 

29) 

30from buildgrid.server.sql.models import ClientIdentityEntry 

31 

32ctx_request_metadata: ContextVar[RequestMetadata] = ContextVar( 

33 "ctx_request_metadata", 

34 default=RequestMetadata( 

35 tool_details=ToolDetails( 

36 tool_name=REQUEST_METADATA_TOOL_NAME, 

37 tool_version=REQUEST_METADATA_TOOL_VERSION, 

38 ), 

39 ), 

40) 

41 

42 

43def metadata_list() -> Tuple[Tuple[str, Union[str, bytes]], ...]: 

44 """Helper function to construct the metadata list from the ContextVar.""" 

45 metadata = ctx_request_metadata.get() 

46 return ((REQUEST_METADATA_HEADER_NAME, metadata.SerializeToString()),) 

47 

48 

49ctx_grpc_request_id: ContextVar[Optional[str]] = ContextVar("grpc_request_id", default=None) 

50 

51 

52def printable_request_metadata(metadata_entries: Any) -> str: 

53 """Given a metadata object, return a human-readable representation 

54 of its `RequestMetadata` entry. 

55 

56 Args: 

57 metadata_entries: tuple of entries obtained from a gRPC context 

58 with, for example, `context.invocation_metadata()`. 

59 

60 Returns: 

61 A string with the metadata contents. 

62 """ 

63 metadata = extract_request_metadata(metadata_entries) 

64 return request_metadata_to_string(metadata) 

65 

66 

67def extract_request_metadata_dict(metadata_entries: Any) -> Dict[str, str]: 

68 metadata = extract_request_metadata(metadata_entries) 

69 return request_metadata_to_dict(metadata) 

70 

71 

72def extract_request_metadata(metadata_entries: Any) -> RequestMetadata: 

73 """Given a list of string tuples, extract the RequestMetadata 

74 header values if they are present. If they were not provided, 

75 returns an empty message. 

76 

77 Args: 

78 metadata_entries: tuple of entries obtained from a gRPC context 

79 with, for example, `context.invocation_metadata()`. 

80 

81 Returns: 

82 A `RequestMetadata` proto. If the metadata is not defined in the 

83 request, the message will be empty. 

84 """ 

85 request_metadata_entry = next( 

86 (entry for entry in metadata_entries if entry.key == REQUEST_METADATA_HEADER_NAME), None 

87 ) 

88 

89 request_metadata = RequestMetadata() 

90 if request_metadata_entry: 

91 request_metadata.ParseFromString(request_metadata_entry.value) 

92 return request_metadata 

93 

94 

95def request_metadata_to_string(request_metadata: RequestMetadata) -> str: 

96 if request_metadata.tool_details: 

97 tool_name = request_metadata.tool_details.tool_name 

98 tool_version = request_metadata.tool_details.tool_version 

99 else: 

100 tool_name = tool_version = "" 

101 

102 return ( 

103 f'tool_name="{tool_name}", tool_version="{tool_version}", ' 

104 f'action_id="{request_metadata.action_id}", ' 

105 f'tool_invocation_id="{request_metadata.tool_invocation_id}", ' 

106 f'correlated_invocations_id="{request_metadata.correlated_invocations_id}", ' 

107 f'action_mnemonic="{request_metadata.action_mnemonic}", ' 

108 f'target_id="{request_metadata.target_id}", ' 

109 f'configuration_id="{request_metadata.configuration_id}"' 

110 ) 

111 

112 

113def request_metadata_to_dict(request_metadata: RequestMetadata) -> Dict[str, str]: 

114 if request_metadata.tool_details: 

115 tool_name = request_metadata.tool_details.tool_name 

116 tool_version = request_metadata.tool_details.tool_version 

117 else: 

118 tool_name = tool_version = "" 

119 

120 return { 

121 "tool_name": tool_name, 

122 "tool_version": tool_version, 

123 "action_id": request_metadata.action_id, 

124 "tool_invocation_id": request_metadata.tool_invocation_id, 

125 "correlated_invocations_id": request_metadata.correlated_invocations_id, 

126 "action_mnemonic": request_metadata.action_mnemonic, 

127 "target_id": request_metadata.target_id, 

128 "configuration_id": request_metadata.configuration_id, 

129 } 

130 

131 

132def extract_client_identity_dict(instance: str, invocation_metadata: Iterable[Tuple[str, Any]]) -> Dict[str, Any]: 

133 client_id = extract_client_identity(instance, invocation_metadata) 

134 return { 

135 "client_id": client_id.id if client_id else "", 

136 "instance": client_id.instance if client_id else "", 

137 "workflow": client_id.workflow if client_id else "", 

138 "actor": client_id.actor if client_id else "", 

139 "subject": client_id.subject if client_id else "", 

140 } 

141 

142 

143def extract_client_identity( 

144 instance: str, invocation_metadata: Iterable[Tuple[str, Any]] 

145) -> Optional[ClientIdentityEntry]: 

146 """Checks for the existence of the client identity in the ClientIdentity 

147 context var. If the context var is not set then extracts 

148 the ClientIdentity from gRPC metadata 

149 

150 Args: 

151 instance (str): the instance where the request was invoked from 

152 invocation_metadata (List[Tuple[str, str]]): grpc metadata 

153 

154 Returns: 

155 Optional[ClientIdentityEntry]: identity of the client if exists 

156 """ 

157 context_client_identity = get_context_client_identity() 

158 if ( 

159 context_client_identity 

160 and context_client_identity.actor 

161 and context_client_identity.subject 

162 and context_client_identity.workflow 

163 ): 

164 return ClientIdentityEntry( 

165 instance=instance, 

166 workflow=context_client_identity.workflow, 

167 actor=context_client_identity.actor, 

168 subject=context_client_identity.subject, 

169 ) 

170 

171 metadata_dict = dict(invocation_metadata) 

172 workflow = metadata_dict.get("x-request-workflow", None) 

173 actor = metadata_dict.get("x-request-actor", None) 

174 subject = metadata_dict.get("x-request-subject", None) 

175 

176 # None or empty strings are invalid 

177 if workflow and actor and subject: 

178 return ClientIdentityEntry(instance=instance, workflow=workflow, actor=actor, subject=subject) 

179 

180 return None 

181 

182 

183def printable_client_identity(instance: str, invocation_metadata: Iterable[Tuple[str, Any]]) -> str: 

184 """Given a metadata object, return a human-readable representation 

185 of its `ClientIdentity` entry. 

186 

187 Args: 

188 instance: REAPI instance name 

189 invocation_metadata: tuple of entries obtained from a gRPC context 

190 with, for example, `context.invocation_metadata()`. 

191 

192 Returns: 

193 A string with the ClientIdentity contents. 

194 """ 

195 client_id = extract_client_identity(instance, invocation_metadata) 

196 return str(client_id) 

197 

198 

199def extract_trailing_client_identity(metadata_entries: Metadata) -> ClientIdentity: 

200 """Given a list of string tuples, extract the ClientIdentity 

201 header values if they are present. If they were not provided, 

202 returns an empty message. 

203 

204 Args: 

205 metadata_entries: Sequence of gRPC trailing metadata headers. 

206 

207 Returns: 

208 A `ClientIdentity` proto. If the metadata is not defined in the 

209 request, the message will be empty. 

210 """ 

211 client_identity_entry = next( 

212 (entry for entry in metadata_entries if entry[0] == CLIENT_IDENTITY_HEADER_NAME), None 

213 ) 

214 

215 client_identity = ClientIdentity() 

216 if client_identity_entry: 

217 client_identity.ParseFromString(cast(bytes, client_identity_entry[1])) 

218 return client_identity