chore: sync current workspace changes

vendor/sqlx-core/src/io/buf.rs

@@ -0,0 +1,58 @@
+use std::str::from_utf8;
+
+use bytes::{Buf, Bytes};
+use memchr::memchr;
+
+use crate::error::Error;
+
+pub trait BufExt: Buf {
+    // Read a nul-terminated byte sequence
+    fn get_bytes_nul(&mut self) -> Result<Bytes, Error>;
+
+    // Read a byte sequence of the exact length
+    fn get_bytes(&mut self, len: usize) -> Bytes;
+
+    // Read a nul-terminated string
+    fn get_str_nul(&mut self) -> Result<String, Error>;
+
+    // Read a string of the exact length
+    fn get_str(&mut self, len: usize) -> Result<String, Error>;
+}
+
+impl BufExt for Bytes {
+    fn get_bytes_nul(&mut self) -> Result<Bytes, Error> {
+        let nul =
+            memchr(b'\0', self).ok_or_else(|| err_protocol!("expected NUL in byte sequence"))?;
+
+        let v = self.slice(0..nul);
+
+        self.advance(nul + 1);
+
+        Ok(v)
+    }
+
+    fn get_bytes(&mut self, len: usize) -> Bytes {
+        let v = self.slice(..len);
+        self.advance(len);
+
+        v
+    }
+
+    fn get_str_nul(&mut self) -> Result<String, Error> {
+        self.get_bytes_nul().and_then(|bytes| {
+            from_utf8(&bytes)
+                .map(ToOwned::to_owned)
+                .map_err(|err| err_protocol!("{}", err))
+        })
+    }
+
+    fn get_str(&mut self, len: usize) -> Result<String, Error> {
+        let v = from_utf8(&self[..len])
+            .map_err(|err| err_protocol!("{}", err))
+            .map(ToOwned::to_owned)?;
+
+        self.advance(len);
+
+        Ok(v)
+    }
+}

vendor/sqlx-core/src/io/buf_mut.rs

@@ -0,0 +1,12 @@
+use bytes::BufMut;
+
+pub trait BufMutExt: BufMut {
+    fn put_str_nul(&mut self, s: &str);
+}
+
+impl BufMutExt for Vec<u8> {
+    fn put_str_nul(&mut self, s: &str) {
+        self.extend(s.as_bytes());
+        self.push(0);
+    }
+}

vendor/sqlx-core/src/io/buf_stream.rs

@@ -0,0 +1,105 @@
+#![allow(dead_code)]
+
+use std::io;
+use std::ops::{Deref, DerefMut};
+
+use bytes::BytesMut;
+use sqlx_rt::{AsyncRead, AsyncReadExt, AsyncWrite};
+
+use crate::error::Error;
+use crate::io::write_and_flush::WriteAndFlush;
+use crate::io::{decode::Decode, encode::Encode};
+use std::io::Cursor;
+
+pub struct BufStream<S>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+{
+    pub(crate) stream: S,
+
+    // writes with `write` to the underlying stream are buffered
+    // this can be flushed with `flush`
+    pub(crate) wbuf: Vec<u8>,
+
+    // we read into the read buffer using 100% safe code
+    rbuf: BytesMut,
+}
+
+impl<S> BufStream<S>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+{
+    pub fn new(stream: S) -> Self {
+        Self {
+            stream,
+            wbuf: Vec::with_capacity(512),
+            rbuf: BytesMut::with_capacity(4096),
+        }
+    }
+
+    pub fn write<'en, T>(&mut self, value: T)
+    where
+        T: Encode<'en, ()>,
+    {
+        self.write_with(value, ())
+    }
+
+    pub fn write_with<'en, T, C>(&mut self, value: T, context: C)
+    where
+        T: Encode<'en, C>,
+    {
+        value.encode_with(&mut self.wbuf, context);
+    }
+
+    pub fn flush(&mut self) -> WriteAndFlush<'_, S> {
+        WriteAndFlush {
+            stream: &mut self.stream,
+            buf: Cursor::new(&mut self.wbuf),
+        }
+    }
+
+    pub async fn read<'de, T>(&mut self, cnt: usize) -> Result<T, Error>
+    where
+        T: Decode<'de, ()>,
+    {
+        self.read_with(cnt, ()).await
+    }
+
+    pub async fn read_with<'de, T, C>(&mut self, cnt: usize, context: C) -> Result<T, Error>
+    where
+        T: Decode<'de, C>,
+    {
+        T::decode_with(self.read_raw(cnt).await?.freeze(), context)
+    }
+
+    pub async fn read_raw(&mut self, cnt: usize) -> Result<BytesMut, Error> {
+        read_raw_into(&mut self.stream, &mut self.rbuf, cnt).await?;
+        let buf = self.rbuf.split_to(cnt);
+
+        Ok(buf)
+    }
+
+    pub async fn read_raw_into(&mut self, buf: &mut BytesMut, cnt: usize) -> Result<(), Error> {
+        read_raw_into(&mut self.stream, buf, cnt).await
+    }
+}
+
+impl<S> Deref for BufStream<S>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+{
+    type Target = S;
+
+    fn deref(&self) -> &Self::Target {
+        &self.stream
+    }
+}
+
+impl<S> DerefMut for BufStream<S>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+{
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.stream
+    }
+}

vendor/sqlx-core/src/io/decode.rs

