refactor code and fix body payloads not being sent

FireTail-io · May 19, 2024 · 6960cbf · 6960cbf
1 parent f18a446
commit 6960cbf
Show file tree

Hide file tree

Showing 8 changed files with 22 additions and 70 deletions.
diff --git a/examples/src/main/java/com/example/demo/DemoApplication.java b/examples/src/main/java/com/example/demo/DemoApplication.java
@@ -4,7 +4,6 @@
 import org.springframework.boot.SpringApplication;
 import org.springframework.boot.autoconfigure.SpringBootApplication;
 import org.springframework.web.bind.annotation.GetMapping;
-import org.springframework.web.bind.annotation.PostMapping;
 import org.springframework.web.bind.annotation.RestController;
 
 import java.time.Clock;
@@ -24,9 +23,4 @@ public static void main(String[] args) {
     public String hello() {
         return String.format("Hello %s, utc: %s!", LocalDateTime.now(), ZonedDateTime.now(Clock.systemUTC()));
     }
-
-    @PostMapping("/hello")
-    public String helloPOST() {
-        return String.format("Zaihan Hello %s, utc: %s!", LocalDateTime.now(), ZonedDateTime.now(Clock.systemUTC()));
-    }
 }
diff --git a/src/main/kotlin/io/firetail/logging/core/FiretailLogger.kt b/src/main/kotlin/io/firetail/logging/core/FiretailLogger.kt
@@ -4,7 +4,6 @@ import io.firetail.logging.spring.FiretailConfig
 import io.firetail.logging.servlet.SpringRequestWrapper
 import io.firetail.logging.servlet.SpringResponseWrapper
 import org.slf4j.LoggerFactory
