Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/decorators/authorize.py: 100.00%

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

15import functools 

16import inspect 

17from typing import Any, Callable, Iterator, TypeVar, cast 

18 

19import grpc 

20 

21from buildgrid.server.auth.manager import authorize_request 

22from buildgrid.server.context import current_instance 

23from buildgrid.server.exceptions import InvalidArgumentError 

24from buildgrid.server.metrics_names import METRIC 

25from buildgrid.server.metrics_utils import timer 

26 

27Func = TypeVar("Func", bound=Callable) # type: ignore[type-arg] 

28 

29 

30def authorize(f: Func) -> Func: 

31 @functools.wraps(f) 

32 def server_stream_wrapper(self: Any, message: Any, context: grpc.ServicerContext) -> Iterator[Any]: 

33 try: 

34 with timer(METRIC.RPC.AUTH_DURATION): 

35 authorize_request(context, current_instance(), f.__name__) 

36 yield from f(self, message, context) 

37 except InvalidArgumentError as e: 

38 context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(e)) 

39 

40 @functools.wraps(f) 

41 def server_unary_wrapper(self: Any, message: Any, context: grpc.ServicerContext) -> Any: 

42 try: 

43 with timer(METRIC.RPC.AUTH_DURATION): 

44 authorize_request(context, current_instance(), f.__name__) 

45 return f(self, message, context) 

46 except InvalidArgumentError as e: 

47 context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(e)) 

48 

49 if inspect.isgeneratorfunction(f): 

50 return cast(Func, server_stream_wrapper) 

51 return cast(Func, server_unary_wrapper)