Move shared response processing code to ohttp.rs (payjoin#660)

spacebear21 · web-flow · commit d8c37c03fd6a · 2025-06-06T17:34:40.000-04:00
The `process_res` methods across send and receive have a lot of shared
code that can be extracted into reusable functions. `process_get_res`
and `process_post_res` essentially do the same thing, and could
conceivably be combined into one function. They were kept separate
because "202 Accepted" is not a valid response for POST requests.
diff --git a/payjoin/src/ohttp.rs b/payjoin/src/ohttp.rs
@@ -52,6 +52,77 @@ pub fn ohttp_encapsulate(
     Ok((buffer, ohttp_ctx))
 }
 
+#[derive(Debug)]
+pub enum DirectoryResponseError {
+    InvalidSize(usize),
+    OhttpDecapsulation(OhttpEncapsulationError),
+    UnexpectedStatusCode(http::StatusCode),
+}
+
+impl fmt::Display for DirectoryResponseError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        use DirectoryResponseError::*;
+
+        match self {
+            OhttpDecapsulation(e) => write!(f, "OHTTP Decapsulation Error: {e}"),
+            InvalidSize(size) => write!(
+                f,
+                "Unexpected response size {}, expected {} bytes",
+                size,
+                crate::directory::ENCAPSULATED_MESSAGE_BYTES
+            ),
+            UnexpectedStatusCode(status) => write!(f, "Unexpected status code: {status}"),
+        }
+    }
+}
+
+impl error::Error for DirectoryResponseError {
+    fn source(&self) -> Option<&(dyn error::Error + 'static)> {
+        use DirectoryResponseError::*;
+
+        match self {
+            OhttpDecapsulation(e) => Some(e),
+            InvalidSize(_) => None,
+            UnexpectedStatusCode(_) => None,
+        }
+    }
+}
+
+pub fn process_get_res(
+    res: &[u8],
+    ohttp_context: ohttp::ClientResponse,
+) -> Result<Option<Vec<u8>>, DirectoryResponseError> {
+    let response = process_ohttp_res(res, ohttp_context)?;
+    match response.status() {
+        http::StatusCode::OK => Ok(Some(response.body().to_vec())),
+        http::StatusCode::ACCEPTED => Ok(None),
+        status_code => Err(DirectoryResponseError::UnexpectedStatusCode(status_code)),
+    }
+}
+
+pub fn process_post_res(
+    res: &[u8],
+    ohttp_context: ohttp::ClientResponse,
+) -> Result<(), DirectoryResponseError> {
+    let response = process_ohttp_res(res, ohttp_context)?;
+    match response.status() {
+        http::StatusCode::OK => Ok(()),
+        status_code => Err(DirectoryResponseError::UnexpectedStatusCode(status_code)),
+    }
+}
+
+fn process_ohttp_res(
+    res: &[u8],
+    ohttp_context: ohttp::ClientResponse,
+) -> Result<http::Response<Vec<u8>>, DirectoryResponseError> {
+    let response_array: &[u8; crate::directory::ENCAPSULATED_MESSAGE_BYTES] =
+        res.try_into().map_err(|_| DirectoryResponseError::InvalidSize(res.len()))?;
+    log::trace!("decapsulating directory response");
+    let res = ohttp_decapsulate(ohttp_context, response_array)
+        .map_err(DirectoryResponseError::OhttpDecapsulation)?;
+    Ok(res)
+}
+
 /// decapsulate ohttp, bhttp response and return http response body and status code
 pub fn ohttp_decapsulate(
     res_ctx: ohttp::ClientResponse,
diff --git a/payjoin/src/receive/v2/error.rs b/payjoin/src/receive/v2/error.rs
@@ -3,7 +3,7 @@ use std::error;
 
 use super::Error::V2;
 use crate::hpke::HpkeError;
-use crate::ohttp::OhttpEncapsulationError;
+use crate::ohttp::{DirectoryResponseError, OhttpEncapsulationError};
 use crate::receive::error::Error;
 
 /// Error that may occur during a v2 session typestate change
@@ -31,10 +31,8 @@ pub(crate) enum InternalSessionError {
     OhttpEncapsulation(OhttpEncapsulationError),
     /// Hybrid Public Key Encryption failed
     Hpke(HpkeError),
-    /// Unexpected response size
-    UnexpectedResponseSize(usize),
-    /// Unexpected status code
-    UnexpectedStatusCode(http::StatusCode),
+    /// The directory returned a bad response
+    DirectoryResponse(DirectoryResponseError),
 }
 
 impl From<OhttpEncapsulationError> for Error {
@@ -56,13 +54,7 @@ impl fmt::Display for SessionError {
             Expired(expiry) => write!(f, "Session expired at {expiry:?}"),
             OhttpEncapsulation(e) => write!(f, "OHTTP Encapsulation Error: {e}"),
             Hpke(e) => write!(f, "Hpke decryption failed: {e}"),
-            UnexpectedResponseSize(size) => write!(
-                f,
-                "Unexpected response size {}, expected {} bytes",
-                size,
-                crate::directory::ENCAPSULATED_MESSAGE_BYTES
-            ),
-            UnexpectedStatusCode(status) => write!(f, "Unexpected status code: {status}"),
+            DirectoryResponse(e) => write!(f, "Directory response error: {e}"),
         }
     }
 }
