Carry cookies across redirects via an opt-in CookieStorage (#2671)

core/src/main/scala/sttp/client4/requestBuilder.scala

@@ -4,6 +4,7 @@ import sttp.client4.internal.SttpFile
 import sttp.client4.internal.Utf8
 import sttp.client4.internal.contentTypeWithCharset
 import sttp.client4.logging.LoggingOptions
+import sttp.client4.wrappers.CookieStorage
 import sttp.client4.wrappers.DigestAuthenticationBackend
 import sttp.model.HasHeaders
 import sttp.model.Header
@@ -161,6 +162,14 @@ trait PartialRequestBuilder[+PR <: PartialRequestBuilder[PR, R], +R]
     onDuplicate = DuplicateHeaderBehavior.Combine
   )
 
+  /** Attaches a [[sttp.client4.wrappers.CookieStorage CookieStorage]] to this request. When following redirects (see
+    * [[sttp.client4.wrappers.FollowRedirectsBackend]], applied to all backends by default), cookies set via
+    * `Set-Cookie` in a redirect chain are then collected into the storage and sent with subsequent requests that they
+    * domain/path-match. Without a storage, cookies are not carried across redirects (the `Cookie` header is sensitive
+    * and stripped). Start with [[sttp.client4.wrappers.CookieStorage.empty]].
+    */
+  def cookieStorage(storage: CookieStorage): PR = attribute(CookieStorage.attributeKey, storage)
+
   private[client4] def hasContentType: Boolean = headers.exists(_.is(HeaderNames.ContentType))
   private[client4] def setContentTypeIfMissing(mt: MediaType): PR =
     if (hasContentType) this else contentType(mt)

core/src/main/scala/sttp/client4/wrappers/CookieStorage.scala