-import java.nio.charset.Charset
 
 class FiretailLogger (val firetailConfig: FiretailConfig) {
     fun logRequest(wrappedRequest: SpringRequestWrapper) =
@@ -34,10 +33,8 @@ class FiretailLogger (val firetailConfig: FiretailConfig) {
         status: Int = wrappedResponse.status,
         duration: Long,
     ) {
-        val charset: Charset = Charsets.UTF_8
-
         LOGGER.info(
-            "${FiretailTemplate.logResponsePrefix} ms: $duration, status: $status, headers: ${wrappedResponse.allHeaders}, body: ${String(wrappedResponse.contentAsByteArray, charset)}",
+            "${FiretailTemplate.logResponsePrefix} ms: $duration, status: $status, headers: ${wrappedResponse.allHeaders}",
         )
     }
     companion object {

diff --git a/src/main/kotlin/io/firetail/logging/core/FiretailTemplate.kt b/src/main/kotlin/io/firetail/logging/core/FiretailTemplate.kt
@@ -29,8 +29,6 @@ class FiretailTemplate(private val firetailConfig: FiretailConfig, private val f
             setRequestProperty("CONTENT-TYPE", "application/nd-json")
         }
 
-        LOGGER.info("Payload: ${jsonBody}")
-
         // Write the JSON body to the request
         val outputStream: OutputStream = connection.outputStream
         with(outputStream) {

diff --git a/src/main/kotlin/io/firetail/logging/servlet/FiretailFilter.kt b/src/main/kotlin/io/firetail/logging/servlet/FiretailFilter.kt
@@ -22,6 +22,7 @@ import org.springframework.web.method.HandlerMethod
 import org.springframework.web.servlet.mvc.method.annotation.RequestMappingHandlerMapping
 import java.util.*
 import java.util.concurrent.CompletableFuture
+import java.nio.charset.Charset
 
 @Service
 @ConditionalOnClass(FiretailConfig::class)
@@ -31,6 +32,7 @@ class FiretailFilter(
     private val firetailConfig: FiretailConfig,
     private val firetailMapper: FiretailMapper,
 ) {
+
     @Autowired
     private lateinit var firetailBuffer: FiretailBuffer
 
@@ -58,19 +60,27 @@ class FiretailFilter(
                     val startTime = System.currentTimeMillis()
                     val wrappedRequest = SpringRequestWrapper(request)
                     firetailLogger.logRequest(wrappedRequest)
-                    val wrappedResponse = SpringResponseWrapper(response)
+		    val wrappedResponse = SpringResponseWrapper(response)
                     try {
                         with(wrappedResponse) {
                             setHeader(REQUEST_ID, firetailLogContext.get(REQUEST_ID))
                             setHeader(CORRELATION_ID, firetailLogContext.get(CORRELATION_ID))
                         }
+
                         chain.doFilter(wrappedRequest, wrappedResponse)
+
+		        val charset: Charset = Charsets.UTF_8
+			val responseArray: ByteArray = wrappedResponse.contentAsByteArray
+			val responseBody: String = responseArray.toString(charset) 
+			wrappedResponse.copyBodyToResponse()
+
                         val duration = System.currentTimeMillis() - startTime
                         firetailLogger.logResponse(wrappedResponse, duration = duration)
                         val firetailLog =
                             firetailMapper.from(
                                 wrappedRequest,
                                 wrappedResponse,
+				responseBody,
                                 duration,
                             )
                         CompletableFuture.runAsync {

diff --git a/src/main/kotlin/io/firetail/logging/servlet/FiretailMapper.kt b/src/main/kotlin/io/firetail/logging/servlet/FiretailMapper.kt
@@ -4,20 +4,16 @@ import com.fasterxml.jackson.databind.ObjectMapper
 import io.firetail.logging.core.FiretailData
 import io.firetail.logging.core.FtRequest
 import io.firetail.logging.core.FtResponse
-import io.firetail.logging.core.FiretailTemplate
 import jakarta.servlet.http.HttpServletRequest
 import jakarta.servlet.http.HttpServletResponse
-import io.firetail.logging.servlet.SpringResponseWrapper
 import java.util.HashMap
 import java.io.StringWriter
 import java.io.InputStreamReader
-import org.slf4j.LoggerFactory
-import java.nio.charset.Charset
 
 class FiretailMapper {
     private val objectMapper = ObjectMapper()
-    fun from(request: HttpServletRequest, response: HttpServletResponse, executionTime: Long): FiretailData {
-        return FiretailData(request = from(request), response = from(response), executionTime = executionTime.toInt())
+    fun from(request: HttpServletRequest, response: HttpServletResponse, responseBody: String, executionTime: Long): FiretailData {
+        return FiretailData(request = from(request), response = from(response, responseBody), executionTime = executionTime.toInt())
     }
 
     fun from(request: HttpServletRequest): FtRequest {
@@ -40,17 +36,14 @@ class FiretailMapper {
         )
     }
 
-    fun from(response: HttpServletResponse): FtResponse {
+    fun from(response: HttpServletResponse, body: String): FtResponse {
         val headers = response.headerNames
             .mapIndexed { _, value -> value to listOf(response.getHeader(value)) }
             .toMap()
-
-        val wrappedResponse = SpringResponseWrapper(response)
-        val charset: Charset = Charsets.UTF_8
 
         return FtResponse(
             statusCode = response.status,
-            body = String(wrappedResponse.contentAsByteArray, charset),
+            body = body,
             headers = headers,
         )
     }
@@ -63,8 +56,4 @@ class FiretailMapper {
     fun from(fireTailData: List<FiretailData>): String {
         return fireTailData.joinToString("\n") { objectMapper.writeValueAsString(it) }
     }
-
-    companion object {
-        private val LOGGER = LoggerFactory.getLogger(FiretailTemplate::class.java)
-    }
 }
diff --git a/src/main/kotlin/io/firetail/logging/servlet/SpringResponseWrapper.kt b/src/main/kotlin/io/firetail/logging/servlet/SpringResponseWrapper.kt
@@ -1,50 +1,16 @@
 package io.firetail.logging.servlet
 
 import io.firetail.logging.core.Constants.Companion.empty
-import io.firetail.logging.core.FiretailTemplate
 import jakarta.servlet.ServletOutputStream
 import jakarta.servlet.http.HttpServletResponse
 import jakarta.servlet.http.HttpServletResponseWrapper
 import java.io.OutputStreamWriter
 import java.io.PrintWriter
 import java.util.function.Consumer
-import java.io.ByteArrayOutputStream
-import java.io.IOException
+import org.springframework.web.util.ContentCachingResponseWrapper
 
-class SpringResponseWrapper(response: HttpServletResponse?) : HttpServletResponseWrapper(response) {
-    private var outputStream: ServletOutputStream? = null
-    private var writer: PrintWriter? = null
-    private var copier: ServletOutputStreamWrapper? = null
+class SpringResponseWrapper(response: HttpServletResponse?) : ContentCachingResponseWrapper(response) {
 
-    override fun getOutputStream(): ServletOutputStream {
-        check(writer == null) { "getWriter() has already been called on this response." }
-        copier = ServletOutputStreamWrapper(response.outputStream)
-        return copier!!
-    }
-
-    override fun getWriter(): PrintWriter {
-        check(outputStream == null) { "getOutputStream() has already been called on this response." }
-        if (writer == null) {
-            copier = ServletOutputStreamWrapper(response.outputStream)
-            writer = PrintWriter(OutputStreamWriter(copier!!, response.characterEncoding), true)
-        }
-        return writer!!
-    }
-
-    override fun flushBuffer() {
-        if (writer != null) {
-            writer!!.flush()
-        } else if (outputStream != null) {
-            copier!!.flush()
-        }
-    }
-
-    val contentAsByteArray: ByteArray
-        get() = if (copier != null) {
-            copier!!.getCopy()
-        } else {
-            empty
-        }
     val allHeaders: Map<String, String>
         get() {
             val headers: MutableMap<String, String> = HashMap()

diff --git a/src/test/kotlin/io/firetail/logging/FiretailMapperTest.kt b/src/test/kotlin/io/firetail/logging/FiretailMapperTest.kt
@@ -16,7 +16,6 @@ class FiretailMapperTest {
 
     private val firetailMapper = FiretailMapper()
 
-    /*
     @Test
     fun fromResponse() {
         val mockResponse: HttpServletResponse = Mockito.mock(HttpServletResponse::class.java)
@@ -26,9 +25,9 @@ class FiretailMapperTest {
         Assertions.assertThat(result.headers)
             .isNotNull
             .hasFieldOrPropertyWithValue(TEST, listOf(TEST_RESULTS))
-    }*/
+    }
 
-    /*@Test
+    @Test
     fun fromRequest() {
         val mockRequest: HttpServletRequest = Mockito.mock(HttpServletRequest::class.java)
 
@@ -47,7 +46,7 @@ class FiretailMapperTest {
         Assertions.assertThat(result.headers)
             .isNotNull
             .hasFieldOrPropertyWithValue(TEST, listOf(TEST_RESULTS))
-    }*/
+    }
 
     @Test
     fun jsonNd() {

diff --git a/src/test/kotlin/io/firetail/logging/RequestInterceptorTests.kt b/src/test/kotlin/io/firetail/logging/RequestInterceptorTests.kt
@@ -85,7 +85,6 @@ class RequestInterceptorTests {
         assertThat(restTemplate.interceptors).isNotEmpty.contains(firetailHeaderInterceptor)
     }
 
-    /*
     @Test
     fun fireTailRequestLoggingAndResponse() {
         firetailMDC.clear()
@@ -112,7 +111,7 @@ class RequestInterceptorTests {
 
         assertThat(firetailBuffer.size() == 1)
         assertThat(firetailBuffer.flush()).isEqualTo("success")
-    } */
+    }
 
     // Emulates a general MVC controller for which we want to
     // assert Firetail calls have been made.