soth-mitm 0.2.2

Rust intercepting proxy crate with deterministic handler/event contracts for SOTH.
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

use super::close_codes::CloseReasonCode;
use super::event_emitters::emit_stream_closed;
use super::flow_hooks::FlowHooks;
use super::io_timeouts::{is_idle_watchdog_timeout, is_stream_stage_timeout};
use super::runtime_governor;
use super::websocket_relay::{relay_websocket_connection, DeflateConfig};
use super::BufferedConn;
use crate::engine::MitmEngine;
use crate::observe::{EventConsumer, FlowContext};
use crate::policy::PolicyEngine;
use crate::protocol::ApplicationProtocol;
use std::io;
use std::sync::Arc;
use tokio::io::{AsyncRead, AsyncWrite};

pub(crate) async fn finalize_websocket_upgrade<P, S, D, U>(
    engine: Arc<MitmEngine<P, S>>,
    runtime_governor: Arc<runtime_governor::RuntimeGovernor>,
    flow_hooks: Arc<dyn FlowHooks>,
    tunnel_context: &FlowContext,
    downstream: BufferedConn<D>,
    upstream: BufferedConn<U>,
    mut bytes_from_client: u64,
    mut bytes_from_server: u64,
    deflate_config: Option<DeflateConfig>,
) -> io::Result<()>
where
    P: PolicyEngine + Send + Sync + 'static,
    S: EventConsumer + Send + Sync + 'static,
    D: AsyncRead + AsyncWrite + Unpin + Send + 'static,
    U: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
    let websocket_context = FlowContext {
        protocol: ApplicationProtocol::WebSocket,
        ..tunnel_context.clone()
    };
    match relay_websocket_connection(
        Arc::clone(&engine),
        runtime_governor,
        flow_hooks,
        websocket_context.clone(),
        downstream,
        upstream,
        deflate_config,
    )
    .await
    {
        Ok(outcome) => {
            bytes_from_client += outcome.bytes_from_client;
            bytes_from_server += outcome.bytes_from_server;
            emit_stream_closed(
                &engine,
                websocket_context,
                CloseReasonCode::WebSocketCompleted,
                None,
                Some(bytes_from_client),
                Some(bytes_from_server),
            );
        }
        Err(error) => {
            let reason = if is_idle_watchdog_timeout(&error) {
                CloseReasonCode::IdleWatchdogTimeout
            } else if is_stream_stage_timeout(&error) {
                CloseReasonCode::StreamStageTimeout
            } else {
                CloseReasonCode::WebSocketError
            };
            emit_stream_closed(
                &engine,
                websocket_context,
                reason,
                Some(error.to_string()),
                Some(bytes_from_client),
                Some(bytes_from_server),
            );
        }
    }
    Ok(())
}