Merge 8383d58236eb8937115ff277ce0d4dbce730643e into 64a629002d078e84bd6bbfb550c1a81a4aa0a8ac

lib/src/main/java/org/asamk/signal/manager/util/AttachmentUtils.java

@@ -1,28 +1,42 @@
 package org.asamk.signal.manager.util;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.whispersystems.signalservice.api.messages.SignalServiceAttachmentStream;
 import org.whispersystems.signalservice.api.push.exceptions.ResumeLocationInvalidException;
 import org.whispersystems.signalservice.api.util.StreamDetails;
 import org.whispersystems.signalservice.internal.push.http.ResumableUploadSpec;
 
+import javax.imageio.ImageIO;
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
 import java.util.Optional;
 import java.util.UUID;
 
 public class AttachmentUtils {
 
+    private static final Logger logger = LoggerFactory.getLogger(AttachmentUtils.class);
+
+    // Images are fully buffered in memory to probe their dimensions, so cap how large a file we'll do this for.
+    private static final long MAX_DIMENSION_PROBE_SIZE = 20 * 1024 * 1024;
+
     public static SignalServiceAttachmentStream createAttachmentStream(
             StreamDetails streamDetails,
             Optional<String> name,
             boolean voiceNote,
             ResumableUploadSpec resumableUploadSpec
-    ) throws ResumeLocationInvalidException {
+    ) throws ResumeLocationInvalidException, IOException {
         final var uploadTimestamp = System.currentTimeMillis();
+        final var probedStream = probeImageDimensions(streamDetails);
         return SignalServiceAttachmentStream.newStreamBuilder()
-                .withStream(streamDetails.getStream())
+                .withStream(probedStream.inputStream())
                 .withContentType(streamDetails.getContentType())
                 .withLength(streamDetails.getLength())
                 .withFileName(name.orElse(null))
                 .withVoiceNote(voiceNote)
+                .withWidth(probedStream.width())
+                .withHeight(probedStream.height())
                 .withUploadTimestamp(uploadTimestamp)
                 .withResumableUploadSpec(resumableUploadSpec)
                 .withUuid(UUID.randomUUID())
@@ -33,7 +47,36 @@ public class AttachmentUtils {
             StreamDetails streamDetails,
             Optional<String> name,
             ResumableUploadSpec resumableUploadSpec
-    ) throws ResumeLocationInvalidException {
+    ) throws ResumeLocationInvalidException, IOException {
         return createAttachmentStream(streamDetails, name, false, resumableUploadSpec);
     }
+
+    /**
+     * Reads the attachment's dimensions if it's an image, so recipients don't get a square-cropped thumbnail.
+     * Falls back to width/height 0 (today's behavior) if the content isn't an image, is too large to probe
+     * cheaply, or fails to parse.
+     */
+    private static ProbedStream probeImageDimensions(StreamDetails streamDetails) throws IOException {
+        final var contentType = streamDetails.getContentType();
+        final var length = streamDetails.getLength();
+        if (contentType == null || !contentType.startsWith("image/") || length <= 0 || length > MAX_DIMENSION_PROBE_SIZE) {
+            return new ProbedStream(streamDetails.getStream(), 0, 0);
+        }
+
+        final var bytes = streamDetails.getStream().readAllBytes();
+        var width = 0;
+        var height = 0;
+        try {
+            final var image = ImageIO.read(new ByteArrayInputStream(bytes));
+            if (image != null) {
+                width = image.getWidth();
+                height = image.getHeight();
+            }
+        } catch (IOException e) {
+            logger.debug("Failed to probe image dimensions, sending without width/height: {}", e.getMessage());
+        }
+        return new ProbedStream(new ByteArrayInputStream(bytes), width, height);
+    }
+
+    private record ProbedStream(InputStream inputStream, int width, int height) {}
 }

lib/src/test/java/org/asamk/signal/manager/util/AttachmentUtilsTest.java

@@ -0,0 +1,47 @@
+package org.asamk.signal.manager.util;
+
+import org.junit.jupiter.api.Test;
+import org.whispersystems.signalservice.api.util.StreamDetails;
+
+import javax.imageio.ImageIO;
+import java.awt.image.BufferedImage;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.util.Optional;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+class AttachmentUtilsTest {
+
+    @Test
+    public void createAttachmentStream_setsWidthAndHeightForImage() throws Exception {
+        final var imageBytes = pngBytes(37, 21);
+        final var streamDetails = new StreamDetails(new ByteArrayInputStream(imageBytes), "image/png", imageBytes.length);
+
+        final var attachment = AttachmentUtils.createAttachmentStream(streamDetails, Optional.of("meme.png"), null);
+
+        assertEquals(37, attachment.getWidth());
+        assertEquals(21, attachment.getHeight());
+        assertArrayEquals(imageBytes, attachment.getInputStream().readAllBytes());
+    }
+
+    @Test
+    public void createAttachmentStream_leavesWidthAndHeightZeroForNonImage() throws Exception {
+        final var bytes = "not an image".getBytes();
+        final var streamDetails = new StreamDetails(new ByteArrayInputStream(bytes), "application/octet-stream", bytes.length);
+
+        final var attachment = AttachmentUtils.createAttachmentStream(streamDetails, Optional.of("file.bin"), null);
+
+        assertEquals(0, attachment.getWidth());
+        assertEquals(0, attachment.getHeight());
+        assertArrayEquals(bytes, attachment.getInputStream().readAllBytes());
+    }
+
+    private static byte[] pngBytes(final int width, final int height) throws Exception {
+        final var image = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
+        final var out = new ByteArrayOutputStream();
+        ImageIO.write(image, "png", out);
+        return out.toByteArray();
+    }
+}