1 /*
2  * Copyright (C) 2018 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 required 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 #ifndef _DNS_DNSTLSQUERYMAP_H
18 #define _DNS_DNSTLSQUERYMAP_H
19 
20 #include <future>
21 #include <map>
22 #include <mutex>
23 #include <vector>
24 
25 #include <android-base/thread_annotations.h>
26 #include <netdutils/Slice.h>
27 
28 #include "DnsTlsServer.h"
29 
30 namespace android {
31 namespace net {
32 
33 // Keeps track of queries and responses.  This class matches responses with queries.
34 // All methods are thread-safe and non-blocking.
35 class DnsTlsQueryMap {
36   public:
37     enum class Response : uint8_t { success, network_error, limit_error, internal_error };
38 
39     struct Query {
40         // The new ID number assigned to this query.
41         uint16_t newId;
42         // A query that has been passed to recordQuery(), with its original ID number.
43         const std::vector<uint8_t> query;
44     };
45 
46     struct Result {
47         Response code;
48         std::vector<uint8_t> response;
49     };
50 
51     struct QueryFuture {
QueryFutureQueryFuture52         QueryFuture(Query query, std::future<Result> result)
53             : query(query), result(std::move(result)) {}
54         Query query;
55         // A future which will resolve to the result of this query.
56         std::future<Result> result;
57     };
58 
59     // Returns an object containing everything needed to complete processing of
60     // this query, or null if the query could not be recorded.
61     std::unique_ptr<QueryFuture> recordQuery(const netdutils::Slice query);
62 
63     // Process a response, including a new ID.  If the response
64     // is not recognized as matching any query, it will be ignored.
65     void onResponse(std::vector<uint8_t> response);
66 
67     // Clear all map contents.  This causes all pending queries to resolve with failure.
68     void clear();
69 
70     // Get all pending queries.  This returns a shallow copy, mostly for thread-safety.
71     std::vector<Query> getAll();
72 
73     // Mark a query has having been retried.  If the query hits the retry limit, it will
74     // be expired at the next call to cleanup.
75     void markTried(uint16_t newId);
76     void cleanup();
77 
78     // Returns true if there are no pending queries.
79     bool empty();
80 
81     // The maximum number of times we will send a query before abandoning it.
82     static constexpr int kMaxTries = 3;
83 
84   private:
85     std::mutex mLock;
86 
87     struct QueryPromise {
QueryPromiseQueryPromise88         QueryPromise(Query query) : query(query) {}
89         Query query;
90         // Number of times the query has been tried.  Limited to kMaxTries.
91         int tries = 0;
92         // A promise whose future is returned by recordQuery()
93         // It is fulfilled by onResponse().
94         std::promise<Result> result;
95     };
96 
97     // Outstanding queries by newId.
98     std::map<uint16_t, QueryPromise> mQueries GUARDED_BY(mLock);
99 
100     // Get a "newId" number that is not currently in use.  Returns -1 if there are none.
101     int32_t getFreeId() REQUIRES(mLock);
102 
103     // Fulfill the result with an error code.
104     static void expire(QueryPromise* _Nonnull p);
105 };
106 
107 }  // end of namespace net
108 }  // end of namespace android
109 
110 #endif  // _DNS_DNSTLSQUERYMAP_H
111