From 777cca245eab34ccf90e3d95b687854e408e21f6 Mon Sep 17 00:00:00 2001
From: Antoine Gersant <antoine.gersant@lesforges.org>
Date: Sun, 28 Oct 2018 14:00:25 -0700
Subject: [PATCH] Added support for range header when serving files

---
 src/rocket_api.rs |   9 ++--
 src/serve.rs      | 125 ++++++++++++++++++++++++++++++++++++----------
 2 files changed, 105 insertions(+), 29 deletions(-)

diff --git a/src/rocket_api.rs b/src/rocket_api.rs
index f0f923a..ff05a71 100644
--- a/src/rocket_api.rs
+++ b/src/rocket_api.rs
@@ -1,8 +1,8 @@
 use rocket::http::{Cookie, Cookies, Status};
 use rocket::request::{self, FromRequest, Request};
-use rocket::response::NamedFile;
 use rocket::{Outcome, State};
 use rocket_contrib::json::Json;
+use std::fs::File;
 use std::path::PathBuf;
 use std::ops::Deref;
 use std::sync::Arc;
@@ -11,6 +11,7 @@ use config::{self, Config};
 use db::DB;
 use errors;
 use index;
+use serve;
 use thumbnails;
 use user;
 use utils;
@@ -227,7 +228,7 @@ fn search(db: State<DB>, _auth: Auth, query: String) -> Result<Json<Vec<index::C
 }
 
 #[get("/serve/<path..>")]
-fn serve(db: State<DB>, _auth: Auth, path: PathBuf) -> Result<NamedFile, errors::Error> {
+fn serve(db: State<DB>, _auth: Auth, path: PathBuf) -> Result<serve::RangeResponder<File>, errors::Error> {
 	let db: &DB = db.deref();
 	let vfs = db.get_vfs()?;
 	let real_path = vfs.virtual_to_real(&path)?;
@@ -238,6 +239,6 @@ fn serve(db: State<DB>, _auth: Auth, path: PathBuf) -> Result<NamedFile, errors:
 		real_path
 	};
 
-	let serving = NamedFile::open(&serve_path)?;
-	Ok(serving)
+	let file = File::open(serve_path)?;
+	Ok(serve::RangeResponder::new(file))
 }
diff --git a/src/serve.rs b/src/serve.rs
index 4681ceb..d8d41a0 100644
--- a/src/serve.rs
+++ b/src/serve.rs
@@ -6,10 +6,13 @@ use iron::modifiers::Header;
 use iron::prelude::*;
 use iron::response::WriteBody;
 use iron::status::{self, Status};
+use rocket;
+use rocket::response::{self, Responder};
 use std::cmp;
 use std::fs::{self, File};
 use std::io::{self, Read, Seek, SeekFrom, Write};
 use std::path::Path;
+use std::str::FromStr;
 
 use errors::{Error, ErrorKind};
 
@@ -98,33 +101,10 @@ impl PartialFile {
 
 impl Modifier<Response> for PartialFile {
 	fn modify(self, res: &mut Response) {
-		use self::PartialFileRange::*;
 		let metadata: Option<_> = self.file.metadata().ok();
 		let file_length: Option<u64> = metadata.map(|m| m.len());
-		let range: Option<(u64, u64)> = match (self.range, file_length) {
-			(FromTo(from, to), Some(file_length)) => {
-				if from <= to && from < file_length {
-					Some((from, cmp::min(to, file_length - 1)))
-				} else {
-					None
-				}
-			}
-			(AllFrom(from), Some(file_length)) => {
-				if from < file_length {
-					Some((from, file_length - 1))
-				} else {
-					None
-				}
-			}
-			(Last(last), Some(file_length)) => {
-				if last < file_length {
-					Some((file_length - last, file_length - 1))
-				} else {
-					Some((0, file_length - 1))
-				}
-			}
-			(_, None) => None,
-		};
+		let range: Option<(u64, u64)> = truncate_range(&self.range, &file_length);
+
 		if let Some(range) = range {
 			let content_range = ContentRange(ContentRangeSpec::Bytes {
 				range: Some(range),
@@ -165,3 +145,98 @@ impl WriteBody for PartialContentBody {
 		io::copy(&mut limiter, res).map(|_| ())
 	}
 }
+
+pub struct RangeResponder<R> {
+	original: R,
+}
+
+impl<'r, R: Responder<'r>> RangeResponder<R> {
+	pub fn new(original: R) -> RangeResponder<R> {
+		RangeResponder{ original }
+	}
+
+	fn ignore_range(self, request: &rocket::request::Request) -> response::Result<'r> {
+		let mut response = self.original.respond_to(request)?;
+		response.set_status(rocket::http::Status::RangeNotSatisfiable);
+		Ok(response)
+	}
+}
+
+fn truncate_range(range: &PartialFileRange, file_length: &Option<u64>) -> Option<(u64, u64)> {
+	use self::PartialFileRange::*;
+
+	match (range, file_length) {
+		(FromTo(from, to), Some(file_length)) => {
+			if from <= to && from < file_length {
+				Some((*from, cmp::min(*to, file_length - 1)))
+			} else {
+				None
+			}
+		}
+		(AllFrom(from), Some(file_length)) => {
+			if from < file_length {
+				Some((*from, file_length - 1))
+			} else {
+				None
+			}
+		}
+		(Last(last), Some(file_length)) => {
+			if last < file_length {
+				Some((file_length - last, file_length - 1))
+			} else {
+				Some((0, file_length - 1))
+			}
+		}
+		(_, None) => None,
+	}
+}
+
+impl<'r> Responder<'r> for RangeResponder<File> {
+
+	fn respond_to(mut self, request: &rocket::request::Request) -> response::Result<'r> {
+
+		use rocket::http::hyper::header::*;
+
+		let range_header = request.headers().get_one("Range");
+		let range_header = match range_header {
+			None => return Ok(self.original.respond_to(request)?),
+			Some(h) => h,
+		};
+
+		let vec_range = match Range::from_str(range_header) {
+			Ok(Range::Bytes(v)) => v,
+			_ => return self.ignore_range(request),
+		};
+
+		let partial_file_range = match vec_range.into_iter().next() {
+			None => PartialFileRange::AllFrom(0),
+			Some(byte_range) => PartialFileRange::from(byte_range),
+		};
+
+		let metadata: Option<_> = self.original.metadata().ok();
+		let file_length: Option<u64> = metadata.map(|m| m.len());
+		let range: Option<(u64, u64)> = truncate_range(&partial_file_range, &file_length);
+
+		if let Some((from, to)) = range {
+			let content_range = ContentRange(ContentRangeSpec::Bytes {
+				range: range,
+				instance_length: file_length,
+			});
+			let content_len = to - from + 1;
+
+			match self.original.seek(SeekFrom::Start(from)) {
+				Ok(_) => (),
+				Err(_) => return Err(rocket::http::Status::InternalServerError),
+			}
+			let partial_original = self.original.take(content_len).into_inner();
+			let mut response = partial_original.respond_to(request)?;
+			response.set_header(ContentLength(content_len));
+			response.set_header(content_range);
+			response.set_status(rocket::http::Status::PartialContent);
+
+			Ok(response)
+		} else {
+			self.ignore_range(request)
+		}
+    }
+}