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.6KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101
  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. fun readLines(coroutineScope: CoroutineScope): ReceiveChannel<ByteArray>
  24. suspend fun writeLines(channel: ReceiveChannel<ByteArray>)
  25. }
  26. /**
  27. * Asynchronous socket that buffers incoming data and emits individual lines.
  28. */
  29. // TODO: Expose advanced TLS options
  30. internal class KtorLineBufferedSocket(private val host: String, private val port: Int, private val tls: Boolean = false): LineBufferedSocket {
  31. companion object {
  32. const val CARRIAGE_RETURN = '\r'.toByte()
  33. const val LINE_FEED = '\n'.toByte()
  34. }
  35. var tlsTrustManager: X509TrustManager? = null
  36. private val log by logger()
  37. private lateinit var socket: Socket
  38. private lateinit var readChannel: ByteReadChannel
  39. private lateinit var writeChannel: ByteWriteChannel
  40. @Suppress("EXPERIMENTAL_API_USAGE")
  41. override suspend fun connect() {
  42. log.info { "Connecting..." }
  43. socket = aSocket(ActorSelectorManager(Dispatchers.IO)).tcp().connect(InetSocketAddress(host, port))
  44. if (tls) {
  45. // TODO: Figure out how exactly scopes work...
  46. socket = socket.tls(GlobalScope.coroutineContext, randomAlgorithm = SecureRandom.getInstanceStrong().algorithm, trustManager = tlsTrustManager)
  47. }
  48. readChannel = socket.openReadChannel()
  49. writeChannel = socket.openWriteChannel()
  50. }
  51. override fun disconnect() {
  52. log.info { "Disconnecting..." }
  53. socket.close()
  54. }
  55. @ExperimentalCoroutinesApi
  56. override fun readLines(coroutineScope: CoroutineScope) = coroutineScope.produce {
  57. val lineBuffer = ByteArray(4096)
  58. var index = 0
  59. while (!readChannel.isClosedForRead) {
  60. var start = index
  61. val count = readChannel.readAvailable(lineBuffer, index, lineBuffer.size - index)
  62. for (i in index until index + count) {
  63. if (lineBuffer[i] == CARRIAGE_RETURN || lineBuffer[i] == LINE_FEED) {
  64. if (start < i) {
  65. val line = lineBuffer.sliceArray(start until i)
  66. log.fine { "<<< ${String(line)}" }
  67. send(line)
  68. }
  69. start = i + 1
  70. }
  71. }
  72. lineBuffer.copyInto(lineBuffer, 0, start)
  73. index = count + index - start
  74. }
  75. }
  76. override suspend fun writeLines(channel: ReceiveChannel<ByteArray>) {
  77. for (line in channel) {
  78. with(writeChannel) {
  79. log.fine { ">>> ${String(line)}" }
  80. writeAvailable(line, 0, line.size)
  81. writeByte(CARRIAGE_RETURN)
  82. writeByte(LINE_FEED)
  83. flush()
  84. }
  85. }
  86. }
  87. }