From 15843aca33c1a730c9ddadb0ad632ad2539668df Mon Sep 17 00:00:00 2001
From: Andy Wilkinson <wilkinsona@vmware.com>
Date: Mon, 12 Sep 2022 14:50:17 +0100
Subject: [PATCH] Add runtime hints for WebFlux's WebSocket upgrade strategies

---
 .../RequestUpgradeStrategyRuntimeHints.java   | 47 +++++++++++++++++++
 .../resources/META-INF/spring/aot.factories   |  2 +
 2 files changed, 49 insertions(+)
 create mode 100644 spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/RequestUpgradeStrategyRuntimeHints.java
 create mode 100644 spring-webflux/src/main/resources/META-INF/spring/aot.factories

diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/RequestUpgradeStrategyRuntimeHints.java b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/RequestUpgradeStrategyRuntimeHints.java
new file mode 100644
index 000000000000..47e022cbb0d1
--- /dev/null
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/server/upgrade/RequestUpgradeStrategyRuntimeHints.java
@@ -0,0 +1,47 @@
+/*
+ * Copyright 2002-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.web.reactive.socket.server.upgrade;
+
+import org.springframework.aot.hint.MemberCategory;
+import org.springframework.aot.hint.RuntimeHints;
+import org.springframework.aot.hint.RuntimeHintsRegistrar;
+import org.springframework.aot.hint.TypeReference;
+
+/**
+ * {@link RuntimeHintsRegistrar} for {@link RequestUpgradeStrategy} implementations.
+ *
+ * @author Andy Wilkinson
+ */
+class RequestUpgradeStrategyRuntimeHints implements RuntimeHintsRegistrar {
+
+	@Override
+	public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
+		hints.reflection().registerType(JettyRequestUpgradeStrategy.class,
+				hint -> hint.withMembers(MemberCategory.INVOKE_PUBLIC_CONSTRUCTORS).onReachableType(
+						TypeReference.of("org.eclipse.jetty.websocket.server.JettyWebSocketServerContainer")));
+		hints.reflection().registerType(ReactorNettyRequestUpgradeStrategy.class,
+				hint -> hint.withMembers(MemberCategory.INVOKE_PUBLIC_CONSTRUCTORS)
+						.onReachableType(TypeReference.of("reactor.netty.http.server.HttpServerResponse")));
+		hints.reflection().registerType(TomcatRequestUpgradeStrategy.class,
+				hint -> hint.withMembers(MemberCategory.INVOKE_PUBLIC_CONSTRUCTORS)
+						.onReachableType(TypeReference.of("org.apache.tomcat.websocket.server.WsHttpUpgradeHandler")));
+		hints.reflection().registerType(UndertowRequestUpgradeStrategy.class,
+				hint -> hint.withMembers(MemberCategory.INVOKE_PUBLIC_CONSTRUCTORS)
+						.onReachableType(TypeReference.of("io.undertow.websockets.WebSocketProtocolHandshakeHandler")));
+	}
+
+}
diff --git a/spring-webflux/src/main/resources/META-INF/spring/aot.factories b/spring-webflux/src/main/resources/META-INF/spring/aot.factories
new file mode 100644
index 000000000000..48bfb76b50db
--- /dev/null
+++ b/spring-webflux/src/main/resources/META-INF/spring/aot.factories
@@ -0,0 +1,2 @@
+org.springframework.aot.hint.RuntimeHintsRegistrar= \
+org.springframework.web.reactive.socket.server.upgrade.RequestUpgradeStrategyRuntimeHints
\ No newline at end of file