Create a buffer dynamically for read operations

vyfor · Apr 6, 2024 · 99bbef4 · 99bbef4
2 parents 93199bd + b694fcf
commit 99bbef4
Show file tree

Hide file tree

Showing 6 changed files with 101 additions and 49 deletions.
diff --git a/src/commonMain/kotlin/io/github/vyfor/kpresence/ipc/Connection.kt b/src/commonMain/kotlin/io/github/vyfor/kpresence/ipc/Connection.kt
@@ -2,7 +2,7 @@ package io.github.vyfor.kpresence.ipc
 
 expect class Connection() {
   fun open()
-  fun read(bufferSize: Int = 4096): ByteArray
+  fun read(): ByteArray
   fun write(opcode: Int, data: String)
   fun close()
 }
diff --git a/src/commonMain/kotlin/io/github/vyfor/kpresence/utils/Utils.kt b/src/commonMain/kotlin/io/github/vyfor/kpresence/utils/Utils.kt
@@ -14,6 +14,15 @@ internal fun Int.reverseBytes(): Int {
     (this and 0x000000ff shl 24)
 }
 
+internal fun ByteArray.byteArrayToInt(): Int {
+  require(size == 4) { "ByteArray size must be 4" }
+
+  return (get(0).toInt() and 0xFF shl 24) or
+    (get(1).toInt() and 0xFF shl 16) or
+    (get(2).toInt() and 0xFF shl 8) or
+    (get(3).toInt() and 0xFF)
+}
+
 expect fun epochMillis(): Long
 
 expect fun getProcessId(): Int
diff --git a/src/jvmMain/kotlin/io/github/vyfor/kpresence/ipc/Connection.kt b/src/jvmMain/kotlin/io/github/vyfor/kpresence/ipc/Connection.kt
@@ -2,6 +2,7 @@ package io.github.vyfor.kpresence.ipc
 
 import io.github.vyfor.kpresence.utils.putInt
 import io.github.vyfor.kpresence.utils.reverseBytes
+import java.io.EOFException
 import java.io.InputStream
 import java.io.OutputStream
 import java.io.RandomAccessFile
@@ -21,8 +22,8 @@ actual class Connection {
     con.open()
   }
 
