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
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
// Copyright 2015-2018 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! This module contains all the types for demuxing DNS oriented streams.

use std::marker::PhantomData;
use std::pin::Pin;
use std::task::{Context, Poll};

use futures::channel::mpsc::{unbounded, UnboundedReceiver};
use futures::stream::{Peekable, Stream, StreamExt};
use futures::{Future, FutureExt};
use log::{debug, warn};

use crate::error::*;
use crate::xfer::dns_handle::DnsHandle;
use crate::xfer::OneshotDnsResponseReceiver;
use crate::xfer::{
    BufDnsRequestStreamHandle, DnsRequest, DnsRequestSender, DnsRequestStreamHandle, DnsResponse,
    OneshotDnsRequest,
};
use crate::Time;

/// This is a generic Exchange implemented over multiplexed DNS connection providers.
///
/// The underlying `DnsRequestSender` is expected to multiplex any I/O connections. DnsExchange assumes that the underlying stream is responsible for this.
#[must_use = "futures do nothing unless polled"]
pub struct DnsExchange<R>
where
    R: Future<Output = Result<DnsResponse, ProtoError>> + 'static + Send + Unpin,
{
    sender: BufDnsRequestStreamHandle<R>,
}

impl<R> DnsExchange<R>
where
    R: Future<Output = Result<DnsResponse, ProtoError>> + 'static + Send + Unpin,
{
    /// Initializes a TcpStream with an existing tcp::TcpStream.
    ///
    /// This is intended for use with a TcpListener and Incoming.
    ///
    /// # Arguments
    ///
    /// * `stream` - the established IO stream for communication
    pub fn from_stream<S, TE>(stream: S) -> (Self, DnsExchangeBackground<S, R, TE>)
    where
        S: DnsRequestSender<DnsResponseFuture = R> + 'static + Send + Unpin,
    {
        let (message_sender, outbound_messages) = unbounded();
        let message_sender = DnsRequestStreamHandle::<R>::new(message_sender);

        Self::from_stream_with_receiver(stream, outbound_messages, message_sender)
    }

    /// Wraps a stream where a sender and receiver have already been established
    pub fn from_stream_with_receiver<S, TE>(
        stream: S,
        receiver: UnboundedReceiver<OneshotDnsRequest<R>>,
        sender: DnsRequestStreamHandle<R>,
    ) -> (Self, DnsExchangeBackground<S, R, TE>)
    where
        S: DnsRequestSender<DnsResponseFuture = R> + 'static + Send + Unpin,
    {
        let background = DnsExchangeBackground {
            io_stream: stream,
            outbound_messages: receiver.peekable(),
            marker: PhantomData,
        };

        let sender = BufDnsRequestStreamHandle::new(sender);

        (Self { sender }, background)
    }

    /// Returns a future, which itself wraps a future which is awaiting connection.
    ///
    /// The connect_future should be lazy.
    pub fn connect<F, S, TE>(connect_future: F) -> DnsExchangeConnect<F, S, R, TE>
    where
        F: Future<Output = Result<S, ProtoError>> + 'static + Send + Unpin,
        S: DnsRequestSender<DnsResponseFuture = R> + 'static + Send + Unpin,
        TE: Time + Unpin,
    {
        let (message_sender, outbound_messages) = unbounded();
        let message_sender = DnsRequestStreamHandle::<R>::new(message_sender);

        DnsExchangeConnect::connect(connect_future, outbound_messages, message_sender)
    }
}

impl<R> Clone for DnsExchange<R>
where
    R: Future<Output = Result<DnsResponse, ProtoError>> + 'static + Send + Unpin,
{
    fn clone(&self) -> Self {
        Self {
            sender: self.sender.clone(),
        }
    }
}

impl<Resp> DnsHandle for DnsExchange<Resp>
where
    Resp: Future<Output = Result<DnsResponse, ProtoError>> + 'static + Send + Unpin,
{
    type Response = DnsExchangeSend<Resp>;

    fn send<R: Into<DnsRequest> + Unpin + Send + 'static>(&mut self, request: R) -> Self::Response {
        DnsExchangeSend {
            result: self.sender.send(request),
            _sender: self.sender.clone(), // TODO: this shouldn't be necessary, currently the presence of Senders is what allows the background to track current users, it generally is dropped right after send, this makes sure that there is at least one active after send
        }
    }
}

/// A Future that will resolve to a Response after sending the request
#[must_use = "futures do nothing unless polled"]
pub struct DnsExchangeSend<R>
where
    R: Future<Output = Result<DnsResponse, ProtoError>> + 'static + Send + Unpin,
{
    result: OneshotDnsResponseReceiver<R>,
    _sender: BufDnsRequestStreamHandle<R>,
}

