From 6f83a175da61702fc7cd245cc320b235121d3fcb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gina=20H=C3=A4u=C3=9Fge?= <osd@foosel.net>
Date: Thu, 26 Nov 2015 19:51:57 +0100
Subject: [PATCH] Replaced way too lax SourceRootFileLoader with
 SelectedFileLoader

---
 src/octoprint/server/__init__.py | 26 ++------------------
 src/octoprint/util/jinja.py      | 41 ++++++++++++++++++++++++++++++--
 2 files changed, 41 insertions(+), 26 deletions(-)

diff --git a/src/octoprint/server/__init__.py b/src/octoprint/server/__init__.py
index c5ef4b60..864e156b 100644
--- a/src/octoprint/server/__init__.py
+++ b/src/octoprint/server/__init__.py
@@ -635,30 +635,8 @@ class Server():
 		if octoprint.util.is_running_from_source():
 			root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
 			allowed = ["AUTHORS.md", "CHANGELOG.md", "THIRDPARTYLICENSES.md"]
-
-			class SourceRootFilesystemLoader(jinja2.FileSystemLoader):
-				def __init__(self, template_filter, prefix, *args, **kwargs):
-					jinja2.FileSystemLoader.__init__(self, *args, **kwargs)
-					self._filter = template_filter
-					if not prefix.endswith("/"):
-						prefix += "/"
-					self._prefix = prefix
-
-				def get_source(self, environment, template):
-					if not template.startswith(self._prefix):
-						raise jinja2.TemplateNotFound(template)
-
-					template = template[len(self._prefix):]
-					if not self._filter(template):
-						raise jinja2.TemplateNotFound(template)
-
-					return jinja2.FileSystemLoader.get_source(self, environment, template)
-
-				def list_templates(self):
-					templates = jinja2.FileSystemLoader.list_templates(self)
-					return map(lambda t: self._prefix + t, filter(self._filter, templates))
-
-			loaders.append(SourceRootFilesystemLoader(lambda t: t in allowed, "_data/", root))
+			files = {name: os.path.join(root, name) for name in allowed}
+			loaders.append(octoprint.util.jinja.SelectedFilesLoader(files))
 
 		jinja_loader = jinja2.ChoiceLoader(loaders)
 		app.jinja_loader = jinja_loader
diff --git a/src/octoprint/util/jinja.py b/src/octoprint/util/jinja.py
index 1c8539a8..34954b88 100644
--- a/src/octoprint/util/jinja.py
+++ b/src/octoprint/util/jinja.py
@@ -10,7 +10,7 @@ import os
 from jinja2 import nodes
 from jinja2.ext import Extension
 from jinja2.loaders import FileSystemLoader, PrefixLoader, ChoiceLoader, \
-	TemplateNotFound, split_template_path
+	BaseLoader, TemplateNotFound, split_template_path
 
 class FilteredFileSystemLoader(FileSystemLoader):
 	"""
@@ -54,6 +54,39 @@ class FilteredFileSystemLoader(FileSystemLoader):
 		return all(filter_results)
 
 
+class SelectedFilesLoader(BaseLoader):
+	def __init__(self, files, encoding="utf-8"):
+		self.files = files
+		self.encoding = encoding
+
+	def get_source(self, environment, template):
+		if not template in self.files:
+			raise TemplateNotFound(template)
+
+		from jinja2.loaders import open_if_exists
+
+		path = self.files[template]
+		f = open_if_exists(path)
+		if f is None:
+			raise TemplateNotFound(template)
+		try:
+			contents = f.read().decode(self.encoding)
+		finally:
+			f.close()
+
+		mtime = os.path.getmtime(path)
+
+		def uptodate():
+			try:
+				return os.path.getmtime(path) == mtime
+			except OSError:
+				return False
+		return contents, path, uptodate
+
+	def list_templates(self):
+		return self.files.keys()
+
+
 class ExceptionHandlerExtension(Extension):
 	tags = {"try"}
 
@@ -115,7 +148,11 @@ def get_all_template_paths(loader):
 		return files
 
 	def collect_templates_for_loader(loader):
-		if isinstance(loader, FilteredFileSystemLoader):
+		if isinstance(loader, SelectedFilesLoader):
+			import copy
+			return copy.copy(loader.files.values())
+
+		elif isinstance(loader, FilteredFileSystemLoader):
 			result = []
 			for folder in loader.searchpath:
 				result += walk_folder(folder)