zio.http.multipart.mixed.MultipartMixed.scala Maven / Gradle / Ivy
The newest version!
package zio.http.multipart.mixed
import java.nio.charset.{CharacterCodingException, StandardCharsets}
import scala.annotation.tailrec
import zio.{Chunk, ZIO, ZNothing}
import zio.stream.{ZChannel, ZPipeline, ZStream}
import zio.http._
import zio.http.internal.FormAST
import zio.http.multipart.mixed.MultipartMixed.Parser
final case class MultipartMixed(source: ZStream[Any, Throwable, Byte], boundary: Boundary, bufferSize: Int = 8192) {
def parts: ZStream[Any, Throwable, MultipartMixed.Part] = ZStream
.unwrapScoped[Any] {
source.toChannel.toPull.map { pull =>
new Parser(boundary, pull, bufferSize).result
}
}
}
object MultipartMixed {
final case class Part(headers: Headers, bytes: ZStream[Any, Throwable, Byte]) {
def contentType: Option[Header.ContentType] =
headers.header(Header.ContentType)
def mediaType: Option[MediaType] =
contentType.map(_.mediaType)
// each part may be a multipart entity on its own
def boundary: Option[Boundary] =
contentType.flatMap(_.boundary)
def toBody: Body = {
val base = Body.fromStreamChunked(bytes)
(this.mediaType, this.boundary) match {
case (Some(mt), Some(boundary)) =>
base.contentType(mt, boundary)
case (Some(mt), None) =>
base.contentType(mt)
case _ =>
base
}
}
}
private[http] val crlf = Chunk[Byte]('\r', '\n')
private[http] final class Parser(
boundary: Boundary,
pull: ZIO[Any, Throwable, Either[Any, Chunk[Byte]]],
bufferSize: Int,
) {
val upstream: ZChannel[Any, Any, Any, Any, Throwable, Chunk[Byte], Any] = ZChannel
.fromZIO(pull)
.foldCauseChannel(
ZChannel.refailCause(_),
{
case Left(_) => ZChannel.unit
case Right(chunk) => ZChannel.write(chunk) *> upstream
},
)
def preamble(
buff: Chunk[Byte] = Chunk.empty,
): ZChannel[Any, ZNothing, Chunk[Byte], Any, Throwable, Nothing, (Chunk[Byte], Boolean)] =
buff.indexOfSlice(crlf) match {
case -1 =>
// have to read more input, can discard most of the buffer but must account to the case where current buffer's tail is the prefix of the next encapsulating/closing token
val keep = buff.takeRight(crlf.size + boundary.closingBoundaryBytes.size - 1)
ZChannel
.readWithCause(
in => preamble(keep ++ in),
ZChannel.refailCause,
_ =>
if (boundary.isClosing(buff))
ZChannel.succeed((Chunk.empty, true))
else if (boundary.isEncapsulating(buff))
ZChannel.succeed((Chunk.empty, false))
else
ZChannel.fail(new IllegalStateException("multipart/chunked body ended with no boundary")),
)
case idx =>
val h = buff.take(idx)
val rest = buff.drop(idx + crlf.size)
if (boundary.isClosing(h))
ZChannel.succeed((rest, true))
else if (boundary.isEncapsulating(h))
ZChannel.succeed((rest, false))
else
preamble(rest)
}
def cont(buff: Chunk[Byte]): ZChannel[Any, ZNothing, Chunk[Byte], Any, Throwable, Part, Any] =
parseHeaders(buff, Headers.empty).flatMap { case (headers, rest, closed) =>
parseContent(headers, rest, closed)
}
def parseHeaders(
buff: Chunk[Byte],
res: Headers,
): ZChannel[Any, ZNothing, Chunk[Byte], Any, Throwable, Nothing, (Headers, Chunk[Byte], Boolean)] = {
buff.indexOfSlice(crlf) match {
case -1 =>
// read more:
ZChannel
.readWithCause(
in => parseHeaders(buff ++ in, res),
ZChannel.refailCause(_),
_ => ZChannel.fail(new IllegalStateException("multipart/chunked body ended while parsing part's headers")),
)
case 0 =>
// empty line, hence no more headers
ZChannel.succeed((res, buff.drop(crlf.size), false))
case idx =>
val h = buff.take(idx)
val rest = buff.drop(idx + crlf.size)
// not sure this is actually valid
if (boundary.isClosing(h))
ZChannel.succeed((res, rest, true))
else if (boundary.isEncapsulating(h)) // ditto
ZChannel.succeed((res, rest, false))
else {
// todo: this is a private class, either avoid using it or refactor
FormAST.Header
.fromBytes(h.toArray, StandardCharsets.UTF_8)
.map { case FormAST.Header(name, value) =>
parseHeaders(rest, res.addHeader(name, value))
}
.getOrElse(parseHeaders(rest, res))
}
}
}
def subStream(
pr: zio.Promise[Throwable, (Chunk[Byte], Boolean)],
initialBuff: Chunk[Byte],
): ZStream[Any, Throwable, Byte] = {
// when at beginning of line
@tailrec def parseBody(
buff: Chunk[Byte],
pendingCrlf: Chunk[Byte],
currLine: Chunk[Byte],
seekingBoundary: Boolean,
): ZChannel[Any, ZNothing, Chunk[Byte], Any, Throwable, Chunk[Byte], (Chunk[Byte], Boolean)] = {
if (buff.size >= bufferSize)
ZChannel.write(buff) *> parseBodyAux(Chunk.empty, pendingCrlf, currLine, seekingBoundary)
else {
currLine.indexOfSlice(crlf) match {
case -1 =>
// at this point we have a partial line, if this line is a prefix of the closing boundary we must keep it
// otherwise we can stash it away or even emit it, making sure we keep enough bytes to match a CRLF exactly on the boundary
// notice we're not actually looking for a prefix, but instead make sure we read enough bytes to rule out the possibility of matching a boundary (encapsulating or closing) followed by a crlf.
// this approach ensures we can handle the case of a boundary followed by a split crlf (\r in this read, \n in the next), we have to make sure to keep buffering in this case as we don't have enough bytes to rule out a boundary.
if (seekingBoundary && currLine.size < (boundary.closingBoundaryBytes.size + crlf.size)) {
ZChannel
.readWithCause(
in => parseBodyAux(buff, pendingCrlf, currLine ++ in, true),
err => ZChannel.write(buff) *> ZChannel.refailCause(err),
_ => {
// still possible that the current line is encapsulating or closing boundary
if (boundary.isClosing(currLine))
ZChannel.write(buff) *> ZChannel.succeed((Chunk.empty, true))
else if (boundary.isEncapsulating(currLine))
ZChannel.write(buff) *> ZChannel.succeed((Chunk.empty, false))
else
ZChannel.write(buff) *> ZChannel.fail(
new IllegalStateException("multipart/chunked body ended with no boundary"),
)
},
)
} else {
// we're no longer at beginning of a line, hence no need to look for boundary until we encounter a new line
val (h, t) = currLine.splitAt(currLine.size - crlf.size + 1)
if (t != currLine) {
// also if we had a pending crlf we now know it's part of the content so we move it to the buffered part
parseBody(buff ++ pendingCrlf ++ h, Chunk.empty, t, seekingBoundary = false)
} else {
ZChannel.readWithCause(
in => parseBodyAux(buff ++ pendingCrlf ++ h, Chunk.empty, t ++ in, seekingBoundary = false),
err => ZChannel.write(buff ++ crlf ++ currLine) *> ZChannel.refailCause(err),
_ =>
ZChannel.write(buff ++ crlf ++ currLine) *> ZChannel.fail(
new IllegalStateException("multipart/chunked body ended with no boundary"),
),
)
}
}
case idx if seekingBoundary => // potential boundary
val (h, rest) = currLine.splitAt(idx)
// if we found a boundary it 'consumes' both the pending and trailing crlf, notice pending crlf is optional (i.e. empty part)
if (boundary.isClosing(h))
ZChannel.write(buff) *> ZChannel.succeed((rest.drop(crlf.size), true))
else if (boundary.isEncapsulating(h))
ZChannel.write(buff) *> ZChannel.succeed((rest.drop(crlf.size), false))
else {
// the crlf we just found can either be part of a following boundary or part of the content
val nextLine = rest.drop(crlf.size)
parseBody(buff ++ pendingCrlf ++ h, crlf, nextLine, seekingBoundary = true)
}
case idx => // plain content
// no need to check for boundary, just buffer and continue with parseBOL
val (h, t) = currLine.splitAt(idx)
parseBody(buff ++ h, crlf, t.drop(crlf.size), seekingBoundary = true)
}
}
}
// escape the tailrec compilation error
def parseBodyAux(
buff: Chunk[Byte],
pendingCrlf: Chunk[Byte],
currLine: Chunk[Byte],
seekingBoundary: Boolean,
): ZChannel[Any, ZNothing, Chunk[Byte], Any, Throwable, Chunk[Byte], (Chunk[Byte], Boolean)] =
parseBody(buff, pendingCrlf, currLine, seekingBoundary)
val ch: ZChannel[Any, ZNothing, Chunk[Byte], Any, Throwable, Chunk[Byte], Unit] =
parseBody(Chunk.empty, Chunk.empty, initialBuff, seekingBoundary = true)
.foldCauseChannel(
err => ZChannel.fromZIO(pr.failCause(err)) *> ZChannel.refailCause(err),
tup => ZChannel.fromZIO(pr.succeed(tup)).unit,
)
val pl: ZPipeline[Any, Throwable, Byte, Byte] = ch.toPipeline
upstream.toStream >>> pl
}
def parseContent(
headers: Headers,
buff: Chunk[Byte],
closed: Boolean,
): ZChannel[Any, ZNothing, Chunk[Byte], Any, Throwable, Part, Any] = {
if (closed)
ZChannel.write(Part(headers, ZStream.empty)) *>
epilogue // ignore the epilogue
else {
ZChannel.unwrap {
zio.Promise
.make[Throwable, (Chunk[Byte], Boolean)]
.map { pr =>
val part = Part(headers, subStream(pr, buff))
ZChannel.write(part) *>
ZChannel
.fromZIO(pr.await)
.flatMap {
case (_, true) =>
epilogue
case (rest, false) =>
cont(rest)
}
}
}
}
}
val epilogue: ZChannel[Any, ZNothing, Chunk[Byte], Any, ZNothing, Nothing, Any] =
ZChannel.identity[ZNothing, Chunk[Byte], Any].drain
val startCh: ZChannel[Any, ZNothing, Chunk[Byte], Any, Throwable, Chunk[Part], Any] =
preamble(Chunk.empty).flatMap {
case (leftover, false) =>
cont(leftover)
case (_, true) =>
epilogue
}
.mapOut(Chunk.single)
val startPl: ZPipeline[Any, Throwable, Byte, Part] = startCh.toPipeline
val result: ZStream[Any, Throwable, Part] = upstream.toStream >>> startPl
}
def fromBody(body: Body, bufferSize: Int = 8192): Option[MultipartMixed] = {
body.boundary.map { b =>
MultipartMixed(body.asStream, b, bufferSize)
}
}
def fromParts(parts: ZStream[Any, Throwable, Part], boundary: Boundary, bufferSize: Int = 8192): MultipartMixed = {
val sep = boundary.encapsulationBoundaryBytes ++ crlf
val term = boundary.closingBoundaryBytes
val bytes =
parts.flatMap { case Part(headers, bytes) =>
val headersBytes: ZStream[Any, CharacterCodingException, Byte] = ZStream
.fromIterable(headers)
.map { h =>
s"${h.headerName}: ${h.renderedValue}\r\n"
} >>>
ZPipeline.utf8Encode
ZStream
.concatAll(
Chunk(ZStream.fromChunk(sep), headersBytes, ZStream.fromChunk(crlf), bytes, ZStream.fromChunk(crlf)),
)
}
.concat(ZStream.fromChunk(term))
MultipartMixed(bytes, boundary, bufferSize)
}
}