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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
//!  SPDX-License-Identifier: MIT
//!
//! Copyright (c) 2023, eunomia-bpf
//! All rights reserved.
//!

//! # Event exporter for bpf-loader-rs
//!
//! ## What does this module do?
//! In a nut shell, this module provides a struct `EventExporter`, which accepts exported data from a ebpf program, verifies the data, then output them to stdout or the user-provided callback function.
//!
//! ## Constructors
//! Usually, ebpf program will export data sampled from the kernel in one of the two ways:
//! - ringbuf
//! - bpf map
//! There are also two constructors for EventExporter, which corresponds to the two ways, `EventExporterBuilder::build_for_ringbuf` and `EventExporterBuilder::build_for_map_sampling`
//!
//! ## What will be produced
//!
//! You can get the data you want in one of the three formats:
//! - Json
//! - PlainText
//! - RawEvent
//!
//! Each time the EventExporter received data from the ebpf program, it will convert the data to the format you want, and call the callback function to acknowledge you the data, or print that to stdout
//! ## Json
//! Convert the received data to a JsonValue. With the help of BTF, we can know at which offset exists which type of data. So It's ok to construct a JsonValue representing the type we want to see. The JsonValue will be passed to the callback in string format, or printed to stdout if you don't want to provide a callback
//!
//! ## PlainText
//! It's similar to JSON, except that it's not structured, only human readable texts
//!
//! ## RawEvent
//! It will call the callback with the original data received from ebpf program. If no callback was provided, it will do nothing.

use crate::{
    btf_container::BtfContainer,
    export_event::checker::check_sample_types_btf,
    meta::{BufferValueInterpreter, ExportedTypesStructMeta, MapSampleMeta, SampleMapType},
};
use anyhow::{anyhow, bail, Context, Result};
use log::debug;
use std::{any::Any, fmt::Display, sync::Arc};

use self::{
    checker::check_export_types_btf,
    event_handlers::{buffer, get_plain_text_checked_types_header, sample_map},
    type_descriptor::{CheckedExportedMember, TypeDescriptor},
};

pub(crate) mod checker;
pub(crate) mod data_dumper;
pub(crate) mod event_handlers;
#[cfg(test)]
mod tests;
/// Contains utilities to describe where to obtain the export type of a map
pub mod type_descriptor;
#[derive(Clone, Copy)]
/// Describe the export format type
pub enum ExportFormatType {
    /// Use human-readable texts to output
    PlainText,
    /// Use machine-readable json to output
    Json,
    /// Only call the callback with raw buffer
    RawEvent,
}
#[derive(Debug)]
/// Represents a sample data that the user will receive
pub enum ReceivedEventData<'a> {
    /// Raw buffer. will be used on simple value sampling pairing with `ExportFormatType::RawEvent`
    Buffer(&'a [u8]),
    /// KeyValue Buffer. will be used on key-value sampling pairing with ` ExportFormatType::RawEvent`
    KeyValueBuffer {
        key: &'a [u8],
        value: &'a [u8],
    },
    // Plain string. will be used on ` ExportFormatType::PlainText`
    PlainText(&'a str),
    // Json string. Will be used on `ExportFormatType::Json`
    JsonText(&'a str),
}
impl<'a> Display for ReceivedEventData<'a> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ReceivedEventData::Buffer(b) => {
                write!(f, "{b:?}")?;
            }
            ReceivedEventData::KeyValueBuffer { key, value } => {
                write!(f, "key: {key:?} value: {value:?}")?;
            }
            ReceivedEventData::PlainText(s) | ReceivedEventData::JsonText(s) => {
                write!(f, "{s}")?;
            }
        }
        Ok(())
    }
}

impl<'a> ReceivedEventData<'a> {
    /// Simply get a &[u8] from this received data, which is the same as the one that the C++ version callback receives
    pub fn trivally_to_plain_bytes(&self) -> &'a [u8] {
        match self {
            ReceivedEventData::Buffer(buf) => buf,
            ReceivedEventData::KeyValueBuffer { value, .. } => value,
            ReceivedEventData::PlainText(txt) => txt.as_bytes(),
            ReceivedEventData::JsonText(txt) => txt.as_bytes(),
        }
    }
}

/// A handler to receive events provided by ebpf kernel program
pub trait EventHandler {
    fn handle_event(&self, context: Option<Arc<dyn Any>>, data: ReceivedEventData);
}

pub(crate) enum ExporterInternalImplementation {
    BufferValueProcessor {
        /// internal handler to process export data to a given format
        event_processor: Box<dyn InternalBufferValueEventProcessor>,
        /// exported value types
        checked_types: Vec<CheckedExportedMember>,
    },
    KeyValueMapProcessor {
        /// internal handler to sample map data to a given format
        event_processor: Box<dyn InternalSampleMapProcessor>,
        /// export map key types meta data
        checked_key_types: Vec<CheckedExportedMember>,
        /// export map value types meta data
        checked_value_types: Vec<CheckedExportedMember>,
        /// Config of the sampling map
        sample_map_config: MapSampleMeta,
    },
}

