DnsRuleResolver.kt 12.9 KB
Newer Older
Daniel Wolf's avatar
Daniel Wolf committed
1
2
3
4
package com.frostnerd.smokescreen.util.proxy

import android.content.Context
import com.frostnerd.dnstunnelproxy.LocalResolver
5
import com.frostnerd.smokescreen.database.entities.DnsRule
Daniel Wolf's avatar
Daniel Wolf committed
6
7
8
import com.frostnerd.smokescreen.database.getDatabase
import com.frostnerd.smokescreen.dialog.DnsRuleDialog
import com.frostnerd.smokescreen.getPreferences
9
import com.frostnerd.smokescreen.util.MaxSizeMap
10
import kotlinx.coroutines.GlobalScope
11
import kotlinx.coroutines.Job
12
import kotlinx.coroutines.launch
13
import org.minidns.dnsmessage.DnsMessage
Daniel Wolf's avatar
Daniel Wolf committed
14
import org.minidns.dnsmessage.Question
15
import org.minidns.record.*
Lysanne Fröhlich's avatar
Lysanne Fröhlich committed
16
17
18
19
import java.util.*
import kotlin.collections.HashSet
import kotlin.math.abs

20
class DnsRuleResolver(context: Context) : LocalResolver(false) {
Lysanne Fröhlich's avatar
Lysanne Fröhlich committed
21
    private val maxWhitelistCacheSize = 250
22
    private val maxResolvedCacheSize = 500
23
    private val maxWildcardResolvedCacheSize = 250
Daniel Wolf's avatar
Daniel Wolf committed
24
25

    private val dao = context.getDatabase().dnsRuleDao()
26
    private val resolveResults = mutableMapOf<Int, String>()
Daniel Wolf's avatar
Daniel Wolf committed
27
28
    private val wwwRegex = Regex("^www\\.")
    private val useUserRules = context.getPreferences().customHostsEnabled
Lysanne Fröhlich's avatar
Lysanne Fröhlich committed
29
30
31
    private var ruleCount: Int? = null
    private var wildcardCount: Int? = null
    private var whitelistCount: Int? = null
32
    private var nonWildcardCount:Int? = null
33
34
    private var wildcardWhitelistCount:Int? = null
    private var nonWildcardWhitelistCount:Int? = null
Lysanne Fröhlich's avatar
Lysanne Fröhlich committed
35
36