@@ -0,0 +1,29 @@
+use bytes::Bytes;
+
+use crate::error::Error;
+
+pub trait ProtocolDecode<'de, Context = ()>
+where
+    Self: Sized,
+{
+    fn decode(buf: Bytes) -> Result<Self, Error>
+    where
+        Self: ProtocolDecode<'de, ()>,
+    {
+        Self::decode_with(buf, ())
+    }
+
+    fn decode_with(buf: Bytes, context: Context) -> Result<Self, Error>;
+}
+
+impl ProtocolDecode<'_> for Bytes {
+    fn decode_with(buf: Bytes, _: ()) -> Result<Self, Error> {
+        Ok(buf)
+    }
+}
+
+impl ProtocolDecode<'_> for () {
+    fn decode_with(_: Bytes, _: ()) -> Result<(), Error> {
+        Ok(())
+    }
+}

vendor/sqlx-core/src/io/encode.rs

@@ -0,0 +1,17 @@
+pub trait ProtocolEncode<'en, Context = ()> {
+    fn encode(&self, buf: &mut Vec<u8>) -> Result<(), crate::Error>
+    where
+        Self: ProtocolEncode<'en, ()>,
+    {
+        self.encode_with(buf, ())
+    }
+
+    fn encode_with(&self, buf: &mut Vec<u8>, context: Context) -> Result<(), crate::Error>;
+}
+
+impl<'en, C> ProtocolEncode<'en, C> for &'_ [u8] {
+    fn encode_with(&self, buf: &mut Vec<u8>, _context: C) -> Result<(), crate::Error> {
+        buf.extend_from_slice(self);
+        Ok(())
+    }
+}

vendor/sqlx-core/src/io/mod.rs

@@ -0,0 +1,26 @@
+mod buf;
+mod buf_mut;
+// mod buf_stream;
+mod decode;
+mod encode;
+mod read_buf;
+// mod write_and_flush;
+
+pub use buf::BufExt;
+pub use buf_mut::BufMutExt;
+//pub use buf_stream::BufStream;
+pub use decode::ProtocolDecode;
+pub use encode::ProtocolEncode;
+pub use read_buf::ReadBuf;
+
+#[cfg(not(feature = "_rt-tokio"))]
+pub use futures_io::AsyncRead;
+
+#[cfg(feature = "_rt-tokio")]
+pub use tokio::io::AsyncRead;
+
+#[cfg(not(feature = "_rt-tokio"))]
+pub use futures_util::io::AsyncReadExt;
+
+#[cfg(feature = "_rt-tokio")]
+pub use tokio::io::AsyncReadExt;

vendor/sqlx-core/src/io/read_buf.rs

@@ -0,0 +1,35 @@
+use bytes::{BufMut, BytesMut};
+
+/// An extension for [`BufMut`] for getting a writeable buffer in safe code.
+pub trait ReadBuf: BufMut {
+    /// Get the full capacity of this buffer as a safely initialized slice.
+    fn init_mut(&mut self) -> &mut [u8];
+}
+
+impl ReadBuf for &'_ mut [u8] {
+    #[inline(always)]
+    fn init_mut(&mut self) -> &mut [u8] {
+        self
+    }
+}
+
+impl ReadBuf for BytesMut {
+    #[inline(always)]
+    fn init_mut(&mut self) -> &mut [u8] {
+        // `self.remaining_mut()` returns `usize::MAX - self.len()`
+        let remaining = self.capacity() - self.len();
+
+        // I'm hoping for most uses that this operation is elided by the optimizer.
+        self.put_bytes(0, remaining);
+
+        self
+    }
+}
+
+#[test]
+fn test_read_buf_bytes_mut() {
+    let mut buf = BytesMut::with_capacity(8);
+    buf.put_u32(0x12345678);
+
+    assert_eq!(buf.init_mut(), [0x12, 0x34, 0x56, 0x78, 0, 0, 0, 0]);
+}

vendor/sqlx-core/src/io/write_and_flush.rs

@@ -0,0 +1,44 @@
+use crate::error::Error;
+use sqlx_rt::AsyncWrite;
+use std::future::Future;
+use std::io::{BufRead, Cursor};
+use std::pin::Pin;
+use std::task::{ready, Context, Poll};
+
+// Atomic operation that writes the full buffer to the stream, flushes the stream, and then
+// clears the buffer (even if either of the two previous operations failed).
+pub struct WriteAndFlush<'a, S> {
+    pub(super) stream: &'a mut S,
+    pub(super) buf: Cursor<&'a mut Vec<u8>>,
+}
+
+impl<S: AsyncWrite + Unpin> Future for WriteAndFlush<'_, S> {
+    type Output = Result<(), Error>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let Self {
+            ref mut stream,
+            ref mut buf,
+        } = *self;
+
+        loop {
+            let read = buf.fill_buf()?;
+
+            if !read.is_empty() {
+                let written = ready!(Pin::new(&mut *stream).poll_write(cx, read)?);
+                buf.consume(written);
+            } else {
+                break;
+            }
+        }
+
+        Pin::new(stream).poll_flush(cx).map_err(Error::Io)
+    }
+}
+
+impl<'a, S> Drop for WriteAndFlush<'a, S> {
+    fn drop(&mut self) {
+        // clear the buffer regardless of whether the flush succeeded or not
+        self.buf.get_mut().clear();
+    }
+}