/// dump export event in user space
pub struct EventExporter {
    /// user defined handler to process export data
    pub(crate) user_export_event_handler: Option<Arc<dyn EventHandler>>,
    pub(crate) internal_impl: ExporterInternalImplementation,
    /// user-defined context
    pub(crate) user_ctx: Option<Arc<dyn Any>>,
    pub(crate) btf_container: Arc<BtfContainer>,
}

impl EventExporter {
    pub(crate) fn dump_data_to_user_callback_or_stdout(&self, data: ReceivedEventData) {
        dump_data_to_user_callback_or_stdout(
            self.user_export_event_handler.clone(),
            self.user_ctx.clone(),
            data,
        );
    }
}
pub(crate) fn dump_data_to_user_callback_or_stdout(
    user_export_event_handler: Option<Arc<dyn EventHandler>>,
    user_ctx: Option<Arc<dyn Any>>,
    data: ReceivedEventData,
) {
    if let Some(callback) = user_export_event_handler.as_ref() {
        callback.handle_event(user_ctx, data);
    } else {
        println!("{data}");
    }
}
pub(crate) trait InternalBufferValueEventProcessor {
    fn handle_event(&self, data: &[u8]) -> Result<()>;
}

pub(crate) trait InternalSampleMapProcessor {
    fn handle_event(&self, key_buffer: &[u8], value_buffer: &[u8]) -> Result<()>;
}

/// The builder of the EventExporter
pub struct EventExporterBuilder {
    export_format: ExportFormatType,
    export_event_handler: Option<Arc<dyn EventHandler>>,
    user_ctx: Option<Arc<dyn Any>>,
}

impl Default for EventExporterBuilder {
    fn default() -> Self {
        Self {
            export_format: ExportFormatType::PlainText,
            export_event_handler: None,
            user_ctx: None,
        }
    }
}

