zio.http.netty.NettyBody.scala Maven / Gradle / Ivy

Go to download
/*
 * Copyright 2021 - 2023 Sporta Technologies PVT LTD & the ZIO HTTP contributors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package zio.http.netty

import java.nio.charset.Charset

import zio._
import zio.stacktracer.TracingImplicits.disableAutoTrace

import zio.stream.{Take, ZChannel, ZStream}

import zio.http.Body.UnsafeBytes
import zio.http.internal.BodyEncoding
import zio.http.{Body, Header}

import io.netty.buffer.{ByteBuf, ByteBufUtil}
import io.netty.util.AsciiString

object NettyBody extends BodyEncoding {

  /**
   * Helper to create Body from AsciiString
   */
  def fromAsciiString(asciiString: AsciiString): Body = AsciiStringBody(asciiString)

  private[zio] def fromAsync(
    unsafeAsync: UnsafeAsync => Unit,
    knownContentLength: Option[Long],
    contentTypeHeader: Option[Header.ContentType] = None,
  ): Body = {
    AsyncBody(
      unsafeAsync,
      knownContentLength,
      contentTypeHeader.map(Body.ContentType.fromHeader),
    )
  }

  /**
   * Helper to create Body from ByteBuf
   */
  private[zio] def fromByteBuf(byteBuf: ByteBuf, contentTypeHeader: Option[Header.ContentType]): Body = {
    if (byteBuf.readableBytes() == 0) Body.EmptyBody
    else {
      Body.ArrayBody(ByteBufUtil.getBytes(byteBuf), contentTypeHeader.map(Body.ContentType.fromHeader))
    }
  }

  override def fromCharSequence(charSequence: CharSequence, charset: Charset): Body =
    fromAsciiString(new AsciiString(charSequence, charset))

  private[zio] final case class AsciiStringBody(
    asciiString: AsciiString,
    override val contentType: Option[Body.ContentType] = None,
  ) extends UnsafeBytes {

    override def asArray(implicit trace: Trace): Task[Array[Byte]] = ZIO.succeed(asciiString.array())

    override def isComplete: Boolean = true

    override def isEmpty: Boolean = asciiString.isEmpty()

    override def asChunk(implicit trace: Trace): Task[Chunk[Byte]] =
      ZIO.succeed(Chunk.fromArray(asciiString.array()))

    override def asStream(implicit trace: Trace): ZStream[Any, Throwable, Byte] =
      ZStream.unwrap(asChunk.map(ZStream.fromChunk(_)))

    override def toString: String = s"Body.fromAsciiString($asciiString)"

    private[zio] override def unsafeAsArray(implicit unsafe: Unsafe): Array[Byte] = asciiString.array()

    override def contentType(newContentType: Body.ContentType): Body = copy(contentType = Some(newContentType))

    override def knownContentLength: Option[Long] = Some(asciiString.length().toLong)
  }

  private[zio] final case class AsyncBody(
    unsafeAsync: UnsafeAsync => Unit,
    knownContentLength: Option[Long],
    override val contentType: Option[Body.ContentType] = None,
  ) extends Body {

    override def asArray(implicit trace: Trace): Task[Array[Byte]] = asChunk.map {
      case b: Chunk.ByteArray => b.array
      case other              => other.toArray
    }

    override def asChunk(implicit trace: Trace): Task[Chunk[Byte]] =
      ZIO.async { cb =>
        try {
          // Cap at 100kB as a precaution in case the server sends an invalid content length
          unsafeAsync(UnsafeAsync.Aggregating(bufferSize(1024 * 100))(cb))
        } catch {
          case e: Throwable => cb(ZIO.fail(e))
        }
      }

    override def asStream(implicit trace: Trace): ZStream[Any, Throwable, Byte] = {
      asyncUnboundedStream[Any, Throwable, Byte](emit =>
        try {
          unsafeAsync(new UnsafeAsync.Streaming(emit))
        } catch {
          case e: Throwable => emit(ZIO.fail(Option(e)))
        },
      )
    }

    // No need to create a large buffer when we know the response is small
    private[this] def bufferSize(maxSize: Int): Int = {
      val cl = knownContentLength.getOrElse(4096L)
      if (cl <= 16L) 16
      else if (cl >= maxSize) maxSize
      else Integer.highestOneBit(cl.toInt - 1) << 1 // Round to next power of 2
    }

    override def isComplete: Boolean = false

    override def isEmpty: Boolean = false

    override def toString: String = s"AsyncBody($unsafeAsync)"

    override def contentType(newContentType: Body.ContentType): Body = copy(contentType = Some(newContentType))
  }

  /**
   * Code ported from zio.stream to use an unbounded queue
   */
  private def asyncUnboundedStream[R, E, A](
    register: ZStream.Emit[R, E, A, Unit] => Unit,
  )(implicit trace: Trace): ZStream[R, E, A] =
    ZStream.unwrapScoped[R](for {
      queue   <- ZIO.acquireRelease(Queue.unbounded[Take[E, A]])(_.shutdown)
      runtime <- ZIO.runtime[R]
    } yield {
      val rtm = runtime.unsafe
      register { k =>
        try {
          rtm
            .run(Take.fromPull(k).flatMap(queue.offer))(trace, Unsafe)
            .getOrThrowFiberFailure()(Unsafe)
          ()
        } catch {
          case FiberFailure(c) if c.isInterrupted =>
        }
      }

      lazy val loop: ZChannel[Any, Any, Any, Any, E, Chunk[A], Unit] =
        ZChannel.unwrap(
          queue.take
            .flatMap(_.done)
            .fold(
              maybeError =>
                ZChannel.fromZIO(queue.shutdown) *>
                  maybeError.fold[ZChannel[Any, Any, Any, Any, E, Chunk[A], Unit]](ZChannel.unit)(ZChannel.fail(_)),
              a => ZChannel.write(a) *> loop,
            ),
        )

      ZStream.fromChannel(loop)
    })

  private[zio] trait UnsafeAsync {
    def apply(message: Chunk[Byte], isLast: Boolean): Unit
    def fail(cause: Throwable): Unit
  }

  private[zio] object UnsafeAsync {
    private val FailNone = Exit.fail(None)

    final case class Aggregating(bufferInitialSize: Int)(callback: Task[Chunk[Byte]] => Unit)(implicit trace: Trace)
        extends UnsafeAsync {

      def apply(message: Chunk[Byte], isLast: Boolean): Unit = {
        assert(isLast)
        callback(Exit.succeed(message))
      }

      def fail(cause: Throwable): Unit =
        callback(ZIO.fail(cause))
    }

    final class Streaming(emit: ZStream.Emit[Any, Throwable, Byte, Unit])(implicit trace: Trace) extends UnsafeAsync {
      def apply(message: Chunk[Byte], isLast: Boolean): Unit = {
        if (message.nonEmpty) emit(Exit.succeed(message))
        if (isLast) emit(FailNone)
      }

      def fail(cause: Throwable): Unit =
        emit(ZIO.fail(Some(cause)))
    }
  }
}