titan-api-codec 1.2.9

//! Codecs for version 1 of the Titan WebSocket API.

use std::str::FromStr;

use titan_api_types::ws::v1;

use crate::{
    codec::{Codec, CodecLoadError, TypedDecoder, TypedEncoder, WrappedDecoder, WrappedEncoder},
    dec::{messagepack::MessagePackDecoder, DecodeError, Decoder},
    enc::{messagepack::MessagePackEncoder, EncodeError, Encoder},
    transform::{
        BrotliCompressor, BrotliDecompressor, GzipCompressor, GzipDecompressor, ZstdCompressor,
        ZstdDecompressor,
    },
};

/// Codec for clients to the Titan WebSocket API version 1.
#[derive(Debug, Default, PartialEq, Eq)]
#[non_exhaustive]
pub enum ClientCodec {
    /// V1 messages encoded via MessagePack with no compression.
    #[default]
    Uncompressed,
    /// V1 messages encoded via MessagePack with zstd compression.
    #[cfg(feature = "zstd")]
    Zstd,
    /// V1 messages encoded via MessagePack with brotli compression.
    #[cfg(feature = "brotli")]
    Brotli,
    /// V1 messages encoded via MessagePack with gzip compression.
    #[cfg(feature = "gzip")]
    Gzip,
}

/// Codec for servers providing the Titan WebSocket API version 1.
#[derive(Debug, Default, PartialEq, Eq)]
#[non_exhaustive]
pub enum ServerCodec {
    /// V1 messages encoded via MessagePack with no compression.
    #[default]
    Uncompressed,
    /// V1 messages encoded via MessagePack with zstd compression.
    #[cfg(feature = "zstd")]
    Zstd,
    /// V1 messages encoded via MessagePack with brotli compression.
    #[cfg(feature = "brotli")]
    Brotli,
    /// V1 messages encoded via MessagePack with gzip compression.
    #[cfg(feature = "gzip")]
    Gzip,
}

impl FromStr for ClientCodec {
    type Err = CodecLoadError;

    /// Attempts to load a [`ClientCodec`] from a WebSocket subprotocol string.
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            v1::WEBSOCKET_SUBPROTO_BASE => Ok(Self::Uncompressed),
            #[cfg(feature = "zstd")]
            v1::WEBSOCKET_SUBPROTO_ZSTD => Ok(Self::Zstd),
            #[cfg(not(feature = "zstd"))]
            v1::WEBSOCKET_SUBPROTO_ZSTD => Err(CodecLoadError::DisabledEncoding("zstd")),
            #[cfg(feature = "brotli")]
            v1::WEBSOCKET_SUBPROTO_BROTLI => Ok(Self::Brotli),
            #[cfg(not(feature = "brotli"))]
            v1::WEBSOCKET_SUBPROTO_ZSTD => Err(CodecLoadError::DisabledEncoding("brotli")),
            #[cfg(feature = "gzip")]
            v1::WEBSOCKET_SUBPROTO_GZIP => Ok(Self::Gzip),
            #[cfg(not(feature = "gzip"))]
            v1::WEBSOCKET_SUBPROTO_ZSTD => Err(CodecLoadError::DisabledEncoding("gzip")),
            _ => Err(CodecLoadError::UnsupportedProtocol(s.into())),
        }
    }
}

impl Codec for ClientCodec {
    type SendItem = v1::ClientRequest;
    type SendError = EncodeError;
    type RecvItem = v1::ServerMessage;
    type RecvError = DecodeError;

    fn encoder(
        &self,
    ) -> Box<dyn TypedEncoder<Self::SendItem, Error = Self::SendError> + Send + Sync> {
        match self {
            Self::Uncompressed => Box::new(WrappedEncoder::new(MessagePackEncoder::default())),
            #[cfg(feature = "zstd")]
            Self::Zstd => Box::new(WrappedEncoder::new(
                MessagePackEncoder::default().transform(ZstdCompressor::default()),
            )),
            #[cfg(feature = "brotli")]
            Self::Brotli => Box::new(WrappedEncoder::new(
                MessagePackEncoder::default().transform(BrotliCompressor::default()),
            )),
            #[cfg(feature = "gzip")]
            Self::Gzip => Box::new(WrappedEncoder::new(
                MessagePackEncoder::default().transform(GzipCompressor::default()),
            )),
        }
    }