impl<R> Future for DnsExchangeSend<R>
where
    R: Future<Output = Result<DnsResponse, ProtoError>> + 'static + Send + Unpin,
{
    type Output = Result<DnsResponse, ProtoError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        // as long as there is no result, poll the exchange
        self.result.poll_unpin(cx)
    }
}

/// This background future is responsible for driving all network operations for the DNS protocol.
///
/// It must be spawned before any DNS messages are sent.
#[must_use = "futures do nothing unless polled"]
pub struct DnsExchangeBackground<S, R, TE>
where
    S: DnsRequestSender<DnsResponseFuture = R> + 'static + Send + Unpin,
    R: Future<Output = Result<DnsResponse, ProtoError>> + 'static + Send + Unpin,
{
    io_stream: S,
    outbound_messages: Peekable<UnboundedReceiver<OneshotDnsRequest<R>>>,
    marker: PhantomData<TE>,
}

impl<S, R, TE> DnsExchangeBackground<S, R, TE>
where
    S: DnsRequestSender<DnsResponseFuture = R> + 'static + Send + Unpin,
    R: Future<Output = Result<DnsResponse, ProtoError>> + 'static + Send + Unpin,
{
    fn pollable_split(
        &mut self,
    ) -> (
        &mut S,
        &mut Peekable<UnboundedReceiver<OneshotDnsRequest<R>>>,
    ) {
        (&mut self.io_stream, &mut self.outbound_messages)
    }
}

impl<S, R, TE> Future for DnsExchangeBackground<S, R, TE>
where
    S: DnsRequestSender<DnsResponseFuture = R> + 'static + Send + Unpin,
    R: Future<Output = Result<DnsResponse, ProtoError>> + 'static + Send + Unpin,
    TE: Time + Unpin,
{
    type Output = Result<(), ProtoError>;

    #[allow(clippy::unused_unit)]
    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        let (io_stream, outbound_messages) = self.pollable_split();
        let mut io_stream = Pin::new(io_stream);
        let mut outbound_messages = Pin::new(outbound_messages);

        // this will not accept incoming data while there is data to send
        //  makes this self throttling.
        loop {
            // poll the underlying stream, to drive it...
            match io_stream.as_mut().poll_next(cx) {
                // The stream is ready
                Poll::Ready(Some(Ok(()))) => (),
                Poll::Pending => {
                    if io_stream.is_shutdown() {
                        // the io_stream is in a shutdown state, we are only waiting for final results...
                        return Poll::Pending;
                    }

                    // NotReady and not shutdown, see if there are more messages to send
                    ()
                } // underlying stream is complete.
                Poll::Ready(None) => {
                    debug!("io_stream is done, shutting down");
                    // TODO: return shutdown error to anything in the stream?

                    return Poll::Ready(Ok(()));
                }
                Poll::Ready(Some(Err(err))) => {
                    warn!("io_stream hit an error, shutting down: {}", err);

                    return Poll::Ready(Err(err));
                }
            }

            // then see if there is more to send
            match outbound_messages.as_mut().poll_next(cx) {
                // already handled above, here to make sure the poll() pops the next message
                Poll::Ready(Some(dns_request)) => {
                    // if there is no peer, this connection should die...
                    let (dns_request, serial_response): (DnsRequest, _) = dns_request.unwrap();

                    match serial_response
                        .send_response(io_stream.send_message::<TE>(dns_request, cx))
                    {
                        Ok(()) => (),
                        Err(_) => {
                            warn!("failed to associate send_message response to the sender");

                            return Poll::Ready(Err(
                                "failed to associate send_message response to the sender".into(),
                            ));
                        }
                    }
                }
                // On not ready, this is our time to return...
                Poll::Pending => return Poll::Pending,
                Poll::Ready(None) => {
                    // if there is nothing that can use this connection to send messages, then this is done...
                    io_stream.shutdown();

                    // now we'll await the stream to shutdown... see io_stream poll above
                }
            }

            // else we loop to poll on the outbound_messages
        }
    }
}

/// A wrapper for a future DnsExchange connection.
///
/// DnsExchangeConnect is clonable, making it possible to share this if the connection
///  will be shared across threads.
///
/// The future will return a tuple of the DnsExchange (for sending messages) and a background
///  for running the background tasks. The background is optional as only one thread should run
///  the background. If returned, it must be spawned before any dns requests will function.
pub struct DnsExchangeConnect<F, S, R, TE>(DnsExchangeConnectInner<F, S, R, TE>)
where
    F: Future<Output = Result<S, ProtoError>> + 'static + Send + Unpin,
    S: DnsRequestSender<DnsResponseFuture = R> + 'static,
    R: Future<Output = Result<DnsResponse, ProtoError>> + 'static + Send + Unpin,
    TE: Time + Unpin;