@@ -0,0 +1,144 @@
+package sttp.client4.wrappers
+
+import sttp.attributes.AttributeKey
+import sttp.model.Uri
+
+/** An immutable cookie jar.
+  *
+  * Collects cookies received in `Set-Cookie` response headers and determines which of them should be sent with a
+  * request to a given URI, applying a subset of the [[https://www.rfc-editor.org/rfc/rfc6265 RFC 6265]] rules:
+  * domain-matching, path-matching and the `Secure` attribute.
+  *
+  * Intended use: attach a storage to a request using [[sttp.client4.PartialRequestBuilder.cookieStorage]]. The
+  * [[FollowRedirectsBackend]] (applied to all backends by default) then, for each request in a redirect chain, sends
+  * the matching stored cookies and threads an updated storage through to the next request. This makes cookies set via
+  * `Set-Cookie` during a redirect chain visible to subsequent requests in that chain - which otherwise doesn't happen,
+  * as the `Cookie` header is a sensitive header, stripped when following redirects.
+  *
+  * Cookies are represented as plain name/value pairs rather than [[sttp.model.headers.CookieWithMeta]]. That type is
+  * available on all platforms, but its `Set-Cookie` rendering and parsing reach `java.time` date formatting (for
+  * `Expires`, via `ZoneId`/`DateTimeFormatter`), a subset of `java.time` not supported on Scala Native; referencing it
+  * from this shared code pulls those symbols in and breaks the Native link. Time-based expiry (`Max-Age` > 0, `Expires`)
+  * is not tracked anyway, as the storage has no notion of the current time; a `Set-Cookie` with `Max-Age` <= 0 removes a
+  * matching cookie, so a server can still clear cookies within a chain.
+  */
+final class CookieStorage private (private val entries: Map[CookieStorage.Key, CookieStorage.Stored]) {
+  import CookieStorage._
+
+  /** A new storage updated with the cookies parsed from the `Set-Cookie` header values received from `setBy`.
+    * Following RFC 6265, a cookie whose `Domain` attribute does not domain-match `setBy` is rejected (to prevent a
+    * host setting cookies for unrelated domains). A cookie with `Max-Age` <= 0 removes a matching stored cookie.
+    */
+  def setFromSetCookieHeaders(setBy: Uri, setCookieHeaders: Iterable[String]): CookieStorage = {
+    val host = hostOf(setBy)
+    val updated = setCookieHeaders.flatMap(parseSetCookie).foldLeft(entries) { (acc, c) =>
+      val hostOnly = c.domain.isEmpty
+      val domain = c.domain.map(normalizeDomain).getOrElse(host)
+      if (domain.isEmpty || !domainMatches(host, domain)) acc
+      else {
+        val key = Key(c.name, domain, c.path.getOrElse(defaultPathOf(setBy)))
+        if (c.removed) acc - key
+        else acc.updated(key, Stored(c.value, c.secure, hostOnly))
+      }
+    }
+    new CookieStorage(updated)
+  }
+
+  /** The cookies, as `name -> value` pairs, that should be sent with a request to `uri`, according to domain-matching,
+    * path-matching and the `Secure` attribute (secure cookies are only sent over `https`).
+    */
+  def cookiesFor(uri: Uri): Seq[(String, String)] = {
+    val host = hostOf(uri)
+    val path = pathOf(uri)
+    val secure = uri.scheme.exists(_.equalsIgnoreCase("https"))
+    entries.collect { case (key, stored) if matches(key, stored, host, path, secure) => key.name -> stored.value }.toSeq
+  }
+
+  private def matches(key: Key, stored: Stored, host: String, path: String, secure: Boolean): Boolean = {
+    val domainMatch = if (stored.hostOnly) host == key.domain else domainMatches(host, key.domain)
+    domainMatch && pathMatches(path, key.path) && (!stored.secure || secure)
+  }
+
+  def isEmpty: Boolean = entries.isEmpty
+}
+
+object CookieStorage {
+
+  /** An empty storage. */
+  val empty: CookieStorage = new CookieStorage(Map.empty)
+
+  /** The attribute key under which a [[CookieStorage]] is attached to a request; see
+    * [[sttp.client4.PartialRequestBuilder.cookieStorage]].
+    */
+  val attributeKey: AttributeKey[CookieStorage] =
+    new AttributeKey[CookieStorage]("sttp.client4.wrappers.CookieStorage")
+
+  private val DefaultPath = "/"
+
+  // a cookie is identified by its name, the domain it's scoped to and its path
+  private case class Key(name: String, domain: String, path: String)
+  private case class Stored(value: String, secure: Boolean, hostOnly: Boolean)
+
+  // a `Set-Cookie` cookie before its domain is resolved against the setting host; `removed` marks a Max-Age <= 0
+  // deletion
+  private case class Parsed(
+      name: String,
+      value: String,
+      domain: Option[String],
+      path: Option[String],
+      secure: Boolean,
+      removed: Boolean
+  )
+
+  private def hostOf(uri: Uri): String = uri.host.getOrElse("").toLowerCase
+  private def pathOf(uri: Uri): String = "/" + uri.path.mkString("/")
+  private def normalizeDomain(d: String): String = d.stripPrefix(".").toLowerCase
+
+  // RFC 6265, 5.1.4: the default-path of a cookie without a `Path` attribute is the setting request's directory - the
+  // path up to, but not including, the rightmost "/" (or "/" if there is none beyond the leading one)
+  private def defaultPathOf(uri: Uri): String = {
+    val p = pathOf(uri)
+    val lastSlash = p.lastIndexOf('/')
+    if (lastSlash <= 0) DefaultPath else p.substring(0, lastSlash)
+  }
+
+  // RFC 6265, 5.1.3: equal, or `host` is a subdomain of `domain`
+  private def domainMatches(host: String, domain: String): Boolean =
+    host == domain || host.endsWith("." + domain)
+
+  // RFC 6265, 5.1.4
+  private def pathMatches(requestPath: String, cookiePath: String): Boolean =
+    requestPath == cookiePath ||
+      (requestPath.startsWith(cookiePath) &&
+        (cookiePath.endsWith("/") || requestPath.charAt(cookiePath.length) == '/'))
+
+  // A minimal `Set-Cookie` parser reading only the attributes used for storage. CookieWithMeta.parse isn't reused
+  // because its `Expires` handling relies on java.time date formatting, which doesn't work on Scala Native; `Expires`
+  // is ignored here anyway, as the storage doesn't track time-based expiry.
+  private def parseSetCookie(headerValue: String): Option[Parsed] =
+    headerValue.split(";").iterator.map(_.trim).filter(_.nonEmpty).toList match {
+      case nameValue :: directives =>
+        val eq = nameValue.indexOf('=')
+        val name = if (eq < 0) nameValue else nameValue.substring(0, eq).trim
+        if (name.isEmpty) None
+        else {
+          val value = if (eq < 0) "" else nameValue.substring(eq + 1).trim
+          val attrs = directives.map { d =>
+            val i = d.indexOf('=')
+            if (i < 0) (d.toLowerCase, "") else (d.substring(0, i).trim.toLowerCase, d.substring(i + 1).trim)
+          }.toMap
+          val maxAge = attrs.get("max-age").flatMap(s => scala.util.Try(s.toLong).toOption)
+          Some(
+            Parsed(
+              name = name,
+              value = value,
+              domain = attrs.get("domain").filter(_.nonEmpty),
+              path = attrs.get("path").filter(_.nonEmpty),
+              secure = attrs.contains("secure"),
+              removed = maxAge.exists(_ <= 0)
+            )
+          )
+        }
+      case Nil => None
+    }
+}

