001    package org.granite.gravity.tomcat;
002    
003    import java.io.IOException;
004    import java.security.MessageDigest;
005    import java.security.NoSuchAlgorithmException;
006    import java.util.ArrayList;
007    import java.util.Collections;
008    import java.util.Enumeration;
009    import java.util.List;
010    import java.util.Queue;
011    import java.util.concurrent.ConcurrentLinkedQueue;
012    
013    import javax.servlet.ServletConfig;
014    import javax.servlet.ServletException;
015    import javax.servlet.ServletRequest;
016    import javax.servlet.ServletRequestWrapper;
017    import javax.servlet.http.HttpServletRequest;
018    import javax.servlet.http.HttpServletResponse;
019    import javax.servlet.http.HttpSession;
020    
021    import org.apache.catalina.connector.RequestFacade;
022    import org.apache.catalina.util.Base64;
023    import org.apache.catalina.websocket.Constants;
024    import org.apache.catalina.websocket.StreamInbound;
025    import org.apache.catalina.websocket.WebSocketServlet;
026    import org.apache.tomcat.util.buf.B2CConverter;
027    import org.apache.tomcat.util.res.StringManager;
028    import org.granite.context.GraniteContext;
029    import org.granite.gravity.Gravity;
030    import org.granite.gravity.GravityManager;
031    import org.granite.gravity.GravityServletUtil;
032    import org.granite.logging.Logger;
033    import org.granite.messaging.webapp.ServletGraniteContext;
034    
035    import flex.messaging.messages.CommandMessage;
036    import flex.messaging.messages.Message;
037    
038    
039    public class TomcatWebSocketServlet extends WebSocketServlet {
040            
041            private static final long serialVersionUID = 1L;
042            
043            private static final Logger log = Logger.getLogger(TomcatWebSocketServlet.class);
044            
045        private static final byte[] WS_ACCEPT = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11".getBytes(B2CConverter.ISO_8859_1);
046        private static final StringManager sm = StringManager.getManager(Constants.Package);
047    
048        private final Queue<MessageDigest> sha1Helpers = new ConcurrentLinkedQueue<MessageDigest>();
049            
050            @Override
051            public void init(ServletConfig config) throws ServletException {
052                    super.init(config);
053                    
054                    GravityServletUtil.init(config);
055            }
056            
057            @Override
058            protected String selectSubProtocol(List<String> subProtocols) {
059                    return subProtocols != null && subProtocols.contains("org.granite.gravity") ? "org.granite.gravity" : null;
060            }
061            
062            @Override
063            protected StreamInbound createWebSocketInbound(String protocol, HttpServletRequest request) {
064                    Gravity gravity = GravityManager.getGravity(getServletContext());
065                    TomcatWebSocketChannelFactory channelFactory = new TomcatWebSocketChannelFactory(gravity, getServletContext());
066                    
067                    try {
068                            String connectMessageId = request.getHeader("connectId") != null ? request.getHeader("connectId") : request.getParameter("connectId");
069                            String clientId = request.getHeader("GDSClientId") != null ? request.getHeader("GDSClientId") : request.getParameter("GDSClientId");
070                            String clientType = request.getHeader("GDSClientType") != null ? request.getHeader("GDSClientType") : request.getParameter("GDSClientType");
071                            String sessionId = null;
072                            HttpSession session = request.getSession(false);
073                            if (session != null) {
074                            ServletGraniteContext.createThreadInstance(gravity.getGraniteConfig(), gravity.getServicesConfig(), 
075                                    getServletContext(), session, clientType);
076                            
077                                    sessionId = session.getId();
078                            }
079                            else {
080                                    for (int i = 0; i < request.getCookies().length; i++) {
081                                            if ("JSESSIONID".equals(request.getCookies()[i].getName())) {
082                                                    sessionId = request.getCookies()[i].getValue();
083                                                    break;
084                                            }
085                                    }                               
086                                    
087                                    ServletGraniteContext.createThreadInstance(gravity.getGraniteConfig(), gravity.getServicesConfig(), 
088                                    getServletContext(), sessionId, clientType); 
089                            }
090                            
091                            log.info("WebSocket connection started %s clientId %s sessionId %s", protocol, clientId, sessionId);
092                            
093                            CommandMessage pingMessage = new CommandMessage();
094                            pingMessage.setMessageId(connectMessageId != null ? connectMessageId : "OPEN_CONNECTION");
095                            pingMessage.setOperation(CommandMessage.CLIENT_PING_OPERATION);
096                            if (clientId != null)
097                                    pingMessage.setClientId(clientId);
098                            
099                            Message ackMessage = gravity.handleMessage(channelFactory, pingMessage);
100                            
101                            TomcatWebSocketChannel channel = gravity.getChannel(channelFactory, (String)ackMessage.getClientId());
102                            
103                            if (!ackMessage.getClientId().equals(clientId))
104                                    channel.setConnectAckMessage(ackMessage);
105                            
106                            return channel.getStreamInbound();
107                    }
108                    finally {
109                            GraniteContext.release();
110                    }
111            }
112            
113        @Override
114        protected void doGet(HttpServletRequest req, HttpServletResponse resp) throws ServletException, IOException {
115    
116            // Information required to send the server handshake message
117            String key;
118            String subProtocol = null;
119            List<String> extensions = Collections.emptyList();
120    
121            if (!headerContainsToken(req, "upgrade", "websocket")) {
122                resp.sendError(HttpServletResponse.SC_BAD_REQUEST);
123                return;
124            }
125    
126            if (!headerContainsToken(req, "connection", "upgrade")) {
127                resp.sendError(HttpServletResponse.SC_BAD_REQUEST);
128                return;
129            }
130    
131            if (!headerContainsToken(req, "sec-websocket-version", "13")) {
132                resp.setStatus(426);
133                resp.setHeader("Sec-WebSocket-Version", "13");
134                return;
135            }
136    
137            key = req.getHeader("Sec-WebSocket-Key");
138            if (key == null) {
139                resp.sendError(HttpServletResponse.SC_BAD_REQUEST);
140                return;
141            }
142    
143            String origin = req.getHeader("Origin");
144            if (!verifyOrigin(origin)) {
145                resp.sendError(HttpServletResponse.SC_FORBIDDEN);
146                return;
147            }
148            
149            // Fix for Tomcat-7.0.29 bad header name (was Sec-WebSocket-Protocol-Client")
150            List<String> subProtocols = getTokensFromHeader(req, "Sec-WebSocket-Protocol");
151            if (!subProtocols.isEmpty())
152                subProtocol = selectSubProtocol(subProtocols);
153    
154            // TODO Read client handshake - Sec-WebSocket-Extensions
155    
156            // TODO Extensions require the ability to specify something (API TBD)
157            //      that can be passed to the Tomcat internals and process extension
158            //      data present when the frame is fragmented.
159    
160            // If we got this far, all is good. Accept the connection.
161            resp.setHeader("Upgrade", "websocket");
162            resp.setHeader("Connection", "upgrade");
163            resp.setHeader("Sec-WebSocket-Accept", getWebSocketAccept(key));
164            if (subProtocol != null)
165                resp.setHeader("Sec-WebSocket-Protocol", subProtocol);
166            
167            if (!extensions.isEmpty()) {
168                // TODO
169            }
170    
171            WsHttpServletRequestWrapper wrapper = new WsHttpServletRequestWrapper(req);
172            StreamInbound inbound = createWebSocketInbound(subProtocol, wrapper);
173            wrapper.invalidate();
174            
175            // Hack to avoid chunked transfer
176            resp.setContentLength(((TomcatWebSocketChannel.MessageInboundImpl)inbound).getAckLength());
177            
178            // Small hack until the Servlet API provides a way to do this.
179            ServletRequest inner = req;
180            // Unwrap the request
181            while (inner instanceof ServletRequestWrapper)
182                inner = ((ServletRequestWrapper)inner).getRequest();
183            
184            if (inner instanceof RequestFacade)
185                ((RequestFacade)inner).doUpgrade(inbound);
186            else
187                resp.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR, sm.getString("servlet.reqUpgradeFail"));
188        }
189        
190        
191        private boolean headerContainsToken(HttpServletRequest req,
192                String headerName, String target) {
193            Enumeration<String> headers = req.getHeaders(headerName);
194            while (headers.hasMoreElements()) {
195                String header = headers.nextElement();
196                String[] tokens = header.split(",");
197                for (String token : tokens) {
198                    if (target.equalsIgnoreCase(token.trim())) {
199                        return true;
200                    }
201                }
202            }
203            return false;
204        }
205        
206        private List<String> getTokensFromHeader(HttpServletRequest req,
207                String headerName) {
208            List<String> result = new ArrayList<String>();
209    
210            Enumeration<String> headers = req.getHeaders(headerName);
211            while (headers.hasMoreElements()) {
212                String header = headers.nextElement();
213                String[] tokens = header.split(",");
214                for (String token : tokens) {
215                    result.add(token.trim());
216                }
217            }
218            return result;
219        }
220        
221        private String getWebSocketAccept(String key) throws ServletException {
222    
223            MessageDigest sha1Helper = sha1Helpers.poll();
224            if (sha1Helper == null) {
225                try {
226                    sha1Helper = MessageDigest.getInstance("SHA1");
227                } catch (NoSuchAlgorithmException e) {
228                    throw new ServletException(e);
229                }
230            }
231    
232            sha1Helper.reset();
233            sha1Helper.update(key.getBytes(B2CConverter.ISO_8859_1));
234            String result = Base64.encode(sha1Helper.digest(WS_ACCEPT));
235    
236            sha1Helpers.add(sha1Helper);
237    
238            return result;
239        }
240    }