play body in memory

package controllers

import play.api.libs.concurrent.Execution.Implicits._
import play.api.libs.iteratee.Iteratee
import play.api.mvc.MultipartFormData.FilePart
import play.api.mvc.{ MultipartFormData, BodyParser, BodyParsers }
import scala.collection.mutable.ArrayBuilder

case class MaxSizeExceeded(length: Long)

object InMemoryBodyParser {

  /**
   * Parse the content as multipart/form-data and store the content in memory
   */
  def multipartFormData(maxSize: Long): BodyParser[MultipartFormData[Either[MaxSizeExceeded, Array[Byte]]]] = BodyParsers.parse.multipartFormData(handleFilePartAsArrayByte(maxSize))

  type PartHandler[A] = PartialFunction[Map[String, String], Iteratee[Array[Byte], A]]

  def handleFilePartAsArrayByte(maxSize: Long): PartHandler[FilePart[Either[MaxSizeExceeded, Array[Byte]]]] = {
    BodyParsers.parse.Multipart.handleFilePart {
      case BodyParsers.parse.Multipart.FileInfo(partName, filename, contentType) ⇒
        // left -> error, the max size was exceeded
        // right -> tupple:
        //          - 1: size of consumed data
        //          - 2: builder for an array of byte
        type StepType = Either[MaxSizeExceeded, (Long, ArrayBuilder[Byte])]
        val initialStep: StepType = Right(0l -> new ArrayBuilder.ofByte())

        Iteratee.fold[Array[Byte], StepType](initialStep) { (errorOrBuf, data) ⇒
          errorOrBuf match {
            case Left(maxSizeExceeded) ⇒ Left(maxSizeExceeded.copy(length = maxSizeExceeded.length + data.size))
            case Right((size, buf)) ⇒ {
              val newSize = size + data.size
              if (newSize > maxSize) {
                Left(MaxSizeExceeded(newSize))
              } else {
                buf ++= data
                Right(newSize -> buf)
              }
            }
          }
        }.map { errorOrBuf ⇒
          errorOrBuf.right.map { builder ⇒
            val content: Array[Byte] = builder._2.result()
            content
          }
        }
    }
  }
}