From fd7fde0518562d60a7a85679715e2b15ee464d37 Mon Sep 17 00:00:00 2001
From: Neng Lu <nlu@twitter.com>
Date: Tue, 18 Sep 2018 11:58:21 -0700
Subject: [PATCH] heron shell access directory

---
 heron/shell/src/python/handlers/browsehandler.py    | 4 +++-
 heron/shell/src/python/handlers/downloadhandler.py  | 2 +-
 heron/shell/src/python/handlers/filedatahandler.py  | 4 +++-
 heron/shell/src/python/handlers/filehandler.py      | 4 +++-
 heron/shell/src/python/handlers/filestatshandler.py | 5 +++--
 heron/shell/src/python/utils.py                     | 8 ++++++++
 6 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/heron/shell/src/python/handlers/browsehandler.py b/heron/shell/src/python/handlers/browsehandler.py
index 80606e6f010..d0f050a99e4 100644
--- a/heron/shell/src/python/handlers/browsehandler.py
+++ b/heron/shell/src/python/handlers/browsehandler.py
@@ -34,11 +34,13 @@ def get(self, path):
     ''' get method '''
     if not path:
       path = "."
-    if path.startswith("/"):
+
+    if not utils.check_path(path):
       self.write("Only relative paths are allowed")
       self.set_status(403)
       self.finish()
       return
+
     t = Template(utils.get_asset("browse.html"))
     args = dict(
         path=path,
diff --git a/heron/shell/src/python/handlers/downloadhandler.py b/heron/shell/src/python/handlers/downloadhandler.py
index ed30ade2f28..dcf652b5070 100644
--- a/heron/shell/src/python/handlers/downloadhandler.py
+++ b/heron/shell/src/python/handlers/downloadhandler.py
@@ -37,7 +37,7 @@ def get(self, path):
     self.connection_closed = False
 
     self.set_header("Content-Disposition", "attachment")
-    if path.startswith("/"):
+    if not utils.check_path(path):
       self.write("Only relative paths are allowed")
       self.set_status(403)
       self.finish()
diff --git a/heron/shell/src/python/handlers/filedatahandler.py b/heron/shell/src/python/handlers/filedatahandler.py
index 271d152cef9..9c57a419c9b 100644
--- a/heron/shell/src/python/handlers/filedatahandler.py
+++ b/heron/shell/src/python/handlers/filedatahandler.py
@@ -32,11 +32,13 @@ def get(self, path):
     """ get method """
     if path is None:
       return {}
-    if path.startswith("/"):
+
+    if not utils.check_path(path):
       self.write("Only relative paths are allowed")
       self.set_status(403)
       self.finish()
       return
+
     offset = self.get_argument("offset", default=-1)
     length = self.get_argument("length", default=-1)
     if not os.path.isfile(path):
diff --git a/heron/shell/src/python/handlers/filehandler.py b/heron/shell/src/python/handlers/filehandler.py
index 51abaac8945..bb35ddf2999 100644
--- a/heron/shell/src/python/handlers/filehandler.py
+++ b/heron/shell/src/python/handlers/filehandler.py
@@ -35,11 +35,13 @@ def get(self, path):
       self.write("No such file")
       self.finish()
       return
-    if path.startswith("/"):
+
+    if not utils.check_path(path):
       self.write("Only relative paths are allowed")
       self.set_status(403)
       self.finish()
       return
+
     args = dict(
         filename=path,
         jquery=utils.get_asset("jquery.js"),
diff --git a/heron/shell/src/python/handlers/filestatshandler.py b/heron/shell/src/python/handlers/filestatshandler.py
index 084bbfd7d56..2f19240c1dd 100644
--- a/heron/shell/src/python/handlers/filestatshandler.py
+++ b/heron/shell/src/python/handlers/filestatshandler.py
@@ -38,11 +38,12 @@ def get(self, path):
     # of the dir that heron-shell is running in. This ensures
     # sandboxing. So we don't allow absolute paths and parent
     # accessing.
-    if path.startswith("/") or ".." in path:
-      self.write("Only relative paths inside job dir are allowed")
+    if not utils.check_path(path):
+      self.write("Only relative paths are allowed")
       self.set_status(403)
       self.finish()
       return
+
     listing = utils.get_listing(path)
     file_stats = {}
     for fn in listing:
diff --git a/heron/shell/src/python/utils.py b/heron/shell/src/python/utils.py
index 6070227f12a..13eb8290e84 100644
--- a/heron/shell/src/python/utils.py
+++ b/heron/shell/src/python/utils.py
@@ -194,3 +194,11 @@ def get_container_id(instance_id):
 def get_asset(asset_name):
   ''' get assset '''
   return pkgutil.get_data("heron.shell", os.path.join("assets", asset_name))
+
+def check_path(path):
+  """
+  file path should be a relative path without ".." in it
+  :param path: file path
+  :return: true if the path is relative and doesn't contain ".."
+  """
+  return not path.startswith("/") and ".." not in path