integration_tests_sv2 0.1.1

Sv2 Integration Tests Framework
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

use std::fmt;

use crate::types::MsgType;
use stratum_apps::stratum_core::parsers_sv2::AnyMessage;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MessageDirection {
    ToDownstream,
    ToUpstream,
}

impl fmt::Display for MessageDirection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            MessageDirection::ToDownstream => write!(f, "downstream"),
            MessageDirection::ToUpstream => write!(f, "upstream"),
        }
    }
}

/// Represents an action that [`Sniffer`] can take on intercepted messages.
#[derive(Debug, Clone)]
pub enum InterceptAction {
    /// Prevents a message from being forwarded and stored into the message aggregator.
    IgnoreMessage(IgnoreMessage),
    /// Intercepts and modifies a message before forwarding it.
    ReplaceMessage(Box<ReplaceMessage>),
}

impl InterceptAction {
    /// Returns the action if it is `IgnoreMessage` or `ReplaceMessage`
    /// with the specified message type.
    pub fn find_matching_action(
        &self,
        msg_type: MsgType,
        direction: MessageDirection,
    ) -> Option<&Self> {
        match self {
            InterceptAction::IgnoreMessage(bm)
                if bm.direction == direction && bm.expected_message_type == msg_type =>
            {
                Some(self)
            }

            InterceptAction::ReplaceMessage(im)
                if im.direction == direction && im.expected_message_type == msg_type =>
            {
                Some(self)
            }

            _ => None,
        }
    }
}
/// Defines an action that prevents a message from being forwarded.
///
/// When a message matching the specified type and direction is intercepted,
/// it will not be added to the message aggregator for inspection and will not be
/// forwarded to the destination. All other messages will continue to be forwarded normally.
#[derive(Debug, Clone)]
pub struct IgnoreMessage {
    direction: MessageDirection,
    expected_message_type: MsgType,
}

impl IgnoreMessage {
    /// Creates a new [`IgnoreMessage`] action.
    ///
    /// - `direction`: The direction of the message to be ignored.
    /// - `expected_message_type`: The type of message to be ignored.
    pub fn new(direction: MessageDirection, expected_message_type: MsgType) -> Self {
        IgnoreMessage {
            direction,
            expected_message_type,
        }
    }
}

impl From<IgnoreMessage> for InterceptAction {
    fn from(value: IgnoreMessage) -> Self {
        InterceptAction::IgnoreMessage(value)
    }
}

/// Allows [`Sniffer`] to replace some intercepted message before forwarding it.
#[derive(Debug, Clone)]
pub struct ReplaceMessage {
    direction: MessageDirection,
    expected_message_type: MsgType,
    pub(crate) replacement_message: AnyMessage<'static>,
}

impl ReplaceMessage {
    /// Constructor of `ReplaceMessage`
    /// - `direction`: direction of message to be intercepted and replaced
    /// - `expected_message_type`: type of message to be intercepted and replaced
    /// - `replacement_message`: message to replace the intercepted one
    /// - `replacement_message_type`: type of message to replace the intercepted one
    pub fn new(
        direction: MessageDirection,
        expected_message_type: MsgType,
        replacement_message: AnyMessage<'static>,
    ) -> Self {
        Self {
            direction,
            expected_message_type,
            replacement_message,
        }
    }
}

impl From<ReplaceMessage> for InterceptAction {
    fn from(value: ReplaceMessage) -> Self {
        InterceptAction::ReplaceMessage(Box::new(value))
    }
}