From 16066479b4ba6544a2eab169ffcb00c0ff7d17bc Mon Sep 17 00:00:00 2001
From: Sam Denty <sam@samdenty.com>
Date: Sat, 8 Oct 2022 13:27:29 +0100
Subject: [PATCH] support decoding without mime type

---
 src/icon_info.rs      | 53 ++++++++++++++++++++++++++-----------------
 src/icon_size/ico.rs  | 12 ++++------
 src/icon_size/jpeg.rs |  7 ------
 src/icon_size/png.rs  |  8 +++----
 4 files changed, 41 insertions(+), 39 deletions(-)

diff --git a/src/icon_info.rs b/src/icon_info.rs
index e1894dc..b46c91d 100644
--- a/src/icon_info.rs
+++ b/src/icon_info.rs
@@ -8,12 +8,12 @@ use std::{
   cmp::Ordering,
   error::Error,
   fmt::{self, Display},
-  io::{self},
+  io,
 };
 
 #[derive(Serialize, Deserialize)]
 #[serde(rename_all = "lowercase")]
-enum IconType {
+enum IconKind {
   PNG,
   JPEG,
   ICO,
@@ -30,6 +30,30 @@ pub enum IconInfo {
 }
 
 impl IconInfo {
+  async fn decode<R: AsyncRead + Unpin>(
+    reader: &mut R,
+    kind: Option<IconKind>,
+  ) -> Result<IconInfo, Box<dyn Error>> {
+    let mut header = [0; 2];
+    reader.read_exact(&mut header).await?;
+
+    match (kind, header) {
+      (Some(IconKind::PNG), _) | (_, [0xC2, 0x89]) => {
+        let size = get_png_size(reader).await?;
+        Ok(IconInfo::PNG { size })
+      }
+      (Some(IconKind::ICO), _) | (_, [0x00, 0x00]) => {
+        let sizes = get_ico_sizes(reader).await?;
+        Ok(IconInfo::ICO { sizes })
+      }
+      (Some(IconKind::JPEG), _) | (_, [0xFF, 0xD8]) => {
+        let size = get_jpeg_size(reader).await?;
+        Ok(IconInfo::JPEG { size })
+      }
+      _ => Err("unknown icon type".into()),
+    }
+  }
+
   pub async fn load(
     url: Url,
     headers: HeaderMap,
@@ -85,7 +109,7 @@ impl IconInfo {
             size: *sizes.largest(),
           });
         }
-        IconType::PNG
+        Some(IconKind::PNG)
       }
 
       (mime::IMAGE, mime::JPEG) => {
@@ -94,7 +118,7 @@ impl IconInfo {
             size: *sizes.largest(),
           });
         }
-        IconType::JPEG
+        Some(IconKind::JPEG)
       }
 
       (mime::IMAGE, "x-icon") | (mime::IMAGE, "vnd.microsoft.icon") => {
@@ -102,28 +126,15 @@ impl IconInfo {
           return Ok(IconInfo::ICO { sizes });
         }
 
-        IconType::ICO
+        Some(IconKind::ICO)
       }
 
-      (mime::IMAGE, mime::SVG) => return Ok(IconInfo::SVG),
+      (mime::IMAGE, mime::SVG) | (mime::TEXT, mime::PLAIN) => return Ok(IconInfo::SVG),
 
-      _ => return Err(format!("unsupported mime type {}", mime).into()),
+      _ => None,
     };
 
-    Ok(match kind {
-      IconType::PNG => {
-        let size = get_png_sizes(&mut body).await?;
-        IconInfo::PNG { size }
-      }
-      IconType::ICO => {
-        let sizes = get_ico_sizes(&mut body).await?;
-        IconInfo::ICO { sizes }
-      }
-      IconType::JPEG => {
-        let size = get_jpeg_size(&mut body).await?;
-        IconInfo::JPEG { size }
-      }
-    })
+    IconInfo::decode(&mut body, kind).await
   }
 
   pub fn size(&self) -> Option<&IconSize> {
diff --git a/src/icon_size/ico.rs b/src/icon_size/ico.rs
index 3c7d64f..0437fc8 100644
--- a/src/icon_size/ico.rs
+++ b/src/icon_size/ico.rs
@@ -1,4 +1,4 @@
-use super::{png::get_png_sizes, IconSize, IconSizes};
+use super::{png::get_png_size, IconSize, IconSizes};
 use byteorder::{LittleEndian, ReadBytesExt};
 use futures::prelude::*;
 use std::{
@@ -13,16 +13,14 @@ const INDEX_SIZE: u16 = 16;
 pub async fn get_ico_sizes<R: AsyncRead + Unpin>(
   reader: &mut R,
 ) -> Result<IconSizes, Box<dyn Error>> {
-  let mut offset = 0;
-  let mut header = [0; 6];
+  let mut offset = 4;
+  let mut header = [0; 4];
   reader.read_exact(&mut header).await?;
-  offset += header.len();
   let mut header = Cursor::new(header);
 
-  let header_type = header.read_u16::<LittleEndian>()?;
   let icon_type = header.read_u16::<LittleEndian>()?;
 
-  if header_type != 0 || icon_type != ICO_TYPE {
+  if icon_type != ICO_TYPE {
     return Err("bad header".into());
   }
 
@@ -48,7 +46,7 @@ pub async fn get_ico_sizes<R: AsyncRead + Unpin>(
       reader.read_exact(&mut data).await?;
       offset += data.len();
 
-      let size = get_png_sizes(reader).await;
+      let size = get_png_size(reader).await;
       if let Ok(size) = size {
         sizes.push(size);
       }
diff --git a/src/icon_size/jpeg.rs b/src/icon_size/jpeg.rs
index 6c32faa..7bd22ff 100644
--- a/src/icon_size/jpeg.rs
+++ b/src/icon_size/jpeg.rs
@@ -7,13 +7,6 @@ use tokio_futures_byteorder::AsyncReadBytesExt;
 pub async fn get_jpeg_size<R: AsyncRead + Unpin>(
   reader: &mut R,
 ) -> Result<IconSize, Box<dyn Error>> {
-  let mut data = [0; 2];
-  reader.read_exact(&mut data).await?;
-  let data = &mut Cursor::new(data);
-
-  // first marker of the file MUST be 0xFFD8
-  assert_slice_eq!(data, 0, &[0xFF, 0xD8], "bad header");
-
   let mut marker = [0; 2];
   let mut depth = 0i32;
 
diff --git a/src/icon_size/png.rs b/src/icon_size/png.rs
index 8223b91..2f80831 100644
--- a/src/icon_size/png.rs
+++ b/src/icon_size/png.rs
@@ -3,15 +3,15 @@ use byteorder::{BigEndian, ReadBytesExt};
 use futures::prelude::*;
 use std::{error::Error, io::Cursor};
 
-pub async fn get_png_sizes<R: AsyncRead + Unpin>(
+pub async fn get_png_size<R: AsyncRead + Unpin>(
   reader: &mut R,
 ) -> Result<IconSize, Box<dyn Error>> {
-  let mut header = [0; 24];
+  let mut header = [0; 22];
   reader.read_exact(&mut header).await?;
   let header = &mut Cursor::new(header);
 
-  assert_slice_eq!(header, 0, b"\x89PNG\r\n\x1a\n", "bad header");
-  assert_slice_eq!(header, 12, b"IHDR", "bad header");
+  assert_slice_eq!(header, 0, b"PNG\r\n\x1a\n", "bad header");
+  assert_slice_eq!(header, 10, b"IHDR", "bad header");
 
   let width = header.read_u32::<BigEndian>()?;
   let height = header.read_u32::<BigEndian>()?;