miasma 0.4.1

Trap malicious AI web scrapers in an endless pit of poisoned training data.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

use std::pin::pin;
use std::sync::{Arc, Mutex};

use async_stream::{stream, try_stream};
use bytes::Bytes;
use futures::{Stream, StreamExt};
use tokio::sync::OwnedSemaphorePermit;
use uuid::Uuid;

use super::{LinkSettings, LinkSettingsInner};
use crate::metrics::{Metrics, UserAgent};
use crate::poison::PoisonClient;
use crate::utils::{cow_helpers, stream_size};
use crate::{MiasmaStream, QueryParams, templating::TemplateBuilder};

pub struct PoisonResponseStreamArgs {
    pub permit: OwnedSemaphorePermit,
    pub poison_client: Arc<PoisonClient>,
    pub link_settings: LinkSettings,
    pub metrics: Option<(UserAgent, Arc<Mutex<Metrics>>)>,
}

/// Build the poison response.
pub fn build_response_stream(
    PoisonResponseStreamArgs {
        permit,
        poison_client,
        link_settings,
        metrics,
    }: PoisonResponseStreamArgs,
) -> impl MiasmaStream {
    let metrics_for_total = metrics.clone();

    let template = TemplateBuilder::with_random_template();

    let stream = try_stream! {
        // carry the semaphore permit through the life of this stream.
        let _permit = permit;

        for chunk in template.start_to_body() {
            yield cow_helpers::as_bytes(chunk);
        }

        for body_section in template.body_sections() {
            yield Bytes::from_static(body_section.pre_poison().as_bytes());

            let mut poison = pin!(stream_size::with_bytes_counted(
                poison_client.stream_poison().await,
                async |n| {
                    if let Some((user_agent, metrics)) = &metrics {
                        metrics.lock().expect("metrics mutex poisoned").record_poison_bytes(user_agent, n);
                    }
                },
            ));
            while let Some(chunk) = poison.next().await {
                yield chunk?;
            }
            for chunk in body_section.post_poison() {
                yield cow_helpers::as_bytes(chunk);
            }
        }

        yield Bytes::from_static(template.body_to_links().as_bytes());

        match link_settings {
            LinkSettings::NoLinks => yield Bytes::from_static(b"None"),
            LinkSettings::Links(l) => {
                let mut links = pin!(build_links_stream(&template, &l));
                while let Some(chunk) = links.next().await {
                    yield chunk;
                }
            },
        }

        for chunk in template.links_to_end() {
            yield cow_helpers::as_bytes(chunk);
        }
    };

    stream_size::with_bytes_counted(stream, async move |n| {
        if let Some((user_agent, metrics)) = &metrics_for_total {
            metrics
                .lock()
                .expect("metrics mutex poisoned")
                .record_total_bytes(user_agent, n);
        }
    })
}

fn build_links_stream(
    template: &TemplateBuilder,
    link_settings: &LinkSettingsInner,
) -> impl Stream<Item = Bytes> {
    let params = match link_settings.next_depth {
        None => String::new(),
        Some(c) => format!("?{}={}", QueryParams::CURRENT_DEPTH_QUERY_PARAM, c),
    };

    stream! {
        for _ in 0..link_settings.count {
            let link = format!(
                "<li><a href=\"{prefix}{id}{params}\">{link_title}</a></li>",
                prefix = link_settings.prefix,
                id = Uuid::new_v4(),
                link_title = template.rand_link_title()
            );
            yield Bytes::from(link);
        }
    }
}