~nickbp/kapiti

ref: 9d460d0b7ab39ac1e39576b34fa6217049c65edb kapiti/src/listen_udp.rs -rw-r--r-- 2.2 KiB
9d460d0bNick Parker Update dependencies to latest along with corresponding code changes 6 months ago
                                                                                
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
use std::net::SocketAddr;

use anyhow::Result;
use async_net::UdpSocket;
use bytes::BytesMut;
use tracing::{self, trace, warn};

use crate::{lookup, runner};

/// Listens for UDP requests until an accept fails or the channel sender is closed.
pub async fn listen_udp(
    udp_recv: &mut UdpSocket,
    server_query_tx: &mut async_channel::Sender<runner::RequestMsg>,
) -> Result<()> {
    let mut packet_buffer = BytesMut::with_capacity(4096);
    loop {
        // Ensure that the buffer has a SIZE suitable for socket.recv_from().
        // If we just leave it with the CAPACITY then it drops data.
        packet_buffer.resize(packet_buffer.capacity(), 0);

        let (recvsize, request_source) = udp_recv.recv_from(&mut packet_buffer).await?;

        // Got a request from somewhere, shorten to actual size received (doesnt affect malloc)
        packet_buffer.truncate(recvsize);

        trace!(
            "Queueing raw UDP request from {:?} ({}b): {:02X?}",
            request_source,
            packet_buffer.len(),
            &packet_buffer[..]
        );
        let msg = runner::RequestMsg {
            src: request_source,
            data: runner::RequestData::Udp(packet_buffer.clone()),
        };

        // Forward request info to be processed and then replied
        server_query_tx.send(msg).await?;
    }
}

pub async fn handle_udp_request(
    lookup: &mut lookup::Lookup,
    request_source: SocketAddr,
    mut buf: BytesMut,
    udp_sock: &UdpSocket,
) {
    match lookup.handle_query(&mut buf).await {
        Ok(()) => {
            // Send the response back to the original UDP requestor.
            trace!("Raw response to {:?} ({}b): {:02X?}", request_source, buf.len(), &buf[..]);
            if let Err(e) = udp_sock.send_to(&mut buf, &request_source).await {
                warn!(
                    "Failed to send UDP DNS response to client={:?}: {:?}",
                    request_source, e
                );
            }
        }
        Err(ioerr) => {
            warn!(
                "Failed to handle request from client={:?}: {:02X?} ({:?})",
                request_source,
                &buf[..],
                ioerr
            );
        }
    }
}