1package com.peachpasswords.sync23import java.io.InputStream4import java.io.OutputStream5import java.net.Socket6import java.security.MessageDigest7import java.util.Base648910111213141516171819202122232425private const val WS_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"2627sealed interface WsMessage {28 data class Text(val text: String) : WsMessage29 data class Binary(val bytes: ByteArray) : WsMessage30 data object Closed : WsMessage31}3233class WebSocketConnection(34 private val socket: Socket,35 private val input: InputStream,36 private val output: OutputStream,37) {38 3940414243 fun accept(): Boolean {44 val request = readHttpHeaders() ?: return false45 val key = request.entries.firstOrNull { it.key.equals("sec-websocket-key", true) }?.value46 ?: return false47 val accept = Base64.getEncoder().encodeToString(48 MessageDigest.getInstance("SHA-1").digest((key + WS_GUID).toByteArray(Charsets.US_ASCII)),49 )50 output.write(51 (52 "HTTP/1.1 101 Switching Protocols\r\n" +53 "Upgrade: websocket\r\n" +54 "Connection: Upgrade\r\n" +55 "Sec-WebSocket-Accept: $accept\r\n\r\n"56 ).toByteArray(Charsets.US_ASCII),57 )58 output.flush()59 return true60 }6162 fun sendText(text: String) = writeFrame(0x1, text.toByteArray(Charsets.UTF_8))6364 fun send(bytes: ByteArray) = writeFrame(0x2, bytes)6566 67 fun receive(): WsMessage {68 while (true) {69 val b0 = input.read()70 if (b0 < 0) return WsMessage.Closed71 val fin = (b0 and 0x80) != 072 when (val opcode = b0 and 0x0f) {73 0x8 -> return WsMessage.Closed74 0x9 -> { writeFrame(0xA, readPayload()); continue } 75 0xA -> { readPayload(); continue } 76 0x1, 0x2 -> {77 if (!fin) return WsMessage.Closed78 val payload = readPayload()79 return if (opcode == 0x1) WsMessage.Text(payload.toString(Charsets.UTF_8))80 else WsMessage.Binary(payload)81 }82 else -> return WsMessage.Closed83 }84 }85 }8687 888990919293 fun closeSilently() {94 runCatching { socket.close() }95 }9697 private fun readHttpHeaders(): Map<String, String>? {98 val head = StringBuilder()99 100 while (head.length < 8192 && !head.endsWith("\r\n\r\n")) {101 val c = input.read()102 if (c < 0) return null103 head.append(c.toChar())104 }105 val lines = head.toString().split("\r\n")106 if (lines.firstOrNull()?.startsWith("GET ") != true) return null107 return lines.drop(1)108 .mapNotNull { line ->109 val i = line.indexOf(':')110 if (i <= 0) null else line.substring(0, i).trim() to line.substring(i + 1).trim()111 }112 .toMap()113 }114115 private fun readPayload(): ByteArray {116 val b1 = input.read()117 if (b1 < 0) return ByteArray(0)118 val masked = (b1 and 0x80) != 0119 var len = (b1 and 0x7f).toLong()120 when (len) {121 126L -> len = ((readByte() shl 8) or readByte()).toLong()122 127L -> {123 len = 0124 repeat(8) { len = (len shl 8) or readByte().toLong() }125 }126 }127 128 129 require(len in 0..(16L * 1024 * 1024)) { "ws frame too large: $len" }130131 val mask = if (masked) ByteArray(4) { readByte().toByte() } else null132 val payload = ByteArray(len.toInt())133 var read = 0134 while (read < payload.size) {135 val n = input.read(payload, read, payload.size - read)136 if (n < 0) break137 read += n138 }139 if (mask != null) {140 for (i in payload.indices) payload[i] = (payload[i].toInt() xor mask[i % 4].toInt()).toByte()141 }142 return payload143 }144145 private fun readByte(): Int {146 val b = input.read()147 if (b < 0) throw java.io.EOFException("websocket closed mid-frame")148 return b149 }150151 152 private fun writeFrame(opcode: Int, payload: ByteArray) {153 output.write(0x80 or opcode)154 when {155 payload.size < 126 -> output.write(payload.size)156 payload.size <= 0xffff -> {157 output.write(126)158 output.write((payload.size shr 8) and 0xff)159 output.write(payload.size and 0xff)160 }161 else -> {162 output.write(127)163 repeat(4) { output.write(0) }164 output.write((payload.size shr 24) and 0xff)165 output.write((payload.size shr 16) and 0xff)166 output.write((payload.size shr 8) and 0xff)167 output.write(payload.size and 0xff)168 }169 }170 output.write(payload)171 output.flush()172 }173}