-
Notifications
You must be signed in to change notification settings - Fork 819
/
StreamWrapper.java
196 lines (172 loc) · 5.76 KB
/
StreamWrapper.java
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
/*
* Copyright (c) 2004, PostgreSQL Global Development Group
* See the LICENSE file in the project root for more information.
*/
// Copyright (c) 2004, Open Cloud Limited.
package org.postgresql.util;
import static org.postgresql.util.internal.Nullness.castNonNull;
import org.checkerframework.checker.nullness.qual.Nullable;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.file.Files;
/**
* Wrapper around a length-limited InputStream.
*
* @author Oliver Jowett (oliver@opencloud.com)
*/
public class StreamWrapper {
private static final int MAX_MEMORY_BUFFER_BYTES = 51200;
private static final String TEMP_FILE_PREFIX = "postgres-pgjdbc-stream";
public StreamWrapper(byte[] data, int offset, int length) {
this.stream = null;
this.rawData = data;
this.offset = offset;
this.length = length;
}
public StreamWrapper(InputStream stream, int length) {
this.stream = stream;
this.rawData = null;
this.offset = 0;
this.length = length;
}
public StreamWrapper(InputStream stream) throws PSQLException {
try {
ByteArrayOutputStream memoryOutputStream = new ByteArrayOutputStream();
final int memoryLength = copyStream(stream, memoryOutputStream, MAX_MEMORY_BUFFER_BYTES);
byte[] rawData = memoryOutputStream.toByteArray();
if (memoryLength == -1) {
final int diskLength;
final File tempFile = Files.createTempFile(TEMP_FILE_PREFIX, null).toFile();
FileOutputStream diskOutputStream = new FileOutputStream(tempFile);
diskOutputStream.write(rawData);
try {
diskLength = copyStream(stream, diskOutputStream, Integer.MAX_VALUE - rawData.length);
if (diskLength == -1) {
throw new PSQLException(GT.tr("Object is too large to send over the protocol."),
PSQLState.NUMERIC_CONSTANT_OUT_OF_RANGE);
}
diskOutputStream.flush();
} finally {
diskOutputStream.close();
}
this.offset = 0;
this.length = rawData.length + diskLength;
this.rawData = null;
this.stream = new FileInputStream(tempFile) {
/*
* Usually, closing stream should be done by pgjdbc clients. Here it's an internally
* managed stream so we need to auto-close it and be sure to delete the temporary file
* when doing so. Auto-closing will be done when the first occurs: reaching EOF or Garbage
* Collection
*/
private boolean closed = false;
private int position = 0;
/**
* Check if we should auto-close this stream
*/
private void checkShouldClose(int readResult) throws IOException {
if (readResult == -1) {
close();
} else {
position += readResult;
if (position >= length) {
close();
}
}
}
public int read(byte[] b) throws IOException {
if (closed) {
return -1;
}
int result = super.read(b);
checkShouldClose(result);
return result;
}
public int read(byte[] b, int off, int len) throws IOException {
if (closed) {
return -1;
}
int result = super.read(b, off, len);
checkShouldClose(result);
return result;
}
public void close() throws IOException {
if (!closed) {
super.close();
tempFile.delete();
closed = true;
}
}
@SuppressWarnings({"deprecation", "removal"})
protected void finalize() throws IOException {
// forcibly close it because super.finalize() may keep the FD open, which may prevent
// file deletion
close();
// javac 13 assumes it can throw Throwable
try {
super.finalize();
} catch (RuntimeException e) {
throw e;
} catch (Error e) {
throw e;
} catch (IOException e) {
throw e;
} catch (Throwable e) {
throw new RuntimeException("Unexpected exception from finalize", e);
}
}
};
} else {
this.rawData = rawData;
this.stream = null;
this.offset = 0;
this.length = rawData.length;
}
} catch (IOException e) {
throw new PSQLException(GT.tr("An I/O error occurred while sending to the backend."),
PSQLState.IO_ERROR, e);
}
}
public InputStream getStream() {
if (stream != null) {
return stream;
}
return new java.io.ByteArrayInputStream(castNonNull(rawData), offset, length);
}
public int getLength() {
return length;
}
public int getOffset() {
return offset;
}
public byte @Nullable [] getBytes() {
return rawData;
}
public String toString() {
return "<stream of " + length + " bytes>";
}
private static int copyStream(InputStream inputStream, OutputStream outputStream, int limit)
throws IOException {
int totalLength = 0;
byte[] buffer = new byte[2048];
int readLength = inputStream.read(buffer);
while (readLength > 0) {
totalLength += readLength;
outputStream.write(buffer, 0, readLength);
if (totalLength >= limit) {
return -1;
}
readLength = inputStream.read(buffer);
}
return totalLength;
}
private final @Nullable InputStream stream;
private final byte @Nullable [] rawData;
private final int offset;
private final int length;
}