feat: expose record count config (#786)

* expose record config

* update default record counts

* make fields optional

* override record count in integration test

Author: sinui0

crates/common/src/config.rs

@@ -7,6 +7,10 @@ use std::error::Error;
 
 // Default is 32 bytes to decrypt the TLS protocol messages.
 const DEFAULT_MAX_RECV_ONLINE: usize = 32;
+// Default maximum number of TLS records to allow.
+//
+// This would allow for up to 50Mb upload from prover to verifier.
+const DEFAULT_RECORDS_LIMIT: usize = 256;
 
 // Current version that is running.
 static VERSION: Lazy<Version> = Lazy::new(|| {
@@ -21,12 +25,18 @@ static VERSION: Lazy<Version> = Lazy::new(|| {
 pub struct ProtocolConfig {
     /// Maximum number of bytes that can be sent.
     max_sent_data: usize,
+    /// Maximum number of application data records that can be sent.
+    #[builder(setter(strip_option), default)]
+    max_sent_records: Option<usize>,
     /// Maximum number of bytes that can be decrypted online, i.e. while the
     /// MPC-TLS connection is active.
     #[builder(default = "DEFAULT_MAX_RECV_ONLINE")]
     max_recv_data_online: usize,
     /// Maximum number of bytes that can be received.
     max_recv_data: usize,
+    /// Maximum number of application data records that can be received.
+    #[builder(setter(strip_option), default)]
+    max_recv_records: Option<usize>,
     /// Version that is being run by prover/verifier.
     #[builder(setter(skip), default = "VERSION.clone()")]
     version: Version,
@@ -54,6 +64,12 @@ impl ProtocolConfig {
         self.max_sent_data
     }
 
+    /// Returns the maximum number of application data records that can
+    /// be sent.
+    pub fn max_sent_records(&self) -> Option<usize> {
+        self.max_sent_records
+    }
+
     /// Returns the maximum number of bytes that can be decrypted online.
     pub fn max_recv_data_online(&self) -> usize {
         self.max_recv_data_online
@@ -63,6 +79,12 @@ impl ProtocolConfig {
     pub fn max_recv_data(&self) -> usize {
         self.max_recv_data
     }
+
+    /// Returns the maximum number of application data records that can
+    /// be received.
+    pub fn max_recv_records(&self) -> Option<usize> {
+        self.max_recv_records
+    }
 }
 
 /// Protocol configuration validator used by checker (i.e. verifier) to perform
@@ -71,8 +93,14 @@ impl ProtocolConfig {
 pub struct ProtocolConfigValidator {
     /// Maximum number of bytes that can be sent.
     max_sent_data: usize,
+    /// Maximum number of application data records that can be sent.
+    #[builder(default = "DEFAULT_RECORDS_LIMIT")]
+    max_sent_records: usize,
     /// Maximum number of bytes that can be received.
     max_recv_data: usize,
+    /// Maximum number of application data records that can be received.
+    #[builder(default = "DEFAULT_RECORDS_LIMIT")]
+    max_recv_records: usize,
     /// Version that is being run by checker.
     #[builder(setter(skip), default = "VERSION.clone()")]
     version: Version,
@@ -89,15 +117,28 @@ impl ProtocolConfigValidator {
         self.max_sent_data
     }
 
+    /// Returns the maximum number of application data records that can
+    /// be sent.
+    pub fn max_sent_records(&self) -> usize {
+        self.max_sent_records
+    }
+
     /// Returns the maximum number of bytes that can be received.
     pub fn max_recv_data(&self) -> usize {
         self.max_recv_data
     }
 
+    /// Returns the maximum number of application data records that can
+    /// be received.
+    pub fn max_recv_records(&self) -> usize {
+        self.max_recv_records
+    }
+
     /// Performs compatibility check of the protocol configuration between
     /// prover and verifier.
     pub fn validate(&self, config: &ProtocolConfig) -> Result<(), ProtocolConfigError> {
         self.check_max_transcript_size(config.max_sent_data, config.max_recv_data)?;
+        self.check_max_records(config.max_sent_records, config.max_recv_records)?;
         self.check_version(&config.version)?;
         Ok(())
     }
@@ -125,6 +166,32 @@ impl ProtocolConfigValidator {
         Ok(())
     }
 
+    fn check_max_records(
+        &self,
+        max_sent_records: Option<usize>,
+        max_recv_records: Option<usize>,
+    ) -> Result<(), ProtocolConfigError> {
+        if let Some(max_sent_records) = max_sent_records {
+            if max_sent_records > self.max_sent_records {
+                return Err(ProtocolConfigError::max_record_count(format!(
+                    "max_sent_records {} is greater than the configured limit {}",
+                    max_sent_records, self.max_sent_records,
+                )));
+            }
+        }
+
+        if let Some(max_recv_records) = max_recv_records {
+            if max_recv_records > self.max_recv_records {
+                return Err(ProtocolConfigError::max_record_count(format!(
+                    "max_recv_records {} is greater than the configured limit {}",
+                    max_recv_records, self.max_recv_records,
+                )));
+            }
+        }
+
+        Ok(())
+    }
+
     // Checks if both versions are the same (might support check for different but
     // compatible versions in the future).
     fn check_version(&self, peer_version: &Version) -> Result<(), ProtocolConfigError> {
@@ -165,6 +232,13 @@ impl ProtocolConfigError {
         }
     }
 
+    fn max_record_count(msg: impl Into<String>) -> Self {
+        Self {
+            kind: ErrorKind::MaxRecordCount,
+            source: Some(msg.into().into()),
+        }
+    }
+
     fn version(msg: impl Into<String>) -> Self {
         Self {
             kind: ErrorKind::Version,
@@ -176,7 +250,8 @@ impl ProtocolConfigError {
 impl fmt::Display for ProtocolConfigError {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self.kind {
-            ErrorKind::MaxTranscriptSize => write!(f, "max transcript size error")?,
+            ErrorKind::MaxTranscriptSize => write!(f, "max transcript size exceeded")?,
+            ErrorKind::MaxRecordCount => write!(f, "max record count exceeded")?,
             ErrorKind::Version => write!(f, "version error")?,
         }
 
@@ -191,6 +266,7 @@ impl fmt::Display for ProtocolConfigError {
 #[derive(Debug)]
 enum ErrorKind {
     MaxTranscriptSize,
+    MaxRecordCount,
     Version,
 }
 

crates/mpc-tls/src/config.rs

@@ -1,9 +1,34 @@
 use derive_builder::Builder;
 
-const MIN_SENT: usize = 32;
-const MIN_SENT_RECORDS: usize = 8;
-const MIN_RECV: usize = 32;
-const MIN_RECV_RECORDS: usize = 8;
+/// Number of TLS protocol bytes that will be sent.
+const PROTOCOL_DATA_SENT: usize = 32;
+/// Number of TLS protocol bytes that will be received.
+const PROTOCOL_DATA_RECV: usize = 32;
+
+/// Number of TLS protocol records that we need to allocate in addition
+/// to the application data records.
+const PROTOCOL_RECORD_COUNT_SENT: usize = 2;
+/// Number of TLS protocol records that we need to allocate in addition
+/// to the application data records.
+const PROTOCOL_RECORD_COUNT_RECV: usize = 2;
+
+/// Computes the record count configuration given the data volume.
+///
+/// Accurately estimating a good default is challenging as we do not
+/// know exactly how much data will be packed into each record in advance.
+fn default_record_count(max_data: usize) -> usize {
+    // We assume a minimum of 8 records for the first 4KB.
+    const MIN: usize = 8;
+
+    // Then we estimate that after 4KB of data is transmitted that they will
+    // average 4KB in size.
+    let remainder = max_data.saturating_sub(4096);
+    let count = remainder.div_ceil(4096);
+
+    // For example, if max_data=32KB then this will return 15. That will result
+    // in about 3MB upload from prover to verifier.
+    MIN + count
+}
 
 /// MPC-TLS configuration.
 #[derive(Debug, Clone, Builder)]
@@ -43,7 +68,7 @@ impl ConfigBuilder {
     /// Builds the configuration.
     pub fn build(&self) -> Result<Config, ConfigBuilderError> {
         let defer_decryption = self.defer_decryption.unwrap_or(true);
-        let max_sent = MIN_SENT
+        let max_sent = PROTOCOL_DATA_SENT
             + self
                 .max_sent
                 .ok_or(ConfigBuilderError::UninitializedField("max_sent"))?;
@@ -60,15 +85,15 @@ impl ConfigBuilder {
             ));
         }
 
-        max_recv_online += MIN_RECV;
-        max_recv += MIN_RECV;
+        max_recv_online += PROTOCOL_DATA_RECV;
+        max_recv += PROTOCOL_DATA_RECV;
 
         let max_sent_records = self
             .max_sent_records
-            .unwrap_or_else(|| MIN_SENT_RECORDS + max_sent.div_ceil(16384));
+            .unwrap_or_else(|| PROTOCOL_RECORD_COUNT_SENT + default_record_count(max_sent));
         let max_recv_records = self
             .max_recv_records
-            .unwrap_or_else(|| MIN_RECV_RECORDS + max_recv.div_ceil(16384));
+            .unwrap_or_else(|| PROTOCOL_RECORD_COUNT_RECV + default_record_count(max_recv));
 
         Ok(Config {
             defer_decryption,
@@ -80,3 +105,13 @@ impl ConfigBuilder {
         })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_default_record_count() {
+        assert_eq!(default_record_count(1 << 15), 15);
+    }
+}

crates/prover/src/config.rs

@@ -49,12 +49,22 @@ impl ProverConfig {
     }
 
     pub(crate) fn build_mpc_tls_config(&self) -> Config {
-        Config::builder()
+        let mut builder = Config::builder();
+
+        builder
             .defer_decryption(self.defer_decryption_from_start)
             .max_sent(self.protocol_config.max_sent_data())
             .max_recv_online(self.protocol_config.max_recv_data_online())
-            .max_recv(self.protocol_config.max_recv_data())
-            .build()
-            .unwrap()
+            .max_recv(self.protocol_config.max_recv_data());
+
+        if let Some(max_sent_records) = self.protocol_config.max_sent_records() {
+            builder.max_sent_records(max_sent_records);
+        }
+
+        if let Some(max_recv_records) = self.protocol_config.max_recv_records() {
+            builder.max_recv_records(max_recv_records);
+        }
+
+        builder.build().unwrap()
     }
 }

crates/tests-integration/tests/defer_decryption.rs

@@ -16,8 +16,12 @@ use tracing::instrument;
 
 // Maximum number of bytes that can be sent from prover to server
 const MAX_SENT_DATA: usize = 1 << 12;
+// Maximum number of application records sent from prover to server
+const MAX_SENT_RECORDS: usize = 4;
 // Maximum number of bytes that can be received by prover from server
 const MAX_RECV_DATA: usize = 1 << 14;
+// Maximum number of application records received by prover from server
+const MAX_RECV_RECORDS: usize = 6;
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[ignore]
@@ -51,7 +55,9 @@ async fn prover<T: AsyncWrite + AsyncRead + Send + Unpin + 'static>(notary_socke
             .protocol_config(
                 ProtocolConfig::builder()
                     .max_sent_data(MAX_SENT_DATA)
+                    .max_sent_records(MAX_SENT_RECORDS)
                     .max_recv_data(MAX_RECV_DATA)
+                    .max_recv_records(MAX_RECV_RECORDS)
                     .build()
                     .unwrap(),
             )

crates/verifier/src/config.rs

@@ -43,11 +43,21 @@ impl VerifierConfig {
     }
 
     pub(crate) fn build_mpc_tls_config(&self, protocol_config: &ProtocolConfig) -> Config {
-        Config::builder()
+        let mut builder = Config::builder();
+
+        builder
             .max_sent(protocol_config.max_sent_data())
             .max_recv_online(protocol_config.max_recv_data_online())
-            .max_recv(protocol_config.max_recv_data())
-            .build()
-            .unwrap()
+            .max_recv(protocol_config.max_recv_data());
+
+        if let Some(max_sent_records) = protocol_config.max_sent_records() {
+            builder.max_sent_records(max_sent_records);
+        }
+
+        if let Some(max_recv_records) = protocol_config.max_recv_records() {
+            builder.max_recv_records(max_recv_records);
+        }
+
+        builder.build().unwrap()
     }
 }