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) {65 try {66 requireSyncFrameSize(bytes.size)67 } catch (error: SyncFrameTooLargeException) {68 close(1009, "sync frame too large")69 throw error70 }71 writeFrame(0x2, bytes)72 }7374 75 fun receive(): WsMessage {76 while (true) {77 val b0 = input.read()78 if (b0 < 0) return WsMessage.Closed79 val fin = (b0 and 0x80) != 080 when (val opcode = b0 and 0x0f) {81 0x8 -> {82 val payload = readPayload()83 val code = if (payload.size >= 2) {84 ((payload[0].toInt() and 0xff) shl 8) or (payload[1].toInt() and 0xff)85 } else {86 null87 }88 if (code == 1009) throw SyncFrameTooLargeException()89 return WsMessage.Closed90 }91 0x9 -> { writeFrame(0xA, readPayload()); continue } 92 0xA -> { readPayload(); continue } 93 0x1, 0x2 -> {94 if (!fin) return WsMessage.Closed95 val payload = readPayload()96 return if (opcode == 0x1) WsMessage.Text(payload.toString(Charsets.UTF_8))97 else WsMessage.Binary(payload)98 }99 else -> return WsMessage.Closed100 }101 }102 }103104 105106107108109110 fun closeSilently() {111 runCatching { socket.close() }112 }113114 private fun readHttpHeaders(): Map<String, String>? {115 val head = StringBuilder()116 117 while (head.length < 8192 && !head.endsWith("\r\n\r\n")) {118 val c = input.read()119 if (c < 0) return null120 head.append(c.toChar())121 }122 val lines = head.toString().split("\r\n")123 if (lines.firstOrNull()?.startsWith("GET ") != true) return null124 return lines.drop(1)125 .mapNotNull { line ->126 val i = line.indexOf(':')127 if (i <= 0) null else line.substring(0, i).trim() to line.substring(i + 1).trim()128 }129 .toMap()130 }131132 private fun readPayload(): ByteArray {133 val b1 = input.read()134 if (b1 < 0) return ByteArray(0)135 val masked = (b1 and 0x80) != 0136 var len = (b1 and 0x7f).toLong()137 when (len) {138 126L -> len = ((readByte() shl 8) or readByte()).toLong()139 127L -> {140 len = 0141 repeat(8) { len = (len shl 8) or readByte().toLong() }142 }143 }144 145 146 if (len !in 0..MAX_SYNC_FRAME_BYTES.toLong()) {147 close(1009, "sync frame too large")148 throw SyncFrameTooLargeException("sync frame exceeds $MAX_SYNC_FRAME_BYTES bytes")149 }150151 val mask = if (masked) ByteArray(4) { readByte().toByte() } else null152 val payload = ByteArray(len.toInt())153 var read = 0154 while (read < payload.size) {155 val n = input.read(payload, read, payload.size - read)156 if (n < 0) break157 read += n158 }159 if (mask != null) {160 for (i in payload.indices) payload[i] = (payload[i].toInt() xor mask[i % 4].toInt()).toByte()161 }162 return payload163 }164165 private fun readByte(): Int {166 val b = input.read()167 if (b < 0) throw java.io.EOFException("websocket closed mid-frame")168 return b169 }170171 fun close(code: Int = 1000, reason: String = "") {172 val reasonBytes = reason.toByteArray(Charsets.UTF_8).take(123).toByteArray()173 val payload = byteArrayOf((code ushr 8).toByte(), code.toByte()) + reasonBytes174 runCatching { writeFrame(0x8, payload) }175 runCatching { socket.close() }176 }177178 179 private fun writeFrame(opcode: Int, payload: ByteArray) {180 output.write(0x80 or opcode)181 when {182 payload.size < 126 -> output.write(payload.size)183 payload.size <= 0xffff -> {184 output.write(126)185 output.write((payload.size shr 8) and 0xff)186 output.write(payload.size and 0xff)187 }188 else -> {189 output.write(127)190 repeat(4) { output.write(0) }191 output.write((payload.size shr 24) and 0xff)192 output.write((payload.size shr 16) and 0xff)193 output.write((payload.size shr 8) and 0xff)194 output.write(payload.size and 0xff)195 }196 }197 output.write(payload)198 output.flush()199 }200}