Spring Reactive WebSocket Extension

main.kt

import org.springframework.boot.SpringApplication
import org.springframework.boot.autoconfigure.SpringBootApplication
import org.springframework.context.ApplicationContextInitializer
import org.springframework.context.support.GenericApplicationContext
import org.springframework.web.reactive.config.EnableWebFlux
import org.springframework.web.reactive.handler.SimpleUrlHandlerMapping
import org.springframework.web.reactive.socket.WebSocketHandler
import org.springframework.web.reactive.socket.WebSocketSession
import org.springframework.context.support.beans
import reactor.core.publisher.Mono
import java.util.HashMap

@SpringBootApplication
@EnableWebFlux
class App

fun main(args: Array<String>) {
    SpringApplication(App::class.java).apply {
        addInitializers(ApplicationContextInitializer<GenericApplicationContext> {
            beans {
                bean {
                    wsRouter {
                        WS("/webSocket") { session ->
                            session.send(
                                    session.receive().map { session.textMessage(it.payloadAsText) }
                            )
                        }
                    }
                }

                bean { webSocketHandlerAdapter() }
            }
        })
    }.run(*args)
}

fun wsRouter(f: WsRoutes.() -> Unit): SimpleUrlHandlerMapping =
        SimpleUrlHandlerMapping().apply {
            WsRoutes().let {
                f(it)
                urlMap = it.urls
                order = it.order
            }
        }

@Suppress("FunctionName")
fun WsRoutes.WS(url: String, f: (session: WebSocketSession) -> Mono<Void>) {
    urls.put(url, WebSocketHandler { session -> f(session) })
}

class WsRoutes(
        val urls: HashMap<String, WebSocketHandler> = hashMapOf(),
        var order: Int = Int.MAX_VALUE
)