From 3143921a1aabc1a54d89c5690bcf72e789d4718c Mon Sep 17 00:00:00 2001
From: dvkt <c@dvkt.io>
Date: Thu, 26 Dec 2019 21:48:29 -0800
Subject: [PATCH] check file type

---
 Cargo.lock    | 10 ++++++++++
 Cargo.toml    |  3 ++-
 src/server.rs | 52 ++++++++++++++++++++++++++++++++++++---------------
 3 files changed, 49 insertions(+), 16 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index b20f49c..cb70728 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -47,6 +47,14 @@ name = "cfg-if"
 version = "0.1.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 
+[[package]]
+name = "content_inspector"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+dependencies = [
+ "memchr 2.2.1 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
 [[package]]
 name = "crossbeam-channel"
 version = "0.4.0"
@@ -257,6 +265,7 @@ name = "phd"
 version = "0.1.0"
 dependencies = [
  "async-std 1.4.0 (registry+https://github.com/rust-lang/crates.io-index)",
+ "content_inspector 0.2.4 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
 [[package]]
@@ -344,6 +353,7 @@ dependencies = [
 "checksum autocfg 0.1.7 (registry+https://github.com/rust-lang/crates.io-index)" = "1d49d90015b3c36167a20fe2810c5cd875ad504b39cff3d4eae7977e6b7c1cb2"
 "checksum bitflags 1.2.1 (registry+https://github.com/rust-lang/crates.io-index)" = "cf1de2fe8c75bc145a2f577add951f8134889b4795d47466a54a5c846d691693"
 "checksum cfg-if 0.1.10 (registry+https://github.com/rust-lang/crates.io-index)" = "4785bdd1c96b2a846b2bd7cc02e86b6b3dbf14e7e53446c4f54c92a361040822"
+"checksum content_inspector 0.2.4 (registry+https://github.com/rust-lang/crates.io-index)" = "b7bda66e858c683005a53a9a60c69a4aca7eeaa45d124526e389f7aec8e62f38"
 "checksum crossbeam-channel 0.4.0 (registry+https://github.com/rust-lang/crates.io-index)" = "acec9a3b0b3559f15aee4f90746c4e5e293b701c0f7d3925d24e01645267b68c"
 "checksum crossbeam-deque 0.7.2 (registry+https://github.com/rust-lang/crates.io-index)" = "c3aa945d63861bfe624b55d153a39684da1e8c0bc8fba932f7ee3a3c16cea3ca"
 "checksum crossbeam-epoch 0.8.0 (registry+https://github.com/rust-lang/crates.io-index)" = "5064ebdbf05ce3cb95e45c8b086f72263f4166b29b97f6baff7ef7fe047b55ac"
diff --git a/Cargo.toml b/Cargo.toml
index 1577afe..c5f1db3 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -5,4 +5,5 @@ authors = ["dvkt <c@dvkt.io>"]
 edition = "2018"
 
 [dependencies]
-async-std = "1.4.0"
\ No newline at end of file
+async-std = "1.4.0"
+content_inspector = "0.2.4"
\ No newline at end of file
diff --git a/src/server.rs b/src/server.rs
index e31eaf6..483f324 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -5,10 +5,13 @@ use async_std::{
     prelude::*,
     task,
 };
+use content_inspector::{inspect, ContentType};
 use std::path::PathBuf;
 
 type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;
 
+const MAX_PEEK_SIZE: usize = 1024;
+
 pub fn start(addr: &str, root: &str) -> Result<()> {
     task::block_on(async {
         let listener = TcpListener::bind(addr).await?;
@@ -41,22 +44,41 @@ async fn respond(stream: &mut TcpStream, selector: &str, root: &str) -> Result<(
     let mut dir = fs::read_dir(path.clone()).await?;
 
     while let Some(Ok(entry)) = dir.next().await {
-        if let Ok(metadata) = entry.metadata().await {
-            let file_type = if metadata.is_file() {
-                '0'
-            } else if metadata.is_dir() {
-                '1'
-            } else {
-                '3'
-            };
-            response.push_str(&format!(
-                "{}{}\t{}\tlocalhost\t7070\r\n",
-                file_type,
-                entry.file_name().into_string().unwrap(),
-                entry.path().to_string_lossy(),
-            ));
-        }
+        let file_type = file_type(&entry).await;
+        response.push_str(&format!(
+            "{}{}\t{}\tlocalhost\t7070\r\n",
+            file_type,
+            entry.file_name().into_string().unwrap(),
+            entry.path().to_string_lossy(),
+        ));
     }
     stream.write_all(response.as_bytes()).await?;
     Ok(())
 }
+
+async fn file_type(dir: &fs::DirEntry) -> char {
+    if let Ok(metadata) = dir.metadata().await {
+        if metadata.is_file() {
+            if let Ok(file) = fs::File::open(&dir.path()).await {
+                let mut buffer: Vec<u8> = vec![];
+                let _ = file
+                    .take(MAX_PEEK_SIZE as u64)
+                    .read_to_end(&mut buffer)
+                    .await;
+                if content_inspector::inspect(&buffer).is_binary() {
+                    '9'
+                } else {
+                    '0'
+                }
+            } else {
+                '9'
+            }
+        } else if metadata.is_dir() {
+            '1'
+        } else {
+            '3'
+        }
+    } else {
+        '3'
+    }
+}