1 /*
2  * Copyright (C) 2022 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 package com.android.testutils
18 
19 import android.net.DnsResolver.CLASS_IN
20 import com.android.net.module.util.DnsPacket
21 import com.android.net.module.util.DnsPacket.ANSECTION
22 import java.net.InetAddress
23 import java.util.concurrent.ConcurrentHashMap
24 
25 const val DEFAULT_TTL_S = 5L
26 
27 /**
28  * Helper class to store the mapping of DNS queries.
29  *
30  * DnsAnswerProvider is built atop a ConcurrentHashMap and as such it provides the same
31  * guarantees as ConcurrentHashMap between writing and reading elements. Specifically :
32  * - Setting an answer happens-before reading the same answer.
33  * - Callers can read and write concurrently from DnsAnswerProvider and expect no
34  *   ConcurrentModificationException.
35  * Freshness of the answers depends on ordering of the threads ; if callers need a
36  * freshness guarantee, they need to provide the happens-before relationship from a
37  * write that they want to observe to the read that they need to be observed.
38  */
39 class DnsAnswerProvider {
40     private val mDnsKeyToRecords = ConcurrentHashMap<String, List<DnsPacket.DnsRecord>>()
41 
42     /**
43      * Get answer for the specified hostname.
44      *
45      * @param query the target hostname.
46      * @param type type of record, could be A or AAAA.
47      *
48      * @return list of [DnsPacket.DnsRecord] associated to the query. Empty if no record matches.
49      */
getAnswernull50     fun getAnswer(query: String, type: Int) = mDnsKeyToRecords[query]
51             .orEmpty().filter { it.nsType == type }
52 
53     /** Set answer for the specified {@code query}.
54      *
55      * @param query the target hostname
56      * @param addresses [List<InetAddress>] which could be used to generate multiple A or AAAA
57      *                  RRs with the corresponding addresses.
58      */
<lambda>null59     fun setAnswer(query: String, hosts: List<InetAddress>) = mDnsKeyToRecords.put(query, hosts.map {
60             DnsPacket.DnsRecord.makeAOrAAAARecord(ANSECTION, query, CLASS_IN, DEFAULT_TTL_S, it)
61         })
62 
clearAnswernull63     fun clearAnswer(query: String) = mDnsKeyToRecords.remove(query)
64     fun clearAll() = mDnsKeyToRecords.clear()
65 }