Commit e058f2bb authored by Mygod's avatar Mygod

Shorten resolve

parent f8142a6f
...@@ -61,6 +61,12 @@ class LocalDnsServer(private val localResolver: suspend (String) -> Array<InetAd ...@@ -61,6 +61,12 @@ class LocalDnsServer(private val localResolver: suspend (String) -> Array<InetAd
*/ */
private const val TTL = 120L private const val TTL = 120L
private const val UDP_PACKET_SIZE = 512 private const val UDP_PACKET_SIZE = 512
private fun prepareDnsResponse(request: Message) = Message(request.header.id).apply {
header.setFlag(Flags.QR.toInt()) // this is a response
if (request.header.getFlag(Flags.RD.toInt())) header.setFlag(Flags.RD.toInt())
addRecord(request.question, Section.QUESTION)
}
} }
private val monitor = ChannelMonitor() private val monitor = ChannelMonitor()
...@@ -109,28 +115,21 @@ class LocalDnsServer(private val localResolver: suspend (String) -> Array<InetAd ...@@ -109,28 +115,21 @@ class LocalDnsServer(private val localResolver: suspend (String) -> Array<InetAd
if (localResults.isEmpty()) return@coroutineScope remote.await() if (localResults.isEmpty()) return@coroutineScope remote.await()
if (localIpMatcher.isEmpty() || localIpMatcher.any { subnet -> localResults.any(subnet::matches) }) { if (localIpMatcher.isEmpty() || localIpMatcher.any { subnet -> localResults.any(subnet::matches) }) {
remote.cancel() remote.cancel()
val response = Message(request.header.id) ByteBuffer.wrap(prepareDnsResponse(request).apply {
response.header.setFlag(Flags.QR.toInt()) // this is a response header.setFlag(Flags.RA.toInt()) // recursion available
if (request.header.getFlag(Flags.RD.toInt())) response.header.setFlag(Flags.RD.toInt()) for (address in localResults) addRecord(when (address) {
response.header.setFlag(Flags.RA.toInt()) // recursion available is Inet4Address -> ARecord(request.question.name, DClass.IN, TTL, address)
response.addRecord(request.question, Section.QUESTION) is Inet6Address -> AAAARecord(request.question.name, DClass.IN, TTL, address)
for (address in localResults) response.addRecord(when (address) { else -> throw IllegalStateException("Unsupported address $address")
is Inet4Address -> ARecord(request.question.name, DClass.IN, TTL, address) }, Section.ANSWER)
is Inet6Address -> AAAARecord(request.question.name, DClass.IN, TTL, address) }.toWire())
else -> throw IllegalStateException("Unsupported address $address") } else remote.await()
}, Section.ANSWER)
return@coroutineScope ByteBuffer.wrap(response.toWire())
}
return@coroutineScope remote.await()
} catch (e: IOException) { } catch (e: IOException) {
remote.cancel() remote.cancel()
printLog(e) printLog(e)
val response = Message(request.header.id) ByteBuffer.wrap(prepareDnsResponse(request).apply {
response.header.rcode = Rcode.SERVFAIL header.rcode = Rcode.SERVFAIL
response.header.setFlag(Flags.QR.toInt()) }.toWire())
if (request.header.getFlag(Flags.RD.toInt())) response.header.setFlag(Flags.RD.toInt())
response.addRecord(request.question, Section.QUESTION)
return@coroutineScope ByteBuffer.wrap(response.toWire())
} }
} }
} }
...@@ -142,9 +141,7 @@ class LocalDnsServer(private val localResolver: suspend (String) -> Array<InetAd ...@@ -142,9 +141,7 @@ class LocalDnsServer(private val localResolver: suspend (String) -> Array<InetAd
channel.connect(proxy) channel.connect(proxy)
val wrapped = remoteDns.tcpWrap(packet) val wrapped = remoteDns.tcpWrap(packet)
while (!channel.finishConnect()) monitor.wait(channel, SelectionKey.OP_CONNECT) while (!channel.finishConnect()) monitor.wait(channel, SelectionKey.OP_CONNECT)
while (channel.write(wrapped) >= 0 && wrapped.hasRemaining()) { while (channel.write(wrapped) >= 0 && wrapped.hasRemaining()) monitor.wait(channel, SelectionKey.OP_WRITE)
monitor.wait(channel, SelectionKey.OP_WRITE)
}
val result = remoteDns.tcpReceiveBuffer(UDP_PACKET_SIZE) val result = remoteDns.tcpReceiveBuffer(UDP_PACKET_SIZE)
remoteDns.tcpUnwrap(result, channel::read) { monitor.wait(channel, SelectionKey.OP_READ) } remoteDns.tcpUnwrap(result, channel::read) { monitor.wait(channel, SelectionKey.OP_READ) }
result result
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment