Add missing pieces needed to load data files accepted by ldns-testns as input.

src/net/server/invoker.rs

@@ -0,0 +1,188 @@
+/// Common service invoking logic for network servers.
+///
+/// Used by [`stream::Connection`][net::server::stream::Connection] and
+/// [`dgram::Dgram`][net::server::dgram::Dgram].
+use core::clone::Clone;
+use core::default::Default;
+use core::future::Future;
+use core::pin::Pin;
+use core::time::Duration;
+use std::boxed::Box;
+
+use futures_util::StreamExt;
+use octseq::Octets;
+use tracing::trace;
+
+use crate::base::message_builder::AdditionalBuilder;
+use crate::base::wire::Composer;
+use crate::base::{Message, StreamTarget};
+
+use super::message::Request;
+use super::service::{Service, ServiceFeedback, ServiceResult};
+use super::util::mk_error_response;
+
+//------------ InvokerStatus --------------------------------------------------
+
+/// The current status of the service invoker.
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum InvokerStatus {
+    /// Processing independent responses.
+    Normal,
+
+    /// Processing related responses.
+    InTransaction,
+
+    /// No more responses to the current request will be processed.
+    Aborting,
+}
+
+//------------ ServiceInvoker -------------------------------------------------
+
+/// Dispatch requests to a [`Service`] and do common response processing.
+///
+/// Response streams will be split into individual responses and passed to the
+/// trait implementer for writing back to the network.
+///
+/// If the [`Service`] impl returns a [`ServiceError`] a corresponding DNS
+/// error response will be created and no further responses from the service
+/// for the current request will be processed and the service response stream
+/// will be dropped.
+///
+/// Also handles [`ServiceFeedback`] by invoking fn impls on the trait
+/// implementing type.
+pub trait ServiceInvoker<RequestOctets, Svc, EnqueueMeta>
+where
+    Svc: Service<RequestOctets> + Send + Sync + 'static,
+    Svc::Target: Composer + Default,
+    RequestOctets: Octets + Send + Sync + 'static,
+    EnqueueMeta: Send + Sync + 'static,
+{
+    /// Dispatch a request and process the responses.
+    ///
+    /// Dispatches the given request to the given [`Service`] impl and
+    /// processes the stream of resulting responses, passing them to the trait
+    /// impl'd [`enqueue_response`] function with the provided metadata for
+    /// writing back to the network. until no more responses exist or the
+    /// trait impl'd [`status`] function reports that the state is
+    /// [`InvokerStatus::Aborting`].
+    ///
+    /// On [`ServiceFeedback::Reconfigure`] passes the new configuration data
+    /// to the trait impl'd [`reconfugure`] function.
+    fn dispatch(
+        &mut self,
+        request: Request<RequestOctets>,
+        svc: Svc,
+        enqueue_meta: EnqueueMeta,
+    ) -> Pin<Box<dyn Future<Output = ()> + Send + '_>>
+    where
+        Self: Send + Sync,
+        Svc::Target: Send,
+        Svc::Stream: Send,
+        Svc::Future: Send,
+    {
+        Box::pin(async move {
+            let req_msg = request.message().clone();
+            let request_id = request.message().header().id();
+
+            // Dispatch the request to the service for processing.
+            trace!("Calling service for request id {request_id}");
+            let mut stream = svc.call(request).await;
+
+            // Handle the resulting stream of responses, most likely just one as
+            // only XFR requests potentially result in multiple responses.
+            trace!(
+                "Awaiting service call results for request id {request_id}"
+            );
+            while let Some(item) = stream.next().await {
+                trace!(
+                    "Processing service call result for request id {request_id}"
+                );
+
+                let response =
+                    self.process_response_stream_item(item, &req_msg);
+
+                if let Some(response) = response {
+                    self.enqueue_response(response, &enqueue_meta).await;
+                }
+
+                if matches!(self.status(), InvokerStatus::Aborting) {
+                    trace!("Aborting response stream processing for request id {request_id}");
+                    break;
+                }
+            }
+            trace!("Finished processing service call results for request id {request_id}");
+        })
+    }
+
+    /// Processing a single response stream item.
+    ///
+    /// Calls [`process_feedback`] if necessary. Extracts any response for
+    /// further processing by the caller.
+    ///
+    /// On [`ServiceError`] calls the trait impl'd [`set_status`] function
+    /// with `InvokerStatus::Aborting` and returns a generated error response
+    /// instead of the response from the service.
+    fn process_response_stream_item(
+        &mut self,
+        stream_item: ServiceResult<Svc::Target>,
+        req_msg: &Message<RequestOctets>,
+    ) -> Option<AdditionalBuilder<StreamTarget<Svc::Target>>> {
+        match stream_item {
+            Ok(call_result) => {
+                let (response, feedback) = call_result.into_inner();
+                if let Some(feedback) = feedback {
+                    self.process_feedback(feedback);
+                }
+                response
+            }
+
+            Err(err) => {
+                self.set_status(InvokerStatus::Aborting);
+                Some(mk_error_response(req_msg, err.rcode().into()))
+            }
+        }
+    }
+
+    //// Acts on [`ServiceFeedback`] received from the [`Service`].
+    ///
+    /// Calls the trait impl'd [`reconfigure`] on
+    /// [`ServiceFeedback::Reconfigure`].
+    ///
+    /// Calls the trait impl'd [`set_status`] on
+    /// [`ServiceFeedback::BeginTransaction`] with
+    /// [`InvokerStatus::InTransaction`].
+    ///
+    /// Calls the trait impl'd [`set_status`] on
+    /// [`ServiceFeedback::EndTransaction`] with [`InvokerStatus::Normal`].
+    fn process_feedback(&mut self, feedback: ServiceFeedback) {
+        match feedback {
+            ServiceFeedback::Reconfigure { idle_timeout } => {
+                self.reconfigure(idle_timeout);
+            }
+
+            ServiceFeedback::BeginTransaction => {
+                self.set_status(InvokerStatus::InTransaction);
+            }
+
+            ServiceFeedback::EndTransaction => {
+                self.set_status(InvokerStatus::Normal);
+            }
+        }
+    }
+
+    /// Returns the current status of the service invoker.
+    fn status(&self) -> InvokerStatus;
+
+    /// Sets the status of the service invoker to the given status.
+    fn set_status(&mut self, status: InvokerStatus);
+
+    /// Reconfigures the network server with new settings.
+    fn reconfigure(&self, idle_timeout: Option<Duration>);
+
+    /// Enqueues a response for writing back to the client.
+    fn enqueue_response<'a>(
+        &'a self,
+        response: AdditionalBuilder<StreamTarget<Svc::Target>>,
+        meta: &'a EnqueueMeta,
+    ) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
+}

