pool_sv2 0.2.1

SV2 pool role
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

use crate::{
    downstream::Downstream,
    error::{self, PoolError, PoolErrorKind},
};
use std::convert::TryInto;
use stratum_apps::{
    stratum_core::{
        binary_sv2::Seq064K,
        extensions_sv2::{RequestExtensions, RequestExtensionsError, RequestExtensionsSuccess},
        handlers_sv2::HandleExtensionsFromClientAsync,
        parsers_sv2::{AnyMessage, Tlv},
    },
    utils::types::Sv2Frame,
};
use tracing::{error, info};

#[cfg_attr(not(test), hotpath::measure_all)]
impl HandleExtensionsFromClientAsync for Downstream {
    type Error = PoolError<error::Downstream>;

    fn get_negotiated_extensions_with_client(
        &self,
        _client_id: Option<usize>,
    ) -> Result<Vec<u16>, Self::Error> {
        Ok(self
            .downstream_data
            .super_safe_lock(|data| data.negotiated_extensions.clone()))
    }

    async fn handle_request_extensions(
        &mut self,
        _client_id: Option<usize>,
        msg: RequestExtensions<'_>,
        _tlv_fields: Option<&[Tlv]>,
    ) -> Result<(), Self::Error> {
        let requested: Vec<u16> = msg.requested_extensions.clone().into_inner();

        info!(
            "Downstream {}: Received RequestExtensions: request_id={}, requested={:?}",
            self.downstream_id, msg.request_id, requested
        );

        // Get supported and required extensions from downstream data
        let (supported_extensions, required_extensions) = (
            self.supported_extensions.clone(),
            self.required_extensions.clone(),
        );

        // Determine which requested extensions we support
        let mut supported: Vec<u16> = Vec::new();
        let mut unsupported: Vec<u16> = Vec::new();

        for ext in &requested {
            if supported_extensions.contains(ext) {
                supported.push(*ext);
            } else {
                unsupported.push(*ext);
            }
        }

        // Check which required extensions the client didn't request
        let missing_required: Vec<u16> = required_extensions
            .iter()
            .filter(|ext| !requested.contains(ext))
            .copied()
            .collect();

        // Determine response based on spec rules:
        // - Success: If at least one extension is supported AND all required extensions are present
        // - Error: If no extensions are supported OR required extensions are missing
        let should_send_error = supported.is_empty() || !missing_required.is_empty();

        if should_send_error {
            // Send error response
            error!(
                "Downstream {}: Extension negotiation error: requested={:?}, supported={:?}, unsupported={:?}, missing_required={:?}",
                self.downstream_id, requested, supported, unsupported, missing_required
            );

            let error = RequestExtensionsError {
                request_id: msg.request_id,
                unsupported_extensions: Seq064K::new(unsupported).map_err(PoolError::shutdown)?,
                required_extensions: Seq064K::new(missing_required.clone())
                    .map_err(PoolError::shutdown)?,
            };

            let frame: Sv2Frame = AnyMessage::Extensions(error.into_static().into())
                .try_into()
                .map_err(PoolError::shutdown)?;
            self.downstream_channel
                .downstream_sender
                .send(frame)
                .await
                .map_err(|_| {
                    PoolError::disconnect(PoolErrorKind::ChannelErrorSender, self.downstream_id)
                })?;

            // If required extensions are missing, the server SHOULD disconnect the client
            if !missing_required.is_empty() {
                error!(
                    "Downstream {}: Client does not support required extensions {:?}. Server MUST disconnect.",
                    self.downstream_id, missing_required
                );
                Err(PoolError::disconnect(
                    PoolErrorKind::ClientDoesNotSupportRequiredExtensions(missing_required),
                    self.downstream_id,
                ))?;
            }
        } else {
            // Send success response with the subset of extensions we both support
            info!(
                "Downstream {}: Extension negotiation success: requested={:?}, negotiated={:?}",
                self.downstream_id, requested, supported
            );

            // Store the negotiated extensions in the shared downstream data
            self.downstream_data.super_safe_lock(|data| {
                data.negotiated_extensions = supported.clone();
            });

            let success = RequestExtensionsSuccess {
                request_id: msg.request_id,
                supported_extensions: Seq064K::new(supported.clone())
                    .map_err(PoolError::shutdown)?,
            };

            let frame: Sv2Frame = AnyMessage::Extensions(success.into_static().into())
                .try_into()
                .map_err(PoolError::shutdown)?;
            self.downstream_channel
                .downstream_sender
                .send(frame)
                .await
                .map_err(|_| {
                    PoolError::disconnect(PoolErrorKind::ChannelErrorSender, self.downstream_id)
                })?;

            info!(
                "Downstream {}: Stored negotiated extensions: {:?}",
                self.downstream_id, supported
            );
        }

        Ok(())
    }
}