feat: support OAuth2 token authentication (#360)

* feat: support OAuth2 token authentication Add support for authentication with simple OAuth2 tokens. * chore: move error message to hint + more tests * fix: clirr check + credentials error message

Author: olavloite

clirr-ignored-differences.xml

@@ -55,6 +55,18 @@
  <className>com/google/cloud/spanner/pgadapter/wireprotocol/AbstractQueryProtocolMessage</className>
  <method>java.lang.String getSql()</method>
  </difference>
+ <difference>
+ <differenceType>7005</differenceType>
+ <className>com/google/cloud/spanner/connection/ConnectionOptionsHelper</className>
+ <method>com.google.cloud.spanner.connection.ConnectionOptions$Builder setCredentials(com.google.cloud.spanner.connection.ConnectionOptions$Builder, com.google.auth.oauth2.GoogleCredentials)</method>
+ <to>com.google.cloud.spanner.connection.ConnectionOptions$Builder setCredentials(com.google.cloud.spanner.connection.ConnectionOptions$Builder, com.google.auth.Credentials)</to>
+ </difference>
+ <difference>
+ <differenceType>7005</differenceType>
+ <className>com/google/cloud/spanner/pgadapter/ConnectionHandler</className>
+ <method>void connectToSpanner(java.lang.String, com.google.auth.oauth2.GoogleCredentials)</method>
+ <to>void connectToSpanner(java.lang.String, com.google.auth.Credentials)</to>
+ </difference>
 
  <!-- Add ignores for all sub packages, as these are considered internal. -->
  <difference>

docs/authentication.md

@@ -33,9 +33,10 @@ java -jar pgadapter.jar -p my-project -i my-instance
 PGAdapter will require all connections to supply credentials for each connection if it is started
 with the `-a` command line argument. The credentials are sent as a password message. The password
 message must contain one of the following:
-1. The password field can contain the JSON payload of a credentials file, for example the contents
+1. The password field can contain a valid OAuth2 token. The username must be `oauth2`.
+2. The password field can contain the JSON payload of a credentials file, for example the contents
  of a service account key file. The username will be ignored in this case.
-2. The username field can contain the email address of a service account and the password field can
+3. The username field can contain the email address of a service account and the password field can
  contain a private key for that service account. Note: The password must include the
  `-----BEGIN PRIVATE KEY-----` and `-----END PRIVATE KEY-----` header and footer.
 
@@ -45,6 +46,12 @@ message must contain one of the following:
 # The -a argument instructs PGAdapter to require authentication.
 java -jar pgadapter.jar -p my-project -i my-instance -a
 
+# Set the username to 'oauth2' and the password to a valid OAuth2 token.
+# Note that PGAdapter will not be able to refresh the OAuth2 token, which means that the connection
+# will expire when the token has expired.
+PGPASSWORD=$(gcloud auth application-default print-access-token --quiet) \
+ psql -U oauth2 -h /tmp -d my-database
+
 # Set the PG password to the contents of a credentials file. This can be both a service account or a
 # user account file. The username will be ignored by PGAdapter.
 PGPASSWORD=$(cat /path/to/credentials.json) psql -h /tmp -d my-database
@@ -73,3 +80,13 @@ PGPASSWORD=$(cat /path/to/credentials.json) \
  psql -h /tmp -d "projects/my-project/instances/my-instance/databases/my-database"
 ```
 
+`psql` will by default show the name of the connected database on the prompt. You can shorten this
+by changing the default prompt of `psql` to for example show `'~'` when connected to the default
+database.
+
+```shell
+PGPASSWORD=$(gcloud auth application-default print-access-token --quiet) \
+PGDATABASE=projects/my-project/instances/my-instance/databases/my-database \
+ psql -U oauth2 -h /tmp \
+ -v "PROMPT1=%~%R%x%#" -v "PROMPT2=%~%R%x%#"
+```

src/main/java/com/google/cloud/spanner/connection/ConnectionOptionsHelper.java

@@ -15,16 +15,15 @@
 package com.google.cloud.spanner.connection;
 
 import com.google.api.core.InternalApi;
-import com.google.auth.oauth2.GoogleCredentials;
+import com.google.auth.Credentials;
 import com.google.cloud.spanner.Spanner;
 import com.google.cloud.spanner.connection.ConnectionOptions.Builder;
 
 /** Simple helper class to get access to a package-private method in the ConnectionOptions. */
 @InternalApi
 public class ConnectionOptionsHelper {
  // TODO: Remove when Builder.setCredentials(..) has been made public.
- public static Builder setCredentials(
- Builder connectionOptionsBuilder, GoogleCredentials credentials) {
+ public static Builder setCredentials(Builder connectionOptionsBuilder, Credentials credentials) {
  return connectionOptionsBuilder.setCredentials(credentials);
  }
 

src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java

@@ -15,7 +15,7 @@
 package com.google.cloud.spanner.pgadapter;
 
 import com.google.api.core.InternalApi;
-import com.google.auth.oauth2.GoogleCredentials;
+import com.google.auth.Credentials;
 import com.google.cloud.spanner.Database;
 import com.google.cloud.spanner.DatabaseAdminClient;
 import com.google.cloud.spanner.DatabaseId;
@@ -128,7 +128,7 @@ void createSSLSocket() throws IOException {
  }
 
  @InternalApi
- public void connectToSpanner(String database, @Nullable GoogleCredentials credentials) {
+ public void connectToSpanner(String database, @Nullable Credentials credentials) {
  OptionsMetadata options = getServer().getOptions();
  String uri =
  options.hasDefaultConnectionUrl()

src/main/java/com/google/cloud/spanner/pgadapter/error/PGException.java

@@ -27,6 +27,7 @@ public static class Builder {
  private Severity severity;
  private SQLState sqlState;
  private String message;
+ private String hints;
 
  private Builder() {}
 
@@ -45,8 +46,13 @@ public Builder setMessage(String message) {
  return this;
  }
 
+ public Builder setHints(String hints) {
+ this.hints = Preconditions.checkNotNull(hints);
+ return this;
+ }
+
  public PGException build() {
- return new PGException(severity, sqlState, message);
+ return new PGException(severity, sqlState, message, hints);
  }
  }
 
@@ -56,11 +62,13 @@ public static Builder newBuilder() {
 
  private final Severity severity;
  private final SQLState sqlState;
+ private final String hints;
 
- private PGException(Severity severity, SQLState sqlState, String message) {
+ private PGException(Severity severity, SQLState sqlState, String message, String hints) {
  super(Preconditions.checkNotNull(message));
  this.severity = severity;
  this.sqlState = sqlState;
+ this.hints = hints;
  }
 
  public Severity getSeverity() {
@@ -70,4 +78,8 @@ public Severity getSeverity() {
  public SQLState getSQLState() {
  return sqlState;
  }
+
+ public String getHints() {
+ return hints;
+ }
 }

src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/ErrorResponse.java

@@ -16,6 +16,7 @@
 
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.pgadapter.error.PGException;
+import com.google.common.base.Strings;
 import java.io.DataOutputStream;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
@@ -32,15 +33,27 @@ public class ErrorResponse extends WireOutput {
  private static final byte CODE_FLAG = 'C';
  private static final byte MESSAGE_FLAG = 'M';
  private static final byte SEVERITY_FLAG = 'S';
+ private static final byte HINT_FLAG = 'H';
  private static final byte NULL_TERMINATOR = 0;
 
  private final byte[] severity;
  private final byte[] errorMessage;
  private final byte[] errorState;
+ private final byte[] hints;
 
  public ErrorResponse(DataOutputStream output, PGException pgException) {
- super(
- output,
+ super(output, calculateLength(pgException));
+ this.errorMessage = pgException.getMessage().getBytes(StandardCharsets.UTF_8);
+ this.errorState = pgException.getSQLState().getBytes();
+ this.severity = pgException.getSeverity().name().getBytes(StandardCharsets.UTF_8);
+ this.hints =
+ Strings.isNullOrEmpty(pgException.getHints())
+ ? null
+ : pgException.getHints().getBytes(StandardCharsets.UTF_8);
+ }
+
+ static int calculateLength(PGException pgException) {
+ int length =
  HEADER_LENGTH
  + FIELD_IDENTIFIER_LENGTH
  + pgException.getSeverity().name().getBytes(StandardCharsets.UTF_8).length
@@ -51,10 +64,14 @@ public ErrorResponse(DataOutputStream output, PGException pgException) {
  + FIELD_IDENTIFIER_LENGTH
  + pgException.getMessage().getBytes(StandardCharsets.UTF_8).length
  + NULL_TERMINATOR_LENGTH
- + NULL_TERMINATOR_LENGTH);
- this.errorMessage = pgException.getMessage().getBytes(StandardCharsets.UTF_8);
- this.errorState = pgException.getSQLState().getBytes();
- this.severity = pgException.getSeverity().name().getBytes(StandardCharsets.UTF_8);
+ + NULL_TERMINATOR_LENGTH;
+ if (!Strings.isNullOrEmpty(pgException.getHints())) {
+ length +=
+ FIELD_IDENTIFIER_LENGTH
+ + pgException.getHints().getBytes(StandardCharsets.UTF_8).length
+ + NULL_TERMINATOR_LENGTH;
+ }
+ return length;
  }
 
  @Override
@@ -68,6 +85,11 @@ protected void sendPayload() throws IOException {
  this.outputStream.writeByte(MESSAGE_FLAG);
  this.outputStream.write(this.errorMessage);
  this.outputStream.writeByte(NULL_TERMINATOR);
+ if (this.hints != null) {
+ this.outputStream.writeByte(HINT_FLAG);
+ this.outputStream.write(this.hints);
+ this.outputStream.writeByte(NULL_TERMINATOR);
+ }
  this.outputStream.writeByte(NULL_TERMINATOR);
  }
 
@@ -83,7 +105,12 @@ protected String getMessageName() {
 
  @Override
  protected String getPayloadString() {
- return new MessageFormat("Length: {0}, " + "Error Message: {1}")
- .format(new Object[] {this.length, new String(this.errorMessage, UTF8)});
+ return new MessageFormat("Length: {0}, Error Message: {1}, Hints: {2}")
+ .format(
+ new Object[] {
+ this.length,
+ new String(this.errorMessage, UTF8),
+ this.hints == null ? "" : new String(this.hints, UTF8)
+ });
  }
 }

src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/PasswordMessage.java

@@ -21,7 +21,10 @@
 import com.google.api.client.util.PemReader.Section;
 import com.google.api.client.util.Strings;
 import com.google.api.core.InternalApi;
+import com.google.auth.Credentials;
+import com.google.auth.oauth2.AccessToken;
 import com.google.auth.oauth2.GoogleCredentials;
+import com.google.auth.oauth2.OAuth2Credentials;
 import com.google.auth.oauth2.ServiceAccountCredentials;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.error.PGException;
@@ -36,8 +39,15 @@
 import org.postgresql.util.ReaderInputStream;
 
 /**
- * A Password Message takes a username and password and input and supposedly handles auth. Here,
- * however, since connections are through localhost, we do not do so.
+ * PGAdapter will convert a password message into gRPC authentication in the following ways:
+ *
+ * <ol>
+ * <li>If the username is 'oauth2' the password will be interpreted as an OAuth2 token.
+ * <li>If the username is an email address and the password contains private key section,
+ * PGAdapter will construct a service account from the email address and private key.
+ * <li>Otherwise, PGAdapter will try to construct a Google credentials instance from the string in
+ * the password message. The username will be ignored.
+ * </ol>
  */
 @InternalApi
 public class PasswordMessage extends ControlMessage {
@@ -71,16 +81,17 @@ protected void sendPayload() throws Exception {
  return;
  }
 
- GoogleCredentials credentials = checkCredentials(this.username, this.password);
+ Credentials credentials = checkCredentials(this.username, this.password);
  if (credentials == null) {
  new ErrorResponse(
  this.outputStream,
  PGException.newBuilder()
- .setMessage(
- "Invalid credentials received. "
- + "PGAdapter expects the password to contain the JSON payload of a credentials file. "
- + "Alternatively, the password may contain only the private key of a service account. "
- + "The user name must in that case contain the service account email address.")
+ .setMessage("Invalid credentials received.")
+ .setHints(
+ "PGAdapter expects credentials to be one of the following:\n"
+ + "1. Username contains the fixed string 'oauth2' and the password field contains a valid OAuth2 token.\n"
+ + "2. Username contains any string and the password field contains the JSON payload of a credentials file (e.g. a service account file).\n"
+ + "3. Username contains the email address of a service account and the password contains the corresponding private key for the service account.")
  .setSQLState(SQLState.InvalidPassword)
  .setSeverity(Severity.ERROR)
  .build())
@@ -96,7 +107,7 @@ private boolean useAuthentication() {
  return this.connection.getServer().getOptions().shouldAuthenticate();
  }
 
- private GoogleCredentials checkCredentials(String username, String password) {
+ private Credentials checkCredentials(String username, String password) {
  if (Strings.isNullOrEmpty(password)) {
  return null;
  }
@@ -126,6 +137,11 @@ private GoogleCredentials checkCredentials(String username, String password) {
  }
  }
 
+ if (!Strings.isNullOrEmpty(username) && username.equalsIgnoreCase("oauth2")) {
+ // Interpret the password as an OAuth2 token.
+ return OAuth2Credentials.create(new AccessToken(password, null));
+ }
+
  // Try to parse the password field as a JSON string that contains a credentials object.
  try {
  return GoogleCredentials.fromStream(new ReaderInputStream(new StringReader(password)));

src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/StartupMessage.java

@@ -15,7 +15,7 @@
 package com.google.cloud.spanner.pgadapter.wireprotocol;
 
 import com.google.api.core.InternalApi;
-import com.google.auth.oauth2.GoogleCredentials;
+import com.google.auth.Credentials;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler.ConnectionStatus;
 import com.google.cloud.spanner.pgadapter.utils.ClientAutoDetector;
@@ -79,7 +79,7 @@ static void createConnectionAndSendStartupMessage(
  ConnectionHandler connection,
  String database,
  Map<String, String> parameters,
- @Nullable GoogleCredentials credentials)
+ @Nullable Credentials credentials)
  throws Exception {
  connection.connectToSpanner(database, credentials);
  for (Entry<String, String> parameter : parameters.entrySet()) {