kernel_sidecar/jupyter/
wire_protocol.rs

1/*
2Python kernel-sidecar uses the jupyter_client.py library to handle the low-level work of converting
3Jupyter Message spec into over-the-wire protocol for ZMQ. There wasn't a jupyter-client crate
4that I wanted to use here, so WireProtocol serialization and deserialization is handled here.
5
6Ref: https://jupyter-client.readthedocs.io/en/latest/messaging.html#the-wire-protocol
7*/
8use crate::jupyter::constants::EMPTY_DICT_BYTES;
9use crate::jupyter::header::Header;
10use bytes::Bytes;
11use ring::hmac;
12use serde::Serialize;
13use zeromq::ZmqMessage;
14
15#[derive(Debug)]
16pub struct WireProtocol {
17    identity: Bytes,
18    delimiter: Bytes,
19    hmac_signature: Bytes,
20    pub header: Bytes,
21    pub parent_header: Bytes,
22    pub metadata: Bytes,
23    pub content: Bytes,
24}
25
26impl WireProtocol {
27    pub fn new<T: Serialize>(header: Header, content: T, hmac_signing_key: &str) -> Self {
28        // Serialize header to JSON then bytes
29        let header = Bytes::from(serde_json::to_vec(&header).expect("Failed to serialize header"));
30        // Make parent_header and metadata both empty dicts serialized to json and then bytes
31        let parent_header = EMPTY_DICT_BYTES.clone();
32        let metadata = EMPTY_DICT_BYTES.clone();
33
34        // If content is passed in, serialize and turn into bytes. Otherwise same as parent_header
35        let content =
36            Bytes::from(serde_json::to_vec(&content).expect("Failed to serialize content"));
37
38        let identity = Bytes::from("kernel");
39        let delimiter = Bytes::from("<IDS|MSG>");
40        let key = Bytes::from(hmac_signing_key.to_owned());
41        let hmac_signature = Self::gen_hmac_sig(&key, &header, &parent_header, &metadata, &content);
42        WireProtocol {
43            identity,
44            delimiter,
45            hmac_signature,
46            header,
47            parent_header,
48            metadata,
49            content,
50        }
51    }
52
53    fn gen_hmac_sig(
54        key: &Bytes,
55        header: &Bytes,
56        parent_header: &Bytes,
57        metadata: &Bytes,
58        content: &Bytes,
59    ) -> Bytes {
60        let key = hmac::Key::new(hmac::HMAC_SHA256, key);
61
62        let mut ctx = hmac::Context::with_key(&key);
63        ctx.update(header);
64        ctx.update(parent_header);
65        ctx.update(metadata);
66        ctx.update(content);
67
68        let tag = ctx.sign();
69        let signature = hex::encode(tag.as_ref());
70        Bytes::from(signature)
71    }
72}
73
74impl From<WireProtocol> for ZmqMessage {
75    fn from(wire_protocol: WireProtocol) -> Self {
76        let mut zmq_message = ZmqMessage::from(wire_protocol.identity);
77        zmq_message.push_back(wire_protocol.delimiter);
78        zmq_message.push_back(wire_protocol.hmac_signature);
79        zmq_message.push_back(wire_protocol.header);
80        zmq_message.push_back(wire_protocol.parent_header);
81        zmq_message.push_back(wire_protocol.metadata);
82        zmq_message.push_back(wire_protocol.content);
83        zmq_message
84    }
85}
86
87impl From<ZmqMessage> for WireProtocol {
88    fn from(zmq_message: ZmqMessage) -> Self {
89        let mut frames = zmq_message.into_vecdeque();
90        // The number of frames coming back from different types of Kernels has been frustratingly
91        // inconsistent. evcxr_jupyter (Rust) sometimes sends back 6 frames, skipping the identity
92        // frame. irkernel (R) sometimes sends two frames worth of identity.
93        let identity = match frames.len() {
94            8 => {
95                // only seen this extra frame in irkernel so far, first frame is bytes and
96                // second frame is the string 'kernel'
97                frames.pop_front().expect("Missing identity frame");
98                frames.pop_front().expect("Missing identity frame")
99            }
100            7 => frames.pop_front().expect("Missing identity frame"),
101            _ => Bytes::from("missing identity header"),
102        };
103        let delimiter = frames.pop_front().expect("Missing delimiter frame");
104        let hmac_signature = frames.pop_front().expect("Missing hmac_signature frame");
105        let header = frames.pop_front().expect("Missing header frame");
106        let parent_header = frames.pop_front().expect("Missing parent_header frame");
107        let metadata = frames.pop_front().expect("Missing metadata frame");
108        let content = frames.pop_front().expect("Missing content frame");
109
110        WireProtocol {
111            identity,
112            delimiter,
113            hmac_signature,
114            header,
115            parent_header,
116            metadata,
117            content,
118        }
119    }
120}
kernel_sidecar/jupyter/wire_protocol.rs

kernel_sidecar/jupyter/
wire_protocol.rs
