Let tornado handle the file downloads

octoprint/server.py

@@ -2,10 +2,9 @@
 __author__ = "Gina Häußge <osd@foosel.net>"
 __license__ = 'GNU Affero General Public License http://www.gnu.org/licenses/agpl.html'
 
-from werkzeug.utils import secure_filename
-from werkzeug.datastructures import Headers
+from werkzeug.utils import secure_filename, redirect
 from sockjs.tornado import SockJSRouter, SockJSConnection
-from flask import Flask, request, render_template, jsonify, send_from_directory, url_for, current_app, session, abort, make_response, Response
+from flask import Flask, request, render_template, jsonify, send_from_directory, url_for, current_app, session, abort, make_response
 from flask.ext.login import LoginManager, login_user, logout_user, login_required, current_user
 from flask.ext.principal import Principal, Permission, RoleNeed, Identity, identity_changed, AnonymousIdentity, identity_loaded, UserNeed
 
@@ -347,8 +346,7 @@ def readGcodeFiles():
 
 @app.route(BASEURL + "gcodefiles/<path:filename>", methods=["GET"])
 def readGcodeFile(filename):
-	path = os.path.join(settings().getBaseFolder("uploads"), filename)
-	return _streamBinaryFile(path, mimeType="text/plain", asAttachment=True)
+	return redirectToTornado(request, "/downloads/gcode/" + filename)
 
 @app.route(BASEURL + "gcodefiles/upload", methods=["POST"])
 @restricted_access
@@ -517,9 +515,7 @@ def getTimelapseData():
 
 @app.route(BASEURL + "timelapse/<filename>", methods=["GET"])
 def downloadTimelapse(filename):
-	if util.isAllowedFile(filename, set(["mpg"])):
-		path = os.path.join(settings().getBaseFolder("timelapse"), filename)
-		return _streamBinaryFile(path, mimeType="video/mpeg", asAttachment=True)
+	return redirectToTornado(request, "/downloads/timelapse/" + filename)
 
 @app.route(BASEURL + "timelapse/<filename>", methods=["DELETE"])
 @restricted_access
@@ -930,29 +926,15 @@ def load_user(id):
 		return userManager.findUser(id)
 	return users.DummyUser()
 
-def _streamBinaryFile(filename, mimeType=None, asAttachment=False, attachmentFilename=None):
-	if not os.path.exists(filename) or not os.path.isfile(filename):
-		return app.make_response(("No such file: %s" % filename, 404, []))
+def redirectToTornado(request, target):
+	requestUrl = request.url
+	appBaseUrl = requestUrl[:requestUrl.find(BASEURL)]
 
-	def generator(path, chunkSize=4096):
-		with open(path, "rb") as f:
-			while True:
-				data = f.read(chunkSize)
-				if not data:
-					break
-				yield data
-
-	headers = Headers()
-	if asAttachment:
-		if attachmentFilename is None:
-			attachmentFilename = os.path.basename(filename)
-		headers.add("Content-Disposition", "attachment", filename=attachmentFilename)
-	headers.add("Content-Length", os.stat(filename).st_size)
-
-	if mimeType is None:
-		mimeType = "application/octet-stream"
-
-	return Response(generator(filename), mimetype=mimeType, headers=headers)
+	redirectUrl = appBaseUrl + target
+	if "?" in requestUrl:
+		fragment = requestUrl[requestUrl.rfind("?"):]
+		redirectUrl += fragment
+	return redirect(redirectUrl)
 
 #~~ startup code
 class Server():
@@ -980,7 +962,7 @@ class Server():
 		from tornado.wsgi import WSGIContainer
 		from tornado.httpserver import HTTPServer
 		from tornado.ioloop import IOLoop
-		from tornado.web import Application, FallbackHandler
+		from tornado.web import Application, FallbackHandler, StaticFileHandler
 
 		debug = self._debug
 
@@ -1029,7 +1011,9 @@ class Server():
 		self._router = SockJSRouter(self._createSocketConnection, "/sockjs")
 
 		self._tornado_app = Application(self._router.urls + [
-			(".*", FallbackHandler, {"fallback": WSGIContainer(app)})
+			(r"/downloads/timelapse/([^/]*\.mpg)", StaticFileHandler, {"path": settings().getBaseFolder("timelapse")}),
+			(r"/downloads/gcode/([^/]*\.(gco|gcode))", StaticFileHandler, {"path": settings().getBaseFolder("uploads")}),
+			(r".*", FallbackHandler, {"fallback": WSGIContainer(app)})
 		])
 		self._server = HTTPServer(self._tornado_app)
 		self._server.listen(self._port, address=self._host)