shadowsocks: write salt along with the first payload chunk

This commit is contained in:
eric
2023-09-05 20:51:00 +08:00
parent af24da946d
commit c14eb4b911

View File

@@ -27,8 +27,7 @@ enum ReadState {
enum WriteState {
WaitingSalt,
PendingSalt(usize, usize),
WaitingChunk,
WaitingChunk(usize),
PendingChunk(usize, (usize, usize)),
}
@@ -251,51 +250,35 @@ where
self.enc.replace(enc);
// ready to write salt
self.write_state = WriteState::PendingSalt(salt_size, 0);
self.write_state = WriteState::WaitingChunk(salt_size);
}
WriteState::PendingSalt(total, written) => {
WriteState::WaitingChunk(salt_size) => {
let me = &mut *self;
// write salt
// TODO write salt together with payload
let nw = ready!(poll_write_buf(
Pin::new(&mut me.inner),
cx,
&mut me.write_buf
))?;
if nw == 0 {
return Err(early_eof()).into();
}
let mut length_and_payload = me.write_buf.split_off(salt_size);
if written + nw >= total {
self.write_state = WriteState::WaitingChunk;
} else {
self.write_state = WriteState::PendingSalt(total, written + nw);
}
}
WriteState::WaitingChunk => {
let me = &mut *self;
// 0x3fff is the mandatory maximum size in ss spec
let consume_len = min(buf.len(), 0x3fff);
let enc = me.enc.as_mut().expect("uninitialized cipher");
// seal payload length
let piece1_size = 2 + me.cipher.tag_len();
me.write_buf.reserve(piece1_size);
unsafe { me.write_buf.set_len(2) };
BigEndian::write_u16(&mut me.write_buf[..2], consume_len as u16);
enc.encrypt(&mut me.write_buf).map_err(|_| crypto_err())?;
let mut piece2 = me.write_buf.split_off(piece1_size);
let length_size = 2 + me.cipher.tag_len();
length_and_payload.reserve(length_size);
unsafe { length_and_payload.set_len(2) };
BigEndian::write_u16(&mut length_and_payload[..2], consume_len as u16);
enc.encrypt(&mut length_and_payload)
.map_err(|_| crypto_err())?;
let mut payload = length_and_payload.split_off(length_size);
// seal payload
let piece2_size = consume_len + me.cipher.tag_len();
piece2.reserve(piece2_size);
piece2.put_slice(&buf[..consume_len]);
enc.encrypt(&mut piece2).map_err(|_| crypto_err())?;
let payload_size = consume_len + me.cipher.tag_len();
payload.reserve(payload_size);
payload.put_slice(&buf[..consume_len]);
enc.encrypt(&mut payload).map_err(|_| crypto_err())?;
// merge length and payload pieces
me.write_buf.unsplit(piece2);
// merge salt, length and payload
length_and_payload.unsplit(payload);
me.write_buf.unsplit(length_and_payload);
// ready to write data
self.write_state =
@@ -321,7 +304,7 @@ where
if written + nw >= total {
// data chunk written, go to next chunk
me.write_state = WriteState::WaitingChunk;
me.write_state = WriteState::WaitingChunk(0);
return Poll::Ready(Ok(consumed));
}