impl<F, S, R, TE> DnsExchangeConnect<F, S, R, TE>
where
    F: Future<Output = Result<S, ProtoError>> + 'static + Send + Unpin,
    S: DnsRequestSender<DnsResponseFuture = R> + 'static,
    R: Future<Output = Result<DnsResponse, ProtoError>> + 'static + Send + Unpin,
    TE: Time + Unpin,
{
    fn connect(
        connect_future: F,
        outbound_messages: UnboundedReceiver<OneshotDnsRequest<R>>,
        sender: DnsRequestStreamHandle<R>,
    ) -> Self {
        DnsExchangeConnect(DnsExchangeConnectInner::Connecting {
            connect_future,
            outbound_messages: Some(outbound_messages),
            sender: Some(sender),
        })
    }
}

#[allow(clippy::type_complexity)]
impl<F, S, R, TE> Future for DnsExchangeConnect<F, S, R, TE>
where
    F: Future<Output = Result<S, ProtoError>> + 'static + Send + Unpin,
    S: DnsRequestSender<DnsResponseFuture = R> + 'static + Send + Unpin,
    R: Future<Output = Result<DnsResponse, ProtoError>> + 'static + Send + Unpin,
    TE: Time + Unpin,
{
    type Output = Result<(DnsExchange<R>, DnsExchangeBackground<S, R, TE>), ProtoError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        self.0.poll_unpin(cx)
    }
}

enum DnsExchangeConnectInner<F, S, R, TE>
where
    F: Future<Output = Result<S, ProtoError>> + 'static + Send,
    S: DnsRequestSender<DnsResponseFuture = R> + 'static + Send,
    R: Future<Output = Result<DnsResponse, ProtoError>> + 'static + Send + Unpin,
    TE: Time + Unpin,
{
    Connecting {
        connect_future: F,
        outbound_messages: Option<UnboundedReceiver<OneshotDnsRequest<R>>>,
        sender: Option<DnsRequestStreamHandle<R>>,
    },
    Connected {
        exchange: DnsExchange<R>,
        background: Option<DnsExchangeBackground<S, R, TE>>,
    },
    FailAll {
        error: ProtoError,
        outbound_messages: UnboundedReceiver<OneshotDnsRequest<R>>,
    },
}

#[allow(clippy::type_complexity)]
impl<F, S, R, TE> Future for DnsExchangeConnectInner<F, S, R, TE>
where
    F: Future<Output = Result<S, ProtoError>> + 'static + Send + Unpin,
    S: DnsRequestSender<DnsResponseFuture = R> + 'static + Send + Unpin,
    R: Future<Output = Result<DnsResponse, ProtoError>> + 'static + Send + Unpin,
    TE: Time + Unpin,
{
    type Output = Result<(DnsExchange<R>, DnsExchangeBackground<S, R, TE>), ProtoError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        loop {
            let next;
            match *self {
                DnsExchangeConnectInner::Connecting {
                    ref mut connect_future,
                    ref mut outbound_messages,
                    ref mut sender,
                } => {
                    let connect_future = Pin::new(connect_future);
                    match connect_future.poll(cx) {
                        Poll::Ready(Ok(stream)) => {
                            //debug!("connection established: {}", stream);

                            let (exchange, background) = DnsExchange::from_stream_with_receiver(
                                stream,
                                outbound_messages
                                    .take()
                                    .expect("cannot poll after complete"),
                                sender.take().expect("cannot poll after complete"),
                            );

                            next = DnsExchangeConnectInner::Connected {
                                exchange,
                                background: Some(background),
                            };
                        }
                        Poll::Pending => return Poll::Pending,
                        Poll::Ready(Err(error)) => {
                            debug!("stream errored while connecting: {:?}", error);
                            next = DnsExchangeConnectInner::FailAll {
                                error,
                                outbound_messages: outbound_messages
                                    .take()
                                    .expect("cannot poll after complete"),
                            }
                        }
                    };
                }
                DnsExchangeConnectInner::Connected {
                    ref exchange,
                    ref mut background,
                } => {
                    let exchange = exchange.clone();
                    let background = background.take().expect("cannot poll after complete");

                    return Poll::Ready(Ok((exchange, background)));
                }
                DnsExchangeConnectInner::FailAll {
                    ref error,
                    ref mut outbound_messages,
                } => {
                    while let Some(outbound_message) = match outbound_messages.poll_next_unpin(cx) {
                        Poll::Ready(opt) => opt,
                        Poll::Pending => return Poll::Pending,
                    } {
                        let response = S::error_response::<TE>(error.clone());
                        // ignoring errors... best effort send...
                        outbound_message.unwrap().1.send_response(response).ok();
                    }

                    return Poll::Ready(Err(error.clone()));
                }
            }

            *self = next;
        }
    }
}