Skip to content

Commit

Permalink
Activate request context for websockets
Browse files Browse the repository at this point in the history 
Fixes quarkusio#9994
  • Loading branch information
@stuartwdouglas
stuartwdouglas committed Jun 15, 2020
1 parent 399e16b commit 0c4cca6
Show file tree
Hide file tree
Showing 5 changed files with 94 additions and 90 deletions.
11 changes: 11 additions & 0 deletions ...-websockets/deployment/src/test/java/io/quarkus/undertow/websockets/test/EchoService.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
package io.quarkus.undertow.websockets.test;

import javax.enterprise.context.RequestScoped;

@RequestScoped
public class EchoService {

public String echo(String msg) {
return msg;
}
}
6 changes: 5 additions & 1 deletion ...ebsockets/deployment/src/test/java/io/quarkus/undertow/websockets/test/EchoWebSocket.java
View file
Original file line number Diff line number Diff line change
@@ -1,14 +1,18 @@
package io.quarkus.undertow.websockets.test;

import javax.inject.Inject;
import javax.websocket.OnMessage;
import javax.websocket.server.ServerEndpoint;

@ServerEndpoint("/echo")
public class EchoWebSocket {

@Inject
EchoService echoService;

@OnMessage
String echo(String msg) {
return msg;
return echoService.echo(msg);
}

}
4 changes: 2 additions & 2 deletions ...eployment/src/test/java/io/quarkus/undertow/websockets/test/WebsocketDevModeTestCase.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ public class WebsocketDevModeTestCase {
@Override
public JavaArchive get() {
return ShrinkWrap.create(JavaArchive.class)
.addClass(EchoWebSocket.class);
.addClasses(EchoWebSocket.class, EchoService.class);
}
});

Expand All @@ -62,7 +62,7 @@ public void onMessage(String s) {
session.close();
}

test.modifySourceFile(EchoWebSocket.class, (s) -> s.replace("return msg;", "return \"changed:\" + msg;"));
test.modifySourceFile(EchoService.class, (s) -> s.replace("return msg;", "return \"changed:\" + msg;"));

