AEAD-2022 TCP writer reports retry buffer length after Pending, dropping appended data

Cargo.toml:

[package]
name = "ss-pending-write-repro"
version = "0.1.0"
edition = "2024"

[dependencies]
shadowsocks = { version = "=1.25.0", default-features = false, features = ["aead-cipher-2022"] }
tokio = { version = "=1.53.1", features = ["io-util"] }

src/main.rs:

use shadowsocks::{
    config::ServerType,
    context::Context as SsContext,
    crypto::CipherKind,
    relay::tcprelay::crypto_io::{CryptoStream, CryptoWrite, StreamType},
};
use std::{
    io,
    pin::Pin,
    task::{Context, Poll, Waker},
};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};

#[derive(Default)]
struct Sink {
    blocked: bool,
    bytes: Vec<u8>,
    waker: Option<Waker>,
}

impl AsyncRead for Sink {
    fn poll_read(
        self: Pin<&mut Self>,
        _: &mut Context<'_>,
        _: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        Poll::Ready(Ok(()))
    }
}

impl AsyncWrite for Sink {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        data: &[u8],
    ) -> Poll<io::Result<usize>> {
        if self.blocked {
            self.waker = Some(cx.waker().clone());
            return Poll::Pending;
        }
        self.bytes.extend_from_slice(data);
        Poll::Ready(Ok(data.len()))
    }

    fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
        Poll::Ready(Ok(()))
    }

    fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
        Poll::Ready(Ok(()))
    }
}

fn main() {
    let context = SsContext::new_shared(ServerType::Local);
    let mut stream = CryptoStream::from_stream(
        &context,
        Sink {
            blocked: true,
            ..Sink::default()
        },
        StreamType::Client,
        CipherKind::AEAD2022_BLAKE3_AES_128_GCM,
        &[0; 16],
    );
    let mut cx = Context::from_waker(Waker::noop());

    assert!(
        Pin::new(&mut stream)
            .poll_write_encrypted(&mut cx, b"a")
            .is_pending()
    );

    let sink = stream.get_mut();
    sink.blocked = false;
    sink.waker.take().unwrap().wake();

    let Poll::Ready(Ok(written)) = Pin::new(&mut stream).poll_write_encrypted(&mut cx, b"ab")
    else {
        panic!("write did not complete");
    };
    println!("Reported plaintext bytes: {written}");
    println!("Ciphertext bytes: {}", stream.get_ref().bytes.len());
    // The cached frame contains only 'a'. The appended 'b' has not been sent.
    assert_eq!(written, 1);
}