    // These sets contain hashes of the hosts, the most significant bit of the hash it 1 for IPv6 and 0 for IPv4
37
    // Hashes are stored because they are shorter than strings (Int=4 Bytes, String=2-3 per char)
38
39
    private var cachedWildcardWhitelisted = HashSet<Int>(15)
    private var cachedNonWildcardWhitelisted = HashSet<Int>(15)
40
    private var cachedResolved = MaxSizeMap<Int, String>(maxResolvedCacheSize, 40)
41
    private var cachedWildcardResolved = MaxSizeMap<Int, String>(maxWildcardResolvedCacheSize, 30)
42
    private var cachedNonIncluded = HashSet<Int>(15)
43

44
45
    private var previousRefreshJob:Job? = null

46
    init {
Lysanne Fröhlich's avatar
Lysanne Fröhlich committed
47
        refreshRuleCount()
48
49
    }

Lysanne Fröhlich's avatar
Lysanne Fröhlich committed
50
    fun refreshRuleCount() {
51
52
        previousRefreshJob?.cancel()
        previousRefreshJob = GlobalScope.launch {
53
            val previousRuleCount = ruleCount
54
55
            ruleCount = dao.getActiveCount().toInt()
            wildcardCount = dao.getActiveWildcardCount().toInt()
56
            val previousWhitelistCount = whitelistCount
57
            whitelistCount = dao.getActiveWhitelistCount().toInt()
58
            nonWildcardCount = ruleCount!! - wildcardCount!!
59

60
            wildcardWhitelistCount = dao.getActiveWildcardWhitelistCount().toInt()
61
62
            nonWildcardWhitelistCount = whitelistCount!! - wildcardWhitelistCount!!

63
64
65
            if(previousWhitelistCount != whitelistCount) {
                cachedWildcardWhitelisted.clear()
                cachedNonWildcardWhitelisted.clear()
66
                preloadWhitelistEntries()
67
            }
68
69
70
            if(previousRuleCount != ruleCount){
                cachedResolved.clear()
                cachedWildcardResolved.clear()
71
                cachedNonIncluded.clear()
72
            }
73
74
        }
    }
Daniel Wolf's avatar
Daniel Wolf committed
75

Daniel Wolf's avatar
Daniel Wolf committed
76
77
78
79
80
81
    private suspend fun preloadWhitelistEntries() {
        cachedNonWildcardWhitelisted.addAll(dao.getRandomNonWildcardWhitelistEntries(100).map {
            hashHost(it.toLowerCase(), Record.TYPE.ANY)
        })
    }

82
83
84
85
    private fun findRuleTarget(question: String, type:Record.TYPE):String? {
        val uniformQuestion = question.replace(wwwRegex, "").toLowerCase(Locale.ROOT)
        val hostHash = hashHost(uniformQuestion, type)
        val wildcardHostHash = hashHost(uniformQuestion, Record.TYPE.ANY)
86

87
88
89
90
91
92
93
94
95
        if(cachedNonIncluded.size != 0 && cachedNonIncluded.contains(hostHash)) return null
        if(whitelistCount != 0) {
            if(cachedNonWildcardWhitelisted.size != 0 && cachedNonWildcardWhitelisted.contains(wildcardHostHash)) return null
            else if(cachedWildcardWhitelisted.size != 0 && cachedWildcardWhitelisted.contains(wildcardHostHash)) return null
        }
        if(nonWildcardCount != 0 && cachedResolved.size != 0) {
            val res = cachedResolved[hostHash]
            if(res != null) {
                return res
96
            }
97
98
99
100
101
        }
        if(wildcardCount != 0 && cachedWildcardResolved.size != 0) {
            val res = cachedWildcardResolved[hostHash]
            if(res != null) {
                return res
102
            }
103
        }
104

105
106
107
108
109
110
111
112
113
114
115
116
117
118
        val whitelistEntry: DnsRule? = if (whitelistCount != 0) {
            val normal = if(nonWildcardWhitelistCount != 0 && (nonWildcardWhitelistCount == null || nonWildcardWhitelistCount != whitelistCount)) dao.findNonWildcardWhitelistEntry(
                uniformQuestion,
                useUserRules
            ).firstOrNull() else null
            normal ?: if(wildcardWhitelistCount != 0) dao.findPossibleWildcardRuleTarget(
                uniformQuestion,
                type,
                useUserRules,
                true,
                false
            ).firstOrNull {
                DnsRuleDialog.databaseHostToMatcher(it.host).reset(uniformQuestion)
                    .matches()
119
            } else null
120
        } else null
121

122
123
124
        if (whitelistEntry != null) {
            if(whitelistEntry.isWildcard) cachedWildcardWhitelisted.add(wildcardHostHash)
            else cachedNonWildcardWhitelisted.add(wildcardHostHash)
125

126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
            if(cachedWildcardWhitelisted.size >= maxWhitelistCacheSize*2) cachedWildcardWhitelisted.clear()
            if(cachedNonWildcardWhitelisted.size >= maxWhitelistCacheSize) cachedNonWildcardWhitelisted.clear()
            return null
        }
        else {
            val resolveResult = if(nonWildcardCount != 0) {
                if(nonWildcardCount == cachedResolved.size) {
                    null // We would have hit cache otherwise
                } else {
                    dao.findRuleTarget(uniformQuestion, type, useUserRules)
                        ?.let {
                            when (it) {
                                "0" -> "0.0.0.0"
                                "1" -> {
                                    if (type == Record.TYPE.AAAA) "::1"
                                    else "127.0.0.1"
Daniel Wolf's avatar
Daniel Wolf committed
142
                                }
143
                                else -> it
Daniel Wolf's avatar
Daniel Wolf committed
144
                            }
145
146
147
148
149
                        }
                }
            } else null
            when {
                resolveResult != null -> {
150
                    cachedResolved[hostHash] = resolveResult
151
152
153
                    return resolveResult
                }
                wildcardCount != 0 -> {
154
                    val wildcardResolveResult = dao.findPossibleWildcardRuleTarget(
Daniel Wolf's avatar
Daniel Wolf committed
155
                        uniformQuestion,
156
                        type,
Daniel Wolf's avatar
Daniel Wolf committed
157
158
159
                        useUserRules,
                        false,
                        true
160
                    ).firstOrNull {
Daniel Wolf's avatar
Daniel Wolf committed
161
162
                        DnsRuleDialog.databaseHostToMatcher(it.host)
                            .reset(uniformQuestion).matches()
163
                    }?.let {
164
                        if (type == Record.TYPE.AAAA) it.ipv6Target
165
166
167
168
169
170
                            ?: it.target
                        else it.target
                    }?.let {
                        when (it) {
                            "0" -> "0.0.0.0"
                            "1" -> {
171
                                if (type == Record.TYPE.AAAA) "::1"
172
                                else "127.0.0.1"
Daniel Wolf's avatar
Daniel Wolf committed
173
                            }
174
                            else -> it
Daniel Wolf's avatar
Daniel Wolf committed
175
                        }
176
                    }
177
                    return if (wildcardResolveResult != null) {
Daniel Wolf's avatar
Daniel Wolf committed
178
                        cachedWildcardResolved[hostHash] = wildcardResolveResult
179
                        wildcardResolveResult
Daniel Wolf's avatar
Daniel Wolf committed
180
                    } else {
181
182
                        if(cachedNonIncluded.size >= maxWhitelistCacheSize) cachedNonIncluded.clear()
                        cachedNonIncluded.add(hostHash)
183
                        null
Daniel Wolf's avatar
Daniel Wolf committed
184
                    }
185
186
                }
                else -> {
187
188
                    if(cachedNonIncluded.size >= maxWhitelistCacheSize) cachedNonIncluded.clear()
                    cachedNonIncluded.add(hostHash)
189
                    return null
Daniel Wolf's avatar
Daniel Wolf committed
190
191
192
193
194
                }
            }
        }
    }

195
196
197
198
199
200
201
202
203
204
205
206
    override suspend fun canResolve(question: Question): Boolean {
        return if ((ruleCount == 0 || (ruleCount != null && ruleCount == whitelistCount)) || (question.type != Record.TYPE.A && question.type != Record.TYPE.AAAA)) {
            false
        } else {
            val res = findRuleTarget(question.name.toString(), question.type)
            return if(res != null) {
                resolveResults[question.hashCode()] = res
                true
            } else false
        }
    }

Lysanne Fröhlich's avatar
Lysanne Fröhlich committed
207
208
209
210
211
212
213
214
215
216
217
    // A fast hashing function with high(er) collision rate
    // As only a few hosts are stored at the same time the collision rate is not important.
    // The effective room is 2^31
    private fun hashHost(host: String, type: Record.TYPE): Int {
        return if (type == Record.TYPE.AAAA) {
            abs(host.hashCode()) * -1 // Sets first bit to 1
        } else {
            abs(host.hashCode()) // Sets first bit to 0
        }
    }

Daniel Wolf's avatar
Daniel Wolf committed
218
    override suspend fun resolve(question: Question): List<Record<*>> {
219
        val result = resolveResults.remove(question.hashCode())
Daniel Wolf's avatar
Daniel Wolf committed
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
        return result?.let {
            val data = if (question.type == Record.TYPE.A) {
                A(it)
            } else {
                AAAA(it)
            }
            listOf(
                Record(
                    question.name.toString(),
                    question.type,
                    question.clazz.value,
                    9999,
                    data
                )
            )
        } ?: throw IllegalStateException()
    }