-  actual fun read(bufferSize: Int): ByteArray {
-    return con.read(bufferSize)
+  actual fun read(): ByteArray {
+    return con.read()
   }
 
   actual fun write(opcode: Int, data: String) {
@@ -35,7 +36,7 @@ actual class Connection {
 
   interface IConnection {
     fun open()
-    fun read(bufferSize: Int): ByteArray
+    fun read(): ByteArray
     fun write(opcode: Int, data: String)
     fun close()
   }
@@ -54,10 +55,13 @@ actual class Connection {
       throw RuntimeException("Could not connect to the pipe!")
     }
 
-    override fun read(bufferSize: Int): ByteArray {
+    override fun read(): ByteArray {
       pipe?.let { stream ->
-        val buffer = ByteArray(bufferSize)
-        stream.read(buffer, 0, bufferSize)
+        stream.readInt()
+        val length = stream.readInt().reverseBytes()
+        val buffer = ByteArray(length)
+
+        stream.read(buffer, 0, length)
         return buffer
       } ?: throw IllegalStateException("Not connected")
     }
@@ -105,10 +109,13 @@ actual class Connection {
       throw RuntimeException("Could not connect to the pipe!")
     }
 
-    override fun read(bufferSize: Int): ByteArray {
+    override fun read(): ByteArray {
       inputStream?.let { stream ->
-        val buffer = ByteArray(bufferSize)
-        stream.read(buffer, 0, bufferSize)
+        stream.readInt()
+        val length = stream.readInt()
+        val buffer = ByteArray(length)
+
+        stream.read(buffer, 0, length)
         return buffer
       } ?: throw IllegalStateException("Not connected")
     }
@@ -131,5 +138,17 @@ actual class Connection {
       inputStream?.close()
       outputStream?.close()
     }
+
+    private fun InputStream.readInt(): Int {
+      val ch1: Int = read()
+      val ch2: Int = read()
+      val ch3: Int = read()
+      val ch4: Int = read()
+      return if (ch1 or ch2 or ch3 or ch4 < 0) {
+        throw EOFException()
+      } else {
+        (ch1 shl 24) + (ch2 shl 16) + (ch3 shl 8) + (ch4 shl 0)
+      }
+    }
   }
 }
diff --git a/src/linuxMain/kotlin/io/github/vyfor/kpresence/ipc/Connection.kt b/src/linuxMain/kotlin/io/github/vyfor/kpresence/ipc/Connection.kt
@@ -1,5 +1,6 @@
 package io.github.vyfor.kpresence.ipc
 
+import io.github.vyfor.kpresence.utils.byteArrayToInt
 import io.github.vyfor.kpresence.utils.putInt
 import io.github.vyfor.kpresence.utils.reverseBytes
 import kotlinx.cinterop.*
@@ -39,23 +40,15 @@ actual class Connection {
     throw RuntimeException("Could not connect to the pipe!")
   }
 
-  actual fun read(bufferSize: Int): ByteArray {
+  actual fun read(): ByteArray {
     if (pipe == -1) throw IllegalStateException("Not connected")
 
-    val buffer = ByteArray(bufferSize)
+    readBytes(4)
+    val length = readBytes(4).first.byteArrayToInt().reverseBytes()
+    val buffer = ByteArray(length)
+    val bytesRead = readBytes(4).second
 
-    val bytesRead = recv(pipe, buffer.refTo(0), bufferSize.convert(), 0).toInt()
-    if (bytesRead < 0) {
-      if (errno == EAGAIN) return buffer
-
-      throw RuntimeException("Error reading from socket")
-    } else if (bytesRead == 0) {
-      close()
-
-      throw RuntimeException("Connection closed")
-    }
-
-    return buffer.copyOf(bytesRead)
+    return buffer.copyOf(bytesRead.toInt())
   }
 
   actual fun write(opcode: Int, data: String) {
@@ -79,4 +72,18 @@ actual class Connection {
   actual fun close() {
     close(pipe)
   }
+
+  private fun readBytes(size: Int): Pair<ByteArray, Long> {
+    val bytes = ByteArray(size)
+    recv(pipe, bytes.refTo(0), bytes.size.convert(), 0).let { bytesRead ->
+      if (bytesRead < 0L) {
+        if (errno == EAGAIN) return bytes to 0
+        throw RuntimeException("Error reading from socket")
+      } else if (bytesRead == 0L) {
+        close()
+        throw RuntimeException("Connection closed")
+      }
+      return bytes to bytesRead
+    }
+  }
 }
diff --git a/src/macosMain/kotlin/io/github/vyfor/kpresence/ipc/Connection.kt b/src/macosMain/kotlin/io/github/vyfor/kpresence/ipc/Connection.kt
@@ -1,5 +1,6 @@
 package io.github.vyfor.kpresence.ipc
 
+import io.github.vyfor.kpresence.utils.byteArrayToInt
 import io.github.vyfor.kpresence.utils.putInt
 import io.github.vyfor.kpresence.utils.reverseBytes
 import kotlinx.cinterop.*
@@ -39,23 +40,15 @@ actual class Connection {
     throw RuntimeException("Could not connect to the pipe!")
   }
 
-  actual fun read(bufferSize: Int): ByteArray {
+  actual fun read(): ByteArray {
     if (pipe == -1) throw IllegalStateException("Not connected")
 
-    val buffer = ByteArray(bufferSize)
+    readBytes(4)
+    val length = readBytes(4).first.byteArrayToInt().reverseBytes()
+    val buffer = ByteArray(length)
+    val bytesRead = readBytes(4).second
 
-    val bytesRead = recv(pipe, buffer.refTo(0), bufferSize.convert(), 0).toInt()
-    if (bytesRead < 0) {
-      if (errno == EAGAIN) return buffer
-
-      throw RuntimeException("Error reading from socket")
-    } else if (bytesRead == 0) {
-      close()
-
-      throw RuntimeException("Connection closed")
-    }
-
-    return buffer.copyOf(bytesRead)
+    return buffer.copyOf(bytesRead.toInt())
   }
 
   actual fun write(opcode: Int, data: String) {
@@ -79,4 +72,18 @@ actual class Connection {
   actual fun close() {
     close(pipe)
   }
+
+  private fun readBytes(size: Int): Pair<ByteArray, Long> {
+    val bytes = ByteArray(size)
+    recv(pipe, bytes.refTo(0), bytes.size.convert(), 0).let { bytesRead ->
+      if (bytesRead < 0L) {
+        if (errno == EAGAIN) return bytes to 0
+        throw RuntimeException("Error reading from socket")
+      } else if (bytesRead == 0L) {
+        close()
+        throw RuntimeException("Connection closed")
+      }
+      return bytes to bytesRead
+    }
+  }
 }
diff --git a/src/mingwX64Main/kotlin/io/github/vyfor/kpresence/ipc/Connection.kt b/src/mingwX64Main/kotlin/io/github/vyfor/kpresence/ipc/Connection.kt
@@ -1,8 +1,11 @@
 package io.github.vyfor.kpresence.ipc
 
+import io.github.vyfor.kpresence.utils.byteArrayToInt
 import io.github.vyfor.kpresence.utils.putInt
 import io.github.vyfor.kpresence.utils.reverseBytes
 import kotlinx.cinterop.*
+import platform.posix.EAGAIN
+import platform.posix.errno
 import platform.windows.*
 
 actual class Connection {
@@ -22,18 +25,14 @@ actual class Connection {
     throw RuntimeException("Could not connect to the pipe!")
   }
 
-  actual fun read(bufferSize: Int): ByteArray = memScoped {
-    pipe?.let { handle ->
-      val buffer = ByteArray(bufferSize)
-      val bytesRead = alloc<UIntVar>()
-
-      val success = buffer.usePinned {
-        ReadFile(handle, it.addressOf(0), bufferSize.convert(), bytesRead.ptr, null)
-      }
-
-      if (success == FALSE) throw RuntimeException("Error reading from socket")
+  actual fun read(): ByteArray = memScoped {
+    pipe?.let { _ ->
+      readBytes(4)
+      val length = readBytes(4).first.byteArrayToInt().reverseBytes()
+      val buffer = ByteArray(length)
+      val bytesRead = readBytes(4).second
 
-      return buffer.copyOf(bytesRead.value.toInt())
+      return buffer.copyOf(bytesRead)
     } ?: throw IllegalStateException("Not connected")
   }
 
@@ -57,4 +56,15 @@ actual class Connection {
   actual fun close() {
     CloseHandle(pipe)
   }
+
+  private fun readBytes(size: Int): Pair<ByteArray, Int> = memScoped {
+    val bytes = ByteArray(size)
+    val bytesRead = alloc<UIntVar>()
+    ReadFile(pipe, bytes.pin().addressOf(0), size.convert(), bytesRead.ptr, null).let { success ->
+      if (success == FALSE) {
+        throw RuntimeException("Error reading from socket")
+      }
+    }
+    return bytes to bytesRead.value.toInt()
+  }
 }