@@ -76,8 +68,7 @@ impl error::Error for SessionError {
             Expired(_) => None,
             OhttpEncapsulation(e) => Some(e),
             Hpke(e) => Some(e),
-            UnexpectedResponseSize(_) => None,
-            UnexpectedStatusCode(_) => None,
+            DirectoryResponse(e) => Some(e),
         }
     }
 }
diff --git a/payjoin/src/receive/v2/mod.rs b/payjoin/src/receive/v2/mod.rs
@@ -17,7 +17,9 @@ use super::{
     v1, InternalPayloadError, JsonReply, OutputSubstitutionError, ReplyableError, SelectionError,
 };
 use crate::hpke::{decrypt_message_a, encrypt_message_b, HpkeKeyPair, HpkePublicKey};
-use crate::ohttp::{ohttp_decapsulate, ohttp_encapsulate, OhttpEncapsulationError, OhttpKeys};
+use crate::ohttp::{
+    ohttp_encapsulate, process_get_res, process_post_res, OhttpEncapsulationError, OhttpKeys,
+};
 use crate::output_substitution::OutputSubstitution;
 use crate::persist::Persister;
 use crate::receive::{parse_payload, InputPair};
@@ -180,23 +182,17 @@ impl Receiver<WithContext> {
         body: &[u8],
         context: ohttp::ClientResponse,
     ) -> Result<Option<Receiver<UncheckedProposal>>, Error> {
-        let response_array: &[u8; crate::directory::ENCAPSULATED_MESSAGE_BYTES] =
-            body.try_into()
-                .map_err(|_| InternalSessionError::UnexpectedResponseSize(body.len()))?;
-        log::trace!("decapsulating directory response");
-        let response = ohttp_decapsulate(context, response_array)
-            .map_err(InternalSessionError::OhttpEncapsulation)?;
-        if response.body().is_empty() {
-            log::debug!("response is empty");
-            return Ok(None);
-        }
-        match String::from_utf8(response.body().to_vec()) {
+        let body = match process_get_res(body, context)
+            .map_err(InternalSessionError::DirectoryResponse)?
+        {
+            Some(body) => body,
+            None => return Ok(None),
+        };
+        match String::from_utf8(body.clone()) {
             // V1 response bodies are utf8 plaintext
             Ok(response) => Ok(Some(Receiver { state: self.extract_proposal_from_v1(response)? })),
             // V2 response bodies are encrypted binary
-            Err(_) => Ok(Some(Receiver {
-                state: self.extract_proposal_from_v2(response.body().to_vec())?,
-            })),
+            Err(_) => Ok(Some(Receiver { state: self.extract_proposal_from_v2(body)? })),
         }
     }
 
@@ -350,16 +346,8 @@ impl Receiver<UncheckedProposal> {
         body: &[u8],
         context: ohttp::ClientResponse,
     ) -> Result<(), SessionError> {
-        let response_array: &[u8; crate::directory::ENCAPSULATED_MESSAGE_BYTES] =
-            body.try_into()
-                .map_err(|_| InternalSessionError::UnexpectedResponseSize(body.len()))?;
-        let response = ohttp_decapsulate(context, response_array)
-            .map_err(InternalSessionError::OhttpEncapsulation)?;
-
-        match response.status() {
-            http::StatusCode::OK => Ok(()),
-            _ => Err(InternalSessionError::UnexpectedStatusCode(response.status()).into()),
-        }
+        process_post_res(body, context)
+            .map_err(|e| InternalSessionError::DirectoryResponse(e).into())
     }
 }
 
