1 /*
2  *  Licensed to the Apache Software Foundation (ASF) under one or more
3  *  contributor license agreements.  See the NOTICE file distributed with
4  *  this work for additional information regarding copyright ownership.
5  *  The ASF licenses this file to You under the Apache License, Version 2.0
6  *  (the "License"); you may not use this file except in compliance with
7  *  the License.  You may obtain a copy of the License at
8  *
9  *      http://www.apache.org/licenses/LICENSE-2.0
10  *
11  *  Unless required by applicable law or agreed to in writing, software
12  *  distributed under the License is distributed on an "AS IS" BASIS,
13  *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14  *  See the License for the specific language governing permissions and
15  *  limitations under the License.
16  */
17 package org.apache.tomcat.util.net;
18
19 import java.nio.BufferOverflowException;
20 import java.nio.ByteBuffer;
21
22 import org.apache.tomcat.util.buf.ByteBufferUtils;
23
24 public class SocketBufferHandler {
25
26     static SocketBufferHandler EMPTY = new SocketBufferHandler(0, 0, false) {
27         @Override
28         public void expand(int newSize) {
29         }
30     };
31
32     private volatile boolean readBufferConfiguredForWrite = true;
33     private volatile ByteBuffer readBuffer;
34
35     private volatile boolean writeBufferConfiguredForWrite = true;
36     private volatile ByteBuffer writeBuffer;
37
38     private final boolean direct;
39
40     public SocketBufferHandler(int readBufferSize, int writeBufferSize,
41             boolean direct) {
42         this.direct = direct;
43         if (direct) {
44             readBuffer = ByteBuffer.allocateDirect(readBufferSize);
45             writeBuffer = ByteBuffer.allocateDirect(writeBufferSize);
46         } else {
47             readBuffer = ByteBuffer.allocate(readBufferSize);
48             writeBuffer = ByteBuffer.allocate(writeBufferSize);
49         }
50     }
51
52
53     public void configureReadBufferForWrite() {
54         setReadBufferConfiguredForWrite(true);
55     }
56
57
58     public void configureReadBufferForRead() {
59         setReadBufferConfiguredForWrite(false);
60     }
61
62
63     private void setReadBufferConfiguredForWrite(boolean readBufferConFiguredForWrite) {
64         // NO-OP if buffer is already in correct state
65         if (this.readBufferConfiguredForWrite != readBufferConFiguredForWrite) {
66             if (readBufferConFiguredForWrite) {
67                 // Switching to write
68                 int remaining = readBuffer.remaining();
69                 if (remaining == 0) {
70                     readBuffer.clear();
71                 } else {
72                     readBuffer.compact();
73                 }
74             } else {
75                 // Switching to read
76                 readBuffer.flip();
77             }
78             this.readBufferConfiguredForWrite = readBufferConFiguredForWrite;
79         }
80     }
81
82
83     public ByteBuffer getReadBuffer() {
84         return readBuffer;
85     }
86
87
88     public boolean isReadBufferEmpty() {
89         if (readBufferConfiguredForWrite) {
90             return readBuffer.position() == 0;
91         } else {
92             return readBuffer.remaining() == 0;
93         }
94     }
95
96
97     public void unReadReadBuffer(ByteBuffer returnedData) {
98         if (isReadBufferEmpty()) {
99             configureReadBufferForWrite();
100             readBuffer.put(returnedData);
101         } else {
102             int bytesReturned = returnedData.remaining();
103             if (readBufferConfiguredForWrite) {
104                 // Writes always start at position zero
105                 if ((readBuffer.position() + bytesReturned) > readBuffer.capacity()) {
106                     throw new BufferOverflowException();
107                 } else {
108                     // Move the bytes up to make space for the returned data
109                     for (int i = 0; i < readBuffer.position(); i++) {
110                         readBuffer.put(i + bytesReturned, readBuffer.get(i));
111                     }
112                     // Insert the bytes returned
113                     for (int i = 0; i < bytesReturned; i++) {
114                         readBuffer.put(i, returnedData.get());
115                     }
116                     // Update the position
117                     readBuffer.position(readBuffer.position() + bytesReturned);
118                 }
119             } else {
120                 // Reads will start at zero but may have progressed
121                 int shiftRequired = bytesReturned - readBuffer.position();
122                 if (shiftRequired > 0) {
123                     if ((readBuffer.capacity() - readBuffer.limit()) < shiftRequired) {
124                         throw new BufferOverflowException();
125                     }
126                     // Move the bytes up to make space for the returned data
127                     int oldLimit = readBuffer.limit();
128                     readBuffer.limit(oldLimit + shiftRequired);
129                     for (int i = readBuffer.position(); i < oldLimit; i++) {
130                         readBuffer.put(i + shiftRequired, readBuffer.get(i));
131                     }
132                 } else {
133                     shiftRequired = 0;
134                 }
135                 // Insert the returned bytes
136                 int insertOffset = readBuffer.position() + shiftRequired - bytesReturned;
137                 for (int i = insertOffset; i < bytesReturned + insertOffset; i++) {
138                     readBuffer.put(i, returnedData.get());
139                 }
140                 readBuffer.position(insertOffset);
141             }
142         }
143     }
144
145
146     public void configureWriteBufferForWrite() {
147         setWriteBufferConfiguredForWrite(true);
148     }
149
150
151     public void configureWriteBufferForRead() {
152         setWriteBufferConfiguredForWrite(false);
153     }
154
155
156     private void setWriteBufferConfiguredForWrite(boolean writeBufferConfiguredForWrite) {
157         // NO-OP if buffer is already in correct state
158         if (this.writeBufferConfiguredForWrite != writeBufferConfiguredForWrite) {
159             if (writeBufferConfiguredForWrite) {
160                 // Switching to write
161                 int remaining = writeBuffer.remaining();
162                 if (remaining == 0) {
163                     writeBuffer.clear();
164                 } else {
165                     writeBuffer.compact();
166                     writeBuffer.position(remaining);
167                     writeBuffer.limit(writeBuffer.capacity());
168                 }
169             } else {
170                 // Switching to read
171                 writeBuffer.flip();
172             }
173             this.writeBufferConfiguredForWrite = writeBufferConfiguredForWrite;
174         }
175     }
176
177
178     public boolean isWriteBufferWritable() {
179         if (writeBufferConfiguredForWrite) {
180             return writeBuffer.hasRemaining();
181         } else {
182             return writeBuffer.remaining() == 0;
183         }
184     }
185
186
187     public ByteBuffer getWriteBuffer() {
188         return writeBuffer;
189     }
190
191
192     public boolean isWriteBufferEmpty() {
193         if (writeBufferConfiguredForWrite) {
194             return writeBuffer.position() == 0;
195         } else {
196             return writeBuffer.remaining() == 0;
197         }
198     }
199
200
201     public void reset() {
202         readBuffer.clear();
203         readBufferConfiguredForWrite = true;
204         writeBuffer.clear();
205         writeBufferConfiguredForWrite = true;
206     }
207
208
209     public void expand(int newSize) {
210         configureReadBufferForWrite();
211         readBuffer = ByteBufferUtils.expand(readBuffer, newSize);
212         configureWriteBufferForWrite();
213         writeBuffer = ByteBufferUtils.expand(writeBuffer, newSize);
214     }
215
216     public void free() {
217         if (direct) {
218             ByteBufferUtils.cleanDirectBuffer(readBuffer);
219             ByteBufferUtils.cleanDirectBuffer(writeBuffer);
220         }
221     }
222
223 }
224