/*
 *  Licensed to the Apache Software Foundation (ASF) under one or more
 *  contributor license agreements.  See the NOTICE file distributed with
 *  this work for additional information regarding copyright ownership.
 *  The ASF licenses this file to You 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.apache.coyote.http11.upgrade;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.Writer;
import java.net.Socket;

import javax.net.SocketFactory;
import javax.servlet.ReadListener;
import javax.servlet.ServletException;
import javax.servlet.ServletInputStream;
import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.ProtocolHandler;
import javax.servlet.http.WebConnection;

import junit.framework.Assert;

import org.junit.Test;

import static org.apache.catalina.startup.SimpleHttpClient.CRLF;

import org.apache.catalina.Context;
import org.apache.catalina.startup.Tomcat;
import org.apache.catalina.startup.TomcatBaseTest;

public class TestUpgradeServletInputStream extends TomcatBaseTest {

    private static final String MESSAGE = "This is a test.\n";

    @Test
    public void testSimpleUpgrade() throws Exception {
        doUpgrade();
    }

    @Test
    public void testSingleMessage() throws Exception {
        UpgradeConnection conn = doUpgrade();
        Writer writer = conn.getWriter();
        BufferedReader reader = conn.getReader();

        writer.write(MESSAGE);
        writer.flush();

        Thread.sleep(500);

        writer.write(MESSAGE);
        writer.flush();

        String response = reader.readLine();

        Assert.assertEquals(MESSAGE, response);
    }


    private UpgradeConnection doUpgrade() throws Exception {
        // Setup Tomcat instance
        Tomcat tomcat = getTomcatInstance();

        // Must have a real docBase - just use temp
        Context ctx =
                tomcat.addContext("", System.getProperty("java.io.tmpdir"));

        UpgradeServlet servlet = new UpgradeServlet();
        Tomcat.addServlet(ctx, "servlet", servlet);
        ctx.addServletMapping("/", "servlet");

        tomcat.start();

        // Use raw socket so the necessary control is available after the HTTP
        // upgrade
        Socket socket =
                SocketFactory.getDefault().createSocket("localhost", getPort());

        socket.setSoTimeout(300000);

        InputStream is = socket.getInputStream();
        OutputStream os = socket.getOutputStream();

        BufferedReader reader = new BufferedReader(new InputStreamReader(is));
        Writer writer = new OutputStreamWriter(os);

        writer.write("GET / HTTP/1.1" + CRLF);
        writer.write("Host: whatever" + CRLF);
        writer.write(CRLF);
        writer.flush();

        String status = reader.readLine();

        Assert.assertEquals("HTTP/1.1 101 Switching Protocols",
                status.substring(0, 32));

        // Skip the remaining response headers
        while (reader.readLine().length() > 0) {
            // Skip
        }

        return new UpgradeConnection(writer, reader);
    }

    private static class UpgradeServlet extends HttpServlet {

        private static final long serialVersionUID = 1L;

        @Override
        protected void doGet(HttpServletRequest req, HttpServletResponse resp)
                throws ServletException, IOException {

            req.upgrade(new EchoNonBlocking());
        }
    }

    private static class UpgradeConnection {
        private final Writer writer;
        private final BufferedReader reader;

        public UpgradeConnection(Writer writer, BufferedReader reader) {
            this.writer = writer;
            this.reader = reader;
        }

        public Writer getWriter() {
            return writer;
        }

        public BufferedReader getReader() {
            return reader;
        }
    }

    private static class EchoNonBlocking implements ProtocolHandler {

        private ServletInputStream sis;
        private ServletOutputStream sos;

        @Override
        public void init(WebConnection connection) {

            try {
                sis = connection.getInputStream();
                sos = connection.getOutputStream();
            } catch (IOException ioe) {
                throw new IllegalStateException(ioe);
            }

            sis.setReadListener(new EchoReadListener());
            sos.setWriteListener(new EchoWriteListener());
        }

        private class EchoReadListener implements ReadListener {

            private byte[] buffer = new byte[8096];

            @Override
            public void onDataAvailable() {
                try {
                    while (sis.isReady()) {
                        int read = sis.read(buffer);
                        if (read > 0) {
                            System.out.print(new String(buffer, 0, read));
                            if (sos.canWrite()) {
                                sos.write(buffer, 0, read);
                            } else {
                                throw new IOException("Unable to echo data. " +
                                        "canWrite() returned false");
                            }
                        }
                    }
                } catch (IOException ioe) {
                    throw new RuntimeException(ioe);
                }
            }

            @Override
            public void onAllDataRead() {
                System.out.println("All data read");
            }

            @Override
            public void onError(Throwable throwable) {
                // TODO Auto-generated method stub

            }
        }

        private class EchoWriteListener implements WriteListener {

            @Override
            public void onWritePossible() {
                // TODO Auto-generated method stub

            }

            @Override
            public void onError(Throwable throwable) {
                // TODO Auto-generated method stub

            }
        }
    }
}

@@ -1051,32 +1051,61 @@
 tStreamWriter;%0D%0A
+import java.io.PrintWriter;%0D%0A
 import java.io.W
@@ -1989,10 +1989,8 @@
 est.
-%5Cn
 %22;%0D%0A
@@ -2211,32 +2211,49 @@
         
+Print
 Writer 
-w
+pw = new PrintW
 riter
- = 
+(
 conn.get
@@ -2260,16 +2260,17 @@
 Writer()
+)
 ;%0D%0A     
@@ -2321,36 +2321,34 @@
 %0D%0A%0D%0A        
-writer.write
+pw.println
 (MESSAGE);%0D%0A
@@ -2347,38 +2347,34 @@
 SAGE);%0D%0A        
+p
 w
-riter
 .flush();%0D%0A%0D%0A   
@@ -2408,28 +2408,26 @@
         
-writer.write
+pw.println
 (MESSAGE
@@ -2430,38 +2430,34 @@
 SAGE);%0D%0A        
+p
 w
-riter
 .flush();%0D%0A%0D%0A   
@@ -2493,32 +2493,166 @@
 .readLine();%0D%0A%0D%0A
+        // Note: BufferedReader.readLine() strips new lines%0D%0A        //       ServletInputStream.readLine() does not strip new lines%0D%0A
         Assert.a
@@ -2680,16 +2680,17 @@
 response
+ 
 );%0D%0A    
@@ -5946,84 +5946,8 @@
  %7B%0D%0A
-                            System.out.print(new String(buffer, 0, read));%0D%0A
     

Fix test