gqlrs_axum/
subscription.rs

1use std::{convert::Infallible, future::Future, str::FromStr, time::Duration};
2
3use async_graphql::{
4    Data, Executor, Result,
5    futures_util::task::{Context, Poll},
6    http::{
7        ALL_WEBSOCKET_PROTOCOLS, DefaultOnConnInitType, DefaultOnPingType, WebSocketProtocols,
8        WsMessage, default_on_connection_init, default_on_ping,
9    },
10    runtime::TokioTimer,
11};
12use axum::{
13    Error,
14    body::{Body, HttpBody},
15    extract::{
16        FromRequestParts, WebSocketUpgrade,
17        ws::{CloseFrame, Message},
18    },
19    http::{self, Request, Response, StatusCode, request::Parts},
20    response::IntoResponse,
21};
22use futures_util::{
23    Sink, SinkExt, Stream, StreamExt, future,
24    future::BoxFuture,
25    stream::{SplitSink, SplitStream},
26};
27use tower_service::Service;
28#[cfg(feature = "tracing")]
29use tracing::{Instrument, Span};
30
31/// A GraphQL protocol extractor.
32///
33/// It extract GraphQL protocol from `SEC_WEBSOCKET_PROTOCOL` header.
34#[derive(Debug, Copy, Clone, PartialEq, Eq)]
35pub struct GraphQLProtocol(WebSocketProtocols);
36
37impl<S> FromRequestParts<S> for GraphQLProtocol
38where
39    S: Send + Sync,
40{
41    type Rejection = StatusCode;
42
43    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
44        parts
45            .headers
46            .get(http::header::SEC_WEBSOCKET_PROTOCOL)
47            .and_then(|value| value.to_str().ok())
48            .and_then(|protocols| {
49                protocols
50                    .split(',')
51                    .find_map(|p| WebSocketProtocols::from_str(p.trim()).ok())
52            })
53            .map(Self)
54            .ok_or(StatusCode::BAD_REQUEST)
55    }
56}
57
58/// A GraphQL subscription service.
59pub struct GraphQLSubscription<E> {
60    executor: E,
61}
62
63impl<E> Clone for GraphQLSubscription<E>
64where
65    E: Executor,
66{
67    fn clone(&self) -> Self {
68        Self {
69            executor: self.executor.clone(),
70        }
71    }
72}
73
74impl<E> GraphQLSubscription<E>
75where
76    E: Executor,
77{
78    /// Create a GraphQL subscription service.
79    pub fn new(executor: E) -> Self {
80        Self { executor }
81    }
82}
83
84impl<B, E> Service<Request<B>> for GraphQLSubscription<E>
85where
86    B: HttpBody + Send + 'static,
87    E: Executor,
88{
89    type Response = Response<Body>;
90    type Error = Infallible;
91    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
92
93    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
94        Poll::Ready(Ok(()))
95    }
96
97    fn call(&mut self, req: Request<B>) -> Self::Future {
98        let executor = self.executor.clone();
99
100        Box::pin(async move {
101            let (mut parts, _body) = req.into_parts();
102
103            let protocol = match GraphQLProtocol::from_request_parts(&mut parts, &()).await {
104                Ok(protocol) => protocol,
105                Err(err) => return Ok(err.into_response()),
106            };
107            let upgrade = match WebSocketUpgrade::from_request_parts(&mut parts, &()).await {
108                Ok(upgrade) => upgrade,
109                Err(err) => return Ok(err.into_response()),
110            };
111
112            let executor = executor.clone();
113
114            #[cfg(feature = "tracing")]
115            let span = Span::current();
116
117            let resp = upgrade
118                .protocols(ALL_WEBSOCKET_PROTOCOLS)
119                .on_upgrade(move |stream| {
120                    let task = GraphQLWebSocket::new(stream, executor, protocol).serve();
121                    #[cfg(feature = "tracing")]
122                    let task = task.instrument(span);
123                    task
124                });
125            Ok(resp.into_response())
126        })
127    }
128}
129
130/// A Websocket connection for GraphQL subscription.
131pub struct GraphQLWebSocket<Sink, Stream, E, OnConnInit, OnPing> {
132    sink: Sink,
133    stream: Stream,
134    executor: E,
135    data: Data,
136    on_connection_init: OnConnInit,
137    on_ping: OnPing,
138    protocol: GraphQLProtocol,
139    keepalive_timeout: Option<Duration>,
140}
141
142impl<S, E>
143    GraphQLWebSocket<
144        SplitSink<S, Message>,
145        SplitStream<S>,
146        E,
147        DefaultOnConnInitType,
148        DefaultOnPingType,
149    >
150where
151    S: Stream<Item = Result<Message, Error>> + Sink<Message>,
152    E: Executor,
153{
154    /// Create a [`GraphQLWebSocket`] object.
155    pub fn new(stream: S, executor: E, protocol: GraphQLProtocol) -> Self {
156        let (sink, stream) = stream.split();
157        GraphQLWebSocket::new_with_pair(sink, stream, executor, protocol)
158    }
159}
160
161impl<Sink, Stream, E> GraphQLWebSocket<Sink, Stream, E, DefaultOnConnInitType, DefaultOnPingType>
162where
163    Sink: futures_util::sink::Sink<Message>,
164    Stream: futures_util::stream::Stream<Item = Result<Message, Error>>,
165    E: Executor,
166{
167    /// Create a [`GraphQLWebSocket`] object with sink and stream objects.
168    pub fn new_with_pair(
169        sink: Sink,
170        stream: Stream,
171        executor: E,
172        protocol: GraphQLProtocol,
173    ) -> Self {
174        GraphQLWebSocket {
175            sink,
176            stream,
177            executor,
178            data: Data::default(),
179            on_connection_init: default_on_connection_init,
180            on_ping: default_on_ping,
181            protocol,
182            keepalive_timeout: None,
183        }
184    }
185}
186
187impl<Sink, Stream, E, OnConnInit, OnConnInitFut, OnPing, OnPingFut>
188    GraphQLWebSocket<Sink, Stream, E, OnConnInit, OnPing>
189where
190    Sink: futures_util::sink::Sink<Message>,
191    Stream: futures_util::stream::Stream<Item = Result<Message, Error>>,
192    E: Executor,
193    OnConnInit: FnOnce(serde_json::Value) -> OnConnInitFut + Send + 'static,
194    OnConnInitFut: Future<Output = async_graphql::Result<Data>> + Send + 'static,
195    OnPing: FnOnce(Option<&Data>, Option<serde_json::Value>) -> OnPingFut + Clone + Send + 'static,
196    OnPingFut: Future<Output = async_graphql::Result<Option<serde_json::Value>>> + Send + 'static,
197{
198    /// Specify the initial subscription context data, usually you can get
199    /// something from the incoming request to create it.
200    #[must_use]
201    pub fn with_data(self, data: Data) -> Self {
202        Self { data, ..self }
203    }
204
205    /// Specify a callback function to be called when the connection is
206    /// initialized.
207    ///
208    /// You can get something from the payload of [`GQL_CONNECTION_INIT` message](https://github.com/apollographql/subscriptions-transport-ws/blob/master/PROTOCOL.md#gql_connection_init) to create [`Data`].
209    /// The data returned by this callback function will be merged with the data
210    /// specified by [`with_data`].
211    #[must_use]
212    pub fn on_connection_init<F, R>(
213        self,
214        callback: F,
215    ) -> GraphQLWebSocket<Sink, Stream, E, F, OnPing>
216    where
217        F: FnOnce(serde_json::Value) -> R + Send + 'static,
218        R: Future<Output = async_graphql::Result<Data>> + Send + 'static,
219    {
220        GraphQLWebSocket {
221            sink: self.sink,
222            stream: self.stream,
223            executor: self.executor,
224            data: self.data,
225            on_connection_init: callback,
226            on_ping: self.on_ping,
227            protocol: self.protocol,
228            keepalive_timeout: self.keepalive_timeout,
229        }
230    }
231
232    /// Specify a ping callback function.
233    ///
234    /// This function if present, will be called with the data sent by the
235    /// client in the [`Ping` message](https://github.com/enisdenjo/graphql-ws/blob/master/PROTOCOL.md#ping).
236    ///
237    /// The function should return the data to be sent in the [`Pong` message](https://github.com/enisdenjo/graphql-ws/blob/master/PROTOCOL.md#pong).
238    ///
239    /// NOTE: Only used for the `graphql-ws` protocol.
240    #[must_use]
241    pub fn on_ping<F, R>(self, callback: F) -> GraphQLWebSocket<Sink, Stream, E, OnConnInit, F>
242    where
243        F: FnOnce(Option<&Data>, Option<serde_json::Value>) -> R + Clone + Send + 'static,
244        R: Future<Output = Result<Option<serde_json::Value>>> + Send + 'static,
245    {
246        GraphQLWebSocket {
247            sink: self.sink,
248            stream: self.stream,
249            executor: self.executor,
250            data: self.data,
251            on_connection_init: self.on_connection_init,
252            on_ping: callback,
253            protocol: self.protocol,
254            keepalive_timeout: self.keepalive_timeout,
255        }
256    }
257
258    /// Sets a timeout for receiving an acknowledgement of the keep-alive ping.
259    ///
260    /// If the ping is not acknowledged within the timeout, the connection will
261    /// be closed.
262    ///
263    /// NOTE: Only used for the `graphql-ws` protocol.
264    #[must_use]
265    pub fn keepalive_timeout(self, timeout: impl Into<Option<Duration>>) -> Self {
266        Self {
267            keepalive_timeout: timeout.into(),
268            ..self
269        }
270    }
271
272    /// Processing subscription requests.
273    pub async fn serve(self) {
274        let input = self
275            .stream
276            .take_while(|res| future::ready(res.is_ok()))
277            .map(Result::unwrap)
278            .filter_map(|msg| {
279                if let Message::Text(_) | Message::Binary(_) = msg {
280                    future::ready(Some(msg))
281                } else {
282                    future::ready(None)
283                }
284            })
285            .map(Message::into_data);
286
287        let stream =
288            async_graphql::http::WebSocket::new(self.executor.clone(), input, self.protocol.0)
289                .connection_data(self.data)
290                .on_connection_init(self.on_connection_init)
291                .on_ping(self.on_ping.clone())
292                .keepalive_timeout(TokioTimer::default(), self.keepalive_timeout)
293                .map(|msg| match msg {
294                    WsMessage::Text(text) => Message::Text(text.into()),
295                    WsMessage::Close(code, status) => Message::Close(Some(CloseFrame {
296                        code,
297                        reason: status.into(),
298                    })),
299                });
300
301        let sink = self.sink;
302        futures_util::pin_mut!(stream, sink);
303
304        while let Some(item) = stream.next().await {
305            if sink.send(item).await.is_err() {
306                break;
307            }
308        }
309    }
310}
gqlrs_axum/subscription.rs

gqlrs_axum/
subscription.rs
