domain 0.12.0

//! RFC 8495 TSIG message authentication middleware.
//!
//! This module provides a TSIG request validation and response signing
//! middleware service. The underlying TSIG RR processing is implemented using
//! the [`rdata::tsig`][crate::rdata::tsig] module.
//!
//! Signed requests that fail signature verification will be rejected.
//!
//! Unsigned requests and correctly signed requests will pass through this
//! middleware unchanged.
//!
//! For requests which were correctly signed the corresponding response(s)
//! will be signed using the same key as the request.
//!
//! # Determining the key that a request was signed with
//!
//! The key that signed a request is output by this middleware via the request
//! metadata in the form [`Option<KS::Key>`], where `KS` denotes the type of
//! [`KeyStore`] that was used to construct this middleware. Upstream services
//! can choose to ignore the metadata by being generic over any kind of
//! metadata, or may offer a [`Service`] impl that specifically accepts the
//! `Option<KS::Key>` metadata type, enabling the upstream service to use
//! the request metadata to determine the key that the request was signed
//! with.
//!
//! # Limitations
//!
//! * RFC 8945 5.2.3 Time Check and Error Handling states: _"The server SHOULD
//!   also cache the most recent Time Signed value in a message generated by a
//!   key and SHOULD return BADTIME if a message received later has an earlier
//!   Time Signed value."_. This is not implemented.

use core::convert::Infallible;
use core::future::{ready, Ready};
use core::marker::PhantomData;
use core::ops::ControlFlow;

use std::vec::Vec;

use octseq::{Octets, OctetsFrom};
use tracing::{error, trace, warn};

use crate::base::iana::Rcode;
use crate::base::message_builder::AdditionalBuilder;
use crate::base::wire::Composer;
use crate::base::{Message, StreamTarget};
use crate::net::server::message::Request;
use crate::net::server::service::{
    CallResult, Service, ServiceError, ServiceFeedback, ServiceResult,
};
use crate::net::server::util::mk_builder_for_target;
use crate::rdata::tsig::Time48;
use crate::tsig::{self, KeyStore, ServerSequence, ServerTransaction};

use super::stream::{MiddlewareStream, PostprocessingStream};
use futures_util::stream::{once, Once};
use futures_util::Stream;

//------------ TsigMiddlewareSvc ----------------------------------------------

/// RFC 8495 TSIG message authentication middleware.
///
/// This middleware service validates TSIG signatures on incoming requests, if
/// any, and adds TSIG signatures to responses to signed requests.
///
/// Upstream services can detect whether a request is signed and with which
/// key by consuming the `Option<KS::Key>` metadata output by this service.
#[derive(Clone, Debug)]
pub struct TsigMiddlewareSvc<RequestOctets, NextSvc, KS, IgnoredRequestMeta>
where
    Infallible: From<<RequestOctets as octseq::OctetsFrom<Vec<u8>>>::Error>,
    KS: Clone + KeyStore,
    KS::Key: Clone,
    NextSvc: Service<RequestOctets, Option<KS::Key>>,
    NextSvc::Target: Composer + Default,
    RequestOctets: Octets + OctetsFrom<Vec<u8>> + Send + Sync + Unpin,
{
    next_svc: NextSvc,

    key_store: KS,

    _phantom: PhantomData<(RequestOctets, IgnoredRequestMeta)>,
}

impl<RequestOctets, NextSvc, KS, IgnoredRequestMeta>
    TsigMiddlewareSvc<RequestOctets, NextSvc, KS, IgnoredRequestMeta>
where
    IgnoredRequestMeta: Default + Clone + Send + Sync + Unpin + 'static,
    Infallible: From<<RequestOctets as octseq::OctetsFrom<Vec<u8>>>::Error>,
    KS: Clone + KeyStore + Unpin + Send + Sync + 'static,
    KS::Key: Clone + Unpin + Send + Sync,
    NextSvc: Service<RequestOctets, Option<KS::Key>>,
    NextSvc::Future: Unpin,
    RequestOctets:
        Octets + OctetsFrom<Vec<u8>> + Send + Sync + 'static + Unpin + Clone,
{
    /// Creates an instance of this middleware service.
    ///
    /// Keys in the provided [`KeyStore`] will be used to verify received signed
    /// requests and to sign the corresponding responses.
    #[must_use]
    pub fn new(next_svc: NextSvc, key_store: KS) -> Self {
        Self {
            next_svc,
            key_store,
            _phantom: PhantomData,
        }
    }
}

impl<RequestOctets, NextSvc, KS, IgnoredRequestMeta>
    TsigMiddlewareSvc<RequestOctets, NextSvc, KS, IgnoredRequestMeta>
