feat: login functionality for experimental host (DRAFT)

pom.xml

@@ -145,6 +145,7 @@
     <module>proto-google-cloud-spanner-executor-v1</module>
     <module>google-cloud-spanner-executor</module>
     <module>google-cloud-spanner-bom</module>
+    <module>proto-google-cloud-spanner-host</module>
   </modules>
 
   <build>

proto-google-cloud-spanner-host/pom.xml

@@ -0,0 +1,77 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+    xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+    xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+  <modelVersion>4.0.0</modelVersion>
+  <parent>
+    <groupId>com.google.cloud</groupId>
+    <artifactId>google-cloud-spanner-parent</artifactId>
+    <version>6.88.0</version>
+  </parent>
+
+  <artifactId>proto-google-cloud-spanner-host</artifactId>
+
+  <properties>
+    <maven.compiler.source>17</maven.compiler.source>
+    <maven.compiler.target>17</maven.compiler.target>
+    <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
+  </properties>
+  <dependencies>
+    <dependency>
+      <groupId>io.grpc</groupId>
+      <artifactId>grpc-protobuf</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.grpc</groupId>
+      <artifactId>grpc-stub</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>com.google.protobuf</groupId>
+      <artifactId>protobuf-java</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>org.bouncycastle</groupId>
+      <artifactId>bcprov-jdk15on</artifactId>
+      <version>1.70</version>
+    </dependency>
+    <dependency>
+      <groupId>io.grpc</groupId>
+      <artifactId>grpc-netty-shaded</artifactId>
+    </dependency>
+    <!-- Test dependencies -->
+    <dependency>
+      <groupId>junit</groupId>
+      <artifactId>junit</artifactId>
+      <scope>test</scope>
+    </dependency>
+  </dependencies>
+  <build>
+    <plugins>
+      <plugin>
+        <groupId>org.xolstice.maven.plugins</groupId>
+        <artifactId>protobuf-maven-plugin</artifactId>
+        <version>0.6.1</version>
+        <configuration>
+          <protocArtifact>com.google.protobuf:protoc:LATEST:exe:${os.detected.classifier}</protocArtifact>
+          <pluginId>grpc-java</pluginId>
+          <pluginArtifact>io.grpc:protoc-gen-grpc-java:LATEST:exe:${os.detected.classifier}</pluginArtifact>
+        </configuration>
+        <executions>
+          <execution>
+            <goals>
+              <goal>compile</goal>
+              <goal>compile-custom</goal>
+            </goals>
+          </execution>
+        </executions>
+      </plugin>
+    </plugins>
+    <extensions>
+      <extension>
+        <groupId>kr.motd.maven</groupId>
+        <artifactId>os-maven-plugin</artifactId>
+        <version>1.7.0</version>
+      </extension>
+    </extensions>
+  </build>
+</project>

proto-google-cloud-spanner-host/src/main/java/com/google/cloud/GrpcClient.java

