diff --git a/docs/Config_en.md b/docs/Config_en.md index d351c724..999f12c3 100644 --- a/docs/Config_en.md +++ b/docs/Config_en.md @@ -132,6 +132,11 @@ value is the agent's configuration. A single agent configuration requires the fo - `module`: defines the corresponding agent client module. - `parameters`: defines the parameters to be passed to the corresponding module. +`HTTPAgent` verifies HTTPS certificates by default. Its optional `verify` parameter +accepts `true` (the default) or a path to a CA bundle for a private certificate +authority. `verify: false` explicitly disables certificate verification for that +agent's requests; it does not change other HTTP sessions. + ### tasks The `tasks` directory contains all task configuration files. The `task_assembly.yaml` collects all task definitions. If diff --git a/src/client/agents/http_agent.py b/src/client/agents/http_agent.py index 677f668c..4c2e9242 100644 --- a/src/client/agents/http_agent.py +++ b/src/client/agents/http_agent.py @@ -1,48 +1,11 @@ -import contextlib import time -import warnings import requests -from urllib3.exceptions import InsecureRequestWarning from src.typings import * from src.utils import * from ..agent import AgentClient -old_merge_environment_settings = requests.Session.merge_environment_settings - - -@contextlib.contextmanager -def no_ssl_verification(): - opened_adapters = set() - - def merge_environment_settings(self, url, proxies, stream, verify, cert): - # Verification happens only once per connection so we need to close - # all the opened adapters once we're done. Otherwise, the effects of - # verify=False persist beyond the end of this context manager. - opened_adapters.add(self.get_adapter(url)) - - settings = old_merge_environment_settings(self, url, proxies, stream, verify, cert) - settings['verify'] = False - - return settings - - requests.Session.merge_environment_settings = merge_environment_settings - - try: - with warnings.catch_warnings(): - warnings.simplefilter('ignore', InsecureRequestWarning) - yield - finally: - requests.Session.merge_environment_settings = old_merge_environment_settings - - for adapter in opened_adapters: - try: - adapter.close() - except: - pass - - class Prompter: @staticmethod def get_prompter(prompter: Union[Dict[str, Any], None]): @@ -170,11 +133,13 @@ def __init__( headers=None, return_format="{response}", prompter=None, + verify=True, **kwargs, ) -> None: super().__init__(**kwargs) self.url = url self.proxies = proxies or {} + self.verify = verify self.headers = headers or {} self.body = body or {} self.return_format = return_format @@ -190,10 +155,10 @@ def inference(self, history: List[dict]) -> str: try: body = self.body.copy() body.update(self._handle_history(history)) - with no_ssl_verification(): - resp = requests.post( - self.url, json=body, headers=self.headers, proxies=self.proxies, timeout=120 - ) + resp = requests.post( + self.url, json=body, headers=self.headers, proxies=self.proxies, + timeout=120, verify=self.verify, + ) # print(resp.status_code, resp.text) if resp.status_code != 200: # print(resp.text) diff --git a/tests/test_http_agent_tls.py b/tests/test_http_agent_tls.py new file mode 100644 index 00000000..f6ff2933 --- /dev/null +++ b/tests/test_http_agent_tls.py @@ -0,0 +1,93 @@ +import json +from pathlib import Path +import shutil +import ssl +import subprocess +import tempfile +import threading +import unittest +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from unittest.mock import patch + +import requests + +from src.client.agents.http_agent import HTTPAgent + + +@unittest.skipUnless(shutil.which("openssl"), "openssl is required for the local TLS fixture") +class HTTPAgentTLS(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.directory = tempfile.TemporaryDirectory() + cls.certificate = str(Path(cls.directory.name) / "certificate.pem") + key = str(Path(cls.directory.name) / "key.pem") + subprocess.run( + ["openssl", "req", "-x509", "-newkey", "rsa:2048", "-nodes", + "-keyout", key, "-out", cls.certificate, "-days", "1", + "-subj", "/CN=localhost", "-addext", "subjectAltName=DNS:localhost,IP:127.0.0.1"], + check=True, capture_output=True, + ) + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *args): + pass + + def do_POST(self): + body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + cls.bodies.append(body) + payload = json.dumps({"reply": "fixture reply"}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + cls.bodies = [] + cls.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.load_cert_chain(cls.certificate, key) + cls.server.socket = context.wrap_socket(cls.server.socket, server_side=True) + cls.worker = threading.Thread(target=cls.server.serve_forever, daemon=True) + cls.worker.start() + cls.url = "https://127.0.0.1:%s/chat" % cls.server.server_port + + @classmethod + def tearDownClass(cls): + cls.server.shutdown() + cls.server.server_close() + cls.worker.join() + cls.directory.cleanup() + + def test_untrusted_certificate_is_rejected_by_default(self): + before = len(self.bodies) + agent = HTTPAgent(url=self.url, return_format="{response[reply]}") + with patch("src.client.agents.http_agent.time.sleep"): + with self.assertRaisesRegex(Exception, "Failed"): + agent.inference([{"role": "user", "content": "Hello"}]) + self.assertEqual(len(self.bodies), before) + + def test_custom_ca_is_scoped_to_this_request(self): + original = requests.Session.merge_environment_settings + agent = HTTPAgent( + url=self.url, verify=self.certificate, return_format="{response[reply]}", + prompter={"name": "role_content_dict", "args": {"agent_role": "assistant"}}, + ) + self.assertEqual(agent.inference([ + {"role": "user", "content": "Hello"}, + {"role": "agent", "content": "Previous answer"}, + ]), "fixture reply") + self.assertEqual(self.bodies[-1]["messages"][-1]["role"], "assistant") + self.assertIs(requests.Session.merge_environment_settings, original) + with self.assertRaises(requests.exceptions.SSLError): + requests.post(self.url, json={}, timeout=5) + + def test_explicit_opt_out_does_not_disable_other_sessions(self): + agent = HTTPAgent(url=self.url, verify=False, return_format="{response[reply]}") + with self.assertWarns(requests.packages.urllib3.exceptions.InsecureRequestWarning): + self.assertEqual(agent.inference([{"role": "user", "content": "Hello"}]), "fixture reply") + with self.assertRaises(requests.exceptions.SSLError): + requests.post(self.url, json={}, timeout=5) + + +if __name__ == "__main__": + unittest.main()