/*
 * Copyright 2002-2013 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
 *
 * http://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.socket.server.support;

import java.io.IOException;
import java.lang.reflect.Method;
import java.util.Collections;

import javax.servlet.ServletContext;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.websocket.Endpoint;
import javax.websocket.server.ServerEndpointConfig;

import org.apache.tomcat.websocket.server.WsHandshakeRequest;
import org.apache.tomcat.websocket.server.WsHttpUpgradeHandler;
import org.apache.tomcat.websocket.server.WsServerContainer;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.util.Assert;
import org.springframework.util.ReflectionUtils;
import org.springframework.web.socket.server.HandshakeFailureException;
import org.springframework.web.socket.server.endpoint.ServerEndpointRegistration;

/**
 * Tomcat support for upgrading an {@link HttpServletRequest} during a WebSocket handshake.
 *
 * @author Rossen Stoyanchev
 * @since 4.0
 */
public class TomcatRequestUpgradeStrategy extends AbstractStandardUpgradeStrategy {

	@Override
	public String[] getSupportedVersions() {
		return new String[] { "13" };
	}

	@Override
	public void upgradeInternal(ServerHttpRequest request, ServerHttpResponse response,
			String acceptedProtocol, Endpoint endpoint) throws IOException {

		Assert.isTrue(request instanceof ServletServerHttpRequest);
		HttpServletRequest servletRequest = ((ServletServerHttpRequest) request).getServletRequest();

		WsHttpUpgradeHandler upgradeHandler;
		try {
			upgradeHandler = servletRequest.upgrade(WsHttpUpgradeHandler.class);
		}
		catch (ServletException e) {
			throw new HandshakeFailureException("Unable to create UpgardeHandler", e);
		}

		WsHandshakeRequest webSocketRequest = new WsHandshakeRequest(servletRequest);
		try {
			Method method = ReflectionUtils.findMethod(WsHandshakeRequest.class, "finished");
			ReflectionUtils.makeAccessible(method);
			method.invoke(webSocketRequest);
		}
		catch (Exception ex) {
			throw new HandshakeFailureException("Failed to upgrade HttpServletRequest", ex);
		}

		String attribute = "javax.websocket.server.ServerContainer";
		ServletContext servletContext = servletRequest.getServletContext();
		WsServerContainer serverContainer = (WsServerContainer) servletContext.getAttribute(attribute);

		ServerEndpointConfig endpointConfig = new ServerEndpointRegistration("/shouldntmatter", endpoint);

		upgradeHandler.preInit(endpoint, endpointConfig, serverContainer, webSocketRequest,
				acceptedProtocol, Collections.<String, String> emptyMap(), servletRequest.isSecure());
	}

}

@@ -682,55 +682,19 @@
 ava.
-io.IOException;%0Aimport java.lang.reflect.Method
+util.Arrays
 ;%0Aim
@@ -720,16 +720,38 @@
 ections;
+%0Aimport java.util.Map;
 %0A%0Aimport
@@ -785,47 +785,8 @@
 xt;%0A
-import javax.servlet.ServletException;%0A
 impo
@@ -844,206 +844,75 @@
 vax.
-websocket.Endpoint;%0Aimport javax.websocket.server.ServerEndpointConfig;%0A%0Aimport org.apache.tomcat.websocket.server.WsHandshakeRequest;%0Aimport org.apache.tomcat.websocket.server.WsHttpUpgradeHandler;
+servlet.http.HttpServletResponse;%0Aimport javax.websocket.Endpoint;%0A
 %0Aimp
@@ -1182,19 +1182,45 @@
 ork.
-util.Assert
+http.server.ServletServerHttpResponse
 ;%0Aim
@@ -1249,31 +1249,22 @@
 rk.util.
-ReflectionUtils
+Assert
 ;%0Aimport
@@ -1643,16 +1643,17 @@
 tegy %7B%0A%0A
+%0A
 %09@Overri
@@ -1883,18 +1883,32 @@
  throws 
-IO
+HandshakeFailure
 Exceptio
@@ -2077,311 +2077,204 @@
 %0A%0A%09%09
-WsHttpUpgradeHandler upgradeHandler;%0A%09%09try %7B%0A%09%09%09upgradeHandler = servletRequest.upgrade(WsHttpUpgradeHandler.class);%0A%09%09%7D%0A%09%09catch (ServletException e) %7B%0A%09%09%09throw new HandshakeFailureException(%22Unable to create UpgardeHandler%22, e);%0A%09%09%7D%0A%0A%09%09WsHandshakeRequest webSocketRequest = new WsHandshake
+Assert.isTrue(response instanceof ServletServerHttpResponse);%0A%09%09HttpServletResponse servletResponse = ((ServletServerHttpResponse) response).getServletResponse();%0A%0A%09%09String path = servlet
 Request
-(servl
+.g
 etRe
@@ -2282,180 +2282,383 @@
 uest
-);%0A%09%09try %7B%0A%09%09%09Method method = ReflectionUtils.findMethod(WsHandshakeRequest.class, %22finished%22);%0A%09%09%09ReflectionUtils.makeAccessible(method
+URI(); // shouldn't matter%0A%09%09Map%3CString, String%3E pathParams = Collections.%3CString, String%3E emptyMap();%0A%0A%09%09ServerEndpointRegistration endpointConfig = new ServerEndpointRegistration(path, endpoint);%0A%09%09endpointConfig.setSubprotocols(Arrays.asList(acceptedProtocol)
 );%0A
+%0A
 %09%09
-%09method.invoke(webSocketRequest
+try %7B%0A%09%09%09getContainer(servletRequest).doUpgrade(servletRequest, servletResponse, endpointConfig, pathParams
 );%0A%09
@@ -2773,16 +2773,96 @@
 x);%0A%09%09%7D%0A
+%09%7D%0A%0A%09private WsServerContainer getContainer(HttpServletRequest servletRequest) %7B
 %0A%09%09Strin
@@ -2993,43 +2993,14 @@
 ;%0A%09%09
-WsServerContainer serverContainer =
+return
  (Ws
@@ -3059,288 +3059,8 @@
 te);
-%0A%0A%09%09ServerEndpointConfig endpointConfig = new ServerEndpointRegistration(%22/shouldntmatter%22, endpoint);%0A%0A%09%09upgradeHandler.preInit(endpoint, endpointConfig, serverContainer, webSocketRequest,%0A%09%09%09%09acceptedProtocol, Collections.%3CString, String%3E emptyMap(), servletRequest.isSecure());
 %0A%09%7D%0A

Make use of Tomcat's new WsServerContainer.doUpgrade