Home | History | Annotate | Download | only in tests 
      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
    139