Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[PLATFORM-860]: Veil memory optimisations and refactoring #60

Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
157 changes: 82 additions & 75 deletions src/private.rs
Original file line number Diff line number Diff line change
@@ -1,35 +1,24 @@
use std::fmt::{Debug, Display};
use std::fmt::{Debug, Display, Write};

#[repr(transparent)]
pub struct DisplayDebug(String);
impl std::fmt::Debug for DisplayDebug {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.0.as_str())
}
}
impl AsRef<str> for DisplayDebug {
#[inline(always)]
fn as_ref(&self) -> &str {
self.0.as_str()
}
}

pub struct RedactFlags {
pub enum RedactSpecialization {
/// Whether the type we're redacting is an Option<T> or not. Poor man's specialization! This is detected
/// by the proc macro reading the path to the type, so it's not perfect.
///
/// This could be improved & rid of in a number of different ways in the future:
///
/// * Once specialization is stabilized, we can use a trait to override redacting behaviour for some types,
/// * Once specialization is stabilized, we can use a trait to override redacting behavior for some types,
/// one of which would be Option<T>.
///
/// * Once std::ptr::metadata and friends are stabilized, we could use it to unsafely cast the dyn Debug pointer
/// to a concrete Option<T> and redact it directly. Probably not the best idea.
///
/// * Once trait upcasting is stabilized, we could use it to upcast the dyn Debug pointer to a dyn Any and then
/// downcast it to a concrete Option<T> and redact it directly.
pub is_option: bool,
Option,
}

#[derive(Clone, Copy)]
pub struct RedactFlags {
/// Whether to only partially redact the data.
///
/// Incompatible with `fixed`.
Expand All @@ -52,14 +41,14 @@ impl RedactFlags {
/// Maximum number of characters to expose at the beginning and end of a partial redact.
const MAX_PARTIAL_EXPOSE: usize = 3;

fn redact_partial(&self, str: &str, redacted: &mut String) {
pub(crate) fn redact_partial(&self, fmt: &mut std::fmt::Formatter, str: &str) -> std::fmt::Result {
WilliamVenner marked this conversation as resolved.
Show resolved Hide resolved
let count = str.chars().filter(|char| char.is_alphanumeric()).count();
if count < Self::MIN_PARTIAL_CHARS {
for char in str.chars() {
if char.is_alphanumeric() {
redacted.push(self.redact_char);
fmt.write_char(self.redact_char)?;
} else {
redacted.push(char);
fmt.write_char(char)?;
}
}
} else {
Expand All @@ -72,100 +61,118 @@ impl RedactFlags {
if char.is_alphanumeric() {
if prefix_gas > 0 {
prefix_gas -= 1;
redacted.push(char);
fmt.write_char(char)?;
MaeIsBad marked this conversation as resolved.
Show resolved Hide resolved
} else if middle_gas > 0 {
middle_gas -= 1;
redacted.push(self.redact_char);
fmt.write_char(self.redact_char)?;
} else {
redacted.push(char);
fmt.write_char(char)?;
}
} else {
redacted.push(char);
fmt.write_char(char)?;
}
}
}
Ok(())
}

fn redact_full(&self, str: &str, redacted: &mut String) {
pub(crate) fn redact_full(&self, fmt: &mut std::fmt::Formatter, str: &str) -> std::fmt::Result {
for char in str.chars() {
if char.is_whitespace() || !char.is_alphanumeric() {
redacted.push(char);
fmt.write_char(char)?;
} else {
redacted.push(self.redact_char);
fmt.write_char(self.redact_char)?;
}
}
Ok(())
}

fn redact_fixed(&self, width: usize, redacted: &mut String) {
redacted.reserve_exact(width);
pub(crate) fn redact_fixed(fmt: &mut std::fmt::Formatter, width: usize, char: char) -> std::fmt::Result {
let mut buf = String::with_capacity(width);
for _ in 0..width {
redacted.push(self.redact_char);
buf.push(char);
}
fmt.write_str(&buf)
}
}

pub enum RedactionTarget<'a> {
/// Redact the output of the type's [`std::fmt::Debug`] implementation.
/// Redact the output of the type's [`Debug`] implementation.
Debug {
this: &'a dyn Debug,

/// Sourced from [`std::fmt::Formatter::alternate`]
alternate: bool,
},

/// Redact the output of the type's [`std::fmt::Display`] implementation.
/// Redact the output of the type's [`Display`] implementation.
Display(&'a dyn Display),
}

pub fn redact(this: RedactionTarget, flags: RedactFlags) -> DisplayDebug {
let mut redacted = String::new();

let to_redactable_string = || match this {
RedactionTarget::Debug { this, alternate: false } => format!("{:?}", this),
RedactionTarget::Debug { this, alternate: true } => format!("{:#?}", this),
RedactionTarget::Display(this) => this.to_string(),
};

impl RedactionTarget<'_> {
/// Pass through directly to the formatter.
#[cfg(feature = "toggle")]
if crate::toggle::get_redaction_behavior().is_plaintext() {
return DisplayDebug(to_redactable_string());
pub(crate) fn passthrough(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
match self {
RedactionTarget::Debug { this, .. } => std::fmt::Debug::fmt(this, fmt),
RedactionTarget::Display(this) => std::fmt::Display::fmt(this, fmt),
}
}
}
impl ToString for RedactionTarget<'_> {
fn to_string(&self) -> String {
match self {
RedactionTarget::Debug { this, alternate: false } => format!("{:?}", this),
RedactionTarget::Debug { this, alternate: true } => format!("{:#?}", this),
RedactionTarget::Display(this) => this.to_string(),
}
}
}

(|| {
if flags.fixed > 0 {
flags.redact_fixed(flags.fixed as usize, &mut redacted);
return;
pub struct RedactionFormatter<'a> {
pub this: RedactionTarget<'a>,
pub flags: RedactFlags,
pub specialization: Option<RedactSpecialization>,
}
impl std::fmt::Debug for RedactionFormatter<'_> {
fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
#[cfg(feature = "toggle")]
if crate::toggle::get_redaction_behavior().is_plaintext() {
return self.this.passthrough(fmt);
}

let redactable_string = to_redactable_string();

redacted.reserve(redactable_string.len());

// Specialize for Option<T>
if flags.is_option {
if redactable_string == "None" {
// We don't need to do any redacting
// https://prima.slack.com/archives/C03URH9N43U/p1661423554871499
} else if let Some(inner) = redactable_string
.strip_prefix("Some(")
.and_then(|inner| inner.strip_suffix(')'))
{
redacted.push_str("Some(");
flags.redact_partial(inner, &mut redacted);
redacted.push(')');
} else {
// This should never happen, but just in case...
flags.redact_full(&redactable_string, &mut redacted);
if self.flags.fixed > 0 {
return RedactFlags::redact_fixed(fmt, self.flags.fixed as usize, self.flags.redact_char);
}

let redactable_string = self.this.to_string();

#[allow(clippy::single_match)]
match self.specialization {
Some(RedactSpecialization::Option) => {
if redactable_string == "None" {
// We don't need to do any redacting
// https://prima.slack.com/archives/C03URH9N43U/p1661423554871499
return fmt.write_str("None");
} else if let Some(inner) = redactable_string
.strip_prefix("Some(")
.and_then(|inner| inner.strip_suffix(')'))
{
fmt.write_str("Some(")?;
self.flags.redact_partial(fmt, inner)?;
return fmt.write_char(')');
} else {
// This should never happen, but just in case...
return self.flags.redact_full(fmt, &redactable_string);
}
}
return;

_ => {}
WilliamVenner marked this conversation as resolved.
Show resolved Hide resolved
}

if flags.partial {
flags.redact_partial(&redactable_string, &mut redacted);
if self.flags.partial {
self.flags.redact_partial(fmt, &redactable_string)
} else {
flags.redact_full(&redactable_string, &mut redacted);
self.flags.redact_full(fmt, &redactable_string)
}
})();

DisplayDebug(redacted)
}
}
32 changes: 24 additions & 8 deletions veil-macros/src/enums.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use crate::{
flags::FieldFlags,
flags::{ExtractFlags, FieldFlags, FieldFlagsParse},
fmt::{self, FormatData},
UnusedDiagnostic,
redact::UnusedDiagnostic,
};
use proc_macro::TokenStream;
use quote::ToTokens;
Expand All @@ -21,7 +21,7 @@ pub(super) fn derive_redact(
unused: &mut UnusedDiagnostic,
) -> Result<TokenStream, syn::Error> {
// Parse #[redact(all, variant, ...)] from the enum attributes, if present.
let top_level_flags = match FieldFlags::extract::<1>(&attrs, false)? {
let top_level_flags = match FieldFlags::extract::<1>("Redact", &attrs, FieldFlagsParse { skip_allowed: false })? {
[Some(flags)] => {
if !flags.all || !flags.variant {
return Err(syn::Error::new(
Expand All @@ -41,7 +41,13 @@ pub(super) fn derive_redact(
// Collect each variant's flags
let mut variant_flags = Vec::with_capacity(e.variants.len());
for variant in &e.variants {
let mut flags = match FieldFlags::extract::<2>(&variant.attrs, top_level_flags.is_some())? {
let mut flags = match FieldFlags::extract::<2>(
"Redact",
&variant.attrs,
FieldFlagsParse {
skip_allowed: top_level_flags.is_some(),
},
)? {
[None, None] => EnumVariantFieldFlags::default(),

[Some(flags), None] => {
Expand Down Expand Up @@ -163,7 +169,17 @@ pub(super) fn derive_redact(
// Variant name redacting
let variant_name = variant.ident.to_string();
let variant_name = if let Some(flags) = &flags.variant_flags {
fmt::generate_redact_call(quote! { &#variant_name }, false, flags, unused)
// The variant name must always be formatted with the Display impl.
let flags = FieldFlags {
display: true,
..*flags
};

// Generate the RedactionFormatter expression for the variant name
let redact = fmt::generate_redact_call(quote! { &#variant_name }, false, &flags, unused);

// Because the other side is expecting a &str, we need to convert the RedactionFormatter to a String (and then to a &str)
quote! { format!("{:?}", #redact).as_str() }
} else {
variant_name.into_token_stream()
};
Expand All @@ -182,7 +198,7 @@ pub(super) fn derive_redact(
"unit structs do not need redacting as they contain no data",
));
} else {
quote! { write!(f, "{:?}", #variant_name)? }
quote! { fmt.write_str(#variant_name)? }
}
}
});
Expand All @@ -191,9 +207,9 @@ pub(super) fn derive_redact(
let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
Ok(quote! {
impl #impl_generics ::std::fmt::Debug for #name_ident #ty_generics #where_clause {
fn fmt(&self, f: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result {
fn fmt(&self, fmt: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result {
#[allow(unused)] // Suppresses unused warning with `#[redact(display)]`
let alternate = f.alternate();
let alternate = fmt.alternate();

match self {
#(Self::#variant_idents #variant_destructures => { #variant_bodies; },)*
Expand Down
Loading