Coverage Summary for Class: WiretapOkHttpWebSocketListener (dev.skymansandy.wiretap.okhttp.ws)

 /*
  * Copyright (c) 2026 skymansandy. All rights reserved.
  */
 
 package dev.skymansandy.wiretap.okhttp.ws
 
 import dev.skymansandy.wiretap.di.WiretapDi
 import dev.skymansandy.wiretap.domain.model.SocketConnection
 import dev.skymansandy.wiretap.domain.model.SocketContentType
 import dev.skymansandy.wiretap.domain.model.SocketMessage
 import dev.skymansandy.wiretap.domain.model.SocketMessageType
 import dev.skymansandy.wiretap.domain.model.SocketStatus
 import dev.skymansandy.wiretap.domain.model.config.ws.WiretapWsConfig
 import dev.skymansandy.wiretap.domain.orchestrator.SocketLogManager
 import dev.skymansandy.wiretap.domain.orchestrator.withConsoleLogging
 import dev.skymansandy.wiretap.helper.logger.WiretapLogger
 import dev.skymansandy.wiretap.helper.markers.InternalWiretapApi
 import dev.skymansandy.wiretap.helper.util.currentTimeMillis
 import dev.skymansandy.wiretap.helper.ws.decodeBinaryFrame
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import okhttp3.Response
 import okhttp3.WebSocket
 import okhttp3.WebSocketListener
 import okio.ByteString
 import org.koin.core.Koin
 import org.koin.core.component.KoinComponent
 import org.koin.core.component.inject
 import kotlin.concurrent.Volatile
 
 /**
  * Wraps a consumer's WebSocketListener to log all WebSocket events via Wiretap.
  *
  * Use the [wiretapped] extension to create an instance:
  * ```kotlin
  * client.newWebSocket(request, myListener.wiretapped())
  * ```
  */
 @OptIn(InternalWiretapApi::class)
 internal class WiretapOkHttpWebSocketListener(
     private val delegate: WebSocketListener,
     private val config: WiretapWsConfig,
 ) : WebSocketListener(), KoinComponent {
 
     override fun getKoin(): Koin = WiretapDi.getKoin()
 
     private val wiretapLogger: WiretapLogger by inject()
     private val sharedSocketLogManager: SocketLogManager by inject()
 
     private val socketLogManager: SocketLogManager by lazy {
         if (config.consoleLogging) sharedSocketLogManager.withConsoleLogging(wiretapLogger)
         else sharedSocketLogManager
     }
 
     private val logScope = CoroutineScope(SupervisorJob() + Dispatchers.IO)
 
     @Volatile
     private var socketId: Long = -1
     private val isSocketActive
         get() = socketId >= 0
 
     override fun onOpen(webSocket: WebSocket, response: Response) {
         if (!config.enabled) {
             delegate.onOpen(webSocket, response)
             return
         }
         // runBlocking required here: socketId must be set before delegate.onOpen
         runBlocking {
             val url = webSocket.request().url.toString()
             val reqHeaders = webSocket.request().headers.toMap()
 
             socketId = socketLogManager.createSocket(
                 SocketConnection(
                     url = url,
                     requestHeaders = reqHeaders,
                     status = SocketStatus.Open,
                     timestamp = currentTimeMillis(),
                     protocol = response.protocol.toString(),
                 ),
             )
         }
 
         val wiretapSocket = WiretapWebSocket(webSocket, socketId, socketLogManager, config)
         delegate.onOpen(wiretapSocket, response)
     }
 
     override fun onMessage(webSocket: WebSocket, text: String) {
         if (isSocketActive) {
             logScope.launch {
                 socketLogManager.logSocketMsg(
                     SocketMessage(
                         socketId = socketId,
                         direction = SocketMessageType.Received,
                         contentType = SocketContentType.Text,
                         content = text,
                         byteCount = text.encodeToByteArray().size.toLong(),
                         timestamp = currentTimeMillis(),
                     ),
                 )
             }
         }
 
         delegate.onMessage(webSocket, text)
     }
 
     override fun onMessage(webSocket: WebSocket, bytes: ByteString) {
         if (isSocketActive) {
             val payload = bytes.toByteArray()
             logScope.launch {
                 socketLogManager.logSocketMsg(
                     SocketMessage(
                         socketId = socketId,
                         direction = SocketMessageType.Received,
                         contentType = SocketContentType.Binary,
                         content = decodeBinaryFrame(payload, config.binaryDecoding),
                         byteCount = payload.size.toLong(),
                         timestamp = currentTimeMillis(),
                     ),
                 )
             }
         }
 
         delegate.onMessage(webSocket, bytes)
     }
 
     override fun onClosing(webSocket: WebSocket, code: Int, reason: String) {
         if (isSocketActive) {
             logScope.launch {
                 socketLogManager.updateSocket(
                     SocketConnection(
                         id = socketId,
                         url = webSocket.request().url.toString(),
                         status = SocketStatus.Closing,
                         closeCode = code,
                         closeReason = reason,
                         timestamp = currentTimeMillis(),
                     ),
                 )
             }
         }
 
         delegate.onClosing(webSocket, code, reason)
     }
 
     override fun onClosed(webSocket: WebSocket, code: Int, reason: String) {
         if (isSocketActive) {
             logScope.launch {
                 socketLogManager.updateSocket(
                     SocketConnection(
                         id = socketId,
                         url = webSocket.request().url.toString(),
                         status = SocketStatus.Closed,
                         closeCode = code,
                         closeReason = reason,
                         closedAt = currentTimeMillis(),
                         timestamp = currentTimeMillis(),
                     ),
                 )
             }
         }
 
         delegate.onClosed(webSocket, code, reason)
     }
 
     override fun onFailure(webSocket: WebSocket, t: Throwable, response: Response?) {
         if (isSocketActive) {
             logScope.launch {
                 socketLogManager.updateSocket(
                     SocketConnection(
                         id = socketId,
                         url = webSocket.request().url.toString(),
                         status = SocketStatus.Failed,
                         failureMessage = t.message ?: t::class.simpleName ?: "Unknown error",
                         closedAt = currentTimeMillis(),
                         timestamp = currentTimeMillis(),
                     ),
                 )
             }
         }
 
         delegate.onFailure(webSocket, t, response)
     }
 }
 
 /**
  * Wraps this [WebSocketListener] with Wiretap logging.
  *
  * ```kotlin
  * client.newWebSocket(request, myListener.wiretapped())
  * ```
  */
 @Deprecated(
     message = "Use wiretapped() with config builder instead.",
     replaceWith = ReplaceWith("wiretapped { enabled = true }"),
 )
 fun WebSocketListener.wiretapped(
     enabled: Boolean = true,
 ): WebSocketListener = WiretapOkHttpWebSocketListener(
     delegate = this,
     config = WiretapWsConfig().apply { this.enabled = enabled },
 )
 
 fun WebSocketListener.wiretapped(
     configure: WiretapWsConfig.() -> Unit,
 ): WebSocketListener = WiretapOkHttpWebSocketListener(
     delegate = this,
     config = WiretapWsConfig().apply(configure),
 )