wehub-resource-sync
102 lines
3.4 KiB
Python
102 lines
3.4 KiB
Python
"""TLS utilities shared across Ray libraries (e.g. Serve)."""
|
|
|
|
import datetime
|
|
import os
|
|
import socket
|
|
from typing import Tuple
|
|
|
|
from ray._common.network_utils import (
|
|
get_localhost_ip,
|
|
node_ip_address_from_perspective,
|
|
)
|
|
|
|
|
|
def generate_self_signed_tls_certs() -> Tuple[str, str]:
|
|
"""Create self-signed key/cert pair for testing.
|
|
|
|
Returns:
|
|
Tuple of (cert_pem_contents, key_pem_contents).
|
|
|
|
Raises:
|
|
ImportError: If the ``cryptography`` library is not installed.
|
|
"""
|
|
try:
|
|
from cryptography import x509
|
|
from cryptography.hazmat.backends import default_backend
|
|
from cryptography.hazmat.primitives import hashes, serialization
|
|
from cryptography.hazmat.primitives.asymmetric import rsa
|
|
from cryptography.x509.oid import NameOID
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
"Using self-signed TLS certs requires `cryptography`. "
|
|
"Install it with: pip install cryptography"
|
|
) from e
|
|
|
|
key = rsa.generate_private_key(
|
|
public_exponent=65537, key_size=2048, backend=default_backend()
|
|
)
|
|
key_contents = key.private_bytes(
|
|
encoding=serialization.Encoding.PEM,
|
|
format=serialization.PrivateFormat.PKCS8,
|
|
encryption_algorithm=serialization.NoEncryption(),
|
|
).decode()
|
|
|
|
subject = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "ray-internal")])
|
|
altnames = x509.SubjectAlternativeName(
|
|
[
|
|
x509.DNSName(socket.gethostbyname(socket.gethostname())),
|
|
x509.DNSName(get_localhost_ip()),
|
|
x509.DNSName(node_ip_address_from_perspective()),
|
|
x509.DNSName("localhost"),
|
|
]
|
|
)
|
|
now = datetime.datetime.utcnow()
|
|
cert = (
|
|
x509.CertificateBuilder()
|
|
.subject_name(subject)
|
|
.issuer_name(subject)
|
|
.add_extension(altnames, critical=False)
|
|
.public_key(key.public_key())
|
|
.serial_number(x509.random_serial_number())
|
|
.not_valid_before(now)
|
|
.not_valid_after(now + datetime.timedelta(days=365))
|
|
.sign(key, hashes.SHA256(), default_backend())
|
|
)
|
|
|
|
cert_contents = cert.public_bytes(serialization.Encoding.PEM).decode()
|
|
return cert_contents, key_contents
|
|
|
|
|
|
def add_port_to_grpc_server(server, address):
|
|
import grpc
|
|
|
|
if os.environ.get("RAY_USE_TLS", "0").lower() in ("1", "true"):
|
|
server_cert_chain, private_key, ca_cert = load_certs_from_env()
|
|
credentials = grpc.ssl_server_credentials(
|
|
[(private_key, server_cert_chain)],
|
|
root_certificates=ca_cert,
|
|
require_client_auth=ca_cert is not None,
|
|
)
|
|
return server.add_secure_port(address, credentials)
|
|
else:
|
|
return server.add_insecure_port(address)
|
|
|
|
|
|
def load_certs_from_env():
|
|
tls_env_vars = ["RAY_TLS_SERVER_CERT", "RAY_TLS_SERVER_KEY", "RAY_TLS_CA_CERT"]
|
|
if any(v not in os.environ for v in tls_env_vars):
|
|
raise RuntimeError(
|
|
"If the environment variable RAY_USE_TLS is set to true "
|
|
"then RAY_TLS_SERVER_CERT, RAY_TLS_SERVER_KEY and "
|
|
"RAY_TLS_CA_CERT must also be set."
|
|
)
|
|
|
|
with open(os.environ["RAY_TLS_SERVER_CERT"], "rb") as f:
|
|
server_cert_chain = f.read()
|
|
with open(os.environ["RAY_TLS_SERVER_KEY"], "rb") as f:
|
|
private_key = f.read()
|
|
with open(os.environ["RAY_TLS_CA_CERT"], "rb") as f:
|
|
ca_cert = f.read()
|
|
|
|
return server_cert_chain, private_key, ca_cert
|