kcptun_rust/
qppstream.rs

1//! QPP stream wrapper - encrypts on write, decrypts on read (compatible with Go QPPPort).
2
3use crate::qpp::{create_prng, QuantumPermutationPad};
4use std::sync::Arc;
5use std::sync::Mutex;
6use tokio::io::{AsyncRead, AsyncWrite};
7
8/// Wraps the read half of a stream and decrypts data with QPP.
9pub struct QppReadHalf<S> {
10    inner: S,
11    pad: Arc<QuantumPermutationPad>,
12    rprng: Arc<Mutex<crate::qpp::Rand>>,
13}
14
15/// Wraps the write half of a stream and encrypts data with QPP.
16pub struct QppWriteHalf<S> {
17    inner: S,
18    pad: Arc<QuantumPermutationPad>,
19    wprng: Arc<Mutex<crate::qpp::Rand>>,
20}
21
22impl<S: AsyncRead + Unpin> QppReadHalf<S> {
23    pub fn new(inner: S, pad: Arc<QuantumPermutationPad>, seed: &[u8]) -> Self {
24        QppReadHalf {
25            inner,
26            pad,
27            rprng: Arc::new(Mutex::new(create_prng(seed))),
28        }
29    }
30}
31
32impl<S: AsyncWrite + Unpin> QppWriteHalf<S> {
33    pub fn new(inner: S, pad: Arc<QuantumPermutationPad>, seed: &[u8]) -> Self {
34        QppWriteHalf {
35            inner,
36            pad,
37            wprng: Arc::new(Mutex::new(create_prng(seed))),
38        }
39    }
40}
41
42impl<S: AsyncRead + Unpin> AsyncRead for QppReadHalf<S> {
43    fn poll_read(
44        mut self: std::pin::Pin<&mut Self>,
45        cx: &mut std::task::Context<'_>,
46        buf: &mut tokio::io::ReadBuf<'_>,
47    ) -> std::task::Poll<std::io::Result<()>> {
48        let filled_before = buf.filled().len();
49        match std::pin::Pin::new(&mut self.inner).poll_read(cx, buf) {
50            std::task::Poll::Ready(Ok(())) => {
51                let filled_len = buf.filled().len();
52                let n = filled_len - filled_before;
53                if n > 0 {
54                    let mut rprng = self.rprng.lock().unwrap();
55                    let filled = buf.filled_mut();
56                    self.pad.decrypt_with_prng(&mut filled[filled_before..], &mut rprng);
57                }
58                std::task::Poll::Ready(Ok(()))
59            }
60            other => other,
61        }
62    }
63}
64
65impl<S: AsyncWrite + Unpin> AsyncWrite for QppWriteHalf<S> {
66    fn poll_write(
67        mut self: std::pin::Pin<&mut Self>,
68        cx: &mut std::task::Context<'_>,
69        buf: &[u8],
70    ) -> std::task::Poll<std::io::Result<usize>> {
71        let mut data = buf.to_vec();
72        {
73            let mut wprng = self.wprng.lock().unwrap();
74            self.pad.encrypt_with_prng(&mut data, &mut wprng);
75        }
76        std::pin::Pin::new(&mut self.inner).poll_write(cx, &data)
77    }
78
79    fn poll_flush(
80        mut self: std::pin::Pin<&mut Self>,
81        cx: &mut std::task::Context<'_>,
82    ) -> std::task::Poll<std::io::Result<()>> {
83        std::pin::Pin::new(&mut self.inner).poll_flush(cx)
84    }
85
86    fn poll_shutdown(
87        mut self: std::pin::Pin<&mut Self>,
88        cx: &mut std::task::Context<'_>,
89    ) -> std::task::Poll<std::io::Result<()>> {
90        std::pin::Pin::new(&mut self.inner).poll_shutdown(cx)
91    }
92}
93
94/// Build QPP read/write halves from already-split stream halves.
95/// Use the same seed (key bytes) on both client and server for interoperability.
96pub fn wrap_with_qpp<R, W>(
97    read_half: R,
98    write_half: W,
99    seed: &[u8],
100    qpp_count: u32,
101) -> (QppReadHalf<R>, QppWriteHalf<W>)
102where
103    R: AsyncRead + Unpin,
104    W: AsyncWrite + Unpin,
105{
106    let pad = Arc::new(QuantumPermutationPad::new(seed, qpp_count as u16));
107    (
108        QppReadHalf::new(read_half, pad.clone(), seed),
109        QppWriteHalf::new(write_half, pad, seed),
110    )
111}
kcptun_rust/qppstream.rs

kcptun_rust/
qppstream.rs
