1 /*
2  * JBoss, Home of Professional Open Source.
3  * Copyright 2014 Red Hat, Inc., and individual contributors
4  * as indicated by the @author tags.
5  *
6  * Licensed under the Apache License, Version 2.0 (the "License");
7  * you may not use this file except in compliance with the License.
8  * You may obtain a copy of the License at
9  *
10  *     http://www.apache.org/licenses/LICENSE-2.0
11  *
12  *  Unless required by applicable law or agreed to in writing, software
13  *  distributed under the License is distributed on an "AS IS" BASIS,
14  *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15  *  See the License for the specific language governing permissions and
16  *  limitations under the License.
17  */

18
19 package io.undertow.websockets.jsr;
20
21 import java.io.IOException;
22 import java.security.AccessController;
23 import java.security.PrivilegedAction;
24 import java.util.Collections;
25 import java.util.Map;
26 import java.util.concurrent.ConcurrentHashMap;
27 import java.util.function.Supplier;
28
29 import javax.websocket.ContainerProvider;
30 import javax.websocket.WebSocketContainer;
31
32 import org.xnio.OptionMap;
33 import org.xnio.Options;
34 import org.xnio.Xnio;
35 import org.xnio.XnioWorker;
36 import io.undertow.connector.ByteBufferPool;
37 import io.undertow.server.DefaultByteBufferPool;
38 import io.undertow.servlet.api.ClassIntrospecter;
39 import io.undertow.servlet.api.InstanceFactory;
40 import io.undertow.servlet.util.DefaultClassIntrospector;
41
42 /**
43  * @author Stuart Douglas
44  */

45 public class UndertowContainerProvider extends ContainerProvider {
46
47     private static final boolean directBuffers = Boolean.getBoolean("io.undertow.websockets.direct-buffers");
48     private static final boolean invokeInIoThread = Boolean.getBoolean("io.undertow.websockets.invoke-in-io-thread");
49
50     private static final RuntimePermission PERMISSION = new RuntimePermission("io.undertow.websockets.jsr.MODIFY_WEBSOCKET_CONTAINER");
51
52     private static final Map<ClassLoader, WebSocketContainer> webSocketContainers = new ConcurrentHashMap<>();
53
54     private static volatile ServerWebSocketContainer defaultContainer;
55     private static volatile boolean defaultContainerDisabled = false;
56
57     private static final SwitchableClassIntrospector defaultIntrospector = new SwitchableClassIntrospector();
58
59     @Override
60     protected WebSocketContainer getContainer() {
61         ClassLoader tccl;
62         if (System.getSecurityManager() == null) {
63             tccl = Thread.currentThread().getContextClassLoader();
64         } else {
65             tccl = AccessController.doPrivileged(new PrivilegedAction<ClassLoader>() {
66                 @Override
67                 public ClassLoader run() {
68                     return Thread.currentThread().getContextClassLoader();
69                 }
70             });
71         }
72         WebSocketContainer webSocketContainer = webSocketContainers.get(tccl);
73         if (webSocketContainer == null) {
74             return getDefaultContainer();
75         }
76         return webSocketContainer;
77     }
78
79     static ServerWebSocketContainer getDefaultContainer() {
80         if (defaultContainerDisabled) {
81             return null;
82         }
83         if (defaultContainer != null) {
84             return defaultContainer;
85         }
86         synchronized (UndertowContainerProvider.class) {
87             if (defaultContainer == null) {
88                 //this is not great, as we have no way to control the lifecycle
89                 //but there is not much we can do
90                 //todo: what options should we use here?
91                 ByteBufferPool buffers = new DefaultByteBufferPool(directBuffers, 1024, 100, 12);
92                 defaultContainer = new ServerWebSocketContainer(defaultIntrospector, UndertowContainerProvider.class.getClassLoader(), new Supplier<XnioWorker>() {
93                     volatile XnioWorker worker;
94
95                     @Override
96                     public XnioWorker get() {
97                         if(worker == null) {
98                             synchronized (this) {
99                                 if(worker == null) {
100                                     try {
101                                         worker = Xnio.getInstance().createWorker(OptionMap.create(Options.THREAD_DAEMON, true));
102                                     } catch (IOException e) {
103                                         throw new RuntimeException(e);
104                                     }
105                                 }
106                             }
107                         }
108                         return worker;
109                     }
110                 }, buffers, Collections.EMPTY_LIST, !invokeInIoThread);
111             }
112             return defaultContainer;
113         }
114     }
115
116     public static void addContainer(final ClassLoader classLoader, final WebSocketContainer webSocketContainer) {
117         SecurityManager sm = System.getSecurityManager();
118         if (sm != null) {
119             sm.checkPermission(PERMISSION);
120         }
121         webSocketContainers.put(classLoader, webSocketContainer);
122     }
123
124     public static void removeContainer(final ClassLoader classLoader) {
125         SecurityManager sm = System.getSecurityManager();
126         if (sm != null) {
127             sm.checkPermission(PERMISSION);
128         }
129         webSocketContainers.remove(classLoader);
130     }
131
132     public void setDefaultClassIntrospector(ClassIntrospecter classIntrospector) {
133         if (classIntrospector == null) {
134             throw new IllegalArgumentException();
135         }
136         defaultIntrospector.setIntrospecter(classIntrospector);
137     }
138
139     public static void disableDefaultContainer() {
140         SecurityManager sm = System.getSecurityManager();
141         if (sm != null) {
142             sm.checkPermission(PERMISSION);
143         }
144         defaultContainerDisabled = true;
145     }
146
147
148     private static class SwitchableClassIntrospector implements ClassIntrospecter {
149
150         private volatile ClassIntrospecter introspecter = DefaultClassIntrospector.INSTANCE;
151
152         @Override
153         public <T> InstanceFactory<T> createInstanceFactory(Class<T> clazz) throws NoSuchMethodException {
154             return introspecter.createInstanceFactory(clazz);
155         }
156
157         public void setIntrospecter(ClassIntrospecter introspecter) {
158             this.introspecter = introspecter;
159         }
160     }
161 }
162