santa/Source/common/SNTPrefixTree.cc

/// 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.

#include "Source/common/SNTPrefixTree.h"

#include <string.h>

#include <mutex>

#define LOGD(format, ...)  // NOP
#define LOGE(format, ...)  // NOP

#define lck_rw_lock_shared(l) pthread_rwlock_rdlock(&l)
#define lck_rw_unlock_shared(l) pthread_rwlock_unlock(&l)
#define lck_rw_lock_exclusive(l) pthread_rwlock_wrlock(&l)
#define lck_rw_unlock_exclusive(l) pthread_rwlock_unlock(&l)

#define lck_rw_lock_shared_to_exclusive(l) \
  ({                                       \
    pthread_rwlock_unlock(&l);             \
    false;                                 \
  })
#define lck_rw_lock_exclusive_to_shared(l) \
  ({                                       \
    pthread_rwlock_unlock(&l);             \
    pthread_rwlock_rdlock(&l);             \
  })

#define lck_mtx_lock(l) l->lock()
#define lck_mtx_unlock(l) l->unlock()

SNTPrefixTree::SNTPrefixTree(uint32_t max_nodes) {
  root_ = new SantaPrefixNode();
  node_count_ = 0;
  max_nodes_ = max_nodes;

  pthread_rwlock_init(&spt_lock_, nullptr);
  spt_add_lock_ = new std::mutex;
}

IOReturn SNTPrefixTree::AddPrefix(const char *prefix, uint64_t *node_count) {
  // Serialize requests to AddPrefix. Otherwise one AddPrefix thread could
  // overwrite whole branches of another. HasPrefix is still free to read the
  // tree, until AddPrefix needs to modify it.
  lck_mtx_lock(spt_add_lock_);

  // Don't allow an empty prefix.
  if (prefix[0] == '\0') return kIOReturnBadArgument;

  LOGD("Trying to add prefix: %s", prefix);

  // Enforce max tree depth.
  size_t len = strnlen(prefix, max_nodes_);

  // Grab a shared lock until a new branch is required.
  lck_rw_lock_shared(spt_lock_);

  SantaPrefixNode *node = root_;
  for (size_t i = 0; i < len; ++i) {
    // If there is a node in the path that is considered a prefix, stop adding.
    // For our purposes we only care about the shortest path that matches.
    if (node->isPrefix) break;

    // Only process a byte at a time.
    uint8_t value = (uint8_t)prefix[i];

    // Create the child if it does not exist.
    if (!node->children[value]) {
      // Upgrade the shared lock.
      // If the upgrade fails, the shared lock is released.
      if (!lck_rw_lock_shared_to_exclusive(spt_lock_)) {
        // Grab a new exclusive lock.
        lck_rw_lock_exclusive(spt_lock_);
      }

      // Is there enough room for the rest of the prefix?
      if ((node_count_ + (len - i)) > max_nodes_) {
        LOGE("Prefix tree is full, can not add: %s", prefix);

        if (node_count) *node_count = node_count_;
        lck_rw_unlock_exclusive(spt_lock_);
        lck_mtx_unlock(spt_add_lock_);
        return kIOReturnNoResources;
      }

      // Create the rest of the prefix.
      while (i < len) {
        value = (uint8_t)prefix[i++];

        SantaPrefixNode *new_node = new SantaPrefixNode();
        node->children[value] = new_node;
        ++node_count_;

        node = new_node;
      }

      // This is the end, mark the node as a prefix.
      LOGD("Added prefix: %s", prefix);

      node->isPrefix = true;

      // Downgrade the exclusive lock
      lck_rw_lock_exclusive_to_shared(spt_lock_);
    } else if (i + 1 == len) {
      // If the child does exist and it is the end...
      // Set the new, higher prefix and prune the now dead nodes.

      if (!lck_rw_lock_shared_to_exclusive(spt_lock_)) {
        lck_rw_lock_exclusive(spt_lock_);
      }

      PruneNode(node->children[value]);

      SantaPrefixNode *new_node = new SantaPrefixNode();
      new_node->isPrefix = true;

      node->children[value] = new_node;
      ++node_count_;

      LOGD("Added prefix: %s", prefix);

      lck_rw_lock_exclusive_to_shared(spt_lock_);
    }

    // Get ready for the next iteration.
    node = node->children[value];
  }

  if (node_count) *node_count = node_count_;

  lck_rw_unlock_shared(spt_lock_);
  lck_mtx_unlock(spt_add_lock_);

  return kIOReturnSuccess;
}

bool SNTPrefixTree::HasPrefix(const char *string) {
  lck_rw_lock_shared(spt_lock_);

  auto found = false;

  SantaPrefixNode *node = root_;

  // A well formed tree will always break this loop. Even if string doesn't
  // terminate.
  const char *p = string;
  while (*p) {
    // Only process a byte at a time.
    node = node->children[(uint8_t)*p++];

    // If it doesn't exist in the tree, no match.
    if (!node) break;

    // If it does exist, is it a prefix?
    if (node->isPrefix) {
      found = true;
      break;
    }
  }

  lck_rw_unlock_shared(spt_lock_);

  return found;
}

void SNTPrefixTree::Reset() {
  lck_rw_lock_exclusive(spt_lock_);

  PruneNode(root_);
  root_ = new SantaPrefixNode();
  node_count_ = 0;

  lck_rw_unlock_exclusive(spt_lock_);
}

void SNTPrefixTree::PruneNode(SantaPrefixNode *target) {
  if (!target) return;

  // For deep trees, a recursive approach will generate too many stack frames.
  // Make a "stack" and walk the tree.
  auto stack = new SantaPrefixNode *[node_count_ + 1];
  if (!stack) {
    LOGE("Unable to prune tree!");

    return;
  }
  auto count = 0;

  // Seed the "stack" with a starting node.
  stack[count++] = target;

  // Start at the target node and walk the tree to find and delete all the
  // sub-nodes.
  while (count) {
    auto node = stack[--count];

    for (int i = 0; i < 256; ++i) {
      if (!node->children[i]) continue;
      stack[count++] = node->children[i];
    }

    delete node;
    --node_count_;
  }

  delete[] stack;
}

SNTPrefixTree::~SNTPrefixTree() {
  lck_rw_lock_exclusive(spt_lock_);
  PruneNode(root_);
  root_ = nullptr;
  lck_rw_unlock_exclusive(spt_lock_);

  pthread_rwlock_destroy(&spt_lock_);
}