Fix unrestricted reads from keeper

This commit is contained in:
Raúl Marín 2024-02-26 19:38:00 +01:00
parent 1a9ee8c318
commit 8a161bdc0a
2 changed files with 28 additions and 16 deletions

View File

@ -169,13 +169,13 @@ public:
/// Type of path to be fetched
enum class ZkPathType
{
Exact, /// Fetch all nodes under this path
Prefix, /// Fetch all nodes starting with this prefix, recursively (multiple paths may match prefix)
Recurse, /// Fatch all nodes under this path, recursively
Exact, /// Fetch all nodes under this path
Prefix, /// Fetch all nodes starting with this prefix, recursively (multiple paths may match prefix)
Recurse, /// Fetch all nodes under this path, recursively
};
/// List of paths to be feched from zookeeper
using Paths = std::deque<std::pair<String, ZkPathType>>;
/// List of paths to be fetched from zookeeper
using Paths = std::unordered_map<String, ZkPathType>;
class ReadFromSystemZooKeeper final : public SourceStepWithFilter
{
@ -222,6 +222,7 @@ private:
ContextPtr context;
ZooKeeperWithFaultInjection::Ptr zookeeper;
bool started = false;
std::unordered_set<String> visited;
};
@ -365,7 +366,8 @@ static void extractPathImpl(const ActionsDAG::Node & node, Paths & res, ContextP
size_t size = values->size();
for (size_t row = 0; row < size; ++row)
res.emplace_back(values->getDataAt(row).toString(), ZkPathType::Exact);
/// Only inserted if the key doesn't exists already
res.insert({values->getDataAt(row).toString(), ZkPathType::Exact});
}
else if (function_name == "equals")
{
@ -385,7 +387,8 @@ static void extractPathImpl(const ActionsDAG::Node & node, Paths & res, ContextP
if (value->column->size() != 1)
return;
res.emplace_back(value->column->getDataAt(0).toString(), ZkPathType::Exact);
/// Only inserted if the key doesn't exists already
res.insert({value->column->getDataAt(0).toString(), ZkPathType::Exact});
}
else if (allow_unrestricted && function_name == "like")
{
@ -404,7 +407,7 @@ static void extractPathImpl(const ActionsDAG::Node & node, Paths & res, ContextP
String pattern = value->column->getDataAt(0).toString();
bool has_metasymbol = false;
String prefix; // pattern prefix before the first metasymbol occurrence
String prefix{}; // pattern prefix before the first metasymbol occurrence
for (size_t i = 0; i < pattern.size(); i++)
{
char c = pattern[i];
@ -430,7 +433,7 @@ static void extractPathImpl(const ActionsDAG::Node & node, Paths & res, ContextP
prefix.append(1, c);
}
res.emplace_back(prefix, has_metasymbol ? ZkPathType::Prefix : ZkPathType::Exact);
res.insert_or_assign(prefix, has_metasymbol ? ZkPathType::Prefix : ZkPathType::Exact);
}
}
@ -443,8 +446,17 @@ static Paths extractPath(const ActionsDAG::NodeRawConstPtrs & filter_nodes, Cont
for (const auto * node : filter_nodes)
extractPathImpl(*node, res, context, allow_unrestricted);
auto node1 = res.find("/");
auto node2 = res.find("");
if ((node1 != res.end() && node1->second != ZkPathType::Exact) || (node2 != res.end() && node2->second != ZkPathType::Exact))
{
/// If we are already searching everything recursively, remove all other nodes
res.clear();
res.insert({"/", ZkPathType::Recurse});
}
if (res.empty() && allow_unrestricted)
res.emplace_back("/", ZkPathType::Recurse);
res.insert({"/", ZkPathType::Recurse});
return res;
}
@ -510,7 +522,6 @@ Chunk SystemZooKeeperSource::generate()
String path_part;
};
std::vector<ListTask> list_tasks;
std::unordered_set<String> added;
while (!paths.empty())
{
if (query_status)
@ -530,8 +541,9 @@ Chunk SystemZooKeeperSource::generate()
std::vector<String> paths_to_list;
while (!paths.empty() && static_cast<Int64>(list_tasks.size()) < max_inflight_requests)
{
auto [path, path_type] = std::move(paths.front());
paths.pop_front();
auto node = paths.extract(paths.begin());
auto & path = node.key();
auto & path_type = node.mapped();
ListTask task;
task.path = path;
@ -612,7 +624,7 @@ Chunk SystemZooKeeperSource::generate()
// Deduplication
String key = list_task.path_part + '/' + get_task.node;
if (auto [it, inserted] = added.emplace(key); !inserted)
if (auto [it, inserted] = visited.emplace(key); !inserted)
continue;
const Coordination::Stat & stat = res.stat;
@ -638,7 +650,7 @@ Chunk SystemZooKeeperSource::generate()
if (list_task.path_type != ZkPathType::Exact && res.stat.numChildren > 0)
{
paths.emplace_back(key, ZkPathType::Recurse);
paths.insert_or_assign(key, ZkPathType::Recurse);
}
}
}

View File

@ -13,7 +13,7 @@ ${CLICKHOUSE_CLIENT} -n --query="CREATE TABLE sample_table (
)
ENGINE ReplicatedMergeTree('/clickhouse/$CLICKHOUSE_TEST_ZOOKEEPER_PREFIX/02221_system_zookeeper_unrestricted_like', '1')
ORDER BY tuple();
DROP TABLE IF EXISTS sample_table;"
DROP TABLE IF EXISTS sample_table SYNC;"
${CLICKHOUSE_CLIENT} -n --query "CREATE TABLE sample_table_2 (