A better Rust ATProto crate
raw and loosely typed data decode options
+132
-29
+132
-29
crates/jacquard-common/src/xrpc/subscription.rs
+132
-29
crates/jacquard-common/src/xrpc/subscription.rs
···
3
3
//! This module defines traits and types for typed WebSocket subscriptions,
4
4
//! mirroring the request/response pattern used for HTTP XRPC endpoints.
5
5
6
+
use n0_future::stream::Boxed;
6
7
use serde::{Deserialize, Serialize};
7
8
use std::error::Error;
8
9
use std::future::Future;
···
10
11
use url::Url;
11
12
12
13
use crate::stream::StreamError;
13
-
use crate::websocket::{WebSocketClient, WebSocketConnection};
14
-
use crate::{CowStr, IntoStatic};
14
+
use crate::websocket::{WebSocketClient, WebSocketConnection, WsSink, WsStream};
15
+
use crate::{CowStr, Data, IntoStatic, RawData, WsMessage};
15
16
16
17
/// Encoding format for subscription messages
17
18
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
···
173
174
pub fn into_stream(
174
175
self,
175
176
) -> (
176
-
crate::websocket::WsSink,
177
-
n0_future::stream::Boxed<Result<StreamMessage<'static, S>, StreamError>>,
177
+
WsSink,
178
+
Boxed<Result<StreamMessage<'static, S>, StreamError>>,
178
179
)
179
180
where
180
181
for<'a> StreamMessage<'a, S>: IntoStatic<Output = StreamMessage<'static, S>>,
···
183
184
184
185
let (tx, rx) = self.connection.split();
185
186
186
-
let stream: n0_future::stream::Boxed<_> = match S::ENCODING {
187
-
MessageEncoding::Json => {
188
-
Box::pin(rx.into_inner().filter_map(|msg| decode_json_msg::<S>(msg)))
189
-
}
190
-
MessageEncoding::DagCbor => {
191
-
Box::pin(rx.into_inner().filter_map(|msg| decode_cbor_msg::<S>(msg)))
192
-
}
187
+
let stream = match S::ENCODING {
188
+
MessageEncoding::Json => rx
189
+
.into_inner()
190
+
.filter_map(|msg| decode_json_msg::<S>(msg))
191
+
.boxed(),
192
+
MessageEncoding::DagCbor => rx
193
+
.into_inner()
194
+
.filter_map(|msg| decode_cbor_msg::<S>(msg))
195
+
.boxed(),
196
+
};
197
+
198
+
(tx, stream)
199
+
}
200
+
201
+
/// Converts the subscription into a stream of raw atproto data.
202
+
pub fn into_raw_data_stream(self) -> (WsSink, Boxed<Result<RawData<'static>, StreamError>>) {
203
+
use n0_future::StreamExt as _;
204
+
205
+
let (tx, rx) = self.connection.split();
206
+
207
+
fn parse_msg<'a>(bytes: &'a [u8]) -> Result<RawData<'a>, serde_json::Error> {
208
+
serde_json::from_slice(bytes)
209
+
}
210
+
fn parse_cbor<'a>(
211
+
bytes: &'a [u8],
212
+
) -> Result<RawData<'a>, serde_ipld_dagcbor::DecodeError<std::convert::Infallible>>
213
+
{
214
+
serde_ipld_dagcbor::from_slice(bytes)
215
+
}
216
+
217
+
let stream = match S::ENCODING {
218
+
MessageEncoding::Json => rx
219
+
.into_inner()
220
+
.filter_map(|msg_result| match msg_result {
221
+
Ok(WsMessage::Text(text)) => Some(
222
+
parse_msg(text.as_ref())
223
+
.map(|v| v.into_static())
224
+
.map_err(StreamError::decode),
225
+
),
226
+
Ok(WsMessage::Binary(bytes)) => Some(
227
+
parse_msg(&bytes)
228
+
.map(|v| v.into_static())
229
+
.map_err(StreamError::decode),
230
+
),
231
+
Ok(WsMessage::Close(_)) => Some(Err(StreamError::closed())),
232
+
Err(e) => Some(Err(e)),
233
+
})
234
+
.boxed(),
235
+
MessageEncoding::DagCbor => rx
236
+
.into_inner()
237
+
.filter_map(|msg_result| match msg_result {
238
+
Ok(WsMessage::Binary(bytes)) => Some(
239
+
parse_cbor(&bytes)
240
+
.map(|v| v.into_static())
241
+
.map_err(|e| StreamError::decode(crate::error::DecodeError::from(e))),
242
+
),
243
+
Ok(WsMessage::Text(_)) => Some(Err(StreamError::wrong_message_format(
244
+
"expected binary frame for CBOR, got text",
245
+
))),
246
+
Ok(WsMessage::Close(_)) => Some(Err(StreamError::closed())),
247
+
Err(e) => Some(Err(e)),
248
+
})
249
+
.boxed(),
250
+
};
251
+
252
+
(tx, stream)
253
+
}
254
+
255
+
/// Converts the subscription into a stream of loosely-typed atproto data.
256
+
pub fn into_data_stream(self) -> (WsSink, Boxed<Result<Data<'static>, StreamError>>) {
257
+
use n0_future::StreamExt as _;
258
+
259
+
let (tx, rx) = self.connection.split();
260
+
261
+
fn parse_msg<'a>(bytes: &'a [u8]) -> Result<Data<'a>, serde_json::Error> {
262
+
serde_json::from_slice(bytes)
263
+
}
264
+
fn parse_cbor<'a>(
265
+
bytes: &'a [u8],
266
+
) -> Result<Data<'a>, serde_ipld_dagcbor::DecodeError<std::convert::Infallible>> {
267
+
serde_ipld_dagcbor::from_slice(bytes)
268
+
}
269
+
270
+
let stream = match S::ENCODING {
271
+
MessageEncoding::Json => rx
272
+
.into_inner()
273
+
.filter_map(|msg_result| match msg_result {
274
+
Ok(WsMessage::Text(text)) => Some(
275
+
parse_msg(text.as_ref())
276
+
.map(|v| v.into_static())
277
+
.map_err(StreamError::decode),
278
+
),
279
+
Ok(WsMessage::Binary(bytes)) => Some(
280
+
parse_msg(&bytes)
281
+
.map(|v| v.into_static())
282
+
.map_err(StreamError::decode),
283
+
),
284
+
Ok(WsMessage::Close(_)) => Some(Err(StreamError::closed())),
285
+
Err(e) => Some(Err(e)),
286
+
})
287
+
.boxed(),
288
+
MessageEncoding::DagCbor => rx
289
+
.into_inner()
290
+
.filter_map(|msg_result| match msg_result {
291
+
Ok(WsMessage::Binary(bytes)) => Some(
292
+
parse_cbor(&bytes)
293
+
.map(|v| v.into_static())
294
+
.map_err(|e| StreamError::decode(crate::error::DecodeError::from(e))),
295
+
),
296
+
Ok(WsMessage::Text(_)) => Some(Err(StreamError::wrong_message_format(
297
+
"expected binary frame for CBOR, got text",
298
+
))),
299
+
Ok(WsMessage::Close(_)) => Some(Err(StreamError::closed())),
300
+
Err(e) => Some(Err(e)),
301
+
})
302
+
.boxed(),
193
303
};
194
304
195
305
(tx, stream)
···
205
315
/// Replaces the internal WebSocket stream with one copy and returns a typed decoded
206
316
/// stream. Both streams receive all messages. Useful for observing raw messages
207
317
/// while also processing typed messages.
208
-
pub fn tee(
209
-
&mut self,
210
-
) -> n0_future::stream::Boxed<Result<StreamMessage<'static, S>, StreamError>>
318
+
pub fn tee(&mut self) -> Boxed<Result<StreamMessage<'static, S>, StreamError>>
211
319
where
212
320
for<'a> StreamMessage<'a, S>: IntoStatic<Output = StreamMessage<'static, S>>,
213
321
{
214
322
use n0_future::StreamExt as _;
215
323
216
324
let rx = self.connection.receiver_mut();
217
-
let (raw_rx, typed_rx_source) = std::mem::replace(
218
-
rx,
219
-
crate::websocket::WsStream::new(futures::stream::empty()),
220
-
)
221
-
.tee();
325
+
let (raw_rx, typed_rx_source) =
326
+
std::mem::replace(rx, WsStream::new(n0_future::stream::empty())).tee();
222
327
223
328
// Put the raw stream back
224
329
*rx = raw_rx;
225
330
226
331
match S::ENCODING {
227
-
MessageEncoding::Json => Box::pin(
228
-
typed_rx_source
229
-
.into_inner()
230
-
.filter_map(|msg| decode_json_msg::<S>(msg)),
231
-
),
232
-
MessageEncoding::DagCbor => Box::pin(
233
-
typed_rx_source
234
-
.into_inner()
235
-
.filter_map(|msg| decode_cbor_msg::<S>(msg)),
236
-
),
332
+
MessageEncoding::Json => typed_rx_source
333
+
.into_inner()
334
+
.filter_map(|msg| decode_json_msg::<S>(msg))
335
+
.boxed(),
336
+
MessageEncoding::DagCbor => typed_rx_source
337
+
.into_inner()
338
+
.filter_map(|msg| decode_cbor_msg::<S>(msg))
339
+
.boxed(),
237
340
}
238
341
}
239
342
}