From 34ef25565748482a43eb1dc553231b5415a0d721 Mon Sep 17 00:00:00 2001
From: Will Song <will.song@trailofbits.com>
Date: Thu, 18 Apr 2024 14:53:04 -0400
Subject: [PATCH] add documentation around the DhKem trait and the crate itself

---
 dhkem/src/lib.rs | 26 +++++++++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/dhkem/src/lib.rs b/dhkem/src/lib.rs
index 97b0c54..2d50572 100644
--- a/dhkem/src/lib.rs
+++ b/dhkem/src/lib.rs
@@ -1,3 +1,14 @@
+//! # Diffie-Hellman (DH) based Key Encapsulation Mechanisms (KEM)
+//!
+//! This crate provides a KEM interface for DH protocols as specified in
+//! [RFC9180](https://datatracker.ietf.org/doc/html/rfc9180#name-dh-based-kem-dhkem)
+//! without the shared secret extraction process. In particular, `Encaps(pk)` in the
+//! RFC returns the encapsulated key and an extracted shared secret, while our
+//! implementation leaves the extraction process up to the user. This type of KEM
+//! construction is currently being used in HPKE, as per the RFC, and in the current
+//! draft of the [TLS KEM
+//! combiner](https://datatracker.ietf.org/doc/html/draft-ietf-tls-hybrid-design-10).
+
 use kem::{Decapsulate, Encapsulate};
 use rand_core::CryptoRngCore;
 
@@ -11,23 +22,36 @@ pub trait SecretBytes {
     fn as_slice(&self) -> &[u8];
 }
 
+/// This is a trait that all KEM models should implement, and should probably be
+/// promoted to the kem crate itself. It specifies the types of encapsulating and
+/// decapsulating keys created by key generation, the shared secret type, and the
+/// encapsulated key type
 pub trait DhKem {
+    /// The type that will implement [`Decapsulate`]
     type DecapsulatingKey: Decapsulate<Self::EncapsulatedKey, Self::SharedSecret>;
+
+    /// The type that will implement [`Encapsulate`]
     type EncapsulatingKey: Encapsulate<Self::EncapsulatedKey, Self::SharedSecret>;
+
+    /// The type of the encapsulated key
     type EncapsulatedKey;
+
     #[cfg(not(test))]
+    /// The type of the shared secret
     type SharedSecret;
 
     #[cfg(test)]
     type SharedSecret: SecretBytes;
 
+    /// Generates a new (decapsulating key, encapsulating key) keypair for the KEM
+    /// model
     fn random_keypair(
         rng: &mut impl CryptoRngCore,
     ) -> (Self::DecapsulatingKey, Self::EncapsulatingKey);
 }
 
 #[cfg(feature = "arithmetic")]
-mod arithmetic;
+pub mod arithmetic;
 
 #[cfg(feature = "x25519")]
 mod x25519_kem;