org.finos.legend.engine.postgres.PostgresWireProtocol Maven / Gradle / Ivy
/*
* Licensed to Crate.io GmbH ("Crate") under one or more contributor
* license agreements. See the NOTICE file distributed with this work for
* additional information regarding copyright ownership. Crate licenses
* this file to you 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.
*
* However, if you have executed another commercial license agreement
* with Crate these terms will supersede the license and you may use the
* software solely pursuant to the terms of the relevant commercial agreement.
*/
package org.finos.legend.engine.postgres;
import com.google.common.net.InetAddresses;
import com.sun.security.jgss.GSSUtil;
import io.netty.buffer.ByteBuf;
import io.netty.channel.Channel;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.ssl.SslContext;
import io.netty.handler.ssl.SslHandler;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.StatusCode;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.context.Scope;
import org.finos.legend.engine.postgres.auth.AuthenticationMethod;
import org.finos.legend.engine.postgres.auth.AuthenticationMethodType;
import org.finos.legend.engine.postgres.auth.AuthenticationProvider;
import org.finos.legend.engine.postgres.auth.KerberosIdentityProvider;
import org.finos.legend.engine.postgres.config.GSSConfig;
import org.finos.legend.engine.postgres.handler.PostgresResultSetMetaData;
import org.finos.legend.engine.postgres.types.PGType;
import org.finos.legend.engine.postgres.types.PGTypes;
import org.finos.legend.engine.postgres.utils.OpenTelemetryUtil;
import org.finos.legend.engine.shared.core.identity.Identity;
import org.finos.legend.engine.shared.core.kerberos.SubjectTools;
import org.finos.legend.engine.shared.core.operational.Assert;
import org.ietf.jgss.GSSContext;
import org.ietf.jgss.GSSCredential;
import org.ietf.jgss.GSSException;
import org.ietf.jgss.GSSManager;
import org.ietf.jgss.GSSName;
import org.ietf.jgss.Oid;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import javax.net.ssl.SSLSession;
import javax.security.auth.Subject;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.SocketException;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.PrivilegedActionException;
import java.security.PrivilegedExceptionAction;
import java.sql.ParameterMetaData;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Locale;
import java.util.Properties;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.function.BiConsumer;
import java.util.function.Supplier;
import static org.finos.legend.engine.postgres.FormatCodes.getFormatCode;
/**
* Netty Handler/FrameDecoder for the Postgres wire protocol.
This class handles the message
* flow and dispatching
*
*
*
* Client Server
*
* (optional ssl negotiation)
*
*
* | SSLRequest |
* |--------------------------------->|
* | |
* | 'S' | 'N' | error | (supported in Enterprise version)
* |<---------------------------------|
*
*
* startup:
* The authentication flow is handled by implementations of {@link AuthenticationMethod}.
*
* | |
* | StartupMessage |
* |--------------------------------->|
* | |
* | Authentication Method |
* | or |
* | AuthenticationOK |
* | or |
* | ErrorResponse |
* |<---------------------------------|
* | |
* | ParameterStatus |
* |<---------------------------------|
* | |
* | ReadyForQuery |
* |<---------------------------------|
*
*
* Simple Query:
*
* + +
* | Q (query) |
* |--------------------------------->|
* | |
* | RowDescription |
* |<---------------------------------|
* | |
* | DataRow |
* |<---------------------------------|
* | DataRow |
* |<---------------------------------|
* | CommandComplete |
* |<---------------------------------|
* | ReadyForQuery |
* |<---------------------------------|
*
* Extended Query
*
* + +
* | Parse |
* |--------------------------------->|
* | |
* | ParseComplete or ErrorResponse |
* |<---------------------------------|
* | |
* | Describe Statement (optional) |
* |--------------------------------->|
* | |
* | ParameterDescription (optional) |
* |<-------------------------------- |
* | |
* | RowDescription (optional) |
* |<-------------------------------- |
* | |
* | Bind |
* |--------------------------------->|
* | |
* | BindComplete or ErrorResponse |
* |<---------------------------------|
* | |
* | Describe Portal (optional) |
* |--------------------------------->|
* | |
* | RowDescription (optional) |
* |<-------------------------------- |
* | |
* | Execute |
* |--------------------------------->|
* | |
* | DataRow | |
* | CommandComplete | |
* | EmptyQueryResponse | |
* | ErrorResponse |
* |<---------------------------------|
* | |
* | Sync |
* |--------------------------------->|
* | |
* | ReadyForQuery |
* |<---------------------------------|
*
*
* Take a look at {@link Messages} to see how the messages are structured.
*
* See postgresql docs for a more detailed
* description of the message flow
*/
public class PostgresWireProtocol
{
private static final Logger LOGGER = LoggerFactory.getLogger(
PostgresWireProtocol.class);
public static int SERVER_VERSION_NUM = 100500;
public static String PG_SERVER_VERSION = "10.5";
final PgDecoder decoder;
final MessageHandler handler;
private final SessionsFactory sessions;
/* private final Function getAccessControl;*/
private final AuthenticationProvider authService;
private final GSSConfig gssConfig;
/* private final Consumer addTransportHandler;
*/
private DelayableWriteChannel channel;
Session session;
private boolean ignoreTillSync = false;
private AuthenticationContext authContext;
private Properties properties;
private Messages messages;
public PostgresWireProtocol(SessionsFactory sessions,
/*Function getAcessControl,*/
/*Consumer addTransportHandler,*/
AuthenticationProvider authService,
GSSConfig gssConfig, Supplier getSslContext,
Messages messages)
{
this.sessions = sessions;
//this.getAccessControl = getAcessControl;
//this.addTransportHandler = addTransportHandler;
this.authService = authService;
this.decoder = new PgDecoder(getSslContext);
this.handler = new MessageHandler();
this.gssConfig = gssConfig;
this.messages = messages;
}
static String readCString(ByteBuf buffer)
{
byte[] bytes = new byte[buffer.bytesBefore((byte) 0) + 1];
if (bytes.length == 0)
{
return null;
}
buffer.readBytes(bytes);
return new String(bytes, 0, bytes.length - 1, StandardCharsets.UTF_8);
}
private static char[] readCharArray(ByteBuf buffer)
{
byte[] bytes = new byte[buffer.bytesBefore((byte) 0) + 1];
if (bytes.length == 0)
{
return null;
}
buffer.readBytes(bytes);
return StandardCharsets.UTF_8.decode(ByteBuffer.wrap(bytes)).array();
}
private static byte[] readByteArray(ByteBuf buffer, int payloadLength)
{
if (payloadLength == 0)
{
return null;
}
byte[] bytes = new byte[payloadLength];
buffer.readBytes(bytes);
return bytes;
}
private Properties readStartupMessage(ByteBuf buffer)
{
Properties properties = new Properties();
while (true)
{
String key = readCString(buffer);
if (key == null)
{
break;
}
String value = readCString(buffer);
LOGGER.trace("payload: key={} value={}", key, value);
if (!"".equals(key) && !"".equals(value))
{
properties.setProperty(key, value);
}
}
return properties;
}
private static class ReadyForQueryCallback implements BiConsumer
{
private final Channel channel;
private final Messages messages;
private ReadyForQueryCallback(Channel channel, Messages messages)
{
this.channel = channel;
this.messages = messages;
}
@Override
public void accept(Object result, Throwable t)
{
if (t instanceof CompletionException)
{
Throwable actualCause = t.getCause();
PostgresServerException postgresServerException = PostgresServerException.wrapException(actualCause);
messages.sendErrorResponse(channel, postgresServerException);
}
boolean clientInterrupted = t instanceof ClientInterrupted
|| (t != null && t.getCause() instanceof ClientInterrupted);
if (!clientInterrupted)
{
messages.sendReadyForQuery(channel);
}
}
}
private class MessageHandler extends SimpleChannelInboundHandler
{
@Override
public void channelRegistered(ChannelHandlerContext ctx) throws Exception
{
channel = new DelayableWriteChannel(ctx.channel());
}
@Override
public boolean acceptInboundMessage(Object msg) throws Exception
{
return true;
}
@Override
public void channelRead0(ChannelHandlerContext ctx, ByteBuf buffer) throws Exception
{
Assert.assertTrue(channel != null, () -> "Channel must be initialized");
try
{
dispatchState(buffer, channel);
}
catch (Throwable t)
{
ignoreTillSync = true;
try
{
/*AccessControl accessControl = session == null
? AccessControl.DISABLED
: getAccessControl.apply(session.sessionSettings());
messages.sendErrorResponse(channel, accessControl, t);
*/
LOGGER.error("Unable to handle query", t);
messages.sendErrorResponse(channel, t);
}
catch (Throwable ti)
{
LOGGER.error("Error trying to send error to client: {}", t, ti);
}
}
}
private void dispatchState(ByteBuf buffer, DelayableWriteChannel channel) throws Exception
{
switch (decoder.state())
{
case STARTUP_PARAMETERS:
handleStartupBody(buffer, channel);
decoder.startupDone();
return;
case CANCEL:
handleCancelRequestBody(buffer, channel);
return;
case MSG:
LOGGER.trace("msg={} msgLength={} readableBytes={}", ((char) decoder.msgType()),
decoder.payloadLength(), buffer.readableBytes());
if (ignoreTillSync && decoder.msgType() != 'S')
{
buffer.skipBytes(decoder.payloadLength());
return;
}
dispatchMessage(buffer, channel);
return;
default:
throw new IllegalStateException("Illegal state: " + decoder.state());
}
}
private void dispatchMessage(ByteBuf buffer, DelayableWriteChannel channel) throws Exception
{
switch (decoder.msgType())
{
case 'Q': // Query (simple)
LOGGER.trace("Dispatching simple query");
handleSimpleQuery(buffer, channel);
return;
case 'P':
LOGGER.trace("Dispatching parse");
handleParseMessage(buffer, channel);
return;
case 'p':
LOGGER.trace("Dispatching password");
handlePassword(buffer, channel, decoder.payloadLength());
return;
case 'B':
LOGGER.trace("Dispatching bind");
handleBindMessage(buffer, channel);
return;
case 'D':
LOGGER.trace("Dispatching describe");
handleDescribeMessage(buffer, channel);
return;
case 'E':
LOGGER.trace("Dispatching execute");
handleExecute(buffer, channel);
return;
case 'H':
LOGGER.trace("Dispatching flush");
handleFlush(channel);
return;
case 'S':
LOGGER.trace("Dispatching sync");
handleSync(channel);
return;
case 'C':
LOGGER.trace("Dispatching close ");
handleClose(buffer, channel);
return;
case 'X': // Terminate (called when jdbc connection is closed)
LOGGER.trace("Dispatching close session");
closeSession();
channel.close();
return;
default:
/* messages.sendErrorResponse(
channel,
session == null
? AccessControl.DISABLED
: getAccessControl.apply(session.sessionSettings()),
new UnsupportedOperationException("Unsupported messageType: " + decoder.msgType()));*/
messages.sendErrorResponse(channel,
new UnsupportedOperationException("Unsupported messageType: " + decoder.msgType()));
}
}
private void closeSession()
{
if (session != null)
{
session.close();
session = null;
}
}
@Override
public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception
{
if (cause instanceof SocketException && cause.getMessage().equals("Connection reset"))
{
LOGGER.info("Connection reset. Client likely terminated connection");
closeSession();
}
else
{
LOGGER.error("Uncaught exception: ", cause);
}
}
@Override
public void channelUnregistered(ChannelHandlerContext ctx) throws Exception
{
LOGGER.trace("channelDisconnected");
channel = null;
closeSession();
super.channelUnregistered(ctx);
}
}
private void handleStartupBody(ByteBuf buffer, Channel channel)
{
properties = readStartupMessage(buffer);
initAuthentication(channel);
}
public static InetAddress getRemoteAddress(Channel channel)
{
if (channel.remoteAddress() instanceof InetSocketAddress)
{
return ((InetSocketAddress) channel.remoteAddress()).getAddress();
}
// In certain cases the channel is an EmbeddedChannel (e.g. in tests)
// and this type of channel has an EmbeddedSocketAddress instance as remoteAddress
// which does not have an address.
// An embedded socket address is handled like a local connection via loopback.
return InetAddresses.forString("127.0.0.1");
}
public static SSLSession getSession(Channel channel)
{
SslHandler sslHandler = channel.pipeline().get(SslHandler.class);
if (sslHandler != null)
{
return sslHandler.engine().getSession();
}
return null;
}
private void initAuthentication(Channel channel)
{
String userName = properties.getProperty("user");
InetAddress address = getRemoteAddress(channel);
SSLSession sslSession = getSession(channel);
ConnectionProperties connProperties = new ConnectionProperties(address, sslSession);
AuthenticationMethod authMethod = authService.resolveAuthenticationType(userName,
connProperties);
if (authMethod == null)
{
String errorMessage = String.format(
Locale.ENGLISH,
"No valid auth.host_based entry found for host \"%s\", user \"%s\". Did you enable TLS in your client?",
address.getHostAddress(), userName
);
messages.sendAuthenticationError(channel, errorMessage);
}
else
{
authContext = new AuthenticationContext(authMethod, connProperties, userName, LOGGER);
if (authMethod.name() == AuthenticationMethodType.PASSWORD)
{
messages.sendAuthenticationCleartextPassword(channel);
return;
}
if (authMethod.name() == AuthenticationMethodType.GSS)
{
if (gssConfig == null)
{
messages.sendAuthenticationError(channel, "GSS Auth not configured in this server");
return;
}
messages.sendAuthenticationKerberos(channel);
return;
}
finishAuthentication(channel);
}
}
private void finishAuthentication(Channel channel)
{
Assert.assertTrue(authContext != null, () -> "finishAuthentication() requires an authContext instance");
try
{
Identity authenticatedUser = authContext.authenticate();
handleAuthSuccess(channel, authenticatedUser);
}
catch (Exception e)
{
messages.sendAuthenticationError(channel, e.getMessage());
LOGGER.error("Auth Error", e);
}
finally
{
authContext.close();
authContext = null;
}
}
private void finishAuthentication(Channel channel, Subject delegSubject)
{
Assert.assertTrue(authContext != null, () -> "finishAuthentication() requires an authContext instance");
try
{
Identity authenticatedUser = KerberosIdentityProvider.getIdentityForSubject(delegSubject);
handleAuthSuccess(channel, authenticatedUser);
}
catch (Exception e)
{
messages.sendAuthenticationError(channel, e.getMessage());
LOGGER.error("Auth Error", e);
}
finally
{
authContext.close();
authContext = null;
}
}
private void handleAuthSuccess(Channel channel, Identity authenticatedUser) throws Exception
{
String database = properties.getProperty("database");
session = sessions.createSession(database, authenticatedUser);
MDC.put("user", authenticatedUser.getName());
messages.sendAuthenticationOK(channel)
.addListener(f -> sendParams(channel))
//.addListener(f -> messages.sendKeyData(channel, session.id(), session.secret()))
.addListener(f ->
{
messages.sendReadyForQuery(channel);
/*if (properties.containsKey("CrateDBTransport")) {
switchToTransportProtocol(channel);
}*/
});
}
/* private void switchToTransportProtocol(Channel channel) {
var pipeline = channel.pipeline();
pipeline.remove("frame-decoder");
pipeline.remove("handler");
// SSL is already done via PostgreSQL handshake/auth
addTransportHandler.accept(pipeline);
}*/
private void sendParams(Channel channel)
{
/* messages.sendParameterStatus(channel, "crate_version", Version.CURRENT.externalNumber());
*/
messages.sendParameterStatus(channel, "server_version", PG_SERVER_VERSION);
messages.sendParameterStatus(channel, "server_encoding", "UTF8");
messages.sendParameterStatus(channel, "client_encoding", "UTF8");
messages.sendParameterStatus(channel, "datestyle", "ISO");
messages.sendParameterStatus(channel, "TimeZone", "UTC");
messages.sendParameterStatus(channel, "integer_datetimes", "on");
}
/**
* Flush Message | 'H' | int32 len
*
* Flush forces the backend to deliver any data pending in it's output buffers.
*/
private void handleFlush(Channel channel)
{
try
{
/* // If we have deferred any executions we need to trigger a sync now because the client is expecting data
// (That we've been holding back, as we don't eager react to `execute` requests. (We do that to optimize batch inserts))
// The sync will also trigger a flush eventually if there are deferred executions.
if (session.hasDeferredExecutions()) {
session.flush();
} else {
channel.flush();
}*/
//since we don't handle buffering we should flash right away
//TODO understand when this is called
session.sync();
}
catch (Throwable t)
{
//messages.sendErrorResponse(channel, getAccessControl.apply(session.sessionSettings()), t);
messages.sendErrorResponse(channel, t);
}
}
/**
* Parse Message header: | 'P' | int32 len
*
* body: | string statementName | string query | int16 numParamTypes | foreach param: | int32
* type_oid (zero = unspecified)
*/
private void handleParseMessage(ByteBuf buffer, final Channel channel)
{
String statementName = readCString(buffer);
final String query = readCString(buffer);
short numParams = buffer.readShort();
List paramTypes = new ArrayList<>(numParams);
for (int i = 0; i < numParams; i++)
{
int oid = buffer.readInt();
int dataType = PGTypes.fromOID(oid);
/* if (dataType == null) {
throw new IllegalArgumentException(
String.format(Locale.ENGLISH, "Can't map PGType with oid=%d to Crate type", oid));
}*/
paramTypes.add(dataType);
}
CompletableFuture parseCompletionFuture = session.parseAsync(statementName, query, paramTypes);
CompletableFuture parseMsgSentFuture = parseCompletionFuture.thenRun(() -> messages.sendParseComplete(channel));
session.setActiveExecution(parseMsgSentFuture);
}
private void handlePassword(ByteBuf buffer, final Channel channel, int payloadLength)
{
switch (authContext.getAuthenticationMethodType())
{
case GSS:
byte[] inputToken = readByteArray(buffer, payloadLength);
if (inputToken == null)
{
messages.sendErrorResponse(channel, new IllegalStateException("GSS Token cannot be empty"));
return;
}
Subject serverSubject = SubjectTools.getSubjectFromKeytab(gssConfig.getKerberosKeytabFile(), gssConfig.getKerberosUserPrincipal(), false);
GSSManager manager = GSSManager.getInstance();
try
{
GSSCredential gssCredential = Subject.doAs(serverSubject, new AcceptorCreator(manager, gssConfig.getKerberosUserPrincipal()));
GSSContext gssContext = manager.createContext(gssCredential);
gssContext.requestCredDeleg(true);
gssContext.requestMutualAuth(true);
byte[] outputToken;
if (!gssContext.isEstablished())
{
outputToken = gssContext.acceptSecContext(inputToken, 0, inputToken.length);
if (outputToken != null)
{
messages.sendGssOutToken(channel, outputToken);
}
}
Subject delegatedSubject = GSSUtil.createSubject(gssContext.getSrcName(), gssContext.getDelegCred());
finishAuthentication(channel, delegatedSubject);
}
catch (PrivilegedActionException | GSSException e)
{
throw new RuntimeException(e);
}
break;
case PASSWORD:
default:
char[] passwd = readCharArray(buffer);
if (passwd != null)
{
authContext.setSecurePassword(passwd);
}
finishAuthentication(channel);
}
}
/**
* Bind Message Header: | 'B' | int32 len
*
* Body:
*
* | string portalName | string statementName
* | int16 numFormatCodes
* foreach
* | int16 formatCode
* | int16 numParams
* foreach
* | int32 valueLength
* | byteN value
* | int16 numResultColumnFormatCodes
* foreach
* | int16 formatCode
*
*/
private void handleBindMessage(ByteBuf buffer, Channel channel)
{
String portalName = readCString(buffer);
String statementName = readCString(buffer);
FormatCodes.FormatCode[] formatCodes = FormatCodes.fromBuffer(buffer);
short numParams = buffer.readShort();
List