1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

use crate::http::{GQLError, GQLRequest, GQLResponse};
use crate::{
    ObjectType, QueryResult, Result, Schema, SubscriptionStubs, SubscriptionTransport,
    SubscriptionType, Variables,
};
use bytes::Bytes;
use std::collections::HashMap;

#[derive(Serialize, Deserialize)]
struct OperationMessage {
    #[serde(rename = "type")]
    ty: String,
    id: Option<String>,
    payload: Option<serde_json::Value>,
}

/// WebSocket transport
#[derive(Default)]
pub struct WebSocketTransport {
    id_to_sid: HashMap<String, usize>,
    sid_to_id: HashMap<usize, String>,
}

impl SubscriptionTransport for WebSocketTransport {
    fn handle_request<Query, Mutation, Subscription>(
        &mut self,
        schema: &Schema<Query, Mutation, Subscription>,
        stubs: &mut SubscriptionStubs<Query, Mutation, Subscription>,
        data: Bytes,
    ) -> Result<Option<Bytes>>
    where
        Query: ObjectType + Sync + Send + 'static,
        Mutation: ObjectType + Sync + Send + 'static,
        Subscription: SubscriptionType + Sync + Send + 'static,
    {
        match serde_json::from_slice::<OperationMessage>(&data) {
            Ok(msg) => match msg.ty.as_str() {
                "connection_init" => Ok(Some(
                    serde_json::to_vec(&OperationMessage {
                        ty: "connection_ack".to_string(),
                        id: None,
                        payload: None,
                    })
                    .unwrap()
                    .into(),
                )),
                "start" => {
                    if let (Some(id), Some(payload)) = (msg.id, msg.payload) {
                        if let Ok(request) = serde_json::from_value::<GQLRequest>(payload) {
                            let variables = if let Some(value) = request.variables {
                                match Variables::parse_from_json(value) {
                                    Ok(variables) => variables,
                                    Err(_) => Default::default(),
                                }
                            } else {
                                Default::default()
                            };

                            match schema.create_subscription_stub(
                                &request.query,
                                request.operation_name.as_deref(),
                                variables,
                            ) {
                                Ok(stub) => {
                                    let stub_id = stubs.add(stub);
                                    self.id_to_sid.insert(id.clone(), stub_id);
                                    self.sid_to_id.insert(stub_id, id);
                                    Ok(None)
                                }
                                Err(err) => Ok(Some(
                                    serde_json::to_vec(&OperationMessage {
                                        ty: "error".to_string(),
                                        id: Some(id),
                                        payload: Some(
                                            serde_json::to_value(GQLError(&err)).unwrap(),
                                        ),
                                    })
                                    .unwrap()
                                    .into(),
                                )),
                            }
                        } else {
                            Ok(None)
                        }
                    } else {
                        Ok(None)
                    }
                }
                "stop" => {
                    if let Some(id) = msg.id {
                        if let Some(id) = self.id_to_sid.remove(&id) {
                            self.sid_to_id.remove(&id);
                            stubs.remove(id);
                        }
                    }
                    Ok(None)
                }
                "connection_terminate" => Err(anyhow::anyhow!("connection_terminate")),
                _ => Err(anyhow::anyhow!("unknown op")),
            },
            Err(err) => Err(err.into()),
        }
    }

    fn handle_response(&mut self, id: usize, result: Result<serde_json::Value>) -> Option<Bytes> {
        if let Some(id) = self.sid_to_id.get(&id) {
            Some(
                serde_json::to_vec(&OperationMessage {
                    ty: "data".to_string(),
                    id: Some(id.clone()),
                    payload: Some(
                        serde_json::to_value(GQLResponse(result.map(|data| QueryResult {
                            data,
                            extensions: None,
                        })))
                        .unwrap(),
                    ),
                })
                .unwrap()
                .into(),
            )
        } else {
            None
        }
    }
}