    override fun cleanup() {}

240
241
242
    // Handle CNAME Cloaking
    // Does not need to handle whitelist as the query has already been forwarded
    override suspend fun mapResponse(message: DnsMessage): DnsMessage {
243
244
245
246
        if(ruleCount == 0 || (ruleCount != null && ruleCount == whitelistCount) || message.questions.size == 0) return message // No rules or only whitelist rules present
        else if(whitelistCount != 0 && hashHost(message.question.name.toString().replace(wwwRegex, "").toLowerCase(Locale.ROOT), message.question.type).let {
                cachedWildcardWhitelisted.contains(it) || cachedNonWildcardWhitelisted.contains(it)
            }) return message
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
        else if(!message.answerSection.any {
                it.type == Record.TYPE.CNAME
            }) return message
        else if(!message.answerSection.any {
                it.type == Record.TYPE.A
            } && ! message.answerSection.any {
                it.type == Record.TYPE.AAAA
            }) return message // The Dns rules only have IPv6 and IPv4

        val ordered = message.answerSection.sortedByDescending { it.type.value } // CNAME at the front

        val mappedAnswers = mutableListOf<Record<*>>()
        val mappedTargets = mutableSetOf<String>()
        for(record in ordered) {
            if(record.type == Record.TYPE.CNAME) {
                val target = (record.payloadData as CNAME).target.toString()
                if(mappedTargets.contains(record.name.toString())) { // Continue skipping the whole CNAME tree
                    mappedTargets.add(target)
                    continue
                }
                mappedAnswers.add(record)
                val originalTargetRecord = followCnameChainToFirstRecord(target, ordered)
                val type = originalTargetRecord?.type ?: message.questions.firstOrNull()?.type ?: Record.TYPE.A
                val ruleData = resolveForCname(target, type)
                if(ruleData != null) {
                    mappedTargets.add(target)
                    mappedAnswers.add(Record(target, type, Record.CLASS.IN, originalTargetRecord?.ttl ?: message.answerSection.minBy {
                        it.ttl
                    }?.ttl ?: record.ttl, ruleData, originalTargetRecord?.unicastQuery ?: false))
                }
            } else if(!mappedTargets.contains(record.name.toString())) mappedAnswers.add(record)
        }

        return if(mappedTargets.size != 0) {
            message.asBuilder().setAnswers(mappedAnswers).build()
        } else {
            message
        }
    }

    private fun followCnameChainToFirstRecord(name:String, records:List<Record<*>>):Record<*>? {
        var target = name
        var recursionDepth = 0
        while(recursionDepth++ < 50) {
            val nextTarget = records.firstOrNull {
                it.name.toString() == target
            }
            if(nextTarget != null && nextTarget.type != Record.TYPE.CNAME) {
                return nextTarget
            } else if(nextTarget == null) return null
            else target = (nextTarget.payloadData as CNAME).target.toString()
        }
        return null
    }

    private fun resolveForCname(host:String, type:Record.TYPE): Data? {
303
        val entry = findRuleTarget(host, type)
304
305
306
307
308
309
310

        return if(entry != null) {
            if(type == Record.TYPE.A) A(entry)
            else AAAA(entry)
        } else null
    }

Daniel Wolf's avatar
Daniel Wolf committed
311
}