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);
setFailOnEdns(bool fail)52     void setFailOnEdns(bool fail) { fail_on_edns_ = fail; }
53     bool running() const;
54     bool startServer();
55     bool stopServer();
listen_address()56     const std::string& listen_address() const {
57         return listen_address_;
58     }
listen_service()59     const std::string& listen_service() const {
60         return listen_service_;
61     }
62     std::vector<std::pair<std::string, ns_type>> queries() const;
63     void clearQueries();
64 
65 private:
66     // Key used for accessing mappings.
67     struct QueryKey {
68         std::string name;
69         unsigned type;
QueryKeyQueryKey70         QueryKey(std::string n, unsigned t) : name(n), type(t) {}
71         bool operator == (const QueryKey& o) const {
72             return name == o.name && type == o.type;
73         }
74         bool operator < (const QueryKey& o) const {
75             if (name < o.name) return true;
76             if (name > o.name) return false;
77             return type < o.type;
78         }
79     };
80 
81     struct QueryKeyHash {
operatorQueryKeyHash82         size_t operator() (const QueryKey& key) const {
83             return std::hash<std::string>()(key.name) +
84                    static_cast<size_t>(key.type);
85         }
86     };
87 
88     // DNS request handler.
89     void requestHandler();
90 
91     // Parses and generates a response message for incoming DNS requests.
92     // Returns false on parsing errors.
93     bool handleDNSRequest(const char* buffer, ssize_t buffer_len,
94                           char* response, size_t* response_len) const;
95 
96     bool addAnswerRecords(const DNSQuestion& question,
97                           std::vector<DNSRecord>* answers) const;
98 
99     bool generateErrorResponse(DNSHeader* header, ns_rcode rcode,
100                                char* response, size_t* response_len) const;
101     bool makeErrorResponse(DNSHeader* header, ns_rcode rcode, char* response,
102                            size_t* response_len) const;
103 
104 
105     // Address and service to listen on, currently limited to UDP.
106     const std::string listen_address_;
107     const std::string listen_service_;
108     // epoll_wait() timeout in ms.
109     const int poll_timeout_ms_;
110     // Error code to return for requests for an unknown name.
111     const uint16_t error_rcode_;
112     // Probability that a valid response is being sent instead of being sent
113     // instead of returning error_rcode_.
114     std::atomic<double> response_probability_;
115 
116     // If true, behave like an old DNS server that doesn't support EDNS.
117     // Default false.
118     std::atomic<bool> fail_on_edns_;
119 
120     // Mappings from (name, type) to registered response and the
121     // mutex protecting them.
122     std::unordered_map<QueryKey, std::string, QueryKeyHash> mappings_
123         GUARDED_BY(mappings_mutex_);
124     mutable std::mutex mappings_mutex_;
125     // Query names received so far and the corresponding mutex.
126     mutable std::vector<std::pair<std::string, ns_type>> queries_
127         GUARDED_BY(queries_mutex_);
128     mutable std::mutex queries_mutex_;
129     // Socket on which the server is listening.
130     int socket_;
131     // File descriptor for epoll.
132     int epoll_fd_;
133     // Signal for request handler termination.
134     std::atomic<bool> terminate_;
135     // Thread for handling incoming threads.
136     std::thread handler_thread_ GUARDED_BY(update_mutex_);
137     std::mutex update_mutex_;
138 };
139 
140 }  // namespace test
141 
142 #endif  // DNS_RESPONDER_H
143