package org.apache.catalina.servlets;

import org.apache.catalina.startup.Tomcat;
import org.apache.catalina.startup.TomcatBaseTest;
import org.junit.Test;

import javax.servlet.AsyncContext;
import javax.servlet.ServletException;
import javax.servlet.http.Cookie;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.File;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.io.Writer;
import java.net.Socket;
import java.nio.charset.Charset;
import java.util.concurrent.CountDownLatch;

import static org.junit.Assert.fail;

/**
 *
 */
public class TestAsyncServlet extends TomcatBaseTest {

    private static CountDownLatch latch1 = new CountDownLatch(1);
    private static CountDownLatch latch2 = new CountDownLatch(1);
    private static CountDownLatch latch3 = new CountDownLatch(1);
    private static boolean isSecondRequest = false;
    private static boolean requestStateLeaked = false;

    @Test
    public void testRequestStateLeak() throws Exception {
        Tomcat tomcat = getTomcatInstance();
        tomcat.getConnector().setProperty("processorCache", "1");
        tomcat.getConnector().setProperty("maxThreads", "1");

        final String contextPath = "/async";
        File appDir = new File(getBuildDirectory(), "webapps" + contextPath);
        org.apache.catalina.Context ctx = tomcat.addWebapp(null, "/", appDir.getAbsolutePath());
        Tomcat.addServlet(ctx, "async", new AsyncFlushServlet());
        ctx.addServletMapping("/*", "async");

        tomcat.start();

        String request1 = "GET /async HTTP/1.1\r\n" +
                "Host: localhost:" + getPort() + "\r\n" +
                "Connection: keep-alive\r\n" +
                "Cache-Control: max-age=0\r\n" +
                "Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8\r\n" +
                "User-Agent: Request1\r\n" +
                "Accept-Encoding: gzip,deflate,sdch\r\n" +
                "Accept-Language: en-US,en;q=0.8,fr;q=0.6,es;q=0.4\r\n" +
                "Cookie: something.that.should.not.leak=true\r\n" +
                "\r\n";

        String request2 = "GET /async HTTP/1.1\r\n" +
                "Host: localhost:" + getPort() + "\r\n" +
                "Connection: keep-alive\r\n" +
                "Cache-Control: max-age=0\r\n" +
                "Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8\r\n" +
                "User-Agent: Request2\r\n" +
                "Accept-Encoding: gzip,deflate,sdch\r\n" +
                "Accept-Language: en-US,en;q=0.8,fr;q=0.6,es;q=0.4\r\n" +
                "\r\n";

        try (final Socket connection = new Socket("localhost", getPort())) {
            connection.setSoLinger(true, 0);
            Writer writer = new OutputStreamWriter(connection.getOutputStream(), Charset.forName("US-ASCII"));
            writer.write(request1);
            writer.flush();

            latch1.await();
            connection.close();
        }

        latch2.await();
        isSecondRequest = true;

        try (final Socket connection = new Socket("localhost", getPort())) {
            connection.setSoLinger(true, 0);
            Writer writer = new OutputStreamWriter(connection.getOutputStream(), Charset.forName("US-ASCII"));
            writer.write(request2);
            writer.flush();
            connection.getInputStream().read();
        }

        latch3.await();
        if (requestStateLeaked) {
            fail("State leaked between requests!");
        }
    }

    @Override
    protected String getProtocol() {
        return "org.apache.coyote.http11.Http11NioProtocol";
    }

    private static class AsyncFlushServlet extends HttpServlet {

        @Override
        protected void doGet(HttpServletRequest req, HttpServletResponse resp) throws ServletException, IOException {
            if (isSecondRequest) {
                if (req.getCookies().length > 0) {
                    for (Cookie cookie : req.getCookies()) {
                        if (cookie.getName().equalsIgnoreCase("something.that.should.not.leak")) {
                            requestStateLeaked = true;
                        }
                    }
                }
                latch3.countDown();
            } else {
                req.getCookies(); // We have to do this so Tomcat will actually parse the cookies from the request
            }

            req.setAttribute("org.apache.catalina.ASYNC_SUPPORTED", true);
            AsyncContext asyncContext = req.startAsync();

            latch1.countDown();

            PrintWriter writer = asyncContext.getResponse().getWriter();
            writer.print('\n');
            writer.flush();

            latch2.countDown();
        }
    }

}