You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

LineBufferedSocket.kt 3.7KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  1. package com.dmdirc.ktirc.io
  2. import com.dmdirc.ktirc.util.logger
  3. import io.ktor.network.selector.ActorSelectorManager
  4. import io.ktor.network.sockets.Socket
  5. import io.ktor.network.sockets.aSocket
  6. import io.ktor.network.sockets.openReadChannel
  7. import io.ktor.network.sockets.openWriteChannel
  8. import io.ktor.network.tls.tls
  9. import kotlinx.coroutines.CoroutineScope
  10. import kotlinx.coroutines.Dispatchers
  11. import kotlinx.coroutines.ExperimentalCoroutinesApi
  12. import kotlinx.coroutines.GlobalScope
  13. import kotlinx.coroutines.channels.ReceiveChannel
  14. import kotlinx.coroutines.channels.produce
  15. import kotlinx.coroutines.io.ByteReadChannel
  16. import kotlinx.coroutines.io.ByteWriteChannel
  17. import java.net.InetSocketAddress
  18. import java.security.SecureRandom
  19. import javax.net.ssl.X509TrustManager
  20. internal interface LineBufferedSocket {
  21. suspend fun connect()
  22. fun disconnect()
  23. suspend fun sendLine(line: ByteArray, offset: Int = 0, length: Int = line.size)
  24. suspend fun sendLine(line: String)
  25. fun readLines(coroutineScope: CoroutineScope): ReceiveChannel<ByteArray>
  26. }
  27. /**
  28. * Asynchronous socket that buffers incoming data and emits individual lines.
  29. */
  30. // TODO: Expose advanced TLS options
  31. internal class KtorLineBufferedSocket(private val host: String, private val port: Int, private val tls: Boolean = false): LineBufferedSocket {
  32. companion object {
  33. const val CARRIAGE_RETURN = '\r'.toByte()
  34. const val LINE_FEED = '\n'.toByte()
  35. }
  36. var tlsTrustManager: X509TrustManager? = null
  37. private val log by logger()
  38. private lateinit var socket: Socket
  39. private lateinit var readChannel: ByteReadChannel
  40. private lateinit var writeChannel: ByteWriteChannel
  41. @Suppress("EXPERIMENTAL_API_USAGE")
  42. override suspend fun connect() {
  43. log.info { "Connecting..." }
  44. socket = aSocket(ActorSelectorManager(Dispatchers.IO)).tcp().connect(InetSocketAddress(host, port))
  45. if (tls) {
  46. // TODO: Figure out how exactly scopes work...
  47. socket = socket.tls(GlobalScope.coroutineContext, randomAlgorithm = SecureRandom.getInstanceStrong().algorithm, trustManager = tlsTrustManager)
  48. }
  49. readChannel = socket.openReadChannel()
  50. writeChannel = socket.openWriteChannel()
  51. }
  52. override fun disconnect() {
  53. log.info { "Disconnecting..." }
  54. socket.close()
  55. }
  56. override suspend fun sendLine(line: ByteArray, offset: Int, length: Int) {
  57. with (writeChannel) {
  58. log.fine { ">>> ${String(line, offset, length)}" }
  59. writeAvailable(line, offset, length)
  60. writeByte(CARRIAGE_RETURN)
  61. writeByte(LINE_FEED)
  62. flush()
  63. }
  64. }
  65. override suspend fun sendLine(line: String) = sendLine(line.toByteArray())
  66. @ExperimentalCoroutinesApi
  67. override fun readLines(coroutineScope: CoroutineScope) = coroutineScope.produce {
  68. val lineBuffer = ByteArray(4096)
  69. var index = 0
  70. while (!readChannel.isClosedForRead) {
  71. var start = index
  72. val count = readChannel.readAvailable(lineBuffer, index, lineBuffer.size - index)
  73. for (i in index until index + count) {
  74. if (lineBuffer[i] == CARRIAGE_RETURN || lineBuffer[i] == LINE_FEED) {
  75. if (start < i) {
  76. val line = lineBuffer.sliceArray(start until i)
  77. log.fine { "<<< ${String(line)}" }
  78. send(line)
  79. }
  80. start = i + 1
  81. }
  82. }
  83. lineBuffer.copyInto(lineBuffer, 0, start)
  84. index = count + index - start
  85. }
  86. }
  87. }