Use the same WasiTlsView types for both p2 & p3 (#12896)

Author: badeend

crates/wasi-tls/Cargo.toml

@@ -11,6 +11,9 @@ description = "Wasmtime implementation of the wasi-tls API"
 [lints]
 workspace = true
 
+[package.metadata.docs.rs]
+all-features = true
+
 [features]
 default = ["p2", "rustls"]
 p2 = ["wasmtime-wasi/p2"]

crates/wasi-tls/src/lib.rs

@@ -13,8 +13,8 @@
 //!     component::{Linker, ResourceTable},
 //!     Store, Engine, Result,
 //! };
-//! use wasmtime_wasi_tls::{WasiTlsCtx, WasiTlsCtxBuilder};
-//! use wasmtime_wasi_tls::p2::{LinkOptions, WasiTls};
+//! use wasmtime_wasi_tls::{WasiTlsCtx, WasiTlsCtxBuilder, WasiTlsView, WasiTlsCtxView};
+//! use wasmtime_wasi_tls::p2::LinkOptions;
 //!
 //! struct Ctx {
 //!     table: ResourceTable,
@@ -28,6 +28,12 @@
 //!     }
 //! }
 //!
+//! impl WasiTlsView for Ctx {
+//!     fn tls(&mut self) -> WasiTlsCtxView<'_> {
+//!         WasiTlsCtxView { ctx: &mut self.wasi_tls_ctx, table: &mut self.table }
+//!     }
+//! }
+//!
 //! #[tokio::main]
 //! async fn main() -> Result<()> {
 //!     let ctx = Ctx {
@@ -49,15 +55,13 @@
 //!
 //!     // Set up wasi-cli
 //!     let mut store = Store::new(&engine, ctx);
-//!     let mut linker = Linker::new(&engine);
+//!     let mut linker: Linker<Ctx> = Linker::new(&engine);
 //!     wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
 //!
 //!     // Add wasi-tls types and turn on the feature in linker
 //!     let mut opts = LinkOptions::default();
 //!     opts.tls(true);
-//!     wasmtime_wasi_tls::p2::add_to_linker(&mut linker, &mut opts, |h: &mut Ctx| {
-//!         WasiTls::new(&h.wasi_tls_ctx, &mut h.table)
-//!     })?;
+//!     wasmtime_wasi_tls::p2::add_to_linker(&mut linker, &opts)?;
 //!
 //!     // ... use `linker` to instantiate within `store` ...
 //!     Ok(())
@@ -72,7 +76,6 @@
 #![doc(test(attr(allow(dead_code, unused_variables, unused_mut))))]
 
 use tokio::io::{AsyncRead, AsyncWrite};
-
 mod error;
 mod providers;
 
@@ -86,6 +89,9 @@ pub mod p3;
 pub use error::Error;
 pub use providers::*;
 
+#[cfg(any(feature = "p2", feature = "p3"))]
+use wasmtime::component::{HasData, ResourceTable};
+
 /// Builder-style structure used to create a [`WasiTlsCtx`].
 #[cfg(any(feature = "p2", feature = "p3"))]
 pub struct WasiTlsCtxBuilder {
@@ -131,6 +137,31 @@ pub struct WasiTlsCtx {
     pub(crate) provider: Box<dyn TlsProvider>,
 }
 
+/// The type for which this crate implements the `wasi:tls` interfaces.
+#[cfg(any(feature = "p2", feature = "p3"))]
+pub(crate) struct WasiTls;
+#[cfg(any(feature = "p2", feature = "p3"))]
+impl HasData for WasiTls {
+    type Data<'a> = WasiTlsCtxView<'a>;
+}
+
+/// View into [`WasiTlsCtx`] implementation and [`ResourceTable`].
+#[cfg(any(feature = "p2", feature = "p3"))]
+pub struct WasiTlsCtxView<'a> {
+    /// Mutable reference to table used to manage resources.
+    pub table: &'a mut ResourceTable,
+
+    /// Mutable reference to the WASI TLS context.
+    pub ctx: &'a mut WasiTlsCtx,
+}
+
+/// A trait which provides internal WASI TLS state.
+#[cfg(any(feature = "p2", feature = "p3"))]
+pub trait WasiTlsView: Send {
+    /// Return a [`WasiTlsCtxView`] from mutable reference to self.
+    fn tls(&mut self) -> WasiTlsCtxView<'_>;
+}
+
 /// The data stream that carries the encrypted TLS data.
 /// Typically this is a TCP stream.
 pub trait TlsTransport: AsyncRead + AsyncWrite + Send + Unpin + 'static {}

crates/wasi-tls/src/p2/host.rs

@@ -5,23 +5,23 @@ use wasmtime_wasi::p2::Pollable;
 use wasmtime_wasi::p2::{DynInputStream, DynOutputStream, DynPollable, IoError};
 
 use crate::p2::{
-    WasiTls, bindings,
+    bindings,
     io::{
         AsyncReadStream, AsyncWriteStream, FutureOutput, WasiFuture, WasiStreamReader,
         WasiStreamWriter,
     },
 };
-use crate::{TlsStream, TlsTransport};
+use crate::{TlsStream, TlsTransport, WasiTlsCtxView};
 
-impl<'a> bindings::types::Host for WasiTls<'a> {}
+impl<'a> bindings::types::Host for WasiTlsCtxView<'a> {}
 
 /// Represents the ClientHandshake which will be used to configure the handshake
 pub struct HostClientHandshake {
     server_name: String,
     transport: Box<dyn TlsTransport>,
 }
 
-impl<'a> bindings::types::HostClientHandshake for WasiTls<'a> {
+impl<'a> bindings::types::HostClientHandshake for WasiTlsCtxView<'a> {
     fn new(
         &mut self,
         server_name: String,
@@ -86,7 +86,7 @@ impl Pollable for HostFutureClientStreams {
     }
 }
 
-impl<'a> bindings::types::HostFutureClientStreams for WasiTls<'a> {
+impl<'a> bindings::types::HostFutureClientStreams for WasiTlsCtxView<'a> {
     fn subscribe(
         &mut self,
         this: Resource<HostFutureClientStreams>,
@@ -145,7 +145,7 @@ pub struct HostClientConnection(
     crate::p2::io::AsyncWriteStream<tokio::io::WriteHalf<Box<dyn TlsStream>>>,
 );
 
-impl<'a> bindings::types::HostClientConnection for WasiTls<'a> {
+impl<'a> bindings::types::HostClientConnection for WasiTlsCtxView<'a> {
     fn close_output(&mut self, this: Resource<HostClientConnection>) -> wasmtime::Result<()> {
         self.table.get_mut(&this)?.0.close()
     }

crates/wasi-tls/src/p2/mod.rs

@@ -1,6 +1,4 @@
-use wasmtime::component::{HasData, ResourceTable};
-
-use crate::WasiTlsCtx;
+use crate::{WasiTls, WasiTlsView};
 
 pub mod bindings;
 mod host;
@@ -9,30 +7,14 @@ mod io;
 pub use bindings::types::LinkOptions;
 pub use host::{HostClientConnection, HostClientHandshake, HostFutureClientStreams};
 
-/// Capture the state necessary for use in the `wasi-tls` API implementation.
-pub struct WasiTls<'a> {
-    pub(crate) ctx: &'a WasiTlsCtx,
-    pub(crate) table: &'a mut ResourceTable,
-}
-
-impl<'a> WasiTls<'a> {
-    /// Create a new Wasi TLS context.
-    pub fn new(ctx: &'a WasiTlsCtx, table: &'a mut ResourceTable) -> Self {
-        Self { ctx, table }
-    }
-}
-
 /// Add the `wasi-tls` world's types to a [`wasmtime::component::Linker`].
-pub fn add_to_linker<T: Send + 'static>(
+pub fn add_to_linker<T>(
     l: &mut wasmtime::component::Linker<T>,
-    opts: &mut LinkOptions,
-    f: fn(&mut T) -> WasiTls<'_>,
-) -> wasmtime::Result<()> {
-    bindings::types::add_to_linker::<_, HasWasiTls>(l, &opts, f)?;
+    opts: &LinkOptions,
+) -> wasmtime::Result<()>
+where
+    T: WasiTlsView + 'static,
+{
+    bindings::types::add_to_linker::<_, WasiTls>(l, &opts, T::tls)?;
     Ok(())
 }
-
-struct HasWasiTls;
-impl HasData for HasWasiTls {
-    type Data<'a> = WasiTls<'a>;
-}

crates/wasi-tls/src/p3/host.rs

@@ -1,8 +1,8 @@
 //! p3 host implementation for `wasi:tls`.
 
 use crate::p3::util::{AsyncReadProducer, AsyncWriteConsumer, Closed, Deferred, Shared, pipe};
-use crate::p3::{WasiTls, WasiTlsCtxView, bindings};
-use crate::{BoxFutureTlsStream, Error, TlsStream};
+use crate::p3::{WasiTls, bindings};
+use crate::{BoxFutureTlsStream, Error, TlsStream, WasiTlsCtxView};
 use std::pin::Pin;
 use std::task::{Context, Poll};
 use tokio::{io::AsyncWriteExt as _, sync::oneshot};

crates/wasi-tls/src/p3/mod.rs

@@ -12,31 +12,9 @@ pub mod bindings;
 pub mod host;
 pub(crate) mod util;
 
-use crate::WasiTlsCtx;
+use crate::{WasiTls, WasiTlsView};
 use bindings::tls::{client, types};
-use wasmtime::component::{HasData, Linker, ResourceTable};
-
-/// The type for which this crate implements the `wasi:tls` interfaces.
-pub struct WasiTls;
-
-impl HasData for WasiTls {
-    type Data<'a> = WasiTlsCtxView<'a>;
-}
-
-/// View into [`WasiTlsCtx`] implementation and [`ResourceTable`].
-pub struct WasiTlsCtxView<'a> {
-    /// Mutable reference to table used to manage resources.
-    pub table: &'a mut ResourceTable,
-
-    /// Mutable reference to the WASI TLS context.
-    pub ctx: &'a mut WasiTlsCtx,
-}
-
-/// A trait which provides internal WASI TLS state.
-pub trait WasiTlsView: Send {
-    /// Return a [`WasiTlsCtxView`] from mutable reference to self.
-    fn tls(&mut self) -> WasiTlsCtxView<'_>;
-}
+use wasmtime::component::Linker;
 
 /// Add all interfaces from this module into the `linker` provided.
 pub fn add_to_linker<T>(linker: &mut Linker<T>) -> wasmtime::Result<()>

crates/wasi-tls/tests/p2/mod.rs

@@ -4,7 +4,9 @@ use wasmtime::{
     format_err,
 };
 use wasmtime_wasi::{WasiCtx, WasiCtxView, WasiView, p2::bindings::Command};
-use wasmtime_wasi_tls::{TlsProvider, WasiTlsCtx, WasiTlsCtxBuilder, p2};
+use wasmtime_wasi_tls::{
+    TlsProvider, WasiTlsCtx, WasiTlsCtxBuilder, WasiTlsCtxView, WasiTlsView, p2,
+};
 
 struct Ctx {
     table: ResourceTable,
@@ -21,6 +23,15 @@ impl WasiView for Ctx {
     }
 }
 
+impl WasiTlsView for Ctx {
+    fn tls(&mut self) -> WasiTlsCtxView<'_> {
+        WasiTlsCtxView {
+            ctx: &mut self.wasi_tls_ctx,
+            table: &mut self.table,
+        }
+    }
+}
+
 async fn run_test(provider: Box<dyn TlsProvider>, path: &str) -> Result<()> {
     let ctx = Ctx {
         table: ResourceTable::new(),
@@ -41,9 +52,7 @@ async fn run_test(provider: Box<dyn TlsProvider>, path: &str) -> Result<()> {
     wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
     let mut opts = p2::LinkOptions::default();
     opts.tls(true);
-    p2::add_to_linker(&mut linker, &mut opts, |h: &mut Ctx| {
-        p2::WasiTls::new(&h.wasi_tls_ctx, &mut h.table)
-    })?;
+    wasmtime_wasi_tls::p2::add_to_linker(&mut linker, &opts)?;
 
     let command = Command::instantiate_async(&mut store, &component, &linker).await?;
     command

crates/wasi-tls/tests/p3/mod.rs

@@ -4,10 +4,7 @@ use wasmtime::{
     format_err,
 };
 use wasmtime_wasi::{WasiCtx, WasiCtxView, WasiView, p3::bindings::Command};
-use wasmtime_wasi_tls::{
-    TlsProvider, WasiTlsCtx, WasiTlsCtxBuilder,
-    p3::{WasiTlsCtxView, WasiTlsView},
-};
+use wasmtime_wasi_tls::{TlsProvider, WasiTlsCtx, WasiTlsCtxBuilder, WasiTlsCtxView, WasiTlsView};
 
 struct Ctx {
     table: ResourceTable,

src/commands/run.rs

@@ -30,8 +30,6 @@ use wasmtime_wasi_keyvalue::{WasiKeyValue, WasiKeyValueCtx, WasiKeyValueCtxBuild
 use wasmtime_wasi_nn::wit::WasiNnView;
 #[cfg(feature = "wasi-threads")]
 use wasmtime_wasi_threads::WasiThreadsCtx;
-#[cfg(feature = "wasi-tls")]
-use wasmtime_wasi_tls::{WasiTlsCtx, p2::WasiTls};
 
 fn parse_preloads(s: &str) -> Result<(String, PathBuf)> {
     let parts: Vec<&str> = s.splitn(2, '=').collect();
@@ -1319,14 +1317,7 @@ impl RunCommand {
                     CliLinker::Component(linker) => {
                         let mut opts = wasmtime_wasi_tls::p2::LinkOptions::default();
                         opts.tls(true);
-                        wasmtime_wasi_tls::p2::add_to_linker(linker, &mut opts, |h| {
-                            let ctx = h.wasip1_ctx.as_mut().expect("wasi is not configured");
-                            let ctx = Arc::get_mut(ctx).unwrap().get_mut().unwrap();
-                            WasiTls::new(
-                                Arc::get_mut(h.wasi_tls.as_mut().unwrap()).unwrap(),
-                                ctx.ctx().table,
-                            )
-                        })?;
+                        wasmtime_wasi_tls::p2::add_to_linker(linker, &opts)?;
 
                         #[cfg(feature = "component-model-async")]
                         if self.run.common.wasi.p3.unwrap_or(crate::common::P3_DEFAULT) {
@@ -1502,7 +1493,7 @@ pub struct Host {
     #[cfg(feature = "wasi-keyvalue")]
     wasi_keyvalue: Option<Arc<WasiKeyValueCtx>>,
     #[cfg(feature = "wasi-tls")]
-    wasi_tls: Option<Arc<WasiTlsCtx>>,
+    wasi_tls: Option<Arc<wasmtime_wasi_tls::WasiTlsCtx>>,
 }
 
 impl Host {
@@ -1551,10 +1542,10 @@ impl wasmtime_wasi_http::p3::WasiHttpView for Host {
     }
 }
 
-#[cfg(all(feature = "wasi-tls", feature = "component-model-async"))]
-impl wasmtime_wasi_tls::p3::WasiTlsView for Host {
-    fn tls(&mut self) -> wasmtime_wasi_tls::p3::WasiTlsCtxView<'_> {
-        wasmtime_wasi_tls::p3::WasiTlsCtxView {
+#[cfg(all(feature = "wasi-tls"))]
+impl wasmtime_wasi_tls::WasiTlsView for Host {
+    fn tls(&mut self) -> wasmtime_wasi_tls::WasiTlsCtxView<'_> {
+        wasmtime_wasi_tls::WasiTlsCtxView {
             table: WasiView::ctx(unwrap_singlethread_context(&mut self.wasip1_ctx)).table,
             ctx: Arc::get_mut(self.wasi_tls.as_mut().unwrap()).unwrap(),
         }