Add support for setting/removing headers

examples/api/src/Main.scala

@@ -2,7 +2,7 @@ package api
 
 import java.util.UUID
 import io.undertow.Undertow
-import ba.sake.sharaf.*,  routing.*
+import ba.sake.sharaf.*, routing.*
 
 @main def main: Unit =
   val module = JsonApiModule(8181)

formson/src/ba/sake/formson/FormDataRW.scala

@@ -12,7 +12,6 @@ import scala.collection.mutable
 import scala.util.Try
 import ba.sake.formson.FormData.*
 
-
 /** Maps a `T` to/from form data map
   */
 trait FormDataRW[T] {

sharaf/src/ba/sake/sharaf/HeaderUpdates.scala

@@ -0,0 +1,26 @@
+package ba.sake.sharaf
+
+import io.undertow.util.HttpString
+
+/** Headers represented as a series of immutable transformations. This is handy when you dynamically remove header(s),
+  * maybe set by a previous Undertow handler.
+  *
+  * @param updates
+  *   Series of header transformations
+  */
+private[sharaf] final case class HeaderUpdates(updates: Seq[HeaderUpdate]) {
+
+  def setting(name: HttpString, values: Seq[String]) =
+    copy(updates = updates.appended(HeaderUpdate.Set(name, values)))
+
+  def setting(name: HttpString, value: String) =
+    copy(updates = updates.appended(HeaderUpdate.Set(name, Seq(value))))
+
+  def removing(name: HttpString) =
+    copy(updates = updates.appended(HeaderUpdate.Remove(name)))
+
+}
+
+private[sharaf] enum HeaderUpdate:
+  case Set(name: HttpString, values: Seq[String])
+  case Remove(name: HttpString)

sharaf/src/ba/sake/sharaf/Response.scala

@@ -5,42 +5,51 @@ import io.undertow.util.HttpString
 
 final class Response[T] private (
     val status: Int,
-    val headers: Map[HttpString, Seq[String]],
+    private[sharaf] val headerUpdates: HeaderUpdates,
     val body: Option[T]
 )(using val rw: ResponseWritable[T]) {
 
-  def withStatus(status: Int) =
+  def withStatus(status: Int): Response[T] =
     copy(status = status)
 
-  def withHeader(name: HttpString, values: Seq[String]) =
-    copy(headers = headers + (name -> values))
-  def withHeader(name: HttpString, value: String) =
-    copy(headers = headers + (name -> Seq(value)))
+  def settingHeader(name: HttpString, values: Seq[String]): Response[T] =
+    copy(headerUpdates = headerUpdates.setting(name, values))
+  def settingHeader(name: String, values: Seq[String]): Response[T] =
+    settingHeader(HttpString(name), values)
+  def settingHeader(name: HttpString, value: String): Response[T] =
+    copy(headerUpdates = headerUpdates.setting(name, value))
+  def settingHeader(name: String, value: String): Response[T] =
+    settingHeader(HttpString(name), value)
+
+  def removingHeader(name: HttpString): Response[T] =
+    copy(headerUpdates = headerUpdates.removing(name))
+  def removingHeader(name: String): Response[T] =
+    removingHeader(HttpString(name))
 
   def withBody[T2: ResponseWritable](body: T2): Response[T2] =
     copy(body = Some(body))
 
   private def copy[T2](
       status: Int = status,
-      headers: Map[HttpString, Seq[String]] = headers,
+      headerUpdates: HeaderUpdates = headerUpdates,
       body: Option[T2] = body
-  )(using ResponseWritable[T2]) = new Response(status, headers, body)
+  )(using ResponseWritable[T2]) = new Response(status, headerUpdates, body)
 }
 
 object Response {
 
-  private val defaultRes = new Response[String](StatusCodes.OK, Map.empty, None)
+  private val defaultRes = new Response[String](StatusCodes.OK, HeaderUpdates(Seq.empty), None)
 
   def apply[T: ResponseWritable] = defaultRes
 
   def withStatus(status: Int) =
     defaultRes.withStatus(status)
 
   def withHeader(name: HttpString, values: Seq[String]) =
-    defaultRes.withHeader(name, values)
+    defaultRes.settingHeader(name, values)
 
   def withHeader(name: HttpString, value: String) =
-    defaultRes.withHeader(name, Seq(value))
+    defaultRes.settingHeader(name, Seq(value))
 
   def withBody[T: ResponseWritable](body: T): Response[T] =
     defaultRes.withBody(body)
@@ -50,6 +59,6 @@ object Response {
     case None        => throw exceptions.NotFoundException(name)
 
   def redirect(location: String): Response[String] =
-    withStatus(StatusCodes.MOVED_PERMANENTLY).withHeader(HttpString("Location"), location)
+    withStatus(StatusCodes.MOVED_PERMANENTLY).settingHeader(HttpString("Location"), location)
 
 }

sharaf/src/ba/sake/sharaf/ResponseWritable.scala

@@ -16,10 +16,19 @@ object ResponseWritable {
 
   private[sharaf] def writeResponse(response: Response[?], exchange: HttpServerExchange): Unit = {
     // headers
-    val allHeaders = response.body.flatMap(response.rw.headers) ++ response.headers
-    allHeaders.foreach { case (name, values) =>
+    val bodyContentHeaders = response.body.flatMap(response.rw.headers)
+    bodyContentHeaders.foreach { case (name, values) =>
       exchange.getResponseHeaders.putAll(name, values.asJava)
     }
+
+    response.headerUpdates.updates.foreach {
+      case HeaderUpdate.Set(name, values) =>
+        exchange.getResponseHeaders.remove(name)
+        exchange.getResponseHeaders.addAll(name, values.asJava)
+      case HeaderUpdate.Remove(name) =>
+        exchange.getResponseHeaders.remove(name)
+    }
+
     // status code
     exchange.setStatusCode(response.status)
     // body