/*
Portions Copyright 2019-2021 ZomboDB, LLC.
Portions Copyright 2021-2022 Technology Concepts & Design, Inc. <support@tcdi.com>
All rights reserved.
Use of this source code is governed by the MIT license that can be found in the LICENSE file.
*/
//! A safe wrapper around Postgres `StringInfo` structure
#![allow(dead_code, non_snake_case)]
use crate::{pg_sys, AllocatedByPostgres, AllocatedByRust, PgBox, WhoAllocated};
use core::fmt::{Display, Formatter};
use core::str::Utf8Error;
use std::io::Error;
/// StringInfoData holds information about an extensible string that is allocated by Postgres'
/// memory system, but generally follows Rust's drop semantics
pub struct StringInfo<AllocatedBy: WhoAllocated = AllocatedByRust> {
inner: PgBox<pg_sys::StringInfoData, AllocatedBy>,
}
impl<AllocatedBy: WhoAllocated> From<StringInfo<AllocatedBy>> for &'static std::ffi::CStr {
fn from(val: StringInfo<AllocatedBy>) -> Self {
let len = val.len();
let ptr = val.into_char_ptr();
unsafe {
std::ffi::CStr::from_bytes_with_nul_unchecked(std::slice::from_raw_parts(
ptr as *const u8,
(len + 1) as usize, // +1 to get the trailing null byte
))
}
}
}
impl<AllocatedBy: WhoAllocated> From<StringInfo<AllocatedBy>> for &'static crate::cstr_core::CStr {
fn from(val: StringInfo<AllocatedBy>) -> Self {
let len = val.len();
let ptr = val.into_char_ptr();
unsafe {
crate::cstr_core::CStr::from_bytes_with_nul_unchecked(std::slice::from_raw_parts(
ptr as *const u8,
(len + 1) as usize, // +1 to get the trailing null byte
))
}
}
}
impl<AllocatedBy: WhoAllocated> std::io::Write for StringInfo<AllocatedBy> {
fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
self.push_bytes(buf);
Ok(buf.len())
}
fn flush(&mut self) -> Result<(), Error> {
Ok(())
}
}
impl<AllocatedBy: WhoAllocated> Display for StringInfo<AllocatedBy> {
/// Convert this [`StringInfo`] into a Rust string. This uses [`String::from_utf8_lossy`] as
/// it's fine for a Postgres [`StringInfo`] to contain null bytes and also not even be proper
/// UTF8.
fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
f.write_str(&String::from_utf8_lossy(self.as_bytes()))
}
}
impl StringInfo<AllocatedByRust> {
/// Construct a new `StringInfo` of its default size, allocated in `CurrentMemoryContext`
///
/// Note that Postgres can only represent up to 1 gigabyte of data in a `StringInfo`
pub fn new() -> Self {
StringInfo {
inner: unsafe {
// SAFETY: makeStringInfo() always returns a valid StringInfoData pointer. It'll
// ereport if it can't
PgBox::<_, AllocatedByRust>::from_rust(pg_sys::makeStringInfo())
},
}
}
/// Construct a new `StringInfo`, allocated by Postgres in `CurrentMemoryContext`, ensuring it
/// has a capacity of the specified `len`.
///
/// Note that Postgres can only represent up to 1 gigabyte of data in a `StringInfo`
pub fn with_capacity(len: i32) -> Self {
let mut si = StringInfo::default();
si.enlarge(len);
si
}
}
impl StringInfo<AllocatedByPostgres> {
/// Construct a `StringInfo` from a Postgres-allocated `pg_sys::StringInfo`.
///
/// The backing `pg_sys::StringInfo` structure will be freed whenever the memory context in which
/// it was originally allocated is reset.
///
/// # Safety
///
/// This function is unsafe as it cannot confirm the provided [`pg_sys::StringInfo`] pointer is
/// valid
pub unsafe fn from_pg(sid: pg_sys::StringInfo) -> Option<Self> {
if sid.is_null() {
None
} else {
Some(StringInfo { inner: PgBox::from_pg(sid) })
}
}
}
impl<AllocatedBy: WhoAllocated> StringInfo<AllocatedBy> {
/// What is the length, excluding the trailing null byte
#[inline]
pub fn len(&self) -> usize {
self.inner.len as _
}
/// Do we have any characters?
#[inline]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
/// Push a Rust character onto the end. A Rust `char` could be 4 bytes in total, so it
/// is converted into a String first to ensure unicode correctness
#[inline]
pub fn push(&mut self, ch: char) {
self.push_str(&ch.to_string());
}
/// Push a String reference onto the end
#[inline]
pub fn push_str(&mut self, s: &str) {
self.push_bytes(s.as_bytes())
}
/// Push arbitrary bytes onto the end. Any byte sequence is allowed, including those with
/// embedded NULLs
///
/// # Panics
///
/// This function will panic if the length of bytes is larger than an `i32`
#[inline]
pub fn push_bytes(&mut self, bytes: &[u8]) {
unsafe {
// SAFETY: self.inner will always be valid
pg_sys::appendBinaryStringInfo(
self.inner.as_ptr(),
bytes.as_ptr() as _,
bytes.len().try_into().expect("len of bytes doesn't fit in an i32"),
)
}
}
/// Push the bytes behind a raw pointer of a given length onto the end
///
/// # Safety
///
/// This function is unsafe as we cannot ensure the specified `ptr` and `len` arguments
/// are what you say that are and otherwise in agreement
#[inline]
pub unsafe fn push_raw(&mut self, ptr: *const std::os::raw::c_char, len: i32) {
unsafe {
// SAFETY: self.inner will always be a valid StringInfoData pointer
// and the caller gets to decide if `ptr` and `len` line up
pg_sys::appendBinaryStringInfo(self.inner.as_ptr(), ptr, len)
}
}
/// Reset the size of the `StringInfo` back to zero-length. This does/// *not** free any
/// previously-allocated memory
#[inline]
pub fn reset(&mut self) {
unsafe {
// SAFETY: self.inner will always be a valid StringInfoData pointer
pg_sys::resetStringInfo(self.inner.as_ptr())
}
}
/// Ensure that this `StringInfo` is at least `needed` bytes long
#[inline]
pub fn enlarge(&mut self, needed: i32) {
unsafe {
// SAFETY: self.inner will always be a valid StringInfoData pointer
pg_sys::enlargeStringInfo(self.inner.as_ptr(), needed)
}
}
/// A `&str` representation.
///
/// # Errors
///
/// If the contained bytes aren't valid UTF8, a [UTF8Error] is returned. Postgres StringInfo
/// is allowed to contain non-UTF8 byte sequences, so this is a real possibility.
#[inline]
pub fn as_str(&self) -> Result<&str, Utf8Error> {
std::str::from_utf8(self.as_bytes())
}
/// A pointer to the backing bytes
#[inline]
pub fn as_ptr(&self) -> *const std::os::raw::c_char {
self.inner.data
}
/// A mutable pointer to the backing bytes
#[inline]
pub fn as_mut_ptr(&self) -> *mut std::os::raw::c_char {
self.inner.data
}
/// A `&[u8]` byte slice representation
#[inline]
pub fn as_bytes(&self) -> &[u8] {
unsafe {
// SAFETY: self.inner will always be a valid StringInfoData pointer, and Postgres will
// never have self.inner.data be invalid
std::slice::from_raw_parts(self.inner.data as *const u8, self.len())
}
}
/// A mutable `&[u8]` byte slice representation
#[inline]
pub fn as_bytes_mut(&mut self) -> &mut [u8] {
unsafe {
// SAFETY: self.inner will always be a valid StringInfoData
// pointer, and Postgres will never have self.inner.data be invalid
std::slice::from_raw_parts_mut(self.inner.data as *mut u8, self.len())
}
}
/// Convert this `StringInfo` into one that is wholly owned and now managed by Postgres
#[inline]
pub fn into_pg(mut self) -> *mut pg_sys::StringInfoData {
// NB: We are replacing self.inner with a PgBox containing the null pointer. However,
// `self` will be dropped as soon as this function ends and we account for this case in our
// drop implementation
let inner = std::mem::replace(&mut self.inner, PgBox::<_, AllocatedBy>::null());
inner.into_pg()
}
/// Convert this `StringInfo` into a `"char *"` that is wholly owned and now managed by Postgres
#[inline]
pub fn into_char_ptr(self) -> *const std::os::raw::c_char {
// in case we're AllocatedByRust, we don't want drop to free `self.inner.data` now that we've
// consumed `self` and are returning a raw pointer to some memory we allocated, so we "round-trip"
// to ensure Rust thinks we're now `AllocatedByPostgres`, which has an empty drop impl
let sid_ptr = self.into_pg();
unsafe {
// SAFETY: we just made the StringInfoData pointer so we know it's valid and properly
// initialized throughout
sid_ptr.as_ref().unwrap_unchecked().data
}
}
}
impl Default for StringInfo<AllocatedByRust> {
fn default() -> Self {
Self::new()
}
}
impl From<String> for StringInfo<AllocatedByRust> {
fn from(s: String) -> Self {
StringInfo::from(s.as_str())
}
}
impl From<&str> for StringInfo<AllocatedByRust> {
fn from(s: &str) -> Self {
let mut rc = StringInfo::new();
rc.push_str(s);
rc
}
}
impl From<Vec<u8>> for StringInfo<AllocatedByRust> {
fn from(v: Vec<u8>) -> Self {
let mut rc = StringInfo::new();
rc.push_bytes(v.as_slice());
rc
}
}
impl From<&[u8]> for StringInfo<AllocatedByRust> {
fn from(v: &[u8]) -> Self {
let mut rc = StringInfo::new();
rc.push_bytes(v);
rc
}
}
impl<AllocatedBy: WhoAllocated> Drop for StringInfo<AllocatedBy> {
fn drop(&mut self) {
unsafe {
// SAFETY: self.inner could represent the null pointer, but if it doesn't, then it's
// one we can use as self.inner.data will always be allocated if self.inner is.
//
// It's also prescribed by Postgres that to fully deallocate a StringInfo pointer, the
// owner is responsible for freeing its .data member... and that's us
if !self.inner.is_null() {
AllocatedBy::maybe_pfree(self.inner.data.cast())
}
}
}
}