use arbitrary::{Arbitrary, Unstructured};
use chrono::{FixedOffset, TimeZone};
use crate::{
auth::AuthMechanism,
body::{
BasicFields, Body, BodyExtension, BodyStructure, MultiPartExtensionData,
SinglePartExtensionData, SpecificFields,
},
core::{
AString, Atom, AtomExt, IString, Literal, LiteralMode, NString, NonEmptyVec, Quoted,
QuotedChar, Tag, Text,
},
datetime::{DateTime, NaiveDate},
envelope::Envelope,
extensions::{enable::CapabilityEnable, quota::Resource},
flag::{Flag, FlagNameAttribute},
mailbox::{ListCharString, Mailbox, MailboxOther},
response::{
Capability, Code, CodeOther, CommandContinuationRequestBasic, Greeting, GreetingKind,
Status,
},
search::SearchKey,
sequence::SequenceSet,
};
macro_rules! implement_tryfrom {
($target:ty, $from:ty) => {
impl<'a> Arbitrary<'a> for $target {
fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
match <$target>::try_from(<$from>::arbitrary(u)?) {
Ok(passed) => Ok(passed),
Err(_) => Err(arbitrary::Error::IncorrectFormat),
}
}
}
};
}
macro_rules! implement_tryfrom_t {
($target:ty, $from:ty) => {
impl<'a, T> Arbitrary<'a> for $target
where
T: Arbitrary<'a>,
{
fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
match <$target>::try_from(<$from>::arbitrary(u)?) {
Ok(passed) => Ok(passed),
Err(_) => Err(arbitrary::Error::IncorrectFormat),
}
}
}
};
}
implement_tryfrom! { Atom<'a>, &str }
implement_tryfrom! { AtomExt<'a>, &str }
implement_tryfrom! { Quoted<'a>, &str }
implement_tryfrom! { Tag<'a>, &str }
implement_tryfrom! { Text<'a>, &str }
implement_tryfrom! { ListCharString<'a>, &str }
implement_tryfrom! { QuotedChar, char }
implement_tryfrom! { Mailbox<'a>, &str }
implement_tryfrom! { Capability<'a>, Atom<'a> }
implement_tryfrom! { Flag<'a>, &str }
implement_tryfrom! { FlagNameAttribute<'a>, Atom<'a> }
implement_tryfrom! { MailboxOther<'a>, AString<'a> }
implement_tryfrom! { CapabilityEnable<'a>, &str }
implement_tryfrom! { Resource<'a>, &str }
implement_tryfrom! { AuthMechanism<'a>, &str }
implement_tryfrom_t! { NonEmptyVec<T>, Vec<T> }
impl<'a> Arbitrary<'a> for CommandContinuationRequestBasic<'a> {
fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
Self::new(Option::<Code>::arbitrary(u)?, Text::arbitrary(u)?)
.map_err(|_| arbitrary::Error::IncorrectFormat)
}
}
impl<'a> Arbitrary<'a> for Greeting<'a> {
fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
Ok(Greeting {
kind: GreetingKind::arbitrary(u)?,
code: Option::<Code>::arbitrary(u)?,
text: {
let text = Text::arbitrary(u)?;
if text.as_ref().starts_with('[') {
Text::unvalidated("...")
} else {
text
}
},
})
}
}
impl<'a> Arbitrary<'a> for Status<'a> {
fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
let code = Option::<Code>::arbitrary(u)?;
let text = if code.is_some() {
Arbitrary::arbitrary(u)?
} else {
let text = Text::arbitrary(u)?;
if text.as_ref().starts_with('[') {
Text::unvalidated("...")
} else {
text
}
};
Ok(match u.int_in_range(0u8..=3)? {
0 => Status::Ok {
tag: Arbitrary::arbitrary(u)?,
code,
text,
},
1 => Status::No {
tag: Arbitrary::arbitrary(u)?,
code,
text,
},
2 => Status::Bad {
tag: Arbitrary::arbitrary(u)?,
code,
text,
},
3 => Status::Bye { code, text },
_ => unreachable!(),
})
}
}
impl<'a> Arbitrary<'a> for Literal<'a> {
fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
match Literal::try_from(<&[u8]>::arbitrary(u)?) {
Ok(mut passed) => {
passed.mode = LiteralMode::arbitrary(u)?;
Ok(passed)
}
Err(_) => Err(arbitrary::Error::IncorrectFormat),
}
}
}
impl<'a> Arbitrary<'a> for CodeOther<'a> {
fn arbitrary(_: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
Ok(CodeOther::unvalidated(b"IMAP-CODEC-CODE-OTHER>".as_ref()))
}
}
impl<'a> Arbitrary<'a> for SearchKey<'a> {
fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
fn make_search_key<'a>(u: &mut Unstructured<'a>) -> arbitrary::Result<SearchKey<'a>> {
Ok(match u.int_in_range(0u8..=33)? {
0 => SearchKey::SequenceSet(SequenceSet::arbitrary(u)?),
1 => SearchKey::All,
2 => SearchKey::Answered,
3 => SearchKey::Bcc(AString::arbitrary(u)?),
4 => SearchKey::Before(NaiveDate::arbitrary(u)?),
5 => SearchKey::Body(AString::arbitrary(u)?),
6 => SearchKey::Cc(AString::arbitrary(u)?),
7 => SearchKey::Deleted,
8 => SearchKey::Draft,
9 => SearchKey::Flagged,
10 => SearchKey::From(AString::arbitrary(u)?),
11 => SearchKey::Header(AString::arbitrary(u)?, AString::arbitrary(u)?),
12 => SearchKey::Keyword(Atom::arbitrary(u)?),
13 => SearchKey::Larger(u32::arbitrary(u)?),
14 => SearchKey::New,
15 => SearchKey::Old,
16 => SearchKey::On(NaiveDate::arbitrary(u)?),
17 => SearchKey::Recent,
18 => SearchKey::Seen,
19 => SearchKey::SentBefore(NaiveDate::arbitrary(u)?),
20 => SearchKey::SentOn(NaiveDate::arbitrary(u)?),
21 => SearchKey::SentSince(NaiveDate::arbitrary(u)?),
22 => SearchKey::Since(NaiveDate::arbitrary(u)?),
23 => SearchKey::Smaller(u32::arbitrary(u)?),
24 => SearchKey::Subject(AString::arbitrary(u)?),
25 => SearchKey::Text(AString::arbitrary(u)?),
26 => SearchKey::To(AString::arbitrary(u)?),
27 => SearchKey::Uid(SequenceSet::arbitrary(u)?),
28 => SearchKey::Unanswered,
29 => SearchKey::Undeleted,
30 => SearchKey::Undraft,
31 => SearchKey::Unflagged,
32 => SearchKey::Unkeyword(Atom::arbitrary(u)?),
33 => SearchKey::Unseen,
_ => unreachable!(),
})
}
fn make_search_key_rec<'a>(
u: &mut Unstructured<'a>,
depth: u8,
) -> arbitrary::Result<SearchKey<'a>> {
if depth == 0 {
return make_search_key(u);
}
Ok(match u.int_in_range(0u8..=36)? {
0 => SearchKey::And({
let keys = {
let len = u.arbitrary_len::<SearchKey>()?;
let mut tmp = Vec::with_capacity(len);
for _ in 0..len {
tmp.push(make_search_key_rec(u, depth - 1)?);
}
tmp
};
if !keys.is_empty() {
NonEmptyVec::try_from(keys).unwrap()
} else {
NonEmptyVec::from(make_search_key(u)?)
}
}),
1 => SearchKey::SequenceSet(SequenceSet::arbitrary(u)?),
2 => SearchKey::All,
3 => SearchKey::Answered,
4 => SearchKey::Bcc(AString::arbitrary(u)?),
5 => SearchKey::Before(NaiveDate::arbitrary(u)?),
6 => SearchKey::Body(AString::arbitrary(u)?),
7 => SearchKey::Cc(AString::arbitrary(u)?),
8 => SearchKey::Deleted,
9 => SearchKey::Draft,
10 => SearchKey::Flagged,
11 => SearchKey::From(AString::arbitrary(u)?),
12 => SearchKey::Header(AString::arbitrary(u)?, AString::arbitrary(u)?),
13 => SearchKey::Keyword(Atom::arbitrary(u)?),
14 => SearchKey::Larger(u32::arbitrary(u)?),
15 => SearchKey::New,
16 => SearchKey::Not(Box::new(make_search_key_rec(u, depth - 1)?)),
17 => SearchKey::Old,
18 => SearchKey::On(NaiveDate::arbitrary(u)?),
19 => SearchKey::Or(
Box::new(make_search_key_rec(u, depth - 1)?),
Box::new(make_search_key_rec(u, depth - 1)?),
),
20 => SearchKey::Recent,
21 => SearchKey::Seen,
22 => SearchKey::SentBefore(NaiveDate::arbitrary(u)?),
23 => SearchKey::SentOn(NaiveDate::arbitrary(u)?),
24 => SearchKey::SentSince(NaiveDate::arbitrary(u)?),
25 => SearchKey::Since(NaiveDate::arbitrary(u)?),
26 => SearchKey::Smaller(u32::arbitrary(u)?),
27 => SearchKey::Subject(AString::arbitrary(u)?),
28 => SearchKey::Text(AString::arbitrary(u)?),
29 => SearchKey::To(AString::arbitrary(u)?),
30 => SearchKey::Uid(SequenceSet::arbitrary(u)?),
31 => SearchKey::Unanswered,
32 => SearchKey::Undeleted,
33 => SearchKey::Undraft,
34 => SearchKey::Unflagged,
35 => SearchKey::Unkeyword(Atom::arbitrary(u)?),
36 => SearchKey::Unseen,
_ => unreachable!(),
})
}
make_search_key_rec(u, 7)
}
}
impl<'a> Arbitrary<'a> for BodyStructure<'a> {
fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
fn make_body_structure_terminator<'a>(
u: &mut Unstructured<'a>,
) -> arbitrary::Result<BodyStructure<'a>> {
Ok(BodyStructure::Single {
body: Body {
basic: BasicFields::arbitrary(u)?,
specific: match u.int_in_range(1..=2)? {
1 => SpecificFields::Basic {
r#type: IString::arbitrary(u)?,
subtype: IString::arbitrary(u)?,
},
2 => SpecificFields::Text {
subtype: IString::arbitrary(u)?,
number_of_lines: u32::arbitrary(u)?,
},
_ => unreachable!(),
},
},
extension_data: Option::<SinglePartExtensionData>::arbitrary(u)?,
})
}
fn make_body_structure_rec<'a>(
u: &mut Unstructured<'a>,
depth: u8,
) -> arbitrary::Result<BodyStructure<'a>> {
if depth == 0 {
return make_body_structure_terminator(u);
}
Ok(match u.int_in_range(1..=2)? {
1 => BodyStructure::Single {
body: Body {
basic: BasicFields::arbitrary(u)?,
specific: match u.int_in_range(1..=3)? {
1 => SpecificFields::Basic {
r#type: IString::arbitrary(u)?,
subtype: IString::arbitrary(u)?,
},
2 => SpecificFields::Message {
envelope: Box::<Envelope>::arbitrary(u)?,
body_structure: Box::new(make_body_structure_rec(u, depth - 1)?),
number_of_lines: u32::arbitrary(u)?,
},
3 => SpecificFields::Text {
subtype: IString::arbitrary(u)?,
number_of_lines: u32::arbitrary(u)?,
},
_ => unreachable!(),
},
},
extension_data: Option::<SinglePartExtensionData>::arbitrary(u)?,
},
2 => BodyStructure::Multi {
bodies: {
let bodies = {
let len = u.arbitrary_len::<BodyStructure>()?;
let mut tmp = Vec::with_capacity(len);
for _ in 0..len {
tmp.push(make_body_structure_rec(u, depth - 1)?);
}
tmp
};
if !bodies.is_empty() {
NonEmptyVec::try_from(bodies).unwrap()
} else {
NonEmptyVec::from(make_body_structure_terminator(u)?)
}
},
subtype: IString::arbitrary(u)?,
extension_data: Option::<MultiPartExtensionData>::arbitrary(u)?,
},
_ => unreachable!(),
})
}
make_body_structure_rec(u, 3)
}
}
impl<'a> Arbitrary<'a> for BodyExtension<'a> {
fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
fn make_body_extension_terminator<'a>(
u: &mut Unstructured<'a>,
) -> arbitrary::Result<BodyExtension<'a>> {
Ok(match u.int_in_range(1..=2)? {
1 => BodyExtension::NString(NString::arbitrary(u)?),
2 => BodyExtension::Number(u32::arbitrary(u)?),
_ => unreachable!(),
})
}
fn make_body_extension_rec<'a>(
u: &mut Unstructured<'a>,
depth: u8,
) -> arbitrary::Result<BodyExtension<'a>> {
if depth == 0 {
return make_body_extension_terminator(u);
}
Ok(match u.int_in_range(1..=2)? {
1 => BodyExtension::NString(NString::arbitrary(u)?),
2 => BodyExtension::Number(u32::arbitrary(u)?),
3 => BodyExtension::List({
let body_extensions = {
let len = u.arbitrary_len::<BodyExtension>()?;
let mut tmp = Vec::with_capacity(len);
for _ in 0..len {
tmp.push(make_body_extension_rec(u, depth - 1)?);
}
tmp
};
if !body_extensions.is_empty() {
NonEmptyVec::try_from(body_extensions).unwrap()
} else {
NonEmptyVec::from(make_body_extension_terminator(u)?)
}
}),
_ => unreachable!(),
})
}
make_body_extension_rec(u, 3)
}
}
impl<'a> Arbitrary<'a> for DateTime {
fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
let local_datetime = chrono::NaiveDateTime::new(
chrono::NaiveDate::from_ymd_opt(
u.int_in_range(0..=9999)?,
u.int_in_range(1..=12)?,
u.int_in_range(1..=31)?,
)
.ok_or(arbitrary::Error::IncorrectFormat)?,
chrono::NaiveTime::arbitrary(u)?,
);
let hours = u.int_in_range(0..=23 * 3600)?;
let minutes = u.int_in_range(0..=59)? * 60;
DateTime::try_from(
FixedOffset::east_opt(hours + minutes)
.unwrap()
.from_local_datetime(&local_datetime)
.unwrap(),
)
.map_err(|_| arbitrary::Error::IncorrectFormat)
}
}
impl<'a> Arbitrary<'a> for NaiveDate {
fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
NaiveDate::try_from(chrono::NaiveDate::arbitrary(u)?)
.map_err(|_| arbitrary::Error::IncorrectFormat)
}
}
#[cfg(test)]
mod tests {
use arbitrary::{Arbitrary, Error, Unstructured};
#[cfg(feature = "bounded-static")]
use bounded_static::{IntoBoundedStatic, ToBoundedStatic};
use rand::{rngs::SmallRng, Rng, SeedableRng};
use crate::{
command::Command,
response::{Greeting, Response},
};
macro_rules! impl_test_arbitrary {
($object:ty) => {
let mut rng = SmallRng::seed_from_u64(1337);
let mut data = [0u8; 256];
rng.try_fill(&mut data).unwrap();
let mut unstructured = Unstructured::new(&data);
let mut count = 0;
loop {
match <$object>::arbitrary(&mut unstructured) {
Ok(_out) => {
count += 1;
#[cfg(feature = "bounded-static")]
{
let out_to_static = _out.to_static();
assert_eq!(_out, out_to_static);
let out_into_static = _out.into_static();
assert_eq!(out_to_static, out_into_static);
}
if count >= 1_000 {
break;
}
}
Err(Error::NotEnoughData | Error::IncorrectFormat) => {
rng.try_fill(&mut data).unwrap();
unstructured = Unstructured::new(&data);
}
Err(Error::EmptyChoose) => {
unreachable!();
}
Err(_) => {
unimplemented!()
}
}
}
};
}
#[test]
fn test_arbitrary_greeting() {
impl_test_arbitrary! {Greeting};
}
#[test]
fn test_arbitrary_command() {
impl_test_arbitrary! {Command};
}
#[test]
fn test_arbitrary_response() {
impl_test_arbitrary! {Response};
}
}