DnsRuleResolver.kt 13.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.Dispatchers
11
import kotlinx.coroutines.GlobalScope
12
import kotlinx.coroutines.Job
13
import kotlinx.coroutines.launch
14
import org.minidns.dnsmessage.DnsMessage
Daniel Wolf's avatar
Daniel Wolf committed
15
import org.minidns.dnsmessage.Question
16
import org.minidns.record.*
Lysanne Fröhlich's avatar
Lysanne Fröhlich committed
17
18
19
20
import java.util.*
import kotlin.collections.HashSet
import kotlin.math.abs

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

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

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

45
46
    private var previousRefreshJob:Job? = null

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

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

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

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

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

83
84
85
86
    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)
87

88
89
90
91
92
93
94
95
96
        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
97
            }
98
99
100
101
102
        }
        if(wildcardCount != 0 && cachedWildcardResolved.size != 0) {
            val res = cachedWildcardResolved[hostHash]
            if(res != null) {
                return res
103
            }
104
        }
105

106
107
108
109
110
111
112
113
114
        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,
115
116
                includeWhitelistEntries = true,
                includeNonWhitelistEntries = false
117
118
119
            ).firstOrNull {
                DnsRuleDialog.databaseHostToMatcher(it.host).reset(uniformQuestion)
                    .matches()
120
            } else null
121
        } else null
122

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

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

202
    override fun canResolve(question: Question): Boolean {
203
        if(StaticDnsRuleResolver.staticRules.containsKey(question.name?.toString()?.lowercase())) {
204
205
            return true
        }
206
207
208
209
210
211
212
213
214
215
216
        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
217
218
219
220
221
222
223
224
225
226
227
    // 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
        }
    }

228
    override fun resolve(question: Question): List<Record<*>> {
229
230
        if(StaticDnsRuleResolver.staticRules.containsKey(question.name?.toString()?.lowercase())) {
            return StaticDnsRuleResolver.staticRules[question.name?.toString()?.lowercase()] ?: emptyList()
231
        }
232
        val result = resolveResults.remove(question.hashCode())
Daniel Wolf's avatar
Daniel Wolf committed
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
        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()
    }

251
252
253
254
255
256
257
    override fun cleanup() {
        cachedWildcardWhitelisted.clear()
        cachedNonWildcardWhitelisted.clear()
        cachedResolved.clear()
        cachedWildcardResolved.clear()
        cachedNonIncluded.clear()
    }
Daniel Wolf's avatar
Daniel Wolf committed
258

259
260
    // Handle CNAME Cloaking
    // Does not need to handle whitelist as the query has already been forwarded
261
    override fun mapResponse(message: DnsMessage): DnsMessage {
262
263
264
265
        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
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
        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)
Daniel Wolf's avatar
Daniel Wolf committed
292
                    mappedAnswers.add(Record(target, type, Record.CLASS.IN, originalTargetRecord?.ttl ?: message.answerSection.minByOrNull {
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
                        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? {
322
        val entry = findRuleTarget(host, type)
323
324
325
326
327
328
329

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

Daniel Wolf's avatar
Daniel Wolf committed
330
}