From 886cd1bbb8ea75ab90137322857ab6943700040d Mon Sep 17 00:00:00 2001
From: Alexander Shparun <alexander.shparun@jetbrains.com>
Date: Thu, 20 Feb 2025 01:12:27 +0100
Subject: [PATCH] [fleet] FL-14501 CSRF protection on workspace server

GitOrigin-RevId: 8a60ac56e05f43fef240efefdbd3095c343f4928
---
 fleet/rpc/src/fleet/rpc/client/FleetClient.kt |  2 +-
 fleet/rpc/src/fleet/rpc/core/Transport.kt     | 20 ++++++--
 fleet/util/core/src/fleet/util/Either.kt      | 47 ++++++++++++-------
 fleet/util/core/src/fleet/util/Result.kt      |  2 +-
 4 files changed, 49 insertions(+), 22 deletions(-)

diff --git a/fleet/rpc/src/fleet/rpc/client/FleetClient.kt b/fleet/rpc/src/fleet/rpc/client/FleetClient.kt
index 16c30f794ace..17a9f603c3d2 100644
--- a/fleet/rpc/src/fleet/rpc/client/FleetClient.kt
+++ b/fleet/rpc/src/fleet/rpc/client/FleetClient.kt
@@ -86,4 +86,4 @@ suspend fun withFleetClient(
       body(fleetClient)
     }
   }
-}
+}
\ No newline at end of file
diff --git a/fleet/rpc/src/fleet/rpc/core/Transport.kt b/fleet/rpc/src/fleet/rpc/core/Transport.kt
index f25b3e1ebf35..81ef0ff490bc 100644
--- a/fleet/rpc/src/fleet/rpc/core/Transport.kt
+++ b/fleet/rpc/src/fleet/rpc/core/Transport.kt
@@ -15,8 +15,10 @@ import kotlin.time.TimeSource
 
 private val logger = logger<Transport<*>>()
 
-class Transport<T>(val outgoing: SendChannel<T>,
-                   val incoming: ReceiveChannel<T>)
+class Transport<T>(
+  val outgoing: SendChannel<T>,
+  val incoming: ReceiveChannel<T>,
+)
 
 fun interface FleetTransportFactory {
   /*
@@ -24,10 +26,20 @@ fun interface FleetTransportFactory {
    * If the connection isn't possible, rethrow the cause as [TransportDisconnectedException].
    * When the underlying transport is broken, e.g. a socket is closed, both channels should be closed with [TransportDisconnectedException]
    */
-  suspend fun connect(transportStats: MutableStateFlow<TransportStats>?,
-                      body: suspend CoroutineScope.(Transport<TransportMessage>) -> Unit)
+  suspend fun connect(
+    transportStats: MutableStateFlow<TransportStats>?,
+    body: suspend CoroutineScope.(Transport<TransportMessage>) -> Unit,
+  )
 }
 
+/**
+ * Use this function when you want to re-create your factory on each connection attempt.
+ */
+fun dynamicTransportFactory(f: suspend () -> FleetTransportFactory): FleetTransportFactory =
+  FleetTransportFactory { socketStats, body ->
+    f().connect(socketStats, body)
+  }
+
 enum class DebugConnectionState {
   Connect,
   Disconnect
diff --git a/fleet/util/core/src/fleet/util/Either.kt b/fleet/util/core/src/fleet/util/Either.kt
index b739e8d8451e..00b8d8b5d08e 100644
--- a/fleet/util/core/src/fleet/util/Either.kt
+++ b/fleet/util/core/src/fleet/util/Either.kt
@@ -24,7 +24,7 @@ import kotlin.error
 @JvmInline
 @Serializable(with = EitherSerializer::class)
 value class Either<out T, E> internal constructor(
-  private val value: Any?
+  private val value: Any?,
 ) {
   /**
    * Returns `true` if this instance represents a value outcome.
@@ -102,20 +102,26 @@ val <T, E> Either<T, E>.value: T
 val <T, E> Either<T, E>.error: E
   get() = requireNotNull(errorOrNull) { "Not a Error" }
 
-inline fun <T, E, R> Either<T, E>.flatMap(f: (T) -> Either<R, E>): Either<R, E> = when (this.isValue) {
-  false -> this as Either<R, E>
-  else -> f(valueOrNull!!)
-}
+inline fun <T, E, R> Either<T, E>.flatMap(f: (T) -> Either<R, E>): Either<R, E> =
+  @Suppress("UNCHECKED_CAST")
+  when (this.isValue) {
+    false -> this as Either<R, E>
+    else -> f(valueOrNull!!)
+  }
 
-inline fun <T, E, R> Either<T, E>.map(f: (T) -> R): Either<R, E> = when (this.isValue) {
-  false -> this as Either<R, E>
-  true -> Either.value(f(valueOrNull!!))
-}
+inline fun <T, E, R> Either<T, E>.map(f: (T) -> R): Either<R, E> =
+  @Suppress("UNCHECKED_CAST")
+  when (this.isValue) {
+    false -> this as Either<R, E>
+    true -> Either.value(f(valueOrNull!!))
+  }
 
-inline fun <T, E, R> Either<T, E>.mapError(f: (E) -> R): Either<T, R> = when (this.isError) {
-  false -> this as Either<T, R>
-  true -> Either.error(f(error))
-}
+inline fun <T, E, R> Either<T, E>.mapError(f: (E) -> R): Either<T, R> =
+  @Suppress("UNCHECKED_CAST")
+  when (this.isError) {
+    false -> this as Either<T, R>
+    true -> Either.error(f(error))
+  }
 
 /**
  * Performs the given [action] on the encapsulated [E] if this instance represents [error][Either.isError].
@@ -129,6 +135,15 @@ inline fun <T, E> Either<T, E>.onError(action: (error: E) -> Unit): Either<T, E>
   return this
 }
 
+/**
+ * Unwraps the value of an `Either` instance.
+ *
+ * If the instance is a `error`, it will call the provided [action] function that is supposed to throw or unwind.
+ * If the instance is a `value`, it will return the contained value.
+ */
+inline fun <T, E> Either<T, E>.unwrap(action: (error: E) -> Nothing): T =
+  onError { action(it) }.value
+
 /**
  * Performs the given [action] on the encapsulated value if this instance represents [value][Either.isValue].
  * Returns the original `Either` unchanged.
@@ -142,8 +157,8 @@ inline fun <T, E> Either<T, E>.onValue(action: (value: T) -> Unit): Either<T, E>
 }
 
 internal class EitherSerializer<T, E>(
-  val valueSerializer: KSerializer<T>,
-  val errorSerializer: KSerializer<E>,
+  private val valueSerializer: KSerializer<T>,
+  private val errorSerializer: KSerializer<E>,
 ) : KSerializer<Either<T, E>> {
 
   override val descriptor: SerialDescriptor
@@ -172,4 +187,4 @@ internal class EitherSerializer<T, E>(
       }
     }
   }
-}
+}
\ No newline at end of file
diff --git a/fleet/util/core/src/fleet/util/Result.kt b/fleet/util/core/src/fleet/util/Result.kt
index 6c1ce90b4a31..98593eb9b007 100644
--- a/fleet/util/core/src/fleet/util/Result.kt
+++ b/fleet/util/core/src/fleet/util/Result.kt
@@ -25,4 +25,4 @@ fun <T, R> Result.Companion.flatMap(source: Iterable<T>, transform: (T) -> Resul
     onFailure = { err -> return failure(err) }
   ) }
   return success(successes)
-}
+}
\ No newline at end of file