pdns/nod.hh

   1 /*
   2  * This file is part of PowerDNS or dnsdist.
   3  * Copyright -- PowerDNS.COM B.V. and its contributors
   4  *
   5  * This program is free software; you can redistribute it and/or modify
   6  * it under the terms of version 2 of the GNU General Public License as
   7  * published by the Free Software Foundation.
   8  *
   9  * In addition, for the avoidance of any doubt, permission is granted to
  10  * link this program with OpenSSL and to (re)distribute the binaries
  11  * produced as the result of such linking.
  12  *
  13  * This program is distributed in the hope that it will be useful,
  14  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  15  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  16  * GNU General Public License for more details.
  17  *
  18  * You should have received a copy of the GNU General Public License
  19  * along with this program; if not, write to the Free Software
  20  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
  21  */
  22 #pragma once
  23 #include <atomic>
  24 #include <mutex>
  25 #include <thread>
  26 #include "dnsname.hh"
  27 #include "stable-bloom.hh"
  28
  29 namespace nod {
  30   const float c_fp_rate = 0.01;
  31   const size_t c_num_cells = 67108864;
  32   const uint8_t c_num_dec = 10;
  33   const unsigned int snapshot_interval_default = 600;
  34   const std::string bf_suffix = "bf";
  35   const std::string sbf_prefix = "sbf";
  36
  37   // Theses classes are not designed to be shared between threads
  38   // Use a new instance per-thread, e.g. using thread local storage
  39   // Synchronization (at the class level) is still needed for reading from
  40   // and writing to the cache dir
  41   // Synchronization (at the instance level) is needed when snapshotting
  42   class PersistentSBF {
  43   public:
  44     PersistentSBF() : d_sbf{c_fp_rate, c_num_cells, c_num_dec} {}
  45     PersistentSBF(uint32_t num_cells) : d_sbf{c_fp_rate, num_cells, c_num_dec} {}
  46     bool init(bool ignore_pid=false);
  47     void setPrefix(const std::string& prefix) { d_prefix = prefix; } // Added to filenames in cachedir
  48     void setCacheDir(const std::string& cachedir);
  49     bool snapshotCurrent(std::thread::id tid); // Write the current file out to disk
  50     void add(const std::string& data) {
  51       // The only time this should block is when snapshotting
  52       std::lock_guard<std::mutex> lock(d_sbf_mutex);
  53       d_sbf.add(data);
  54     }
  55     bool test(const std::string& data) { return d_sbf.test(data); }
  56     bool testAndAdd(const std::string& data) {
  57       // The only time this should block is when snapshotting
  58       std::lock_guard<std::mutex> lock(d_sbf_mutex);
  59       return d_sbf.testAndAdd(data);
  60     }
  61   private:
  62     bool d_init{false};
  63     bf::stableBF d_sbf; // Stable Bloom Filter
  64     std::string d_cachedir;
  65     std::string d_prefix = sbf_prefix;
  66     std::mutex d_sbf_mutex; // Per-instance mutex for snapshots
  67     static std::mutex d_cachedir_mutex; // One mutex for all instances of this class
  68   };
  69
  70   class NODDB {
  71   public:
  72     NODDB() : d_psbf{} {}
  73     NODDB(uint32_t num_cells) : d_psbf{num_cells} {}
  74     // Set ignore_pid to true if you don't mind loading files
  75     // created by the current process
  76     bool init(bool ignore_pid=false) {
  77       d_psbf.setPrefix("nod");
  78       return d_psbf.init(ignore_pid);
  79     }
  80     bool isNewDomain(const std::string& domain); // Returns true if newly observed domain
  81     bool isNewDomain(const DNSName& dname); // As above
  82     bool isNewDomainWithParent(const std::string& domain, std::string& observed); // Returns true if newly observed domain, in which case "observed" contains the parent domain which *was* observed (or "" if domain is . or no parent domains observed)
  83     bool isNewDomainWithParent(const DNSName& dname, std::string& observed); // As above
  84     void addDomain(const DNSName& dname); // You need to add this to refresh frequently used domains
  85     void addDomain(const std::string& domain); // As above
  86     void setSnapshotInterval(unsigned int secs) { d_snapshot_interval = secs; }
  87     void setCacheDir(const std::string& cachedir) { d_psbf.setCacheDir(cachedir); }
  88     bool snapshotCurrent(std::thread::id tid) { return d_psbf.snapshotCurrent(tid); }
  89     static void startHousekeepingThread(std::shared_ptr<NODDB> noddbp, std::thread::id tid) {
  90       noddbp->housekeepingThread(tid);
  91     }
  92   private:
  93     PersistentSBF d_psbf;
  94     unsigned int d_snapshot_interval{snapshot_interval_default}; // Number seconds between snapshots
  95     void housekeepingThread(std::thread::id tid);
  96   };
  97
  98   class UniqueResponseDB {
  99   public:
 100     UniqueResponseDB() : d_psbf{} {}
 101     UniqueResponseDB(uint32_t num_cells) : d_psbf{num_cells} {}
 102     bool init(bool ignore_pid=false) {
 103       d_psbf.setPrefix("udr");
 104       return d_psbf.init(ignore_pid);
 105     }
 106     bool isUniqueResponse(const std::string& response);
 107     void addResponse(const std::string& response);
 108     void setSnapshotInterval(unsigned int secs) { d_snapshot_interval = secs; }
 109     void setCacheDir(const std::string& cachedir) { d_psbf.setCacheDir(cachedir); }
 110     bool snapshotCurrent(std::thread::id tid) { return d_psbf.snapshotCurrent(tid); }
 111     static void startHousekeepingThread(std::shared_ptr<UniqueResponseDB> udrdbp, std::thread::id tid) {
 112       udrdbp->housekeepingThread(tid);
 113     }
 114   private:
 115     PersistentSBF d_psbf;
 116     unsigned int d_snapshot_interval{snapshot_interval_default}; // Number seconds between snapshots
 117     void housekeepingThread(std::thread::id tid);
 118   };
 119
 120 }