Files
santa/Source/common/SNTPrefixTree.h
Kent Ma 1c04c3a257 Remove code guarded by #ifdef kernel macros (#752)
* Remove code guarded by #ifdef kernel macros
2022-03-15 14:38:40 -04:00

92 lines
2.9 KiB
C++

/// Copyright 2018 Google Inc. All rights reserved.
///
/// Licensed under the Apache License, Version 2.0 (the "License");
/// you may not use this file except in compliance with the License.
/// You may obtain a copy of the License at
///
/// http://www.apache.org/licenses/LICENSE-2.0
///
/// Unless required by applicable law or agreed to in writing, software
/// distributed under the License is distributed on an "AS IS" BASIS,
/// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
/// See the License for the specific language governing permissions and
/// limitations under the License.
#ifndef SANTA__SANTA_DRIVER__SANTAPREFIXTREE_H
#define SANTA__SANTA_DRIVER__SANTAPREFIXTREE_H
#include <IOKit/IOReturn.h>
#include <sys/param.h>
// Support for unit testing.
#include <pthread.h>
#include <stdint.h>
#include <mutex>
///
/// SantaPrefixTree is a simple prefix tree implementation.
/// Operations are thread safe.
///
class SNTPrefixTree {
public:
// Add a prefix to the tree.
// Optionally pass node_count to get the number of nodes after the add.
IOReturn AddPrefix(const char *, uint64_t *node_count = nullptr);
// Check if the tree has a prefix for string.
bool HasPrefix(const char *string);
// Reset the tree.
void Reset();
SNTPrefixTree(uint32_t max_nodes = kDefaultMaxNodes);
~SNTPrefixTree();
private:
///
/// SantaPrefixNode is a wrapper class that represents one byte.
/// 1 node can represent a whole ASCII character.
/// For example a pointer to the 'A' node will be stored at children[0x41].
/// It takes 1-4 nodes to represent a UTF-8 encoded Unicode character.
///
/// The path for "/🤘" would look like this:
/// children[0x2f] -> children[0xf0] -> children[0x9f] -> children[0xa4]
/// -> children[0x98]
///
/// The path for "/dev" is:
/// children[0x2f] -> children[0x64] -> children[0x65] -> children[0x76]
///
/// Lookups of children are O(1).
///
/// Having the nodes represented by a smaller width, such as a nibble (1/2
/// byte), would drastically decrease the memory footprint but would double
/// required dereferences.
///
/// TODO(bur): Potentially convert this into a full on radix tree.
///
class SantaPrefixNode {
public:
bool isPrefix;
SantaPrefixNode *children[256];
};
// PruneNode will remove the passed in node from the tree.
// The passed in node and all subnodes will be deleted.
// It is the caller's responsibility to reset the pointer to this node (held
// by the parent). If the tree is in use grab the exclusive lock.
void PruneNode(SantaPrefixNode *);
SantaPrefixNode *root_;
// Each node takes up ~2k, assuming MAXPATHLEN is 1024 max out at ~2MB.
static const uint32_t kDefaultMaxNodes = MAXPATHLEN;
uint32_t max_nodes_;
uint32_t node_count_;
pthread_rwlock_t spt_lock_;
std::mutex *spt_add_lock_;
};
#endif /* SANTA__SANTA_DRIVER__SANTAPREFIXTREE_H */