src/net/server/mod.rs

@@ -278,6 +278,7 @@ pub mod batcher;
 pub mod buf;
 pub mod dgram;
 pub mod error;
+pub mod invoker;
 pub mod message;
 pub mod metrics;
 pub mod middleware;

src/stelline/client.rs

@@ -1,4 +1,5 @@
 #![allow(clippy::type_complexity)]
+use core::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
 use core::ops::Deref;
 
 use std::boxed::Box;
@@ -14,6 +15,7 @@ use std::vec::Vec;
 use bytes::Bytes;
 #[cfg(all(feature = "std", test))]
 use mock_instant::thread_local::MockClock;
+use tokio::time::Instant;
 use tracing::{debug, info_span, trace};
 use tracing_subscriber::EnvFilter;
 
@@ -25,6 +27,9 @@ use crate::net::client::request::{
     GetResponseMulti, RequestMessage, RequestMessageMulti, SendRequest,
     SendRequestMulti,
 };
+use crate::net::server::message::{
+    Request, TransportSpecificContext, UdpTransportContext,
+};
 use crate::stelline::matches::match_multi_msg;
 use crate::zonefile::inplace::Entry::Record;
 
@@ -127,6 +132,10 @@ pub async fn do_client_simple<R: SendRequest<RequestMessage<Vec<u8>>>>(
         request: R,
     ) -> Result<(), StellineErrorCause> {
         let mut resp: Option<Message<Bytes>> = None;
+        let mock_client_addr =
+            SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0));
+        let mock_transport_ctx =
+            TransportSpecificContext::Udp(UdpTransportContext::new(None));
 
         // Assume steps are in order. Maybe we need to define that.
         for step in &stelline.scenario.steps {
@@ -156,7 +165,18 @@ pub async fn do_client_simple<R: SendRequest<RequestMessage<Vec<u8>>>>(
                         .entry
                         .as_ref()
                         .ok_or(StellineErrorCause::MissingStepEntry)?;
-                    if !match_msg(entry, &answer, true) {
+                    let client_addr = entry
+                        .client_addr
+                        .map(|ip| SocketAddr::new(ip, 0))
+                        .unwrap_or(mock_client_addr);
+                    let req = Request::new(
+                        client_addr,
+                        Instant::now(),
+                        answer,
+                        mock_transport_ctx.clone(),
+                        (),
+                    );
+                    if !match_msg(entry, &req, true) {
                         return Err(StellineErrorCause::MismatchedAnswer);
                     }
                 }
@@ -459,6 +479,11 @@ pub async fn do_client<'a, T: ClientFactory>(
             MockClock::set_system_time(Duration::ZERO);
         }
 
+        let mock_client_addr =
+            SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0));
+        let mock_transport_ctx =
+            TransportSpecificContext::Udp(UdpTransportContext::new(None));
+
         // Assume steps are in order. Maybe we need to define that.
         for step in &stelline.scenario.steps {
             let span =
@@ -501,6 +526,11 @@ pub async fn do_client<'a, T: ClientFactory>(
                         return Err(StellineErrorCause::MissingResponse);
                     };
 