    fn decoder(
        &self,
    ) -> Box<dyn TypedDecoder<Item = Self::RecvItem, Error = Self::RecvError> + Send + Sync> {
        match self {
            Self::Uncompressed => Box::new(WrappedDecoder::new(MessagePackDecoder::default())),
            #[cfg(feature = "zstd")]
            Self::Zstd => Box::new(WrappedDecoder::new(
                MessagePackDecoder::default().transformed(ZstdDecompressor::default()),
            )),
            #[cfg(feature = "brotli")]
            Self::Brotli => Box::new(WrappedDecoder::new(
                MessagePackDecoder::default().transformed(BrotliDecompressor::default()),
            )),
            #[cfg(feature = "gzip")]
            Self::Gzip => Box::new(WrappedDecoder::new(
                MessagePackDecoder::default().transformed(GzipDecompressor::default()),
            )),
        }
    }
}

impl FromStr for ServerCodec {
    type Err = CodecLoadError;

    /// Attempts to load a [`ServerCodec`] from a WebSocket subprotocol string.
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            v1::WEBSOCKET_SUBPROTO_BASE => Ok(Self::Uncompressed),
            #[cfg(feature = "zstd")]
            v1::WEBSOCKET_SUBPROTO_ZSTD => Ok(Self::Zstd),
            #[cfg(not(feature = "zstd"))]
            v1::WEBSOCKET_SUBPROTO_ZSTD => Err(CodecLoadError::DisabledEncoding("zstd")),
            #[cfg(feature = "brotli")]
            v1::WEBSOCKET_SUBPROTO_BROTLI => Ok(Self::Brotli),
            #[cfg(not(feature = "brotli"))]
            v1::WEBSOCKET_SUBPROTO_ZSTD => Err(CodecLoadError::DisabledEncoding("brotli")),
            #[cfg(feature = "gzip")]
            v1::WEBSOCKET_SUBPROTO_GZIP => Ok(Self::Gzip),
            #[cfg(not(feature = "gzip"))]
            v1::WEBSOCKET_SUBPROTO_ZSTD => Err(CodecLoadError::DisabledEncoding("gzip")),
            _ => Err(CodecLoadError::UnsupportedProtocol(s.into())),
        }
    }
}

impl Codec for ServerCodec {
    type SendItem = v1::ServerMessage;
    type SendError = EncodeError;
    type RecvItem = v1::ClientRequest;
    type RecvError = DecodeError;

    fn encoder(
        &self,
    ) -> Box<dyn TypedEncoder<Self::SendItem, Error = Self::SendError> + Send + Sync> {
        match self {
            Self::Uncompressed => Box::new(WrappedEncoder::new(MessagePackEncoder::default())),
            #[cfg(feature = "zstd")]
            Self::Zstd => Box::new(WrappedEncoder::new(
                MessagePackEncoder::default().transform(ZstdCompressor::default()),
            )),
            #[cfg(feature = "brotli")]
            Self::Brotli => Box::new(WrappedEncoder::new(
                MessagePackEncoder::default().transform(BrotliCompressor::default()),
            )),
            #[cfg(feature = "gzip")]
            Self::Gzip => Box::new(WrappedEncoder::new(
                MessagePackEncoder::default().transform(GzipCompressor::default()),
            )),
        }
    }

    fn decoder(
        &self,
    ) -> Box<dyn TypedDecoder<Item = Self::RecvItem, Error = Self::RecvError> + Send + Sync> {
        match self {
            Self::Uncompressed => Box::new(WrappedDecoder::new(MessagePackDecoder::default())),
            #[cfg(feature = "zstd")]
            Self::Zstd => Box::new(WrappedDecoder::new(
                MessagePackDecoder::default().transformed(ZstdDecompressor::default()),
            )),
            #[cfg(feature = "brotli")]
            Self::Brotli => Box::new(WrappedDecoder::new(
                MessagePackDecoder::default().transformed(BrotliDecompressor::default()),
            )),
            #[cfg(feature = "gzip")]
            Self::Gzip => Box::new(WrappedDecoder::new(
                MessagePackDecoder::default().transformed(GzipDecompressor::default()),
            )),
        }
    }
}