session = ContainerProvider.getWebSocketContainer().connectToServer(new Endpoint() {
@Override
Expand Down
11 changes: 0 additions & 11 deletions ...loyment/src/main/java/io/quarkus/undertow/deployment/UndertowArcIntegrationBuildStep.java
View file
Original file line number Diff line number Diff line change
@@ -1,22 +1,11 @@
package io.quarkus.undertow.deployment;

import io.quarkus.arc.deployment.BeanContainerBuildItem;
import io.quarkus.arc.deployment.BeanDefiningAnnotationBuildItem;
import io.quarkus.deployment.annotations.BuildProducer;
import io.quarkus.deployment.annotations.BuildStep;
import io.quarkus.deployment.annotations.ExecutionTime;
import io.quarkus.deployment.annotations.Record;
import io.quarkus.undertow.runtime.UndertowDeploymentRecorder;

public class UndertowArcIntegrationBuildStep {

@BuildStep
@Record(ExecutionTime.STATIC_INIT)
ServletExtensionBuildItem integrateRequestContext(BeanContainerBuildItem beanContainerBuildItem,
UndertowDeploymentRecorder recorder) {
return new ServletExtensionBuildItem(recorder.setupRequestScope(beanContainerBuildItem.getValue()));
}

@BuildStep
void beanDefiningAnnotations(BuildProducer<BeanDefiningAnnotationBuildItem> annotations) {
annotations.produce(new BeanDefiningAnnotationBuildItem(UndertowBuildStep.WEB_FILTER));
Expand Down
152 changes: 76 additions & 76 deletions ...ndertow/runtime/src/main/java/io/quarkus/undertow/runtime/UndertowDeploymentRecorder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -429,6 +429,7 @@ public DeploymentManager bootServletContainer(RuntimeValue<DeploymentInfo> info,
.addInitParam(QuarkusErrorServlet.SHOW_STACK, Boolean.toString(launchMode.isDevOrTest())));
}
}
setupRequestScope(info.getValue(), beanContainer);

try {
ClassIntrospecter defaultVal = info.getValue().getClassIntrospecter();
Expand Down Expand Up @@ -519,91 +520,90 @@ public void addServletExtension(RuntimeValue<DeploymentInfo> deployment, Servlet
deployment.getValue().addServletExtension(extension);
}

public ServletExtension setupRequestScope(BeanContainer beanContainer) {
return new ServletExtension() {
public void setupRequestScope(DeploymentInfo deploymentInfo, BeanContainer beanContainer) {
CurrentVertxRequest currentVertxRequest = CDI.current().select(CurrentVertxRequest.class).get();
Instance<CurrentIdentityAssociation> identityAssociations = CDI.current()
.select(CurrentIdentityAssociation.class);
CurrentIdentityAssociation association;
if (identityAssociations.isResolvable()) {
association = identityAssociations.get();
} else {
association = null;
}
deploymentInfo.addThreadSetupAction(new ThreadSetupHandler() {
@Override
public void handleDeployment(DeploymentInfo deploymentInfo, ServletContext servletContext) {
CurrentVertxRequest currentVertxRequest = CDI.current().select(CurrentVertxRequest.class).get();
Instance<CurrentIdentityAssociation> identityAssociations = CDI.current()
.select(CurrentIdentityAssociation.class);
CurrentIdentityAssociation association;
if (identityAssociations.isResolvable()) {
association = identityAssociations.get();
} else {
association = null;
}
deploymentInfo.addThreadSetupAction(new ThreadSetupHandler() {
public <T, C> ThreadSetupHandler.Action<T, C> create(Action<T, C> action) {
return new Action<T, C>() {
@Override
public <T, C> ThreadSetupHandler.Action<T, C> create(Action<T, C> action) {
return new Action<T, C>() {
@Override
public T call(HttpServerExchange exchange, C context) throws Exception {
// Not sure what to do here
if (exchange == null) {
return action.call(exchange, context);
public T call(HttpServerExchange exchange, C context) throws Exception {
// Not sure what to do here
ManagedContext requestContext = beanContainer.requestContext();
if (requestContext.isActive()) {
return action.call(exchange, context);
} else if (exchange == null) {
requestContext.activate();
try {
return action.call(exchange, context);
} finally {
requestContext.destroy();
}
} else {
InjectableContext.ContextState existingRequestContext = exchange
.getAttachment(REQUEST_CONTEXT);
try {
requestContext.activate(existingRequestContext);

VertxHttpExchange delegate = (VertxHttpExchange) exchange.getDelegate();
RoutingContext rc = (RoutingContext) delegate.getContext();
currentVertxRequest.setCurrent(rc);

if (association != null) {
association
.setIdentity(QuarkusHttpUser.getSecurityIdentity(rc, null));
}
ManagedContext requestContext = beanContainer.requestContext();
if (requestContext.isActive()) {
return action.call(exchange, context);
} else {
InjectableContext.ContextState existingRequestContext = exchange
.getAttachment(REQUEST_CONTEXT);
try {
requestContext.activate(existingRequestContext);

VertxHttpExchange delegate = (VertxHttpExchange) exchange.getDelegate();
RoutingContext rc = (RoutingContext) delegate.getContext();
currentVertxRequest.setCurrent(rc);

if (association != null) {
association
.setIdentity(QuarkusHttpUser.getSecurityIdentity(rc, null));
}

return action.call(exchange, context);
} finally {
ServletRequestContext src = exchange
.getAttachment(ServletRequestContext.ATTACHMENT_KEY);
HttpServletRequestImpl req = src.getOriginalRequest();
if (req.isAsyncStarted()) {
exchange.putAttachment(REQUEST_CONTEXT, requestContext.getState());
requestContext.deactivate();
if (existingRequestContext == null) {
req.getAsyncContextInternal().addListener(new AsyncListener() {
@Override
public void onComplete(AsyncEvent event) throws IOException {
requestContext.activate(exchange
.getAttachment(REQUEST_CONTEXT));
requestContext.terminate();
}

@Override
public void onTimeout(AsyncEvent event) throws IOException {
onComplete(event);
}

@Override
public void onError(AsyncEvent event) throws IOException {
onComplete(event);
}

@Override
public void onStartAsync(AsyncEvent event) throws IOException {

}
});

return action.call(exchange, context);
} finally {
ServletRequestContext src = exchange
.getAttachment(ServletRequestContext.ATTACHMENT_KEY);
HttpServletRequestImpl req = src.getOriginalRequest();
if (req.isAsyncStarted()) {
exchange.putAttachment(REQUEST_CONTEXT, requestContext.getState());
requestContext.deactivate();
if (existingRequestContext == null) {
req.getAsyncContextInternal().addListener(new AsyncListener() {
@Override
public void onComplete(AsyncEvent event) throws IOException {
requestContext.activate(exchange
.getAttachment(REQUEST_CONTEXT));
requestContext.terminate();
}

@Override
public void onTimeout(AsyncEvent event) throws IOException {
onComplete(event);
}
} else {
requestContext.terminate();
}

@Override
public void onError(AsyncEvent event) throws IOException {
onComplete(event);
}

@Override
public void onStartAsync(AsyncEvent event) throws IOException {

}
});
}
} else {
requestContext.terminate();
}
}
};
}
}
});
};
}
};
});
}

public void addServletContainerInitializer(RuntimeValue<DeploymentInfo> deployment,
Expand Down

0 comments on commit 0c4cca6

Please sign in to comment.