spin_factor_outbound_mqtt/
host.rs

1use std::{sync::Arc, time::Duration};
2
3use anyhow::Result;
4use spin_core::{
5    async_trait,
6    wasmtime::component::{Accessor, Resource},
7};
8use spin_factor_otel::OtelFactorState;
9use spin_factor_outbound_networking::config::allowed_hosts::OutboundAllowedHosts;
10use spin_factor_outbound_networking::{ConnectionPermit, ConnectionSemaphore};
11use spin_world::spin::mqtt::mqtt as v3;
12use spin_world::v2::mqtt as v2;
13use tracing::{Level, instrument};
14
15use crate::{ClientCreator, allowed_hosts::AllowedHostChecker};
16
17pub struct InstanceState {
18    allowed_hosts: AllowedHostChecker,
19    connections: spin_resource_table::Table<(Arc<dyn MqttClient>, ConnectionPermit)>,
20    create_client: Arc<dyn ClientCreator>,
21    semaphore: ConnectionSemaphore,
22    otel: OtelFactorState,
23    max_payload_size_bytes: Option<usize>,
24}
25
26impl InstanceState {
27    pub fn new(
28        allowed_hosts: OutboundAllowedHosts,
29        create_client: Arc<dyn ClientCreator>,
30        semaphore: ConnectionSemaphore,
31        otel: OtelFactorState,
32        max_payload_size_bytes: Option<usize>,
33    ) -> Self {
34        Self {
35            allowed_hosts: AllowedHostChecker::new(allowed_hosts),
36            create_client,
37            connections: spin_resource_table::Table::new(1024),
38            semaphore,
39            otel,
40            max_payload_size_bytes,
41        }
42    }
43}
44
45#[async_trait]
46pub trait MqttClient: Send + Sync {
47    async fn publish_bytes(
48        &self,
49        topic: String,
50        qos: v3::Qos,
51        payload: Vec<u8>,
52    ) -> Result<(), v3::Error>;
53}
54
55impl InstanceState {
56    async fn is_address_allowed(&self, address: &str) -> Result<bool> {
57        self.allowed_hosts.is_address_allowed(address).await
58    }
59
60    async fn establish_connection(
61        &mut self,
62        address: String,
63        username: String,
64        password: String,
65        keep_alive_interval: Duration,
66    ) -> Result<Resource<v2::Connection>, v2::Error> {
67        let permit = self
68            .semaphore
69            .acquire()
70            .await
71            .map_err(|_| v2::Error::TooManyConnections)?;
72        let client =
73            (self.create_client).create(address, username, password, keep_alive_interval)?;
74        self.connections
75            .push((client, permit))
76            .map(Resource::new_own)
77            .map_err(|_| v2::Error::TooManyConnections)
78    }
79
80    fn get_conn(&self, connection: Resource<v2::Connection>) -> Result<&dyn MqttClient, v2::Error> {
81        self.connections
82            .get(connection.rep())
83            .ok_or(v2::Error::Other(
84                "could not find connection for resource".into(),
85            ))
86            .map(|(c, _permit)| c.as_ref())
87    }
88
89    fn get_conn_v3(
90        &self,
91        connection: Resource<v3::Connection>,
92    ) -> Result<Arc<dyn MqttClient>, v3::Error> {
93        self.connections
94            .get(connection.rep())
95            .map(|(c, _permit)| c.clone())
96            .ok_or(v3::Error::Other(
97                "could not find connection for resource".into(),
98            ))
99    }
100}
101
102impl v3::Host for InstanceState {
103    fn convert_error(&mut self, err: v3::Error) -> anyhow::Result<v3::Error> {
104        Ok(err)
105    }
106}
107
108impl v3::HostConnection for InstanceState {
109    async fn drop(&mut self, connection: Resource<v3::Connection>) -> anyhow::Result<()> {
110        self.connections.remove(connection.rep());
111        Ok(())
112    }
113}
114
115impl<T: Send> v3::HostConnectionWithStore<T> for crate::MqttFactorData {
116    #[instrument(name = "spin_outbound_mqtt.open_connection", skip(accessor, password), err(level = Level::INFO), fields(otel.kind = "client"))]
117    async fn open(
118        accessor: &Accessor<T, Self>,
119        address: String,
120        username: String,
121        password: String,
122        keep_alive_interval_in_secs: u64,
123    ) -> Result<Resource<v3::Connection>, v3::Error> {
124        let (allowed_host_checker, create_client, semaphore) = accessor.with(|mut access| {
125            let host = access.get();
126            host.otel.reparent_tracing_span();
127            (
128                host.allowed_hosts.clone(),
129                host.create_client.clone(),
130                host.semaphore.clone(),
131            )
132        });
133
134        if !allowed_host_checker
135            .is_address_allowed(&address)
136            .await
137            .map_err(|e| v3::Error::Other(e.to_string()))?
138        {
139            return Err(v3::Error::ConnectionFailed(format!(
140                "address {address} is not permitted"
141            )));
142        }
143
144        let permit = semaphore
145            .acquire()
146            .await
147            .map_err(|_| v3::Error::TooManyConnections)?;
148
149        let client = create_client.create(
150            address,
151            username,
152            password,
153            Duration::from_secs(keep_alive_interval_in_secs),
154        )?;
155
156        accessor.with(|mut access| {
157            let host = access.get();
158            host.connections
159                .push((client, permit))
160                .map(Resource::new_own)
161                .map_err(|_| v3::Error::TooManyConnections)
162        })
163    }
164
165    #[instrument(name = "spin_outbound_mqtt.publish", skip(accessor, connection, payload), err(level = Level::INFO),
166        fields(
167            otel.kind = "producer",
168            messaging.operation = "publish",
169            messaging.system = "mqtt",
170            messaging.destination.name = topic,
171        ))]
172    async fn publish(
173        accessor: &Accessor<T, Self>,
174        connection: Resource<v3::Connection>,
175        topic: String,
176        payload: v3::Payload,
177        qos: v3::Qos,
178    ) -> Result<(), v3::Error> {
179        let (conn, max_payload_size_bytes) = accessor.with(|mut access| {
180            let host = access.get();
181            host.otel.reparent_tracing_span();
182            host.get_conn_v3(connection)
183                .map(|c| (c, host.max_payload_size_bytes))
184        })?;
185
186        if let Some(limit) = max_payload_size_bytes
187            && payload.len() > limit
188        {
189            return Err(v3::Error::Other(format!(
190                "payload size {} exceeds the maximum allowed size of {} bytes",
191                payload.len(),
192                limit
193            )));
194        }
195
196        conn.publish_bytes(topic, qos, payload).await?;
197
198        Ok(())
199    }
200}
201
202impl v2::Host for InstanceState {
203    fn convert_error(&mut self, error: v2::Error) -> Result<v2::Error> {
204        Ok(error)
205    }
206}
207
208impl v2::HostConnection for InstanceState {
209    #[instrument(name = "spin_outbound_mqtt.open_connection", skip(self, password), err(level = Level::INFO), fields(otel.kind = "client"))]
210    async fn open(
211        &mut self,
212        address: String,
213        username: String,
214        password: String,
215        keep_alive_interval: u64,
216    ) -> Result<Resource<v2::Connection>, v2::Error> {
217        self.otel.reparent_tracing_span();
218
219        if !self
220            .is_address_allowed(&address)
221            .await
222            .map_err(|e| v2::Error::Other(e.to_string()))?
223        {
224            return Err(v2::Error::ConnectionFailed(format!(
225                "address {address} is not permitted"
226            )));
227        }
228        self.establish_connection(
229            address,
230            username,
231            password,
232            Duration::from_secs(keep_alive_interval),
233        )
234        .await
235    }
236
237    /// Publish a message to the MQTT broker.
238    ///
239    /// OTEL trace propagation is not directly supported in MQTT V3. You will need to embed the
240    /// current trace context into the payload yourself.
241    /// https://w3c.github.io/trace-context-mqtt/#mqtt-v3-recommendation.
242    #[instrument(name = "spin_outbound_mqtt.publish", skip(self, connection, payload), err(level = Level::INFO),
243        fields(
244            otel.kind = "producer",
245            messaging.operation = "publish",
246            messaging.system = "mqtt",
247            messaging.destination.name = topic,
248        ))]
249    async fn publish(
250        &mut self,
251        connection: Resource<v2::Connection>,
252        topic: String,
253        payload: Vec<u8>,
254        qos: v2::Qos,
255    ) -> Result<(), v2::Error> {
256        self.otel.reparent_tracing_span();
257
258        if let Some(limit) = self.max_payload_size_bytes
259            && payload.len() > limit
260        {
261            return Err(v2::Error::Other(format!(
262                "payload size {} exceeds the maximum allowed size of {} bytes",
263                payload.len(),
264                limit
265            )));
266        }
267
268        let conn = self.get_conn(connection)?;
269
270        let qos = match qos {
271            v2::Qos::AtMostOnce => v3::Qos::AtMostOnce,
272            v2::Qos::AtLeastOnce => v3::Qos::AtLeastOnce,
273            v2::Qos::ExactlyOnce => v3::Qos::ExactlyOnce,
274        };
275
276        conn.publish_bytes(topic, qos, payload).await?;
277
278        Ok(())
279    }
280
281    async fn drop(&mut self, connection: Resource<v2::Connection>) -> anyhow::Result<()> {
282        self.connections.remove(connection.rep());
283        Ok(())
284    }
285}
286
287pub fn other_error_v3(e: impl std::fmt::Display) -> v3::Error {
288    v3::Error::Other(e.to_string())
289}
spin_factor_outbound_mqtt/host.rs

spin_factor_outbound_mqtt/
host.rs
