tests/dns_responder/dns_responder.h

   1 /*
   2  * Copyright (C) 2016 The Android Open Source Project
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *      http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless requied by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  *
  16  */
  17
  18 #ifndef DNS_RESPONDER_H
  19 #define DNS_RESPONDER_H
  20
  21 #include <arpa/nameser.h>
  22
  23 #include <atomic>
  24 #include <mutex>
  25 #include <string>
  26 #include <thread>
  27 #include <unordered_map>
  28 #include <vector>
  29
  30 #include <android-base/thread_annotations.h>
  31
  32 namespace test {
  33
  34 struct DNSHeader;
  35 struct DNSQuestion;
  36 struct DNSRecord;
  37
  38 /*
  39  * Simple DNS responder, which replies to queries with the registered response
  40  * for that type. Class is assumed to be IN. If no response is registered, the
  41  * default error response code is returned.
  42  */
  43 class DNSResponder {
  44 public:
  45     DNSResponder(std::string listen_address, std::string listen_service,
  46                  int poll_timeout_ms, uint16_t error_rcode,
  47                  double response_probability);
  48     ~DNSResponder();
  49     void addMapping(const char* name, ns_type type, const char* addr);
  50     void removeMapping(const char* name, ns_type type);
  51     void setResponseProbability(double response_probability);
  52     bool running() const;
  53     bool startServer();
  54     bool stopServer();
  55     const std::string& listen_address() const {
  56         return listen_address_;
  57     }
  58     const std::string& listen_service() const {
  59         return listen_service_;
  60     }
  61     std::vector<std::pair<std::string, ns_type>> queries() const;
  62     void clearQueries();
  63
  64 private:
  65     // Key used for accessing mappings.
  66     struct QueryKey {
  67         std::string name;
  68         unsigned type;
  69         QueryKey(std::string n, unsigned t) : name(n), type(t) {}
  70         bool operator == (const QueryKey& o) const {
  71             return name == o.name && type == o.type;
  72         }
  73         bool operator < (const QueryKey& o) const {
  74             if (name < o.name) return true;
  75             if (name > o.name) return false;
  76             return type < o.type;
  77         }
  78     };
  79
  80     struct QueryKeyHash {
  81         size_t operator() (const QueryKey& key) const {
  82             return std::hash<std::string>()(key.name) +
  83                    static_cast<size_t>(key.type);
  84         }
  85     };
  86
  87     // DNS request handler.
  88     void requestHandler();
  89
  90     // Parses and generates a response message for incoming DNS requests.
  91     // Returns false on parsing errors.
  92     bool handleDNSRequest(const char* buffer, ssize_t buffer_len,
  93                           char* response, size_t* response_len) const;
  94
  95     bool addAnswerRecords(const DNSQuestion& question,
  96                           std::vector<DNSRecord>* answers) const;
  97
  98     bool generateErrorResponse(DNSHeader* header, ns_rcode rcode,
  99                                char* response, size_t* response_len) const;
 100     bool makeErrorResponse(DNSHeader* header, ns_rcode rcode, char* response,
 101                            size_t* response_len) const;
 102
 103
 104     // Address and service to listen on, currently limited to UDP.
 105     const std::string listen_address_;
 106     const std::string listen_service_;
 107     // epoll_wait() timeout in ms.
 108     const int poll_timeout_ms_;
 109     // Error code to return for requests for an unknown name.
 110     const uint16_t error_rcode_;
 111     // Probability that a valid response is being sent instead of being sent
 112     // instead of returning error_rcode_.
 113     std::atomic<double> response_probability_;
 114
 115     // Mappings from (name, type) to registered response and the
 116     // mutex protecting them.
 117     std::unordered_map<QueryKey, std::string, QueryKeyHash> mappings_
 118         GUARDED_BY(mappings_mutex_);
 119     // TODO(imaipi): enable GUARDED_BY(mappings_mutex_);
 120     std::mutex mappings_mutex_;
 121     // Query names received so far and the corresponding mutex.
 122     mutable std::vector<std::pair<std::string, ns_type>> queries_
 123         GUARDED_BY(queries_mutex_);
 124     mutable std::mutex queries_mutex_;
 125     // Socket on which the server is listening.
 126     int socket_;
 127     // File descriptor for epoll.
 128     int epoll_fd_;
 129     // Signal for request handler termination.
 130     std::atomic<bool> terminate_ GUARDED_BY(update_mutex_);
 131     // Thread for handling incoming threads.
 132     std::thread handler_thread_ GUARDED_BY(update_mutex_);
 133     std::mutex update_mutex_;
 134 };
 135
 136 }  // namespace test
 137
 138 #endif  // DNS_RESPONDER_H