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

mod tcp;
mod udp;

use std::{io, net::SocketAddr};

use bytes::Bytes;
use tokio::net::{lookup_host, ToSocketAddrs};
use tokio_stream::wrappers::ReceiverStream;
use tracing_core::subscriber::NoSubscriber;
use tracing_futures::WithSubscriber;

pub use tcp::*;
pub use udp::*;

/// A sequence of [errors](std::io::Error) which occurred during a connection attempt.
///
/// These are paired with a [`SocketAddr`] because the connection attempt might fail multiple times
/// during DNS resolution.
#[derive(Debug)]
pub struct ConnectionErrors(pub Vec<(SocketAddr, io::Error)>);

/// Provides an interface for connecting (and reconnecting) to Graylog. Without an established
/// connection logs will not be sent. Messages logged without an established connection will sit in
/// the buffer until they can be drained.
#[derive(Debug)]
#[must_use]
pub struct ConnectionHandle<A, Conn> {
    pub(crate) addr: A,
    pub(crate) receiver: ReceiverStream<Bytes>,
    pub(crate) conn: Conn,
}

impl<A, Conn> ConnectionHandle<A, Conn> {
    /// Returns the connection address.
    pub fn address(&self) -> &A {
        &self.addr
    }
}

impl<A> ConnectionHandle<A, TcpConnection>
where
    A: ToSocketAddrs,
{
    /// Connects to Graylog via TCP using the address provided.
    ///
    /// This will perform DNS resolution and attempt to connect to each [`SocketAddr`] provided.
    pub async fn connect(&mut self) -> ConnectionErrors {
        // Do a DNS lookup if `addr` is a hostname
        let addrs = lookup_host(&self.addr).await.into_iter().flatten();

        // Loop through the IP addresses that the hostname resolved to
        let mut errors = Vec::new();
        for addr in addrs {
            let fut = self
                .conn
                .handle(addr, &mut self.receiver)
                .with_subscriber(NoSubscriber::default());
            if let Err(err) = fut.await {
                errors.push((addr, err));
            }
        }
        ConnectionErrors(errors)
    }
}

#[cfg(feature = "rustls-tls")]
impl<A> ConnectionHandle<A, TlsConnection>
where
    A: ToSocketAddrs,
{
    /// Connects to Graylog via TLS using the address provided.
    ///
    /// This will perform DNS resolution and attempt to connect to each [`SocketAddr`] provided.
    pub async fn connect(&mut self) -> ConnectionErrors {
        // Do a DNS lookup if `addr` is a hostname
        let addrs = lookup_host(&self.addr).await.into_iter().flatten();

        // Loop through the IP addresses that the hostname resolved to
        let mut errors = Vec::new();
        for addr in addrs {
            let fut = self
                .conn
                .handle(addr, &mut self.receiver)
                .with_subscriber(NoSubscriber::default());
            if let Err(err) = fut.await {
                errors.push((addr, err));
            }
        }
        ConnectionErrors(errors)
    }
}

impl<A> ConnectionHandle<A, UdpConnection>
where
    A: ToSocketAddrs,
{
    /// Connects to Graylog via UDP using the address provided.
    ///
    /// This will perform DNS resolution and attempt to connect to each [`SocketAddr`] provided.
    pub async fn connect(&mut self) -> ConnectionErrors {
        // Do a DNS lookup if `addr` is a hostname
        let addrs = lookup_host(&self.addr).await.into_iter().flatten();

        // Loop through the IP addresses that the hostname resolved to
        let mut errors = Vec::new();
        for addr in addrs {
            let fut = self
                .conn
                .handle(addr, &mut self.receiver)
                .with_subscriber(NoSubscriber::default());
            if let Err(err) = fut.await {
                errors.push((addr, err));
            }
        }
        ConnectionErrors(errors)
    }
}