hickory_server/server/
response_handler.rs

1// Copyright 2015-2021 Benjamin Fry <benjaminfry@me.com>
2//
3// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
4// https://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
5// https://opensource.org/licenses/MIT>, at your option. This file may not be
6// copied, modified, or distributed except according to those terms.
7
8use std::{io, net::SocketAddr};
9
10use hickory_proto::rr::Record;
11use tracing::{debug, trace};
12
13use crate::server::Protocol;
14use crate::{
15    authority::MessageResponse,
16    proto::{
17        serialize::binary::BinEncoder, xfer::SerialMessage, BufDnsStreamHandle, DnsStreamHandle,
18    },
19    server::ResponseInfo,
20};
21
22/// A handler for send a response to a client
23#[async_trait::async_trait]
24pub trait ResponseHandler: Clone + Send + Sync + Unpin + 'static {
25    // TODO: add associated error type
26    //type Error;
27
28    /// Serializes and sends a message to to the wrapped handle
29    ///
30    /// self is consumed as only one message should ever be sent in response to a Request
31    async fn send_response<'a>(
32        &mut self,
33        response: MessageResponse<
34            '_,
35            'a,
36            impl Iterator<Item = &'a Record> + Send + 'a,
37            impl Iterator<Item = &'a Record> + Send + 'a,
38            impl Iterator<Item = &'a Record> + Send + 'a,
39            impl Iterator<Item = &'a Record> + Send + 'a,
40        >,
41    ) -> io::Result<ResponseInfo>;
42}
43
44/// A handler for wrapping a BufStreamHandle, which will properly serialize the message and add the
45///  associated destination.
46#[derive(Clone)]
47pub struct ResponseHandle {
48    dst: SocketAddr,
49    stream_handle: BufDnsStreamHandle,
50    protocol: Protocol,
51}
52
53impl ResponseHandle {
54    /// Returns a new `ResponseHandle` for sending a response message
55    pub fn new(dst: SocketAddr, stream_handle: BufDnsStreamHandle, protocol: Protocol) -> Self {
56        Self {
57            dst,
58            stream_handle,
59            protocol,
60        }
61    }
62
63    /// Selects an appropriate maximum serialized size for the given response.
64    fn max_size_for_response<'a>(
65        &self,
66        response: &MessageResponse<
67            '_,
68            'a,
69            impl Iterator<Item = &'a Record> + Send + 'a,
70            impl Iterator<Item = &'a Record> + Send + 'a,
71            impl Iterator<Item = &'a Record> + Send + 'a,
72            impl Iterator<Item = &'a Record> + Send + 'a,
73        >,
74    ) -> u16 {
75        match self.protocol {
76            Protocol::Udp => {
77                // Use EDNS, if available.
78                if let Some(edns) = response.get_edns() {
79                    edns.max_payload()
80                } else {
81                    // No EDNS, use the recommended max from RFC6891.
82                    hickory_proto::udp::MAX_RECEIVE_BUFFER_SIZE as u16
83                }
84            }
85            _ => u16::MAX,
86        }
87    }
88}
89
90#[async_trait::async_trait]
91impl ResponseHandler for ResponseHandle {
92    /// Serializes and sends a message to to the wrapped handle
93    ///
94    /// self is consumed as only one message should ever be sent in response to a Request
95    async fn send_response<'a>(
96        &mut self,
97        response: MessageResponse<
98            '_,
99            'a,
100            impl Iterator<Item = &'a Record> + Send + 'a,
101            impl Iterator<Item = &'a Record> + Send + 'a,
102            impl Iterator<Item = &'a Record> + Send + 'a,
103            impl Iterator<Item = &'a Record> + Send + 'a,
104        >,
105    ) -> io::Result<ResponseInfo> {
106        debug!(
107            "response: {} response_code: {}",
108            response.header().id(),
109            response.header().response_code(),
110        );
111        let mut buffer = Vec::with_capacity(512);
112        let encode_result = {
113            let mut encoder = BinEncoder::new(&mut buffer);
114
115            // Set an appropriate maximum on the encoder.
116            let max_size = self.max_size_for_response(&response);
117            trace!(
118                "setting response max size: {max_size} for protocol: {:?}",
119                self.protocol
120            );
121            encoder.set_max_size(max_size);
122
123            response.destructive_emit(&mut encoder)
124        };
125
126        let info = encode_result.map_err(|e| {
127            io::Error::new(io::ErrorKind::Other, format!("error encoding message: {e}"))
128        })?;
129
130        self.stream_handle
131            .send(SerialMessage::new(buffer, self.dst))
132            .map_err(|_| io::Error::new(io::ErrorKind::Other, "unknown"))?;
133
134        Ok(info)
135    }
136}