spin_factor_wasi/
sockets.rs

1//! Socket quota tracking and WASI socket host implementations.
2//!
3//! This module provides [`SocketPermitState`], [`SpinSocketsView`], and
4//! [`SpinSockets`] — the types needed to intercept WASI TCP/UDP socket
5//! creation and enforce a per-app cap on the number of concurrently open
6//! sockets.
7
8use std::{
9    collections::HashMap,
10    marker::PhantomData,
11    sync::{Arc, Mutex},
12};
13
14use spin_connection_semaphore::{ConnectionPermit, ConnectionSemaphore};
15use wasmtime::component::{HasData, Resource};
16use wasmtime_wasi::p2::bindings::sockets::network::{
17    ErrorCode as SocketErrorCode, Host as NetworkHost, Network,
18};
19use wasmtime_wasi::p2::bindings::sockets::tcp::{self as p2_tcp, IpSocketAddress, ShutdownType};
20use wasmtime_wasi::p2::bindings::sockets::tcp_create_socket as p2_tcp_create;
21use wasmtime_wasi::p2::bindings::sockets::udp as p2_udp;
22use wasmtime_wasi::p2::bindings::sockets::udp_create_socket as p2_udp_create;
23use wasmtime_wasi::p2::{DynInputStream, DynOutputStream, DynPollable};
24use wasmtime_wasi::sockets::{TcpSocket, UdpSocket, WasiSockets, WasiSocketsCtxView};
25
26/// Shared state for tracking per-socket semaphore permits. Permits are
27/// acquired when a socket is allocated (at `start_connect` for TCP, at
28/// `create_udp_socket` for UDP) and released when the socket resource is dropped.
29pub struct SocketPermitState {
30    semaphore: ConnectionSemaphore,
31    /// Active permits keyed by socket resource rep, released when the resource is dropped.
32    active: Mutex<HashMap<u32, ConnectionPermit>>,
33}
34
35impl SocketPermitState {
36    pub fn new(semaphore: ConnectionSemaphore) -> Arc<Self> {
37        Arc::new(Self {
38            semaphore,
39            active: Mutex::new(HashMap::new()),
40        })
41    }
42}
43
44/// A view over WASI socket state that carries an optional per-instance socket
45/// permit store, enabling per-connection quota tracking.
46pub struct SpinSocketsView<'a, T> {
47    pub(crate) inner: WasiSocketsCtxView<'a>,
48    pub(crate) permit_state: Option<Arc<SocketPermitState>>,
49    pub(crate) getter: fn(&mut T) -> WasiSocketsCtxView<'_>,
50}
51
52impl<'a, T> std::ops::Deref for SpinSocketsView<'a, T> {
53    type Target = WasiSocketsCtxView<'a>;
54    fn deref(&self) -> &Self::Target {
55        &self.inner
56    }
57}
58
59impl<T> std::ops::DerefMut for SpinSocketsView<'_, T> {
60    fn deref_mut(&mut self) -> &mut Self::Target {
61        &mut self.inner
62    }
63}
64
65/// [`HasData`] accessor for [`SpinSocketsView`], used in place of [`WasiSockets`]
66/// when registering TCP socket bindings so that `start_connect` and `drop` can
67/// participate in socket quota tracking.
68pub struct SpinSockets<T>(PhantomData<fn() -> T>);
69
70impl<T: 'static> HasData for SpinSockets<T> {
71    type Data<'a> = SpinSocketsView<'a, T>;
72}
73
74impl<'a, T> SpinSocketsView<'a, T> {
75    /// Consumes this view and returns the inner [`WasiSocketsCtxView`].
76    pub fn into_wasi(self) -> WasiSocketsCtxView<'a> {
77        self.inner
78    }
79}
80
81impl<T> SpinSocketsView<'_, T> {
82    /// Attempts to acquire a connection permit from the semaphore.
83    ///
84    /// Returns `Ok(None)` when no quota is configured, `Ok(Some(permit))` on
85    /// success, or `Err(())` when the quota is exhausted.
86    ///
87    /// The returned permit is unregistered — call [`Self::register_permit`] once
88    /// the socket resource rep is known to tie its lifetime to the socket.
89    pub(crate) fn try_acquire(&self) -> Result<Option<ConnectionPermit>, ()> {
90        let Some(state) = &self.permit_state else {
91            return Ok(None);
92        };
93        state.semaphore.try_acquire().map(Some).ok_or(())
94    }
95
96    /// Registers `permit` under `socket_rep` so it is held until the socket is
97    /// dropped. No-op when `permit` is `None` (no quota configured).
98    pub(crate) fn register_permit(&self, socket_rep: u32, permit: Option<ConnectionPermit>) {
99        let (Some(state), Some(permit)) = (&self.permit_state, permit) else {
100            return;
101        };
102        state
103            .active
104            .lock()
105            .unwrap_or_else(|e| e.into_inner())
106            .insert(socket_rep, permit);
107    }
108
109    /// Releases the connection permit for `socket_rep`, if any.
110    pub(crate) fn release_permit(&self, socket_rep: u32) {
111        if let Some(state) = &self.permit_state {
112            state
113                .active
114                .lock()
115                .unwrap_or_else(|e| e.into_inner())
116                .remove(&socket_rep);
117        }
118    }
119}
120
121impl<T> p2_tcp::Host for SpinSocketsView<'_, T> {}
122
123impl<T> p2_tcp::HostTcpSocket for SpinSocketsView<'_, T> {
124    async fn start_bind(
125        &mut self,
126        this: Resource<TcpSocket>,
127        network: Resource<Network>,
128        local_address: IpSocketAddress,
129    ) -> wasmtime_wasi::p2::SocketResult<()> {
130        p2_tcp::HostTcpSocket::start_bind(&mut self.inner, this, network, local_address).await
131    }
132
133    fn finish_bind(&mut self, this: Resource<TcpSocket>) -> wasmtime_wasi::p2::SocketResult<()> {
134        p2_tcp::HostTcpSocket::finish_bind(&mut self.inner, this)
135    }
136
137    async fn start_connect(
138        &mut self,
139        this: Resource<TcpSocket>,
140        network: Resource<Network>,
141        remote_address: IpSocketAddress,
142    ) -> wasmtime_wasi::p2::SocketResult<()> {
143        let socket_rep = this.rep();
144        // Unlike outbound HTTP (which queues when its permit pool is exhausted),
145        // sockets fail immediately. Waiting would risk deadlock if a component
146        // holds sockets open across async yield points, and raw-socket callers
147        // are better positioned to implement their own retry logic.
148        let Ok(permit) = self.try_acquire() else {
149            tracing::warn!("TCP socket connection refused: connection quota exhausted");
150            return Err(SocketErrorCode::NewSocketLimit.into());
151        };
152        let result =
153            p2_tcp::HostTcpSocket::start_connect(&mut self.inner, this, network, remote_address)
154                .await;
155        if result.is_ok() {
156            self.register_permit(socket_rep, permit);
157        }
158        // On error, `permit` is dropped here, automatically releasing the semaphore slot.
159        result
160    }
161
162    fn finish_connect(
163        &mut self,
164        this: Resource<TcpSocket>,
165    ) -> wasmtime_wasi::p2::SocketResult<(Resource<DynInputStream>, Resource<DynOutputStream>)>
166    {
167        p2_tcp::HostTcpSocket::finish_connect(&mut self.inner, this)
168    }
169
170    fn start_listen(&mut self, this: Resource<TcpSocket>) -> wasmtime_wasi::p2::SocketResult<()> {
171        p2_tcp::HostTcpSocket::start_listen(&mut self.inner, this)
172    }
173
174    fn finish_listen(&mut self, this: Resource<TcpSocket>) -> wasmtime_wasi::p2::SocketResult<()> {
175        p2_tcp::HostTcpSocket::finish_listen(&mut self.inner, this)
176    }
177
178    fn accept(
179        &mut self,
180        this: Resource<TcpSocket>,
181    ) -> wasmtime_wasi::p2::SocketResult<(
182        Resource<TcpSocket>,
183        Resource<DynInputStream>,
184        Resource<DynOutputStream>,
185    )> {
186        p2_tcp::HostTcpSocket::accept(&mut self.inner, this)
187    }
188
189    fn local_address(
190        &mut self,
191        this: Resource<TcpSocket>,
192    ) -> wasmtime_wasi::p2::SocketResult<IpSocketAddress> {
193        p2_tcp::HostTcpSocket::local_address(&mut self.inner, this)
194    }
195
196    fn remote_address(
197        &mut self,
198        this: Resource<TcpSocket>,
199    ) -> wasmtime_wasi::p2::SocketResult<IpSocketAddress> {
200        p2_tcp::HostTcpSocket::remote_address(&mut self.inner, this)
201    }
202
203    fn is_listening(&mut self, this: Resource<TcpSocket>) -> wasmtime::Result<bool> {
204        p2_tcp::HostTcpSocket::is_listening(&mut self.inner, this)
205    }
206
207    fn address_family(
208        &mut self,
209        this: Resource<TcpSocket>,
210    ) -> wasmtime::Result<wasmtime_wasi::p2::bindings::sockets::network::IpAddressFamily> {
211        p2_tcp::HostTcpSocket::address_family(&mut self.inner, this)
212    }
213
214    fn set_listen_backlog_size(
215        &mut self,
216        this: Resource<TcpSocket>,
217        value: u64,
218    ) -> wasmtime_wasi::p2::SocketResult<()> {
219        p2_tcp::HostTcpSocket::set_listen_backlog_size(&mut self.inner, this, value)
220    }
221
222    fn keep_alive_enabled(
223        &mut self,
224        this: Resource<TcpSocket>,
225    ) -> wasmtime_wasi::p2::SocketResult<bool> {
226        p2_tcp::HostTcpSocket::keep_alive_enabled(&mut self.inner, this)
227    }
228
229    fn set_keep_alive_enabled(
230        &mut self,
231        this: Resource<TcpSocket>,
232        value: bool,
233    ) -> wasmtime_wasi::p2::SocketResult<()> {
234        p2_tcp::HostTcpSocket::set_keep_alive_enabled(&mut self.inner, this, value)
235    }
236
237    fn keep_alive_idle_time(
238        &mut self,
239        this: Resource<TcpSocket>,
240    ) -> wasmtime_wasi::p2::SocketResult<u64> {
241        p2_tcp::HostTcpSocket::keep_alive_idle_time(&mut self.inner, this)
242    }
243
244    fn set_keep_alive_idle_time(
245        &mut self,
246        this: Resource<TcpSocket>,
247        value: u64,
248    ) -> wasmtime_wasi::p2::SocketResult<()> {
249        p2_tcp::HostTcpSocket::set_keep_alive_idle_time(&mut self.inner, this, value)
250    }
251
252    fn keep_alive_interval(
253        &mut self,
254        this: Resource<TcpSocket>,
255    ) -> wasmtime_wasi::p2::SocketResult<u64> {
256        p2_tcp::HostTcpSocket::keep_alive_interval(&mut self.inner, this)
257    }
258
259    fn set_keep_alive_interval(
260        &mut self,
261        this: Resource<TcpSocket>,
262        value: u64,
263    ) -> wasmtime_wasi::p2::SocketResult<()> {
264        p2_tcp::HostTcpSocket::set_keep_alive_interval(&mut self.inner, this, value)
265    }
266
267    fn keep_alive_count(
268        &mut self,
269        this: Resource<TcpSocket>,
270    ) -> wasmtime_wasi::p2::SocketResult<u32> {
271        p2_tcp::HostTcpSocket::keep_alive_count(&mut self.inner, this)
272    }
273
274    fn set_keep_alive_count(
275        &mut self,
276        this: Resource<TcpSocket>,
277        value: u32,
278    ) -> wasmtime_wasi::p2::SocketResult<()> {
279        p2_tcp::HostTcpSocket::set_keep_alive_count(&mut self.inner, this, value)
280    }
281
282    fn hop_limit(&mut self, this: Resource<TcpSocket>) -> wasmtime_wasi::p2::SocketResult<u8> {
283        p2_tcp::HostTcpSocket::hop_limit(&mut self.inner, this)
284    }
285
286    fn set_hop_limit(
287        &mut self,
288        this: Resource<TcpSocket>,
289        value: u8,
290    ) -> wasmtime_wasi::p2::SocketResult<()> {
291        p2_tcp::HostTcpSocket::set_hop_limit(&mut self.inner, this, value)
292    }
293
294    fn receive_buffer_size(
295        &mut self,
296        this: Resource<TcpSocket>,
297    ) -> wasmtime_wasi::p2::SocketResult<u64> {
298        p2_tcp::HostTcpSocket::receive_buffer_size(&mut self.inner, this)
299    }
300
301    fn set_receive_buffer_size(
302        &mut self,
303        this: Resource<TcpSocket>,
304        value: u64,
305    ) -> wasmtime_wasi::p2::SocketResult<()> {
306        p2_tcp::HostTcpSocket::set_receive_buffer_size(&mut self.inner, this, value)
307    }
308
309    fn send_buffer_size(
310        &mut self,
311        this: Resource<TcpSocket>,
312    ) -> wasmtime_wasi::p2::SocketResult<u64> {
313        p2_tcp::HostTcpSocket::send_buffer_size(&mut self.inner, this)
314    }
315
316    fn set_send_buffer_size(
317        &mut self,
318        this: Resource<TcpSocket>,
319        value: u64,
320    ) -> wasmtime_wasi::p2::SocketResult<()> {
321        p2_tcp::HostTcpSocket::set_send_buffer_size(&mut self.inner, this, value)
322    }
323
324    fn subscribe(&mut self, this: Resource<TcpSocket>) -> wasmtime::Result<Resource<DynPollable>> {
325        p2_tcp::HostTcpSocket::subscribe(&mut self.inner, this)
326    }
327
328    fn shutdown(
329        &mut self,
330        this: Resource<TcpSocket>,
331        shutdown_type: ShutdownType,
332    ) -> wasmtime_wasi::p2::SocketResult<()> {
333        p2_tcp::HostTcpSocket::shutdown(&mut self.inner, this, shutdown_type)
334    }
335
336    fn drop(&mut self, this: Resource<TcpSocket>) -> wasmtime::Result<()> {
337        self.release_permit(this.rep());
338        p2_tcp::HostTcpSocket::drop(&mut self.inner, this)
339    }
340}
341
342impl<T> NetworkHost for SpinSocketsView<'_, T> {
343    fn convert_error_code(
344        &mut self,
345        error: wasmtime_wasi::p2::SocketError,
346    ) -> wasmtime::Result<wasmtime_wasi::p2::bindings::sockets::network::ErrorCode> {
347        NetworkHost::convert_error_code(&mut self.inner, error)
348    }
349
350    fn network_error_code(
351        &mut self,
352        err: Resource<wasmtime::Error>,
353    ) -> wasmtime::Result<Option<wasmtime_wasi::p2::bindings::sockets::network::ErrorCode>> {
354        NetworkHost::network_error_code(&mut self.inner, err)
355    }
356}
357
358impl<T> wasmtime_wasi::p2::bindings::sockets::network::HostNetwork for SpinSocketsView<'_, T> {
359    fn drop(&mut self, this: Resource<Network>) -> wasmtime::Result<()> {
360        wasmtime_wasi::p2::bindings::sockets::network::HostNetwork::drop(&mut self.inner, this)
361    }
362}
363
364impl<T> p2_tcp_create::Host for SpinSocketsView<'_, T> {
365    fn create_tcp_socket(
366        &mut self,
367        address_family: wasmtime_wasi::p2::bindings::sockets::network::IpAddressFamily,
368    ) -> wasmtime_wasi::p2::SocketResult<Resource<TcpSocket>> {
369        p2_tcp_create::Host::create_tcp_socket(&mut self.inner, address_family)
370    }
371}
372
373impl<T> p2_udp::Host for SpinSocketsView<'_, T> {}
374
375impl<T> p2_udp::HostUdpSocket for SpinSocketsView<'_, T> {
376    async fn start_bind(
377        &mut self,
378        this: Resource<p2_udp::UdpSocket>,
379        network: Resource<p2_udp::Network>,
380        local_address: p2_udp::IpSocketAddress,
381    ) -> wasmtime_wasi::p2::SocketResult<()> {
382        p2_udp::HostUdpSocket::start_bind(&mut self.inner, this, network, local_address).await
383    }
384
385    fn finish_bind(
386        &mut self,
387        this: Resource<p2_udp::UdpSocket>,
388    ) -> wasmtime_wasi::p2::SocketResult<()> {
389        p2_udp::HostUdpSocket::finish_bind(&mut self.inner, this)
390    }
391
392    async fn stream(
393        &mut self,
394        this: Resource<p2_udp::UdpSocket>,
395        remote_address: Option<p2_udp::IpSocketAddress>,
396    ) -> wasmtime_wasi::p2::SocketResult<(
397        Resource<p2_udp::IncomingDatagramStream>,
398        Resource<p2_udp::OutgoingDatagramStream>,
399    )> {
400        p2_udp::HostUdpSocket::stream(&mut self.inner, this, remote_address).await
401    }
402
403    fn local_address(
404        &mut self,
405        this: Resource<p2_udp::UdpSocket>,
406    ) -> wasmtime_wasi::p2::SocketResult<p2_udp::IpSocketAddress> {
407        p2_udp::HostUdpSocket::local_address(&mut self.inner, this)
408    }
409
410    fn remote_address(
411        &mut self,
412        this: Resource<p2_udp::UdpSocket>,
413    ) -> wasmtime_wasi::p2::SocketResult<p2_udp::IpSocketAddress> {
414        p2_udp::HostUdpSocket::remote_address(&mut self.inner, this)
415    }
416
417    fn address_family(
418        &mut self,
419        this: Resource<p2_udp::UdpSocket>,
420    ) -> wasmtime::Result<p2_udp::IpAddressFamily> {
421        p2_udp::HostUdpSocket::address_family(&mut self.inner, this)
422    }
423
424    fn unicast_hop_limit(
425        &mut self,
426        this: Resource<p2_udp::UdpSocket>,
427    ) -> wasmtime_wasi::p2::SocketResult<u8> {
428        p2_udp::HostUdpSocket::unicast_hop_limit(&mut self.inner, this)
429    }
430
431    fn set_unicast_hop_limit(
432        &mut self,
433        this: Resource<p2_udp::UdpSocket>,
434        value: u8,
435    ) -> wasmtime_wasi::p2::SocketResult<()> {
436        p2_udp::HostUdpSocket::set_unicast_hop_limit(&mut self.inner, this, value)
437    }
438
439    fn receive_buffer_size(
440        &mut self,
441        this: Resource<p2_udp::UdpSocket>,
442    ) -> wasmtime_wasi::p2::SocketResult<u64> {
443        p2_udp::HostUdpSocket::receive_buffer_size(&mut self.inner, this)
444    }
445
446    fn set_receive_buffer_size(
447        &mut self,
448        this: Resource<p2_udp::UdpSocket>,
449        value: u64,
450    ) -> wasmtime_wasi::p2::SocketResult<()> {
451        p2_udp::HostUdpSocket::set_receive_buffer_size(&mut self.inner, this, value)
452    }
453
454    fn send_buffer_size(
455        &mut self,
456        this: Resource<p2_udp::UdpSocket>,
457    ) -> wasmtime_wasi::p2::SocketResult<u64> {
458        p2_udp::HostUdpSocket::send_buffer_size(&mut self.inner, this)
459    }
460
461    fn set_send_buffer_size(
462        &mut self,
463        this: Resource<p2_udp::UdpSocket>,
464        value: u64,
465    ) -> wasmtime_wasi::p2::SocketResult<()> {
466        p2_udp::HostUdpSocket::set_send_buffer_size(&mut self.inner, this, value)
467    }
468
469    fn subscribe(
470        &mut self,
471        this: Resource<p2_udp::UdpSocket>,
472    ) -> wasmtime::Result<Resource<DynPollable>> {
473        p2_udp::HostUdpSocket::subscribe(&mut self.inner, this)
474    }
475
476    fn drop(&mut self, this: Resource<p2_udp::UdpSocket>) -> wasmtime::Result<()> {
477        self.release_permit(this.rep());
478        p2_udp::HostUdpSocket::drop(&mut self.inner, this)
479    }
480}
481
482impl<T> p2_udp::HostIncomingDatagramStream for SpinSocketsView<'_, T> {
483    fn receive(
484        &mut self,
485        this: Resource<p2_udp::IncomingDatagramStream>,
486        max_results: u64,
487    ) -> wasmtime_wasi::p2::SocketResult<Vec<p2_udp::IncomingDatagram>> {
488        p2_udp::HostIncomingDatagramStream::receive(&mut self.inner, this, max_results)
489    }
490
491    fn subscribe(
492        &mut self,
493        this: Resource<p2_udp::IncomingDatagramStream>,
494    ) -> wasmtime::Result<Resource<DynPollable>> {
495        p2_udp::HostIncomingDatagramStream::subscribe(&mut self.inner, this)
496    }
497
498    fn drop(&mut self, this: Resource<p2_udp::IncomingDatagramStream>) -> wasmtime::Result<()> {
499        p2_udp::HostIncomingDatagramStream::drop(&mut self.inner, this)
500    }
501}
502
503impl<T> p2_udp::HostOutgoingDatagramStream for SpinSocketsView<'_, T> {
504    fn check_send(
505        &mut self,
506        this: Resource<p2_udp::OutgoingDatagramStream>,
507    ) -> wasmtime_wasi::p2::SocketResult<u64> {
508        p2_udp::HostOutgoingDatagramStream::check_send(&mut self.inner, this)
509    }
510
511    async fn send(
512        &mut self,
513        this: Resource<p2_udp::OutgoingDatagramStream>,
514        datagrams: Vec<p2_udp::OutgoingDatagram>,
515    ) -> wasmtime_wasi::p2::SocketResult<u64> {
516        p2_udp::HostOutgoingDatagramStream::send(&mut self.inner, this, datagrams).await
517    }
518
519    fn subscribe(
520        &mut self,
521        this: Resource<p2_udp::OutgoingDatagramStream>,
522    ) -> wasmtime::Result<Resource<DynPollable>> {
523        p2_udp::HostOutgoingDatagramStream::subscribe(&mut self.inner, this)
524    }
525
526    fn drop(&mut self, this: Resource<p2_udp::OutgoingDatagramStream>) -> wasmtime::Result<()> {
527        p2_udp::HostOutgoingDatagramStream::drop(&mut self.inner, this)
528    }
529}
530
531impl<T> p2_udp_create::Host for SpinSocketsView<'_, T> {
532    fn create_udp_socket(
533        &mut self,
534        address_family: wasmtime_wasi::p2::bindings::sockets::network::IpAddressFamily,
535    ) -> wasmtime_wasi::p2::SocketResult<Resource<UdpSocket>> {
536        // Check quota before allocating the socket resource.
537        // See the analogous comment in `start_connect` for why we fail
538        // immediately rather than waiting (as outbound HTTP does).
539        let Ok(permit) = self.try_acquire() else {
540            tracing::warn!("UDP socket creation refused: connection quota exhausted");
541            return Err(SocketErrorCode::NewSocketLimit.into());
542        };
543        let sock = p2_udp_create::Host::create_udp_socket(&mut self.inner, address_family)?;
544        self.register_permit(sock.rep(), permit);
545        Ok(sock)
546    }
547}
548
549// ===== p3 impls =====
550
551use wasmtime::AsContextMut as _;
552use wasmtime::component::{Access, Accessor};
553use wasmtime_wasi::p3::bindings::sockets::types::{
554    self as p3_types, Duration as p3_Duration, ErrorCode as p3_ErrorCode, Host as p3_Host,
555    HostTcpSocket as p3_HostTcpSocket, HostTcpSocketWithStore, HostUdpSocket as p3_HostUdpSocket,
556    HostUdpSocketWithStore, IpAddressFamily as p3_IpAddressFamily,
557    IpSocketAddress as p3_IpSocketAddress,
558};
559use wasmtime_wasi::p3::sockets::SocketResult as P3SocketResult;
560
561impl<T> p3_Host for SpinSocketsView<'_, T> {
562    fn convert_error_code(
563        &mut self,
564        error: wasmtime_wasi::p3::sockets::SocketError,
565    ) -> wasmtime::Result<p3_ErrorCode> {
566        p3_Host::convert_error_code(&mut self.inner, error)
567    }
568}
569
570impl<T> p3_HostTcpSocket for SpinSocketsView<'_, T> {
571    async fn bind(
572        &mut self,
573        socket: Resource<p3_types::TcpSocket>,
574        local_address: p3_IpSocketAddress,
575    ) -> P3SocketResult<()> {
576        p3_HostTcpSocket::bind(&mut self.inner, socket, local_address).await
577    }
578
579    fn create(
580        &mut self,
581        address_family: p3_IpAddressFamily,
582    ) -> P3SocketResult<Resource<p3_types::TcpSocket>> {
583        p3_HostTcpSocket::create(&mut self.inner, address_family)
584    }
585
586    fn get_local_address(
587        &mut self,
588        socket: Resource<p3_types::TcpSocket>,
589    ) -> P3SocketResult<p3_IpSocketAddress> {
590        p3_HostTcpSocket::get_local_address(&mut self.inner, socket)
591    }
592
593    fn get_remote_address(
594        &mut self,
595        socket: Resource<p3_types::TcpSocket>,
596    ) -> P3SocketResult<p3_IpSocketAddress> {
597        p3_HostTcpSocket::get_remote_address(&mut self.inner, socket)
598    }
599
600    fn get_is_listening(
601        &mut self,
602        socket: Resource<p3_types::TcpSocket>,
603    ) -> wasmtime::Result<bool> {
604        p3_HostTcpSocket::get_is_listening(&mut self.inner, socket)
605    }
606
607    fn get_address_family(
608        &mut self,
609        socket: Resource<p3_types::TcpSocket>,
610    ) -> wasmtime::Result<p3_IpAddressFamily> {
611        p3_HostTcpSocket::get_address_family(&mut self.inner, socket)
612    }
613
614    fn set_listen_backlog_size(
615        &mut self,
616        socket: Resource<p3_types::TcpSocket>,
617        value: u64,
618    ) -> P3SocketResult<()> {
619        p3_HostTcpSocket::set_listen_backlog_size(&mut self.inner, socket, value)
620    }
621
622    fn get_keep_alive_enabled(
623        &mut self,
624        socket: Resource<p3_types::TcpSocket>,
625    ) -> P3SocketResult<bool> {
626        p3_HostTcpSocket::get_keep_alive_enabled(&mut self.inner, socket)
627    }
628
629    fn set_keep_alive_enabled(
630        &mut self,
631        socket: Resource<p3_types::TcpSocket>,
632        value: bool,
633    ) -> P3SocketResult<()> {
634        p3_HostTcpSocket::set_keep_alive_enabled(&mut self.inner, socket, value)
635    }
636
637    fn get_keep_alive_idle_time(
638        &mut self,
639        socket: Resource<p3_types::TcpSocket>,
640    ) -> P3SocketResult<p3_Duration> {
641        p3_HostTcpSocket::get_keep_alive_idle_time(&mut self.inner, socket)
642    }
643
644    fn set_keep_alive_idle_time(
645        &mut self,
646        socket: Resource<p3_types::TcpSocket>,
647        value: p3_Duration,
648    ) -> P3SocketResult<()> {
649        p3_HostTcpSocket::set_keep_alive_idle_time(&mut self.inner, socket, value)
650    }
651
652    fn get_keep_alive_interval(
653        &mut self,
654        socket: Resource<p3_types::TcpSocket>,
655    ) -> P3SocketResult<p3_Duration> {
656        p3_HostTcpSocket::get_keep_alive_interval(&mut self.inner, socket)
657    }
658
659    fn set_keep_alive_interval(
660        &mut self,
661        socket: Resource<p3_types::TcpSocket>,
662        value: p3_Duration,
663    ) -> P3SocketResult<()> {
664        p3_HostTcpSocket::set_keep_alive_interval(&mut self.inner, socket, value)
665    }
666
667    fn get_keep_alive_count(
668        &mut self,
669        socket: Resource<p3_types::TcpSocket>,
670    ) -> P3SocketResult<u32> {
671        p3_HostTcpSocket::get_keep_alive_count(&mut self.inner, socket)
672    }
673
674    fn set_keep_alive_count(
675        &mut self,
676        socket: Resource<p3_types::TcpSocket>,
677        value: u32,
678    ) -> P3SocketResult<()> {
679        p3_HostTcpSocket::set_keep_alive_count(&mut self.inner, socket, value)
680    }
681
682    fn get_hop_limit(&mut self, socket: Resource<p3_types::TcpSocket>) -> P3SocketResult<u8> {
683        p3_HostTcpSocket::get_hop_limit(&mut self.inner, socket)
684    }
685
686    fn set_hop_limit(
687        &mut self,
688        socket: Resource<p3_types::TcpSocket>,
689        value: u8,
690    ) -> P3SocketResult<()> {
691        p3_HostTcpSocket::set_hop_limit(&mut self.inner, socket, value)
692    }
693
694    fn get_receive_buffer_size(
695        &mut self,
696        socket: Resource<p3_types::TcpSocket>,
697    ) -> P3SocketResult<u64> {
698        p3_HostTcpSocket::get_receive_buffer_size(&mut self.inner, socket)
699    }
700
701    fn set_receive_buffer_size(
702        &mut self,
703        socket: Resource<p3_types::TcpSocket>,
704        value: u64,
705    ) -> P3SocketResult<()> {
706        p3_HostTcpSocket::set_receive_buffer_size(&mut self.inner, socket, value)
707    }
708
709    fn get_send_buffer_size(
710        &mut self,
711        socket: Resource<p3_types::TcpSocket>,
712    ) -> P3SocketResult<u64> {
713        p3_HostTcpSocket::get_send_buffer_size(&mut self.inner, socket)
714    }
715
716    fn set_send_buffer_size(
717        &mut self,
718        socket: Resource<p3_types::TcpSocket>,
719        value: u64,
720    ) -> P3SocketResult<()> {
721        p3_HostTcpSocket::set_send_buffer_size(&mut self.inner, socket, value)
722    }
723
724    fn drop(&mut self, sock: Resource<p3_types::TcpSocket>) -> wasmtime::Result<()> {
725        self.release_permit(sock.rep());
726        p3_HostTcpSocket::drop(&mut self.inner, sock)
727    }
728}
729
730impl<T> p3_HostUdpSocket for SpinSocketsView<'_, T> {
731    async fn bind(
732        &mut self,
733        socket: Resource<p3_types::UdpSocket>,
734        local_address: p3_IpSocketAddress,
735    ) -> P3SocketResult<()> {
736        p3_HostUdpSocket::bind(&mut self.inner, socket, local_address).await
737    }
738
739    async fn connect(
740        &mut self,
741        socket: Resource<p3_types::UdpSocket>,
742        remote_address: p3_IpSocketAddress,
743    ) -> P3SocketResult<()> {
744        p3_HostUdpSocket::connect(&mut self.inner, socket, remote_address).await
745    }
746
747    fn create(
748        &mut self,
749        address_family: p3_IpAddressFamily,
750    ) -> P3SocketResult<Resource<p3_types::UdpSocket>> {
751        // Check quota before allocating the socket resource.
752        // See the analogous comment in `start_connect` for why we fail
753        // immediately rather than waiting (as outbound HTTP does).
754        let Ok(permit) = self.try_acquire() else {
755            tracing::warn!("UDP socket creation refused: connection quota exhausted");
756            return Err(p3_ErrorCode::Other(Some("connection quota exhausted".into())).into());
757        };
758        let sock = p3_HostUdpSocket::create(&mut self.inner, address_family)?;
759        self.register_permit(sock.rep(), permit);
760        Ok(sock)
761    }
762
763    fn disconnect(&mut self, socket: Resource<p3_types::UdpSocket>) -> P3SocketResult<()> {
764        p3_HostUdpSocket::disconnect(&mut self.inner, socket)
765    }
766
767    fn get_local_address(
768        &mut self,
769        socket: Resource<p3_types::UdpSocket>,
770    ) -> P3SocketResult<p3_IpSocketAddress> {
771        p3_HostUdpSocket::get_local_address(&mut self.inner, socket)
772    }
773
774    fn get_remote_address(
775        &mut self,
776        socket: Resource<p3_types::UdpSocket>,
777    ) -> P3SocketResult<p3_IpSocketAddress> {
778        p3_HostUdpSocket::get_remote_address(&mut self.inner, socket)
779    }
780
781    fn get_address_family(
782        &mut self,
783        socket: Resource<p3_types::UdpSocket>,
784    ) -> wasmtime::Result<p3_IpAddressFamily> {
785        p3_HostUdpSocket::get_address_family(&mut self.inner, socket)
786    }
787
788    fn get_unicast_hop_limit(
789        &mut self,
790        socket: Resource<p3_types::UdpSocket>,
791    ) -> P3SocketResult<u8> {
792        p3_HostUdpSocket::get_unicast_hop_limit(&mut self.inner, socket)
793    }
794
795    fn set_unicast_hop_limit(
796        &mut self,
797        socket: Resource<p3_types::UdpSocket>,
798        value: u8,
799    ) -> P3SocketResult<()> {
800        p3_HostUdpSocket::set_unicast_hop_limit(&mut self.inner, socket, value)
801    }
802
803    fn get_receive_buffer_size(
804        &mut self,
805        socket: Resource<p3_types::UdpSocket>,
806    ) -> P3SocketResult<u64> {
807        p3_HostUdpSocket::get_receive_buffer_size(&mut self.inner, socket)
808    }
809
810    fn set_receive_buffer_size(
811        &mut self,
812        socket: Resource<p3_types::UdpSocket>,
813        value: u64,
814    ) -> P3SocketResult<()> {
815        p3_HostUdpSocket::set_receive_buffer_size(&mut self.inner, socket, value)
816    }
817
818    fn get_send_buffer_size(
819        &mut self,
820        socket: Resource<p3_types::UdpSocket>,
821    ) -> P3SocketResult<u64> {
822        p3_HostUdpSocket::get_send_buffer_size(&mut self.inner, socket)
823    }
824
825    fn set_send_buffer_size(
826        &mut self,
827        socket: Resource<p3_types::UdpSocket>,
828        value: u64,
829    ) -> P3SocketResult<()> {
830        p3_HostUdpSocket::set_send_buffer_size(&mut self.inner, socket, value)
831    }
832
833    fn drop(&mut self, sock: Resource<p3_types::UdpSocket>) -> wasmtime::Result<()> {
834        self.release_permit(sock.rep());
835        p3_HostUdpSocket::drop(&mut self.inner, sock)
836    }
837}
838
839impl<T: Send + 'static> HostTcpSocketWithStore<T> for SpinSockets<T> {
840    async fn connect(
841        store: &Accessor<T, Self>,
842        socket: Resource<p3_types::TcpSocket>,
843        remote_address: p3_IpSocketAddress,
844    ) -> P3SocketResult<()> {
845        let socket_rep = socket.rep();
846        // Unlike outbound HTTP (which queues when its permit pool is exhausted),
847        // sockets fail immediately. See p2 `start_connect` for rationale.
848        let permit = match store.with(|mut access| access.get().try_acquire()) {
849            Ok(p) => p,
850            Err(()) => {
851                tracing::warn!("TCP socket connection refused: connection quota exhausted");
852                return Err(p3_ErrorCode::Other(Some("connection quota exhausted".into())).into());
853            }
854        };
855        let getter = store.with(|mut store| store.get().getter);
856        let wasi_accessor = store.with_getter::<WasiSockets>(getter);
857        let result: P3SocketResult<()> = <WasiSockets as HostTcpSocketWithStore<T>>::connect(
858            &wasi_accessor,
859            socket,
860            remote_address,
861        )
862        .await;
863        if result.is_ok() {
864            store.with(|mut access| {
865                access.get().register_permit(socket_rep, permit);
866            });
867        }
868        result
869    }
870
871    async fn listen(
872        mut store: Access<'_, T, Self>,
873        socket: Resource<p3_types::TcpSocket>,
874    ) -> P3SocketResult<wasmtime::component::StreamReader<Resource<p3_types::TcpSocket>>> {
875        let getter = store.get().getter;
876        let wasi_store = Access::<T, WasiSockets>::new(store.as_context_mut(), getter);
877        <WasiSockets as HostTcpSocketWithStore<T>>::listen(wasi_store, socket).await
878    }
879
880    fn send(
881        mut store: Access<'_, T, Self>,
882        socket: Resource<p3_types::TcpSocket>,
883        data: wasmtime::component::StreamReader<u8>,
884    ) -> wasmtime::Result<wasmtime::component::FutureReader<Result<(), p3_ErrorCode>>> {
885        let getter = store.get().getter;
886        let wasi_store = Access::<T, WasiSockets>::new(store.as_context_mut(), getter);
887        <WasiSockets as HostTcpSocketWithStore<T>>::send(wasi_store, socket, data)
888    }
889
890    fn receive(
891        mut store: Access<'_, T, Self>,
892        socket: Resource<p3_types::TcpSocket>,
893    ) -> wasmtime::Result<(
894        wasmtime::component::StreamReader<u8>,
895        wasmtime::component::FutureReader<Result<(), p3_ErrorCode>>,
896    )> {
897        let getter = store.get().getter;
898        let wasi_store = Access::<T, WasiSockets>::new(store.as_context_mut(), getter);
899        <WasiSockets as HostTcpSocketWithStore<T>>::receive(wasi_store, socket)
900    }
901}
902
903impl<T: 'static> HostUdpSocketWithStore<T> for SpinSockets<T> {
904    async fn send(
905        store: &Accessor<T, Self>,
906        socket: Resource<p3_types::UdpSocket>,
907        data: Vec<u8>,
908        remote_address: Option<p3_IpSocketAddress>,
909    ) -> P3SocketResult<()> {
910        let getter = store.with(|mut store| store.get().getter);
911        let wasi_accessor = store.with_getter::<WasiSockets>(getter);
912        <WasiSockets as HostUdpSocketWithStore<T>>::send(
913            &wasi_accessor,
914            socket,
915            data,
916            remote_address,
917        )
918        .await
919    }
920
921    async fn receive(
922        store: &Accessor<T, Self>,
923        socket: Resource<p3_types::UdpSocket>,
924    ) -> P3SocketResult<(Vec<u8>, p3_IpSocketAddress)> {
925        let getter = store.with(|mut store| store.get().getter);
926        let wasi_accessor = store.with_getter::<WasiSockets>(getter);
927        <WasiSockets as HostUdpSocketWithStore<T>>::receive(&wasi_accessor, socket).await
928    }
929}
spin_factor_wasi/sockets.rs

spin_factor_wasi/
sockets.rs
