org.postgresql.core.PGStream Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of mogdb-jdbc Show documentation
Show all versions of mogdb-jdbc Show documentation
Java JDBC driver for MogDB
/*
* Copyright (c) 2017, PostgreSQL Global Development Group
* See the LICENSE file in the project root for more information.
*/
package io.mogdb.core;
import io.mogdb.util.GT;
import io.mogdb.util.HostSpec;
import io.mogdb.util.PSQLException;
import io.mogdb.util.PSQLState;
import java.io.BufferedOutputStream;
import java.io.Closeable;
import java.io.EOFException;
import java.io.FilterOutputStream;
import java.io.Flushable;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.Writer;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.net.SocketTimeoutException;
import java.sql.SQLException;
import javax.net.SocketFactory;
/**
* Wrapper around the raw connection to the server that implements some basic primitives
* (reading/writing formatted data, doing string encoding, etc).
*
* In general, instances of PGStream are not threadsafe; the caller must ensure that only one thread
* at a time is accessing a particular PGStream instance.
*/
public class PGStream implements Closeable, Flushable {
private final SocketFactory socketFactory;
private final HostSpec hostSpec;
private final byte[] _int4buf;
private final byte[] _int2buf;
private Socket connection;
private VisibleBufferedInputStream pg_input;
private OutputStream pg_output;
private byte[] streamBuffer;
private Encoding encoding;
private Writer encodingWriter;
/**
* Constructor: Connect to the PostgreSQL back end and return a stream connection.
*
* @param socketFactory socket factory to use when creating sockets
* @param hostSpec the host and port to connect to
* @param timeout timeout in milliseconds, or 0 if no timeout set
* @throws IOException if an IOException occurs below it.
*/
public PGStream(SocketFactory socketFactory, HostSpec hostSpec, int timeout) throws IOException {
this.socketFactory = socketFactory;
this.hostSpec = hostSpec;
Socket socket = socketFactory.createSocket();
if (!socket.isConnected()) {
// When using a SOCKS proxy, the host might not be resolvable locally,
// thus we defer resolution until the traffic reaches the proxy. If there
// is no proxy, we must resolve the host to an IP to connect the socket.
InetSocketAddress address = hostSpec.shouldResolve()
? new InetSocketAddress(hostSpec.getHost(), hostSpec.getPort())
: InetSocketAddress.createUnresolved(hostSpec.getHost(), hostSpec.getPort());
socket.connect(address, timeout);
}
changeSocket(socket);
setEncoding(Encoding.getJVMEncoding("UTF-8"));
_int2buf = new byte[2];
_int4buf = new byte[4];
}
/**
* Constructor: Connect to the PostgreSQL back end and return a stream connection.
*
* @param socketFactory socket factory
* @param hostSpec the host and port to connect to
* @throws IOException if an IOException occurs below it.
* @deprecated use {@link #PGStream(SocketFactory, io.mogdb.util.HostSpec, int)}
*/
@Deprecated
public PGStream(SocketFactory socketFactory, HostSpec hostSpec) throws IOException {
this(socketFactory, hostSpec, 0);
}
public HostSpec getHostSpec() {
return hostSpec;
}
public Socket getSocket() {
return connection;
}
public SocketFactory getSocketFactory() {
return socketFactory;
}
/**
* Check for pending backend messages without blocking. Might return false when there actually are
* messages waiting, depending on the characteristics of the underlying socket. This is used to
* detect asynchronous notifies from the backend, when available.
*
* @return true if there is a pending backend message
* @throws IOException if something wrong happens
*/
public boolean hasMessagePending() throws IOException {
if (pg_input.available() > 0) {
return true;
}
// In certain cases, available returns 0, yet there are bytes
int soTimeout = getNetworkTimeout();
setNetworkTimeout(1);
try {
return pg_input.peek() != -1;
} catch (SocketTimeoutException e) {
return false;
} finally {
setNetworkTimeout(soTimeout);
}
}
/**
* Switch this stream to using a new socket. Any existing socket is not closed; it's
* assumed that we are changing to a new socket that delegates to the original socket (e.g. SSL).
*
* @param socket the new socket to change to
* @throws IOException if something goes wrong
*/
public void changeSocket(Socket socket) throws IOException {
this.connection = socket;
// Submitted by Jason Venner . Disable Nagle
// as we are selective about flushing output only when we
// really need to.
connection.setTcpNoDelay(true);
// Buffer sizes submitted by Sverre H Huseby
pg_input = new VisibleBufferedInputStream(connection.getInputStream(), 8192);
pg_output = new BufferedOutputStream(connection.getOutputStream(), 8192);
if (encoding != null) {
setEncoding(encoding);
}
}
public Encoding getEncoding() {
return encoding;
}
/**
* Change the encoding used by this connection.
*
* @param encoding the new encoding to use
* @throws IOException if something goes wrong
*/
public void setEncoding(Encoding encoding) throws IOException {
if (this.encoding != null && this.encoding.name().equals(encoding.name())) {
return;
}
// Close down any old writer.
if (encodingWriter != null) {
encodingWriter.close();
}
this.encoding = encoding;
// Intercept flush() downcalls from the writer; our caller
// will call PGStream.flush() as needed.
OutputStream interceptor = new FilterOutputStream(pg_output) {
public void flush() throws IOException {
}
public void close() throws IOException {
super.flush();
}
};
encodingWriter = encoding.getEncodingWriter(interceptor);
}
/**
* Get a Writer instance that encodes directly onto the underlying stream.
*
* The returned Writer should not be closed, as it's a shared object. Writer.flush needs to be
* called when switching between use of the Writer and use of the PGStream write methods, but it
* won't actually flush output all the way out -- call {@link #flush} to actually ensure all
* output has been pushed to the server.
*
* @return the shared Writer instance
* @throws IOException if something goes wrong.
*/
public Writer getEncodingWriter() throws IOException {
if (encodingWriter == null) {
throw new IOException("No encoding has been set on this connection");
}
return encodingWriter;
}
/**
* Sends a single character to the back end.
*
* @param val the character to be sent
* @throws IOException if an I/O error occurs
*/
public void sendChar(int val) throws IOException {
pg_output.write(val);
}
/**
* Sends a 4-byte integer to the back end.
*
* @param val the integer to be sent
* @throws IOException if an I/O error occurs
*/
public void sendInteger4(int val) throws IOException {
_int4buf[0] = (byte) (val >>> 24);
_int4buf[1] = (byte) (val >>> 16);
_int4buf[2] = (byte) (val >>> 8);
_int4buf[3] = (byte) (val);
pg_output.write(_int4buf);
}
/**
* Sends a 2-byte integer (short) to the back end.
*
* @param val the integer to be sent
* @throws IOException if an I/O error occurs or {@code val} cannot be encoded in 2 bytes
*/
public void sendInteger2(int val) throws IOException {
if (val < Short.MIN_VALUE || val > Short.MAX_VALUE) {
throw new IOException("Tried to send an out-of-range integer as a 2-byte value: " + val);
}
_int2buf[0] = (byte) (val >>> 8);
_int2buf[1] = (byte) val;
pg_output.write(_int2buf);
}
/**
* Send an array of bytes to the backend.
*
* @param buf The array of bytes to be sent
* @throws IOException if an I/O error occurs
*/
public void send(byte[] buf) throws IOException {
pg_output.write(buf);
}
/**
* Send a fixed-size array of bytes to the backend. If {@code buf.length < siz}, pad with zeros.
* If {@code buf.lengh > siz}, truncate the array.
*
* @param buf the array of bytes to be sent
* @param siz the number of bytes to be sent
* @throws IOException if an I/O error occurs
*/
public void send(byte[] buf, int siz) throws IOException {
send(buf, 0, siz);
}
/**
* Send a fixed-size array of bytes to the backend. If {@code length < siz}, pad with zeros. If
* {@code length > siz}, truncate the array.
*
* @param buf the array of bytes to be sent
* @param off offset in the array to start sending from
* @param siz the number of bytes to be sent
* @throws IOException if an I/O error occurs
*/
public void send(byte[] buf, int off, int siz) throws IOException {
int bufamt = buf.length - off;
pg_output.write(buf, off, bufamt < siz ? bufamt : siz);
for (int i = bufamt; i < siz; ++i) {
pg_output.write(0);
}
}
/**
* Receives a single character from the backend, without advancing the current protocol stream
* position.
*
* @return the character received
* @throws IOException if an I/O Error occurs
*/
public int peekChar() throws IOException {
int c = pg_input.peek();
if (c < 0) {
throw new EOFException("EOF Exception");
}
return c;
}
/**
* Receives a single character from the backend.
*
* @return the character received
* @throws IOException if an I/O Error occurs
*/
public int receiveChar() throws IOException {
int c = pg_input.read();
if (c < 0) {
throw new EOFException("EOF Exception");
}
return c;
}
/**
* Receives a four byte integer from the backend.
*
* @return the integer received from the backend
* @throws IOException if an I/O error occurs
*/
public int receiveInteger4() throws IOException {
if (pg_input.read(_int4buf) != 4) {
throw new EOFException("EOF Exception");
}
return (_int4buf[0] & 0xFF) << 24 | (_int4buf[1] & 0xFF) << 16 | (_int4buf[2] & 0xFF) << 8
| _int4buf[3] & 0xFF;
}
/**
* Receives a two byte integer from the backend.
*
* @return the integer received from the backend
* @throws IOException if an I/O error occurs
*/
public int receiveInteger2() throws IOException {
if (pg_input.read(_int2buf) != 2) {
throw new EOFException("EOF Exception");
}
return (_int2buf[0] & 0xFF) << 8 | _int2buf[1] & 0xFF;
}
/**
* Receives a fixed-size string from the backend.
*
* @param len the length of the string to receive, in bytes.
* @return the decoded string
* @throws IOException if something wrong happens
*/
public String receiveString(int len) throws IOException {
if (!pg_input.ensureBytes(len)) {
throw new EOFException("EOF Exception");
}
String res = encoding.decode(pg_input.getBuffer(), pg_input.getIndex(), len);
pg_input.skip(len);
return res;
}
/**
* Receives a fixed-size string from the backend, and tries to avoid "UTF-8 decode failed"
* errors.
*
* @param len the length of the string to receive, in bytes.
* @return the decoded string
* @throws IOException if something wrong happens
*/
public EncodingPredictor.DecodeResult receiveErrorString(int len) throws IOException {
if (!pg_input.ensureBytes(len)) {
throw new EOFException("EOF Exception");
}
EncodingPredictor.DecodeResult res;
try {
String value = encoding.decode(pg_input.getBuffer(), pg_input.getIndex(), len);
// no autodetect warning as the message was converted on its own
res = new EncodingPredictor.DecodeResult(value, null);
} catch (IOException e) {
res = EncodingPredictor.decode(pg_input.getBuffer(), pg_input.getIndex(), len);
if (res == null) {
Encoding enc = Encoding.defaultEncoding();
String value = enc.decode(pg_input.getBuffer(), pg_input.getIndex(), len);
res = new EncodingPredictor.DecodeResult(value, enc.name());
}
}
pg_input.skip(len);
return res;
}
/**
* Receives a null-terminated string from the backend. If we don't see a null, then we assume
* something has gone wrong.
*
* @return string from back end
* @throws IOException if an I/O error occurs, or end of file
*/
public String receiveString() throws IOException {
int len = pg_input.scanCStringLength();
String res = encoding.decode(pg_input.getBuffer(), pg_input.getIndex(), len - 1);
pg_input.skip(len);
return res;
}
/**
* Read a tuple from the back end. A tuple is a two dimensional array of bytes. This variant reads
* the V3 protocol's tuple representation.
*
* @return tuple from the back end
* @throws IOException if a data I/O error occurs
*/
public byte[][] receiveTupleV3() throws IOException, OutOfMemoryError {
// TODO: use l_msgSize
int l_msgSize = receiveInteger4();
int l_nf = receiveInteger2();
byte[][] answer = new byte[l_nf][];
OutOfMemoryError oom = null;
for (int i = 0; i < l_nf; ++i) {
int l_size = receiveInteger4();
if (l_size != -1) {
try {
answer[i] = new byte[l_size];
receive(answer[i], 0, l_size);
} catch (OutOfMemoryError oome) {
oom = oome;
skip(l_size);
}
}
}
if (oom != null) {
throw oom;
}
return answer;
}
/**
* Reads in a given number of bytes from the backend.
*
* @param siz number of bytes to read
* @return array of bytes received
* @throws IOException if a data I/O error occurs
*/
public byte[] receive(int siz) throws IOException {
byte[] answer = new byte[siz];
receive(answer, 0, siz);
return answer;
}
/**
* Reads in a given number of bytes from the backend.
*
* @param buf buffer to store result
* @param off offset in buffer
* @param siz number of bytes to read
* @throws IOException if a data I/O error occurs
*/
public void receive(byte[] buf, int off, int siz) throws IOException {
int s = 0;
while (s < siz) {
int w = pg_input.read(buf, off + s, siz - s);
if (w < 0) {
throw new EOFException("EOF Exception");
}
s += w;
}
}
public void skip(int size) throws IOException {
long s = 0;
while (s < size) {
s += pg_input.skip(size - s);
}
}
/**
* Copy data from an input stream to the connection.
*
* @param inStream the stream to read data from
* @param remaining the number of bytes to copy
* @throws IOException if a data I/O error occurs
*/
public void sendStream(InputStream inStream, int remaining) throws IOException {
int expectedLength = remaining;
if (streamBuffer == null) {
streamBuffer = new byte[8192];
}
while (remaining > 0) {
int count = (remaining > streamBuffer.length ? streamBuffer.length : remaining);
int readCount;
try {
readCount = inStream.read(streamBuffer, 0, count);
if (readCount < 0) {
throw new EOFException(
GT.tr("Premature end of input stream, expected {0} bytes, but only read {1}.",
expectedLength, expectedLength - remaining));
}
} catch (IOException ioe) {
while (remaining > 0) {
send(streamBuffer, count);
remaining -= count;
count = (remaining > streamBuffer.length ? streamBuffer.length : remaining);
}
throw new PGBindException(ioe);
}
send(streamBuffer, readCount);
remaining -= readCount;
}
}
/**
* Flush any pending output to the backend.
*
* @throws IOException if an I/O error occurs
*/
@Override
public void flush() throws IOException {
if (encodingWriter != null) {
encodingWriter.flush();
}
pg_output.flush();
}
/**
* Consume an expected EOF from the backend.
*
* @throws IOException if an I/O error occurs
* @throws SQLException if we get something other than an EOF
*/
public void receiveEOF() throws SQLException, IOException {
int c = pg_input.read();
if (c < 0) {
return;
}
throw new PSQLException(GT.tr("Expected an EOF from server, got: {0}", c),
PSQLState.COMMUNICATION_ERROR);
}
/**
* Closes the connection.
*
* @throws IOException if an I/O Error occurs
*/
@Override
public void close() throws IOException {
if (encodingWriter != null) {
encodingWriter.close();
}
pg_output.close();
pg_input.close();
connection.close();
}
public void setNetworkTimeout(int milliseconds) throws IOException {
connection.setSoTimeout(milliseconds);
}
public int getNetworkTimeout() throws IOException {
return connection.getSoTimeout();
}
/**
* This function returns the String which contains the server socket address and client socket address.
* This function is only used to distinguish different connections when we create dump file.
*/
public String getConnectInfo() {
String clientSocketAddress = connection.getLocalSocketAddress().toString().substring(1);
String serverSocketAddress = connection.getRemoteSocketAddress().toString().substring(1);
return clientSocketAddress + "/" + serverSocketAddress;
}
/**
* This function is only for creating dump file when I/O error happens.
*
*/
public String getInputBufferByHex() {
try {
int bufferSize=pg_input.available();
if (bufferSize <= 0)
return "buffer is empty";
byte []bytes=new byte[bufferSize];
pg_input.read(bytes);
return Utils.toHexString(bytes);
} catch (Exception e) {
return "exception happens";
}
}
/**
* This function is used to return socket status when I/O error happens.
*
*/
public String getSocketStatus() {
StringBuilder status = new StringBuilder();
if (connection.isClosed()) {
status.append("socket is closed; ");
} else {
status.append("socket is not closed; ");
}
try {
this.getSocket().sendUrgentData(1);
status.append("Urgent packet sent to backend successfully; ");
} catch(IOException e) {
status.append("Sending Urgent packet failed, detail: " + e.getMessage() + ". ");
}
return status.toString();
}
}