feat: Adding new samples for the python callouts (#70)

* adding publickey for the jwt sample * adding jwt to requirements and jwt example with RS256 and HS256 * adding tests * removing hmac, refactoring to consider the Authorization Bearer from the header instead of metadata * removing hmac test, adjusting the tests for the rs256 * fix: fixing failing tests, changing async_mode to observability_mode=False based on the envoy proto changes * removing observability_mode and changing the jwt token * changing the jwt token * refactoring the success test and adding cryptography to the req-test for the jwt token generation
GoogleCloudPlatform · Apr 11, 2024 · 4ac042f · 4ac042f
1 parent f6b501a
commit 4ac042f
Show file tree

Hide file tree

Showing 11 changed files with 235 additions and 12 deletions.
diff --git a/callouts/python/extproc/example/jwt_auth/Dockerfile b/callouts/python/extproc/example/jwt_auth/Dockerfile
@@ -0,0 +1,22 @@
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+FROM service-callout-common-python
+
+# Copy over example specific files.
+COPY extproc/example/jwt_auth/*.py ./
+
+# Set up communication ports.
+EXPOSE 443
+EXPOSE 80
+EXPOSE 8080
+
+# Start the service.
+CMD [ "/usr/bin/python3", "-um", "service_callout_example" ]
diff --git a/callouts/python/extproc/example/jwt_auth/__init__.py b/callouts/python/extproc/example/jwt_auth/__init__.py
diff --git a/callouts/python/extproc/example/jwt_auth/service_callout_example.py b/callouts/python/extproc/example/jwt_auth/service_callout_example.py
@@ -0,0 +1,69 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import jwt
+from jwt.exceptions import InvalidTokenError
+
+from grpc import ServicerContext
+from envoy.service.ext_proc.v3 import external_processor_pb2 as service_pb2
+from extproc.service import callout_server
+from extproc.service import callout_tools
+
+def extract_jwt_token(request_headers):
+  jwt_token = next((header.raw_value.decode('utf-8')
+                   for header in request_headers.headers.headers
+                   if header.key == 'Authorization'), None)
+  extracted_jwt = jwt_token.split(' ')[1] if jwt_token and ' ' in jwt_token else jwt_token
+  return extracted_jwt
+
+def validate_jwt_token(key, request_headers, algorithm):
+  jwt_token = extract_jwt_token(request_headers)
+  try:
+    decoded = jwt.decode(jwt_token, key, algorithms=[algorithm])
+    logging.info('Approved - Decoded Values: %s', decoded)
+    return decoded
+  except InvalidTokenError:
+    return None
+
+class CalloutServerExample(callout_server.CalloutServer):
+  """Example callout server.
+
+  For request header callouts we provide a mutation to add multiple headers
+  based on the decoded fields for example '{decoded-name: John Doe}', and to
+  clear the route cache if the JWT Authorization is valid.
+  A valid token example value can be found below.
+
+  Valid Token for RS256:
+  eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWUsImlhdCI6MTcxMjE3MzQ2MSwiZXhwIjoyMDc1NjU4MjYxfQ.Vv-Lwn1z8BbVBGm-T1EKxv6T3XKCeRlvRrRmdu8USFdZUoSBK_aThzwzM2T8hlpReYsX9YFdJ3hMfq6OZTfHvfPLXvAt7iSKa03ZoPQzU8bRGzYy8xrb0ZQfrejGfHS5iHukzA8vtI2UAJ_9wFQiY5_VGHOBv9116efslbg-_gItJ2avJb0A0yr5uUwmE336rYEwgm4DzzfnTqPt8kcJwkONUsjEH__mePrva1qDT4qtfTPQpGa35TW8n9yZqse3h1w3xyxUfJd3BlDmoz6pQp2CvZkhdQpkWA1bnwpdqSDC7bHk4tYX6K5Q19na-2ff7gkmHZHJr0G9e_vAhQiE5w
+
+  """
+
+  def on_request_headers(
+      self, headers: service_pb2.HttpHeaders,
+      context: ServicerContext):
+    """Custom processor on request headers."""
+
+    decoded = validate_jwt_token(self.public_key, headers, "RS256")
+
+    if decoded:
+      decoded_items = [('decoded-' + key, str(value)) for key, value in decoded.items()]
+      return callout_tools.add_header_mutation(add=decoded_items, clear_route_cache=True)
+    else:
+      callout_tools.deny_request(context, 'Authorization token is invalid')
+
+
+if __name__ == '__main__':
+  # Run the gRPC service
+  CalloutServerExample(insecure_address=('0.0.0.0', 8080)).run()
diff --git a/callouts/python/extproc/service/callout_server.py b/callouts/python/extproc/service/callout_server.py
@@ -101,6 +101,7 @@ def __init__(
       default_ip: str | None = None,
       cert_path: str = './extproc/ssl_creds/localhost.crt',
       cert_key_path: str = './extproc/ssl_creds/localhost.key',
+      public_key_path: str = './extproc/ssl_creds/publickey.pem',
       server_thread_count: int = 2,
   ):
     self._setup = False
@@ -138,6 +139,10 @@ def __init__(
     with open(cert_key_path, 'rb') as file:
       self.cert_key = file.read()
       file.close()
+    self.public_key_path = public_key_path
+    with open(public_key_path, 'rb') as file:
+      self.public_key = file.read()
+      file.close()
 
   def run(self):
     """Start all requested servers and listen for new connections; blocking."""

diff --git a/callouts/python/extproc/ssl_creds/publickey.pem b/callouts/python/extproc/ssl_creds/publickey.pem
@@ -0,0 +1,9 @@
+-----BEGIN PUBLIC KEY-----
+MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEArQltE47oJyFLvARFSGL3
+TD/v9hhDUA42YWfHikeQuraAxm4I7A7bfxsSEgIOHIO9HyL8tF6JZxj2aHb3xwzO
+JSe8hwTnWQ/E10KBOqJEex2bi1CdtDNKV4FqrUXsCkVI7mLEg1wZ8y4mcNXXL/Zx
+qaPgamvDaJ6YhkUaRJLIRQ+k/SlPU7P8fsDEASK2tho9FDIe/Yz4VLsXGAEANCUB
+or5zrXcLj4rpqZ1E/I/IerVoOZdNRLAmViXnAhfOSpxqSRFWD8j6gahrMtJxkvaZ
+RqLTlhfO/o8lELfig+t8QGgZk66POpigwyhCuCb841XjBs9lG1XCz4omZxK0xWol
+yQIDAQAB
+-----END PUBLIC KEY-----
diff --git a/callouts/python/extproc/tests/basic_grpc_test.py b/callouts/python/extproc/tests/basic_grpc_test.py
@@ -159,29 +159,29 @@ def test_basic_server_capabilites(self, server: CalloutServerTest) -> None:
       headers = HttpHeaders(end_of_stream=False)
       end_headers = HttpHeaders(end_of_stream=True)
 
-      value = make_request(stub, request_body=body, async_mode=False)
+      value = make_request(stub, request_body=body)
       assert value.HasField('request_body')
       assert value.request_body == add_body_mutation(body='-added-body')
 
-      value = make_request(stub, response_body=body, async_mode=False)
+      value = make_request(stub, response_body=body)
       assert value.HasField('response_body')
       assert value.response_body == add_body_mutation(body='new-body',
                                                       clear_body=True)
 
-      value = make_request(stub, response_headers=headers, async_mode=False)
+      value = make_request(stub, response_headers=headers)
       assert value.HasField('response_headers')
       assert value.response_headers == add_header_mutation(
           add=[('hello', 'service-extensions')])
 
-      value = make_request(stub, request_headers=headers, async_mode=False)
+      value = make_request(stub, request_headers=headers)
       assert value.HasField('request_headers')
       assert value.request_headers == add_header_mutation(
           add=[(':host', 'service-extensions.com'), (':path', '/'),
                ('header-request', 'request')],
           clear_route_cache=True,
           remove=['foo'])
 
-      make_request(stub, request_headers=end_headers, async_mode=False)
+      make_request(stub, request_headers=end_headers)
       channel.close()
 
   @pytest.mark.parametrize('server', [_local_test_args], indirect=True)
@@ -242,8 +242,7 @@ def test_custom_server_config() -> None:
     with grpc.insecure_channel(f'{ip}:{insecure_port}') as channel:
       stub = ExternalProcessorStub(channel)
       value = make_request(stub,
-                           response_headers=HttpHeaders(end_of_stream=True),
-                           async_mode=False)
+                           response_headers=HttpHeaders(end_of_stream=True))
       assert value.HasField('response_headers')
       assert value.response_headers == add_header_mutation(
           add=[('hello', 'service-extensions')])

diff --git a/callouts/python/extproc/tests/jwt_auth_test.py b/callouts/python/extproc/tests/jwt_auth_test.py
@@ -0,0 +1,117 @@
+# Copyright 2024 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import print_function
+
+import datetime
+import re
+
+import jwt
+from envoy.config.core.v3.base_pb2 import HeaderMap
+from envoy.config.core.v3.base_pb2 import HeaderValue
+from envoy.service.ext_proc.v3 import external_processor_pb2 as service_pb2
+from envoy.service.ext_proc.v3 import external_processor_pb2_grpc as service_pb2_grpc
+
+import grpc
+import pytest
+
+from extproc.example.jwt_auth.service_callout_example import (
+    CalloutServerExample as CalloutServerTest)
+from extproc.service.callout_tools import add_header_mutation
+from extproc.tests.basic_grpc_test import (
+    make_request,
+    setup_server,
+    get_insecure_channel,
+    insecure_kwargs,
+)
+
+
+# Import the setup server test fixture.
+_ = setup_server
+_local_test_args = {"kwargs": insecure_kwargs, "test_class": CalloutServerTest}
+
+@pytest.mark.parametrize('server', [_local_test_args], indirect=True)
+def test_jwt_auth_rs256_failure(server: CalloutServerTest) -> None:
+  with get_insecure_channel(server) as channel:
+    stub = service_pb2_grpc.ExternalProcessorStub(channel)
+
+    # Construct the HeaderMap
+    header_map = HeaderMap()
+    header_value = HeaderValue(key="Authorization", raw_value=b"")
+    header_map.headers.extend([header_value])
+
+    # Construct HttpHeaders with the HeaderMap
+    request_headers = service_pb2.HttpHeaders(headers=header_map,
+                                              end_of_stream=True)
+
+    # Use request_headers in the request
+    with pytest.raises(grpc.RpcError) as e:
+      make_request(stub, request_headers=request_headers)
+    assert e.value.code() == grpc.StatusCode.PERMISSION_DENIED
+
+@pytest.mark.parametrize('server', [_local_test_args], indirect=True)
+def test_jwt_auth_rs256_success(server: CalloutServerTest) -> None:
+    with get_insecure_channel(server) as channel:
+        stub = service_pb2_grpc.ExternalProcessorStub(channel)
+
+        # Load the private key
+        private_key_path = './extproc/ssl_creds/localhost.key'
+        with open(private_key_path, 'r') as key_file:
+            private_key = key_file.read()
+
+        # Define the payload for the JWT
+        payload = {
+            "sub": "1234567890",
+            "name": "John Doe",
+            "admin": True,
+            "iat": datetime.datetime.utcnow(),
+            "exp": datetime.datetime.utcnow() + datetime.timedelta(hours=1)
+        }
+
+        # Generate the JWT token
+        jwt_token = jwt.encode(payload, private_key, algorithm="RS256")
+
+        # Authorization header value
+        authorization_header_value = f"Bearer {jwt_token}"
+
+        # Construct the HeaderMap
+        header_map = HeaderMap()
+        header_value = HeaderValue(key="Authorization", raw_value=bytes(authorization_header_value, 'utf-8'))
+        header_map.headers.extend([header_value])
+
+        # Construct HttpHeaders with the HeaderMap
+        request_headers = service_pb2.HttpHeaders(headers=header_map, end_of_stream=True)
+
+        # Construct the decoded items list from the payload
+        decoded_items = [(f'decoded-{key}', str(value)) for key, value in payload.items() if key != 'exp' and key != 'iat']
+        # Adding formatted 'iat' and 'exp' to match the test format
+        decoded_items.extend([
+            ('decoded-iat', str(int(payload['iat'].timestamp()))),
+            ('decoded-exp', str(int(payload['exp'].timestamp())))
+        ])
+
+        value = make_request(stub, request_headers=request_headers)
+        assert value.HasField('request_headers')
+        # Instead of directly comparing the full response, check the presence and basic validation of decoded items
+        assert 'header_mutation' in str(value)
+        for key, expected_value in decoded_items:
+          # Check presence of key
+          assert key in str(value)
+          # For 'iat' and 'exp', check if it matches the pattern since the value will be different
+          if key in ['decoded-iat', 'decoded-exp']:
+            pattern = rf'{key}"\s*raw_value:\s*"\d+"'
+            assert re.search(pattern, str(value)), f"{key} does not match expected pattern"
+          else:
+            # For other keys, check the exact value
+            pattern = rf'{key}"\s*raw_value:\s*"{expected_value}"'
+            assert re.search(pattern, str(value)), f"{key} value {expected_value} not found"
diff --git a/callouts/python/extproc/tests/normalize_header_test.py b/callouts/python/extproc/tests/normalize_header_test.py
@@ -40,11 +40,11 @@ def test_normalize_header(server: CalloutServerTest) -> None:
     headers = service_pb2.HttpHeaders(end_of_stream=False)
     end_headers = service_pb2.HttpHeaders(end_of_stream=True)
 
-    value = make_request(stub, request_headers=headers, async_mode=False)
+    value = make_request(stub, request_headers=headers)
     assert value.HasField('request_headers')
     assert value.request_headers == callout_tools.normalize_header_mutation(
         headers=headers,
         clear_route_cache=True,
     )
 
-    make_request(stub, request_headers=end_headers, async_mode=False)
+    make_request(stub, request_headers=end_headers)
diff --git a/callouts/python/extproc/tests/update_header_test.py b/callouts/python/extproc/tests/update_header_test.py
@@ -45,19 +45,19 @@ def test_append_action(server: CalloutServerTest) -> None:
     headers = service_pb2.HttpHeaders(end_of_stream=False)
     end_headers = service_pb2.HttpHeaders(end_of_stream=True)
 
-    value = make_request(stub, response_headers=headers, async_mode=False)
+    value = make_request(stub, response_headers=headers)
     assert value.HasField('response_headers')
     assert value.response_headers == callout_tools.add_header_mutation(
         add=[('header-response', 'response-new-value')],
         append_action=HeaderValueOption.HeaderAppendAction.
         OVERWRITE_IF_EXISTS_OR_ADD)
 
-    value = make_request(stub, request_headers=headers, async_mode=False)
+    value = make_request(stub, request_headers=headers)
     assert value.HasField('request_headers')
     assert value.request_headers == callout_tools.add_header_mutation(
         add=[('header-request', 'request-new-value')],
         append_action=HeaderValueOption.HeaderAppendAction.
         OVERWRITE_IF_EXISTS_OR_ADD,
         clear_route_cache=True)
 
-    make_request(stub, request_headers=end_headers, async_mode=False)
+    make_request(stub, request_headers=end_headers)
diff --git a/callouts/python/requirements-test.txt b/callouts/python/requirements-test.txt
@@ -1 +1,2 @@
 pytest==7.4.2
+cryptography==42.0.5
diff --git a/callouts/python/requirements.txt b/callouts/python/requirements.txt
@@ -1,2 +1,3 @@
 grpcio==1.59.0
 google-cloud-logging==3.9.0
+pyjwt==2.8.0