1 /*
2  * JBoss, Home of Professional Open Source.
3  * Copyright 2014 Red Hat, Inc., and individual contributors
4  * as indicated by the @author tags.
5  *
6  * Licensed under the Apache License, Version 2.0 (the "License");
7  * you may not use this file except in compliance with the License.
8  * You may obtain a copy of the License at
9  *
10  *     http://www.apache.org/licenses/LICENSE-2.0
11  *
12  *  Unless required by applicable law or agreed to in writing, software
13  *  distributed under the License is distributed on an "AS IS" BASIS,
14  *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15  *  See the License for the specific language governing permissions and
16  *  limitations under the License.
17  */
18
19 package io.undertow.util;
20
21 import io.undertow.UndertowLogger;
22 import org.xnio.ChannelExceptionHandler;
23 import org.xnio.ChannelListener;
24 import org.xnio.ChannelListeners;
25 import org.xnio.IoUtils;
26 import org.xnio.StreamConnection;
27 import org.xnio.XnioExecutor;
28 import org.xnio.conduits.ConduitStreamSinkChannel;
29 import org.xnio.conduits.ConduitStreamSourceChannel;
30
31 import java.io.Closeable;
32 import java.io.IOException;
33 import java.nio.ByteBuffer;
34 import java.util.concurrent.TimeUnit;
35
36 /**
37  * @author Stuart Douglas
38  */
39 public class ConnectionUtils {
40
41     private static final long MAX_DRAIN_TIME = Long.getLong("io.undertow.max-drain-time", 10000);
42
43     private ConnectionUtils() {
44
45     }
46
47     /**
48      * Cleanly close a connection, by shutting down and flushing writes and then draining reads.
49      * <p>
50      * If this fails the connection is forcibly closed.
51      *
52      * @param connection The connection
53      * @param additional Any additional resources to close once the connection has been closed
54      */
55     public static void cleanClose(StreamConnection connection, Closeable... additional) {
56         try {
57             connection.getSinkChannel().shutdownWrites();
58             if (!connection.getSinkChannel().flush()) {
59                 connection.getSinkChannel().setWriteListener(ChannelListeners.flushingChannelListener(new ChannelListener<ConduitStreamSinkChannel>() {
60                     @Override
61                     public void handleEvent(ConduitStreamSinkChannel channel) {
62                         doDrain(connection, additional);
63                     }
64                 }, new ChannelExceptionHandler<ConduitStreamSinkChannel>() {
65                     @Override
66                     public void handleException(ConduitStreamSinkChannel channel, IOException exception) {
67                         UndertowLogger.REQUEST_IO_LOGGER.ioException(exception);
68                         IoUtils.safeClose(connection);
69                         IoUtils.safeClose(additional);
70                     }
71                 }));
72                 connection.getSinkChannel().resumeWrites();
73             } else {
74                 doDrain(connection, additional);
75             }
76
77         } catch (Throwable e) {
78             if (e instanceof IOException) {
79                 UndertowLogger.REQUEST_IO_LOGGER.ioException((IOException) e);
80             } else {
81                 UndertowLogger.REQUEST_IO_LOGGER.ioException(new IOException(e));
82             }
83             IoUtils.safeClose(connection);
84             IoUtils.safeClose(additional);
85         }
86     }
87
88     private static void doDrain(final StreamConnection connection, final Closeable... additional) {
89         if (!connection.getSourceChannel().isOpen()) {
90             IoUtils.safeClose(connection);
91             IoUtils.safeClose(additional);
92             return;
93         }
94         final ByteBuffer b = ByteBuffer.allocate(1);
95         try {
96             int res = connection.getSourceChannel().read(b);
97             b.clear();
98             if (res == 0) {
99                 final XnioExecutor.Key key = WorkerUtils.executeAfter(connection.getIoThread(), new Runnable() {
100                     @Override
101                     public void run() {
102                         IoUtils.safeClose(connection);
103                         IoUtils.safeClose(additional);
104                     }
105                 }, MAX_DRAIN_TIME, TimeUnit.MILLISECONDS);
106                 connection.getSourceChannel().setReadListener(new ChannelListener<ConduitStreamSourceChannel>() {
107                     @Override
108                     public void handleEvent(ConduitStreamSourceChannel channel) {
109                         try {
110                             int res = channel.read(b);
111                             if (res != 0) {
112                                 IoUtils.safeClose(connection);
113                                 IoUtils.safeClose(additional);
114                                 key.remove();
115                             }
116                         } catch (Exception e) {
117                             if (e instanceof IOException) {
118                                 UndertowLogger.REQUEST_IO_LOGGER.ioException((IOException) e);
119                             } else {
120                                 UndertowLogger.REQUEST_IO_LOGGER.ioException(new IOException(e));
121                             }
122                             IoUtils.safeClose(connection);
123                             IoUtils.safeClose(additional);
124                             key.remove();
125                         }
126                     }
127                 });
128                 connection.getSourceChannel().resumeReads();
129             } else {
130                 IoUtils.safeClose(connection);
131                 IoUtils.safeClose(additional);
132             }
133         } catch (Throwable e) {
134             if (e instanceof IOException) {
135                 UndertowLogger.REQUEST_IO_LOGGER.ioException((IOException) e);
136             } else {
137                 UndertowLogger.REQUEST_IO_LOGGER.ioException(new IOException(e));
138             }
139             IoUtils.safeClose(connection);
140             IoUtils.safeClose(additional);
141         }
142     }
143
144
145 }
146