impl EventExporterBuilder {
    /// Create a default Builder, with export_format defaults to `ExportFormat::PlainText`
    pub fn new() -> Self {
        Self::default()
    }
    /// Set the export format that the ebpf program will export
    pub fn set_export_format(self, fmt: ExportFormatType) -> Self {
        Self {
            export_format: fmt,
            ..self
        }
    }
    /// Set a user-defined event handler callback
    pub fn set_export_event_handler(self, handler: Arc<dyn EventHandler>) -> Self {
        Self {
            export_event_handler: Some(handler),
            ..self
        }
    }
    /// Set the user-defined context
    pub fn set_user_context<T: Any>(self, ctx: T) -> Self {
        Self {
            user_ctx: Some(Arc::new(ctx)),
            ..self
        }
    }
    /// Build an exporter use TypeDescriptor. Which can easily specify the source to obtain the value type
    pub fn build_for_single_value_with_type_descriptor(
        self,
        export_type: TypeDescriptor,
        btf_container: Arc<BtfContainer>,
        intepreter: &BufferValueInterpreter,
    ) -> Result<Arc<EventExporter>> {
        let mut checked_exported_members =
            export_type.build_checked_exported_members(btf_container.borrow_btf())?;
        if matches!(intepreter, BufferValueInterpreter::StackTrace { .. })
            && !matches!(self.export_format, ExportFormatType::PlainText)
        {
            bail!("Intepreter `stack_trace` could only be paired with plaintext export format");
        }
        Ok(Arc::new_cyclic(move |me| {
            let internal_event_processor: Box<dyn InternalBufferValueEventProcessor> =
                match (self.export_format, intepreter) {
                    (ExportFormatType::Json, BufferValueInterpreter::DefaultStruct) => {
                        Box::new(buffer::JsonExportEventHandler {
                            exporter: me.clone(),
                        })
                    }
                    (
                        ExportFormatType::PlainText,
                        BufferValueInterpreter::StackTrace {
                            field_map,
                            with_symbols,
                        },
                    ) => {
                        debug!("Using stack trace exporter");
                        Box::new(buffer::PlainTextStackTraceExportEventHandler {
                            exporter: me.clone(),
                            field_mapping: field_map.clone(),
                            with_symbols: *with_symbols,
                        })
                    }
                    (ExportFormatType::PlainText, BufferValueInterpreter::DefaultStruct) => {
                        let header = get_plain_text_checked_types_header(
                            &mut checked_exported_members,
                            "TIME     ",
                        );
                        dump_data_to_user_callback_or_stdout(
                            self.export_event_handler.clone(),
                            self.user_ctx.clone(),
                            ReceivedEventData::PlainText(header.as_str()),
                        );
                        Box::new(buffer::PlainStringExportEventHandler {
                            exporter: me.clone(),
                        })
                    }

                    (ExportFormatType::RawEvent, BufferValueInterpreter::DefaultStruct) => {
                        Box::new(buffer::RawExportEventHandler {
                            exporter: me.clone(),
                        })
                    }
                    (_, _) => unreachable!("Unexpected exportformattype + intepreter"),
                };
            EventExporter {
                user_export_event_handler: self.export_event_handler,
                user_ctx: self.user_ctx,
                btf_container,
                internal_impl: ExporterInternalImplementation::BufferValueProcessor {
                    event_processor: internal_event_processor,
                    checked_types: checked_exported_members,
                },
            }
        }))
    }
    /// Build an EventExporter which is suitable for processing value-only data
    /// export_types: The type of the data that kernel program sends. The types will be verified using BTF
    /// btf_container: The BTF information
    pub fn build_for_single_value(
        self,
        export_type: &ExportedTypesStructMeta,
        btf_container: Arc<BtfContainer>,
        intepreter: &BufferValueInterpreter,
    ) -> Result<Arc<EventExporter>> {
        let checked_members = check_export_types_btf(export_type, btf_container.borrow_btf())?;
        Self::build_for_single_value_with_type_descriptor(
            self,
            TypeDescriptor::CheckedMembers(checked_members),
            btf_container,
            intepreter,
        )
    }
    /// Build an EventExporter, but use TypeDescriptor to indicate where to fetch the value type
    pub fn build_for_key_value_with_type_desc(
        self,
        key_export_type: TypeDescriptor,
        value_export_type: TypeDescriptor,
        sample_config: &MapSampleMeta,
        btf_container: Arc<BtfContainer>,
    ) -> Result<Arc<EventExporter>> {
        let mut checked_key_types =
            key_export_type.build_checked_exported_members(btf_container.borrow_btf())?;
        let mut checked_value_types =
            value_export_type.build_checked_exported_members(btf_container.borrow_btf())?;

        if matches!(self.export_format, ExportFormatType::PlainText)
            && matches!(sample_config.ty, SampleMapType::LinearHist)
        {
            bail!("Linear hist sampling is not supported now");
        }
        Ok(Arc::new_cyclic(move |me| {
            let internal_sample_map_processor: Box<dyn InternalSampleMapProcessor> = match self
                .export_format
            {
                ExportFormatType::PlainText => match sample_config.ty {
                    SampleMapType::Log2Hist => Box::new(sample_map::Log2HistExportEventHandler {
                        exporter: me.clone(),
                    }),
                    SampleMapType::DefaultKV => {
                        let header = String::from("TIME     ");
                        let header =
                            get_plain_text_checked_types_header(&mut checked_key_types, header);
                        let header =
                            get_plain_text_checked_types_header(&mut checked_value_types, header);
                        dump_data_to_user_callback_or_stdout(
                            self.export_event_handler.clone(),
                            self.user_ctx.clone(),
                            ReceivedEventData::PlainText(header.as_str()),
                        );
                        Box::new(sample_map::DefaultKVStringExportEventHandler {
                            exporter: me.clone(),
                        })
                    }
                    SampleMapType::LinearHist => unreachable!(),
                },
                ExportFormatType::Json => Box::new(sample_map::JsonExportEventHandler {
                    exporter: me.clone(),
                }),
                ExportFormatType::RawEvent => Box::new(sample_map::RawExportEventHandler {
                    exporter: me.clone(),
                }),
            };
            EventExporter {
                user_export_event_handler: self.export_event_handler,
                internal_impl: ExporterInternalImplementation::KeyValueMapProcessor {
                    event_processor: internal_sample_map_processor,
                    checked_key_types,
                    checked_value_types,
                    sample_map_config: sample_config.clone(),
                },
                user_ctx: self.user_ctx,
                btf_container,
            }
        }))
    }
    /// Build an EventExporter to process sampling map
    /// key_type_id: The type id of the map key
    /// value_type_id: The type id of the map value
    /// sample_config: Detailed configuration of the sampling map
    /// export_types: Value types of the map, will be verified
    /// btf_container: The btf information
    pub fn build_for_key_value(
        self,
        key_type_id: u32,
        value_type_id: u32,
        sample_config: &MapSampleMeta,
        export_type: &ExportedTypesStructMeta,
        btf_container: Arc<BtfContainer>,
    ) -> Result<Arc<EventExporter>> {
        let btf = btf_container.borrow_btf();
        let checked_key_types = check_sample_types_btf(btf, key_type_id, None)
            .with_context(|| anyhow!("Failed to check key type"))?;
        let checked_value_types =
            check_sample_types_btf(btf, value_type_id, Some(export_type.clone()))
                .with_context(|| anyhow!("Failed to check value type"))?;
        self.build_for_key_value_with_type_desc(
            TypeDescriptor::CheckedMembers(checked_key_types),
            TypeDescriptor::CheckedMembers(checked_value_types),
            sample_config,
            btf_container,
        )
    }
}