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

async_graphql_axum/
subscription.rs
