2022-04-19 04:25:11 +00:00
use std ::{
future ::Future ,
str ::FromStr ,
time ::{ Duration , Instant } ,
} ;
2020-12-04 04:13:52 +00:00
2020-03-17 09:26:59 +00:00
use actix ::{
2022-04-19 04:25:11 +00:00
Actor , ActorContext , ActorFutureExt , ActorStreamExt , AsyncContext , ContextFutureSpawner ,
StreamHandler , WrapFuture , WrapStream ,
2020-03-17 09:26:59 +00:00
} ;
2022-04-19 04:25:11 +00:00
use actix_http ::{ error ::PayloadError , ws } ;
use actix_web ::{ web ::Bytes , Error , HttpRequest , HttpResponse } ;
2021-02-06 03:02:41 +00:00
use actix_web_actors ::ws ::{ CloseReason , Message , ProtocolError , WebsocketContext } ;
2022-04-19 04:25:11 +00:00
use async_graphql ::{
http ::{ WebSocket , WebSocketProtocols , WsMessage , ALL_WEBSOCKET_PROTOCOLS } ,
Data , ObjectType , Result , Schema , SubscriptionType ,
} ;
use futures_util ::{ future ::Ready , stream ::Stream } ;
2021-10-22 14:01:14 +00:00
2020-04-07 06:30:46 +00:00
const HEARTBEAT_INTERVAL : Duration = Duration ::from_secs ( 5 ) ;
const CLIENT_TIMEOUT : Duration = Duration ::from_secs ( 10 ) ;
2021-11-12 13:24:24 +00:00
#[ derive(thiserror::Error, Debug) ]
#[ error( " failed to parse graphql protocol " ) ]
pub struct ParseGraphQLProtocolError ;
type DefaultOnConnInitType = fn ( serde_json ::Value ) -> Ready < async_graphql ::Result < Data > > ;
fn default_on_connection_init ( _ : serde_json ::Value ) -> Ready < async_graphql ::Result < Data > > {
futures_util ::future ::ready ( Ok ( Data ::default ( ) ) )
2021-11-12 08:58:13 +00:00
}
2021-11-12 13:24:24 +00:00
/// A builder for websocket subscription actor.
pub struct GraphQLSubscription < Query , Mutation , Subscription , OnInit > {
2020-12-04 04:13:52 +00:00
schema : Schema < Query , Mutation , Subscription > ,
2021-11-12 13:24:24 +00:00
data : Data ,
on_connection_init : OnInit ,
2020-03-17 09:26:59 +00:00
}
2021-01-16 04:47:32 +00:00
impl < Query , Mutation , Subscription >
2021-11-12 13:24:24 +00:00
GraphQLSubscription < Query , Mutation , Subscription , DefaultOnConnInitType >
2020-03-17 09:26:59 +00:00
where
2021-01-14 04:41:59 +00:00
Query : ObjectType + 'static ,
Mutation : ObjectType + 'static ,
Subscription : SubscriptionType + 'static ,
2020-03-17 09:26:59 +00:00
{
2021-11-12 13:24:24 +00:00
/// Create a GraphQL subscription builder.
pub fn new ( schema : Schema < Query , Mutation , Subscription > ) -> Self {
Self {
schema ,
data : Default ::default ( ) ,
on_connection_init : default_on_connection_init ,
}
2020-12-04 04:35:35 +00:00
}
2021-01-16 04:47:32 +00:00
}
2020-12-04 04:35:35 +00:00
2021-11-12 13:24:24 +00:00
impl < Query , Mutation , Subscription , OnInit , OnInitFut >
GraphQLSubscription < Query , Mutation , Subscription , OnInit >
2021-01-16 04:47:32 +00:00
where
Query : ObjectType + 'static ,
Mutation : ObjectType + 'static ,
Subscription : SubscriptionType + 'static ,
2022-08-15 04:10:03 +00:00
OnInit : FnOnce ( serde_json ::Value ) -> OnInitFut + Unpin + Send + 'static ,
2021-11-12 13:24:24 +00:00
OnInitFut : Future < Output = async_graphql ::Result < Data > > + Send + 'static ,
2021-01-16 04:47:32 +00:00
{
2022-04-19 04:25:11 +00:00
/// Specify the initial subscription context data, usually you can get
/// something from the incoming request to create it.
2022-01-24 06:14:07 +00:00
#[ must_use ]
2021-11-12 13:24:24 +00:00
pub fn with_data ( self , data : Data ) -> Self {
Self { data , .. self }
}
2022-04-19 04:25:11 +00:00
/// Specify a callback function to be called when the connection is
/// initialized.
2021-11-12 13:24:24 +00:00
///
/// 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`].
2022-04-19 04:25:11 +00:00
/// The data returned by this callback function will be merged with the data
/// specified by [`with_data`].
2021-11-12 13:24:24 +00:00
pub fn on_connection_init < OnConnInit2 , Fut > (
self ,
callback : OnConnInit2 ,
) -> GraphQLSubscription < Query , Mutation , Subscription , OnConnInit2 >
where
2022-08-15 04:10:03 +00:00
OnConnInit2 : FnOnce ( serde_json ::Value ) -> Fut + Unpin + Send + 'static ,
2021-11-12 13:24:24 +00:00
Fut : Future < Output = async_graphql ::Result < Data > > + Send + 'static ,
{
GraphQLSubscription {
schema : self . schema ,
data : self . data ,
on_connection_init : callback ,
}
}
/// Start the subscription actor.
pub fn start < S > ( self , request : & HttpRequest , stream : S ) -> Result < HttpResponse , Error >
2020-12-04 04:35:35 +00:00
where
2021-11-12 13:24:24 +00:00
S : Stream < Item = Result < Bytes , PayloadError > > + 'static ,
2020-12-04 04:13:52 +00:00
{
2021-11-12 13:24:24 +00:00
let protocol = request
2020-12-04 04:13:52 +00:00
. headers ( )
. get ( " sec-websocket-protocol " )
. and_then ( | value | value . to_str ( ) . ok ( ) )
2020-12-04 17:16:14 +00:00
. and_then ( | protocols | {
protocols
. split ( ',' )
. find_map ( | p | WebSocketProtocols ::from_str ( p . trim ( ) ) . ok ( ) )
2021-11-12 13:24:24 +00:00
} )
. ok_or_else ( | | actix_web ::error ::ErrorBadRequest ( ParseGraphQLProtocolError ) ) ? ;
let actor = GraphQLSubscriptionActor {
schema : self . schema ,
data : Some ( self . data ) ,
protocol ,
last_heartbeat : Instant ::now ( ) ,
messages : None ,
on_connection_init : Some ( self . on_connection_init ) ,
continuation : Vec ::new ( ) ,
2020-12-04 04:13:52 +00:00
} ;
2021-12-12 02:10:28 +00:00
actix_web_actors ::ws ::WsResponseBuilder ::new ( actor , request , stream )
. protocols ( & ALL_WEBSOCKET_PROTOCOLS )
. start ( )
2020-04-23 07:30:12 +00:00
}
2021-11-12 13:24:24 +00:00
}
struct GraphQLSubscriptionActor < Query , Mutation , Subscription , OnInit > {
schema : Schema < Query , Mutation , Subscription > ,
data : Option < Data > ,
protocol : WebSocketProtocols ,
last_heartbeat : Instant ,
messages : Option < async_channel ::Sender < Vec < u8 > > > ,
on_connection_init : Option < OnInit > ,
continuation : Vec < u8 > ,
}
2020-04-23 07:30:12 +00:00
2021-11-12 13:24:24 +00:00
impl < Query , Mutation , Subscription , OnInit , OnInitFut >
GraphQLSubscriptionActor < Query , Mutation , Subscription , OnInit >
where
Query : ObjectType + 'static ,
Mutation : ObjectType + 'static ,
Subscription : SubscriptionType + 'static ,
OnInit : FnOnce ( serde_json ::Value ) -> OnInitFut + Unpin + Send + 'static ,
OnInitFut : Future < Output = Result < Data > > + Send + 'static ,
{
2020-09-17 18:22:54 +00:00
fn send_heartbeats ( & self , ctx : & mut WebsocketContext < Self > ) {
2020-04-07 06:30:46 +00:00
ctx . run_interval ( HEARTBEAT_INTERVAL , | act , ctx | {
2020-09-17 18:22:54 +00:00
if Instant ::now ( ) . duration_since ( act . last_heartbeat ) > CLIENT_TIMEOUT {
2020-03-19 09:20:12 +00:00
ctx . stop ( ) ;
}
2020-04-07 06:30:46 +00:00
ctx . ping ( b " " ) ;
2020-03-19 09:20:12 +00:00
} ) ;
}
2020-03-17 09:26:59 +00:00
}
2021-11-12 13:24:24 +00:00
impl < Query , Mutation , Subscription , OnInit , OnInitFut > Actor
for GraphQLSubscriptionActor < Query , Mutation , Subscription , OnInit >
2020-03-17 09:26:59 +00:00
where
2021-01-14 04:41:59 +00:00
Query : ObjectType + 'static ,
Mutation : ObjectType + 'static ,
Subscription : SubscriptionType + 'static ,
2021-11-12 13:24:24 +00:00
OnInit : FnOnce ( serde_json ::Value ) -> OnInitFut + Unpin + Send + 'static ,
OnInitFut : Future < Output = Result < Data > > + Send + 'static ,
2020-03-17 09:26:59 +00:00
{
type Context = WebsocketContext < Self > ;
fn started ( & mut self , ctx : & mut Self ::Context ) {
2020-09-17 18:22:54 +00:00
self . send_heartbeats ( ctx ) ;
2020-10-15 09:33:38 +00:00
let ( tx , rx ) = async_channel ::unbounded ( ) ;
2020-09-17 18:22:54 +00:00
2021-11-12 13:24:24 +00:00
WebSocket ::new ( self . schema . clone ( ) , rx , self . protocol )
. connection_data ( self . data . take ( ) . unwrap ( ) )
. on_connection_init ( self . on_connection_init . take ( ) . unwrap ( ) )
. into_actor ( self )
. map ( | response , _act , ctx | match response {
WsMessage ::Text ( text ) = > ctx . text ( text ) ,
WsMessage ::Close ( code , msg ) = > ctx . close ( Some ( CloseReason {
code : code . into ( ) ,
description : Some ( msg ) ,
} ) ) ,
} )
. finish ( )
. spawn ( ctx ) ;
2020-09-17 18:22:54 +00:00
self . messages = Some ( tx ) ;
2020-03-17 09:26:59 +00:00
}
}
2021-11-12 13:24:24 +00:00
impl < Query , Mutation , Subscription , OnInit , OnInitFut > StreamHandler < Result < Message , ProtocolError > >
for GraphQLSubscriptionActor < Query , Mutation , Subscription , OnInit >
2020-03-17 09:26:59 +00:00
where
2021-01-14 04:41:59 +00:00
Query : ObjectType + 'static ,
Mutation : ObjectType + 'static ,
Subscription : SubscriptionType + 'static ,
2021-11-12 13:24:24 +00:00
OnInit : FnOnce ( serde_json ::Value ) -> OnInitFut + Unpin + Send + 'static ,
OnInitFut : Future < Output = Result < Data > > + Send + 'static ,
2020-03-17 09:26:59 +00:00
{
fn handle ( & mut self , msg : Result < Message , ProtocolError > , ctx : & mut Self ::Context ) {
let msg = match msg {
Err ( _ ) = > {
ctx . stop ( ) ;
return ;
}
Ok ( msg ) = > msg ,
} ;
2020-09-17 18:22:54 +00:00
let message = match msg {
2020-03-17 09:26:59 +00:00
Message ::Ping ( msg ) = > {
2020-09-17 18:22:54 +00:00
self . last_heartbeat = Instant ::now ( ) ;
2020-03-17 09:26:59 +00:00
ctx . pong ( & msg ) ;
2020-09-17 18:22:54 +00:00
None
2020-03-17 09:26:59 +00:00
}
Message ::Pong ( _ ) = > {
2020-09-17 18:22:54 +00:00
self . last_heartbeat = Instant ::now ( ) ;
None
2020-03-17 09:26:59 +00:00
}
2020-09-17 18:22:54 +00:00
Message ::Continuation ( item ) = > match item {
ws ::Item ::FirstText ( bytes ) | ws ::Item ::FirstBinary ( bytes ) = > {
self . continuation = bytes . to_vec ( ) ;
None
2020-03-17 09:26:59 +00:00
}
2020-09-17 18:22:54 +00:00
ws ::Item ::Continue ( bytes ) = > {
self . continuation . extend_from_slice ( & bytes ) ;
None
}
ws ::Item ::Last ( bytes ) = > {
self . continuation . extend_from_slice ( & bytes ) ;
Some ( std ::mem ::take ( & mut self . continuation ) )
}
} ,
2021-10-22 14:01:14 +00:00
Message ::Text ( s ) = > Some ( s . into_bytes ( ) . to_vec ( ) ) ,
2020-09-17 18:22:54 +00:00
Message ::Binary ( bytes ) = > Some ( bytes . to_vec ( ) ) ,
Message ::Close ( _ ) = > {
2020-03-17 09:26:59 +00:00
ctx . stop ( ) ;
2020-09-17 18:22:54 +00:00
None
2020-03-17 09:26:59 +00:00
}
2020-09-17 18:22:54 +00:00
Message ::Nop = > None ,
} ;
2020-03-17 09:26:59 +00:00
2020-09-17 18:22:54 +00:00
if let Some ( message ) = message {
2020-10-15 09:33:38 +00:00
let sender = self . messages . as_ref ( ) . unwrap ( ) . clone ( ) ;
2020-09-17 18:22:54 +00:00
async move { sender . send ( message ) . await }
. into_actor ( self )
. map ( | res , _actor , ctx | match res {
Ok ( ( ) ) = > { }
Err ( _ ) = > ctx . stop ( ) ,
} )
. spawn ( ctx )
}
2020-03-17 09:26:59 +00:00
}
2021-11-12 08:58:13 +00:00
}