dnstap 0.1.6

An implementation of the dnstap protocol
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

use crate::dns_message::*;
use framestream::EncoderWriter;
use mio::deprecated::{UnixSocket, UnixStream};
use mio::timer::Timeout;
use mio::*;
use protobuf::*;
use std::io::{self, BufWriter, Write};
use std::path::PathBuf;
use std::time;
use std::usize;

pub const BUFFER_SIZE: usize = 262_144;
pub const CONTENT_TYPE: &str = "protobuf:dnstap.Dnstap";
pub const RETRY_DELAY_SECS: u64 = 1;

pub const NOTIFY_TOK: Token = Token(usize::MAX - 1);
pub const TIMER_TOK: Token = Token(usize::MAX - 2);
pub const UNIX_SOCKET_TOK: Token = Token(usize::MAX - 3);

pub struct Context {
    pub mio_poll: Poll,
    pub mio_timers: timer::Timer<Token>,
    pub retry_timeout: Option<Timeout>,
    pub dnstap_rx: channel::Receiver<DNSMessage>,
    pub unix_socket_path: Option<PathBuf>,
    pub unix_stream: Option<UnixStream>,
    pub frame_stream: Option<EncoderWriter<BufWriter<UnixStream>>>,
}

impl Context {
    pub fn message_cb(&mut self) {
        if let Some(ref unix_stream) = self.unix_stream {
            self.mio_poll
                .reregister(
                    unix_stream,
                    UNIX_SOCKET_TOK,
                    Ready::writable(),
                    PollOpt::edge() | PollOpt::oneshot(),
                )
                .unwrap();
        }
    }

    pub fn write_cb(&mut self, event: Event) {
        if self.frame_stream.is_none() {
            debug_assert!(self.unix_stream.is_none());
            return;
        }
        if event.kind().is_hup() || event.kind().is_error() {
            self.unix_stream = None;
            self.frame_stream = None;
            self.retry_timeout
                .take()
                .and_then(|timeout| self.mio_timers.cancel_timeout(&timeout));
            self.retry_timeout = Some(
                self.mio_timers
                    .set_timeout(time::Duration::from_secs(RETRY_DELAY_SECS), TIMER_TOK)
                    .unwrap(),
            );
            return;
        }
        let frame_stream = self.frame_stream.as_mut().unwrap();
        while let Ok(dns_message) = self.dnstap_rx.try_recv() {
            let dns_message_bytes = dns_message.into_protobuf().write_to_bytes().unwrap();
            match frame_stream.write_all(&dns_message_bytes).or_else(|_| {
                let _ = frame_stream.flush();
                frame_stream.write_all(&dns_message_bytes)
            }) {
                Err(ref e)
                    if e.kind() == io::ErrorKind::WouldBlock
                        || e.kind() == io::ErrorKind::Interrupted =>
                {
                    self.mio_poll
                        .reregister(
                            self.unix_stream.as_ref().unwrap(),
                            UNIX_SOCKET_TOK,
                            Ready::writable(),
                            PollOpt::edge() | PollOpt::oneshot(),
                        )
                        .unwrap();
                    break;
                }
                Err(e) => {
                    let _ = frame_stream.flush();
                    panic!("Cannot write to the frame stream any more: {}", e)
                }
                _ => {}
            }
        }
        let _ = frame_stream.flush();
        self.mio_poll
            .reregister(
                &self.dnstap_rx,
                NOTIFY_TOK,
                Ready::readable(),
                PollOpt::edge() | PollOpt::oneshot(),
            )
            .unwrap();
    }

    pub fn connect(&mut self) {
        if self.frame_stream.is_some() {
            debug_assert!(self.unix_stream.is_some());
            return;
        }
        assert!(self.unix_socket_path.is_some());
        let unix_socket = UnixSocket::stream().unwrap();
        let unix_stream = match unix_socket.connect(&self.unix_socket_path.clone().unwrap()) {
            Ok((unix_stream, _connected)) => unix_stream,
            Err(_) => {
                self.retry_timeout
                    .take()
                    .and_then(|timeout| self.mio_timers.cancel_timeout(&timeout));
                self.retry_timeout = Some(
                    self.mio_timers
                        .set_timeout(time::Duration::from_secs(RETRY_DELAY_SECS), TIMER_TOK)
                        .unwrap(),
                );
                return;
            }
        };
        let frame_stream = EncoderWriter::new(
            BufWriter::with_capacity(BUFFER_SIZE, unix_stream.try_clone().unwrap()),
            Some(CONTENT_TYPE.to_owned()),
        );
        self.mio_poll
            .register(
                &unix_stream,
                UNIX_SOCKET_TOK,
                Ready::writable(),
                PollOpt::edge() | PollOpt::oneshot(),
            )
            .unwrap();
        self.unix_stream = Some(unix_stream);
        self.frame_stream = Some(frame_stream);
    }
}