Handle all CDSI server error codes

Match against all the error codes the documentation says the server can 
produce. Map these to error types in the app languages.

Cargo.lock

@@ -2013,6 +2013,7 @@ dependencies = [
  "libsignal-svr3",
  "log",
  "nonzero_ext",
+ "num_enum",
  "pin-project-lite",
  "proptest",
  "proptest-state-machine",
@@ -2025,6 +2026,7 @@ dependencies = [
  "serde_json",
  "sha2",
  "snow",
+ "strum",
  "thiserror",
  "tokio",
  "tokio-boring",

java/client/src/test/java/org/signal/libsignal/net/CdsiLookupResponseTest.java

@@ -71,12 +71,17 @@ public class CdsiLookupResponseTest {
 
     assertLookupErrorIs(
         "InvalidToken", CdsiInvalidTokenException.class, "Request token was invalid");
+    assertLookupErrorIs(
+        "InvalidArgument",
+        IllegalArgumentException.class,
+        "invalid argument: request was invalid: fake reason");
     assertLookupErrorIs(
         "Parse", CdsiProtocolException.class, "Failed to parse the response from the server");
     assertLookupErrorIs("ConnectDnsFailed", IOException.class, "DNS lookup failed");
     assertLookupErrorIs(
         "WebSocketIdleTooLong", NetworkException.class, "channel was idle for too long");
     assertLookupErrorIs("Timeout", NetworkException.class, "timeout");
+    assertLookupErrorIs("ServerCrashed", CdsiProtocolException.class, "Server error: crashed");
   }
 
   private static <E extends Exception> E assertLookupErrorIs(

node/ts/test/NetTest.ts

@@ -168,6 +168,11 @@ describe('cdsi lookup', () => {
           ErrorCode.CdsiInvalidToken,
           'request token was invalid',
         ],
+        [
+          'InvalidArgument',
+          ErrorCode.Generic,
+          'request was invalid: fake reason',
+        ],
         [
           'Parse',
           ErrorCode.IoError,
@@ -184,6 +189,7 @@ describe('cdsi lookup', () => {
           'websocket error: channel was idle for too long',
         ],
         ['Timeout', ErrorCode.IoError, 'lookup timed out'],
+        ['ServerCrashed', ErrorCode.IoError, 'server error: crashed'],
       ];
       cases.forEach((testCase) => {
         const [name, expectedCode, expectedMessage] = testCase;

rust/bridge/shared/src/ffi/error.rs

@@ -200,7 +200,10 @@ impl From<libsignal_net::cdsi::LookupError> for SignalFfiError {
             LookupError::ConnectTransport(e) => SignalFfiError::Io(e.into()),
             LookupError::WebSocket(e) => SignalFfiError::WebSocket(e),
             LookupError::Timeout => SignalFfiError::Timeout,
-            LookupError::ParseError | LookupError::Protocol | LookupError::InvalidResponse => {
+            LookupError::ParseError
+            | LookupError::Protocol
+            | LookupError::InvalidResponse
+            | LookupError::Server { reason: _ } => {
                 SignalFfiError::NetworkProtocol(value.to_string())
             }
             LookupError::RateLimited {
@@ -209,6 +212,9 @@ impl From<libsignal_net::cdsi::LookupError> for SignalFfiError {
                 retry_after_seconds: retry_after,
             },
             LookupError::InvalidToken => SignalFfiError::CdsiInvalidToken,
+            LookupError::InvalidArgument { server_reason: _ } => {
+                SignalFfiError::Signal(SignalProtocolError::InvalidArgument(value.to_string()))
+            }
         }
     }
 }

rust/bridge/shared/src/jni/error.rs

@@ -255,6 +255,11 @@ impl From<libsignal_net::cdsi::LookupError> for SignalJniError {
             LookupError::AttestationError(e) => return e.into(),
             LookupError::ConnectTransport(e) => return IoError::from(e).into(),
             LookupError::WebSocket(e) => return e.into(),
+            LookupError::InvalidArgument { server_reason: _ } => {
+                return SignalJniError::Protocol(SignalProtocolError::InvalidArgument(
+                    e.to_string(),
+                ))
+            }
             LookupError::InvalidResponse => CdsiError::InvalidResponse,
             LookupError::Protocol => CdsiError::Protocol,
             LookupError::RateLimited {
@@ -264,6 +269,7 @@ impl From<libsignal_net::cdsi::LookupError> for SignalJniError {
             },
             LookupError::ParseError => CdsiError::ParseError,
             LookupError::InvalidToken => CdsiError::InvalidToken,
+            LookupError::Server { reason } => CdsiError::Server { reason },
         })
     }
 }

rust/bridge/shared/src/jni/mod.rs

@@ -572,7 +572,10 @@ where
             error,
         ),
         SignalJniError::Cdsi(
-            CdsiError::InvalidResponse | CdsiError::ParseError | CdsiError::Protocol,
+            CdsiError::InvalidResponse
+            | CdsiError::ParseError
+            | CdsiError::Protocol
+            | CdsiError::Server { reason: _ },
         ) => (
             jni_class_name!(org.signal.libsignal.net.CdsiProtocolException),
             error,

rust/bridge/shared/src/net/cdsi.rs

@@ -34,6 +34,8 @@ pub enum CdsiError {
     ParseError,
     /// Request token was invalid
     InvalidToken,
+    /// Server error: {reason}
+    Server { reason: &'static str },
 }
 
 #[derive(Default)]

rust/bridge/shared/src/node/error.rs

@@ -445,7 +445,7 @@ impl SignalNodeError for libsignal_net::cdsi::LookupError {
             Self::RateLimited {
                 retry_after_seconds,
             } => (
-                RATE_LIMITED_ERROR,
+                Some(RATE_LIMITED_ERROR),
                 Some({
                     let props = cx.empty_object();
                     let retry_after = retry_after_seconds.convert_into(cx)?;
@@ -454,26 +454,21 @@ impl SignalNodeError for libsignal_net::cdsi::LookupError {
                 }),
             ),
             Self::AttestationError(e) => return e.throw(cx, module, operation_name),
-            Self::InvalidToken => ("CdsiInvalidToken", None),
+            Self::InvalidArgument { server_reason: _ } => (None, None),
+            Self::InvalidToken => (Some("CdsiInvalidToken"), None),
             Self::Timeout
             | Self::ConnectTransport(_)
             | Self::WebSocket(_)
             | Self::Protocol
             | Self::InvalidResponse
-            | Self::ParseError => (IO_ERROR, None),
+            | Self::ParseError
+            | Self::Server { reason: _ } => (Some(IO_ERROR), None),
         };
         let message = self.to_string();
-        new_js_error(
-            cx,
-            module,
-            Some(name),
-            &message,
-            operation_name,
-            extra_props,
-        )
-        .map(|e| cx.throw(e))
-        // Make sure we still throw something.
-        .unwrap_or_else(|| cx.throw_error(&message))
+        new_js_error(cx, module, name, &message, operation_name, extra_props)
+            .map(|e| cx.throw(e))
+            // Make sure we still throw something.
+            .unwrap_or_else(|| cx.throw_error(&message))
     }
 }
 

rust/bridge/shared/src/testing/net.rs

@@ -53,12 +53,40 @@ enum TestingCdsiLookupError {
     InvalidResponse,
     RetryAfter42Seconds,
     InvalidToken,
+    InvalidArgument,
     Parse,
     ConnectDnsFailed,
     WebSocketIdleTooLong,
     Timeout,
+    ServerCrashed,
 }
 
+const _: () = {
+    /// This code isn't ever executed. It exists so that when new cases are
+    /// added to `LookupError`, this will fail to compile until corresponding
+    /// cases are added to `TestingCdsiLookupError`
+    #[allow(unused)]
+    fn match_on_lookup_error(value: &'static LookupError) -> TestingCdsiLookupError {
+        match value {
+            LookupError::Protocol => TestingCdsiLookupError::Protocol,
+            LookupError::AttestationError(_) => TestingCdsiLookupError::AttestationDataError,
+            LookupError::InvalidResponse => TestingCdsiLookupError::InvalidResponse,
+            LookupError::RateLimited {
+                retry_after_seconds: _,
+            } => TestingCdsiLookupError::RetryAfter42Seconds,
+            LookupError::InvalidToken => TestingCdsiLookupError::InvalidToken,
+            LookupError::InvalidArgument { server_reason: _ } => {
+                TestingCdsiLookupError::InvalidArgument
+            }
+            LookupError::ParseError => TestingCdsiLookupError::Parse,
+            LookupError::ConnectTransport(_) => TestingCdsiLookupError::ConnectDnsFailed,
+            LookupError::WebSocket(_) => TestingCdsiLookupError::WebSocketIdleTooLong,
+            LookupError::Timeout => TestingCdsiLookupError::Timeout,
+            LookupError::Server { reason } => TestingCdsiLookupError::ServerCrashed,
+        }
+    }
+};
+
 impl TryFrom<String> for TestingCdsiLookupError {
     type Error = <Self as FromStr>::Err;
     fn try_from(value: String) -> Result<Self, Self::Error> {
@@ -84,6 +112,9 @@ fn TESTING_CdsiLookupErrorConvert(
             retry_after_seconds: 42,
         },
         TestingCdsiLookupError::InvalidToken => LookupError::InvalidToken,
+        TestingCdsiLookupError::InvalidArgument => LookupError::InvalidArgument {
+            server_reason: "fake reason".into(),
+        },
         TestingCdsiLookupError::Parse => LookupError::ParseError,
         TestingCdsiLookupError::ConnectDnsFailed => LookupError::ConnectTransport(
             libsignal_net::infra::errors::TransportConnectError::DnsError,
@@ -92,6 +123,7 @@ fn TESTING_CdsiLookupErrorConvert(
             libsignal_net::infra::ws::WebSocketServiceError::ChannelIdleTooLong,
         ),
         TestingCdsiLookupError::Timeout => LookupError::Timeout,
+        TestingCdsiLookupError::ServerCrashed => LookupError::Server { reason: "crashed" },
     })
 }
 

rust/net/Cargo.toml

@@ -29,6 +29,7 @@ hyper = { version = "1.2.0", features = ["http1", "http2", "client"] }
 itertools = "0.12.0"
 lazy_static = "1.4.0"
 log = "0.4.19"
+num_enum = "0.6.1"
 pin-project-lite = "0.2.4"
 prost = "0.12.1"
 rand = "0.8.0"
@@ -37,6 +38,7 @@ rustls-native-certs = "0.6.3"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 sha2 = "0.10.8"
+strum = { version = "0.26", features = ["derive"] }
 thiserror = "1.0.38"
 tokio = { version = "1", features = ["rt", "time", "macros"] }
 tokio-boring = { git = "https://github.com/signalapp/boring", branch = "libsignal" }

rust/net/src/cdsi.rs

@@ -14,6 +14,8 @@ use prost::Message as _;
 use thiserror::Error;
 use tokio::net::TcpStream;
 use tokio_boring::SslStream;
+use tungstenite::protocol::frame::coding::CloseCode;
+use tungstenite::protocol::CloseFrame;
 use uuid::Uuid;
 
 use crate::auth::HttpBasicAuth;
@@ -285,6 +287,10 @@ pub enum LookupError {
     WebSocket(WebSocketServiceError),
     /// lookup timed out
     Timeout,
+    /// request was invalid: {server_reason}
+    InvalidArgument { server_reason: String },
+    /// server error: {reason}
+    Server { reason: &'static str },
 }
 
 impl From<AttestedConnectionError> for LookupError {
@@ -342,13 +348,6 @@ struct RateLimitExceededResponse {
     retry_after_seconds: u32,
 }
 
-impl RateLimitExceededResponse {
-    /// Numeric code set by the server on the websocket close frame.
-    const CLOSE_CODE: u16 = 4008;
-}
-
-const INVALID_TOKEN_CLOSE_CODE: u16 = 4101;
-
 #[cfg_attr(test, derive(Debug))]
 pub struct ClientResponseCollector<S = SslStream<TcpStream>>(CdsiConnection<S>);
 
@@ -372,28 +371,11 @@ impl<S: AsyncDuplexStream> CdsiConnection<S> {
         request: LookupRequest,
     ) -> Result<(Token, ClientResponseCollector<S>), LookupError> {
         self.0.send(request.into_client_request()).await?;
-        let token_response: ClientResponse = match self.0.receive().await? {
-            NextOrClose::Next(response) => response,
-            NextOrClose::Close(close) => {
-                if let Some(close) = close {
-                    match u16::from(close.code) {
-                        RateLimitExceededResponse::CLOSE_CODE => {
-                            if let Ok(RateLimitExceededResponse {
-                                retry_after_seconds,
-                            }) = serde_json::from_str(&close.reason)
-                            {
-                                return Err(LookupError::RateLimited {
-                                    retry_after_seconds,
-                                });
-                            }
-                        }
-                        INVALID_TOKEN_CLOSE_CODE => return Err(LookupError::InvalidToken),
-                        _ => (),
-                    }
-                };
-                return Err(LookupError::Protocol);
-            }
-        };
+        let token_response: ClientResponse = self.0.receive().await?.next_or_else(|close| {
+            close
+                .and_then(err_for_close)
+                .unwrap_or(LookupError::Protocol)
+        })?;
 
         if token_response.token.is_empty() {
             return Err(LookupError::Protocol);
@@ -416,20 +398,76 @@ impl<S: AsyncDuplexStream> ClientResponseCollector<S> {
         };
 
         connection.0.send(token_ack).await?;
-        let mut response: ClientResponse = connection
-            .0
-            .receive()
-            .await?
-            .next_or(LookupError::Protocol)?;
-        while let NextOrClose::Next(decoded) = connection.0.receive_bytes().await? {
-            response
-                .merge(decoded.as_ref())
-                .map_err(LookupError::from)?;
+        let mut response: ClientResponse = connection.0.receive().await?.next_or_else(|close| {
+            close
+                .and_then(err_for_close)
+                .unwrap_or(LookupError::Protocol)
+        })?;
+        loop {
+            match connection.0.receive_bytes().await? {
+                NextOrClose::Next(decoded) => {
+                    response
+                        .merge(decoded.as_ref())
+                        .map_err(LookupError::from)?;
+                }
+                NextOrClose::Close(
+                    None
+                    | Some(CloseFrame {
+                        code: CloseCode::Normal,
+                        reason: _,
+                    }),
+                ) => break,
+                NextOrClose::Close(Some(close)) => {
+                    return Err(err_for_close(close).unwrap_or(LookupError::Protocol))
+                }
+            }
         }
         Ok(response.try_into()?)
     }
 }
 
+/// Produces a [`LookupError`] for the provided [`CloseFrame`].
+///
+/// Returns `Some(err)` if there is a relevant `LookupError` value for the
+/// provided close frame. Otherwise returns `None`.
+fn err_for_close(CloseFrame { code, reason }: CloseFrame<'_>) -> Option<LookupError> {
+    /// Numeric code set by the server on the websocket close frame.
+    #[repr(u16)]
+    #[derive(Copy, Clone, num_enum::TryFromPrimitive, strum::IntoStaticStr)]
+    enum CdsiCloseCode {
+        InvalidArgument = 4003,
+        RateLimitExceeded = 4008,
+        ServerInternalError = 4013,
+        ServerUnavailable = 4014,
+        InvalidToken = 4101,
+    }
+
+    let Ok(code) = CdsiCloseCode::try_from(u16::from(code)) else {
+        log::warn!("got unexpected websocket error code: {code}",);
+        return None;
+    };
+
+    match code {
+        CdsiCloseCode::InvalidArgument => Some(LookupError::InvalidArgument {
+            server_reason: reason.into_owned(),
+        }),
+        CdsiCloseCode::InvalidToken => Some(LookupError::InvalidToken),
+        CdsiCloseCode::RateLimitExceeded => {
+            let RateLimitExceededResponse {
+                retry_after_seconds,
+            } = serde_json::from_str(&reason).ok()?;
+            Some(LookupError::RateLimited {
+                retry_after_seconds,
+            })
+        }
+        CdsiCloseCode::ServerInternalError | CdsiCloseCode::ServerUnavailable => {
+            Some(LookupError::Server {
+                reason: code.into(),
+            })
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::time::Duration;
@@ -760,8 +798,12 @@ mod test {
 
         let response = collector.collect().await;
 
-        // TODO this is wrong, this should be a RateLimited error.
-        assert_matches!(response, Err(LookupError::Protocol));
+        assert_matches!(
+            response,
+            Err(LookupError::RateLimited {
+                retry_after_seconds: RETRY_AFTER_SECS
+            })
+        )
     }
 
     #[tokio::test]

rust/net/src/infra/ws.rs

@@ -495,6 +495,16 @@ impl<T> NextOrClose<T> {
             Self::Next(t) => Ok(t),
         }
     }
+
+    pub fn next_or_else<E>(
+        self,
+        on_close: impl FnOnce(Option<CloseFrame<'static>>) -> E,
+    ) -> Result<T, E> {
+        match self {
+            Self::Next(t) => Ok(t),
+            Self::Close(close) => Err(on_close(close)),
+        }
+    }
 }
 
 impl<S> AttestedConnection<S>

swift/Tests/LibSignalClientTests/NetTests.swift

@@ -71,6 +71,11 @@ final class NetTests: XCTestCase {
         } catch SignalError.cdsiInvalidToken(let message) {
             XCTAssertEqual(message, "CDSI request token was invalid")
         }
+        do {
+            try failWithError("InvalidArgument")
+        } catch SignalError.invalidArgument(let message) {
+            XCTAssertEqual(message, "invalid argument: request was invalid: fake reason")
+        }
         do {
             try failWithError("Parse")
         } catch SignalError.networkProtocolError(let message) {
@@ -91,6 +96,11 @@ final class NetTests: XCTestCase {
         } catch SignalError.timeoutError(let message) {
             XCTAssertEqual(message, "Operation timed out")
         }
+        do {
+            try failWithError("ServerCrashed")
+        } catch SignalError.networkProtocolError(let message) {
+            XCTAssertEqual(message, "Protocol error: server error: crashed")
+        }
     }
 
     func testCdsiLookupCompilation() async throws {