diff --git a/Cargo.lock b/Cargo.lock index d7a7256..9e3b827 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -50,6 +50,12 @@ version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" +[[package]] +name = "autocfg" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" + [[package]] name = "base64" version = "0.22.1" @@ -62,6 +68,17 @@ version = "1.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" +[[package]] +name = "bitfield-struct" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8769c4854c5ada2852ddf6fd09d15cf43d4c2aaeccb4de6432f5402f08a6003b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "bitflags" version = "2.11.1" @@ -231,6 +248,19 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "compio-ktls" +version = "0.1.0" +source = "git+https://github.com/compio-rs/compio-ktls.git?branch=alpn#746d50d8450993b432058a0bea042dbb426e97bf" +dependencies = [ + "compio-buf", + "compio-io", + "ktls-core", + "libc", + "rustls", + "synchrony", +] + [[package]] name = "compio-log" version = "0.1.0" @@ -267,6 +297,7 @@ version = "0.1.0" dependencies = [ "compio", "compio-executor", + "compio-ktls", "compio-log", "once_cell", "pem", @@ -709,6 +740,20 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "ktls-core" +version = "0.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5502ef8386c9b091182a810e10582871abde6f5187ac6418983939443fcd92cb" +dependencies = [ + "bitfield-struct", + "libc", + "nix", + "rustls", + "zeroize", + "zeroize_derive", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -795,12 +840,34 @@ version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + [[package]] name = "mod_use" version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "21909324aa58f5c284d91cac514c6d081210901dc372d7c8ea6a9d7e0406097a" +[[package]] +name = "nix" +version = "0.30.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6" +dependencies = [ + "bitflags", + "cfg-if", + "cfg_aliases", + "libc", + "memoffset", +] + [[package]] name = "no-std-compat" version = "0.4.1" @@ -1749,6 +1816,17 @@ version = "1.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" +[[package]] +name = "zeroize_derive" +version = "1.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85a5b4158499876c763cb03bc4e49185d3cccbabb15b33c627f7884f43db852e" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "zmij" version = "1.0.21" diff --git a/Cargo.toml b/Cargo.toml index da1ed0c..c24bb8b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,5 +17,9 @@ compio = { version = "0.19.0-rc.1", default-features = false } compio-executor = "0.1.0-rc.1" compio-log = "0.1.0" +[workspace.dependencies."compio-ktls"] +git = "https://github.com/compio-rs/compio-ktls.git" +branch = "alpn" + [patch.crates-io] compio-py-dynamic-openssl = { path = "compio-py-dynamic-openssl" } diff --git a/compio-py/Cargo.toml b/compio-py/Cargo.toml index a5fb8da..25827c8 100644 --- a/compio-py/Cargo.toml +++ b/compio-py/Cargo.toml @@ -25,6 +25,9 @@ socket2 = "0.6.2" rustls-webpki = "0" pem = "3.0.6" pkcs8 = { version = "0.10.2", features = ["encryption"] } -compio = { workspace = true, features = ["io-uring", "polling", "io", "py-dynamic-openssl", "rustls", "ring"] } +compio = { workspace = true, features = ["io-uring", "polling", "io", "io-ancillary", "py-dynamic-openssl", "rustls", "ring"] } compio-executor = { workspace = true } compio-log = { workspace = true } + +[target."cfg(target_os = \"linux\")".dependencies] +compio-ktls = { workspace = true } diff --git a/compio-py/src/event_loop.rs b/compio-py/src/event_loop.rs index 2fcb27f..250398b 100644 --- a/compio-py/src/event_loop.rs +++ b/compio-py/src/event_loop.rs @@ -28,12 +28,27 @@ use crate::{ static COMPIO_FUTURE: OnceCell> = OnceCell::new(); +#[derive(Debug, Clone, Copy, Default)] +pub enum KtlsMode { + #[default] + Prefer, + Require, + Disabled, +} + +impl KtlsMode { + pub fn enabled(self) -> bool { + matches!(self, KtlsMode::Prefer | KtlsMode::Require) + } +} + #[pyclass(subclass)] pub struct CompioLoop { runtime: OwnedRefCell, stopping: Arc, debug: Arc, registered: Py, + ktls_mode: KtlsMode, } #[pymethods] @@ -45,6 +60,7 @@ impl CompioLoop { stopping: Default::default(), debug: Default::default(), registered: import::weakref::weak_key_dict(py)?.cast_into()?.unbind(), + ktls_mode: KtlsMode::default(), }) } @@ -268,6 +284,27 @@ impl CompioLoop { slf.borrow() .spawn_py(py, net::PySocket::new(pyloop, domain, ty, protocol)) } + + #[getter] + fn ktls_mode(&self) -> String { + format!("{:?}", self.ktls_mode) + } + + #[setter] + fn set_ktls_mode(&mut self, mode: &Bound) -> PyResult<()> { + self.ktls_mode = match mode.str()?.to_str()? { + "Prefer" => KtlsMode::Prefer, + "Require" => KtlsMode::Require, + "Disabled" => KtlsMode::Disabled, + other => { + return Err(PyValueError::new_err(format!( + "unknown KTLS mode: {}", + other + ))); + } + }; + Ok(()) + } } impl CompioLoop { @@ -332,6 +369,10 @@ impl CompioLoop { Ok(rv) } + pub fn get_ktls_mode(&self) -> KtlsMode { + self.ktls_mode + } + fn socket_to_fd(&self, py: Python, sock: &Py) -> PyResult> { let registered = self.registered.bind(py); let sock = sock.bind(py); diff --git a/compio-py/src/lib.rs b/compio-py/src/lib.rs index fd29ef1..a1fcb77 100644 --- a/compio-py/src/lib.rs +++ b/compio-py/src/lib.rs @@ -60,6 +60,15 @@ enum Either { Right(R), } +impl Either { + fn into + From>(self) -> T { + match self { + Self::Left(l) => T::from(l), + Self::Right(r) => T::from(r), + } + } +} + impl IoBuf for Either { fn as_init(&self) -> &[u8] { match self { diff --git a/compio-py/src/net/mod.rs b/compio-py/src/net/mod.rs index 616aeee..b6eaf28 100644 --- a/compio-py/src/net/mod.rs +++ b/compio-py/src/net/mod.rs @@ -4,12 +4,16 @@ use std::{io, net::Shutdown}; use compio::{ - buf::{BufResult, IntoInner, IoBuf, IoBufMut, IoVectoredBufMut, buf_try}, + buf::{BufResult, IntoInner, IoBuf, IoBufMut, IoVectoredBuf, IoVectoredBufMut, buf_try}, driver::{ ToSharedFd, impl_raw_fd, - op::{self, BufResultExt}, + op::{self, BufResultExt, RecvResultExt, VecBufResultExt}, + }, + io::{ + AsyncRead, AsyncWrite, + ancillary::{AsyncReadAncillary, AsyncWriteAncillary}, + util::Splittable, }, - io::{AsyncRead, AsyncWrite, util::Splittable}, }; use pyo3::{ exceptions::PyTypeError, @@ -39,6 +43,32 @@ impl AsyncRead for &SocketStream { } } +impl AsyncReadAncillary for &SocketStream { + #[inline] + async fn read_with_ancillary( + &mut self, + buffer: T, + control: C, + ) -> BufResult<(usize, usize), (T, C)> { + self.inner + .recv_msg(buffer, control, op::RecvFlags::empty()) + .await + .map_res(|(res, len, _addr)| (res, len)) + } + + #[inline] + async fn read_vectored_with_ancillary( + &mut self, + buffer: T, + control: C, + ) -> BufResult<(usize, usize), (T, C)> { + self.inner + .recv_msg_vectored(buffer, control, op::RecvFlags::empty()) + .await + .map_res(|(res, len, _addr)| (res, len)) + } +} + impl AsyncWrite for &SocketStream { async fn write(&mut self, buf: T) -> BufResult { self.inner.send(buf, op::SendFlags::empty()).await @@ -53,6 +83,30 @@ impl AsyncWrite for &SocketStream { } } +impl AsyncWriteAncillary for &SocketStream { + #[inline] + async fn write_with_ancillary( + &mut self, + buffer: T, + control: C, + ) -> BufResult { + self.inner + .send_msg(buffer, control, None, op::SendFlags::empty()) + .await + } + + #[inline] + async fn write_vectored_with_ancillary( + &mut self, + buffer: T, + control: C, + ) -> BufResult { + self.inner + .send_msg_vectored(buffer, control, None, op::SendFlags::empty()) + .await + } +} + impl AsyncRead for SocketStream { #[inline] async fn read(&mut self, buf: B) -> BufResult { @@ -65,6 +119,24 @@ impl AsyncRead for SocketStream { } } +impl AsyncReadAncillary for SocketStream { + async fn read_with_ancillary( + &mut self, + buffer: T, + control: C, + ) -> BufResult<(usize, usize), (T, C)> { + (&*self).read_with_ancillary(buffer, control).await + } + + async fn read_vectored_with_ancillary( + &mut self, + buffer: T, + control: C, + ) -> BufResult<(usize, usize), (T, C)> { + (&*self).read_vectored_with_ancillary(buffer, control).await + } +} + impl AsyncWrite for SocketStream { async fn write(&mut self, buf: T) -> BufResult { (&*self).write(buf).await @@ -79,6 +151,26 @@ impl AsyncWrite for SocketStream { } } +impl AsyncWriteAncillary for SocketStream { + async fn write_with_ancillary( + &mut self, + buffer: T, + control: C, + ) -> BufResult { + (&*self).write_with_ancillary(buffer, control).await + } + + async fn write_vectored_with_ancillary( + &mut self, + buffer: T, + control: C, + ) -> BufResult { + (&*self) + .write_vectored_with_ancillary(buffer, control) + .await + } +} + pub struct ReadHalf(Socket); pub struct WriteHalf(Socket); @@ -212,12 +304,60 @@ impl Socket { unsafe { res.map_advanced() } } + pub async fn recv_msg( + &self, + buffer: T, + control: C, + flags: op::RecvFlags, + ) -> BufResult<(usize, usize, Option), (T, C)> { + self.recv_msg_vectored([buffer], control, flags) + .await + .map_buffer(|([buffer], control)| (buffer, control)) + } + + pub async fn recv_msg_vectored( + &self, + buffer: T, + control: C, + flags: op::RecvFlags, + ) -> BufResult<(usize, usize, Option), (T, C)> { + let fd = self.to_shared_fd(); + let op = op::RecvMsg::new(fd, buffer, control, flags); + let res = runtime::execute(op).await; + let res = res.into_inner().map_addr(); + unsafe { res.map_vec_advanced() } + } + async fn send(&self, buffer: T, flags: op::SendFlags) -> BufResult { let fd = self.to_shared_fd(); let op = op::Send::new(fd, buffer, flags); runtime::execute(op).await.into_inner() } + pub async fn send_msg( + &self, + buffer: T, + control: C, + addr: Option<&SockAddr>, + flags: op::SendFlags, + ) -> BufResult { + self.send_msg_vectored([buffer], control, addr, flags) + .await + .map_buffer(|([buffer], control)| (buffer, control)) + } + + pub async fn send_msg_vectored( + &self, + buffer: T, + control: C, + addr: Option<&SockAddr>, + flags: op::SendFlags, + ) -> BufResult { + let fd = self.to_shared_fd(); + let op = op::SendMsg::new(fd, buffer, control, addr.cloned(), flags); + runtime::execute(op).await.into_inner() + } + async fn shutdown(&self, how: Shutdown) -> io::Result<()> { #[cfg(unix)] { diff --git a/compio-py/src/net/socket.rs b/compio-py/src/net/socket.rs index 58d8257..2813b7d 100644 --- a/compio-py/src/net/socket.rs +++ b/compio-py/src/net/socket.rs @@ -17,7 +17,7 @@ use compio::{ }; use pyo3::{ IntoPyObjectExt, - exceptions::{PyOSError, PyTypeError, PyValueError}, + exceptions::{PyNotImplementedError, PyOSError, PyRuntimeError, PyTypeError, PyValueError}, prelude::*, types::{PyBytes, PyList}, }; @@ -27,7 +27,11 @@ use super::{ Socket, SocketStream, idna_converter, name_to_ip, ssl::{RustlsContext, SSLImpl, SSLSocket, SSLSocketMetadata}, }; -use crate::{Either, event_loop::CompioLoop, extract_py_err, import, py_any_to_buffer}; +use crate::{ + Either, + event_loop::{CompioLoop, KtlsMode}, + extract_py_err, import, py_any_to_buffer, +}; #[pyclass(unsendable, name = "Socket")] pub struct PySocket { @@ -193,12 +197,13 @@ impl PySocket { ) -> PyResult> { let this = slf.clone().unbind(); let slf = slf.borrow().pyloop.bind(py).borrow(); + let ktls_mode = slf.get_ktls_mode(); slf.spawn_py(py, async move { let mut metadata = SSLSocketMetadata::default(); metadata.server_side = server_side; - // First, verify parameters and prepare either TlsAcceptor or TlsConnector. - let tls = Python::attach(|py| { + // First, verify parameters and prepare either SSLContext or Server/ClientConfig + let ctx = Python::attach(|py| { let has_ossl = py_dynamic_openssl::load_py(py)?; // Coerce sslcontext to either SSLContext or RustlsContext let ctx: Either<_, Bound> = match sslcontext @@ -213,15 +218,20 @@ impl PySocket { None if !server_side => Either::Left(import::ssl::create_default_context(py)?), None => Err(PyValueError::new_err("server_side requires sslcontext"))?, }; - // Build TlsAcceptor or TlsConnector from the context - match ctx { - Either::Left(ctx) if has_ossl => SSLContext::try_from(ctx).map(|ctx| { + Ok(match ctx { + Either::Left(ctx) if has_ossl => { + let ctx = SSLContext::try_from(ctx)?; + if matches!(ktls_mode, KtlsMode::Require) { + return Err(PyNotImplementedError::new_err( + "kTLS is unimplemented for OpenSSL", + )); + } if server_side { - Either::Left(TlsAcceptor::from(ctx)) + Either::Left(Either::Left(ctx)) } else { - Either::Right(TlsConnector::from(ctx)) + Either::Right(Either::Left(ctx)) } - }), + } Either::Left(ctx) => { // This case is guaranteed to be a default client-side SSLContext debug_assert!(!server_side); @@ -236,20 +246,21 @@ impl PySocket { .add(CertificateDer::from(cert)) .map_err(|e| PyValueError::new_err(e.to_string()))?; } - let config = rustls::ClientConfig::builder() + let mut config = rustls::ClientConfig::builder() .with_root_certificates(root_store) .with_no_client_auth(); - Ok(Either::Right(TlsConnector::from(Arc::new(config)))) + config.enable_secret_extraction = ktls_mode.enabled(); + Either::Right(Either::Right(Arc::new(config))) } Either::Right(ctx) => { metadata.implementation = SSLImpl::Rustls; let ctx = ctx.borrow(); - Ok(match ctx.build(py, server_side)? { - Either::Left(c) => Either::Left(TlsAcceptor::from(c)), - Either::Right(c) => Either::Right(TlsConnector::from(c)), - }) + match ctx.build(py, server_side)? { + Either::Left(c) => Either::Left(Either::Right(c)), + Either::Right(c) => Either::Right(Either::Right(c)), + } } - } + }) })?; // Then, do TLS handshake accordingly @@ -258,9 +269,36 @@ impl PySocket { }; metadata.fd = inner.as_raw_fd(); let stream = SocketStream { inner }; - let stream = match tls { - Either::Left(acceptor) => extract_py_err(acceptor.accept(stream).await)?, - Either::Right(connector) => { + let stream = match ctx { + #[cfg(target_os = "linux")] + Either::Left(Either::Right(ctx)) => { + let stream = if ktls_mode.enabled() { + extract_py_err( + compio_ktls::KtlsAcceptor::from(ctx.clone()) + .accept(stream) + .await, + )? + } else { + Err(stream) + }; + match (ktls_mode, stream) { + (KtlsMode::Require, Err(_)) => { + return Err(PyRuntimeError::new_err("kTLS is unsupported")); + } + (_, Ok(stream)) => { + metadata.ktls = true; + stream.into() + } + (_, Err(stream)) => { + extract_py_err(TlsAcceptor::from(ctx).accept(stream).await)?.into() + } + } + } + Either::Left(ctx) => { + debug_assert!(!matches!(ktls_mode, KtlsMode::Require)); + extract_py_err(ctx.into::().accept(stream).await)?.into() + } + Either::Right(ctx) => { let name = match server_hostname { Some(name) => name, None => stream @@ -272,7 +310,36 @@ impl PySocket { .ip() .to_string(), }; - extract_py_err(connector.connect(&name, stream).await)? + #[cfg(target_os = "linux")] + let stream = match (ktls_mode, &ctx) { + (KtlsMode::Require, Either::Left(_)) => unreachable!(), + (KtlsMode::Require | KtlsMode::Prefer, Either::Right(ctx)) => { + extract_py_err( + compio_ktls::KtlsConnector::from(ctx.clone()) + .connect(&name, stream) + .await, + )? + } + _ => Err(stream), + }; + #[cfg(not(target_os = "linux"))] + let stream = Err::(stream); + match (ktls_mode, stream) { + (KtlsMode::Require, Err(_)) => { + return Err(PyRuntimeError::new_err("kTLS is unsupported")); + } + #[cfg(target_os = "linux")] + (_, Ok(stream)) => { + metadata.ktls = true; + stream.into() + } + #[cfg(not(target_os = "linux"))] + (_, Ok(_)) => unreachable!(), + (_, Err(stream)) => { + extract_py_err(ctx.into::().connect(&name, stream).await)? + .into() + } + } } }; diff --git a/compio-py/src/net/ssl.rs b/compio-py/src/net/ssl.rs index 31e51ff..9cc2308 100644 --- a/compio-py/src/net/ssl.rs +++ b/compio-py/src/net/ssl.rs @@ -56,6 +56,7 @@ pub struct SSLSocketMetadata { pub implementation: SSLImpl, pub server_side: bool, pub fd: RawFd, + pub ktls: bool, } #[pyclass(unsendable)] @@ -69,7 +70,7 @@ impl SSLSocket { pub fn new( py: Python, pyloop: &Py, - stream: TlsStream, + stream: InnerStream, metadata: SSLSocketMetadata, ) -> PyResult> { Py::new( @@ -132,8 +133,11 @@ impl SSLSocket { fn __repr__(&self) -> PyResult { Ok(match self.inner.borrow_opt()?.as_ref() { Some(_) => format!( - "", - self.metadata.implementation, self.metadata.server_side, self.metadata.fd, + "", + self.metadata.implementation, + if self.metadata.ktls { "+kTLS" } else { "" }, + self.metadata.server_side, + self.metadata.fd, ), None => "".to_string(), }) @@ -141,40 +145,101 @@ impl SSLSocket { } #[derive(Clone)] -struct SSLSocketInner(Rc>>>); +struct SSLSocketInner(Rc>>); impl SSLSocketInner { - fn new(stream: TlsStream) -> Self { + fn new(stream: InnerStream) -> Self { Self(Rc::new(RefCell::new(Some(stream)))) } #[inline] - fn borrow_opt(&self) -> PyResult>>> { + fn borrow_opt(&self) -> PyResult>> { self.0 .try_borrow() .map_err(|_| PyRuntimeError::new_err("concurrent access to SSLSocket")) } #[inline] - fn borrow(&self) -> PyResult>> { + fn borrow(&self) -> PyResult> { cell::Ref::filter_map(self.borrow_opt()?, |rv| rv.as_ref()) .map_err(|_| PyOSError::new_err("socket is closed")) } #[inline] - fn borrow_mut_opt(&self) -> PyResult>>> { + fn borrow_mut_opt(&self) -> PyResult>> { self.0 .try_borrow_mut() .map_err(|_| PyRuntimeError::new_err("concurrent access to SSLSocket")) } #[inline] - fn borrow_mut(&self) -> PyResult>> { + fn borrow_mut(&self) -> PyResult> { cell::RefMut::filter_map(self.borrow_mut_opt()?, |rv| rv.as_mut()) .map_err(|_| PyOSError::new_err("socket is closed")) } } +#[cfg(target_os = "linux")] +pub enum InnerStream { + UserSpace(TlsStream), + Kernel(compio_ktls::KtlsStream), +} + +#[cfg(target_os = "linux")] +impl InnerStream { + async fn read(&mut self, buf: B) -> compio::BufResult { + match self { + InnerStream::UserSpace(stream) => stream.read(buf).await, + InnerStream::Kernel(stream) => stream.read(buf).await, + } + } + + async fn write(&mut self, buf: T) -> compio::BufResult { + match self { + InnerStream::UserSpace(stream) => stream.write(buf).await, + InnerStream::Kernel(stream) => stream.write(buf).await, + } + } + + async fn flush(&mut self) -> io::Result<()> { + match self { + InnerStream::UserSpace(stream) => stream.flush().await, + InnerStream::Kernel(stream) => stream.flush().await, + } + } + + async fn shutdown(&mut self) -> io::Result<()> { + match self { + InnerStream::UserSpace(stream) => stream.shutdown().await, + InnerStream::Kernel(stream) => stream.shutdown().await, + } + } + + fn negotiated_alpn(&self) -> Option> { + match self { + InnerStream::UserSpace(stream) => stream.negotiated_alpn(), + InnerStream::Kernel(stream) => stream.negotiated_alpn(), + } + } +} + +#[cfg(target_os = "linux")] +impl From> for InnerStream { + fn from(stream: TlsStream) -> Self { + Self::UserSpace(stream) + } +} + +#[cfg(target_os = "linux")] +impl From> for InnerStream { + fn from(stream: compio_ktls::KtlsStream) -> Self { + Self::Kernel(stream) + } +} + +#[cfg(not(target_os = "linux"))] +type InnerStream = TlsStream; + #[pyclass] pub struct RustlsContext { state: ContextState, @@ -429,6 +494,7 @@ impl RustlsContext { if let Some(protocols) = alpn_protocols { config.alpn_protocols = protocols; } + config.enable_secret_extraction = true; Either::Left(Arc::new(config)) } else { let builder = ClientConfig::builder_with_details( @@ -465,6 +531,7 @@ impl RustlsContext { if let Some(protocols) = alpn_protocols { config.alpn_protocols = protocols; } + config.enable_secret_extraction = true; Either::Right(Arc::new(config)) }; state.config = Some(rv.clone()); diff --git a/src/compio/__init__.py b/src/compio/__init__.py index d712b9b..36fc878 100644 --- a/src/compio/__init__.py +++ b/src/compio/__init__.py @@ -2,6 +2,6 @@ # Copyright 2025 Fantix King from ._core import Handle, Socket, SSLSocket, TimerHandle -from .loop import CompioLoop, DriverType +from .loop import CompioLoop, DriverType, KtlsMode -__all__ = ["CompioLoop", "DriverType", "Handle", "Socket", "SSLSocket", "TimerHandle"] +__all__ = ["CompioLoop", "DriverType", "Handle", "Socket", "SSLSocket", "TimerHandle", "KtlsMode"] diff --git a/src/compio/_core.pyi b/src/compio/_core.pyi index 12e0948..4e0ef26 100644 --- a/src/compio/_core.pyi +++ b/src/compio/_core.pyi @@ -71,6 +71,10 @@ class CompioLoop: type: socket.SocketKind | int = -1, proto: int = -1, ) -> Socket: ... + @property + def ktls_mode(self) -> str: ... + @ktls_mode.setter + def ktls_mode(self, mode: Any) -> None: ... class Socket: async def connect(self, address: _Address, /) -> None: ... @@ -101,7 +105,7 @@ class RustlsContext: @property def alpn_protocols(self) -> Optional[tuple[bytes, ...]]: ... @alpn_protocols.setter - def alpn_protocols(self, protocols: Optional[tuple[bytes, ...]]) -> None: ... + def alpn_protocols(self, protocols: Optional[Iterable[bytes]]) -> None: ... @alpn_protocols.deleter def alpn_protocols(self) -> None: ... @property diff --git a/src/compio/loop.py b/src/compio/loop.py index 043daa8..f707781 100644 --- a/src/compio/loop.py +++ b/src/compio/loop.py @@ -23,6 +23,13 @@ class DriverType(enum.StrEnum): IOCP = "IOCP" +class KtlsMode(enum.StrEnum): + DEFAULT = "Default" + PREFER = "Prefer" + REQUIRE = "Require" + DISABLED = "Disabled" + + class CompioLoop(_core.CompioLoop, asyncio.AbstractEventLoop): _exception_handler: Optional[_ExceptionHandler] @@ -35,6 +42,14 @@ def __init__(self) -> None: def get_driver_type(self) -> DriverType: return DriverType(super().get_driver_type()) + @property # type: ignore[override] + def ktls_mode(self) -> KtlsMode: + return KtlsMode(super().ktls_mode) + + @ktls_mode.setter + def ktls_mode(self, mode: KtlsMode) -> None: + _core.CompioLoop.ktls_mode.__set__(self, mode.value) # type: ignore[attr-defined] + def __repr__(self) -> str: try: driver_type = f"driver={self.get_driver_type().value} "