Guest User

Untitled

a guest
Nov 18th, 2018
81
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. package jashkasoft.ithink
  2.  
  3. import java.nio.ByteBuffer
  4.  
  5. enum class FrameType (val controlFrame: Boolean, val opcode: Int) {
  6.     /**
  7.      * Regular application level text frame
  8.      */
  9.     TEXT(false, 1),
  10.  
  11.     /**
  12.      * Regular application level binary frame
  13.      */
  14.     BINARY(false, 2),
  15.  
  16.     /**
  17.      * Low level close frame type
  18.      */
  19.     CLOSE(true, 8),
  20.  
  21.     /**
  22.      * Low level ping frame type
  23.      */
  24.     PING(true, 9),
  25.  
  26.     /**
  27.      * Low level pong frame type
  28.      */
  29.     PONG(true, 0xa);
  30.  
  31.     companion object {
  32.         private val maxOpcode = FrameType.values().maxBy { it.opcode }!!.opcode
  33.  
  34.         private val byOpcodeArray = Array(maxOpcode + 1) { op -> FrameType.values().singleOrNull { it.opcode == op } }
  35.  
  36.         /**
  37.          * Find [FrameType] instance by numeric [opcode]
  38.          * @return a [FrameType] instance or `null` of the [opcode] value is not valid
  39.          */
  40.         operator fun get(opcode: Int): FrameType? = if (opcode in 0..maxOpcode) byOpcodeArray[opcode] else null
  41.     }
  42. }
  43.  
  44.  
  45. sealed class Frame(val fin: Boolean, val frameType: FrameType, val buffer: ByteBuffer) {
  46.     private val initialSize = buffer.remaining()
  47.  
  48.     /**
  49.      * Represents an application level binary frame.
  50.      * In a RAW web socket session a big text frame could be fragmented
  51.      * (separated into several text frames so they have [fin] = false except the last one).
  52.      * Note that usually there is no need to handle fragments unless you have a RAW web socket session.
  53.      */
  54.     class Binary(fin: Boolean, buffer: ByteBuffer) : Frame(fin, FrameType.BINARY, buffer)
  55.  
  56.     override fun toString() = "Frame $frameType (fin=$fin, buffer len = $initialSize)"
  57. }
  58.  
  59.  
  60. fun ByteBuffer.copy(size: Int = remaining()): ByteBuffer {
  61.     return ByteBuffer.allocate(size).apply {
  62.         this@copy.slice().moveTo(this@apply)
  63.         clear()
  64.     }
  65. }
  66.  
  67. fun ByteBuffer.moveTo(destination: ByteBuffer, limit: Int = Int.MAX_VALUE): Int {
  68.     val size = minOf(limit, remaining(), destination.remaining())
  69.     if (size == remaining()) {
  70.         destination.put(this)
  71.     } else {
  72.         val l = limit()
  73.         limit(position() + size)
  74.         destination.put(this)
  75.         limit(l)
  76.     }
  77.     return size
  78. }
  79.  
  80. fun Boolean.flagAt(at: Int) = if (this) 1 shl at else 0
  81.  
  82. infix fun Byte.xor(other: Byte) = toInt().xor(other.toInt()).toByte()
  83.  
  84. fun ByteBuffer.xor(other: ByteBuffer) {
  85.     val bb = slice()
  86.     val mask = other.slice()
  87.     val maskSize = mask.remaining()
  88.  
  89.     for (i in 0 until bb.remaining()) {
  90.         bb.put(i, bb.get(i) xor mask[i % maskSize])
  91.     }
  92. }
  93.  
  94. class Serializer {
  95.  
  96.     private var frameBody: ByteBuffer? = null
  97.     private var maskBuffer: ByteBuffer? = null
  98.  
  99.     var masking: Boolean = false
  100.  
  101.     fun serialize(frame: Frame, buffer: ByteBuffer) {
  102.         while (writeCurrentPayload(buffer)) {
  103.             val mask = masking
  104.             setMaskBuffer(mask)
  105.  
  106.             val headerSize = estimateFrameHeaderSize(frame, mask)
  107.             if (buffer.remaining() < headerSize) {
  108.                 break
  109.             }
  110.  
  111.             serializeHeader(frame, buffer, mask)
  112.             frameBody = frame.buffer.maskedIfNeeded()
  113.  
  114.             frameBody!!.array().iterator().forEach { print("$it ") }
  115.             println()
  116.         }
  117.     }
  118.  
  119.     private fun serializeHeader(frame: Frame, buffer: ByteBuffer, mask: Boolean) {
  120.         val size = frame.buffer.remaining()
  121.         val length1 = when {
  122.             size < 126 -> size
  123.             size <= 0xffff -> 126
  124.             else -> 127
  125.         }
  126.  
  127.         buffer.put(
  128.                 (frame.fin.flagAt(7) or frame.frameType.opcode).toByte()
  129.         )
  130.         buffer.put(
  131.                 (mask.flagAt(7) or length1).toByte()
  132.         )
  133.  
  134.         if (length1 == 126) {
  135.             buffer.putShort(frame.buffer.remaining().toShort())
  136.         } else if (length1 == 127) {
  137.             buffer.putLong(frame.buffer.remaining().toLong())
  138.         }
  139.  
  140.         maskBuffer?.duplicate()?.moveTo(buffer)
  141.     }
  142.  
  143.     private fun estimateFrameHeaderSize(f: Frame, mask: Boolean): Int {
  144.         val size = f.buffer.remaining()
  145.         return when {
  146.             size < 126 -> 2
  147.             size <= Short.MAX_VALUE -> 2 + 2
  148.             else -> 2 + 8
  149.         } + maskSize(mask)
  150.     }
  151.  
  152.  
  153.     private fun writeCurrentPayload(buffer: ByteBuffer): Boolean {
  154.         val frame = frameBody ?: return true
  155.         frame.moveTo(buffer)
  156.         if (!frame.hasRemaining()) {
  157.             frameBody = null
  158.             return true
  159.         }
  160.  
  161.         return false
  162.     }
  163.  
  164.     private fun maskSize(mask: Boolean) = if (mask) 4 else 0
  165.  
  166.     private fun ByteBuffer.maskedIfNeeded() = maskBuffer?.let { mask -> copy().apply { xor(mask) } } ?: this
  167.  
  168.     private fun setMaskBuffer(mask: Boolean) {
  169.         if (mask) {
  170.             maskBuffer = ByteBuffer.allocate(4).apply {
  171.                 putInt(2.hashCode())
  172.                 clear()
  173.             }
  174.         } else {
  175.             maskBuffer = null
  176.         }
  177.     }
  178. }
  179.  
  180. fun main(args: Array<String>) {
  181.  
  182.     Serializer().apply {
  183.         val content = "hello".toByteArray()
  184.         val content1 = "world".toByteArray()
  185.  
  186.         val sendBuffer = ByteBuffer.allocate(content.size + 2)
  187.         val sendBuffer1 = ByteBuffer.allocate(content1.size + 2)
  188.         serialize(Frame.Binary(false, ByteBuffer.wrap(content)), sendBuffer)
  189.         serialize(Frame.Binary(true, ByteBuffer.wrap(content1)), sendBuffer1)
  190.  
  191.         sendBuffer.flip()
  192.         sendBuffer1.flip()
  193.  
  194.         println(sendBuffer.array().iterator().forEach { print("$it ") })
  195.         println(sendBuffer1.array().iterator().forEach { print("$it ") })
  196.     }
  197. }
RAW Paste Data