@@ -0,0 +1,111 @@
+/*
+ * Copyright 2017 Google LLC
+ *
+ * Licensed 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.
+ */
+
+package com.google.cloud;
+
+import io.grpc.ChannelCredentials;
+import io.grpc.Grpc;
+import io.grpc.ManagedChannel;
+import io.grpc.TlsChannelCredentials;
+import io.grpc.netty.shaded.io.grpc.netty.GrpcSslContexts;
+import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder;
+import java.io.File;
+import java.io.IOException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.security.cert.Certificate;
+import java.security.cert.X509Certificate;
+import javax.net.ssl.SSLPeerUnverifiedException;
+import javax.net.ssl.SSLSocket;
+import javax.net.ssl.SSLSocketFactory;
+import spanner.experimental.LoginServiceGrpc;
+
+public class GrpcClient {
+
+  private final ManagedChannel channel;
+  private final LoginServiceGrpc.LoginServiceStub loginService;
+  private X509Certificate serverCertificate; // Store the server certificate
+  private final int DEFAULT_PORT = 15000;
+
+  public GrpcClient(String endpoint) throws IOException {
+    this(endpoint, null, null);
+  }
+
+  public GrpcClient(String endpoint, String clientCertificate, String clientCertificateKey) {
+    try {
+      URI uri = new URI(endpoint);
+      String host = uri.getHost();
+      int port = (uri.getPort() == -1) ? DEFAULT_PORT : uri.getPort();
+
+      if (host == null) {
+        throw new IllegalArgumentException("Invalid endpoint: " + endpoint);
+      }
+
+      // Retrieve the server certificate during handshake
+      this.serverCertificate = fetchServerCertificate(host, port);
+
+      if (clientCertificate != null && clientCertificateKey != null) {
+        this.channel =
+            NettyChannelBuilder.forAddress(host, port)
+                .sslContext(
+                    GrpcSslContexts.forClient()
+                        .keyManager(
+                            new File(clientCertificate),
+                            new File(clientCertificateKey)) // Client auth
+                        .build())
+                .build();
+      } else {
+        // Normal TLS
+        ChannelCredentials credentials = TlsChannelCredentials.newBuilder().build();
+        this.channel = Grpc.newChannelBuilderForAddress(host, port, credentials).build();
+      }
+      this.loginService = LoginServiceGrpc.newStub(channel);
+    } catch (URISyntaxException | IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  /** Establish a TLS connection to fetch the server certificate. */
+  private X509Certificate fetchServerCertificate(String host, int port) throws IOException {
+    try {
+      SSLSocketFactory sslSocketFactory = (SSLSocketFactory) SSLSocketFactory.getDefault();
+      try (SSLSocket socket = (SSLSocket) sslSocketFactory.createSocket(host, port)) {
+        socket.startHandshake();
+        Certificate[] serverCerts = socket.getSession().getPeerCertificates();
+        if (serverCerts.length > 0 && serverCerts[0] instanceof X509Certificate) {
+          return (X509Certificate) serverCerts[0]; // Store the first certificate
+        }
+      }
+    } catch (SSLPeerUnverifiedException e) {
+      throw new IOException("Failed to verify server certificate: " + e.getMessage(), e);
+    }
+    throw new IOException("No server certificate found");
+  }
+
+  public LoginServiceGrpc.LoginServiceStub getLoginService() {
+    return loginService;
+  }
+
+  public X509Certificate getServerCertificate() {
+    return serverCertificate;
+  }
+
+  public void close() {
+    if (channel != null) {
+      channel.shutdown();
+    }
+  }
+}

proto-google-cloud-spanner-host/src/main/java/com/google/cloud/LoginClient.java

@@ -0,0 +1,118 @@
+/*
+ * Copyright 2017 Google LLC
+ *
+ * Licensed 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.
+ */
+
+package com.google.cloud;
+
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.ReentrantLock;
+import spanner.experimental.AccessToken;
+
+public class LoginClient {
+
+  private final String username;
+  private final String password;
+  private final String endpoint;
+  private volatile AccessToken accessToken;
+  private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
+  private final ReentrantLock refreshLock = new ReentrantLock();
+  private ScheduledFuture<?> scheduledTask; // Holds the scheduled task
+  private String clientCertificate = null;
+  private String clientCertificateKey = null;
+
+  private static final long TOKEN_REFRESH_THRESHOLD_SECONDS = 300; // Refresh 5 minutes before expiry
+
+  public LoginClient(String username, String password, String endpoint) throws Exception {
+    this(username, password, endpoint, null, null);
+  }
+
+  public LoginClient(
+      String username,
+      String password,
+      String endpoint,
+      String clientCertificate,
+      String clientCertificateKey) {
+    this.username = username;
+    this.password = password;
+    this.endpoint = endpoint;
+    if (clientCertificate != null && clientCertificateKey != null) {
+      this.clientCertificate = clientCertificate;
+      this.clientCertificateKey = clientCertificateKey;
+    }
+    login();
+    scheduleNextTokenRefresh();
+  }
+
+  private void login() {
+    try {
+      Scram scram =
+          new Scram(
+              this.username,
+              this.password,
+              new GrpcClient(this.endpoint, this.clientCertificate, this.clientCertificateKey));
+      this.accessToken = scram.login();
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  public AccessToken getAccessToken() {
+    return accessToken;
+  }
+
+  private void scheduleNextTokenRefresh() {
+    if (accessToken == null) return;
+    long delay =
+        (accessToken.getExpirationTime().getSeconds() - System.currentTimeMillis() / 1000)
+            - TOKEN_REFRESH_THRESHOLD_SECONDS;
+
+    if (delay <= 0) {
+      refreshToken();
+      return;
+    }
+    if (scheduledTask != null) {
+      scheduledTask.cancel(false);
+    }
+    // Schedule a new token refresh exactly when needed
+    scheduledTask = scheduler.schedule(this::refreshToken, delay, TimeUnit.SECONDS);
+    System.out.println("Next token refresh scheduled in " + delay + " seconds.");
+  }
+
+  private void refreshToken() {
+    if (!refreshLock.tryLock()) return; // Prevent multiple simultaneous refreshes
+
+    try {
+      System.out.println("Refreshing access token...");
+      login();
+      System.out.println("New token acquired.\n" + getAccessToken());
+      scheduleNextTokenRefresh();
+    } catch (Exception e) {
+      System.err.println("Token refresh failed: " + e.getMessage());
+    } finally {
+      refreshLock.unlock();
+    }
+  }
+
+  public void shutdown() {
+    System.out.println("Shutting down LoginClient...");
+    if (scheduledTask != null) {
+      scheduledTask.cancel(false); // Cancel any pending token refresh task
+    }
+    scheduler.shutdown();
+  }
+}