try to fix

zalando · Aug 16, 2021 · 274802c · 274802c
1 parent b9e5537
commit 274802c
Showing 1 changed file with 3 additions and 16 deletions.
diff --git a/logbook-ktor-server/src/main/kotlin/org/zalando/logbook/server/LogbookServer.kt b/logbook-ktor-server/src/main/kotlin/org/zalando/logbook/server/LogbookServer.kt
@@ -29,37 +29,24 @@ class LogbookServer(
     }
 
     companion object : ApplicationFeature<Application, Config, LogbookServer> {
-        private val requestKey: AttributeKey<ServerRequest> = AttributeKey("ServerRequest")
-        private val requestWritingStageKey: AttributeKey<Logbook.RequestWritingStage> = AttributeKey("Logbook.RequestWritingStage")
         private val responseProcessingStageKey: AttributeKey<Logbook.ResponseProcessingStage> = AttributeKey("Logbook.ResponseProcessingStage")
         override val key: AttributeKey<LogbookServer> = AttributeKey("LogbookServer")
         override fun install(pipeline: Application, configure: Config.() -> Unit): LogbookServer {
             val config = Config().apply(configure)
             val feature = LogbookServer(config.logbook)
 
-            pipeline.intercept(ApplicationCallPipeline.Monitoring) {
+            pipeline.receivePipeline.intercept(ApplicationReceivePipeline.Before) {
                 val request = ServerRequest(call.request)
                 val requestWritingStage = feature.logbook.process(request)
                 if (request.shouldBuffer() && call.request.receiveChannel().availableForRead > 0) {
-                    call.attributes.put(requestWritingStageKey, requestWritingStage)
-                    call.attributes.put(requestKey, request)
-                } else {
-                    val responseProcessingStage = requestWritingStage.write()
-                    call.attributes.put(responseProcessingStageKey, responseProcessingStage)
-                }
-                proceed()
-            }
-
-            pipeline.receivePipeline.intercept(ApplicationReceivePipeline.Before) {
-                val requestWritingStage = context.attributes.getOrNull(requestWritingStageKey)
-                if (requestWritingStage != null) {
-                    val request = context.attributes[requestKey]
                     val content = call.request.receiveChannel().readBytes()
                     request.buffer(content)
                     val responseProcessingStage = requestWritingStage.write()
                     call.attributes.put(responseProcessingStageKey, responseProcessingStage)
                     proceedWith(ApplicationReceiveRequest(it.typeInfo, ByteReadChannel(content)))
                 } else {
+                    val responseProcessingStage = requestWritingStage.write()
+                    call.attributes.put(responseProcessingStageKey, responseProcessingStage)
                     proceed()
                 }
             }