From 35ffb61530571618ae601c0b598dd18aa2a76790 Mon Sep 17 00:00:00 2001
From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Fri, 3 May 2024 15:47:54 +0200
Subject: [PATCH] dnsdist: Log the correct amount of bytes sent for DoH w/
 nghttp2

---
 pdns/dnsdistdist/dnsdist-nghttp2-in.cc   |  3 ++-
 pdns/dnsdistdist/dnsdist-tcp-upstream.hh |  2 +-
 pdns/dnsdistdist/dnsdist-tcp.cc          | 10 +++++-----
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/pdns/dnsdistdist/dnsdist-nghttp2-in.cc b/pdns/dnsdistdist/dnsdist-nghttp2-in.cc
index f88294168e..b4e3b4dd3b 100644
--- a/pdns/dnsdistdist/dnsdist-nghttp2-in.cc
+++ b/pdns/dnsdistdist/dnsdist-nghttp2-in.cc
@@ -568,8 +568,9 @@ IOState IncomingHTTP2Connection::sendResponse(const struct timeval& now, TCPResp
     responseBuffer = std::move(response.d_buffer);
   }
 
+  auto sent = responseBuffer.size();
   sendResponse(response.d_idstate.d_streamID, context, statusCode, d_ci.cs->dohFrontend->d_customResponseHeaders, contentType, sendContentType);
-  handleResponseSent(response);
+  handleResponseSent(response, sent);
 
   return hasPendingWrite() ? IOState::NeedWrite : IOState::Done;
 }
diff --git a/pdns/dnsdistdist/dnsdist-tcp-upstream.hh b/pdns/dnsdistdist/dnsdist-tcp-upstream.hh
index 40ca9335a3..7dfcd06282 100644
--- a/pdns/dnsdistdist/dnsdist-tcp-upstream.hh
+++ b/pdns/dnsdistdist/dnsdist-tcp-upstream.hh
@@ -138,7 +138,7 @@ public:
   void handleXFRResponse(const struct timeval& now, TCPResponse&& response) override;
 
   virtual IOState sendResponse(const struct timeval& now, TCPResponse&& response);
-  void handleResponseSent(TCPResponse& currentResponse);
+  void handleResponseSent(TCPResponse& currentResponse, size_t sentBytes);
   virtual IOState handleHandshake(const struct timeval& now);
   void handleHandshakeDone(const struct timeval& now);
   ProxyProtocolResult handleProxyProtocolPayload();
diff --git a/pdns/dnsdistdist/dnsdist-tcp.cc b/pdns/dnsdistdist/dnsdist-tcp.cc
index 27b2a116e0..68d31bf06a 100644
--- a/pdns/dnsdistdist/dnsdist-tcp.cc
+++ b/pdns/dnsdistdist/dnsdist-tcp.cc
@@ -196,7 +196,7 @@ static IOState sendQueuedResponses(std::shared_ptr<IncomingTCPConnectionState>&
   return IOState::Done;
 }
 
-void IncomingTCPConnectionState::handleResponseSent(TCPResponse& currentResponse)
+void IncomingTCPConnectionState::handleResponseSent(TCPResponse& currentResponse, size_t sentBytes)
 {
   if (currentResponse.d_idstate.qtype == QType::AXFR || currentResponse.d_idstate.qtype == QType::IXFR) {
     return;
@@ -208,13 +208,13 @@ void IncomingTCPConnectionState::handleResponseSent(TCPResponse& currentResponse
   if (!currentResponse.d_idstate.selfGenerated && backend) {
     const auto& ids = currentResponse.d_idstate;
     double udiff = ids.queryRealTime.udiff();
-    vinfolog("Got answer from %s, relayed to %s (%s, %d bytes), took %f us", backend->d_config.remote.toStringWithPort(), ids.origRemote.toStringWithPort(), getProtocol().toString(), currentResponse.d_buffer.size(), udiff);
+    vinfolog("Got answer from %s, relayed to %s (%s, %d bytes), took %f us", backend->d_config.remote.toStringWithPort(), ids.origRemote.toStringWithPort(), getProtocol().toString(), sentBytes, udiff);
 
     auto backendProtocol = backend->getProtocol();
     if (backendProtocol == dnsdist::Protocol::DoUDP && !currentResponse.d_idstate.forwardedOverUDP) {
       backendProtocol = dnsdist::Protocol::DoTCP;
     }
-    ::handleResponseSent(ids, udiff, d_ci.remote, backend->d_config.remote, static_cast<unsigned int>(currentResponse.d_buffer.size()), currentResponse.d_cleartextDH, backendProtocol, true);
+    ::handleResponseSent(ids, udiff, d_ci.remote, backend->d_config.remote, static_cast<unsigned int>(sentBytes), currentResponse.d_cleartextDH, backendProtocol, true);
   }
   else {
     const auto& ids = currentResponse.d_idstate;
@@ -316,7 +316,7 @@ IOState IncomingTCPConnectionState::sendResponse(const struct timeval& now, TCPR
     auto iostate = d_handler.tryWrite(d_currentResponse.d_buffer, d_currentPos, d_currentResponse.d_buffer.size());
     if (iostate == IOState::Done) {
       DEBUGLOG("response sent from " << __PRETTY_FUNCTION__);
-      handleResponseSent(d_currentResponse);
+      handleResponseSent(d_currentResponse, d_currentResponse.d_buffer.size());
       return iostate;
     }
     d_lastIOBlocked = true;
@@ -1121,7 +1121,7 @@ void IncomingTCPConnectionState::handleIO()
         iostate = d_handler.tryWrite(d_currentResponse.d_buffer, d_currentPos, d_currentResponse.d_buffer.size());
         if (iostate == IOState::Done) {
           DEBUGLOG("response sent from " << __PRETTY_FUNCTION__);
-          handleResponseSent(d_currentResponse);
+          handleResponseSent(d_currentResponse, d_currentResponse.d_buffer.size());
           d_state = State::idle;
         }
         else {
-- 
2.47.2