PacketEncoding.kt

package com.proximyst.ussrm.packet.lowlevel

import com.proximyst.ussrm.packet.lowlevel.data.VarInt
import io.netty.buffer.ByteBuf
import io.netty.buffer.Unpooled
import java.io.ByteArrayOutputStream
import java.util.zip.DeflaterOutputStream
import java.util.zip.Inflater

/**
 * Defines different encodings to use for packet handling.
 */
enum class PacketEncoding {
  /**
   * Handle the packet with no compression whatsoever.
   */
  UNCOMPRESSED {
    override fun decode(data : ByteBuf) : ByteBuf {
      val length = VarInt.decode(data)?.contained() ?: return Unpooled.buffer()
      return data.readBytes(length)
    }

    override fun encode(id : VarInt, data : ByteBuf) : ByteBuf {
      return id.encode().writeBytes(data)
    }
  },

  /**
   * Conform to the compression format, but handle the packet with no compression.
   * This is used when it doesn't hit the threshold for compression or when
   * the incoming packet is 100% sure to not be compressed.
   */
  UNCOMPRESSED_WITH_COMPRESSION_ENABLED {
    override fun decode(data : ByteBuf) : ByteBuf {
      var length = VarInt.decode(data)?.contained() ?: return Unpooled.buffer()
      val beforeDataLength = data.readableBytes()
      VarInt.decode(data)
      length -= beforeDataLength - data.readableBytes()

      return data.readBytes(length)
    }

    override fun encode(id : VarInt, data : ByteBuf) : ByteBuf {
      val buf = id.encode()
      buf.writeBytes(VarInt(0).encode())
      return buf.writeBytes(data)
    }
  },

  /**
   * Handles the packet with ZLIB compression.
   *
   * The only time it handles it without compression is decoding packets
   * which have an uncompressed length of 0.
   */
  COMPRESSED {
    override fun decode(data : ByteBuf) : ByteBuf {
      var length = VarInt.decode(data)?.contained() ?: return Unpooled.buffer()
      val beforeDataLength = data.readableBytes()
      val uncompressedLength = VarInt.decode(data)?.contained() ?: return Unpooled.buffer()
      length -= beforeDataLength - data.readableBytes()
      val compressed = data.readBytes(length)
      if (uncompressedLength == 0)
        return compressed

      val inflater = Inflater()
      inflater.setInput(compressed.array())
      val result = ByteArray(uncompressedLength)
      if (inflater.inflate(result) != uncompressedLength)
        throw IllegalArgumentException("the uncompressed data is not correctly labelled")

      return Unpooled.copiedBuffer(result)
    }

    override fun encode(id : VarInt, data : ByteBuf) : ByteBuf {
      val encodedId = id.encode().array()
      val size = data.readableBytes() + encodedId.size

      val compressedData = ByteArrayOutputStream().also {
        it.write(size)
        DeflaterOutputStream(it).use {
          it.write(encodedId)
          it.write(data.array())
        }
      }.toByteArray()

      return VarInt(compressedData.size).encode().writeBytes(compressedData)
    }
  },
  ;

  /**
   * Encodes the [data] together with the [id] to a [ByteBuf] per the enum constant.
   *
   * @param id
   *    The packet ID of the packet.
   *
   * @param data
   *    The data of the packet.
   *
   * @return
   *    The encoded [ByteBuf], which may be ZLIB compressed per the enum constant.
   */
  abstract fun encode(id : VarInt, data : ByteBuf) : ByteBuf

  /**
   * Decodes the [data] from the [ByteBuf] and returns its ID and data as [ByteBuf].
   *
   * @throws IllegalArgumentException
   *    If the uncompressed data is not correctly labelled, this is thrown.
   *
   * @param data
   *    The [ByteBuf] of the incoming packet. Its readableBytes will be modified.
   *
   * @return
   *    The [ByteBuf] of the ID and data of the enclosed packet.
   */
  abstract fun decode(data : ByteBuf) : ByteBuf
}