where
    IgnoredRequestMeta: Default + Clone + Send + Sync + Unpin + 'static,
    Infallible: From<<RequestOctets as octseq::OctetsFrom<Vec<u8>>>::Error>,
    KS: Clone + KeyStore + Unpin + Send + Sync + 'static,
    KS::Key: Clone + Unpin + Send + Sync,
    NextSvc: Service<RequestOctets, Option<KS::Key>>,
    NextSvc::Future: Unpin,
    RequestOctets:
        Octets + OctetsFrom<Vec<u8>> + Send + Sync + 'static + Unpin + Clone,
{
    #[allow(clippy::type_complexity)]
    fn preprocess(
        req: &Request<RequestOctets, IgnoredRequestMeta>,
        key_store: &KS,
    ) -> Result<
        ControlFlow<
            AdditionalBuilder<StreamTarget<NextSvc::Target>>,
            Option<(
                Request<RequestOctets, Option<KS::Key>>,
                TsigSigner<KS::Key>,
            )>,
        >,
        ServiceError,
    > {
        let octets = req.message().as_slice().to_vec();
        let mut mut_msg = Message::from_octets(octets).unwrap();

        match tsig::ServerTransaction::request(
            key_store,
            &mut mut_msg,
            Time48::now(),
        ) {
            Ok(None) => {
                // Message is not TSIG signed.
            }

            Ok(Some(tsig)) => {
                // Message is TSIG signed by a known key.
                trace!(
                    "Request is signed with TSIG key '{}'",
                    tsig.key().name()
                );

                // Convert to RequestOctets so that the non-TSIG signed
                // message case can just pass through the RequestOctets.
                let source = mut_msg.into_octets();
                let octets = RequestOctets::octets_from(source);
                let new_msg = Message::from_octets(octets).unwrap();

                let mut new_req = Request::new(
                    req.client_addr(),
                    req.received_at(),
                    new_msg,
                    req.transport_ctx().clone(),
                    Some(tsig.wrapped_key().clone()),
                );

                let num_bytes_to_reserve = tsig.key().compose_len();
                new_req.reserve_bytes(num_bytes_to_reserve);

                return Ok(ControlFlow::Continue(Some((
                    new_req,
                    TsigSigner::Transaction(tsig),
                ))));
            }

            Err(err) => {
                // Message is incorrectly signed or signed with an unknown key.
                warn!(
                    "{} from {} refused: {err}",
                    req.message().header().opcode(),
                    req.client_addr(),
                );

                let builder = mk_builder_for_target();

                let res = match err.build_message(req.message(), builder) {
                    Ok(additional) => Ok(ControlFlow::Break(additional)),
                    Err(err) => {
                        error!("Unable to build TSIG error response: {err}");
                        Err(ServiceError::InternalError)
                    }
                };

                return res;
            }
        }

        Ok(ControlFlow::Continue(None))
    }

    /// Sign the given response, or if necessary construct and return an
    /// alternate response.
    fn postprocess(
        request: &Request<RequestOctets, IgnoredRequestMeta>,
        response: &mut AdditionalBuilder<StreamTarget<NextSvc::Target>>,
        state: &mut PostprocessingState<KS::Key>,
    ) -> Result<
        Option<AdditionalBuilder<StreamTarget<NextSvc::Target>>>,
        ServiceError,
    > {
        // Remove the limit we should have imposed during pre-processing so
        // that we can use the space we reserved for the OPT RR.
        response.clear_push_limit();

        let truncation_ctx;

        let res = match &mut state.signer {
            Some(TsigSigner::Transaction(_)) => {
                // Extract the single response signer and consume it in the
                // signing process.
                let Some(TsigSigner::Transaction(signer)) =
                    state.signer.take()
                else {
                    unreachable!()
                };

                trace!(
                    "Signing single response with TSIG key '{}'",
                    signer.key().name()
                );

                // We have to clone the key here in case the signer produces
                // an error, otherwise we lose access to the key as the signer
                // is consumed by calling answer(). The caller has control
                // over the key type via KS::Key so if cloning cost is a
                // problem the caller can choose to wrap the key in an Arc or
                // such to reduce the cloning cost.
                truncation_ctx = TruncationContext::NoSignerOnlyTheKey(
                    signer.key().clone(),
                );

                signer.answer(response, Time48::now())
            }

            Some(TsigSigner::Sequence(ref mut signer)) => {
                // Use the multi-response signer to sign the response.
                trace!(
                    "Signing response stream with TSIG key '{}'",
                    signer.key().name()
                );

                let res = signer.answer(response, Time48::now());

                truncation_ctx = TruncationContext::HaveSigner(signer);

                res
            }

            None => {
                // Nothing to do as unsigned requests don't require response
                // signing.
                return Ok(None);
            }
        };

        // Handle signing failure due to push error, i.e. there wasn't enough
        // space in the response to add the TSIG RR. This shouldn't happen
        // because we reserve space in preprocess() for the TSIG RR that we
        // add when signing.
        if res.is_err() {
            // 5.3. Generation of TSIG on Answers
            //   "If addition of the TSIG record will cause the message to be
            //   truncated, the server MUST alter the response so that a TSIG
            //   can be included. This response contains only the question and
            //   a TSIG record, has the TC bit set, and has an RCODE of 0
            //   (NOERROR). At this point, the client SHOULD retry the request
            //   using TCP (as per Section 4.2.2 of [RFC1035])."
            Ok(Some(Self::mk_signed_truncated_response(
                request,
                truncation_ctx,
            )?))
        } else {
            Ok(None)
        }
    }

    fn mk_signed_truncated_response(
        request: &Request<RequestOctets, IgnoredRequestMeta>,
        truncation_ctx: TruncationContext<'_, KS::Key, tsig::Key>,
    ) -> Result<AdditionalBuilder<StreamTarget<NextSvc::Target>>, ServiceError>
    {
        let builder = mk_builder_for_target();
        let mut new_response = builder
            .start_answer(request.message(), Rcode::NOERROR)
            .unwrap();
        new_response.header_mut().set_tc(true);
        let mut additional = new_response.additional();

        match truncation_ctx {
            TruncationContext::HaveSigner(signer) => {
                if let Err(err) =
                    signer.answer(&mut additional, Time48::now())
                {
                    error!("Unable to sign truncated TSIG response: {err}");
                    Err(ServiceError::InternalError)
                } else {
                    Ok(additional)
                }
            }

            TruncationContext::NoSignerOnlyTheKey(key) => {
                // We can't use the TSIG signer state we just had as that was
                // consumed in the failed attempt to sign the answer, so we
                // have to create a new TSIG state in order to sign the
                // truncated response.
                let octets = request.message().as_slice().to_vec();
                let mut mut_msg = Message::from_octets(octets).unwrap();

                match ServerTransaction::request(
                    &key,
                    &mut mut_msg,
                    Time48::now(),
                ) {
                    Ok(None) => {
                        error!("Unable to create signer for truncated TSIG response: internal error: request is not signed but was expected to be");
                        Err(ServiceError::InternalError)
                    }

                    Err(err) => {
                        error!("Unable to create signer for truncated TSIG response: {err}");
                        Err(ServiceError::InternalError)
                    }

                    Ok(Some(signer)) => {
                        if let Err(err) =
                            signer.answer(&mut additional, Time48::now())
                        {
                            error!("Unable to sign truncated TSIG response: {err}");
                            Err(ServiceError::InternalError)
                        } else {
                            Ok(additional)
                        }
                    }
                }
            }
        }
    }

    fn map_stream_item(
        request: Request<RequestOctets, IgnoredRequestMeta>,
        stream_item: ServiceResult<NextSvc::Target>,
        pp_config: &mut PostprocessingState<KS::Key>,
    ) -> ServiceResult<NextSvc::Target> {
        if let Ok(mut call_res) = stream_item {
            if matches!(
                call_res.feedback(),
                Some(ServiceFeedback::BeginTransaction)
            ) {
                // Does it need converting from the variant that supports
                // single messages only (ServerTransaction) to the variant
                // that supports signing multiple messages (ServerSequence)?
                // Note: Confusingly BeginTransaction and ServerTransaction
                // use the term "transaction" to mean completely the opposite
                // of each other. With BeginTransaction we mean that the
                // caller should expect a sequence of response messages
                // instead of the usual single response message. With
                // ServerTransaction the TSIG code means handling of single
                // messages only and NOT sequences for which there is a
                // separate ServerSequence type. Sigh.
                if let Some(TsigSigner::Transaction(tsig_txn)) =
                    pp_config.signer.take()
                {
                    // Do the conversion and store the result for future
                    // invocations of this function for subsequent items
                    // in the response stream.
                    pp_config.signer = Some(TsigSigner::Sequence(
                        ServerSequence::from(tsig_txn),
                    ));
                }
            }

            if let Some(response) = call_res.response_mut() {
                if let Some(new_response) =
                    Self::postprocess(&request, response, pp_config)?
                {
                    *response = new_response;
                }
            }

            Ok(call_res)
        } else {
            stream_item
        }
    }
}

