/*
 * JBoss, Home of Professional Open Source.
 * Copyright 2014 Red Hat, Inc., and individual contributors
 * as indicated by the @author tags.
 *
 * 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 io.undertow.util;

import io.undertow.UndertowLogger;
import org.xnio.ChannelExceptionHandler;
import org.xnio.ChannelListener;
import org.xnio.ChannelListeners;
import org.xnio.IoUtils;
import org.xnio.StreamConnection;
import org.xnio.conduits.ConduitStreamSinkChannel;
import org.xnio.conduits.ConduitStreamSourceChannel;

import java.io.Closeable;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.concurrent.TimeUnit;

/**
 * @author Stuart Douglas
 */
public class ConnectionUtils {

    private static final long MAX_DRAIN_TIME = Long.getLong("io.undertow.max-drain-time", 10000);

    private ConnectionUtils() {

    }

    /**
     * Cleanly close a connection, by shutting down and flushing writes and then draining reads.
     * <p/>
     * If this fails the connection is forcibly closed.
     *
     * @param connection The connection
     * @param additional Any additional resources to close once the connection has been closed
     */
    public static void cleanClose(StreamConnection connection, Closeable... additional) {
        try {
            connection.getSinkChannel().shutdownWrites();
            if (!connection.getSinkChannel().flush()) {
                connection.getSinkChannel().setWriteListener(ChannelListeners.flushingChannelListener(new ChannelListener<ConduitStreamSinkChannel>() {
                    @Override
                    public void handleEvent(ConduitStreamSinkChannel channel) {
                        doDrain(connection, additional);
                    }
                }, new ChannelExceptionHandler<ConduitStreamSinkChannel>() {
                    @Override
                    public void handleException(ConduitStreamSinkChannel channel, IOException exception) {
                        UndertowLogger.REQUEST_IO_LOGGER.ioException(exception);
                        IoUtils.safeClose(connection);
                        IoUtils.safeClose(additional);
                    }
                }));
                connection.getSinkChannel().resumeWrites();
            } else {
                doDrain(connection, additional);
            }

        } catch (Exception e) {
            if (e instanceof IOException) {
                UndertowLogger.REQUEST_IO_LOGGER.ioException((IOException) e);
            } else {
                UndertowLogger.REQUEST_IO_LOGGER.ioException(new IOException(e));
            }
            IoUtils.safeClose(connection);
            IoUtils.safeClose(additional);
        }
    }

    private static void doDrain(final StreamConnection connection, final Closeable... additional) {
        if (!connection.getSourceChannel().isOpen()) {
            IoUtils.safeClose(connection);
            IoUtils.safeClose(additional);
            return;
        }
        final ByteBuffer b = ByteBuffer.allocate(1);
        try {
            int res = connection.getSourceChannel().read(b);
            if (res == 0) {
                connection.getSourceChannel().setReadListener(new ChannelListener<ConduitStreamSourceChannel>() {
                    @Override
                    public void handleEvent(ConduitStreamSourceChannel channel) {
                        try {
                            int res = channel.read(b);
                            if (res == 0) {
                                return;
                            } else {
                                IoUtils.safeClose(connection);
                                IoUtils.safeClose(additional);
                            }
                        } catch (Exception e) {
                            if (e instanceof IOException) {
                                UndertowLogger.REQUEST_IO_LOGGER.ioException((IOException) e);
                            } else {
                                UndertowLogger.REQUEST_IO_LOGGER.ioException(new IOException(e));
                            }
                            IoUtils.safeClose(connection);
                            IoUtils.safeClose(additional);
                        }
                    }
                });
                connection.getSourceChannel().resumeReads();
                connection.getIoThread().executeAfter(new Runnable() {
                    @Override
                    public void run() {
                        IoUtils.safeClose(connection);
                        IoUtils.safeClose(additional);
                    }
                }, MAX_DRAIN_TIME, TimeUnit.MILLISECONDS);
            } else {
                IoUtils.safeClose(connection);
                IoUtils.safeClose(additional);
            }
        } catch (Exception e) {
            if (e instanceof IOException) {
                UndertowLogger.REQUEST_IO_LOGGER.ioException((IOException) e);
            } else {
                UndertowLogger.REQUEST_IO_LOGGER.ioException(new IOException(e));
            }
            IoUtils.safeClose(connection);
            IoUtils.safeClose(additional);
        }
    }


}

@@ -931,24 +931,54 @@
 Connection;%0A
+import org.xnio.XnioExecutor;%0A
 import org.x
@@ -3653,39 +3653,423 @@
             
-if (res == 0) %7B
+b.clear();%0A            if (res == 0) %7B%0A                final XnioExecutor.Key key = connection.getIoThread().executeAfter(new Runnable() %7B%0A                    @Override%0A                    public void run() %7B%0A                        IoUtils.safeClose(connection);%0A                        IoUtils.safeClose(additional);%0A                    %7D%0A                %7D, MAX_DRAIN_TIME, TimeUnit.MILLISECONDS);
 %0A           
@@ -4408,93 +4408,16 @@
 res 
-=
+!
 = 0) %7B%0A
-                                return;%0A                            %7D else %7B%0A
     
@@ -4530,32 +4530,78 @@
 se(additional);%0A
+                                key.remove();%0A
                 
@@ -5092,32 +5092,74 @@
 se(additional);%0A
+                            key.remove();%0A
                 
@@ -5210,16 +5210,16 @@
     %7D);%0A
+
         
@@ -5275,340 +5275,8 @@
 ();%0A
-                connection.getIoThread().executeAfter(new Runnable() %7B%0A                    @Override%0A                    public void run() %7B%0A                        IoUtils.safeClose(connection);%0A                        IoUtils.safeClose(additional);%0A                    %7D%0A                %7D, MAX_DRAIN_TIME, TimeUnit.MILLISECONDS);%0A
     

Clear the close timeout key