Merge pull request #115 from xmtp/daniel-compression

feat: support compression on encoded content types
xmtp · Feb 3, 2024 · b7af067 · b7af067
2 parents 634597a + 353663b
commit b7af067
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 2 deletions.
diff --git a/lib/src/content/codec_registry.dart b/lib/src/content/codec_registry.dart
@@ -1,15 +1,26 @@
+import 'dart:convert' as convert;
+import 'dart:io' as io;
+
 import 'package:xmtp_proto/xmtp_proto.dart' as xmtp;
 
 import 'codec.dart';
 import 'decoded.dart';
 import 'text_codec.dart';
 
+typedef Compressor = convert.Codec<List<int>, List<int>>;
+
 /// This is a registry of codecs for particular types.
 ///
 /// It knows how to apply the codecs to [decode] or [encode]
 /// [xmtp.EncodedContent] to [DecodedContent]..
 class CodecRegistry implements Codec<DecodedContent> {
   final Map<String, Codec> _codecs = {};
+  static final Map<xmtp.Compression, Compressor> _compressors = {
+    xmtp.Compression.COMPRESSION_GZIP: io.gzip,
+    xmtp.Compression.COMPRESSION_DEFLATE: io.zlib,
+  }; // TODO: consider supporting custom compressors
+  static final Set<xmtp.Compression> supportedCompressions =
+      _compressors.keys.toSet();
 
   void registerCodec(Codec codec) => _codecs[_key(codec.contentType)] = codec;
 
@@ -20,6 +31,18 @@ class CodecRegistry implements Codec<DecodedContent> {
   /// Use the registered codecs to decode the [encoded] content.
   @override
   Future<DecodedContent> decode(xmtp.EncodedContent encoded) async {
+    if (encoded.hasCompression()) {
+      var compressor = _compressors[encoded.compression];
+      if (compressor == null) {
+        throw StateError(
+            "unable to decode unsupported compression ${encoded.compression}");
+      }
+      var decompressed = compressor.decode(encoded.content);
+      encoded = xmtp.EncodedContent()
+        ..mergeFromMessage(encoded)
+        ..clearCompression()
+        ..content = decompressed;
+    }
     var codec = _codecFor(encoded.type);
     if (codec == null) {
       if (encoded.hasFallback()) {
@@ -34,13 +57,30 @@ class CodecRegistry implements Codec<DecodedContent> {
 
   /// Use the registered codecs to encode the [content].
   @override
-  Future<xmtp.EncodedContent> encode(DecodedContent decoded) async {
+  Future<xmtp.EncodedContent> encode(
+    DecodedContent decoded, {
+    xmtp.Compression? compression,
+  }) async {
     var type = decoded.contentType;
     var codec = _codecFor(type);
     if (codec == null) {
       throw StateError("unable to encode unsupported type ${_key(type)}");
     }
-    return codec.encode(decoded.content);
+    var encoded = await codec.encode(decoded.content);
+    // TODO: consider warning if it isn't compressed but should be
+    if (compression != null) {
+      var compressor = _compressors[compression];
+      if (compressor == null) {
+        throw StateError(
+            "unable to encode unsupported compression $compression");
+      }
+      var compressed = compressor.encode(encoded.content);
+      encoded = xmtp.EncodedContent()
+        ..mergeFromMessage(encoded)
+        ..compression = compression
+        ..content = compressed;
+    }
+    return encoded;
   }
 
   @override

diff --git a/test/content/codec_registry_test.dart b/test/content/codec_registry_test.dart
@@ -18,6 +18,34 @@ void main() {
     expect(decoded.content, "foo bar");
   });
 
+  test('compression should work during encoding and decoding', () async {
+    var registry = CodecRegistry();
+    registry.registerCodec(TextCodec());
+    var someText = "blah blah blah" * 100;
+    for (var compression in CodecRegistry.supportedCompressions) {
+      var encodedSmall = await registry.encode(
+        DecodedContent(contentTypeText, someText),
+        compression: compression,
+      );
+      var encodedLarge = await registry.encode(
+        DecodedContent(contentTypeText, someText),
+      );
+      expect(encodedSmall.type, contentTypeText);
+      expect(encodedLarge.type, contentTypeText);
+      expect(encodedSmall.hasCompression(), true);
+      expect(encodedLarge.hasCompression(), false);
+      expect(encodedSmall.content.isNotEmpty, true);
+      expect(encodedLarge.content.isNotEmpty, true);
+      expect(encodedSmall.content.length < encodedLarge.content.length, true);
+      var decodedSmall = await registry.decode(encodedSmall);
+      var decodedLarge = await registry.decode(encodedLarge);
+      expect(decodedSmall.contentType, contentTypeText);
+      expect(decodedLarge.contentType, contentTypeText);
+      expect(decodedSmall.content, someText);
+      expect(decodedLarge.content, someText);
+    }
+  });
+
   test('unknown types should throw', () async {
     var registry = CodecRegistry();
     registry.registerCodec(TextCodec());