Index: clangd/URI.cpp =================================================================== --- clangd/URI.cpp +++ clangd/URI.cpp @@ -11,8 +11,11 @@ #include "llvm/ADT/Twine.h" #include "llvm/Support/Error.h" #include "llvm/Support/Format.h" +#include "llvm/Support/FormatVariadic.h" #include "llvm/Support/Path.h" +#include #include +#include #include LLVM_INSTANTIATE_REGISTRY(clang::clangd::URISchemeRegistry) @@ -128,6 +131,17 @@ return Result; } +static bool IsValidScheme(llvm::StringRef Scheme) { + if (Scheme.empty()) + return false; + if (!std::isalpha(Scheme[0])) + return false; + return std::all_of(Scheme.begin() + 1, Scheme.end(), [](char C) { + return std::isalpha(C) || std::isdigit(C) || C == '+' || C == '.' || + C == '-'; + }); +} + } // namespace URI::URI(llvm::StringRef Scheme, llvm::StringRef Authority, @@ -158,9 +172,13 @@ llvm::StringRef Uri = OrigUri; auto Pos = Uri.find(':'); - if (Pos == 0 || Pos == llvm::StringRef::npos) + if (Pos == llvm::StringRef::npos) return make_string_error("Scheme must be provided in URI: " + OrigUri); - U.Scheme = percentDecode(Uri.substr(0, Pos)); + auto SchemeStr = Uri.substr(0, Pos); + U.Scheme = percentDecode(SchemeStr); + if (!IsValidScheme(U.Scheme)) + return make_string_error(llvm::formatv("Invalid scheme: {0} (decoded: {1})", + SchemeStr, U.Scheme)); Uri = Uri.substr(Pos + 1); if (Uri.consume_front("//")) { Pos = Uri.find('/'); Index: unittests/clangd/URITests.cpp =================================================================== --- unittests/clangd/URITests.cpp +++ unittests/clangd/URITests.cpp @@ -51,9 +51,9 @@ TEST(PercentEncodingTest, Decode) { EXPECT_EQ(parseOrDie("x:a/b/c").body(), "a/b/c"); - EXPECT_EQ(parseOrDie("%3a://%3a/%3").scheme(), ":"); - EXPECT_EQ(parseOrDie("%3a://%3a/%3").authority(), ":"); - EXPECT_EQ(parseOrDie("%3a://%3a/%3").body(), "/%3"); + EXPECT_EQ(parseOrDie("s%2b://%3a/%3").scheme(), "s+"); + EXPECT_EQ(parseOrDie("s%2b://%3a/%3").authority(), ":"); + EXPECT_EQ(parseOrDie("s%2b://%3a/%3").body(), "/%3"); EXPECT_EQ(parseOrDie("x:a%21b%3ac~").body(), "a!b:c~"); } @@ -132,6 +132,7 @@ // Empty. EXPECT_TRUE(FailedParse("")); EXPECT_TRUE(FailedParse(":/a/b/c")); + EXPECT_TRUE(FailedParse("\"/a/b/c\" IWYU pragma: abc")); } TEST(URITest, Resolve) {