|
| 1 | +""" |
| 2 | +HTTPX client utility with SSL configuration support. |
| 3 | +
|
| 4 | +This module patches httpx.Client.__init__ to automatically apply SSL configuration |
| 5 | +from environment variables. It supports: |
| 6 | +- SSL verification control |
| 7 | +- Custom CA certificates |
| 8 | +- Mutual TLS (mTLS) with client certificates and keys |
| 9 | +
|
| 10 | +The patching is done via monkey patching httpx.Client.__init__, which covers all use cases: |
| 11 | +- Direct Client instantiation: httpx.Client() |
| 12 | +- Convenience methods: httpx.get(), httpx.post(), etc. (they internally use Client) |
| 13 | +
|
| 14 | +No code changes are needed in places that use httpx. |
| 15 | +""" |
| 16 | + |
| 17 | +import base64 |
| 18 | +import ssl |
| 19 | +import tempfile |
| 20 | +from functools import wraps |
| 21 | +from pathlib import Path |
| 22 | +from typing import Any |
| 23 | + |
| 24 | +import httpx |
| 25 | + |
| 26 | +from dify_plugin.config.config import DifyPluginEnv |
| 27 | + |
| 28 | +# Store original method before patching - use getattr to avoid reload issues |
| 29 | +_original_client_init = getattr(httpx.Client.__init__, "__wrapped__", httpx.Client.__init__) |
| 30 | + |
| 31 | + |
| 32 | +def _decode_base64_cert(data: str | None) -> bytes | None: |
| 33 | + """ |
| 34 | + Decode base64 encoded certificate data. |
| 35 | +
|
| 36 | + :param data: Base64 encoded certificate data |
| 37 | + :return: Decoded bytes or None if data is None/empty |
| 38 | + """ |
| 39 | + if not data: |
| 40 | + return None |
| 41 | + try: |
| 42 | + return base64.b64decode(data) |
| 43 | + except Exception as e: |
| 44 | + raise ValueError(f"Failed to decode base64 certificate data: {e}") from e |
| 45 | + |
| 46 | + |
| 47 | +def _create_ssl_context(config: DifyPluginEnv) -> ssl.SSLContext | bool: |
| 48 | + """ |
| 49 | + Create SSL context based on environment configuration. |
| 50 | +
|
| 51 | + :param config: DifyPluginEnv configuration instance |
| 52 | + :return: SSL context, True (verify), or False (no verify) |
| 53 | + """ |
| 54 | + # If SSL verification is disabled, return False |
| 55 | + if not config.HTTP_REQUEST_NODE_SSL_VERIFY: |
| 56 | + return False |
| 57 | + |
| 58 | + # Check if we have custom SSL configuration |
| 59 | + has_ca_cert = bool(config.HTTP_REQUEST_NODE_SSL_CERT_DATA) |
| 60 | + has_client_cert = bool(config.HTTP_REQUEST_NODE_SSL_CLIENT_CERT_DATA) |
| 61 | + has_client_key = bool(config.HTTP_REQUEST_NODE_SSL_CLIENT_KEY_DATA) |
| 62 | + |
| 63 | + # If no custom SSL configuration, use default verification |
| 64 | + if not (has_ca_cert or has_client_cert or has_client_key): |
| 65 | + return True |
| 66 | + |
| 67 | + # Create custom SSL context |
| 68 | + ssl_context = ssl.create_default_context() |
| 69 | + |
| 70 | + # Load custom CA certificate if provided |
| 71 | + if has_ca_cert: |
| 72 | + ca_cert_data = _decode_base64_cert(config.HTTP_REQUEST_NODE_SSL_CERT_DATA) |
| 73 | + if ca_cert_data: |
| 74 | + # Write CA cert to temporary file and load it |
| 75 | + with tempfile.NamedTemporaryFile(mode="wb", suffix=".pem", delete=False) as ca_file: |
| 76 | + ca_file.write(ca_cert_data) |
| 77 | + ca_cert_path = ca_file.name |
| 78 | + try: |
| 79 | + ssl_context.load_verify_locations(cafile=ca_cert_path) |
| 80 | + finally: |
| 81 | + # Clean up temporary file |
| 82 | + Path(ca_cert_path).unlink(missing_ok=True) |
| 83 | + |
| 84 | + # Load client certificate and key for mutual TLS if provided |
| 85 | + if has_client_cert and has_client_key: |
| 86 | + client_cert_data = _decode_base64_cert(config.HTTP_REQUEST_NODE_SSL_CLIENT_CERT_DATA) |
| 87 | + client_key_data = _decode_base64_cert(config.HTTP_REQUEST_NODE_SSL_CLIENT_KEY_DATA) |
| 88 | + |
| 89 | + if client_cert_data and client_key_data: |
| 90 | + # Write client cert and key to temporary files |
| 91 | + with ( |
| 92 | + tempfile.NamedTemporaryFile(mode="wb", suffix=".pem", delete=False) as cert_file, |
| 93 | + tempfile.NamedTemporaryFile(mode="wb", suffix=".pem", delete=False) as key_file, |
| 94 | + ): |
| 95 | + cert_file.write(client_cert_data) |
| 96 | + key_file.write(client_key_data) |
| 97 | + cert_path = cert_file.name |
| 98 | + key_path = key_file.name |
| 99 | + |
| 100 | + try: |
| 101 | + ssl_context.load_cert_chain(certfile=cert_path, keyfile=key_path) |
| 102 | + finally: |
| 103 | + # Clean up temporary files |
| 104 | + Path(cert_path).unlink(missing_ok=True) |
| 105 | + Path(key_path).unlink(missing_ok=True) |
| 106 | + |
| 107 | + return ssl_context |
| 108 | + |
| 109 | + |
| 110 | +@wraps(_original_client_init) |
| 111 | +def _patched_client_init(self, *args: Any, **kwargs: Any) -> None: |
| 112 | + """ |
| 113 | + Patched httpx.Client.__init__ that injects SSL configuration. |
| 114 | +
|
| 115 | + This single patch covers all httpx usage patterns: |
| 116 | + - httpx.Client() - direct instantiation |
| 117 | + - httpx.get(), httpx.post(), etc. - these internally create Client instances |
| 118 | + """ |
| 119 | + if "verify" not in kwargs: |
| 120 | + config = DifyPluginEnv() |
| 121 | + kwargs["verify"] = _create_ssl_context(config) |
| 122 | + return _original_client_init(self, *args, **kwargs) |
| 123 | + |
| 124 | + |
| 125 | +# Apply monkey patch to httpx.Client.__init__ |
| 126 | +# This single patch is sufficient because: |
| 127 | +# - httpx.get/post/put/delete/patch/head/options all call httpx.request() |
| 128 | +# - httpx.request() creates a Client instance internally with the verify parameter |
| 129 | +# - So patching Client.__init__ catches all cases |
| 130 | +httpx.Client.__init__ = _patched_client_init |
0 commit comments