Skip to main content

hatchery_node/
uplink_journal.rs

1//! Persistent uplink journal and JournalReq responder for the A2 restreamer.
2//!
3//! The node appends opaque payloads locally. When C2 asks for a gap range over
4//! the multiplex uplink, this module serves `JournalReply` from disk.
5
6use hatchery_core::FileJournal;
7use hatchery_core::{
8    ChannelKind, Frame, FrameError, FrameType, OpenPayload,
9};
10use std::path::PathBuf;
11use std::sync::Arc;
12use tokio::io::{AsyncReadExt, AsyncWriteExt};
13use tokio::net::TcpStream;
14
15/// Node-side journal bound to a restreamer uplink.
16pub struct UplinkJournal {
17    node_id: String,
18    journal: Arc<FileJournal>,
19}
20
21impl UplinkJournal {
22    pub fn open(node_id: impl Into<String>, path: impl Into<PathBuf>) -> Result<Self, String> {
23        let journal = FileJournal::open(path.into()).map_err(|e| e.to_string())?;
24        Ok(Self {
25            node_id: node_id.into(),
26            journal: Arc::new(journal),
27        })
28    }
29
30    pub fn append(&self, channel_id: u32, payload: &[u8]) -> Result<u64, String> {
31        self.journal
32            .append(channel_id, payload)
33            .map_err(|e| e.to_string())
34    }
35
36    pub fn journal(&self) -> &FileJournal {
37        &self.journal
38    }
39
40    /// Dial the restreamer node listen address, OPEN a channel, and answer JournalReq.
41    pub async fn serve_uplink(
42        self: Arc<Self>,
43        restreamer_addr: std::net::SocketAddr,
44        channel_id: u32,
45        kind: ChannelKind,
46    ) -> Result<(), String> {
47        let mut stream = TcpStream::connect(restreamer_addr)
48            .await
49            .map_err(|e| e.to_string())?;
50        let open = Frame::open(
51            channel_id,
52            &OpenPayload {
53                channel_kind: kind,
54                node_id: self.node_id.clone(),
55            },
56        )
57        .map_err(|e| e.to_string())?;
58        stream
59            .write_all(&open.encode().map_err(|e| e.to_string())?)
60            .await
61            .map_err(|e| e.to_string())?;
62
63        let mut buf = vec![0u8; 64 * 1024];
64        let mut filled = 0usize;
65        loop {
66            let n = stream
67                .read(&mut buf[filled..])
68                .await
69                .map_err(|e| e.to_string())?;
70            if n == 0 {
71                break;
72            }
73            filled += n;
74            while filled > 0 {
75                match Frame::decode(&buf[..filled]) {
76                    Ok((frame, consumed)) => {
77                        buf.copy_within(consumed..filled, 0);
78                        filled -= consumed;
79                        match frame.frame_type {
80                            FrameType::JournalReq => {
81                                let limit = frame.journal_req_limit().unwrap_or(64) as usize;
82                                let body = self
83                                    .journal
84                                    .reply_body(frame.channel_id, frame.seq, limit)
85                                    .map_err(|e| e.to_string())?;
86                                let reply = Frame::journal_reply(frame.channel_id, &body)
87                                    .map_err(|e| e.to_string())?;
88                                stream
89                                    .write_all(&reply.encode().map_err(|e| e.to_string())?)
90                                    .await
91                                    .map_err(|e| e.to_string())?;
92                            }
93                            FrameType::Window | FrameType::Ping | FrameType::Pong => {}
94                            FrameType::GoAway | FrameType::Close => return Ok(()),
95                            _ => {}
96                        }
97                    }
98                    Err(FrameError::Truncated) => break,
99                    Err(e) => return Err(e.to_string()),
100                }
101            }
102        }
103        Ok(())
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    #[test]
112    fn appends_persist() {
113        let dir = std::env::temp_dir().join(format!("node-uj-{}", std::process::id()));
114        let _ = std::fs::remove_dir_all(&dir);
115        let j = UplinkJournal::open("node-a", dir.join("j")).unwrap();
116        assert_eq!(j.append(1, b"x").unwrap(), 1);
117        assert_eq!(j.append(1, b"y").unwrap(), 2);
118        let body = j.journal().reply_body(1, 0, 10).unwrap();
119        assert_eq!(body.entries.len(), 2);
120        let _ = std::fs::remove_dir_all(&dir);
121    }
122}