From 14b49a69dc46f30160e5100265153e98c09de6b1 Mon Sep 17 00:00:00 2001
From: David Venhoek <david@tweedegolf.com>
Date: Thu, 8 Aug 2024 17:20:50 +0200
Subject: [PATCH] Initial implementation of server side POW enforcement.

---
 rustls/src/server/hs.rs    |  7 +++++--
 rustls/src/server/tls13.rs | 30 ++++++++++++++++++++++++++++--
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/rustls/src/server/hs.rs b/rustls/src/server/hs.rs
index 9b86a161..b1a3e23a 100644
--- a/rustls/src/server/hs.rs
+++ b/rustls/src/server/hs.rs
@@ -23,8 +23,8 @@ use crate::msgs::enums::{Compression, ExtensionType, NamedGroup};
 #[cfg(feature = "tls12")]
 use crate::msgs::handshake::SessionId;
 use crate::msgs::handshake::{
-    ClientHelloPayload, ConvertProtocolNameList, ConvertServerNameList, HandshakePayload,
-    KeyExchangeAlgorithm, Random, ServerExtension,
+    ClientHelloPayload, ClientPuzzleChallenge, ConvertProtocolNameList, ConvertServerNameList,
+    HandshakePayload, KeyExchangeAlgorithm, Random, ServerExtension,
 };
 use crate::msgs::message::{Message, MessagePayload};
 use crate::msgs::persist;
@@ -213,6 +213,7 @@ pub(super) struct ExpectClientHello {
     #[cfg(feature = "tls12")]
     pub(super) using_ems: bool,
     pub(super) done_retry: bool,
+    pub(super) challenge: Option<ClientPuzzleChallenge>,
     pub(super) send_tickets: usize,
 }
 
@@ -233,6 +234,7 @@ impl ExpectClientHello {
             #[cfg(feature = "tls12")]
             using_ems: false,
             done_retry: false,
+            challenge: None,
             send_tickets: 0,
         }
     }
@@ -383,6 +385,7 @@ impl ExpectClientHello {
                 suite,
                 randoms,
                 done_retry: self.done_retry,
+                challenge: self.challenge,
                 send_tickets: self.send_tickets,
                 extra_exts: self.extra_exts,
             }
diff --git a/rustls/src/server/tls13.rs b/rustls/src/server/tls13.rs
index 27c65a0a..f01a4438 100644
--- a/rustls/src/server/tls13.rs
+++ b/rustls/src/server/tls13.rs
@@ -43,8 +43,8 @@ mod client_hello {
     use crate::msgs::enums::{Compression, NamedGroup, PSKKeyExchangeMode};
     use crate::msgs::handshake::{
         CertReqExtension, CertificatePayloadTls13, CertificateRequestPayloadTls13,
-        ClientHelloPayload, HelloRetryExtension, HelloRetryRequest, KeyShareEntry, Random,
-        ServerExtension, ServerHelloPayload, SessionId,
+        ClientHelloPayload, ClientPuzzleChallenge, ClientPuzzleExtension, HelloRetryExtension,
+        HelloRetryRequest, KeyShareEntry, Random, ServerExtension, ServerHelloPayload, SessionId,
     };
     use crate::server::common::ActiveCertifiedKey;
     use crate::sign;
@@ -66,6 +66,7 @@ mod client_hello {
         pub(in crate::server) suite: &'static Tls13CipherSuite,
         pub(in crate::server) randoms: ConnectionRandoms,
         pub(in crate::server) done_retry: bool,
+        pub(in crate::server) challenge: Option<ClientPuzzleChallenge>,
         pub(in crate::server) send_tickets: usize,
         pub(in crate::server) extra_exts: Vec<ServerExtension>,
     }
@@ -195,6 +196,19 @@ mod client_hello {
                 (share.group == selected_kxg.name()).then_some((share, selected_kxg))
             });
 
+            // Force hrr if the client hasn't solved a required puzzle yet
+            let chosen_share_and_kxg = if let (Some(challenge), Some(solution)) =
+                (&self.challenge, client_hello.puzzle_solution())
+            {
+                if challenge.check(&solution) {
+                    chosen_share_and_kxg
+                } else {
+                    None
+                }
+            } else {
+                None
+            };
+
             let chosen_share_and_kxg = match chosen_share_and_kxg {
                 Some(s) => s,
                 None => {
@@ -209,12 +223,16 @@ mod client_hello {
                         ));
                     }
 
+                    let challenge =
+                        ClientPuzzleChallenge::new_cookie(self.config.provider.secure_random)?;
+
                     emit_hello_retry_request(
                         &mut self.transcript,
                         self.suite,
                         client_hello.session_id,
                         cx.common,
                         selected_kxg.name(),
+                        Some(challenge.clone()),
                     );
                     emit_fake_ccs(cx.common);
 
@@ -228,6 +246,7 @@ mod client_hello {
                         #[cfg(feature = "tls12")]
                         using_ems: false,
                         done_retry: true,
+                        challenge: Some(challenge),
                         send_tickets: self.send_tickets,
                         extra_exts: self.extra_exts,
                     });
@@ -582,6 +601,7 @@ mod client_hello {
         session_id: SessionId,
         common: &mut CommonState,
         group: NamedGroup,
+        challenge: Option<ClientPuzzleChallenge>,
     ) {
         let mut req = HelloRetryRequest {
             legacy_version: ProtocolVersion::TLSv1_2,
@@ -596,6 +616,12 @@ mod client_hello {
             .push(HelloRetryExtension::SupportedVersions(
                 ProtocolVersion::TLSv1_3,
             ));
+        if let Some(challenge) = challenge {
+            req.extensions
+                .push(HelloRetryExtension::ClientPuzzle(
+                    ClientPuzzleExtension::from_challenge(challenge),
+                ));
+        }
 
         let m = Message {
             version: ProtocolVersion::TLSv1_2,