libp2p_swarm/handler/
select.rs

1// Copyright 2019 Parity Technologies (UK) Ltd.
2//
3// Permission is hereby granted, free of charge, to any person obtaining a
4// copy of this software and associated documentation files (the "Software"),
5// to deal in the Software without restriction, including without limitation
6// the rights to use, copy, modify, merge, publish, distribute, sublicense,
7// and/or sell copies of the Software, and to permit persons to whom the
8// Software is furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
14// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
18// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
19// DEALINGS IN THE SOFTWARE.
20
21use crate::handler::{
22    AddressChange, ConnectionEvent, ConnectionHandler, ConnectionHandlerEvent, DialUpgradeError,
23    FullyNegotiatedInbound, FullyNegotiatedOutbound, InboundUpgradeSend, KeepAlive,
24    ListenUpgradeError, OutboundUpgradeSend, StreamUpgradeError, SubstreamProtocol,
25};
26use crate::upgrade::SendWrapper;
27use either::Either;
28use futures::future;
29use libp2p_core::upgrade::SelectUpgrade;
30use std::{cmp, task::Context, task::Poll};
31
32/// Implementation of [`ConnectionHandler`] that combines two protocols into one.
33#[derive(Debug, Clone)]
34pub struct ConnectionHandlerSelect<TProto1, TProto2> {
35    /// The first protocol.
36    proto1: TProto1,
37    /// The second protocol.
38    proto2: TProto2,
39}
40
41impl<TProto1, TProto2> ConnectionHandlerSelect<TProto1, TProto2> {
42    /// Builds a [`ConnectionHandlerSelect`].
43    pub(crate) fn new(proto1: TProto1, proto2: TProto2) -> Self {
44        ConnectionHandlerSelect { proto1, proto2 }
45    }
46
47    pub fn into_inner(self) -> (TProto1, TProto2) {
48        (self.proto1, self.proto2)
49    }
50}
51
52impl<S1OOI, S2OOI, S1OP, S2OP>
53    FullyNegotiatedOutbound<Either<SendWrapper<S1OP>, SendWrapper<S2OP>>, Either<S1OOI, S2OOI>>
54where
55    S1OP: OutboundUpgradeSend,
56    S2OP: OutboundUpgradeSend,
57    S1OOI: Send + 'static,
58    S2OOI: Send + 'static,
59{
60    pub(crate) fn transpose(
61        self,
62    ) -> Either<FullyNegotiatedOutbound<S1OP, S1OOI>, FullyNegotiatedOutbound<S2OP, S2OOI>> {
63        match self {
64            FullyNegotiatedOutbound {
65                protocol: future::Either::Left(protocol),
66                info: Either::Left(info),
67            } => Either::Left(FullyNegotiatedOutbound { protocol, info }),
68            FullyNegotiatedOutbound {
69                protocol: future::Either::Right(protocol),
70                info: Either::Right(info),
71            } => Either::Right(FullyNegotiatedOutbound { protocol, info }),
72            _ => panic!("wrong API usage: the protocol doesn't match the upgrade info"),
73        }
74    }
75}
76
77impl<S1IP, S1IOI, S2IP, S2IOI>
78    FullyNegotiatedInbound<SelectUpgrade<SendWrapper<S1IP>, SendWrapper<S2IP>>, (S1IOI, S2IOI)>
79where
80    S1IP: InboundUpgradeSend,
81    S2IP: InboundUpgradeSend,
82{
83    pub(crate) fn transpose(
84        self,
85    ) -> Either<FullyNegotiatedInbound<S1IP, S1IOI>, FullyNegotiatedInbound<S2IP, S2IOI>> {
86        match self {
87            FullyNegotiatedInbound {
88                protocol: future::Either::Left(protocol),
89                info: (i1, _i2),
90            } => Either::Left(FullyNegotiatedInbound { protocol, info: i1 }),
91            FullyNegotiatedInbound {
92                protocol: future::Either::Right(protocol),
93                info: (_i1, i2),
94            } => Either::Right(FullyNegotiatedInbound { protocol, info: i2 }),
95        }
96    }
97}
98
99impl<S1OOI, S2OOI, S1OP, S2OP>
100    DialUpgradeError<Either<S1OOI, S2OOI>, Either<SendWrapper<S1OP>, SendWrapper<S2OP>>>
101where
102    S1OP: OutboundUpgradeSend,
103    S2OP: OutboundUpgradeSend,
104    S1OOI: Send + 'static,
105    S2OOI: Send + 'static,
106{
107    pub(crate) fn transpose(
108        self,
109    ) -> Either<DialUpgradeError<S1OOI, S1OP>, DialUpgradeError<S2OOI, S2OP>> {
110        match self {
111            DialUpgradeError {
112                info: Either::Left(info),
113                error: StreamUpgradeError::Apply(Either::Left(err)),
114            } => Either::Left(DialUpgradeError {
115                info,
116                error: StreamUpgradeError::Apply(err),
117            }),
118            DialUpgradeError {
119                info: Either::Right(info),
120                error: StreamUpgradeError::Apply(Either::Right(err)),
121            } => Either::Right(DialUpgradeError {
122                info,
123                error: StreamUpgradeError::Apply(err),
124            }),
125            DialUpgradeError {
126                info: Either::Left(info),
127                error: e,
128            } => Either::Left(DialUpgradeError {
129                info,
130                error: e.map_upgrade_err(|_| panic!("already handled above")),
131            }),
132            DialUpgradeError {
133                info: Either::Right(info),
134                error: e,
135            } => Either::Right(DialUpgradeError {
136                info,
137                error: e.map_upgrade_err(|_| panic!("already handled above")),
138            }),
139        }
140    }
141}
142
143impl<TProto1, TProto2> ConnectionHandlerSelect<TProto1, TProto2>
144where
145    TProto1: ConnectionHandler,
146    TProto2: ConnectionHandler,
147{
148    fn on_listen_upgrade_error(
149        &mut self,
150        ListenUpgradeError {
151            info: (i1, i2),
152            error,
153        }: ListenUpgradeError<
154            <Self as ConnectionHandler>::InboundOpenInfo,
155            <Self as ConnectionHandler>::InboundProtocol,
156        >,
157    ) {
158        match error {
159            Either::Left(error) => {
160                self.proto1
161                    .on_connection_event(ConnectionEvent::ListenUpgradeError(ListenUpgradeError {
162                        info: i1,
163                        error,
164                    }));
165            }
166            Either::Right(error) => {
167                self.proto2
168                    .on_connection_event(ConnectionEvent::ListenUpgradeError(ListenUpgradeError {
169                        info: i2,
170                        error,
171                    }));
172            }
173        }
174    }
175}
176
177impl<TProto1, TProto2> ConnectionHandler for ConnectionHandlerSelect<TProto1, TProto2>
178where
179    TProto1: ConnectionHandler,
180    TProto2: ConnectionHandler,
181{
182    type FromBehaviour = Either<TProto1::FromBehaviour, TProto2::FromBehaviour>;
183    type ToBehaviour = Either<TProto1::ToBehaviour, TProto2::ToBehaviour>;
184    #[allow(deprecated)]
185    type Error = Either<TProto1::Error, TProto2::Error>;
186    type InboundProtocol = SelectUpgrade<
187        SendWrapper<<TProto1 as ConnectionHandler>::InboundProtocol>,
188        SendWrapper<<TProto2 as ConnectionHandler>::InboundProtocol>,
189    >;
190    type OutboundProtocol =
191        Either<SendWrapper<TProto1::OutboundProtocol>, SendWrapper<TProto2::OutboundProtocol>>;
192    type OutboundOpenInfo = Either<TProto1::OutboundOpenInfo, TProto2::OutboundOpenInfo>;
193    type InboundOpenInfo = (TProto1::InboundOpenInfo, TProto2::InboundOpenInfo);
194
195    fn listen_protocol(&self) -> SubstreamProtocol<Self::InboundProtocol, Self::InboundOpenInfo> {
196        let proto1 = self.proto1.listen_protocol();
197        let proto2 = self.proto2.listen_protocol();
198        let timeout = *std::cmp::max(proto1.timeout(), proto2.timeout());
199        let (u1, i1) = proto1.into_upgrade();
200        let (u2, i2) = proto2.into_upgrade();
201        let choice = SelectUpgrade::new(SendWrapper(u1), SendWrapper(u2));
202        SubstreamProtocol::new(choice, (i1, i2)).with_timeout(timeout)
203    }
204
205    fn on_behaviour_event(&mut self, event: Self::FromBehaviour) {
206        match event {
207            Either::Left(event) => self.proto1.on_behaviour_event(event),
208            Either::Right(event) => self.proto2.on_behaviour_event(event),
209        }
210    }
211
212    fn connection_keep_alive(&self) -> KeepAlive {
213        cmp::max(
214            self.proto1.connection_keep_alive(),
215            self.proto2.connection_keep_alive(),
216        )
217    }
218
219    #[allow(deprecated)]
220    fn poll(
221        &mut self,
222        cx: &mut Context<'_>,
223    ) -> Poll<
224        ConnectionHandlerEvent<
225            Self::OutboundProtocol,
226            Self::OutboundOpenInfo,
227            Self::ToBehaviour,
228            Self::Error,
229        >,
230    > {
231        match self.proto1.poll(cx) {
232            Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(event)) => {
233                return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(Either::Left(event)));
234            }
235            #[allow(deprecated)]
236            Poll::Ready(ConnectionHandlerEvent::Close(event)) => {
237                return Poll::Ready(ConnectionHandlerEvent::Close(Either::Left(event)));
238            }
239            Poll::Ready(ConnectionHandlerEvent::OutboundSubstreamRequest { protocol }) => {
240                return Poll::Ready(ConnectionHandlerEvent::OutboundSubstreamRequest {
241                    protocol: protocol
242                        .map_upgrade(|u| Either::Left(SendWrapper(u)))
243                        .map_info(Either::Left),
244                });
245            }
246            Poll::Ready(ConnectionHandlerEvent::ReportRemoteProtocols(support)) => {
247                return Poll::Ready(ConnectionHandlerEvent::ReportRemoteProtocols(support));
248            }
249            Poll::Pending => (),
250        };
251
252        match self.proto2.poll(cx) {
253            Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(event)) => {
254                return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(Either::Right(
255                    event,
256                )));
257            }
258            #[allow(deprecated)]
259            Poll::Ready(ConnectionHandlerEvent::Close(event)) => {
260                return Poll::Ready(ConnectionHandlerEvent::Close(Either::Right(event)));
261            }
262            Poll::Ready(ConnectionHandlerEvent::OutboundSubstreamRequest { protocol }) => {
263                return Poll::Ready(ConnectionHandlerEvent::OutboundSubstreamRequest {
264                    protocol: protocol
265                        .map_upgrade(|u| Either::Right(SendWrapper(u)))
266                        .map_info(Either::Right),
267                });
268            }
269            Poll::Ready(ConnectionHandlerEvent::ReportRemoteProtocols(support)) => {
270                return Poll::Ready(ConnectionHandlerEvent::ReportRemoteProtocols(support));
271            }
272            Poll::Pending => (),
273        };
274
275        Poll::Pending
276    }
277
278    fn on_connection_event(
279        &mut self,
280        event: ConnectionEvent<
281            Self::InboundProtocol,
282            Self::OutboundProtocol,
283            Self::InboundOpenInfo,
284            Self::OutboundOpenInfo,
285        >,
286    ) {
287        match event {
288            ConnectionEvent::FullyNegotiatedOutbound(fully_negotiated_outbound) => {
289                match fully_negotiated_outbound.transpose() {
290                    Either::Left(f) => self
291                        .proto1
292                        .on_connection_event(ConnectionEvent::FullyNegotiatedOutbound(f)),
293                    Either::Right(f) => self
294                        .proto2
295                        .on_connection_event(ConnectionEvent::FullyNegotiatedOutbound(f)),
296                }
297            }
298            ConnectionEvent::FullyNegotiatedInbound(fully_negotiated_inbound) => {
299                match fully_negotiated_inbound.transpose() {
300                    Either::Left(f) => self
301                        .proto1
302                        .on_connection_event(ConnectionEvent::FullyNegotiatedInbound(f)),
303                    Either::Right(f) => self
304                        .proto2
305                        .on_connection_event(ConnectionEvent::FullyNegotiatedInbound(f)),
306                }
307            }
308            ConnectionEvent::AddressChange(address) => {
309                self.proto1
310                    .on_connection_event(ConnectionEvent::AddressChange(AddressChange {
311                        new_address: address.new_address,
312                    }));
313
314                self.proto2
315                    .on_connection_event(ConnectionEvent::AddressChange(AddressChange {
316                        new_address: address.new_address,
317                    }));
318            }
319            ConnectionEvent::DialUpgradeError(dial_upgrade_error) => {
320                match dial_upgrade_error.transpose() {
321                    Either::Left(err) => self
322                        .proto1
323                        .on_connection_event(ConnectionEvent::DialUpgradeError(err)),
324                    Either::Right(err) => self
325                        .proto2
326                        .on_connection_event(ConnectionEvent::DialUpgradeError(err)),
327                }
328            }
329            ConnectionEvent::ListenUpgradeError(listen_upgrade_error) => {
330                self.on_listen_upgrade_error(listen_upgrade_error)
331            }
332            ConnectionEvent::LocalProtocolsChange(supported_protocols) => {
333                self.proto1
334                    .on_connection_event(ConnectionEvent::LocalProtocolsChange(
335                        supported_protocols.clone(),
336                    ));
337                self.proto2
338                    .on_connection_event(ConnectionEvent::LocalProtocolsChange(
339                        supported_protocols,
340                    ));
341            }
342            ConnectionEvent::RemoteProtocolsChange(supported_protocols) => {
343                self.proto1
344                    .on_connection_event(ConnectionEvent::RemoteProtocolsChange(
345                        supported_protocols.clone(),
346                    ));
347                self.proto2
348                    .on_connection_event(ConnectionEvent::RemoteProtocolsChange(
349                        supported_protocols,
350                    ));
351            }
352        }
353    }
354}