Implement CONNECT-UDP toy server backend

2 new commandline flags:
* `connect_udp_proxy_targets` sets the allow list of allowed CONNECT-UDP proxy target servers. Must be set for toy server to allow CONNECT-UDP requests.
* `proxy_server_label` sets the server label for proxy error headers per RFC 9209. If not set, one is randomly generated.

PiperOrigin-RevId: 483394501
diff --git a/quiche/quic/tools/connect_server_backend.cc b/quiche/quic/tools/connect_server_backend.cc
index 646f575..cc46571 100644
--- a/quiche/quic/tools/connect_server_backend.cc
+++ b/quiche/quic/tools/connect_server_backend.cc
@@ -13,6 +13,7 @@
 #include "quiche/quic/core/quic_server_id.h"
 #include "quiche/quic/core/socket_factory.h"
 #include "quiche/quic/tools/connect_tunnel.h"
+#include "quiche/quic/tools/connect_udp_tunnel.h"
 #include "quiche/quic/tools/quic_simple_server_backend.h"
 #include "quiche/common/platform/api/quiche_bug_tracker.h"
 #include "quiche/common/platform/api/quiche_logging.h"
@@ -35,15 +36,23 @@
 
 ConnectServerBackend::ConnectServerBackend(
     std::unique_ptr<QuicSimpleServerBackend> non_connect_backend,
-    absl::flat_hash_set<QuicServerId> acceptable_destinations)
+    absl::flat_hash_set<QuicServerId> acceptable_connect_destinations,
+    absl::flat_hash_set<QuicServerId> acceptable_connect_udp_targets,
+    std::string server_label)
     : non_connect_backend_(std::move(non_connect_backend)),