#[cfg(test)]
mod test {
    use std::str::FromStr;

    use titan_api_types::ws::v1;

    use crate::codec::{Codec, CodecLoadError};

    use super::{ClientCodec, ServerCodec};

    #[test]
    fn construct_client_codec_base() {
        let codec = ClientCodec::from_str(v1::WEBSOCKET_SUBPROTO_BASE)
            .expect("should construct base codec");
        assert_eq!(codec, ClientCodec::Uncompressed);
    }

    #[cfg(feature = "zstd")]
    #[test]
    fn construct_client_codec_zstd() {
        let codec = ClientCodec::from_str(v1::WEBSOCKET_SUBPROTO_ZSTD)
            .expect("should construct zstd codec");
        assert_eq!(codec, ClientCodec::Zstd);
    }

    #[cfg(feature = "brotli")]
    #[test]
    fn construct_client_codec_brotli() {
        let codec = ClientCodec::from_str(v1::WEBSOCKET_SUBPROTO_BROTLI)
            .expect("should construct brotli codec");
        assert_eq!(codec, ClientCodec::Brotli);
    }

    #[cfg(feature = "gzip")]
    #[test]
    fn construct_client_codec_gzip() {
        let codec = ClientCodec::from_str(v1::WEBSOCKET_SUBPROTO_GZIP)
            .expect("should construct gzip codec");
        assert_eq!(codec, ClientCodec::Gzip);
    }

    #[test]
    fn construct_client_codec_invalid() {
        let err =
            ClientCodec::from_str("invalid").expect_err("should have errored on invalid protocol");
        assert_eq!(err, CodecLoadError::UnsupportedProtocol("invalid".into()))
    }

    #[test]
    fn construct_server_codec_base() {
        let codec = ServerCodec::from_str(v1::WEBSOCKET_SUBPROTO_BASE)
            .expect("should construct base codec");
        assert_eq!(codec, ServerCodec::Uncompressed);
    }

    #[cfg(feature = "zstd")]
    #[test]
    fn construct_server_codec_zstd() {
        let codec = ServerCodec::from_str(v1::WEBSOCKET_SUBPROTO_ZSTD)
            .expect("should construct zstd codec");
        assert_eq!(codec, ServerCodec::Zstd);
    }

    #[cfg(feature = "brotli")]
    #[test]
    fn construct_server_codec_brotli() {
        let codec = ServerCodec::from_str(v1::WEBSOCKET_SUBPROTO_BROTLI)
            .expect("should construct brotli codec");
        assert_eq!(codec, ServerCodec::Brotli);
    }

    #[cfg(feature = "gzip")]
    #[test]
    fn construct_server_codec_gzip() {
        let codec = ServerCodec::from_str(v1::WEBSOCKET_SUBPROTO_GZIP)
            .expect("should construct gzip codec");
        assert_eq!(codec, ServerCodec::Gzip);
    }

    #[test]
    fn construct_server_codec_invalid() {
        let err =
            ServerCodec::from_str("invalid").expect_err("should have errored on invalid protocol");
        assert_eq!(err, CodecLoadError::UnsupportedProtocol("invalid".into()))
    }

    #[test]
    fn roundtrip_base() {
        let client_codec = ClientCodec::Uncompressed;
        let mut encoder = client_codec.encoder();
        let server_codec = ServerCodec::Uncompressed;
        let mut decoder = server_codec.decoder();

        let request = v1::ClientRequest {
            id: 1,
            data: v1::RequestData::GetInfo(v1::GetInfoRequest::default()),
        };
        let encoded = encoder
            .encode_mut(&request)
            .expect("should encode client request");
        let decoded = decoder
            .decode_mut(encoded)
            .expect("should decode client request");

        assert_eq!(request, decoded);
    }
}