1use 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#[derive(Debug, Clone)]
34pub struct ConnectionHandlerSelect<TProto1, TProto2> {
35 proto1: TProto1,
37 proto2: TProto2,
39}
40
41impl<TProto1, TProto2> ConnectionHandlerSelect<TProto1, TProto2> {
42 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}