+                    let client_addr = entry
+                        .client_addr
+                        .map(|ip| SocketAddr::new(ip, 0))
+                        .unwrap_or(mock_client_addr);
+
                     if entry
                         .matches
                         .as_ref()
@@ -563,12 +593,19 @@ pub async fn do_client<'a, T: ClientFactory>(
 
                             trace!("Received answer.");
                             trace!(?resp);
+                            let req = Request::new(
+                                client_addr,
+                                Instant::now(),
+                                resp,
+                                mock_transport_ctx.clone(),
+                                (),
+                            );
 
                             let mut out_entry = Some(vec![]);
                             match_multi_msg(
                                 &entry,
                                 0,
-                                &resp,
+                                &req,
                                 true,
                                 &mut out_entry,
                             );
@@ -642,8 +679,15 @@ pub async fn do_client<'a, T: ClientFactory>(
 
                             trace!("Received answer.");
                             trace!(?resp);
+                            let req = Request::new(
+                                client_addr,
+                                Instant::now(),
+                                resp,
+                                mock_transport_ctx.clone(),
+                                (),
+                            );
                             if !match_multi_msg(
-                                entry, idx, &resp, true, &mut None,
+                                entry, idx, &req, true, &mut None,
                             ) {
                                 return Err(
                                     StellineErrorCause::MismatchedAnswer,

src/stelline/connection.rs

@@ -1,17 +1,23 @@
+use core::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
+
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::Waker;
 use std::task::{Context, Poll};
 use std::vec::Vec;
 
-use super::client::CurrStepValue;
-use super::parse_stelline::Stelline;
-use super::server::do_server;
-
 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
+use tokio::time::Instant;
 
 use crate::base::message_builder::AdditionalBuilder;
 use crate::base::Message;
+use crate::net::server::message::{
+    NonUdpTransportContext, Request, TransportSpecificContext,
+};
+
+use super::client::CurrStepValue;
+use super::parse_stelline::Stelline;
+use super::server::do_server;
 
 #[derive(Debug)]
 pub struct Connection {
@@ -20,7 +26,6 @@ pub struct Connection {
     waker: Option<Waker>,
     reply: Option<AdditionalBuilder<Vec<u8>>>,
     send_body: bool,
-
     tmpbuf: Vec<u8>,
 }
 
@@ -85,10 +90,23 @@ impl AsyncWrite for Connection {
         }
         let msg = Message::from_octets(self.tmpbuf[2..].to_vec()).unwrap();
         self.tmpbuf = Vec::new();
-        let opt_reply = do_server(&msg, &self.stelline, &self.step_value);
-        if opt_reply.is_some() {
+
+        let mock_client_addr =
+            SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0));
+        let mock_transport_ctx = TransportSpecificContext::NonUdp(
+            NonUdpTransportContext::new(None),
+        );
+        let req = Request::new(
+            mock_client_addr,
+            Instant::now(),
+            msg,
+            mock_transport_ctx.clone(),
+            (),
+        );
+
+        if let Some((opt_reply, _indices)) = do_server(&req, &self.stelline, &self.step_value) {
             // Do we need to support more than one reply?
-            self.reply = opt_reply;
+            self.reply = Some(opt_reply);
             let opt_waker = self.waker.take();
             if let Some(waker) = opt_waker {
                 waker.wake();