bws_rs/
utils.rs

1use std::io::Write;
2
3use sha1::Digest;
4
5pub struct BaseKv<K: PartialOrd, V> {
6    pub key: K,
7    pub val: V,
8}
9
10pub mod io {
11    use std::ops::{Deref, DerefMut};
12
13    pub trait PollRead {
14        fn poll_read<'a>(
15            &'a mut self,
16        ) -> std::pin::Pin<
17            Box<dyn 'a + Send + std::future::Future<Output = Result<Option<Vec<u8>>, String>>>,
18        >;
19    }
20    pub trait PollWrite {
21        fn poll_write<'a>(
22            &'a mut self,
23            buff: &'a [u8],
24        ) -> std::pin::Pin<
25            Box<dyn 'a + Send + std::future::Future<Output = Result<usize, std::io::Error>>>,
26        >;
27    }
28    pub struct BuffIo<const N: usize> {
29        buff: Vec<u8>,
30    }
31    impl<const N: usize> Deref for BuffIo<N> {
32        type Target = Vec<u8>;
33
34        fn deref(&self) -> &Self::Target {
35            &self.buff
36        }
37    }
38    impl<const N: usize> DerefMut for BuffIo<N> {
39        fn deref_mut(&mut self) -> &mut Self::Target {
40            &mut self.buff
41        }
42    }
43}
44pub enum ChunkParseError {
45    HashNoMatch,
46    IllegalContent,
47    Io(String),
48}
49#[derive(Clone, Copy)]
50enum ParseProcessState {
51    Head,
52    Content,
53    End,
54}
55struct ChunkHead {
56    content_size: usize,
57    ext: String,
58    signature: String,
59}
60///chunk parse, will auto verify sha256 value for every chunk
61pub async fn chunk_parse<R: io::PollRead + Send, W: tokio::io::AsyncWrite + Send + Unpin>(
62    mut src: R,
63    dst: &mut W,
64    circle_hasher: &mut crate::authorization::v4::HmacSha256CircleHasher,
65) -> Result<usize, ChunkParseError> {
66    let mut total_buff = Vec::<u8>::with_capacity(10 << 20);
67    let mut head = None;
68    let mut state = ParseProcessState::Head;
69    let mut total_size = 0;
70    while let Some(mut content) = src.poll_read().await.map_err(ChunkParseError::Io)? {
71        total_buff.append(&mut content);
72        state = parse_buff(
73            &mut total_buff,
74            dst,
75            state,
76            &mut head,
77            &mut total_size,
78            circle_hasher,
79        )
80        .await?;
81        if let ParseProcessState::End = state {
82            return Ok(total_size);
83        }
84    }
85    Ok(todo!())
86}
87async fn parse_buff<W: tokio::io::AsyncWrite + Send + Unpin>(
88    content: &mut Vec<u8>,
89    dst: &mut W,
90    mut state: ParseProcessState,
91    head: &mut Option<ChunkHead>,
92    total_size: &mut usize,
93    circle_hasher: &mut crate::authorization::v4::HmacSha256CircleHasher,
94) -> Result<ParseProcessState, ChunkParseError> {
95    use tokio::io::AsyncWriteExt;
96    while !content.is_empty() {
97        match state {
98            ParseProcessState::Head => {
99                if let Some(pos) = content.windows(2).position(|x| x == b"\r\n") {
100                    *head = Some(parse_chunk_line(&content[0..pos]).map_err(|_| {
101                        log::warn!("parse content line error\n{}", unsafe {
102                            std::str::from_utf8_unchecked(&content[0..pos])
103                        });
104                        ChunkParseError::IllegalContent
105                    })?);
106                    if let Some(hdr) = head {
107                        if hdr.content_size == 0 {
108                            let next=circle_hasher.next(
109                                "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
110                            ).unwrap();
111                            if hdr.signature.as_str() == next.as_str() {
112                                log::info!("all signature verify pass");
113                                state = ParseProcessState::End;
114                                return Ok(state);
115                            } else {
116                                log::info!("hash no match expect {next} got {}", hdr.signature);
117                                return Err(ChunkParseError::HashNoMatch);
118                            }
119                        }
120                    } else {
121                        log::info!("chunk not complete,wait");
122                    }
123                    content.drain(0..pos + 2);
124                    state = ParseProcessState::Content;
125                } else {
126                    return Ok(state);
127                }
128            }
129            ParseProcessState::Content => {
130                if let Some(hdr) = head {
131                    let content_len = content.len();
132                    if content_len >= hdr.content_size + 2 {
133                        if &content[hdr.content_size..hdr.content_size + 2] != b"\r\n" {
134                            log::warn!("content end is not chunk split symbol [{}]", unsafe {
135                                std::str::from_utf8_unchecked(
136                                    &content[hdr.content_size..hdr.content_size + 2],
137                                )
138                            });
139                            return Err(ChunkParseError::IllegalContent);
140                        }
141                        let mut hsh = sha2::Sha256::new();
142                        let _ = hsh.write_all(&content[0..hdr.content_size]);
143                        let hsh = hsh.finalize();
144                        let curr_hash = circle_hasher.next(hex::encode(hsh).as_str()).unwrap();
145                        if curr_hash != hdr.signature {
146                            log::warn!("chunk hash not match, return error");
147                            return Err(ChunkParseError::HashNoMatch);
148                        }
149                        // log::info!(
150                        //     "chunk signature verify pass {curr_hash} content length {}\n{}",
151                        //     content.len(),
152                        //     unsafe { std::str::from_utf8_unchecked(content) }
153                        // );
154                        dst.write_all(&content[0..hdr.content_size])
155                            .await
156                            .map_err(|err| ChunkParseError::Io(err.to_string()))?;
157                        content.drain(0..hdr.content_size + 2);
158                        // log::info!("{}", unsafe { std::str::from_utf8_unchecked(content) });
159                        *total_size += hdr.content_size;
160                        *head = None;
161                        state=ParseProcessState::Head;
162                    } else {
163                        return Ok(state);
164                    }
165                }
166            }
167            ParseProcessState::End => return Ok(state),
168        }
169    }
170    Ok(state)
171}
172fn parse_chunk_line(src: &[u8]) -> Result<ChunkHead, ()> {
173    let ret = src
174        .windows(1)
175        .position(|r| r == b";")
176        .and_then(|p1: usize| {
177            let raw = &src[..p1];
178            usize::from_str_radix(unsafe { std::str::from_utf8_unchecked(raw) }.trim(), 16)
179                .ok()
180                .and_then(|size| {
181                    let raw = &src[p1 + 1..];
182                    let raw = raw.splitn(2, |x| *x == b'=').collect::<Vec<&[u8]>>();
183                    if raw.len() != 2 {
184                        None
185                    } else {
186                        let ext = raw[0];
187                        let signature = raw[1];
188                        Some(ChunkHead {
189                            content_size: size,
190                            ext: unsafe { std::str::from_utf8_unchecked(ext) }.to_string(),
191                            signature: unsafe { std::str::from_utf8_unchecked(signature) }
192                                .to_string(),
193                        })
194                    }
195                })
196        });
197    ret.ok_or(())
198}
199fn parse_chunk_content<'a>(src: &'a [u8], ext: &str, signature: &str) -> Result<&'a [u8], ()> {
200    todo!()
201}
bws_rs/utils.rs

bws_rs/
utils.rs
