Add SSL/TLS support for HTTP and gRPC servers (#18973)
Co-authored-by: guys@spotify.com
This commit is contained in:
Kangyan-Zhou
@@ -1,4 +1,5 @@
|
||||
import json
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
@@ -266,5 +267,185 @@ class TestPortArgs(unittest.TestCase):
|
||||
self.assertIn("expected ':' after ']'", str(context.exception))
|
||||
|
||||
|
||||
class TestSSLArgs(unittest.TestCase):
|
||||
def test_default_ssl_fields_are_none(self):
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
self.assertIsNone(server_args.ssl_keyfile)
|
||||
self.assertIsNone(server_args.ssl_certfile)
|
||||
self.assertIsNone(server_args.ssl_ca_certs)
|
||||
self.assertIsNone(server_args.ssl_keyfile_password)
|
||||
|
||||
def test_ssl_keyfile_without_certfile_raises(self):
|
||||
with self.assertRaises(ValueError) as context:
|
||||
ServerArgs(model_path="dummy", ssl_keyfile="key.pem")
|
||||
self.assertIn("--ssl-certfile", str(context.exception))
|
||||
|
||||
def test_ssl_certfile_without_keyfile_raises(self):
|
||||
with self.assertRaises(ValueError) as context:
|
||||
ServerArgs(model_path="dummy", ssl_certfile="cert.pem")
|
||||
self.assertIn("--ssl-keyfile", str(context.exception))
|
||||
|
||||
@patch("os.path.isfile", return_value=True)
|
||||
def test_ssl_both_keyfile_and_certfile_accepted(self, _mock_isfile):
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy", ssl_keyfile="key.pem", ssl_certfile="cert.pem"
|
||||
)
|
||||
self.assertEqual(server_args.ssl_keyfile, "key.pem")
|
||||
self.assertEqual(server_args.ssl_certfile, "cert.pem")
|
||||
|
||||
def test_url_returns_http_without_ssl(self):
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
self.assertTrue(server_args.url().startswith("http://"))
|
||||
|
||||
def test_url_rewrites_all_interfaces_to_loopback(self):
|
||||
server_args = ServerArgs(model_path="dummy", host="0.0.0.0")
|
||||
self.assertEqual(server_args.url(), "http://127.0.0.1:30000")
|
||||
|
||||
def test_url_rewrites_empty_host_to_loopback(self):
|
||||
server_args = ServerArgs(model_path="dummy", host="")
|
||||
self.assertEqual(server_args.url(), "http://127.0.0.1:30000")
|
||||
|
||||
def test_url_rewrites_ipv6_all_interfaces_to_loopback(self):
|
||||
server_args = ServerArgs(model_path="dummy", host="::")
|
||||
self.assertEqual(server_args.url(), "http://[::1]:30000")
|
||||
|
||||
@patch("os.path.isfile", return_value=True)
|
||||
def test_url_returns_https_with_ssl(self, _mock_isfile):
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy", ssl_keyfile="key.pem", ssl_certfile="cert.pem"
|
||||
)
|
||||
self.assertTrue(server_args.url().startswith("https://"))
|
||||
|
||||
@patch("os.path.isfile", return_value=True)
|
||||
def test_ssl_cli_args_parsed(self, _mock_isfile):
|
||||
server_args = prepare_server_args(
|
||||
[
|
||||
"--model-path",
|
||||
"dummy",
|
||||
"--ssl-keyfile",
|
||||
"key.pem",
|
||||
"--ssl-certfile",
|
||||
"cert.pem",
|
||||
"--ssl-ca-certs",
|
||||
"ca.pem",
|
||||
"--ssl-keyfile-password",
|
||||
"secret",
|
||||
]
|
||||
)
|
||||
self.assertEqual(server_args.ssl_keyfile, "key.pem")
|
||||
self.assertEqual(server_args.ssl_certfile, "cert.pem")
|
||||
self.assertEqual(server_args.ssl_ca_certs, "ca.pem")
|
||||
self.assertEqual(server_args.ssl_keyfile_password, "secret")
|
||||
|
||||
def test_ssl_verify_without_ssl(self):
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
self.assertIs(server_args.ssl_verify(), True)
|
||||
|
||||
@patch("os.path.isfile", return_value=True)
|
||||
def test_ssl_verify_with_ssl_no_ca(self, _mock_isfile):
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy", ssl_keyfile="key.pem", ssl_certfile="cert.pem"
|
||||
)
|
||||
self.assertIs(server_args.ssl_verify(), False)
|
||||
|
||||
@patch("os.path.isfile", return_value=True)
|
||||
def test_ssl_verify_with_ssl_and_ca(self, _mock_isfile):
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy",
|
||||
ssl_keyfile="key.pem",
|
||||
ssl_certfile="cert.pem",
|
||||
ssl_ca_certs="ca.pem",
|
||||
)
|
||||
self.assertEqual(server_args.ssl_verify(), "ca.pem")
|
||||
|
||||
def test_ssl_ca_certs_without_certfile_raises(self):
|
||||
with self.assertRaises(ValueError) as context:
|
||||
ServerArgs(model_path="dummy", ssl_ca_certs="ca.pem")
|
||||
self.assertIn("--ssl-ca-certs", str(context.exception))
|
||||
|
||||
def test_ssl_keyfile_password_without_certfile_raises(self):
|
||||
with self.assertRaises(ValueError) as context:
|
||||
ServerArgs(model_path="dummy", ssl_keyfile_password="secret")
|
||||
self.assertIn("--ssl-keyfile-password", str(context.exception))
|
||||
|
||||
def test_ssl_keyfile_not_found_raises(self):
|
||||
with self.assertRaises(ValueError) as context:
|
||||
ServerArgs(
|
||||
model_path="dummy",
|
||||
ssl_keyfile="/nonexistent/key.pem",
|
||||
ssl_certfile="/nonexistent/cert.pem",
|
||||
)
|
||||
self.assertIn("not found", str(context.exception))
|
||||
|
||||
def test_ssl_certfile_not_found_raises(self):
|
||||
with tempfile.NamedTemporaryFile(suffix=".pem") as keyfile:
|
||||
with self.assertRaises(ValueError) as context:
|
||||
ServerArgs(
|
||||
model_path="dummy",
|
||||
ssl_keyfile=keyfile.name,
|
||||
ssl_certfile="/nonexistent/cert.pem",
|
||||
)
|
||||
self.assertIn("SSL certificate file not found", str(context.exception))
|
||||
|
||||
def test_ssl_ca_certs_not_found_raises(self):
|
||||
with tempfile.NamedTemporaryFile(suffix=".pem") as keyfile:
|
||||
with tempfile.NamedTemporaryFile(suffix=".pem") as certfile:
|
||||
with self.assertRaises(ValueError) as context:
|
||||
ServerArgs(
|
||||
model_path="dummy",
|
||||
ssl_keyfile=keyfile.name,
|
||||
ssl_certfile=certfile.name,
|
||||
ssl_ca_certs="/nonexistent/ca.pem",
|
||||
)
|
||||
self.assertIn(
|
||||
"SSL CA certificates file not found", str(context.exception)
|
||||
)
|
||||
|
||||
@patch("os.path.isfile", return_value=True)
|
||||
def test_url_returns_https_with_ssl_and_ipv6(self, _mock_isfile):
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy",
|
||||
host="::1",
|
||||
ssl_keyfile="key.pem",
|
||||
ssl_certfile="cert.pem",
|
||||
)
|
||||
self.assertEqual(server_args.url(), "https://[::1]:30000")
|
||||
|
||||
def test_enable_ssl_refresh_default_false(self):
|
||||
server_args = ServerArgs(model_path="dummy")
|
||||
self.assertFalse(server_args.enable_ssl_refresh)
|
||||
|
||||
def test_enable_ssl_refresh_without_ssl_raises(self):
|
||||
with self.assertRaises(ValueError) as context:
|
||||
ServerArgs(model_path="dummy", enable_ssl_refresh=True)
|
||||
self.assertIn("--enable-ssl-refresh", str(context.exception))
|
||||
self.assertIn("--ssl-certfile", str(context.exception))
|
||||
|
||||
@patch("os.path.isfile", return_value=True)
|
||||
def test_enable_ssl_refresh_with_ssl_accepted(self, _mock_isfile):
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy",
|
||||
ssl_keyfile="key.pem",
|
||||
ssl_certfile="cert.pem",
|
||||
enable_ssl_refresh=True,
|
||||
)
|
||||
self.assertTrue(server_args.enable_ssl_refresh)
|
||||
|
||||
@patch("os.path.isfile", return_value=True)
|
||||
def test_enable_ssl_refresh_cli_flag(self, _mock_isfile):
|
||||
server_args = prepare_server_args(
|
||||
[
|
||||
"--model-path",
|
||||
"dummy",
|
||||
"--ssl-keyfile",
|
||||
"key.pem",
|
||||
"--ssl-certfile",
|
||||
"cert.pem",
|
||||
"--enable-ssl-refresh",
|
||||
]
|
||||
)
|
||||
self.assertTrue(server_args.enable_ssl_refresh)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
165
test/registered/core/test_ssl_cert_refresher.py
Normal file
165
test/registered/core/test_ssl_cert_refresher.py
Normal file
@@ -0,0 +1,165 @@
|
||||
import asyncio
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from sglang.srt.entrypoints.ssl_utils import SSLCertRefresher
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=9, suite="stage-a-cpu-only")
|
||||
|
||||
|
||||
def _make_temp_pem(content: bytes) -> str:
|
||||
"""Create a temporary PEM file and return its path."""
|
||||
f = tempfile.NamedTemporaryFile(suffix=".pem", delete=False)
|
||||
f.write(content)
|
||||
f.flush()
|
||||
f.close()
|
||||
return f.name
|
||||
|
||||
|
||||
class TestSSLCertRefresher(CustomTestCase):
|
||||
"""Tests for the SSLCertRefresher class."""
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self._temp_files: list[str] = []
|
||||
|
||||
def tearDown(self):
|
||||
for path in self._temp_files:
|
||||
try:
|
||||
os.unlink(path)
|
||||
except OSError:
|
||||
pass
|
||||
super().tearDown()
|
||||
|
||||
def _track(self, path: str) -> str:
|
||||
"""Register a temp file for cleanup."""
|
||||
self._temp_files.append(path)
|
||||
return path
|
||||
|
||||
def _run_async(self, coro):
|
||||
"""Helper to run an async coroutine in tests."""
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
return loop.run_until_complete(coro)
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
def test_reload_cert_key_on_file_change(self):
|
||||
"""SSLCertRefresher calls load_cert_chain when cert/key files change."""
|
||||
mock_ctx = MagicMock()
|
||||
cert_path = self._track(_make_temp_pem(b"CERT_V1"))
|
||||
key_path = self._track(_make_temp_pem(b"KEY_V1"))
|
||||
|
||||
async def _test():
|
||||
refresher = SSLCertRefresher(mock_ctx, key_path, cert_path)
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
with open(cert_path, "w") as f:
|
||||
f.write("CERT_V2")
|
||||
|
||||
await asyncio.sleep(1.5)
|
||||
refresher.stop()
|
||||
return mock_ctx
|
||||
|
||||
result_ctx = self._run_async(_test())
|
||||
result_ctx.load_cert_chain.assert_called_with(cert_path, key_path)
|
||||
|
||||
def test_reload_ca_on_file_change(self):
|
||||
"""SSLCertRefresher calls load_verify_locations when CA file changes."""
|
||||
mock_ctx = MagicMock()
|
||||
cert_path = self._track(_make_temp_pem(b"CERT"))
|
||||
key_path = self._track(_make_temp_pem(b"KEY"))
|
||||
ca_path = self._track(_make_temp_pem(b"CA_V1"))
|
||||
|
||||
async def _test():
|
||||
refresher = SSLCertRefresher(mock_ctx, key_path, cert_path, ca_path)
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
with open(ca_path, "w") as f:
|
||||
f.write("CA_V2")
|
||||
|
||||
await asyncio.sleep(1.5)
|
||||
refresher.stop()
|
||||
return mock_ctx
|
||||
|
||||
result_ctx = self._run_async(_test())
|
||||
result_ctx.load_verify_locations.assert_called_with(ca_path)
|
||||
|
||||
def test_stop_cancels_tasks(self):
|
||||
"""Calling stop() prevents further reloads."""
|
||||
mock_ctx = MagicMock()
|
||||
cert_path = self._track(_make_temp_pem(b"CERT"))
|
||||
key_path = self._track(_make_temp_pem(b"KEY"))
|
||||
|
||||
async def _test():
|
||||
refresher = SSLCertRefresher(mock_ctx, key_path, cert_path)
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
refresher.stop()
|
||||
|
||||
with open(cert_path, "w") as f:
|
||||
f.write("CERT_AFTER_STOP")
|
||||
|
||||
await asyncio.sleep(1.0)
|
||||
return mock_ctx
|
||||
|
||||
result_ctx = self._run_async(_test())
|
||||
result_ctx.load_cert_chain.assert_not_called()
|
||||
|
||||
def test_no_ca_watcher_when_ca_not_provided(self):
|
||||
"""No CA watcher task is created when ca_path is None."""
|
||||
mock_ctx = MagicMock()
|
||||
cert_path = self._track(_make_temp_pem(b"CERT"))
|
||||
key_path = self._track(_make_temp_pem(b"KEY"))
|
||||
|
||||
async def _test():
|
||||
refresher = SSLCertRefresher(mock_ctx, key_path, cert_path)
|
||||
self.assertEqual(len(refresher._tasks), 1)
|
||||
refresher.stop()
|
||||
|
||||
self._run_async(_test())
|
||||
|
||||
def test_ca_watcher_created_when_ca_provided(self):
|
||||
"""A CA watcher task is created when ca_path is provided."""
|
||||
mock_ctx = MagicMock()
|
||||
cert_path = self._track(_make_temp_pem(b"CERT"))
|
||||
key_path = self._track(_make_temp_pem(b"KEY"))
|
||||
ca_path = self._track(_make_temp_pem(b"CA"))
|
||||
|
||||
async def _test():
|
||||
refresher = SSLCertRefresher(mock_ctx, key_path, cert_path, ca_path)
|
||||
self.assertEqual(len(refresher._tasks), 2)
|
||||
refresher.stop()
|
||||
|
||||
self._run_async(_test())
|
||||
|
||||
def test_reload_error_does_not_crash(self):
|
||||
"""A reload error is logged but doesn't crash the watcher."""
|
||||
mock_ctx = MagicMock()
|
||||
mock_ctx.load_cert_chain.side_effect = Exception("bad cert")
|
||||
cert_path = self._track(_make_temp_pem(b"CERT"))
|
||||
key_path = self._track(_make_temp_pem(b"KEY"))
|
||||
|
||||
async def _test():
|
||||
refresher = SSLCertRefresher(mock_ctx, key_path, cert_path)
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
with open(cert_path, "w") as f:
|
||||
f.write("BAD_CERT")
|
||||
|
||||
await asyncio.sleep(1.5)
|
||||
|
||||
for task in refresher._tasks:
|
||||
self.assertFalse(task.done())
|
||||
|
||||
refresher.stop()
|
||||
|
||||
self._run_async(_test())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user