@@ -648,15 +636,8 @@ impl Receiver<PayjoinProposal> {
         res: &[u8],
         ohttp_context: ohttp::ClientResponse,
     ) -> Result<(), Error> {
-        let response_array: &[u8; crate::directory::ENCAPSULATED_MESSAGE_BYTES] =
-            res.try_into().map_err(|_| InternalSessionError::UnexpectedResponseSize(res.len()))?;
-        let res = ohttp_decapsulate(ohttp_context, response_array)
-            .map_err(InternalSessionError::OhttpEncapsulation)?;
-        if res.status().is_success() {
-            Ok(())
-        } else {
-            Err(InternalSessionError::UnexpectedStatusCode(res.status()).into())
-        }
+        process_post_res(res, ohttp_context)
+            .map_err(|e| InternalSessionError::DirectoryResponse(e).into())
     }
 }
 
diff --git a/payjoin/src/send/v2/error.rs b/payjoin/src/send/v2/error.rs
@@ -1,5 +1,6 @@
 use core::fmt;
 
+use crate::ohttp::DirectoryResponseError;
 use crate::uri::url_ext::ParseReceiverPubkeyParamError;
 
 /// Error returned when request could not be created.
@@ -73,25 +74,19 @@ pub struct EncapsulationError(InternalEncapsulationError);
 
 #[derive(Debug)]
 pub(crate) enum InternalEncapsulationError {
-    /// The response size is not the expected size.
-    InvalidSize(usize),
-    /// The status code is not the expected status code.
-    UnexpectedStatusCode(http::StatusCode),
     /// The HPKE failed.
     Hpke(crate::hpke::HpkeError),
-    /// The encapsulation failed.
-    Ohttp(crate::ohttp::OhttpEncapsulationError),
+    /// The directory returned a bad response
+    DirectoryResponse(DirectoryResponseError),
 }
 
 impl fmt::Display for EncapsulationError {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         use InternalEncapsulationError::*;
 
         match &self.0 {
-            InvalidSize(size) => write!(f, "invalid size: {size}"),
-            UnexpectedStatusCode(status) => write!(f, "unexpected status code: {status}"),
-            Ohttp(error) => write!(f, "OHTTP encapsulation error: {error}"),
             Hpke(error) => write!(f, "HPKE error: {error}"),
+            DirectoryResponse(e) => write!(f, "Directory response error: {e}"),
         }
     }
 }
@@ -101,10 +96,8 @@ impl std::error::Error for EncapsulationError {
         use InternalEncapsulationError::*;
 
         match &self.0 {
-            InvalidSize(_) => None,
-            UnexpectedStatusCode(_) => None,
-            Ohttp(error) => Some(error),
             Hpke(error) => Some(error),
+            DirectoryResponse(e) => Some(e),
         }
     }
 }
@@ -115,8 +108,6 @@ impl From<InternalEncapsulationError> for EncapsulationError {
 
 impl From<InternalEncapsulationError> for super::ResponseError {
     fn from(value: InternalEncapsulationError) -> Self {
-        super::ResponseError::Validation(
-            super::InternalValidationError::V2Encapsulation(value.into()).into(),
-        )
+        super::InternalValidationError::V2Encapsulation(value.into()).into()
     }
 }
diff --git a/payjoin/src/send/v2/mod.rs b/payjoin/src/send/v2/mod.rs
@@ -32,7 +32,7 @@ use url::Url;
 use super::error::BuildSenderError;
 use super::*;
 use crate::hpke::{decrypt_message_b, encrypt_message_a, HpkeSecretKey};
-use crate::ohttp::{ohttp_decapsulate, ohttp_encapsulate};
+use crate::ohttp::{ohttp_encapsulate, process_get_res, process_post_res};
 use crate::persist::Persister;
 use crate::send::v1;
 use crate::uri::{ShortId, UrlExt};
@@ -325,24 +325,15 @@ impl Sender<V2PostContext> {
         self,
         response: &[u8],
     ) -> Result<Sender<V2GetContext>, EncapsulationError> {
-        let response_array: &[u8; crate::directory::ENCAPSULATED_MESSAGE_BYTES] = response
-            .try_into()
-            .map_err(|_| InternalEncapsulationError::InvalidSize(response.len()))?;
-        let response = ohttp_decapsulate(self.state.ohttp_ctx, response_array)
-            .map_err(InternalEncapsulationError::Ohttp)?;
-        match response.status() {
-            http::StatusCode::OK => {
-                // return OK with new Typestate
-                Ok(Sender {
-                    state: V2GetContext {
-                        endpoint: self.state.endpoint,
-                        psbt_ctx: self.state.psbt_ctx,
-                        hpke_ctx: self.state.hpke_ctx,
-                    },
-                })
-            }
-            _ => Err(InternalEncapsulationError::UnexpectedStatusCode(response.status()))?,
-        }
+        process_post_res(response, self.state.ohttp_ctx)
+            .map_err(InternalEncapsulationError::DirectoryResponse)?;
+        Ok(Sender {
+            state: V2GetContext {
+                endpoint: self.state.endpoint,
+                psbt_ctx: self.state.psbt_ctx,
+                hpke_ctx: self.state.hpke_ctx,
+            },
+        })
     }
 }
 
