From 74cba0ccae8a6103bc1e2aefc0ebf5f0f30587af Mon Sep 17 00:00:00 2001 From: Manjunath Kudlur Date: Thu, 7 Apr 2016 19:59:40 -0800 Subject: [PATCH] Parse the URI from scheme and allow only [a-zA-Z][0-9a-zA-Z.]* Change: 119338733 --- tensorflow/core/platform/env_test.cc | 11 +++++++++++ tensorflow/core/platform/file_system.cc | 9 ++++++++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/tensorflow/core/platform/env_test.cc b/tensorflow/core/platform/env_test.cc index 6bf69f988c0..35b231d7694 100644 --- a/tensorflow/core/platform/env_test.cc +++ b/tensorflow/core/platform/env_test.cc @@ -116,4 +116,15 @@ TEST(EnvTest, IPFS) { } } +TEST(EnvTest, GetSchemeForURI) { + EXPECT_EQ(GetSchemeFromURI("http://foo"), "http"); + EXPECT_EQ(GetSchemeFromURI("/encrypted/://foo"), ""); + EXPECT_EQ(GetSchemeFromURI("/usr/local/foo"), ""); + EXPECT_EQ(GetSchemeFromURI("file:///usr/local/foo"), "file"); + EXPECT_EQ(GetSchemeFromURI("local.file:///usr/local/foo"), "local.file"); + EXPECT_EQ(GetSchemeFromURI("a-b:///foo"), ""); + EXPECT_EQ(GetSchemeFromURI(":///foo"), ""); + EXPECT_EQ(GetSchemeFromURI("9dfd:///foo"), ""); +} + } // namespace tensorflow diff --git a/tensorflow/core/platform/file_system.cc b/tensorflow/core/platform/file_system.cc index 4f19bff50b5..9943b32b7eb 100644 --- a/tensorflow/core/platform/file_system.cc +++ b/tensorflow/core/platform/file_system.cc @@ -17,6 +17,7 @@ limitations under the License. #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/gtl/map_util.h" #include "tensorflow/core/lib/gtl/stl_util.h" +#include "tensorflow/core/lib/strings/scanner.h" #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/platform/protobuf.h" @@ -66,7 +67,13 @@ FileSystem* FileSystemRegistryImpl::Lookup(const string& scheme) { string GetSchemeFromURI(const string& name) { auto colon_loc = name.find(":"); - if (colon_loc != string::npos) { + // Make sure scheme matches [a-zA-Z][0-9a-zA-Z.]* + // TODO(keveman): Allow "+" and "-" in the scheme. + if (colon_loc != string::npos && + strings::Scanner(StringPiece(name.data(), colon_loc)) + .One(strings::Scanner::LETTER) + .Many(strings::Scanner::LETTER_DIGIT_DOT) + .GetResult()) { return name.substr(0, colon_loc); } return "";