Headers & errors support

3 files changed, 44 insertions, 14 deletions

diff --git a/core/src/main/scala/com/softwaremill/sttp/HttpConnectionSttpHandler.scala b/core/src/main/scala/com/softwaremill/sttp/HttpConnectionSttpHandler.scala
index fdbb322..07025d5 100644
--- a/core/src/main/scala/com/softwaremill/sttp/HttpConnectionSttpHandler.scala
+++ b/core/src/main/scala/com/softwaremill/sttp/HttpConnectionSttpHandler.scala
@@ -1,11 +1,6 @@
 package com.softwaremill.sttp
 
-import java.io.{
-  ByteArrayOutputStream,
-  InputStream,
-  OutputStream,
-  OutputStreamWriter
-}
+import java.io._
 import java.net.HttpURLConnection
 import java.nio.channels.Channels
 import java.nio.file.Files
@@ -14,6 +9,7 @@ import com.softwaremill.sttp.model._
 
 import scala.annotation.tailrec
 import scala.io.Source
+import scala.collection.JavaConverters._
 
 object HttpConnectionSttpHandler extends SttpHandler[Id, Nothing] {
   override def send[T](r: Request,
@@ -24,8 +20,13 @@ object HttpConnectionSttpHandler extends SttpHandler[Id, Nothing] {
     c.setDoInput(true)
     setBody(r.body, c)
 
-    val status = c.getResponseCode
-    Response(status, readResponse(c.getInputStream, responseAs))
+    try {
+      val is = c.getInputStream
+      readResponse(c, is, responseAs)
+    } catch {
+      case _: IOException if c.getResponseCode != -1 =>
+        readResponse(c, c.getErrorStream, responseAs)
+    }
   }
 
   private def setBody(body: RequestBody, c: HttpURLConnection): Unit = {
@@ -76,8 +77,19 @@ object HttpConnectionSttpHandler extends SttpHandler[Id, Nothing] {
     }
   }
 
-  private def readResponse[T](is: InputStream,
-                              responseAs: ResponseAs[T, Nothing]): T =
+  private def readResponse[T](
+      c: HttpURLConnection,
+      is: InputStream,
+      responseAs: ResponseAs[T, Nothing]): Response[T] = {
+
+    val headers = c.getHeaderFields.asScala.toVector
+      .filter(_._1 != null)
+      .flatMap { case (k, vv) => vv.asScala.map((k, _)) }
+    Response(readResponseBody(is, responseAs), c.getResponseCode, headers)
+  }
+
+  private def readResponseBody[T](is: InputStream,
+                                  responseAs: ResponseAs[T, Nothing]): T =
     responseAs match {
       case IgnoreResponse =>
         @tailrec def consume(): Unit = if (is.read() != -1) consume()
diff --git a/core/src/main/scala/com/softwaremill/sttp/Response.scala b/core/src/main/scala/com/softwaremill/sttp/Response.scala
index 4b90259..dbaf71f 100644
--- a/core/src/main/scala/com/softwaremill/sttp/Response.scala
+++ b/core/src/main/scala/com/softwaremill/sttp/Response.scala
@@ -1,3 +1,21 @@
 package com.softwaremill.sttp
 
-case class Response[T](status: Int, body: T)
+import scala.collection.immutable.Seq
+import scala.util.Try
+
+case class Response[T](body: T, code: Int, headers: Seq[(String, String)]) {
+  def is200: Boolean = code == 200
+  def isSuccess: Boolean = code >= 200 && code < 300
+  def isRedirect: Boolean = code >= 300 && code < 400
+  def isClientError: Boolean = code >= 400 && code < 500
+  def isServerError: Boolean = code >= 500 && code < 600
+
+  def header(h: String): Option[String] =
+    headers.find(_._1.equalsIgnoreCase(h)).map(_._2)
+  def headers(h: String): Seq[String] =
+    headers.filter(_._1.equalsIgnoreCase(h)).map(_._2)
+
+  def contentType: Option[String] = header(ContentTypeHeader)
+  def contentLength: Option[Long] =
+    header(ContentLengthHeader).flatMap(cl => Try(cl.toLong).toOption)
+}
diff --git a/core/src/main/scala/com/softwaremill/sttp/package.scala b/core/src/main/scala/com/softwaremill/sttp/package.scala
index 79d9a17..c39f256 100644
--- a/core/src/main/scala/com/softwaremill/sttp/package.scala
+++ b/core/src/main/scala/com/softwaremill/sttp/package.scala
@@ -130,10 +130,9 @@ package object sttp {
     def header(k: String,
                v: String,
                replaceExisting: Boolean = false): RequestTemplate[U] = {
-      val kLower = k.toLowerCase
       val current =
         if (replaceExisting)
-          headers.filterNot(_._1.toLowerCase.contains(kLower))
+          headers.filterNot(_._1.equalsIgnoreCase(k))
         else headers
       this.copy(headers = current :+ (k -> v))
     }
@@ -232,7 +231,8 @@ package object sttp {
 
   val sttp: RequestTemplate[Empty] = RequestTemplate.empty
 
-  private val ContentTypeHeader = "content-type"
+  private[sttp] val ContentTypeHeader = "content-type"
+  private[sttp] val ContentLengthHeader = "content-length"
   private val Utf8 = "utf-8"
 
   private val ApplicationOctetStreamContentType = "application/octet-stream"