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
32namespace test {
33
34struct DNSHeader;
35struct DNSQuestion;
36struct 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 */
43class DNSResponder {
44public:
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
64private:
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