//--- Service

/// This [`Service`] implementation specifies that the upstream service will
/// be passed metadata of type [`Option<KS::Key>`]. The upstream service can
/// optionally use this to learn which TSIG key signed the request.
///
/// This service does not accept downstream metadata, explicitly restricting
/// what it accepts to `()`. This is because (a) the service should be the
/// first layer above the network server, or as near as possible, such that it
/// receives unmodified requests and that the responses it generates are sent
/// over the network without prior modification, and thus it is not very
/// likely that the is a downstream layer that has metadata to supply to us,
/// and (b) because this service does not propagate the metadata it receives
/// from downstream but instead outputs [`Option<KS::Key>`] metadata to
/// upstream services.
impl<RequestOctets, NextSvc, KS, IgnoredRequestMeta>
    Service<RequestOctets, IgnoredRequestMeta>
    for TsigMiddlewareSvc<RequestOctets, NextSvc, KS, IgnoredRequestMeta>
where
    IgnoredRequestMeta: Default + Clone + Send + Sync + Unpin + 'static,
    Infallible: From<<RequestOctets as octseq::OctetsFrom<Vec<u8>>>::Error>,
    KS: Clone + KeyStore + Unpin + Send + Sync + 'static,
    KS::Key: Clone + Unpin + Send + Sync,
    NextSvc: Service<RequestOctets, Option<KS::Key>>,
    NextSvc::Future: Unpin,
    RequestOctets:
        Octets + OctetsFrom<Vec<u8>> + Send + Sync + 'static + Unpin + Clone,
{
    type Target = NextSvc::Target;
    type Stream = MiddlewareStream<
        NextSvc::Future,
        NextSvc::Stream,
        PostprocessingStream<
            RequestOctets,
            NextSvc::Future,
            NextSvc::Stream,
            IgnoredRequestMeta,
            PostprocessingState<KS::Key>,
        >,
        Once<Ready<ServiceResult<Self::Target>>>,
        <NextSvc::Stream as Stream>::Item,
    >;
    type Future = Ready<Self::Stream>;

    fn call(
        &self,
        request: Request<RequestOctets, IgnoredRequestMeta>,
    ) -> Self::Future {
        match Self::preprocess(&request, &self.key_store) {
            Ok(ControlFlow::Continue(Some((modified_req, signer)))) => {
                let pp_config = PostprocessingState::new(signer);

                let svc_call_fut = self.next_svc.call(modified_req);

                let map = PostprocessingStream::new(
                    svc_call_fut,
                    request,
                    pp_config,
                    Self::map_stream_item,
                );

                ready(MiddlewareStream::Map(map))
            }

            Ok(ControlFlow::Continue(None)) => {
                let request = request.with_new_metadata(None);
                let svc_call_fut = self.next_svc.call(request);
                ready(MiddlewareStream::IdentityFuture(svc_call_fut))
            }

            Ok(ControlFlow::Break(additional)) => {
                ready(MiddlewareStream::Result(once(ready(Ok(
                    CallResult::new(additional),
                )))))
            }

            Err(err) => {
                ready(MiddlewareStream::Result(once(ready(Err(err)))))
            }
        }
    }
}