-      acceptable_destinations_(std::move(acceptable_destinations)) {
+      acceptable_connect_destinations_(
+          std::move(acceptable_connect_destinations)),
+      acceptable_connect_udp_targets_(
+          std::move(acceptable_connect_udp_targets)),
+      server_label_(std::move(server_label)) {
   QUICHE_DCHECK(non_connect_backend_);
+  QUICHE_DCHECK(!server_label_.empty());
 }
 
 ConnectServerBackend::~ConnectServerBackend() {
   // Expect all streams to be closed before destroying backend.
-  QUICHE_DCHECK(tunnels_.empty());
+  QUICHE_DCHECK(connect_tunnels_.empty());
+  QUICHE_DCHECK(connect_udp_tunnels_.empty());
 }
 
 bool ConnectServerBackend::InitializeBackend(const std::string&) {
@@ -54,7 +63,8 @@
 
 void ConnectServerBackend::SetSocketFactory(SocketFactory* socket_factory) {
   QUICHE_DCHECK_NE(socket_factory_, socket_factory);
-  QUICHE_DCHECK(tunnels_.empty());
+  QUICHE_DCHECK(connect_tunnels_.empty());
+  QUICHE_DCHECK(connect_udp_tunnels_.empty());
   socket_factory_ = socket_factory;
 }
 
@@ -80,28 +90,42 @@
     return;
   }
 
-  if (request_headers.contains(":protocol")) {
-    // Anything other than normal CONNECT not supported.
-    // TODO(ericorth): Add CONNECT-UDP support.
+  if (!request_headers.contains(":protocol")) {
+    // normal CONNECT
+    auto [tunnel_it, inserted] = connect_tunnels_.emplace(
+        request_handler->stream_id(),
+        std::make_unique<ConnectTunnel>(request_handler, socket_factory_,
+                                        acceptable_connect_destinations_));
+    QUICHE_DCHECK(inserted);
+
+    tunnel_it->second->OpenTunnel(request_headers);
+  } else if (request_headers.find(":protocol")->second == "connect-udp") {
+    // CONNECT-UDP
+    auto [tunnel_it, inserted] = connect_udp_tunnels_.emplace(
+        request_handler->stream_id(),
+        std::make_unique<ConnectUdpTunnel>(request_handler, socket_factory_,
+                                           server_label_,
+                                           acceptable_connect_udp_targets_));
+    QUICHE_DCHECK(inserted);
+
+    tunnel_it->second->OpenTunnel(request_headers);
+  } else {
+    // Not a supported request.
     non_connect_backend_->HandleConnectHeaders(request_headers,
                                                request_handler);
-    return;
   }
-
-  auto [tunnel_it, inserted] = tunnels_.emplace(
-      request_handler->stream_id(),
-      std::make_unique<ConnectTunnel>(request_handler, socket_factory_,
-                                      acceptable_destinations_));
-  QUICHE_DCHECK(inserted);
-
-  tunnel_it->second->OpenTunnel(request_headers);
 }
 
 void ConnectServerBackend::HandleConnectData(absl::string_view data,
                                              bool data_complete,
                                              RequestHandler* request_handler) {
-  auto tunnel_it = tunnels_.find(request_handler->stream_id());
-  if (tunnel_it == tunnels_.end()) {
+  // Expect ConnectUdpTunnels to register a datagram visitor, causing the
+  // stream to process data as capsules.  HandleConnectData() should therefore
+  // never be called for streams with a ConnectUdpTunnel.
+  QUICHE_DCHECK(!connect_udp_tunnels_.contains(request_handler->stream_id()));
+
+  auto tunnel_it = connect_tunnels_.find(request_handler->stream_id());
+  if (tunnel_it == connect_tunnels_.end()) {
     // If tunnel not found, perhaps it's something being handled for
     // non-CONNECT. Possible because this method could be called for anything
     // with a ":method":"CONNECT" header, but this class does not handle such
@@ -116,16 +140,22 @@
   }
   if (data_complete) {
     tunnel_it->second->OnClientStreamClose();
-    tunnels_.erase(tunnel_it);
+    connect_tunnels_.erase(tunnel_it);
   }
 }
 
 void ConnectServerBackend::CloseBackendResponseStream(
     QuicSimpleServerBackend::RequestHandler* request_handler) {
-  auto tunnel_it = tunnels_.find(request_handler->stream_id());
-  if (tunnel_it != tunnels_.end()) {
+  auto tunnel_it = connect_tunnels_.find(request_handler->stream_id());
+  if (tunnel_it != connect_tunnels_.end()) {
     tunnel_it->second->OnClientStreamClose();
-    tunnels_.erase(tunnel_it);
+    connect_tunnels_.erase(tunnel_it);
+  }
+
+  auto udp_tunnel_it = connect_udp_tunnels_.find(request_handler->stream_id());
+  if (udp_tunnel_it != connect_udp_tunnels_.end()) {
+    udp_tunnel_it->second->OnClientStreamClose();
+    connect_udp_tunnels_.erase(udp_tunnel_it);
   }
 
   non_connect_backend_->CloseBackendResponseStream(request_handler);
diff --git a/quiche/quic/tools/connect_server_backend.h b/quiche/quic/tools/connect_server_backend.h
index 10eb7c5..c3dcab6 100644
--- a/quiche/quic/tools/connect_server_backend.h
+++ b/quiche/quic/tools/connect_server_backend.h
@@ -15,17 +15,23 @@
 #include "quiche/quic/core/quic_server_id.h"
 #include "quiche/quic/core/socket_factory.h"
 #include "quiche/quic/tools/connect_tunnel.h"
+#include "quiche/quic/tools/connect_udp_tunnel.h"
 #include "quiche/quic/tools/quic_simple_server_backend.h"
 
 namespace quic {
 
-// QUIC server backend that handles CONNECT requests. Non-CONNECT requests are
-// delegated to a separate backend.
+// QUIC server backend that handles CONNECT and CONNECT-UDP requests.
+// Non-CONNECT requests are delegated to a separate backend.
 class ConnectServerBackend : public QuicSimpleServerBackend {
  public:
+  // `server_label` is an identifier (typically randomly generated) to identify
+  // the server or backend in error headers, per the requirements of RFC 9209,
+  // Section 2.
   ConnectServerBackend(
       std::unique_ptr<QuicSimpleServerBackend> non_connect_backend,
-      absl::flat_hash_set<QuicServerId> acceptable_destinations);
+      absl::flat_hash_set<QuicServerId> acceptable_connect_destinations,
+      absl::flat_hash_set<QuicServerId> acceptable_connect_udp_targets,
+      std::string server_label);
 
   ConnectServerBackend(const ConnectServerBackend&) = delete;
   ConnectServerBackend& operator=(const ConnectServerBackend&) = delete;
@@ -48,10 +54,15 @@
 
  private:
   std::unique_ptr<QuicSimpleServerBackend> non_connect_backend_;
-  const absl::flat_hash_set<QuicServerId> acceptable_destinations_;
+  const absl::flat_hash_set<QuicServerId> acceptable_connect_destinations_;
+  const absl::flat_hash_set<QuicServerId> acceptable_connect_udp_targets_;
+  const std::string server_label_;
 
   SocketFactory* socket_factory_;  // unowned
-  absl::flat_hash_map<QuicStreamId, std::unique_ptr<ConnectTunnel>> tunnels_;
+  absl::flat_hash_map<QuicStreamId, std::unique_ptr<ConnectTunnel>>
+      connect_tunnels_;
+  absl::flat_hash_map<QuicStreamId, std::unique_ptr<ConnectUdpTunnel>>
+      connect_udp_tunnels_;
 };
 
 }  // namespace quic
diff --git a/quiche/quic/tools/connect_udp_tunnel.cc b/quiche/quic/tools/connect_udp_tunnel.cc
index ebeef8b..96612fc 100644
--- a/quiche/quic/tools/connect_udp_tunnel.cc
+++ b/quiche/quic/tools/connect_udp_tunnel.cc
@@ -167,14 +167,15 @@
 
 ConnectUdpTunnel::ConnectUdpTunnel(
     QuicSimpleServerBackend::RequestHandler* client_stream_request_handler,
-    SocketFactory* socket_factory, uint64_t server_label,
+    SocketFactory* socket_factory, std::string server_label,
     absl::flat_hash_set<QuicServerId> acceptable_targets)
     : acceptable_targets_(std::move(acceptable_targets)),
       socket_factory_(socket_factory),
-      server_label_(server_label),
+      server_label_(std::move(server_label)),
       client_stream_request_handler_(client_stream_request_handler) {
   QUICHE_DCHECK(client_stream_request_handler_);
   QUICHE_DCHECK(socket_factory_);
+  QUICHE_DCHECK(!server_label_.empty());
 }
 
 ConnectUdpTunnel::~ConnectUdpTunnel() {
@@ -385,8 +386,7 @@
   spdy::Http2HeaderBlock headers;
   headers[":status"] = status;
 
-  structured_headers::Item proxy_status_item(
-      absl::StrCat("QuicToyServer", server_label_));
+  structured_headers::Item proxy_status_item(server_label_);
   structured_headers::Item proxy_status_error_item(
       std::string{proxy_status_error});
   structured_headers::Item proxy_status_details_item(
diff --git a/quiche/quic/tools/connect_udp_tunnel.h b/quiche/quic/tools/connect_udp_tunnel.h
index f254b61..71819df 100644
--- a/quiche/quic/tools/connect_udp_tunnel.h
+++ b/quiche/quic/tools/connect_udp_tunnel.h
@@ -2,8 +2,8 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
-#ifndef QUICHE_QUIC_TOOLS_CONNECT_TUNNEL_H_
-#define QUICHE_QUIC_TOOLS_CONNECT_TUNNEL_H_
+#ifndef QUICHE_QUIC_TOOLS_CONNECT_UDP_TUNNEL_H_
+#define QUICHE_QUIC_TOOLS_CONNECT_UDP_TUNNEL_H_
 
 #include <cstdint>
 #include <memory>
@@ -36,7 +36,7 @@
   // per the requirements of RFC 9209, Section 2.
   ConnectUdpTunnel(
       QuicSimpleServerBackend::RequestHandler* client_stream_request_handler,
-      SocketFactory* socket_factory, uint64_t server_label,
+      SocketFactory* socket_factory, std::string server_label,
       absl::flat_hash_set<QuicServerId> acceptable_targets);
   ~ConnectUdpTunnel();
   ConnectUdpTunnel(const ConnectUdpTunnel&) = delete;
@@ -80,7 +80,7 @@
 
   const absl::flat_hash_set<QuicServerId> acceptable_targets_;
   SocketFactory* const socket_factory_;
-  const uint64_t server_label_;
+  const std::string server_label_;
 
   // Null when client stream closed.
   QuicSimpleServerBackend::RequestHandler* client_stream_request_handler_;
@@ -94,4 +94,4 @@
 
 }  // namespace quic
 
-#endif  // QUICHE_QUIC_TOOLS_CONNECT_TUNNEL_H_
+#endif  // QUICHE_QUIC_TOOLS_CONNECT_UDP_TUNNEL_H_
diff --git a/quiche/quic/tools/connect_udp_tunnel_test.cc b/quiche/quic/tools/connect_udp_tunnel_test.cc
index 23b1c11..9d7e88c 100644
--- a/quiche/quic/tools/connect_udp_tunnel_test.cc
+++ b/quiche/quic/tools/connect_udp_tunnel_test.cc
@@ -142,7 +142,7 @@
   ConnectUdpTunnel tunnel_{
       &request_handler_,
       &socket_factory_,
-      /*server_label=*/123,
+      "server_label",
       /*acceptable_targets=*/
       {{std::string(kAcceptableTarget), kAcceptablePort},
        {TestLoopback4().ToString(), kAcceptablePort},
diff --git a/quiche/quic/tools/quic_toy_server.cc b/quiche/quic/tools/quic_toy_server.cc
index 132ecce..3f23af6 100644
--- a/quiche/quic/tools/quic_toy_server.cc
+++ b/quiche/quic/tools/quic_toy_server.cc
@@ -4,10 +4,12 @@
 
 #include "quiche/quic/tools/quic_toy_server.h"
 
+#include <string>
 #include <utility>
 #include <vector>
 
 #include "absl/container/flat_hash_set.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/str_split.h"
 #include "absl/strings/string_view.h"
 #include "quiche/quic/core/quic_server_id.h"
@@ -18,6 +20,7 @@
 #include "quiche/quic/tools/quic_memory_cache_backend.h"
 #include "quiche/common/platform/api/quiche_command_line_flags.h"
 #include "quiche/common/platform/api/quiche_logging.h"
+#include "quiche/common/quiche_random.h"
 
 DEFINE_QUICHE_COMMAND_LINE_FLAG(int32_t, port, 6121,
                                 "The port the quic server will listen on.");
@@ -48,7 +51,20 @@
 DEFINE_QUICHE_COMMAND_LINE_FLAG(
     std::string, connect_proxy_destinations, "",
     "Specifies a comma-separated list of destinations (\"hostname:port\") to "
-    "which the quic server will allow tunneling via CONNECT.");
+    "which the QUIC server will allow tunneling via CONNECT.");
+
+DEFINE_QUICHE_COMMAND_LINE_FLAG(
+    std::string, connect_udp_proxy_targets, "",
+    "Specifies a comma-separated list of target servers (\"hostname:port\") to "
+    "which the QUIC server will allow tunneling via CONNECT-UDP.");
+
+DEFINE_QUICHE_COMMAND_LINE_FLAG(
+    std::string, proxy_server_label, "",
+    "Specifies an identifier to identify the server in proxy error headers, "
+    "per the requirements of RFC 9209, Section 2. It should uniquely identify "
+    "the running service between separate running instances of the QUIC toy "
+    "server binary. If not specified, one will be randomly generated as "
+    "\"QuicToyServerN\" where N is a random uint64_t.");
 
 namespace quic {
 
@@ -68,6 +84,8 @@
   }
 
   if (!quiche::GetQuicheCommandLineFlag(FLAGS_connect_proxy_destinations)
+           .empty() ||
+      !quiche::GetQuicheCommandLineFlag(FLAGS_connect_udp_proxy_targets)
            .empty()) {
     absl::flat_hash_set<QuicServerId> connect_proxy_destinations;
     for (absl::string_view destination : absl::StrSplit(
@@ -79,10 +97,31 @@
       connect_proxy_destinations.insert(
           std::move(destination_server_id).value());
     }
-    QUICHE_CHECK(!connect_proxy_destinations.empty());
+
+    absl::flat_hash_set<QuicServerId> connect_udp_proxy_targets;
+    for (absl::string_view target : absl::StrSplit(
+             quiche::GetQuicheCommandLineFlag(FLAGS_connect_udp_proxy_targets),
+             ',', absl::SkipEmpty())) {
+      absl::optional<QuicServerId> target_server_id =
+          QuicServerId::ParseFromHostPortString(target);
+      QUICHE_CHECK(target_server_id.has_value());
+      connect_udp_proxy_targets.insert(std::move(target_server_id).value());
+    }
+
+    QUICHE_CHECK(!connect_proxy_destinations.empty() ||
+                 !connect_udp_proxy_targets.empty());
+
+    std::string proxy_server_label =
+        quiche::GetQuicheCommandLineFlag(FLAGS_proxy_server_label);
+    if (proxy_server_label.empty()) {
+      proxy_server_label = absl::StrCat(
+          "QuicToyServer",
+          quiche::QuicheRandom::GetInstance()->InsecureRandUint64());
+    }
 
     return std::make_unique<ConnectServerBackend>(
-        std::move(memory_cache_backend), std::move(connect_proxy_destinations));
+        std::move(memory_cache_backend), std::move(connect_proxy_destinations),
+        std::move(connect_udp_proxy_targets), std::move(proxy_server_label));
   }
 
   return memory_cache_backend;