@@ -404,16 +395,11 @@ impl Sender<V2GetContext> {
         response: &[u8],
         ohttp_ctx: ohttp::ClientResponse,
     ) -> Result<Option<Psbt>, ResponseError> {
-        let response_array: &[u8; crate::directory::ENCAPSULATED_MESSAGE_BYTES] = response
-            .try_into()
-            .map_err(|_| InternalEncapsulationError::InvalidSize(response.len()))?;
-
-        let response = ohttp_decapsulate(ohttp_ctx, response_array)
-            .map_err(InternalEncapsulationError::Ohttp)?;
-        let body = match response.status() {
-            http::StatusCode::OK => response.body().to_vec(),
-            http::StatusCode::ACCEPTED => return Ok(None),
-            _ => return Err(InternalEncapsulationError::UnexpectedStatusCode(response.status()))?,
+        let body = match process_get_res(response, ohttp_ctx)
+            .map_err(InternalEncapsulationError::DirectoryResponse)?
+        {
+            Some(body) => body,
+            None => return Ok(None),
         };
         let psbt = decrypt_message_b(
             &body,

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,6 @@`
`1`	`1`	`use core::fmt;`
`2`	`2`
	`3`	`+use crate::ohttp::DirectoryResponseError;`
`3`	`4`	`use crate::uri::url_ext::ParseReceiverPubkeyParamError;`
`4`	`5`
`5`	`6`	`/// Error returned when request could not be created.`
`@@ -73,25 +74,19 @@ pub struct EncapsulationError(InternalEncapsulationError);`
`73`	`74`
`74`	`75`	`#[derive(Debug)]`
`75`	`76`	`pub(crate) enum InternalEncapsulationError {`
`76`		`- /// The response size is not the expected size.`
`77`		`- InvalidSize(usize),`
`78`		`- /// The status code is not the expected status code.`
`79`		`- UnexpectedStatusCode(http::StatusCode),`
`80`	`77`	`/// The HPKE failed.`
`81`	`78`	`Hpke(crate::hpke::HpkeError),`
`82`		`- /// The encapsulation failed.`
`83`		`- Ohttp(crate::ohttp::OhttpEncapsulationError),`
	`79`	`+ /// The directory returned a bad response`
	`80`	`+ DirectoryResponse(DirectoryResponseError),`
`84`	`81`	`}`
`85`	`82`
`86`	`83`	`impl fmt::Display for EncapsulationError {`
`87`	`84`	`fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {`
`88`	`85`	`use InternalEncapsulationError::*;`
`89`	`86`
`90`	`87`	`match &self.0 {`
`91`		`- InvalidSize(size) => write!(f, "invalid size: {size}"),`
`92`		`- UnexpectedStatusCode(status) => write!(f, "unexpected status code: {status}"),`
`93`		`- Ohttp(error) => write!(f, "OHTTP encapsulation error: {error}"),`
`94`	`88`	`Hpke(error) => write!(f, "HPKE error: {error}"),`
	`89`	`+ DirectoryResponse(e) => write!(f, "Directory response error: {e}"),`
`95`	`90`	`}`
`96`	`91`	`}`
`97`	`92`	`}`
`@@ -101,10 +96,8 @@ impl std::error::Error for EncapsulationError {`
`101`	`96`	`use InternalEncapsulationError::*;`
`102`	`97`
`103`	`98`	`match &self.0 {`
`104`		`- InvalidSize(_) => None,`
`105`		`- UnexpectedStatusCode(_) => None,`
`106`		`- Ohttp(error) => Some(error),`
`107`	`99`	`Hpke(error) => Some(error),`
	`100`	`+ DirectoryResponse(e) => Some(e),`
`108`	`101`	`}`
`109`	`102`	`}`
`110`	`103`	`}`
`@@ -115,8 +108,6 @@ impl From<InternalEncapsulationError> for EncapsulationError {`
`115`	`108`
`116`	`109`	`impl From<InternalEncapsulationError> for super::ResponseError {`
`117`	`110`	`fn from(value: InternalEncapsulationError) -> Self {`
`118`		`- super::ResponseError::Validation(`
`119`		`- super::InternalValidationError::V2Encapsulation(value.into()).into(),`
`120`		`- )`
	`111`	`+ super::InternalValidationError::V2Encapsulation(value.into()).into()`
`121`	`112`	`}`
`122`	`113`	`}`