/// Data needed to do signing during response post-processing.
pub struct PostprocessingState<K> {
    /// The signer used to verify the request.
    ///
    /// Needed to sign responses.
    ///
    /// We store it as an Option because ServerTransaction::answer() consumes
    /// the signer so have to first take it out of this struct, as a reference
    /// is held to the struct so it iself cannot be consumed.
    signer: Option<TsigSigner<K>>,
}

impl<K> PostprocessingState<K> {
    fn new(signer: TsigSigner<K>) -> Self {
        Self {
            signer: Some(signer),
        }
    }
}

/// A wrapper around [`ServerTransaction`] and [`ServerSequence`].
///
/// This wrapper allows us to write calling code once that invokes methods on
/// the TSIG signer/validator which have the same name and purpose for single
/// response vs multiple response streams, yet have distinct Rust types and so
/// must be called on the correct type, without needing to know at the call
/// site which of the distinct types it actually is.
#[derive(Clone, Debug)]
enum TsigSigner<K> {
    /// A [`ServerTransaction`] for signing a single response.
    Transaction(ServerTransaction<K>),

    /// A [`ServerSequence`] for signing multiple responses.
    Sequence(ServerSequence<K>),
}

//------------ TruncationContext ----------------------------------------------

enum TruncationContext<'a, KSeq, KTxn> {
    HaveSigner(&'a mut tsig::ServerSequence<KSeq>),

    NoSignerOnlyTheKey(KTxn),
}