core/src/main/scala/sttp/client4/wrappers/FollowRedirectsBackend.scala

@@ -16,15 +16,19 @@ abstract class FollowRedirectsBackend[F[_], P] private (
   override def send[T](request: GenericRequest[T, R]): F[Response[T]] = sendWithCounter(request, 0)
 
   protected def sendWithCounter[T](request: GenericRequest[T, R], redirects: Int): F[Response[T]] = {
+    // if a cookie storage is attached, send the cookies matching this request's URI (cookies are otherwise lost
+    // across redirects, as the `Cookie` header is sensitive and stripped)
+    val requestWithCookies = applyStoredCookies(request)
+
     // if there are nested follow redirect backends, disabling them and handling redirects here
     // using a def instead of a val so that errors are properly caught
-    def resp = delegate.send(request.followRedirects(false))
+    def resp = delegate.send(requestWithCookies.followRedirects(false))
 
     if (request.options.followRedirects) {
       resp
         .flatMap { (response: Response[T]) =>
           if (response.isRedirect) {
-            followRedirect(request, response, redirects)
+            followRedirect(updateStoredCookies(request, response), response, redirects)
           } else {
             monad.unit(response)
           }
@@ -34,7 +38,7 @@ abstract class FollowRedirectsBackend[F[_], P] private (
             case Some(re) if re.response.isRedirect =>
               re.response.header(HeaderNames.Location) match {
                 case None      => monad.error(e) // no location header, propagating the exception
-                case Some(loc) => followRedirect(request, re.response, redirects, loc)
+                case Some(loc) => followRedirect(updateStoredCookies(request, re.response), re.response, redirects, loc)
               }
             case _ => monad.error(e)
           }
@@ -44,6 +48,28 @@ abstract class FollowRedirectsBackend[F[_], P] private (
     }
   }
 
+  /** If a [[CookieStorage]] is attached, adds the cookies matching the request's URI as a `Cookie` header. No-op
+    * otherwise, so the default behaviour is unchanged unless a storage is explicitly attached.
+    */
+  private def applyStoredCookies[T](request: GenericRequest[T, R]): GenericRequest[T, R] =
+    request.attribute(CookieStorage.attributeKey) match {
+      case Some(storage) =>
+        val cookies = storage.cookiesFor(request.uri)
+        if (cookies.isEmpty) request else request.cookies(cookies: _*)
+      case None => request
+    }
+
+  /** If a [[CookieStorage]] is attached, returns the request with the storage updated with the response's `Set-Cookie`
+    * cookies, so that the next request in the redirect chain can send them. No-op otherwise.
+    */
+  private def updateStoredCookies[T](request: GenericRequest[T, R], response: ResponseMetadata): GenericRequest[T, R] =
+    request.attribute(CookieStorage.attributeKey) match {
+      case Some(storage) =>
+        val updated = storage.setFromSetCookieHeaders(request.uri, response.headers(HeaderNames.SetCookie))
+        request.attribute(CookieStorage.attributeKey, updated)
+      case None => request
+    }
+
   private def followRedirect[T](
       request: GenericRequest[T, R],
       response: Response[T],

core/src/test/scala/sttp/client4/FollowRedirectsBackendTest.scala

@@ -4,9 +4,9 @@ import org.scalatest.EitherValues
 import org.scalatest.funsuite.AnyFunSuite
 import org.scalatest.matchers.should.Matchers
 import sttp.client4.testing.{BackendStub, ResponseStub}
-import sttp.client4.wrappers.{FollowRedirectsBackend, FollowRedirectsConfig}
+import sttp.client4.wrappers.{CookieStorage, FollowRedirectsBackend, FollowRedirectsConfig}
 import sttp.model.internal.Rfc3986
-import sttp.model.{Header, StatusCode, Uri}
+import sttp.model.{Header, HeaderNames, ResponseMetadata, StatusCode, Uri}
 
 class FollowRedirectsBackendTest extends AnyFunSuite with Matchers with EitherValues {
   val testData = List(
@@ -53,4 +53,73 @@ class FollowRedirectsBackendTest extends AnyFunSuite with Matchers with EitherVa
     result.body.value shouldBe "All good!"
   }
 
+  private def cookiesIn(r: GenericRequest[_, _]): Set[String] =
+    r.header(HeaderNames.Cookie).map(_.split("; ").toSet).getOrElse(Set.empty)
+
+  // a redirect chain example.com/0 -> /1 -> ... -> /n, where each hop sets a cookie `c<id>`; records the cookies
+  // (name=value) seen in the `Cookie` header of each request, by target id
+  private def redirectChainSettingCookies(n: Int): (SyncBackend, collection.Map[Int, Set[String]]) = {
+    def url(id: Int) = uri"https://example.com/$id"
+    val seen = scala.collection.mutable.Map[Int, Set[String]]()
+    val stub = BackendStub.synchronous.whenRequestMatchesPartial {
+      case r if r.uri.host.contains("example.com") =>
+        val id = r.uri.path.last.toInt
+        seen(id) = cookiesIn(r)
+        if (id < n)
+          ResponseStub.adjust(
+            "",
+            StatusCode.Found,
+            Vector(Header.location(url(id + 1)), Header(HeaderNames.SetCookie, s"c$id=$id"))
+          )
+        else ResponseStub.adjust("done", StatusCode.Ok)
+    }
+    (FollowRedirectsBackend(stub), seen)
+  }
+
+  test("should send cookies set during a redirect chain to subsequent requests when a cookie storage is attached") {
+    val (backend, seen) = redirectChainSettingCookies(3)
+
+    val result = basicRequest.get(uri"https://example.com/0").cookieStorage(CookieStorage.empty).send(backend)
+
+    result.code shouldBe StatusCode.Ok
+    seen(0) shouldBe empty // nothing set yet
+    seen(1) shouldBe Set("c0=0")
+    seen(2) shouldBe Set("c0=0", "c1=1")
+    seen(3) shouldBe Set("c0=0", "c1=1", "c2=2")
+  }
+
+  test("should not carry cookies across redirects when no cookie storage is attached") {
+    val (backend, seen) = redirectChainSettingCookies(3)
+
+    val result = basicRequest.get(uri"https://example.com/0").send(backend)
+
+    result.code shouldBe StatusCode.Ok
+    seen.values.foreach(_ shouldBe empty)
+  }
+
+  test("should harvest cookies from a redirect that a backend signals by throwing") {
+    var atTarget = Set.empty[String]
+    val stub = BackendStub.synchronous.whenRequestMatchesPartial {
+      // first hop: signal the redirect (with a Set-Cookie) by throwing, as some backends do
+      case r if r.uri == uri"https://example.com/0" =>
+        throw ResponseException.UnexpectedStatusCode(
+          "",
+          ResponseMetadata(
+            StatusCode.Found,
+            "",
+            Vector(Header.location(uri"https://example.com/1"), Header(HeaderNames.SetCookie, "s=1"))
+          )
+        )
+      case r if r.uri == uri"https://example.com/1" =>
+        atTarget = cookiesIn(r)
+        ResponseStub.adjust("done", StatusCode.Ok)
+    }
+    val backend = FollowRedirectsBackend(stub)
+
+    val result = basicRequest.get(uri"https://example.com/0").cookieStorage(CookieStorage.empty).send(backend)
+
+    result.code shouldBe StatusCode.Ok
+    atTarget shouldBe Set("s=1")
+  }
+
 }

core/src/test/scala/sttp/client4/wrappers/CookieStorageTest.scala

@@ -0,0 +1,69 @@
+package sttp.client4.wrappers
+
+import org.scalatest.funsuite.AnyFunSuite
+import org.scalatest.matchers.should.Matchers
+import sttp.client4._
+
+class CookieStorageTest extends AnyFunSuite with Matchers {
+  private def names(cs: Seq[(String, String)]) = cs.map(_._1).toSet
+
+  test("stores a host-only cookie and sends it back to the same host") {
+    val storage = CookieStorage.empty.setFromSetCookieHeaders(uri"https://example.com/a", List("s=1"))
+    names(storage.cookiesFor(uri"https://example.com/b")) shouldBe Set("s")
+  }
+
+  test("does not send a host-only cookie to a subdomain") {
+    val storage = CookieStorage.empty.setFromSetCookieHeaders(uri"https://example.com/", List("s=1"))
+    storage.cookiesFor(uri"https://sub.example.com/") shouldBe empty
+  }
+
+  test("sends a domain cookie to matching subdomains, but not to unrelated domains") {
+    val storage = CookieStorage.empty.setFromSetCookieHeaders(uri"https://example.com/", List("s=1; Domain=example.com"))
+    names(storage.cookiesFor(uri"https://sub.example.com/")) shouldBe Set("s")
+    storage.cookiesFor(uri"https://other.com/") shouldBe empty
+  }
+
+  test("rejects a cookie whose domain does not match the setting host") {
+    val storage = CookieStorage.empty.setFromSetCookieHeaders(uri"https://example.com/", List("s=1; Domain=evil.com"))
+    storage.isEmpty shouldBe true
+  }
+
+  test("normalizes the cookie domain (leading dot, case)") {
+    val storage =
+      CookieStorage.empty.setFromSetCookieHeaders(uri"https://example.com/", List("s=1; Domain=.EXAMPLE.com"))
+    names(storage.cookiesFor(uri"https://sub.example.com/")) shouldBe Set("s")
+  }
+
+  test("respects the cookie path, requiring a path-segment boundary") {
+    val storage = CookieStorage.empty.setFromSetCookieHeaders(uri"https://example.com/admin", List("s=1; Path=/admin"))
+    names(storage.cookiesFor(uri"https://example.com/admin/x")) shouldBe Set("s")
+    storage.cookiesFor(uri"https://example.com/administrator") shouldBe empty // prefix, but not a segment boundary
+    storage.cookiesFor(uri"https://example.com/public") shouldBe empty
+  }
+
+  test("defaults a path-less cookie to the setting request's directory (RFC 6265 5.1.4)") {
+    val storage = CookieStorage.empty.setFromSetCookieHeaders(uri"https://example.com/admin/page", List("s=1"))
+    names(storage.cookiesFor(uri"https://example.com/admin/other")) shouldBe Set("s")
+    storage.cookiesFor(uri"https://example.com/elsewhere") shouldBe empty
+  }
+
+  test("does not send a secure cookie over http") {
+    val storage = CookieStorage.empty.setFromSetCookieHeaders(uri"https://example.com/", List("s=1; Secure"))
+    storage.cookiesFor(uri"http://example.com/") shouldBe empty
+    names(storage.cookiesFor(uri"https://example.com/")) shouldBe Set("s")
+  }
+
+  test("a later cookie with the same name/domain/path overwrites the earlier one") {
+    val storage = CookieStorage.empty
+      .setFromSetCookieHeaders(uri"https://example.com/", List("s=1"))
+      .setFromSetCookieHeaders(uri"https://example.com/", List("s=2"))
+    storage.cookiesFor(uri"https://example.com/") shouldBe Seq("s" -> "2")
+  }
+
+  test("a cookie with Max-Age <= 0 removes a matching stored cookie") {
+    val storage = CookieStorage.empty
+      .setFromSetCookieHeaders(uri"https://example.com/", List("s=1"))
+      .setFromSetCookieHeaders(uri"https://example.com/", List("s=; Max-Age=0"))
+    storage.isEmpty shouldBe true
+  }
+}