diff --git a/google-auth-library-java/oauth2_http/java/com/google/auth/mtls/MtlsUtils.java b/google-auth-library-java/oauth2_http/java/com/google/auth/mtls/MtlsUtils.java index a5c4c0f86e77..d46d3822f827 100644 --- a/google-auth-library-java/oauth2_http/java/com/google/auth/mtls/MtlsUtils.java +++ b/google-auth-library-java/oauth2_http/java/com/google/auth/mtls/MtlsUtils.java @@ -34,12 +34,17 @@ import com.google.auth.oauth2.EnvironmentProvider; import com.google.auth.oauth2.PropertyProvider; import com.google.common.base.Strings; +import com.google.common.io.BaseEncoding; import java.io.File; import java.io.FileInputStream; import java.io.IOException; import java.io.InputStream; +import java.nio.file.Files; +import java.nio.file.Paths; +import java.security.MessageDigest; import java.util.Locale; import org.jspecify.annotations.NullMarked; +import org.jspecify.annotations.Nullable; /** * Utility class for mTLS related operations. @@ -57,6 +62,168 @@ private MtlsUtils() { // Prevent instantiation for Utility class } + /** + * Returns if mutual TLS client certificate should be used. Returns true if valid workload + * certificates are configured or if GOOGLE_API_USE_CLIENT_CERTIFICATE is explicitly set to true + * (e.g. for Enterprise Certificate Proxy or custom MtlsProviders), unless explicitly disabled via + * GOOGLE_API_USE_CLIENT_CERTIFICATE=false. + */ + public static boolean useMtlsClientCertificate( + EnvironmentProvider envProvider, PropertyProvider propProvider) { + String useClientCertificate = envProvider.getEnv("GOOGLE_API_USE_CLIENT_CERTIFICATE"); + if ("false".equalsIgnoreCase(useClientCertificate)) { + return false; + } + if (getWorkloadCertPath(envProvider, propProvider) != null) { + return true; + } + return "true".equalsIgnoreCase(useClientCertificate); + } + + /** + * Resolves and returns the path to the mutual TLS client certificate, or null if none should be + * used. + * + *

Possible outcomes: + * + *

    + *
  1. Non-null {@link String} (Valid happy path): A valid workload certificate + * configuration was found and both the certificate and private key files exist and are + * readable. + *
  2. {@link IllegalStateException} (Invalid state - fail closed): An explicit {@code + * GOOGLE_API_CERTIFICATE_CONFIG} path or an existing default well-known certificate + * configuration file is missing, unreadable, malformed, or references missing/unreadable + * certificate or private key files. This is treated as an unrecoverable misconfiguration. + *
  3. {@code null} (Safe fallback / fail open): Client certificates are explicitly + * disabled via {@code GOOGLE_API_USE_CLIENT_CERTIFICATE=false}, no explicit configuration + * is set and the default well-known configuration file does not exist on disk, or the + * configuration specifies an non-workload source (e.g., ECP/PKCS11 without a {@code + * workload} section). Callers can proceed without workload certificate file polling. + *
+ */ + public static @Nullable String getWorkloadCertPath( + EnvironmentProvider envProvider, PropertyProvider propProvider) { + String useClientCertificate = envProvider.getEnv("GOOGLE_API_USE_CLIENT_CERTIFICATE"); + if ("false".equalsIgnoreCase(useClientCertificate)) { + return null; + } + + String explicitConfigPath = envProvider.getEnv(CERTIFICATE_CONFIGURATION_ENV_VARIABLE); + + // 1. Explicit Configuration Path (Fail Closed) + if (!Strings.isNullOrEmpty(explicitConfigPath)) { + File configFile = new File(explicitConfigPath); + if (!configFile.exists()) { + throw new IllegalStateException( + "Certificate configuration file specified via GOOGLE_API_CERTIFICATE_CONFIG at '" + + explicitConfigPath + + "' does not exist."); + } + if (!configFile.isFile() || !configFile.canRead()) { + throw new IllegalStateException( + "Failed to read certificate configuration file specified via" + + " GOOGLE_API_CERTIFICATE_CONFIG at '" + + explicitConfigPath + + "'."); + } + WorkloadCertificateConfiguration config; + try { + config = getWorkloadCertificateConfiguration(envProvider, propProvider, explicitConfigPath); + } catch (CertificateSourceUnavailableException e) { + // ECP / PKCS11 configuration without workload section; safe fallback + return null; + } catch (Exception e) { + throw new IllegalStateException( + "Certificate configuration file specified via GOOGLE_API_CERTIFICATE_CONFIG at '" + + explicitConfigPath + + "' is malformed: " + + e.getMessage(), + e); + } + checkCertAndKeyFilesReadable(config, explicitConfigPath, false); + return config.getCertPath(); + } + + // 2. Implicit / Default gcloud Configuration Path + File defaultConfigFile = null; + try { + defaultConfigFile = getWellKnownCertificateConfigFile(envProvider, propProvider); + } catch (IOException e) { + // APPDATA missing on Windows, etc. Safe fallback. + } + if (defaultConfigFile != null && defaultConfigFile.exists()) { + if (!defaultConfigFile.isFile() || !defaultConfigFile.canRead()) { + throw new IllegalStateException( + "Default certificate configuration file at '" + + defaultConfigFile.getAbsolutePath() + + "' exists but could not be read."); + } + WorkloadCertificateConfiguration config = null; + try { + config = getWorkloadCertificateConfiguration(envProvider, propProvider, null); + } catch (CertificateSourceUnavailableException e) { + // ECP-only configuration without workload section; safe fallback + } catch (Exception e) { + throw new IllegalStateException( + "Default certificate configuration file at '" + + defaultConfigFile.getAbsolutePath() + + "' is malformed: " + + e.getMessage(), + e); + } + if (config != null) { + checkCertAndKeyFilesReadable(config, defaultConfigFile.getAbsolutePath(), true); + return config.getCertPath(); + } + } + + return null; + } + + private static void checkCertAndKeyFilesReadable( + WorkloadCertificateConfiguration config, String configPath, boolean isDefaultConfig) { + File certFile = new File(config.getCertPath()); + File keyFile = new File(config.getPrivateKeyPath()); + if (!certFile.isFile() || !certFile.canRead() || !keyFile.isFile() || !keyFile.canRead()) { + String sourcePrefix = + isDefaultConfig + ? "referenced by default configuration '" + : "referenced by configuration '"; + throw new IllegalStateException( + "Failed to read certificate/key file at '" + + config.getCertPath() + + "' or '" + + config.getPrivateKeyPath() + + "' " + + sourcePrefix + + configPath + + "'."); + } + } + + /** + * Computes the lower-case SHA-256 hex fingerprint of the certificate file at {@code certPath}. + * + *

Unlike {@link #getWorkloadCertPath}, which validates configuration at channel initialization + * and fails closed on errors, this method is called dynamically at runtime during active RPCs to + * detect certificate rotations on disk. External certificate rotators may temporarily delete, + * truncate, or rewrite the certificate file mid-RPC. Returning {@code null} on read/digest + * exceptions (which callers normalize to {@code ""}) allows runtime refresh checks to ignore + * transient mid-write states and keep the active healthy channel without failing in-flight RPCs. + */ + public static @Nullable String getCertificateFingerprint(@Nullable String certPath) { + if (certPath == null) { + return null; + } + try { + byte[] certBytes = Files.readAllBytes(Paths.get(certPath)); + byte[] digest = MessageDigest.getInstance("SHA-256").digest(certBytes); + return BaseEncoding.base16().lowerCase().encode(digest); + } catch (Exception e) { + return null; + } + } + /** * Returns the path to the client certificate file specified by the loaded workload certificate * configuration. @@ -65,14 +232,17 @@ private MtlsUtils() { * @throws IOException if the certificate configuration cannot be found or loaded. */ public static String getCertificatePath( - EnvironmentProvider envProvider, PropertyProvider propProvider, String certConfigPathOverride) + EnvironmentProvider envProvider, + PropertyProvider propProvider, + @Nullable String certConfigPathOverride) throws IOException { String certPath = getWorkloadCertificateConfiguration(envProvider, propProvider, certConfigPathOverride) .getCertPath(); if (Strings.isNullOrEmpty(certPath)) { throw new CertificateSourceUnavailableException( - "Certificate configuration loaded successfully, but does not contain a 'certificate_file' path."); + "Certificate configuration loaded successfully, but does not contain a" + + " 'cert_configs.workload.cert_path' path."); } return certPath; } @@ -92,7 +262,9 @@ public static String getCertificatePath( * @throws IOException if the configuration file cannot be found, read, or parsed */ static WorkloadCertificateConfiguration getWorkloadCertificateConfiguration( - EnvironmentProvider envProvider, PropertyProvider propProvider, String certConfigPathOverride) + EnvironmentProvider envProvider, + PropertyProvider propProvider, + @Nullable String certConfigPathOverride) throws IOException { File certConfig; if (certConfigPathOverride != null) { @@ -106,7 +278,7 @@ static WorkloadCertificateConfiguration getWorkloadCertificateConfiguration( } } - if (!certConfig.isFile()) { + if (!certConfig.isFile() || !certConfig.canRead()) { throw new CertificateSourceUnavailableException( "Certificate configuration file does not exist or is not a file: " + certConfig.getAbsolutePath()); diff --git a/google-auth-library-java/oauth2_http/javatests/com/google/auth/mtls/MtlsUtilsTest.java b/google-auth-library-java/oauth2_http/javatests/com/google/auth/mtls/MtlsUtilsTest.java index f3fdf05a4c32..d5fdc840d804 100644 --- a/google-auth-library-java/oauth2_http/javatests/com/google/auth/mtls/MtlsUtilsTest.java +++ b/google-auth-library-java/oauth2_http/javatests/com/google/auth/mtls/MtlsUtilsTest.java @@ -101,6 +101,21 @@ public String getProperty(String name, String def) { () -> MtlsUtils.getCertificatePath(envProvider, propProvider, configFile.toString())); } + @Test + void getCertificatePath_ecpOnlyConfig_throwsCertificateSourceUnavailableException() + throws IOException { + Path configFile = tempDir.resolve("ecp_config.json"); + Files.write( + configFile, "{\"cert_configs\":{\"enterprise_certificates\":{\"libs\":[]}}}".getBytes()); + + EnvironmentProvider envProvider = name -> null; + PropertyProvider propProvider = (name, def) -> def; + + assertThrows( + CertificateSourceUnavailableException.class, + () -> MtlsUtils.getCertificatePath(envProvider, propProvider, configFile.toString())); + } + @Test void getWorkloadCertificateConfiguration_overridePath() throws IOException { Path configFile = tempDir.resolve("custom_config.json"); @@ -243,4 +258,484 @@ public String getProperty(String name, String def) { assertEquals("APPDATA environment variable is not set on Windows.", exception.getMessage()); } + + @Test + void + useMtlsClientCertificate_trueWithNoCertsOnDisk_returnsTrueWhileWorkloadCertPathReturnsNull() { + EnvironmentProvider envProvider = + name -> "GOOGLE_API_USE_CLIENT_CERTIFICATE".equals(name) ? "true" : null; + PropertyProvider propProvider = + (name, def) -> { + if ("user.home".equals(name)) return tempDir.toString(); + if ("os.name".equals(name)) return "Linux"; + return def; + }; + + assertTrue(MtlsUtils.useMtlsClientCertificate(envProvider, propProvider)); + assertNull(MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + } + + @Test + void useMtlsClientCertificate_trueWithEcpOnlyConfig_returnsTrueAndWorkloadCertPathReturnsNull() + throws IOException { + Path configFile = tempDir.resolve("ecp_config.json"); + Files.write(configFile, "{\"cert_configs\":{\"enterprise_certificates\":{}}}".getBytes()); + + EnvironmentProvider envProvider = + name -> { + if ("GOOGLE_API_USE_CLIENT_CERTIFICATE".equals(name)) return "true"; + if ("GOOGLE_API_CERTIFICATE_CONFIG".equals(name)) return configFile.toString(); + return null; + }; + PropertyProvider propProvider = (name, def) -> def; + + assertTrue(MtlsUtils.useMtlsClientCertificate(envProvider, propProvider)); + assertNull(MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + } + + @Test + void useMtlsClientCertificate_false_returnsFalse() { + EnvironmentProvider envProvider = + name -> "GOOGLE_API_USE_CLIENT_CERTIFICATE".equals(name) ? "false" : null; + PropertyProvider propProvider = (name, def) -> def; + + assertFalse(MtlsUtils.useMtlsClientCertificate(envProvider, propProvider)); + assertNull(MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + } + + @Test + void useMtlsClientCertificate_falseEvenWhenWorkloadCertsExist_returnsFalse() throws IOException { + Path certFile = tempDir.resolve("cert.pem"); + Path keyFile = tempDir.resolve("key.pem"); + Files.write(certFile, "dummy cert".getBytes()); + Files.write(keyFile, "dummy key".getBytes()); + + Path configFile = tempDir.resolve("config.json"); + String configJson = + String.format( + "{\"cert_configs\":{\"workload\":{\"cert_path\":\"%s\",\"key_path\":\"%s\"}}}", + certFile.toString().replace("\\", "\\\\"), keyFile.toString().replace("\\", "\\\\")); + Files.write(configFile, configJson.getBytes()); + + EnvironmentProvider envProvider = + name -> { + if ("GOOGLE_API_USE_CLIENT_CERTIFICATE".equals(name)) return "false"; + if ("GOOGLE_API_CERTIFICATE_CONFIG".equals(name)) return configFile.toString(); + return null; + }; + PropertyProvider propProvider = (name, def) -> def; + + assertFalse(MtlsUtils.useMtlsClientCertificate(envProvider, propProvider)); + assertNull(MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + } + + @Test + void useMtlsClientCertificate_unsetWithNoCertsOnDisk_returnsFalse() { + EnvironmentProvider envProvider = name -> null; + PropertyProvider propProvider = + (name, def) -> { + if ("user.home".equals(name)) return tempDir.toString(); + if ("os.name".equals(name)) return "Linux"; + return def; + }; + + assertFalse(MtlsUtils.useMtlsClientCertificate(envProvider, propProvider)); + assertNull(MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + } + + // --- Explicit GOOGLE_API_CERTIFICATE_CONFIG Tests (Fail Closed) --- + + @Test + void getWorkloadCertPath_explicitConfigMissing_throwsIllegalStateException() { + EnvironmentProvider envProvider = + name -> "GOOGLE_API_CERTIFICATE_CONFIG".equals(name) ? "/nonexistent/config.json" : null; + PropertyProvider propProvider = (name, def) -> def; + + IllegalStateException exception = + assertThrows( + IllegalStateException.class, + () -> MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + assertTrue( + exception + .getMessage() + .contains( + "specified via GOOGLE_API_CERTIFICATE_CONFIG at '/nonexistent/config.json' does not" + + " exist")); + } + + @Test + void getWorkloadCertPath_explicitConfigIsDirectory_throwsIllegalStateException() + throws IOException { + Path configDir = tempDir.resolve("config_dir"); + Files.createDirectory(configDir); + + EnvironmentProvider envProvider = + name -> "GOOGLE_API_CERTIFICATE_CONFIG".equals(name) ? configDir.toString() : null; + PropertyProvider propProvider = (name, def) -> def; + + IllegalStateException exception = + assertThrows( + IllegalStateException.class, + () -> MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + assertTrue( + exception + .getMessage() + .contains( + "Failed to read certificate configuration file specified via" + + " GOOGLE_API_CERTIFICATE_CONFIG")); + } + + @Test + void getWorkloadCertPath_explicitConfigUnreadable_throwsIllegalStateException() + throws IOException { + Path configFile = tempDir.resolve("unreadable_config.json"); + Files.write(configFile, "{}".getBytes()); + File file = configFile.toFile(); + if (file.setReadable(false)) { + try { + EnvironmentProvider envProvider = + name -> "GOOGLE_API_CERTIFICATE_CONFIG".equals(name) ? configFile.toString() : null; + PropertyProvider propProvider = (name, def) -> def; + + IllegalStateException exception = + assertThrows( + IllegalStateException.class, + () -> MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + assertTrue( + exception + .getMessage() + .contains( + "Failed to read certificate configuration file specified via" + + " GOOGLE_API_CERTIFICATE_CONFIG")); + } finally { + file.setReadable(true); + } + } + } + + @Test + void getWorkloadCertPath_explicitConfigMalformedJson_throwsIllegalStateException() + throws IOException { + Path configFile = tempDir.resolve("malformed.json"); + Files.write(configFile, "{ invalid json".getBytes()); + + EnvironmentProvider envProvider = + name -> "GOOGLE_API_CERTIFICATE_CONFIG".equals(name) ? configFile.toString() : null; + PropertyProvider propProvider = (name, def) -> def; + + IllegalStateException exception = + assertThrows( + IllegalStateException.class, + () -> MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + assertTrue( + exception + .getMessage() + .contains( + "specified via GOOGLE_API_CERTIFICATE_CONFIG at '" + + configFile.toString() + + "' is malformed")); + } + + @Test + void getWorkloadCertPath_explicitConfigOnlyEcp_returnsNullSafely() throws IOException { + Path configFile = tempDir.resolve("ecp_config.json"); + Files.write(configFile, "{\"cert_configs\":{\"enterprise_certificates\":{}}}".getBytes()); + + EnvironmentProvider envProvider = + name -> "GOOGLE_API_CERTIFICATE_CONFIG".equals(name) ? configFile.toString() : null; + PropertyProvider propProvider = (name, def) -> def; + + assertNull(MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + } + + @Test + void getWorkloadCertPath_explicitConfigCertFileMissing_throwsIllegalStateException() + throws IOException { + Path keyFile = tempDir.resolve("key.pem"); + Files.write(keyFile, "dummy key".getBytes()); + + Path configFile = tempDir.resolve("config.json"); + String configJson = + String.format( + "{\"cert_configs\":{\"workload\":{\"cert_path\":\"/nonexistent/cert.pem\",\"key_path\":\"%s\"}}}", + keyFile.toString().replace("\\", "\\\\")); + Files.write(configFile, configJson.getBytes()); + + EnvironmentProvider envProvider = + name -> "GOOGLE_API_CERTIFICATE_CONFIG".equals(name) ? configFile.toString() : null; + PropertyProvider propProvider = (name, def) -> def; + + IllegalStateException exception = + assertThrows( + IllegalStateException.class, + () -> MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + assertTrue(exception.getMessage().contains("Failed to read certificate/key file")); + assertTrue( + exception + .getMessage() + .contains("referenced by configuration '" + configFile.toString() + "'")); + } + + @Test + void getWorkloadCertPath_explicitConfigCertFileIsDirectory_throwsIllegalStateException() + throws IOException { + Path certDir = tempDir.resolve("cert_dir"); + Files.createDirectory(certDir); + Path keyFile = tempDir.resolve("key.pem"); + Files.write(keyFile, "dummy key".getBytes()); + + Path configFile = tempDir.resolve("config.json"); + String configJson = + String.format( + "{\"cert_configs\":{\"workload\":{\"cert_path\":\"%s\",\"key_path\":\"%s\"}}}", + certDir.toString().replace("\\", "\\\\"), keyFile.toString().replace("\\", "\\\\")); + Files.write(configFile, configJson.getBytes()); + + EnvironmentProvider envProvider = + name -> "GOOGLE_API_CERTIFICATE_CONFIG".equals(name) ? configFile.toString() : null; + PropertyProvider propProvider = (name, def) -> def; + + IllegalStateException exception = + assertThrows( + IllegalStateException.class, + () -> MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + assertTrue(exception.getMessage().contains("Failed to read certificate/key file")); + } + + @Test + void getWorkloadCertPath_explicitConfigKeyFileMissing_throwsIllegalStateException() + throws IOException { + Path certFile = tempDir.resolve("cert.pem"); + Files.write(certFile, "dummy cert".getBytes()); + + Path configFile = tempDir.resolve("config.json"); + String configJson = + String.format( + "{\"cert_configs\":{\"workload\":{\"cert_path\":\"%s\",\"key_path\":\"/nonexistent/key.pem\"}}}", + certFile.toString().replace("\\", "\\\\")); + Files.write(configFile, configJson.getBytes()); + + EnvironmentProvider envProvider = + name -> "GOOGLE_API_CERTIFICATE_CONFIG".equals(name) ? configFile.toString() : null; + PropertyProvider propProvider = (name, def) -> def; + + IllegalStateException exception = + assertThrows( + IllegalStateException.class, + () -> MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + assertTrue(exception.getMessage().contains("Failed to read certificate/key file")); + assertTrue( + exception + .getMessage() + .contains("referenced by configuration '" + configFile.toString() + "'")); + } + + @Test + void getWorkloadCertPath_explicitConfigValid_returnsCertPath() throws IOException { + Path certFile = tempDir.resolve("cert.pem"); + Path keyFile = tempDir.resolve("key.pem"); + Files.write(certFile, "dummy cert".getBytes()); + Files.write(keyFile, "dummy key".getBytes()); + + Path configFile = tempDir.resolve("config.json"); + String configJson = + String.format( + "{\"cert_configs\":{\"workload\":{\"cert_path\":\"%s\",\"key_path\":\"%s\"}}}", + certFile.toString().replace("\\", "\\\\"), keyFile.toString().replace("\\", "\\\\")); + Files.write(configFile, configJson.getBytes()); + + EnvironmentProvider envProvider = + name -> "GOOGLE_API_CERTIFICATE_CONFIG".equals(name) ? configFile.toString() : null; + PropertyProvider propProvider = (name, def) -> def; + + assertTrue(MtlsUtils.useMtlsClientCertificate(envProvider, propProvider)); + assertEquals(certFile.toString(), MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + } + + // --- Implicit / Default gcloud Config Tests --- + + @Test + void getWorkloadCertPath_defaultConfigMissing_returnsNullSafely() { + EnvironmentProvider envProvider = name -> null; + PropertyProvider propProvider = + (name, def) -> { + if ("user.home".equals(name)) return tempDir.toString(); + if ("os.name".equals(name)) return "Linux"; + return def; + }; + + assertNull(MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + } + + @Test + void getWorkloadCertPath_defaultConfigIsDirectory_throwsIllegalStateException() + throws IOException { + Path gcloudDir = tempDir.resolve(".config/gcloud"); + Files.createDirectories(gcloudDir); + Path defaultConfigFile = gcloudDir.resolve("certificate_config.json"); + Files.createDirectory(defaultConfigFile); + + EnvironmentProvider envProvider = name -> null; + PropertyProvider propProvider = + (name, def) -> { + if ("user.home".equals(name)) return tempDir.toString(); + if ("os.name".equals(name)) return "Linux"; + return def; + }; + + IllegalStateException exception = + assertThrows( + IllegalStateException.class, + () -> MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + assertTrue( + exception + .getMessage() + .contains( + "Default certificate configuration file at '" + + defaultConfigFile.toFile().getAbsolutePath() + + "' exists but could not be read")); + } + + @Test + void getWorkloadCertPath_defaultConfigMalformedJson_throwsIllegalStateException() + throws IOException { + Path gcloudDir = tempDir.resolve(".config/gcloud"); + Files.createDirectories(gcloudDir); + Path defaultConfigFile = gcloudDir.resolve("certificate_config.json"); + Files.write(defaultConfigFile, "{ malformed json".getBytes()); + + EnvironmentProvider envProvider = name -> null; + PropertyProvider propProvider = + (name, def) -> { + if ("user.home".equals(name)) return tempDir.toString(); + if ("os.name".equals(name)) return "Linux"; + return def; + }; + + IllegalStateException exception = + assertThrows( + IllegalStateException.class, + () -> MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + assertTrue( + exception + .getMessage() + .contains( + "Default certificate configuration file at '" + + defaultConfigFile.toFile().getAbsolutePath() + + "' is malformed")); + } + + @Test + void getWorkloadCertPath_defaultConfigOnlyEcp_returnsNullSafely() throws IOException { + Path gcloudDir = tempDir.resolve(".config/gcloud"); + Files.createDirectories(gcloudDir); + Path defaultConfigFile = gcloudDir.resolve("certificate_config.json"); + Files.write( + defaultConfigFile, + "{\"cert_configs\":{\"enterprise_certificates\":{\"libs\":[]}}}".getBytes()); + + EnvironmentProvider envProvider = name -> null; + PropertyProvider propProvider = + (name, def) -> { + if ("user.home".equals(name)) return tempDir.toString(); + if ("os.name".equals(name)) return "Linux"; + return def; + }; + + assertNull(MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + } + + @Test + void getWorkloadCertPath_defaultConfigCertFileMissing_throwsIllegalStateException() + throws IOException { + Path keyFile = tempDir.resolve("key.pem"); + Files.write(keyFile, "dummy key".getBytes()); + + Path gcloudDir = tempDir.resolve(".config/gcloud"); + Files.createDirectories(gcloudDir); + Path defaultConfigFile = gcloudDir.resolve("certificate_config.json"); + String configJson = + String.format( + "{\"cert_configs\":{\"workload\":{\"cert_path\":\"/nonexistent/cert.pem\",\"key_path\":\"%s\"}}}", + keyFile.toString().replace("\\", "\\\\")); + Files.write(defaultConfigFile, configJson.getBytes()); + + EnvironmentProvider envProvider = name -> null; + PropertyProvider propProvider = + (name, def) -> { + if ("user.home".equals(name)) return tempDir.toString(); + if ("os.name".equals(name)) return "Linux"; + return def; + }; + + IllegalStateException exception = + assertThrows( + IllegalStateException.class, + () -> MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + assertTrue(exception.getMessage().contains("Failed to read certificate/key file")); + assertTrue( + exception + .getMessage() + .contains( + "referenced by default configuration '" + + defaultConfigFile.toFile().getAbsolutePath() + + "'")); + } + + @Test + void getWorkloadCertPath_defaultConfigValid_returnsCertPath() throws IOException { + Path certFile = tempDir.resolve("cert.pem"); + Path keyFile = tempDir.resolve("key.pem"); + Files.write(certFile, "dummy cert".getBytes()); + Files.write(keyFile, "dummy key".getBytes()); + + Path gcloudDir = tempDir.resolve(".config/gcloud"); + Files.createDirectories(gcloudDir); + Path defaultConfigFile = gcloudDir.resolve("certificate_config.json"); + String configJson = + String.format( + "{\"cert_configs\":{\"workload\":{\"cert_path\":\"%s\",\"key_path\":\"%s\"}}}", + certFile.toString().replace("\\", "\\\\"), keyFile.toString().replace("\\", "\\\\")); + Files.write(defaultConfigFile, configJson.getBytes()); + + EnvironmentProvider envProvider = name -> null; + PropertyProvider propProvider = + (name, def) -> { + if ("user.home".equals(name)) return tempDir.toString(); + if ("os.name".equals(name)) return "Linux"; + return def; + }; + + assertTrue(MtlsUtils.useMtlsClientCertificate(envProvider, propProvider)); + assertEquals(certFile.toString(), MtlsUtils.getWorkloadCertPath(envProvider, propProvider)); + } + + // --- General Helpers & Stubs Tests --- + + @Test + void getCertificateFingerprint_validFile_returnsSha256() throws IOException { + Path file = tempDir.resolve("test.crt"); + Files.write(file, "hello world".getBytes()); + + String fingerprint = MtlsUtils.getCertificateFingerprint(file.toString()); + assertNotNull(fingerprint); + assertEquals(64, fingerprint.length()); // SHA-256 hex string length + assertEquals("b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9", fingerprint); + } + + @Test + void getCertificateFingerprint_emptyFile_returnsValidSha256() throws IOException { + Path emptyFile = tempDir.resolve("empty.crt"); + Files.write(emptyFile, new byte[0]); + + String fingerprint = MtlsUtils.getCertificateFingerprint(emptyFile.toString()); + assertEquals("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", fingerprint); + } + + @Test + void getCertificateFingerprint_invalidOrNull_returnsNull() { + assertNull(MtlsUtils.getCertificateFingerprint(null)); + assertNull(MtlsUtils.getCertificateFingerprint("/nonexistent/file.crt")); + assertNull(MtlsUtils.getCertificateFingerprint(tempDir.toString())); // Directory + } } diff --git a/sdk-platform-java/gax-java/gax-grpc/pom.xml b/sdk-platform-java/gax-java/gax-grpc/pom.xml index d717ac0944b5..0abd59208e3c 100644 --- a/sdk-platform-java/gax-java/gax-grpc/pom.xml +++ b/sdk-platform-java/gax-java/gax-grpc/pom.xml @@ -162,8 +162,7 @@ maven-surefire-plugin - !InstantiatingGrpcChannelProviderTest#testLogDirectPathMisconfig_AttemptDirectPathNotSetAndAttemptDirectPathXdsSetViaEnv_warns,!InstantiatingGrpcChannelProviderTest#canUseDirectPath_directPathEnvVarNotSet_attemptDirectPathIsTrue,InstantiatingGrpcChannelProviderTest#testLogDirectPathMisconfigWrongCredential - + !InstantiatingGrpcChannelProviderTest#testLogDirectPathMisconfig_AttemptDirectPathNotSetAndAttemptDirectPathXdsSetViaEnv_warns diff --git a/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/ChannelPool.java b/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/ChannelPool.java index fab73a55dccf..8004593a2673 100644 --- a/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/ChannelPool.java +++ b/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/ChannelPool.java @@ -31,8 +31,10 @@ import com.google.api.core.InternalApi; import com.google.api.gax.core.FixedExecutorProvider; +import com.google.api.gax.rpc.mtls.CertificateRotationTracker; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; +import com.google.common.base.Strings; import com.google.common.collect.ImmutableList; import io.grpc.CallOptions; import io.grpc.Channel; @@ -53,9 +55,11 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; import java.util.logging.Level; import java.util.logging.Logger; +import javax.annotation.concurrent.GuardedBy; import org.jspecify.annotations.NullMarked; import org.jspecify.annotations.Nullable; @@ -72,18 +76,28 @@ @NullMarked class ChannelPool extends ManagedChannel { static final String CHANNEL_POOL_CONSECUTIVE_RESIZING_WARNING = - "The gRPC ChannelPool used in the client has been flagged to be repeatedly resizing (5+ times). See https://github.com/googleapis/google-cloud-java/blob/main/docs/grpc_channel_pool_guide.md for more information about this behavior."; + "The gRPC ChannelPool used in the client has been flagged to be repeatedly resizing (5+" + + " times). See" + + " https://github.com/googleapis/google-cloud-java/blob/main/docs/grpc_channel_pool_guide.md" + + " for more information about this behavior."; @VisibleForTesting static final Logger LOG = Logger.getLogger(ChannelPool.class.getName()); private static final java.time.Duration REFRESH_PERIOD = java.time.Duration.ofMinutes(50); private final ChannelPoolSettings settings; private final ChannelFactory channelFactory; private final FixedExecutorProvider backgroundExecutorProvider; + private final String workloadCertPath; private @Nullable ScheduledFuture refreshFuture = null; private @Nullable ScheduledFuture resizeFuture = null; + private final CertificateRotationTracker rotationTracker; private final Object entryWriteLock = new Object(); + + @GuardedBy("entryWriteLock") + private boolean isShutdown = false; + + private final AtomicLong generation = new AtomicLong(0); @VisibleForTesting final AtomicReference> entries = new AtomicReference<>(); private final AtomicInteger indexTicker = new AtomicInteger(); private final String authority; @@ -100,14 +114,15 @@ class ChannelPool extends ManagedChannel { static ChannelPool create( ChannelPoolSettings settings, ChannelFactory channelFactory, - @Nullable ScheduledExecutorService backgroundExecutor) + @Nullable ScheduledExecutorService backgroundExecutor, + @Nullable String workloadCertPath) throws IOException { FixedExecutorProvider executorProvider = backgroundExecutor == null ? FixedExecutorProvider.create(Executors.newSingleThreadScheduledExecutor(), true) : FixedExecutorProvider.create(backgroundExecutor, false); - return new ChannelPool(settings, channelFactory, executorProvider); + return new ChannelPool(settings, channelFactory, executorProvider, workloadCertPath); } /** @@ -121,11 +136,14 @@ static ChannelPool create( ChannelPool( ChannelPoolSettings settings, ChannelFactory channelFactory, - FixedExecutorProvider executorProvider) + FixedExecutorProvider executorProvider, + @Nullable String workloadCertPath) throws IOException { this.settings = settings; this.channelFactory = channelFactory; this.backgroundExecutorProvider = executorProvider; + this.workloadCertPath = workloadCertPath; + this.rotationTracker = new CertificateRotationTracker(workloadCertPath); ImmutableList.Builder initialListBuilder = ImmutableList.builder(); @@ -185,19 +203,22 @@ Channel getChannel(int affinity) { public ManagedChannel shutdown() { LOG.fine("Initiating graceful shutdown due to explicit request"); - // Resize and refresh tasks can block on channel priming. We don't need - // to wait for the channels to be ready since we're shutting down the - // pool. Allowing interrupt to speed it up. - if (resizeFuture != null) { - resizeFuture.cancel(true); - } - if (refreshFuture != null) { - refreshFuture.cancel(true); - } + synchronized (entryWriteLock) { + isShutdown = true; + // Resize and refresh tasks can block on channel priming. We don't need + // to wait for the channels to be ready since we're shutting down the + // pool. Allowing interrupt to speed it up. + if (resizeFuture != null) { + resizeFuture.cancel(true); + } + if (refreshFuture != null) { + refreshFuture.cancel(true); + } - List localEntries = entries.get(); - for (Entry entry : localEntries) { - entry.channel.shutdown(); + List localEntries = entries.get(); + for (Entry entry : localEntries) { + entry.channel.shutdown(); + } } if (backgroundExecutorProvider.shouldAutoClose()) { @@ -210,6 +231,11 @@ public ManagedChannel shutdown() { /** {@inheritDoc} */ @Override public boolean isShutdown() { + synchronized (entryWriteLock) { + if (isShutdown) { + return true; + } + } List localEntries = entries.get(); for (Entry entry : localEntries) { if (!entry.channel.isShutdown()) { @@ -236,16 +262,19 @@ public boolean isTerminated() { public ManagedChannel shutdownNow() { LOG.fine("Initiating immediate shutdown due to explicit request"); - if (resizeFuture != null) { - resizeFuture.cancel(true); - } - if (refreshFuture != null) { - refreshFuture.cancel(true); - } + synchronized (entryWriteLock) { + isShutdown = true; + if (resizeFuture != null) { + resizeFuture.cancel(true); + } + if (refreshFuture != null) { + refreshFuture.cancel(true); + } - List localEntries = entries.get(); - for (Entry entry : localEntries) { - entry.channel.shutdownNow(); + List localEntries = entries.get(); + for (Entry entry : localEntries) { + entry.channel.shutdownNow(); + } } if (backgroundExecutorProvider.shouldAutoClose()) { @@ -419,14 +448,38 @@ private void expand(int desiredSize) { entries.set(newEntries.build()); } + /** + * Periodically refreshes all channels when {@link + * ChannelPoolSettings#isPreemptiveRefreshEnabled()} is enabled (to mitigate hourly GFE + * disconnects). This applies to all channels even when {@code workloadCertPath == null}. If + * {@code workloadCertPath} is configured, also updates the tracked certificate fingerprint on + * success (or skips if the certificate file is currently unreadable or mid-write on disk). + */ private void refreshSafely() { try { - refresh(); + synchronized (entryWriteLock) { + String currentDiskFingerprint = rotationTracker.readDiskFingerprint(); + if (workloadCertPath != null && currentDiskFingerprint.isEmpty()) { + return; + } + if (refreshAll() && !currentDiskFingerprint.isEmpty()) { + rotationTracker.markRefreshed(currentDiskFingerprint); + } + } } catch (Exception e) { - LOG.log(Level.WARNING, "Failed to pre-emptively refresh channnels", e); + LOG.log(Level.WARNING, "Failed to pre-emptively refresh channels", e); } } + @VisibleForTesting + void invalidateDiskFingerprintCache() { + rotationTracker.invalidateCache(); + } + + boolean shouldRefresh() { + return rotationTracker.shouldRefresh(); + } + /** * Replace all of the channels in the channel pool with fresh ones. This is meant to mitigate the * hourly GFE disconnects by giving clients the ability to prime the channel on reconnect. @@ -443,28 +496,104 @@ void refresh() { // - then thread2 will shut down channel that thread1 will put back into circulation (after it // replaces the list) synchronized (entryWriteLock) { - LOG.fine("Refreshing all channels"); + if (isShutdown) { + return; + } + if (workloadCertPath == null) { + refreshAll(); + return; + } + String currentDiskFingerprint = rotationTracker.readDiskFingerprint(); + if (currentDiskFingerprint.isEmpty()) { + return; + } + + // Double-check fingerprint inside the lock + if (rotationTracker.isAlreadyActive(currentDiskFingerprint)) { + LOG.fine( + "Channel pool was already refreshed by a concurrent thread, skipping duplicate" + + " refresh"); + return; + } + + if (refreshAll()) { + rotationTracker.markRefreshed(currentDiskFingerprint); + } + } + } + + @InternalApi("Visible for testing") + @Nullable String getWorkloadCertPath() { + return workloadCertPath; + } + + @InternalApi("Visible for testing") + boolean refreshAll() { + synchronized (entryWriteLock) { + if (isShutdown) { + return false; + } + String activeFingerprint = rotationTracker.getActiveCertFingerprint(); + LOG.fine( + "Refreshing all channels" + + (Strings.isNullOrEmpty(activeFingerprint) + ? "" + : " with certificate fingerprint: " + activeFingerprint)); ArrayList newEntries = new ArrayList<>(entries.get()); + boolean anyCreated = false; + boolean allCreated = !newEntries.isEmpty(); + List createdEntries = new ArrayList<>(); - for (int i = 0; i < newEntries.size(); i++) { - try { - newEntries.set(i, new Entry(channelFactory.createSingleChannel())); - } catch (IOException e) { - LOG.log(Level.WARNING, "Failed to refresh channel, leaving old channel", e); + try { + for (int i = 0; i < newEntries.size(); i++) { + try { + Entry newEntry = new Entry(channelFactory.createSingleChannel()); + createdEntries.add(newEntry); + newEntries.set(i, newEntry); + anyCreated = true; + } catch (Exception e) { + allCreated = false; + LOG.log(Level.WARNING, "Failed to refresh channel, leaving old channel", e); + } } - } - ImmutableList replacedEntries = entries.getAndSet(ImmutableList.copyOf(newEntries)); + if (!anyCreated) { + return false; + } - // Shutdown the channels that were cycled out. - for (Entry e : replacedEntries) { - if (!newEntries.contains(e)) { + ImmutableList replacedEntries = entries.getAndSet(ImmutableList.copyOf(newEntries)); + createdEntries.clear(); // Ownership transferred to pool + + // Shutdown the channels that were cycled out. + for (Entry e : replacedEntries) { + if (!newEntries.contains(e)) { + e.requestShutdown(); + } + } + generation.incrementAndGet(); + return allCreated; + } finally { + // If an Error aborted before getAndSet, shut down newly created channels so they don't leak + for (Entry e : createdEntries) { e.requestShutdown(); } } } } + /** + * Returns the current channel pool generation counter. + * + *

The generation is a monotonically increasing counter incremented each time {@link + * #refreshAll()} replaces the channels in the pool. Retry loops ({@code AttemptCallable} and + * {@code ServerStreamingAttemptCallable}) snapshot the generation before starting an RPC attempt + * and compare it after an {@code UNAUTHENTICATED} failure to determine whether the pool rotated + * to a new certificate generation during or after the attempt. + */ + long getGeneration() { + return generation.get(); + } + /** * Get and retain a Channel Entry. The returned Entry will have its rpc count incremented, * preventing it from getting recycled. @@ -616,17 +745,30 @@ public ClientCall newCall( MethodDescriptor methodDescriptor, CallOptions callOptions) { Entry entry = getRetainedEntry(affinity); - - return new ReleasingClientCall<>(entry.channel.newCall(methodDescriptor, callOptions), entry); + try { + return new ReleasingClientCall<>( + entry.channel.newCall(methodDescriptor, callOptions), entry); + } catch (Throwable t) { + entry.release(); + throw t; + } } } - /** ClientCall wrapper that makes sure to decrement the outstanding RPC count on completion. */ + /** + * ClientCall wrapper that makes sure to decrement the outstanding RPC count on completion. + * + *

Contract: Exactly one call to {@link #start(Listener, Metadata)} is required to balance + * reference counts. Early cancellation before {@code start()} is recorded and safely decrements + * the reference count when {@code start()} is subsequently invoked. + */ static class ReleasingClientCall extends SimpleForwardingClientCall { - private @Nullable CancellationException cancellationException; + private final Object callLock = new Object(); + private volatile @Nullable CancellationException cancellationException; final Entry entry; private final AtomicBoolean wasClosed = new AtomicBoolean(); private final AtomicBoolean wasReleased = new AtomicBoolean(); + private final AtomicBoolean wasStarted = new AtomicBoolean(); public ReleasingClientCall(ClientCall delegate, Entry entry) { super(delegate); @@ -635,51 +777,81 @@ public ReleasingClientCall(ClientCall delegate, Entry entry) { @Override public void start(Listener responseListener, Metadata headers) { - if (cancellationException != null) { - throw new IllegalStateException("Call is already cancelled", cancellationException); - } - try { - super.start( - new SimpleForwardingClientCallListener(responseListener) { - @Override - public void onClose(Status status, Metadata trailers) { - if (!wasClosed.compareAndSet(false, true)) { - LOG.log( - Level.WARNING, - "Call is being closed more than once. Please make sure that onClose() is not being manually called."); - return; - } - try { - super.onClose(status, trailers); - } finally { - if (wasReleased.compareAndSet(false, true)) { - entry.release(); - } else { + synchronized (callLock) { + if (!wasStarted.compareAndSet(false, true)) { + throw new IllegalStateException("Call is already started"); + } + if (cancellationException != null) { + if (wasReleased.compareAndSet(false, true)) { + entry.release(); + } + throw new IllegalStateException("Call is already cancelled", cancellationException); + } + try { + super.start( + new SimpleForwardingClientCallListener(responseListener) { + @Override + public void onClose(Status status, Metadata trailers) { + if (!wasClosed.compareAndSet(false, true)) { LOG.log( Level.WARNING, - "Entry was released before the call is closed. This may be due to an exception on start of the call."); + "Call is being closed more than once. Please make sure that onClose() is" + + " not being manually called."); + return; + } + try { + super.onClose(status, trailers); + } finally { + if (wasReleased.compareAndSet(false, true)) { + entry.release(); + } else { + LOG.log( + Level.WARNING, + "Entry was released before the call is closed. This may be due to an" + + " exception on start of the call."); + } } } - } - }, - headers); - } catch (Exception e) { - // In case start failed, make sure to release - if (wasReleased.compareAndSet(false, true)) { - entry.release(); - } else { - LOG.log( - Level.WARNING, - "The entry is already released. This indicates that onClose() has already been called previously"); + }, + headers); + } catch (Throwable t) { + // In case start failed, make sure to release + if (wasReleased.compareAndSet(false, true)) { + entry.release(); + } else { + LOG.log( + Level.WARNING, + "The entry is already released. This indicates that onClose() has already been" + + " called previously"); + } + throw t; } - throw e; } } @Override public void cancel(@Nullable String message, @Nullable Throwable cause) { - this.cancellationException = new CancellationException(message); - super.cancel(message, cause); + boolean releaseImmediately = false; + try { + synchronized (callLock) { + this.cancellationException = new CancellationException(message); + if (!wasStarted.get()) { + releaseImmediately = true; + } + if (delegate() != null) { + super.cancel(message, cause); + } + } + } catch (Throwable t) { + if (!wasStarted.get()) { + releaseImmediately = true; + } + throw t; + } finally { + if (releaseImmediately && wasReleased.compareAndSet(false, true)) { + entry.release(); + } + } } } } diff --git a/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcCallContext.java b/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcCallContext.java index 23f56c5f8951..c0288a1ae382 100644 --- a/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcCallContext.java +++ b/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcCallContext.java @@ -99,6 +99,7 @@ public final class GrpcCallContext implements ApiCallContext { private final ApiCallContextOptions options; private final EndpointContext endpointContext; private final boolean isDirectPath; + @Nullable private final TransportChannel transportChannel; /** Returns an empty instance with a null channel and default {@link CallOptions}. */ public static GrpcCallContext createDefault() { @@ -115,7 +116,8 @@ public static GrpcCallContext createDefault() { null, null, null, - false); + false, + null); } /** Returns an instance with the given channel and {@link CallOptions}. */ @@ -133,7 +135,8 @@ public static GrpcCallContext of(Channel channel, CallOptions callOptions) { null, null, null, - false); + false, + null); } private GrpcCallContext( @@ -149,7 +152,8 @@ private GrpcCallContext( @Nullable RetrySettings retrySettings, @Nullable Set retryableCodes, @Nullable EndpointContext endpointContext, - boolean isDirectPath) { + boolean isDirectPath, + @Nullable TransportChannel transportChannel) { this.channel = channel; this.credentials = credentials; Preconditions.checkNotNull(callOptions); @@ -169,6 +173,7 @@ private GrpcCallContext( this.endpointContext = endpointContext == null ? EndpointContext.getDefaultInstance() : endpointContext; this.isDirectPath = isDirectPath; + this.transportChannel = transportChannel; } /** @@ -210,7 +215,13 @@ public GrpcCallContext withCredentials(Credentials newCredentials) { retrySettings, retryableCodes, endpointContext, - isDirectPath); + isDirectPath, + transportChannel); + } + + @Override + public TransportChannel getTransportChannel() { + return transportChannel; } @Override @@ -234,7 +245,8 @@ public GrpcCallContext withTransportChannel(TransportChannel inputChannel) { retrySettings, retryableCodes, endpointContext, - transportChannel.isDirectPath()); + transportChannel.isDirectPath(), + inputChannel); } @Override @@ -253,7 +265,8 @@ public GrpcCallContext withEndpointContext(EndpointContext endpointContext) { retrySettings, retryableCodes, endpointContext, - isDirectPath); + isDirectPath, + transportChannel); } /** This method is obsolete. Use {@link #withTimeoutDuration(java.time.Duration)} instead. */ @@ -271,7 +284,7 @@ public GrpcCallContext withTimeoutDuration(java.time.@Nullable Duration timeout) } // Prevent expanding timeouts - if (timeout != null && this.timeout != null && this.timeout.compareTo(timeout) <= 0) { + if (this.timeout != null && (timeout == null || this.timeout.compareTo(timeout) <= 0)) { return this; } @@ -288,7 +301,8 @@ public GrpcCallContext withTimeoutDuration(java.time.@Nullable Duration timeout) retrySettings, retryableCodes, endpointContext, - isDirectPath); + isDirectPath, + transportChannel); } @Override @@ -334,7 +348,8 @@ public GrpcCallContext withStreamWaitTimeoutDuration( retrySettings, retryableCodes, endpointContext, - isDirectPath); + isDirectPath, + transportChannel); } /** @@ -369,7 +384,8 @@ public GrpcCallContext withStreamIdleTimeoutDuration( retrySettings, retryableCodes, endpointContext, - isDirectPath); + isDirectPath, + transportChannel); } @BetaApi("The surface for channel affinity is not stable yet and may change in the future.") @@ -387,7 +403,8 @@ public GrpcCallContext withChannelAffinity(@Nullable Integer affinity) { retrySettings, retryableCodes, endpointContext, - isDirectPath); + isDirectPath, + transportChannel); } @BetaApi("The surface for extra headers is not stable yet and may change in the future.") @@ -409,7 +426,8 @@ public GrpcCallContext withExtraHeaders(Map> extraHeaders) retrySettings, retryableCodes, endpointContext, - isDirectPath); + isDirectPath, + transportChannel); } @Override @@ -432,7 +450,8 @@ public GrpcCallContext withRetrySettings(RetrySettings retrySettings) { retrySettings, retryableCodes, endpointContext, - isDirectPath); + isDirectPath, + transportChannel); } @Override @@ -455,7 +474,8 @@ public GrpcCallContext withRetryableCodes(Set retryableCodes) { retrySettings, retryableCodes, endpointContext, - isDirectPath); + isDirectPath, + transportChannel); } @Override @@ -542,6 +562,12 @@ public ApiCallContext merge(ApiCallContext inputCallContext) { newCallOptions = newCallOptions.withOption(TRACER_KEY, newTracer); } + TransportChannel newTransportChannel = grpcCallContext.transportChannel; + if (newTransportChannel == null + && (grpcCallContext.channel == null || grpcCallContext.channel.equals(channel))) { + newTransportChannel = transportChannel; + } + // The EndpointContext is not updated as there should be no reason for a user // to update this. return new GrpcCallContext( @@ -557,7 +583,8 @@ public ApiCallContext merge(ApiCallContext inputCallContext) { newRetrySettings, newRetryableCodes, endpointContext, - newIsDirectPath); + newIsDirectPath, + newTransportChannel); } /** The {@link Channel} set on this context. */ @@ -635,7 +662,8 @@ public GrpcCallContext withChannel(@Nullable Channel newChannel) { retrySettings, retryableCodes, endpointContext, - isDirectPath); + isDirectPath, + (newChannel != null && newChannel.equals(channel)) ? transportChannel : null); } /** Returns a new instance with the call options set to the given call options. */ @@ -653,7 +681,8 @@ public GrpcCallContext withCallOptions(CallOptions newCallOptions) { retrySettings, retryableCodes, endpointContext, - isDirectPath); + isDirectPath, + transportChannel); } public GrpcCallContext withRequestParamsDynamicHeaderOption(String requestParams) { @@ -698,7 +727,8 @@ public GrpcCallContext withOption(Key key, T value) { retrySettings, retryableCodes, endpointContext, - isDirectPath); + isDirectPath, + transportChannel); } /** {@inheritDoc} */ @@ -759,7 +789,8 @@ public int hashCode() { options, retrySettings, retryableCodes, - endpointContext); + endpointContext, + transportChannel); } @Override @@ -783,7 +814,8 @@ public boolean equals(@Nullable Object o) { && Objects.equals(options, that.options) && Objects.equals(retrySettings, that.retrySettings) && Objects.equals(retryableCodes, that.retryableCodes) - && Objects.equals(endpointContext, that.endpointContext); + && Objects.equals(endpointContext, that.endpointContext) + && Objects.equals(transportChannel, that.transportChannel); } Metadata getMetadata() { diff --git a/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcTransportChannel.java b/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcTransportChannel.java index e0a520facb17..63180a8149f6 100644 --- a/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcTransportChannel.java +++ b/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcTransportChannel.java @@ -68,6 +68,32 @@ public Channel getChannel() { return getManagedChannel(); } + @Override + public void refresh() { + Channel channel = getChannel(); + if (channel instanceof ChannelPool) { + ((ChannelPool) channel).refresh(); + } + } + + @Override + public boolean shouldRefresh() { + Channel channel = getChannel(); + if (channel instanceof ChannelPool) { + return ((ChannelPool) channel).shouldRefresh(); + } + return false; + } + + @Override + public long getGeneration() { + Channel channel = getChannel(); + if (channel instanceof ChannelPool) { + return ((ChannelPool) channel).getGeneration(); + } + return 0; + } + @Override public void shutdown() { getManagedChannel().shutdown(); diff --git a/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java b/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java index ac42396f006a..8a1c1fc8b0bb 100644 --- a/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java +++ b/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java @@ -401,12 +401,19 @@ public TransportChannel getTransportChannel() throws IOException { } private TransportChannel createChannel() throws IOException { + String workloadCertPath = + !this.canUseDirectPath() + && mtlsProvider != null + && certificateBasedAccess.useMtlsClientCertificate() + ? certificateBasedAccess.getWorkloadCertPath() + : null; return GrpcTransportChannel.newBuilder() .setManagedChannel( ChannelPool.create( channelPoolSettings, InstantiatingGrpcChannelProvider.this::createSingleChannel, - backgroundExecutor)) + backgroundExecutor, + workloadCertPath)) .setDirectPath(this.canUseDirectPath()) .build(); } @@ -465,8 +472,9 @@ private void logDirectPathMisconfig() { level, "Env var " + DIRECT_PATH_ENV_ENABLE_XDS - + " was found and set to TRUE, but DirectPath was not enabled for this client. If this is intended for " - + "this client, please note that this is a misconfiguration and set the attemptDirectPath option as well."); + + " was found and set to TRUE, but DirectPath was not enabled for this client. If" + + " this is intended for this client, please note that this is a misconfiguration" + + " and set the attemptDirectPath option as well."); } // Case 2: Direct Path xDS was enabled via Builder. Direct Path Traffic Director must be set // (enabled with `setAttemptDirectPath(true)`) along with xDS. @@ -474,7 +482,9 @@ private void logDirectPathMisconfig() { else if (isDirectPathXdsEnabledViaBuilderOption()) { LOG.log( level, - "DirectPath is misconfigured. The DirectPath XDS option was set, but the attemptDirectPath option was not. Please set both the attemptDirectPath and attemptDirectPathXds options."); + "DirectPath is misconfigured. The DirectPath XDS option was set, but the" + + " attemptDirectPath option was not. Please set both the attemptDirectPath and" + + " attemptDirectPathXds options."); } } else { // Case 3: credential is not correctly set @@ -666,7 +676,8 @@ ChannelCredentials createS2ASecuredChannelCredentials() { // Fallback to plaintext connection to S2A. LOG.log( Level.INFO, - "Cannot establish an mTLS connection to S2A because autoconfig endpoint did not return a mtls address to reach S2A."); + "Cannot establish an mTLS connection to S2A because autoconfig endpoint did not" + + " return a mtls address to reach S2A."); s2aChannelCredentials = createPlaintextToS2AChannelCredentials(plaintextAddress); return s2aChannelCredentials; } @@ -685,7 +696,9 @@ ChannelCredentials createS2ASecuredChannelCredentials() { // Fallback to plaintext-to-S2A connection on error. LOG.log( Level.WARNING, - "Cannot establish an mTLS connection to S2A due to error creating MTLS to MDS TlsChannelCredentials credentials, falling back to plaintext connection to S2A: " + "Cannot establish an mTLS connection to S2A due to error creating MTLS to MDS" + + " TlsChannelCredentials credentials, falling back to plaintext connection to" + + " S2A: " + ignore.getMessage()); s2aChannelCredentials = createPlaintextToS2AChannelCredentials(plaintextAddress); return s2aChannelCredentials; @@ -755,6 +768,8 @@ public ManagedChannelBuilder createChannelBuilder() throws IOException { if (channelCredentials != null) { // Create the channel using channel credentials created via DCA. builder = Grpc.newChannelBuilder(endpoint, channelCredentials); + } else if (mtlsProvider != null && certificateBasedAccess.useMtlsClientCertificate()) { + throw new IOException("Failed to initialize mTLS channel credentials"); } else { // Could not create channel credentials via DCA. In accordance with // https://google.aip.dev/auth/4115, if credentials not available through @@ -1403,7 +1418,8 @@ public InstantiatingGrpcChannelProvider build() { "DefaultMtlsProviderFactory encountered unexpected IOException: " + e.getMessage()); LOG.log( Level.WARNING, - "mTLS configuration was detected on the device, but mTLS failed to initialize. Falling back to non-mTLS channel."); + "mTLS configuration was detected on the device, but mTLS failed to initialize." + + " Falling back to non-mTLS channel."); } } } diff --git a/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/ChannelPoolTest.java b/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/ChannelPoolTest.java index 5bfdc7754759..9c408f2d128c 100644 --- a/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/ChannelPoolTest.java +++ b/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/ChannelPoolTest.java @@ -81,13 +81,18 @@ class ChannelPoolTest { private static final int DEFAULT_AWAIT_TERMINATION_SEC = 10; private ChannelPool pool; + private java.nio.file.Path tempCert; @AfterEach - void cleanup() throws InterruptedException { + void cleanup() throws InterruptedException, IOException { if (pool != null) { pool.shutdown(); pool.awaitTermination(DEFAULT_AWAIT_TERMINATION_SEC, TimeUnit.SECONDS); } + if (tempCert != null) { + java.nio.file.Files.deleteIfExists(tempCert); + tempCert = null; + } } @Test @@ -101,6 +106,7 @@ void testAuthority() throws IOException { ChannelPool.create( ChannelPoolSettings.staticallySized(2), new FakeChannelFactory(Arrays.asList(sub1, sub2)), + null, null); assertThat(pool.authority()).isEqualTo("myAuth"); } @@ -117,6 +123,7 @@ void testRoundRobin() throws IOException { ChannelPool.create( ChannelPoolSettings.staticallySized(channels.size()), new FakeChannelFactory(channels), + null, null); verifyTargetChannel(pool, channels, sub1); @@ -195,6 +202,7 @@ void ensureEvenDistribution() throws InterruptedException, IOException { ChannelPool.create( ChannelPoolSettings.staticallySized(numChannels), new FakeChannelFactory(Arrays.asList(channels)), + null, null); int numThreads = 20; @@ -233,6 +241,7 @@ void channelPrimerShouldCallPoolConstruction() throws IOException { .setPreemptiveRefreshEnabled(true) .build(), new FakeChannelFactory(Arrays.asList(channel1, channel2), mockChannelPrimer), + null, null); Mockito.verify(mockChannelPrimer, Mockito.times(2)) .primeChannel(Mockito.any(ManagedChannel.class)); @@ -273,7 +282,8 @@ void channelPrimerIsCalledPeriodically() throws IOException { .setPreemptiveRefreshEnabled(true) .build(), channelFactory, - provider); + provider, + null); // 1 call during the creation Mockito.verify(mockChannelPrimer, Mockito.times(1)) .primeChannel(Mockito.any(ManagedChannel.class)); @@ -297,7 +307,7 @@ void callShouldCompleteAfterCreation() throws IOException { ManagedChannel replacementChannel = mock(ManagedChannel.class); FakeChannelFactory channelFactory = new FakeChannelFactory(ImmutableList.of(underlyingChannel, replacementChannel)); - pool = ChannelPool.create(ChannelPoolSettings.staticallySized(1), channelFactory, null); + pool = ChannelPool.create(ChannelPoolSettings.staticallySized(1), channelFactory, null, null); // create a mock call when new call comes to the underlying channel MockClientCall mockClientCall = new MockClientCall<>(1, Status.OK); @@ -322,7 +332,7 @@ void callShouldCompleteAfterCreation() throws IOException { ClientCall call = pool.newCall(FakeMethodDescriptor.create(), CallOptions.DEFAULT); - pool.refresh(); + pool.refreshAll(); // shutdown is not called because there is still an outstanding call, even if it hasn't started Mockito.verify(underlyingChannel, Mockito.after(200).never()).shutdown(); @@ -346,7 +356,7 @@ void callShouldCompleteAfterStarted() throws IOException { FakeChannelFactory channelFactory = new FakeChannelFactory(ImmutableList.of(underlyingChannel, replacementChannel)); - pool = ChannelPool.create(ChannelPoolSettings.staticallySized(1), channelFactory, null); + pool = ChannelPool.create(ChannelPoolSettings.staticallySized(1), channelFactory, null, null); // create a mock call when new call comes to the underlying channel MockClientCall mockClientCall = new MockClientCall<>(1, Status.OK); @@ -373,7 +383,7 @@ void callShouldCompleteAfterStarted() throws IOException { // start clientCall call.start(listener, new Metadata()); - pool.refresh(); + pool.refreshAll(); // shutdown is not called because there is still an outstanding call Mockito.verify(underlyingChannel, Mockito.after(200).never()).shutdown(); @@ -391,7 +401,7 @@ void channelShouldShutdown() throws IOException { FakeChannelFactory channelFactory = new FakeChannelFactory(ImmutableList.of(underlyingChannel, replacementChannel)); - pool = ChannelPool.create(ChannelPoolSettings.staticallySized(1), channelFactory, null); + pool = ChannelPool.create(ChannelPoolSettings.staticallySized(1), channelFactory, null, null); // create a mock call when new call comes to the underlying channel MockClientCall mockClientCall = new MockClientCall<>(1, Status.OK); @@ -422,11 +432,269 @@ void channelShouldShutdown() throws IOException { call.sendMessage("message"); // shutdown is not called because it has not been shutdown yet Mockito.verify(underlyingChannel, Mockito.after(200).never()).shutdown(); - pool.refresh(); + pool.refreshAll(); // shutdown is called because the outstanding call has completed Mockito.verify(underlyingChannel, Mockito.atLeastOnce()).shutdown(); } + @Test + void testCancelBeforeStartReleasesChannelEntry() throws IOException { + ManagedChannel underlyingChannel = mock(ManagedChannel.class); + ManagedChannel replacementChannel = mock(ManagedChannel.class); + FakeChannelFactory channelFactory = + new FakeChannelFactory(ImmutableList.of(underlyingChannel, replacementChannel)); + pool = ChannelPool.create(ChannelPoolSettings.staticallySized(1), channelFactory, null, null); + + ClientCall call = + pool.newCall(FakeMethodDescriptor.create(), CallOptions.DEFAULT); + + pool.refreshAll(); + Mockito.verify(underlyingChannel, Mockito.never()).shutdown(); + + call.cancel("Cancelled early", null); + Mockito.verify(underlyingChannel, Mockito.times(1)).shutdown(); + } + + @Test + void channelReactiveMTlsRefreshShouldConditionallySwapChannels() + throws IOException, InterruptedException { + ManagedChannel underlyingChannel1 = Mockito.mock(ManagedChannel.class); + ManagedChannel underlyingChannel2 = Mockito.mock(ManagedChannel.class); + + FakeChannelFactory channelFactory = + new FakeChannelFactory(ImmutableList.of(underlyingChannel1, underlyingChannel2)); + + // Create a temp file to act as the cert + tempCert = java.nio.file.Files.createTempFile("cert", ".pem"); + + java.nio.file.Path clientCert = + java.nio.file.Paths.get("src", "test", "resources", "client_cert.pem"); + java.nio.file.Files.copy( + clientCert, tempCert, java.nio.file.StandardCopyOption.REPLACE_EXISTING); + + ChannelPoolSettings channelPoolSettings = + ChannelPoolSettings.builder().setInitialChannelCount(1).build(); + + pool = ChannelPool.create(channelPoolSettings, channelFactory, null, tempCert.toString()); + + // Initially uses channel1 + pool.newCall(FakeMethodDescriptor.create(), CallOptions.DEFAULT); + Mockito.verify(underlyingChannel1, Mockito.times(1)) + .newCall(Mockito.>any(), Mockito.any(CallOptions.class)); + + // Try a reactive refresh *without* changing the cert content (should no-op) + pool.refresh(); + + // Verify it's STILL channel1 + pool.newCall(FakeMethodDescriptor.create(), CallOptions.DEFAULT); + Mockito.verify(underlyingChannel1, Mockito.times(2)) + .newCall(Mockito.>any(), Mockito.any(CallOptions.class)); + + // The ChannelPool caches fingerprints for 1000ms, wait for it to expire + pool.invalidateDiskFingerprintCache(); + + java.nio.file.Path rootCert = + java.nio.file.Paths.get("src", "test", "resources", "root_cert.pem"); + java.nio.file.Files.copy(rootCert, tempCert, java.nio.file.StandardCopyOption.REPLACE_EXISTING); + + // Try a reactive refresh *with* a changed cert content (should swap channels) + pool.refresh(); + + // Verify it is NOW channel2 + pool.newCall(FakeMethodDescriptor.create(), CallOptions.DEFAULT); + Mockito.verify(underlyingChannel2, Mockito.times(1)) + .newCall(Mockito.>any(), Mockito.any(CallOptions.class)); + } + + @Test + void channelReactiveMTlsRefresh_failedCreationDoesNotMutateFingerprintAndAllowsRetry() + throws IOException { + ManagedChannel channel1 = Mockito.mock(ManagedChannel.class); + ManagedChannel channel2 = Mockito.mock(ManagedChannel.class); + ChannelFactory channelFactory = + Mockito.mock(ChannelFactory.class, Mockito.withSettings().withoutAnnotations()); + + // Initial creation returns channel1, refresh attempt 1 throws IOException, refresh attempt 2 + // returns channel2 + Mockito.when(channelFactory.createSingleChannel()) + .thenReturn(channel1) + .thenThrow(new IOException("Transient channel creation error")) + .thenReturn(channel2); + + tempCert = java.nio.file.Files.createTempFile("cert", ".pem"); + java.nio.file.Path clientCert = + java.nio.file.Paths.get("src", "test", "resources", "client_cert.pem"); + java.nio.file.Files.copy( + clientCert, tempCert, java.nio.file.StandardCopyOption.REPLACE_EXISTING); + + ChannelPoolSettings channelPoolSettings = + ChannelPoolSettings.builder().setInitialChannelCount(1).build(); + + pool = ChannelPool.create(channelPoolSettings, channelFactory, null, tempCert.toString()); + + // Initially uses channel1 + pool.newCall(FakeMethodDescriptor.create(), CallOptions.DEFAULT); + Mockito.verify(channel1, Mockito.times(1)) + .newCall(Mockito.>any(), Mockito.any(CallOptions.class)); + + // Rotate cert on disk + pool.invalidateDiskFingerprintCache(); + java.nio.file.Path rootCert = + java.nio.file.Paths.get("src", "test", "resources", "root_cert.pem"); + java.nio.file.Files.copy(rootCert, tempCert, java.nio.file.StandardCopyOption.REPLACE_EXISTING); + + // Refresh attempt 1: createSingleChannel throws IOException. + // Refresh should fail to replace channel and MUST NOT record the new cert fingerprint as + // active. + pool.refresh(); + + // Verify still channel1 + pool.newCall(FakeMethodDescriptor.create(), CallOptions.DEFAULT); + Mockito.verify(channel1, Mockito.times(2)) + .newCall(Mockito.>any(), Mockito.any(CallOptions.class)); + + // Refresh attempt 2: with the same cert file on disk (cache expired), channelFactory now + // succeeds. + // If the fingerprint had been mutated on the failed attempt, this call would be skipped as a + // duplicate! + pool.invalidateDiskFingerprintCache(); + pool.refresh(); + + // Verify it has now swapped to channel2! + pool.newCall(FakeMethodDescriptor.create(), CallOptions.DEFAULT); + Mockito.verify(channel2, Mockito.times(1)) + .newCall(Mockito.>any(), Mockito.any(CallOptions.class)); + } + + @Test + void + channelReactiveMTlsRefresh_partialFailureInMultiChannelPool_retainsShouldRefreshAndCompletesOnSubsequentRefresh() + throws IOException { + ManagedChannel initial1 = Mockito.mock(ManagedChannel.class); + ManagedChannel initial2 = Mockito.mock(ManagedChannel.class); + ManagedChannel rotated1 = Mockito.mock(ManagedChannel.class); + ManagedChannel rotated1SecondPass = Mockito.mock(ManagedChannel.class); + ManagedChannel rotated2 = Mockito.mock(ManagedChannel.class); + ChannelFactory channelFactory = + Mockito.mock(ChannelFactory.class, Mockito.withSettings().withoutAnnotations()); + + // Initial creation: initial1, initial2 + // Refresh pass 1: rotated1 succeeds, second throws IOException + // Refresh pass 2: rotated1SecondPass, rotated2 both succeed + Mockito.when(channelFactory.createSingleChannel()) + .thenReturn(initial1, initial2) + .thenReturn(rotated1) + .thenThrow(new IOException("Transient failure on second sub-channel")) + .thenReturn(rotated1SecondPass, rotated2); + + tempCert = java.nio.file.Files.createTempFile("cert", ".pem"); + java.nio.file.Path clientCert = + java.nio.file.Paths.get("src", "test", "resources", "client_cert.pem"); + java.nio.file.Files.copy( + clientCert, tempCert, java.nio.file.StandardCopyOption.REPLACE_EXISTING); + + pool = + ChannelPool.create( + ChannelPoolSettings.staticallySized(2), channelFactory, null, tempCert.toString()); + + // Rotate cert on disk + pool.invalidateDiskFingerprintCache(); + java.nio.file.Path rootCert = + java.nio.file.Paths.get("src", "test", "resources", "root_cert.pem"); + java.nio.file.Files.copy(rootCert, tempCert, java.nio.file.StandardCopyOption.REPLACE_EXISTING); + + assertThat(pool.shouldRefresh()).isTrue(); + long genBefore = pool.getGeneration(); + + // First refresh: partial failure (channel 0 rotates to rotated1, channel 1 fails and keeps + // initial2) + pool.refresh(); + + // Generation should still increment since partial progress was committed + assertThat(pool.getGeneration()).isGreaterThan(genBefore); + // initial1 should have been shut down, initial2 should NOT be shut down yet + Mockito.verify(initial1).shutdown(); + Mockito.verify(initial2, Mockito.never()).shutdown(); + + // Crucial assertion: shouldRefresh() MUST remain true so subsequent 401s on unrotated channel 1 + // trigger retry/refresh + pool.invalidateDiskFingerprintCache(); + assertThat(pool.shouldRefresh()).isTrue(); + + // Second refresh: both channels succeed + pool.refresh(); + + assertThat(pool.shouldRefresh()).isFalse(); + Mockito.verify(initial2).shutdown(); + } + + @Test + void refreshAll_runtimeExceptionOrError_doesNotLeakCreatedChannels() throws IOException { + ManagedChannel initial1 = Mockito.mock(ManagedChannel.class); + ManagedChannel initial2 = Mockito.mock(ManagedChannel.class); + ManagedChannel createdBeforeRuntimeEx = Mockito.mock(ManagedChannel.class); + ManagedChannel createdBeforeError = Mockito.mock(ManagedChannel.class); + ChannelFactory channelFactory = + Mockito.mock(ChannelFactory.class, Mockito.withSettings().withoutAnnotations()); + + Mockito.when(channelFactory.createSingleChannel()) + .thenReturn(initial1, initial2) + .thenReturn(createdBeforeRuntimeEx) + .thenThrow(new RuntimeException("Unchecked runtime exception")) + .thenReturn(createdBeforeError) + .thenThrow(new AssertionError("Simulated Error during refresh")); + + pool = ChannelPool.create(ChannelPoolSettings.staticallySized(2), channelFactory, null, null); + + // Case 1: RuntimeException on channel 1 after creating channel 0 -> caught as Exception, + // partial progress committed + boolean allCreated = pool.refreshAll(); + assertThat(allCreated).isFalse(); + Mockito.verify(initial1).shutdown(); + + // Case 2: Error on channel 1 after creating channel 0 -> aborts, finally block must shut down + // createdBeforeError + org.junit.jupiter.api.Assertions.assertThrows(AssertionError.class, () -> pool.refreshAll()); + Mockito.verify(createdBeforeError).shutdown(); + } + + @Test + void refresh_onShutdownPool_noOpsAndCreatesNoChannels() throws IOException { + ManagedChannel channel1 = mock(ManagedChannel.class); + ManagedChannel channel2 = mock(ManagedChannel.class); + ChannelFactory channelFactory = + Mockito.mock(ChannelFactory.class, Mockito.withSettings().withoutAnnotations()); + Mockito.when(channelFactory.createSingleChannel()).thenReturn(channel1, channel2); + + pool = ChannelPool.create(ChannelPoolSettings.staticallySized(1), channelFactory, null, null); + Mockito.verify(channelFactory, Mockito.times(1)).createSingleChannel(); + + pool.shutdown(); + assertThat(pool.isShutdown()).isTrue(); + + // Invoking refresh or refreshAll on shut down pool must no-op and never create new subchannels + pool.refresh(); + boolean refreshed = pool.refreshAll(); + assertThat(refreshed).isFalse(); + Mockito.verify(channelFactory, Mockito.times(1)).createSingleChannel(); + assertThat(pool.isShutdown()).isTrue(); + } + + @Test + void generationCounterIncrementsOnRefresh() throws IOException { + ManagedChannel channel1 = mock(ManagedChannel.class); + ManagedChannel channel2 = mock(ManagedChannel.class); + ChannelFactory channelFactory = + Mockito.mock(ChannelFactory.class, Mockito.withSettings().withoutAnnotations()); + Mockito.when(channelFactory.createSingleChannel()).thenReturn(channel1, channel2); + + pool = ChannelPool.create(ChannelPoolSettings.staticallySized(1), channelFactory, null, null); + assertThat(pool.getGeneration()).isEqualTo(0); + + pool.refreshAll(); + assertThat(pool.getGeneration()).isEqualTo(1); + } + @Test void channelRefreshShouldSwapChannels() throws IOException { ManagedChannel underlyingChannel1 = mock(ManagedChannel.class); @@ -450,7 +718,8 @@ void channelRefreshShouldSwapChannels() throws IOException { .setPreemptiveRefreshEnabled(true) .build(), channelFactory, - provider); + provider, + null); Mockito.reset(underlyingChannel1); pool.newCall(FakeMethodDescriptor.create(), CallOptions.DEFAULT); @@ -459,10 +728,41 @@ void channelRefreshShouldSwapChannels() throws IOException { .newCall(Mockito.>any(), Mockito.any(CallOptions.class)); // swap channel - pool.refresh(); + pool.refreshAll(); + + pool.newCall(FakeMethodDescriptor.create(), CallOptions.DEFAULT); + + Mockito.verify(underlyingChannel2, Mockito.only()) + .newCall(Mockito.>any(), Mockito.any(CallOptions.class)); + } + + @Test + void testRefreshWithNullWorkloadCertPathSwapsChannel() throws IOException { + ScheduledExecutorService executor = + Mockito.mock(ScheduledExecutorService.class, Mockito.withSettings().withoutAnnotations()); + FixedExecutorProvider provider = FixedExecutorProvider.create(executor); + ManagedChannel underlyingChannel1 = Mockito.mock(ManagedChannel.class); + ManagedChannel underlyingChannel2 = Mockito.mock(ManagedChannel.class); + FakeChannelFactory channelFactory = + new FakeChannelFactory(ImmutableList.of(underlyingChannel1, underlyingChannel2)); + pool = + new ChannelPool( + ChannelPoolSettings.staticallySized(1).toBuilder() + .setPreemptiveRefreshEnabled(true) + .build(), + channelFactory, + provider, + null); + Mockito.reset(underlyingChannel1); pool.newCall(FakeMethodDescriptor.create(), CallOptions.DEFAULT); + Mockito.verify(underlyingChannel1, Mockito.only()) + .newCall(Mockito.>any(), Mockito.any(CallOptions.class)); + // Calling refresh() when workloadCertPath is null should fall back to refreshAll() + pool.refresh(); + + pool.newCall(FakeMethodDescriptor.create(), CallOptions.DEFAULT); Mockito.verify(underlyingChannel2, Mockito.only()) .newCall(Mockito.>any(), Mockito.any(CallOptions.class)); } @@ -486,7 +786,8 @@ void channelCountShouldNotChangeWhenOutstandingRpcsAreWithinLimits() throws Exce .setMaxRpcsPerChannel(2) .build(), channelFactory, - provider); + provider, + null); assertThat(pool.entries.get()).hasSize(2); // Start the minimum number of @@ -553,7 +854,8 @@ void customResizeDeltaIsRespected() throws Exception { .setMaxResizeDelta(5) .build(), channelFactory, - provider); + provider, + null); assertThat(pool.entries.get()).hasSize(2); // Add 20 RPCs to push expansion @@ -586,7 +888,8 @@ void removedIdleChannelsAreShutdown() throws Exception { .setMaxRpcsPerChannel(2) .build(), channelFactory, - provider); + provider, + null); assertThat(pool.entries.get()).hasSize(2); // With no outstanding RPCs, the pool should shrink @@ -614,7 +917,8 @@ void removedActiveChannelsAreShutdown() throws Exception { .setMaxRpcsPerChannel(2) .build(), channelFactory, - provider); + provider, + null); assertThat(pool.entries.get()).hasSize(2); // Start 2 RPCs @@ -652,7 +956,7 @@ void testReleasingClientCallCancelEarly() throws IOException { Mockito.when(fakeChannel.newCall(Mockito.any(), Mockito.any())).thenReturn(mockClientCall); ChannelPoolSettings channelPoolSettings = ChannelPoolSettings.staticallySized(1); ChannelFactory factory = new FakeChannelFactory(ImmutableList.of(fakeChannel)); - pool = ChannelPool.create(channelPoolSettings, factory, null); + pool = ChannelPool.create(channelPoolSettings, factory, null, null); EndpointContext endpointContext = Mockito.mock(EndpointContext.class, Mockito.withSettings().withoutAnnotations()); @@ -717,7 +1021,8 @@ void repeatedResizingLogsWarningOnExpand() throws Exception { .setMaxChannelCount(10) .build(), channelFactory, - provider); + provider, + null); assertThat(pool.entries.get()).hasSize(1); FakeLogHandler logHandler = new FakeLogHandler(); @@ -769,7 +1074,8 @@ void repeatedResizingLogsWarningOnShrink() throws Exception { .setMaxChannelCount(10) .build(), channelFactory, - provider); + provider, + null); assertThat(pool.entries.get()).hasSize(10); FakeLogHandler logHandler = new FakeLogHandler(); @@ -805,7 +1111,7 @@ void testDoubleRelease() throws Exception { ChannelPoolSettings channelPoolSettings = ChannelPoolSettings.staticallySized(1); ChannelFactory factory = new FakeChannelFactory(ImmutableList.of(fakeChannel)); - pool = ChannelPool.create(channelPoolSettings, factory, null); + pool = ChannelPool.create(channelPoolSettings, factory, null, null); EndpointContext endpointContext = Mockito.mock(EndpointContext.class, Mockito.withSettings().withoutAnnotations()); @@ -843,7 +1149,8 @@ void testDoubleRelease() throws Exception { // Ensure that the channel pool properly logged the double call and kept the refCount correct assertThat(logHandler.getAllMessages()) .contains( - "Call is being closed more than once. Please make sure that onClose() is not being manually called."); + "Call is being closed more than once. Please make sure that onClose() is not being" + + " manually called."); assertThat(pool.entries.get()).hasSize(1); ChannelPool.Entry entry = pool.entries.get().get(0); assertThat(entry.outstandingRpcs.get()).isEqualTo(0); @@ -879,7 +1186,8 @@ void minChannelsClampedToMaxChannelCountUnderHighLoad() throws Exception { .setMaxChannelCount(5) .build(), channelFactory, - provider); + provider, + null); assertThat(pool.entries.get()).hasSize(1); // Add 20 RPCs, which would require 10 channels (20/2) @@ -914,7 +1222,8 @@ void maxChannelsClampedToMinChannelCountUnderLowLoad() throws Exception { .setMaxChannelCount(10) .build(), channelFactory, - provider); + provider, + null); assertThat(pool.entries.get()).hasSize(5); // With no outstanding RPCs, the pool should want to shrink to 0 @@ -924,4 +1233,274 @@ void maxChannelsClampedToMinChannelCountUnderLowLoad() throws Exception { // Should be clamped to minChannelCount = 3 assertThat(pool.entries.get()).hasSize(3); } + + @Test + void shouldRefresh_doesNotCacheNegativeResultAndDetectsSubsequentRotationImmediately() + throws IOException { + ManagedChannel initial = Mockito.mock(ManagedChannel.class); + ManagedChannel rotated = Mockito.mock(ManagedChannel.class); + ChannelFactory channelFactory = + Mockito.mock(ChannelFactory.class, Mockito.withSettings().withoutAnnotations()); + Mockito.when(channelFactory.createSingleChannel()).thenReturn(initial, rotated); + + tempCert = java.nio.file.Files.createTempFile("cert", ".pem"); + java.nio.file.Path clientCert = + java.nio.file.Paths.get("src", "test", "resources", "client_cert.pem"); + java.nio.file.Files.copy( + clientCert, tempCert, java.nio.file.StandardCopyOption.REPLACE_EXISTING); + + pool = + ChannelPool.create( + ChannelPoolSettings.staticallySized(1), channelFactory, null, tempCert.toString()); + + // First check returns false (unchanged disk cert) + assertThat(pool.shouldRefresh()).isFalse(); + + // Immediately rotate cert on disk WITHOUT invalidating the 1-second cache + java.nio.file.Path rootCert = + java.nio.file.Paths.get("src", "test", "resources", "root_cert.pem"); + java.nio.file.Files.copy(rootCert, tempCert, java.nio.file.StandardCopyOption.REPLACE_EXISTING); + + // Must immediately detect rotation because negative/unchanged disk checks are not cached for 1s + assertThat(pool.shouldRefresh()).isTrue(); + + // Refresh should update activeCertFingerprint and clear any cached positive check + pool.refresh(); + assertThat(pool.shouldRefresh()).isFalse(); + } + + @Test + void newCall_whenDelegateThrowsError_releasesEntryAndShutsDownRetiredChannel() + throws IOException { + ManagedChannel initial = Mockito.mock(ManagedChannel.class); + ManagedChannel rotated = Mockito.mock(ManagedChannel.class); + ChannelFactory channelFactory = + Mockito.mock(ChannelFactory.class, Mockito.withSettings().withoutAnnotations()); + Mockito.when(channelFactory.createSingleChannel()).thenReturn(initial, rotated); + Mockito.when(initial.newCall(Mockito.any(), Mockito.any())) + .thenThrow(new LinkageError("Simulated native/JNI linkage error")); + + pool = ChannelPool.create(ChannelPoolSettings.staticallySized(1), channelFactory, null, null); + + assertThrows(LinkageError.class, () -> pool.newCall(METHOD_RECOGNIZE, CallOptions.DEFAULT)); + + // Rotating the pool should immediately shut down initial channel because its ref count is 0 + pool.refresh(); + Mockito.verify(initial).shutdown(); + } + + @Test + @SuppressWarnings("unchecked") + void start_whenDelegateThrowsError_releasesEntryAndShutsDownRetiredChannel() throws IOException { + ManagedChannel initial = Mockito.mock(ManagedChannel.class); + ManagedChannel rotated = Mockito.mock(ManagedChannel.class); + ClientCall mockCall = Mockito.mock(ClientCall.class); + ChannelFactory channelFactory = + Mockito.mock(ChannelFactory.class, Mockito.withSettings().withoutAnnotations()); + Mockito.when(channelFactory.createSingleChannel()).thenReturn(initial, rotated); + Mockito.when(initial.newCall(Mockito.any(), Mockito.any())).thenReturn((ClientCall) mockCall); + Mockito.doThrow(new AssertionError("Simulated Error in start")) + .when(mockCall) + .start(Mockito.any(), Mockito.any()); + + pool = ChannelPool.create(ChannelPoolSettings.staticallySized(1), channelFactory, null, null); + + ClientCall call = pool.newCall(METHOD_RECOGNIZE, CallOptions.DEFAULT); + // Rotate pool while call is retained + pool.refresh(); + Mockito.verify(initial, Mockito.never()).shutdown(); + + // Calling start() throws Error, which must release the retained entry and trigger shutdown + assertThrows( + AssertionError.class, + () -> call.start(new ClientCall.Listener() {}, new io.grpc.Metadata())); + Mockito.verify(initial).shutdown(); + } + + @Test + @SuppressWarnings("unchecked") + void cancel_whenDelegateThrowsException_releasesEntryAndShutsDownRetiredChannel() + throws IOException { + ManagedChannel initial = Mockito.mock(ManagedChannel.class); + ManagedChannel rotated = Mockito.mock(ManagedChannel.class); + ClientCall mockCall = Mockito.mock(ClientCall.class); + ChannelFactory channelFactory = + Mockito.mock(ChannelFactory.class, Mockito.withSettings().withoutAnnotations()); + Mockito.when(channelFactory.createSingleChannel()).thenReturn(initial, rotated); + Mockito.when(initial.newCall(Mockito.any(), Mockito.any())).thenReturn((ClientCall) mockCall); + Mockito.doThrow(new RuntimeException("Simulated cancel exception")) + .when(mockCall) + .cancel(Mockito.any(), Mockito.any()); + + pool = ChannelPool.create(ChannelPoolSettings.staticallySized(1), channelFactory, null, null); + + ClientCall call = pool.newCall(METHOD_RECOGNIZE, CallOptions.DEFAULT); + pool.refresh(); + Mockito.verify(initial, Mockito.never()).shutdown(); + + assertThrows(RuntimeException.class, () -> call.cancel("cancelled", null)); + Mockito.verify(initial).shutdown(); + } + + @Test + @SuppressWarnings("unchecked") + void concurrentStartAndCancel_neverLeaksOrDoubleReleasesEntry() throws Exception { + ManagedChannel initial = Mockito.mock(ManagedChannel.class); + ManagedChannel rotated = Mockito.mock(ManagedChannel.class); + ChannelFactory channelFactory = + Mockito.mock(ChannelFactory.class, Mockito.withSettings().withoutAnnotations()); + Mockito.when(channelFactory.createSingleChannel()).thenReturn(initial, rotated); + + Mockito.when(initial.newCall(Mockito.any(), Mockito.any())) + .thenAnswer( + invocation -> + new ClientCall() { + private Listener listener; + private boolean cancelled; + + @Override + public synchronized void start( + Listener responseListener, io.grpc.Metadata headers) { + this.listener = responseListener; + if (cancelled) { + responseListener.onClose(io.grpc.Status.CANCELLED, new io.grpc.Metadata()); + } + } + + @Override + public synchronized void cancel(String message, Throwable cause) { + cancelled = true; + if (listener != null) { + listener.onClose(io.grpc.Status.CANCELLED, new io.grpc.Metadata()); + } + } + + @Override + public void request(int numMessages) {} + + @Override + public void halfClose() {} + + @Override + public void sendMessage(Color message) {} + }); + + pool = ChannelPool.create(ChannelPoolSettings.staticallySized(1), channelFactory, null, null); + + int iterations = 100; + java.util.concurrent.ExecutorService executor = + java.util.concurrent.Executors.newFixedThreadPool(2); + try { + for (int i = 0; i < iterations; i++) { + ClientCall call = pool.newCall(METHOD_RECOGNIZE, CallOptions.DEFAULT); + java.util.concurrent.CyclicBarrier barrier = new java.util.concurrent.CyclicBarrier(2); + java.util.concurrent.Future f1 = + executor.submit( + () -> { + try { + barrier.await(); + call.start(new ClientCall.Listener() {}, new io.grpc.Metadata()); + } catch (Exception ignored) { + } + }); + java.util.concurrent.Future f2 = + executor.submit( + () -> { + try { + barrier.await(); + call.cancel("cancel", null); + } catch (Exception ignored) { + } + }); + f1.get(5, java.util.concurrent.TimeUnit.SECONDS); + f2.get(5, java.util.concurrent.TimeUnit.SECONDS); + } + } finally { + executor.shutdownNow(); + } + + // Rotate pool: initial channel must shut down cleanly, proving outstandingRpcs == 0 (no leaks + // or negative counts) + pool.refresh(); + Mockito.verify(initial).shutdown(); + } + + @Test + void cancel_whenStartedAndSuperCancelThrows_doesNotReleasePrematurelyUntilOnClose() + throws Exception { + ManagedChannel initial = Mockito.mock(ManagedChannel.class); + ManagedChannel replacement = Mockito.mock(ManagedChannel.class); + @SuppressWarnings("unchecked") + ClientCall delegateCall = Mockito.mock(ClientCall.class); + @SuppressWarnings("unchecked") + ArgumentCaptor> listenerCaptor = + ArgumentCaptor.forClass(ClientCall.Listener.class); + Mockito.doThrow(new RuntimeException("cancel failure")) + .when(delegateCall) + .cancel(Mockito.any(), Mockito.any()); + Mockito.when(initial.newCall(Mockito.eq(METHOD_RECOGNIZE), Mockito.any())) + .thenReturn(delegateCall); + + java.util.concurrent.atomic.AtomicInteger createCount = + new java.util.concurrent.atomic.AtomicInteger(0); + pool = + new ChannelPool( + ChannelPoolSettings.staticallySized(1), + () -> createCount.getAndIncrement() == 0 ? initial : replacement, + FixedExecutorProvider.create(Mockito.mock(ScheduledExecutorService.class)), + null); + + ClientCall call = pool.newCall(METHOD_RECOGNIZE, CallOptions.DEFAULT); + call.start(new ClientCall.Listener() {}, new Metadata()); + Mockito.verify(delegateCall).start(listenerCaptor.capture(), Mockito.any()); + + assertThrows(RuntimeException.class, () -> call.cancel("abort", null)); + + // Rotate pool while call is still active (onClose hasn't fired yet): + // initial channel must NOT be shut down yet because call is still active + pool.refreshAll(); + Mockito.verify(initial, Mockito.never()).shutdown(); + + // Once onClose fires, entry is released and initial channel shuts down + listenerCaptor.getValue().onClose(Status.CANCELLED, new Metadata()); + Mockito.verify(initial).shutdown(); + } + + @Test + void start_whenCalledTwice_throwsIllegalStateExceptionAndDoesNotReleaseFirstCallEntry() + throws Exception { + ManagedChannel initial = Mockito.mock(ManagedChannel.class); + ManagedChannel replacement = Mockito.mock(ManagedChannel.class); + @SuppressWarnings("unchecked") + ClientCall delegateCall = Mockito.mock(ClientCall.class); + @SuppressWarnings("unchecked") + ArgumentCaptor> listenerCaptor = + ArgumentCaptor.forClass(ClientCall.Listener.class); + Mockito.when(initial.newCall(Mockito.eq(METHOD_RECOGNIZE), Mockito.any())) + .thenReturn(delegateCall); + + java.util.concurrent.atomic.AtomicInteger createCount = + new java.util.concurrent.atomic.AtomicInteger(0); + pool = + new ChannelPool( + ChannelPoolSettings.staticallySized(1), + () -> createCount.getAndIncrement() == 0 ? initial : replacement, + FixedExecutorProvider.create(Mockito.mock(ScheduledExecutorService.class)), + null); + + ClientCall call = pool.newCall(METHOD_RECOGNIZE, CallOptions.DEFAULT); + call.start(new ClientCall.Listener() {}, new Metadata()); + Mockito.verify(delegateCall).start(listenerCaptor.capture(), Mockito.any()); + + // Duplicate start() must throw IllegalStateException without releasing the entry + assertThrows( + IllegalStateException.class, + () -> call.start(new ClientCall.Listener() {}, new Metadata())); + + pool.refreshAll(); + Mockito.verify(initial, Mockito.never()).shutdown(); + + listenerCaptor.getValue().onClose(Status.OK, new Metadata()); + Mockito.verify(initial).shutdown(); + } } diff --git a/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcCallContextTest.java b/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcCallContextTest.java index e20767fdb8ed..cbaa7af2475c 100644 --- a/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcCallContextTest.java +++ b/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcCallContextTest.java @@ -494,4 +494,64 @@ private static Map> createTestExtraHeaders(String... keyVal } return extraHeaders; } + + @Test + public void testEqualsAndHashCode() { + ManagedChannel managedChannel1 = org.mockito.Mockito.mock(ManagedChannel.class); + ManagedChannel managedChannel2 = org.mockito.Mockito.mock(ManagedChannel.class); + + GrpcTransportChannel transportChannel1 = GrpcTransportChannel.create(managedChannel1); + GrpcTransportChannel transportChannel2 = GrpcTransportChannel.create(managedChannel2); + + GrpcCallContext context1 = + GrpcCallContext.createDefault().withTransportChannel(transportChannel1); + GrpcCallContext context2 = + GrpcCallContext.createDefault().withTransportChannel(transportChannel1); + GrpcCallContext context3 = + GrpcCallContext.createDefault().withTransportChannel(transportChannel2); + + org.junit.jupiter.api.Assertions.assertEquals(context1, context2); + org.junit.jupiter.api.Assertions.assertEquals(context1.hashCode(), context2.hashCode()); + + org.junit.jupiter.api.Assertions.assertNotEquals(context1, context3); + } + + @Test + public void testMergeWithCustomChannelClearsTransportChannel() { + ManagedChannel defaultChannel = org.mockito.Mockito.mock(ManagedChannel.class); + ManagedChannel customChannel = org.mockito.Mockito.mock(ManagedChannel.class); + GrpcTransportChannel transportChannel = GrpcTransportChannel.create(defaultChannel); + + GrpcCallContext baseContext = + GrpcCallContext.createDefault().withTransportChannel(transportChannel); + GrpcCallContext overrideContext = GrpcCallContext.of(customChannel, CallOptions.DEFAULT); + + GrpcCallContext mergedContext = (GrpcCallContext) baseContext.merge(overrideContext); + assertEquals(customChannel, mergedContext.getChannel()); + assertNull(mergedContext.getTransportChannel()); + } + + @Test + public void testWithChannelWithCustomChannelClearsTransportChannel() { + ManagedChannel defaultChannel = org.mockito.Mockito.mock(ManagedChannel.class); + ManagedChannel customChannel = org.mockito.Mockito.mock(ManagedChannel.class); + GrpcTransportChannel transportChannel = GrpcTransportChannel.create(defaultChannel); + + GrpcCallContext baseContext = + GrpcCallContext.createDefault().withTransportChannel(transportChannel); + GrpcCallContext updatedContext = baseContext.withChannel(customChannel); + + assertEquals(customChannel, updatedContext.getChannel()); + assertNull(updatedContext.getTransportChannel()); + + // Clearing channel via withChannel(null) also clears transportChannel + GrpcCallContext nullChannelContext = baseContext.withChannel(null); + assertNull(nullChannelContext.getChannel()); + assertNull(nullChannelContext.getTransportChannel()); + + // Merging a cleared context into defaultContext falls back to defaultContext's transportChannel + GrpcCallContext mergedWithNullChannel = (GrpcCallContext) baseContext.merge(nullChannelContext); + assertEquals(defaultChannel, mergedWithNullChannel.getChannel()); + assertEquals(transportChannel, mergedWithNullChannel.getTransportChannel()); + } } diff --git a/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcClientCallsTest.java b/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcClientCallsTest.java index 2aa9279e249f..6877eb1dbe1e 100644 --- a/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcClientCallsTest.java +++ b/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcClientCallsTest.java @@ -125,6 +125,7 @@ void testAffinity() throws IOException { ChannelPool.create( ChannelPoolSettings.staticallySized(2), new FakeChannelFactory(Arrays.asList(channel0, channel1)), + null, null); GrpcCallContext context = defaultCallContext.withChannel(pool); diff --git a/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcLoggingInterceptorTest.java b/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcLoggingInterceptorTest.java index fad4cd468b95..c93db599d575 100644 --- a/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcLoggingInterceptorTest.java +++ b/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcLoggingInterceptorTest.java @@ -32,7 +32,6 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.spy; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -83,7 +82,7 @@ void testInterceptor_basic() { void testInterceptor_responseListener() { when(channel.newCall(Mockito.>any(), any(CallOptions.class))) .thenReturn(call); - GrpcLoggingInterceptor interceptor = spy(new GrpcLoggingInterceptor()); + GrpcLoggingInterceptor interceptor = new GrpcLoggingInterceptor(); Channel intercepted = ClientInterceptors.intercept(channel, interceptor); @SuppressWarnings("unchecked") ClientCall.Listener listener = mock(ClientCall.Listener.class); diff --git a/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProviderTest.java b/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProviderTest.java index be0365866615..d083c01ad845 100644 --- a/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProviderTest.java +++ b/sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProviderTest.java @@ -48,6 +48,7 @@ import com.google.api.gax.rpc.internal.EnvironmentProvider; import com.google.api.gax.rpc.mtls.AbstractMtlsTransportChannelTest; import com.google.api.gax.rpc.mtls.CertificateBasedAccess; +import com.google.api.gax.rpc.testing.FakeMtlsProvider; import com.google.auth.ApiKeyCredentials; import com.google.auth.Credentials; import com.google.auth.http.AuthHttpConstants; @@ -664,7 +665,9 @@ private void createAndCloseTransportChannel(InstantiatingGrpcChannelProvider pro createAndCloseTransportChannel(provider); assertThat(logHandler.getAllMessages()) .contains( - "DirectPath is misconfigured. The DirectPath XDS option was set, but the attemptDirectPath option was not. Please set both the attemptDirectPath and attemptDirectPathXds options."); + "DirectPath is misconfigured. The DirectPath XDS option was set, but the" + + " attemptDirectPath option was not. Please set both the attemptDirectPath and" + + " attemptDirectPathXds options."); InstantiatingGrpcChannelProvider.LOG.removeHandler(logHandler); } @@ -682,8 +685,10 @@ void testLogDirectPathMisconfig_AttemptDirectPathNotSetAndAttemptDirectPathXdsSe createAndCloseTransportChannel(provider); assertThat(logHandler.getAllMessages()) .contains( - "Env var GOOGLE_CLOUD_ENABLE_DIRECT_PATH_XDS was found and set to TRUE, but DirectPath was not enabled for this client. If this is intended for " - + "this client, please note that this is a misconfiguration and set the attemptDirectPath option as well."); + "Env var GOOGLE_CLOUD_ENABLE_DIRECT_PATH_XDS was found and set to TRUE, but DirectPath" + + " was not enabled for this client. If this is intended for this client, please" + + " note that this is a misconfiguration and set the attemptDirectPath option as" + + " well."); InstantiatingGrpcChannelProvider.LOG.removeHandler(logHandler); } @@ -711,6 +716,7 @@ void testLogDirectPathMisconfigWrongCredential() throws Exception { InstantiatingGrpcChannelProvider.newBuilder() .setAttemptDirectPathXds() .setAttemptDirectPath(true) + .setEnvProvider(name -> null) .setHeaderProvider( mock(HeaderProvider.class, Mockito.withSettings().withoutAnnotations())) .setExecutor(mock(Executor.class)) @@ -877,12 +883,14 @@ public void canUseDirectPath_directPathEnvVarDisabled() throws IOException { @Test public void canUseDirectPath_directPathEnvVarNotSet_attemptDirectPathIsTrue() { System.setProperty("os.name", "Linux"); + EnvironmentProvider envProvider = name -> null; InstantiatingGrpcChannelProvider.Builder builder = InstantiatingGrpcChannelProvider.newBuilder() .setCertificateBasedAccess(certificateBasedAccess) .setAttemptDirectPath(true) .setCredentials(computeEngineCredentials) - .setEndpoint(DEFAULT_ENDPOINT); + .setEndpoint(DEFAULT_ENDPOINT) + .setEnvProvider(envProvider); InstantiatingGrpcChannelProvider provider = new InstantiatingGrpcChannelProvider(builder, GCE_PRODUCTION_NAME_AFTER_2016); Truth.assertThat(provider.canUseDirectPath()).isTrue(); @@ -891,12 +899,14 @@ public void canUseDirectPath_directPathEnvVarNotSet_attemptDirectPathIsTrue() { @Test public void canUseDirectPath_directPathEnvVarNotSet_attemptDirectPathIsFalse() { System.setProperty("os.name", "Linux"); + EnvironmentProvider envProvider = name -> null; InstantiatingGrpcChannelProvider.Builder builder = InstantiatingGrpcChannelProvider.newBuilder() .setCertificateBasedAccess(certificateBasedAccess) .setAttemptDirectPath(false) .setCredentials(computeEngineCredentials) - .setEndpoint(DEFAULT_ENDPOINT); + .setEndpoint(DEFAULT_ENDPOINT) + .setEnvProvider(envProvider); InstantiatingGrpcChannelProvider provider = new InstantiatingGrpcChannelProvider(builder, GCE_PRODUCTION_NAME_AFTER_2016); Truth.assertThat(provider.canUseDirectPath()).isFalse(); @@ -1201,7 +1211,8 @@ void createS2ASecuredChannelCredentials_bothS2AAddressesNull_returnsNull() { assertThat(provider.createS2ASecuredChannelCredentials()).isNotNull(); assertThat(logHandler.getAllMessages()) .contains( - "Cannot establish an mTLS connection to S2A because autoconfig endpoint did not return a mtls address to reach S2A."); + "Cannot establish an mTLS connection to S2A because autoconfig endpoint did not return" + + " a mtls address to reach S2A."); InstantiatingGrpcChannelProvider.LOG.removeHandler(logHandler); } @@ -1247,7 +1258,8 @@ void createS2ASecuredChannelCredentials_returnsPlaintextToS2AS2AChannelCredentia assertThat(provider.createS2ASecuredChannelCredentials()).isNotNull(); assertThat(logHandler.getAllMessages()) .contains( - "Cannot establish an mTLS connection to S2A because MTLS to MDS credentials do not exist on filesystem, falling back to plaintext connection to S2A"); + "Cannot establish an mTLS connection to S2A because MTLS to MDS credentials do not" + + " exist on filesystem, falling back to plaintext connection to S2A"); InstantiatingGrpcChannelProvider.LOG.removeHandler(logHandler); } @@ -1342,6 +1354,96 @@ void testSettingBackgroundExecutor() { assertThat(provider.getBackgroundExecutor()).isEqualTo(mockExecutor); } + @Test + void createChannel_whenDirectPathEnabled_ignoresWorkloadCertPath() throws Exception { + System.setProperty("os.name", "Linux"); + EnvironmentProvider envProvider = + mock(EnvironmentProvider.class, Mockito.withSettings().withoutAnnotations()); + Mockito.when( + envProvider.getenv( + InstantiatingGrpcChannelProvider.DIRECT_PATH_ENV_DISABLE_DIRECT_PATH)) + .thenReturn("false"); + CertificateBasedAccess mtlsCertificateBasedAccess = + mock(CertificateBasedAccess.class, Mockito.withSettings().withoutAnnotations()); + Mockito.when(mtlsCertificateBasedAccess.useMtlsClientCertificate()).thenReturn(true); + Mockito.when(mtlsCertificateBasedAccess.getWorkloadCertPath()) + .thenReturn("/path/to/workload/cert.pem"); + MtlsProvider mtlsProvider = + new FakeMtlsProvider(FakeMtlsProvider.createTestMtlsKeyStore(), "", false); + + InstantiatingGrpcChannelProvider.Builder builder = + InstantiatingGrpcChannelProvider.newBuilder() + .setCertificateBasedAccess(mtlsCertificateBasedAccess) + .setMtlsProvider(mtlsProvider) + .setAttemptDirectPath(true) + .setCredentials(computeEngineCredentials) + .setEndpoint(DEFAULT_ENDPOINT) + .setEnvProvider(envProvider) + .setHeaderProvider( + mock(HeaderProvider.class, Mockito.withSettings().withoutAnnotations())); + InstantiatingGrpcChannelProvider provider = + new InstantiatingGrpcChannelProvider(builder, GCE_PRODUCTION_NAME_AFTER_2016); + Truth.assertThat(provider.canUseDirectPath()).isTrue(); + + TransportChannel transportChannel = provider.getTransportChannel(); + try { + ChannelPool pool = (ChannelPool) ((GrpcTransportChannel) transportChannel).getChannel(); + assertThat(pool.getWorkloadCertPath()).isNull(); + } finally { + transportChannel.shutdownNow(); + } + } + + @Test + void createChannel_whenMtlsActive_passesWorkloadCertPathToChannelPool() throws Exception { + CertificateBasedAccess mtlsCertificateBasedAccess = + mock(CertificateBasedAccess.class, Mockito.withSettings().withoutAnnotations()); + Mockito.when(mtlsCertificateBasedAccess.useMtlsClientCertificate()).thenReturn(true); + Mockito.when(mtlsCertificateBasedAccess.getWorkloadCertPath()) + .thenReturn("/path/to/workload/cert.pem"); + MtlsProvider mtlsProvider = + new FakeMtlsProvider(FakeMtlsProvider.createTestMtlsKeyStore(), "", false); + + InstantiatingGrpcChannelProvider provider = + InstantiatingGrpcChannelProvider.newBuilder() + .setCertificateBasedAccess(mtlsCertificateBasedAccess) + .setMtlsProvider(mtlsProvider) + .setAttemptDirectPath(false) + .setEndpoint(DEFAULT_ENDPOINT) + .setHeaderProvider( + mock(HeaderProvider.class, Mockito.withSettings().withoutAnnotations())) + .build(); + + TransportChannel transportChannel = provider.getTransportChannel(); + try { + ChannelPool pool = (ChannelPool) ((GrpcTransportChannel) transportChannel).getChannel(); + assertThat(pool.getWorkloadCertPath()).isEqualTo("/path/to/workload/cert.pem"); + } finally { + transportChannel.shutdownNow(); + } + } + + @Test + void createChannelBuilder_whenMtlsActiveAndCredentialsNull_throwsIOException() { + CertificateBasedAccess mtlsCertificateBasedAccess = + mock(CertificateBasedAccess.class, Mockito.withSettings().withoutAnnotations()); + Mockito.when(mtlsCertificateBasedAccess.useMtlsClientCertificate()).thenReturn(true); + MtlsProvider mtlsProviderWithNullKeyStore = new FakeMtlsProvider(null, "", false); + + InstantiatingGrpcChannelProvider provider = + InstantiatingGrpcChannelProvider.newBuilder() + .setCertificateBasedAccess(mtlsCertificateBasedAccess) + .setMtlsProvider(mtlsProviderWithNullKeyStore) + .setAttemptDirectPath(false) + .setEndpoint(DEFAULT_ENDPOINT) + .setHeaderProvider( + mock(HeaderProvider.class, Mockito.withSettings().withoutAnnotations())) + .build(); + + IOException thrown = assertThrows(IOException.class, provider::createChannelBuilder); + assertThat(thrown).hasMessageThat().contains("Failed to initialize mTLS channel credentials"); + } + private static class FakeLogHandler extends Handler { List records = new ArrayList<>(); diff --git a/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonCallContext.java b/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonCallContext.java index 2679b51860df..1944b6c9a6b3 100644 --- a/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonCallContext.java +++ b/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonCallContext.java @@ -82,6 +82,7 @@ public final class HttpJsonCallContext implements ApiCallContext { private final @Nullable RetrySettings retrySettings; private final @Nullable ImmutableSet retryableCodes; private final EndpointContext endpointContext; + @Nullable private final TransportChannel transportChannel; /** Returns an empty instance. */ public static HttpJsonCallContext createDefault() { @@ -96,6 +97,7 @@ public static HttpJsonCallContext createDefault() { null, null, null, + null, null); } @@ -111,6 +113,7 @@ public static HttpJsonCallContext of(HttpJsonChannel channel, HttpJsonCallOption null, null, null, + null, null); } @@ -125,7 +128,8 @@ private HttpJsonCallContext( @Nullable ApiTracer tracer, @Nullable RetrySettings defaultRetrySettings, @Nullable Set defaultRetryableCodes, - @Nullable EndpointContext endpointContext) { + @Nullable EndpointContext endpointContext, + @Nullable TransportChannel transportChannel) { this.channel = channel; this.callOptions = callOptions; this.timeout = timeout; @@ -141,6 +145,7 @@ private HttpJsonCallContext( // a valid EndpointContext with user configurations after the client has been initialized. this.endpointContext = endpointContext == null ? EndpointContext.getDefaultInstance() : endpointContext; + this.transportChannel = transportChannel; } /** @@ -220,6 +225,12 @@ public HttpJsonCallContext merge(ApiCallContext inputCallContext) { newRetryableCodes = this.retryableCodes; } + TransportChannel newTransportChannel = httpJsonCallContext.transportChannel; + if (newTransportChannel == null + && (httpJsonCallContext.channel == null || httpJsonCallContext.channel.equals(channel))) { + newTransportChannel = this.transportChannel; + } + // The EndpointContext is not updated as there should be no reason for a user // to update this. return new HttpJsonCallContext( @@ -233,7 +244,8 @@ public HttpJsonCallContext merge(ApiCallContext inputCallContext) { newTracer, newRetrySettings, newRetryableCodes, - endpointContext); + endpointContext, + newTransportChannel); } @Override @@ -251,7 +263,24 @@ public HttpJsonCallContext withTransportChannel(TransportChannel inputChannel) { "Expected HttpJsonTransportChannel, got " + inputChannel.getClass().getName()); } HttpJsonTransportChannel transportChannel = (HttpJsonTransportChannel) inputChannel; - return withChannel(transportChannel.getChannel()); + return new HttpJsonCallContext( + transportChannel.getChannel(), + this.callOptions, + this.timeout, + this.streamWaitTimeout, + this.streamIdleTimeout, + this.extraHeaders, + this.options, + this.tracer, + this.retrySettings, + this.retryableCodes, + this.endpointContext, + transportChannel); + } + + @Override + public TransportChannel getTransportChannel() { + return transportChannel; } /** This method is obsolete. Use {@link #withTimeoutDuration(java.time.Duration)} instead. */ @@ -275,7 +304,8 @@ public HttpJsonCallContext withEndpointContext(EndpointContext endpointContext) this.tracer, this.retrySettings, this.retryableCodes, - endpointContext); + endpointContext, + this.transportChannel); } @Override @@ -286,7 +316,7 @@ public HttpJsonCallContext withTimeoutDuration(java.time.Duration timeout) { } // Prevent expanding deadlines - if (timeout != null && this.timeout != null && this.timeout.compareTo(timeout) <= 0) { + if (this.timeout != null && (timeout == null || this.timeout.compareTo(timeout) <= 0)) { return this; } @@ -301,7 +331,8 @@ public HttpJsonCallContext withTimeoutDuration(java.time.Duration timeout) { this.tracer, this.retrySettings, this.retryableCodes, - this.endpointContext); + this.endpointContext, + this.transportChannel); } /** This method is obsolete. Use {@link #getTimeoutDuration()} instead. */ @@ -346,7 +377,8 @@ public HttpJsonCallContext withStreamWaitTimeoutDuration( this.tracer, this.retrySettings, this.retryableCodes, - this.endpointContext); + this.endpointContext, + this.transportChannel); } /** This method is obsolete. Use {@link #getStreamWaitTimeoutDuration()} instead. */ @@ -396,7 +428,8 @@ public HttpJsonCallContext withStreamIdleTimeoutDuration( this.tracer, this.retrySettings, this.retryableCodes, - this.endpointContext); + this.endpointContext, + this.transportChannel); } /** This method is obsolete. Use {@link #getStreamIdleTimeoutDuration()} instead. */ @@ -433,7 +466,8 @@ public ApiCallContext withExtraHeaders(Map> extraHeaders) { this.tracer, this.retrySettings, this.retryableCodes, - this.endpointContext); + this.endpointContext, + this.transportChannel); } @BetaApi("The surface for extra headers is not stable yet and may change in the future.") @@ -457,7 +491,8 @@ public ApiCallContext withOption(Key key, T value) { this.tracer, this.retrySettings, this.retryableCodes, - this.endpointContext); + this.endpointContext, + this.transportChannel); } /** {@inheritDoc} */ @@ -527,7 +562,8 @@ public HttpJsonCallContext withRetrySettings(RetrySettings retrySettings) { this.tracer, retrySettings, this.retryableCodes, - this.endpointContext); + this.endpointContext, + this.transportChannel); } @Override @@ -548,7 +584,8 @@ public HttpJsonCallContext withRetryableCodes(Set retryableCode this.tracer, this.retrySettings, retryableCodes, - this.endpointContext); + this.endpointContext, + this.transportChannel); } public HttpJsonCallContext withChannel(@Nullable HttpJsonChannel newChannel) { @@ -563,7 +600,8 @@ public HttpJsonCallContext withChannel(@Nullable HttpJsonChannel newChannel) { this.tracer, this.retrySettings, this.retryableCodes, - this.endpointContext); + this.endpointContext, + (newChannel != null && newChannel.equals(this.channel)) ? this.transportChannel : null); } public HttpJsonCallContext withCallOptions(HttpJsonCallOptions newCallOptions) { @@ -578,7 +616,8 @@ public HttpJsonCallContext withCallOptions(HttpJsonCallOptions newCallOptions) { this.tracer, this.retrySettings, this.retryableCodes, - this.endpointContext); + this.endpointContext, + this.transportChannel); } @Deprecated @@ -614,7 +653,8 @@ public HttpJsonCallContext withTracer(@Nonnull ApiTracer newTracer) { newTracer, this.retrySettings, this.retryableCodes, - this.endpointContext); + this.endpointContext, + this.transportChannel); } @Override @@ -634,7 +674,8 @@ public boolean equals(@Nullable Object o) { && Objects.equals(this.tracer, that.tracer) && Objects.equals(this.retrySettings, that.retrySettings) && Objects.equals(this.retryableCodes, that.retryableCodes) - && Objects.equals(this.endpointContext, that.endpointContext); + && Objects.equals(this.endpointContext, that.endpointContext) + && Objects.equals(this.transportChannel, that.transportChannel); } @Override @@ -648,6 +689,7 @@ public int hashCode() { tracer, retrySettings, retryableCodes, - endpointContext); + endpointContext, + transportChannel); } } diff --git a/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonTransportChannel.java b/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonTransportChannel.java index 813622b6a97e..fde0673f650d 100644 --- a/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonTransportChannel.java +++ b/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonTransportChannel.java @@ -64,6 +64,21 @@ public HttpJsonChannel getChannel() { return getManagedChannel(); } + @Override + public void refresh() { + getManagedChannel().refresh(); + } + + @Override + public boolean shouldRefresh() { + return getManagedChannel().shouldRefresh(); + } + + @Override + public long getGeneration() { + return getManagedChannel().getGeneration(); + } + @Override public void shutdown() { getManagedChannel().shutdown(); diff --git a/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/InstantiatingHttpJsonChannelProvider.java b/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/InstantiatingHttpJsonChannelProvider.java index 92ce4efe36aa..11d1856b5d4b 100644 --- a/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/InstantiatingHttpJsonChannelProvider.java +++ b/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/InstantiatingHttpJsonChannelProvider.java @@ -195,71 +195,93 @@ public TransportChannelProvider withCredentials(Credentials credentials) { "InstantiatingHttpJsonChannelProvider doesn't need credentials"); } - HttpTransport createHttpTransport() throws IOException, GeneralSecurityException { - NetHttpTransport.Builder builder = new NetHttpTransport.Builder(); - configureMtls(builder); - HttpJsonConscryptUtils.configureConscryptSecurityProvider(builder); - return builder.build(); - } - - private NetHttpTransport.Builder configureMtls(NetHttpTransport.Builder builder) - throws IOException, GeneralSecurityException { - if (mtlsProvider == null || !certificateBasedAccess.useMtlsClientCertificate()) { - return builder; - } - KeyStore mtlsKeyStore = mtlsProvider.getKeyStore(); - if (mtlsKeyStore == null) { - return builder; + @Nullable HttpTransport createHttpTransport() throws IOException, GeneralSecurityException { + if (mtlsProvider == null) { + return null; } - builder.trustCertificates(null, mtlsKeyStore, ""); - Provider conscryptProvider = HttpJsonConscryptUtils.getConscryptProvider(); - if (conscryptProvider == null) { - // Fall back to standard JDK JSSE if Conscrypt provider is unavailable - return builder; + if (certificateBasedAccess.useMtlsClientCertificate()) { + KeyStore mtlsKeyStore = mtlsProvider.getKeyStore(); + if (mtlsKeyStore != null) { + NetHttpTransport.Builder builder = new NetHttpTransport.Builder(); + builder.trustCertificates(null, mtlsKeyStore, ""); + Provider conscryptProvider = HttpJsonConscryptUtils.getConscryptProvider(); + if (conscryptProvider != null) { + SSLContext sslContext = SSLContext.getInstance("TLS", conscryptProvider); + SslUtils.initSslContext( + sslContext, + null, + SslUtils.getPkixTrustManagerFactory(), + mtlsKeyStore, + "", + SslUtils.getDefaultKeyManagerFactory()); + builder.setSslSocketFactory(sslContext.getSocketFactory()); + } + HttpJsonConscryptUtils.configureConscryptSecurityProvider(builder); + return builder.build(); + } } - // Explicitly initialize SSLContext with the Conscrypt provider so that the client certificate - // key managers - // and trust manager factory (TMF) are bound to Conscrypt's TLS implementation (supporting PQC - // key exchange). - SSLContext sslContext = SSLContext.getInstance("TLS", conscryptProvider); - SslUtils.initSslContext( - sslContext, - null, - SslUtils.getPkixTrustManagerFactory(), - mtlsKeyStore, - "", - SslUtils.getDefaultKeyManagerFactory()); - builder.setSslSocketFactory(sslContext.getSocketFactory()); - return builder; + return null; } - private HttpJsonTransportChannel createChannel() throws IOException, GeneralSecurityException { + private ManagedHttpJsonChannel createSingleManagedChannel() + throws IOException, GeneralSecurityException { HttpTransport httpTransportToUse = httpTransport; if (httpTransportToUse == null) { httpTransportToUse = createHttpTransport(); + if (httpTransportToUse == null + && mtlsProvider != null + && certificateBasedAccess.useMtlsClientCertificate()) { + throw new IOException("Failed to initialize mTLS HttpTransport"); + } } + return ManagedHttpJsonChannel.newBuilder() + .setEndpoint(endpoint) + .setExecutor(executor) + .setHttpTransport(httpTransportToUse) + .setManageHttpTransport(httpTransport == null) + .build(); + } - // Pass the executor to the ManagedChannel. If no executor was provided (or null), - // the channel will use a default executor for the calls. - ManagedHttpJsonChannel channel = - ManagedHttpJsonChannel.newBuilder() - .setEndpoint(endpoint) - .setExecutor(executor) - .setHttpTransport(httpTransportToUse) - .build(); - - HttpJsonClientInterceptor headerInterceptor = - new HttpJsonHeaderInterceptor(headerProvider.getHeaders()); - - channel = new ManagedHttpJsonInterceptorChannel(channel, new HttpJsonLoggingInterceptor()); - channel = new ManagedHttpJsonInterceptorChannel(channel, headerInterceptor); - if (interceptorProvider != null && interceptorProvider.getInterceptors() != null) { - for (HttpJsonClientInterceptor interceptor : interceptorProvider.getInterceptors()) { - channel = new ManagedHttpJsonInterceptorChannel(channel, interceptor); + private HttpJsonTransportChannel createChannel() throws IOException, GeneralSecurityException { + boolean isMtlsActive = + httpTransport == null + && mtlsProvider != null + && certificateBasedAccess.useMtlsClientCertificate(); + String workloadCertPath = isMtlsActive ? certificateBasedAccess.getWorkloadCertPath() : null; + + ManagedHttpJsonChannel initialChannel = createSingleManagedChannel(); + try { + java.util.function.Supplier channelFactory = + () -> { + try { + return createSingleManagedChannel(); + } catch (Exception e) { + throw new java.lang.RuntimeException( + "Failed to create fresh ManagedHttpJsonChannel", e); + } + }; + + ManagedHttpJsonChannel channel = + workloadCertPath != null + ? new RefreshingHttpJsonChannel(initialChannel, channelFactory, workloadCertPath) + : initialChannel; + + HttpJsonClientInterceptor headerInterceptor = + new HttpJsonHeaderInterceptor(headerProvider.getHeaders()); + + channel = new ManagedHttpJsonInterceptorChannel(channel, new HttpJsonLoggingInterceptor()); + channel = new ManagedHttpJsonInterceptorChannel(channel, headerInterceptor); + if (interceptorProvider != null && interceptorProvider.getInterceptors() != null) { + for (HttpJsonClientInterceptor interceptor : interceptorProvider.getInterceptors()) { + channel = new ManagedHttpJsonInterceptorChannel(channel, interceptor); + } } - } - return HttpJsonTransportChannel.newBuilder().setManagedChannel(channel).build(); + return HttpJsonTransportChannel.newBuilder().setManagedChannel(channel).build(); + } catch (Throwable t) { + initialChannel.shutdownNow(); + throw t; + } } /** The endpoint to be used for the channel. */ diff --git a/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/ManagedHttpJsonChannel.java b/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/ManagedHttpJsonChannel.java index 87767bee5c7f..86863e4ce51c 100644 --- a/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/ManagedHttpJsonChannel.java +++ b/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/ManagedHttpJsonChannel.java @@ -52,11 +52,25 @@ public class ManagedHttpJsonChannel implements HttpJsonChannel, BackgroundResour private final boolean usingDefaultExecutor; private final String endpoint; private final HttpTransport httpTransport; + private final boolean usingDefaultTransport; private final ScheduledExecutorService deadlineScheduledExecutorService; private boolean isTransportShutdown; protected ManagedHttpJsonChannel() { - this(null, true, null, null); + this(null, true, null, null, true); + } + + protected ManagedHttpJsonChannel(boolean isDelegatingWrapper) { + this.executor = null; + this.usingDefaultExecutor = false; + this.endpoint = null; + this.httpTransport = null; + this.usingDefaultTransport = false; + this.deadlineScheduledExecutorService = null; + } + + public long getGeneration() { + return 0; } String getEndpoint() { @@ -72,7 +86,8 @@ private ManagedHttpJsonChannel( @Nullable Executor executor, boolean usingDefaultExecutor, @Nullable String endpoint, - @Nullable HttpTransport httpTransport) { + @Nullable HttpTransport httpTransport, + boolean usingDefaultTransport) { this.executor = executor; this.usingDefaultExecutor = usingDefaultExecutor; this.endpoint = endpoint; @@ -82,6 +97,7 @@ private ManagedHttpJsonChannel( new NetHttpTransport.Builder()) .build() : httpTransport; + this.usingDefaultTransport = usingDefaultTransport || httpTransport == null; this.deadlineScheduledExecutorService = Executors.newSingleThreadScheduledExecutor(); } @@ -98,6 +114,12 @@ public HttpJsonClientCall newCall( deadlineScheduledExecutorService); } + public void refresh() {} + + public boolean shouldRefresh() { + return false; + } + @VisibleForTesting Executor getExecutor() { return executor; @@ -116,7 +138,9 @@ public synchronized void shutdown() { ((ExecutorService) executor).shutdown(); } deadlineScheduledExecutorService.shutdown(); - httpTransport.shutdown(); + if (usingDefaultTransport) { + httpTransport.shutdown(); + } isTransportShutdown = true; } catch (IOException e) { // TODO: Log this scenario once we implemented the Cloud SDK logging. @@ -158,7 +182,9 @@ public void shutdownNow() { ((ExecutorService) executor).shutdownNow(); } deadlineScheduledExecutorService.shutdownNow(); - httpTransport.shutdown(); + if (usingDefaultTransport) { + httpTransport.shutdown(); + } isTransportShutdown = true; } catch (IOException e) { // TODO: Log this scenario once we implemented the Cloud SDK logging. @@ -205,9 +231,11 @@ public static class Builder { private String endpoint; private HttpTransport httpTransport; private boolean usingDefaultExecutor; + private boolean usingDefaultTransport; private Builder() { this.usingDefaultExecutor = false; + this.usingDefaultTransport = false; } public Builder setExecutor(Executor executor) { @@ -225,6 +253,11 @@ public Builder setHttpTransport(HttpTransport httpTransport) { return this; } + Builder setManageHttpTransport(boolean manageHttpTransport) { + this.usingDefaultTransport = manageHttpTransport; + return this; + } + public ManagedHttpJsonChannel build() { Preconditions.checkNotNull(endpoint); @@ -237,14 +270,8 @@ public ManagedHttpJsonChannel build() { usingDefaultExecutor = true; } - if (httpTransport == null) { - httpTransport = - HttpJsonConscryptUtils.configureConscryptSecurityProvider( - new NetHttpTransport.Builder()) - .build(); - } - - return new ManagedHttpJsonChannel(executor, usingDefaultExecutor, endpoint, httpTransport); + return new ManagedHttpJsonChannel( + executor, usingDefaultExecutor, endpoint, httpTransport, usingDefaultTransport); } } } diff --git a/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/ManagedHttpJsonInterceptorChannel.java b/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/ManagedHttpJsonInterceptorChannel.java index eaaa8c3a7c56..9ccbbc4ec1a2 100644 --- a/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/ManagedHttpJsonInterceptorChannel.java +++ b/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/ManagedHttpJsonInterceptorChannel.java @@ -29,7 +29,9 @@ */ package com.google.api.gax.httpjson; +import com.google.api.client.http.HttpTransport; import com.google.common.annotations.VisibleForTesting; +import java.util.concurrent.Executor; import java.util.concurrent.TimeUnit; import org.jspecify.annotations.NullMarked; @@ -41,22 +43,54 @@ class ManagedHttpJsonInterceptorChannel extends ManagedHttpJsonChannel { ManagedHttpJsonInterceptorChannel( ManagedHttpJsonChannel channel, HttpJsonClientInterceptor interceptor) { - super(); + super(true); this.channel = channel; this.interceptor = interceptor; } + @Override + public long getGeneration() { + return channel.getGeneration(); + } + @VisibleForTesting ManagedHttpJsonChannel getChannel() { return channel; } + @Override + String getEndpoint() { + return channel.getEndpoint(); + } + + @Override + @VisibleForTesting + HttpTransport getHttpTransport() { + return channel.getHttpTransport(); + } + + @Override + @VisibleForTesting + Executor getExecutor() { + return channel.getExecutor(); + } + @Override public HttpJsonClientCall newCall( ApiMethodDescriptor methodDescriptor, HttpJsonCallOptions callOptions) { return interceptor.interceptCall(methodDescriptor, callOptions, channel); } + @Override + public void refresh() { + channel.refresh(); + } + + @Override + public boolean shouldRefresh() { + return channel.shouldRefresh(); + } + @Override public synchronized void shutdown() { channel.shutdown(); diff --git a/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/RefreshingHttpJsonChannel.java b/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/RefreshingHttpJsonChannel.java new file mode 100644 index 000000000000..3d9557152845 --- /dev/null +++ b/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/RefreshingHttpJsonChannel.java @@ -0,0 +1,399 @@ +/* + * Copyright 2026 Google LLC + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are + * met: + * + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above + * copyright notice, this list of conditions and the following disclaimer + * in the documentation and/or other materials provided with the + * distribution. + * * Neither the name of Google LLC nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS + * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT + * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR + * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT + * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, + * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT + * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, + * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ +package com.google.api.gax.httpjson; + +import com.google.api.client.http.HttpTransport; +import com.google.api.core.InternalApi; +import com.google.api.gax.httpjson.ForwardingHttpJsonClientCall.SimpleForwardingHttpJsonClientCall; +import com.google.api.gax.httpjson.ForwardingHttpJsonClientCallListener.SimpleForwardingHttpJsonClientCallListener; +import com.google.api.gax.rpc.mtls.CertificateRotationTracker; +import com.google.api.gax.rpc.mtls.WorkloadCertificateUtils; +import com.google.common.annotations.VisibleForTesting; +import java.util.concurrent.CancellationException; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Supplier; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.jspecify.annotations.Nullable; + +/** + * An implementation of {@link ManagedHttpJsonChannel} that supports dynamic mTLS certificate + * rotation by thread-safely hot-swapping the underlying active HTTP/JSON channel while gracefully + * retiring older connections after all active in-flight requests complete. + */ +@InternalApi +public class RefreshingHttpJsonChannel extends ManagedHttpJsonChannel { + + private static final Logger LOG = Logger.getLogger(RefreshingHttpJsonChannel.class.getName()); + + private final CertificateRotationTracker rotationTracker; + private final Supplier channelFactory; + private final String workloadCertPath; + private final AtomicReference activeEntry; + // Keep track of all entries to properly await their termination + private final ConcurrentLinkedQueue allEntries = new ConcurrentLinkedQueue<>(); + private final Object refreshLock = new Object(); + private final AtomicLong generation = new AtomicLong(0); + + public RefreshingHttpJsonChannel( + Supplier channelFactory, String workloadCertPath) { + this(channelFactory.get(), channelFactory, workloadCertPath); + } + + public RefreshingHttpJsonChannel( + ManagedHttpJsonChannel initialChannel, + Supplier channelFactory, + String workloadCertPath) { + super(true); + this.channelFactory = channelFactory; + this.workloadCertPath = workloadCertPath; + ChannelEntry initial = new ChannelEntry(initialChannel); + this.activeEntry = new AtomicReference<>(initial); + this.allEntries.add(initial); + try { + this.rotationTracker = + new CertificateRotationTracker( + this::getWorkloadCertPath, this::getCertificateFingerprint); + } catch (Throwable t) { + initialChannel.shutdownNow(); + throw t; + } + } + + // Visible for testing + String getWorkloadCertPath() { + return workloadCertPath; + } + + // Visible for testing + String getCertificateFingerprint(String certPath) { + return WorkloadCertificateUtils.getCertificateFingerprint(certPath); + } + + @Override + public boolean shouldRefresh() { + return rotationTracker.shouldRefresh(); + } + + @Override + public void refresh() { + synchronized (refreshLock) { + if (isShutdown()) { + return; + } + String currentDiskFingerprint = rotationTracker.readDiskFingerprint(); + if (currentDiskFingerprint.isEmpty()) { + return; + } + + // Double-check inside refreshLock + if (rotationTracker.isAlreadyActive(currentDiskFingerprint)) { + LOG.fine( + "HTTP/JSON channel was already refreshed by a concurrent thread, skipping duplicate" + + " refresh"); + return; + } + + LOG.info("mTLS certificate rotation detected. Triggering HTTP/JSON channel pool refresh."); + + ChannelEntry newEntry = new ChannelEntry(channelFactory.get()); + allEntries.add(newEntry); + // Prune terminated entries after adding newEntry to ensure allEntries is never empty + allEntries.removeIf(entry -> entry != newEntry && entry.channel.isTerminated()); + + ChannelEntry oldEntry = activeEntry.getAndSet(newEntry); + rotationTracker.markRefreshed(currentDiskFingerprint); + generation.incrementAndGet(); + + if (oldEntry != null) { + oldEntry.requestShutdown(); + } + } + } + + @Override + public long getGeneration() { + return generation.get(); + } + + private ChannelEntry getRetainedEntry() { + while (true) { + ChannelEntry entry = activeEntry.get(); + if (entry.retain()) { + return entry; + } + if (entry == activeEntry.get()) { + throw new IllegalStateException("Channel has been shut down"); + } + } + } + + @Override + public HttpJsonClientCall newCall( + ApiMethodDescriptor methodDescriptor, HttpJsonCallOptions callOptions) { + ChannelEntry entry = getRetainedEntry(); + try { + HttpJsonClientCall delegateCall = + entry.channel.newCall(methodDescriptor, callOptions); + return new ReleasingHttpJsonClientCall<>(delegateCall, entry); + } catch (Throwable t) { + entry.release(); + throw t; + } + } + + @Override + java.util.concurrent.Executor getExecutor() { + return activeEntry.get().channel.getExecutor(); + } + + @VisibleForTesting + ManagedHttpJsonChannel getActiveChannel() { + return activeEntry.get().channel; + } + + private volatile boolean isShuttingDown = false; + + @Override + public void shutdown() { + synchronized (refreshLock) { + isShuttingDown = true; + for (ChannelEntry entry : allEntries) { + entry.requestShutdown(); + } + } + } + + @Override + public boolean isShutdown() { + return isShuttingDown; + } + + @Override + public boolean isTerminated() { + if (!isShuttingDown) { + return false; + } + for (ChannelEntry entry : allEntries) { + if (!entry.channel.isTerminated()) { + return false; + } + } + return true; + } + + @Override + public void shutdownNow() { + synchronized (refreshLock) { + isShuttingDown = true; + for (ChannelEntry entry : allEntries) { + entry.shutdownRequested.set(true); + entry.shutdownInitiated.set(true); + entry.channel.shutdownNow(); + } + } + } + + @VisibleForTesting + void invalidateDiskFingerprintCache() { + rotationTracker.invalidateCache(); + } + + @Override + public boolean awaitTermination(long duration, TimeUnit unit) throws InterruptedException { + long endNanos = System.nanoTime() + unit.toNanos(duration); + for (ChannelEntry entry : allEntries) { + if (entry.channel.isTerminated()) { + continue; + } + long remainingNanos = endNanos - System.nanoTime(); + if (remainingNanos <= 0) { + return false; + } + if (!entry.channel.awaitTermination(remainingNanos, TimeUnit.NANOSECONDS) + && !entry.channel.isTerminated()) { + return false; + } + } + return true; + } + + @Override + public void close() { + shutdown(); + } + + @Override + String getEndpoint() { + return activeEntry.get().channel.getEndpoint(); + } + + @Override + @VisibleForTesting + HttpTransport getHttpTransport() { + return activeEntry.get().channel.getHttpTransport(); + } + + /** Internal container to manage request reference-counting and graceful shutdown. */ + private static class ChannelEntry { + private final ManagedHttpJsonChannel channel; + private final AtomicInteger outstandingCalls = new AtomicInteger(0); + private final AtomicBoolean shutdownRequested = new AtomicBoolean(false); + private final AtomicBoolean shutdownInitiated = new AtomicBoolean(false); + + ChannelEntry(ManagedHttpJsonChannel channel) { + this.channel = channel; + } + + boolean retain() { + outstandingCalls.incrementAndGet(); + if (shutdownRequested.get()) { + release(); + return false; + } + return true; + } + + void release() { + int count = outstandingCalls.decrementAndGet(); + if (count < 0) { + LOG.warning("Channel entry reference count dropped below 0"); + } + // Must check outstandingCalls after shutdownRequested (in reverse order of retain()) to + // ensure mutual exclusion. + if (shutdownRequested.get() && outstandingCalls.get() == 0) { + shutdown(); + } + } + + void requestShutdown() { + shutdownRequested.set(true); + if (outstandingCalls.get() == 0) { + shutdown(); + } + } + + private void shutdown() { + if (shutdownInitiated.compareAndSet(false, true)) { + try { + channel.shutdown(); + } catch (Exception e) { + LOG.log(Level.WARNING, "Error shutting down retired HTTP/JSON channel", e); + } + } + } + } + + /** A client call decorator that decrements the entry counter upon call completion. */ + private static class ReleasingHttpJsonClientCall + extends SimpleForwardingHttpJsonClientCall { + + private final Object callLock = new Object(); + private volatile @Nullable CancellationException cancellationException; + private final ChannelEntry entry; + private final AtomicBoolean wasClosed = new AtomicBoolean(false); + private final AtomicBoolean wasReleased = new AtomicBoolean(false); + private final AtomicBoolean wasStarted = new AtomicBoolean(false); + + ReleasingHttpJsonClientCall(HttpJsonClientCall delegate, ChannelEntry entry) { + super(delegate); + this.entry = entry; + } + + @Override + public void start(Listener responseListener, HttpJsonMetadata requestHeaders) { + synchronized (callLock) { + if (!wasStarted.compareAndSet(false, true)) { + throw new IllegalStateException("Call is already started"); + } + if (cancellationException != null) { + if (wasReleased.compareAndSet(false, true)) { + entry.release(); + } + throw new IllegalStateException("Call is already cancelled", cancellationException); + } + try { + super.start( + new SimpleForwardingHttpJsonClientCallListener(responseListener) { + @Override + public void onClose(int statusCode, HttpJsonMetadata trailers) { + if (!wasClosed.compareAndSet(false, true)) { + return; + } + try { + super.onClose(statusCode, trailers); + } finally { + if (wasReleased.compareAndSet(false, true)) { + entry.release(); + } + } + } + }, + requestHeaders); + } catch (Throwable t) { + if (wasReleased.compareAndSet(false, true)) { + entry.release(); + } + throw t; + } + } + } + + @Override + public void cancel(@Nullable String message, @Nullable Throwable cause) { + boolean releaseImmediately = false; + try { + synchronized (callLock) { + this.cancellationException = new CancellationException(message); + if (!wasStarted.get()) { + releaseImmediately = true; + } + if (delegate() != null) { + super.cancel(message, cause); + } + } + } catch (Throwable t) { + if (!wasStarted.get()) { + releaseImmediately = true; + } + throw t; + } finally { + if (releaseImmediately && wasReleased.compareAndSet(false, true)) { + entry.release(); + } + } + } + } +} diff --git a/sdk-platform-java/gax-java/gax-httpjson/src/test/java/com/google/api/gax/httpjson/HttpJsonCallContextTest.java b/sdk-platform-java/gax-java/gax-httpjson/src/test/java/com/google/api/gax/httpjson/HttpJsonCallContextTest.java index 08044522e729..7ca2de0775c0 100644 --- a/sdk-platform-java/gax-java/gax-httpjson/src/test/java/com/google/api/gax/httpjson/HttpJsonCallContextTest.java +++ b/sdk-platform-java/gax-java/gax-httpjson/src/test/java/com/google/api/gax/httpjson/HttpJsonCallContextTest.java @@ -334,4 +334,54 @@ void testMergeOptions() { assertEquals(testContext2, mergedContext.getOption(contextKey2)); assertEquals(testContext3, mergedContext.getOption(contextKey3)); } + + @Test + void testWithChannelClearsStaleTransportChannel() { + ManagedHttpJsonChannel channel1 = + mock(ManagedHttpJsonChannel.class, Mockito.withSettings().withoutAnnotations()); + ManagedHttpJsonChannel channel2 = + mock(ManagedHttpJsonChannel.class, Mockito.withSettings().withoutAnnotations()); + + HttpJsonTransportChannel transportChannel1 = + HttpJsonTransportChannel.newBuilder().setManagedChannel(channel1).build(); + + HttpJsonCallContext context = + HttpJsonCallContext.createDefault().withTransportChannel(transportChannel1); + Truth.assertThat(context.getTransportChannel()).isSameInstanceAs(transportChannel1); + + // Retains transportChannel when setting same channel + Truth.assertThat(context.withChannel(channel1).getTransportChannel()) + .isSameInstanceAs(transportChannel1); + + // Clears transportChannel to null when setting null or a different channel + HttpJsonCallContext nullChannelContext = context.withChannel(null); + Truth.assertThat(nullChannelContext.getChannel()).isNull(); + Truth.assertThat(nullChannelContext.getTransportChannel()).isNull(); + Truth.assertThat(context.withChannel(channel2).getTransportChannel()).isNull(); + + // Merging a cleared context into default context preserves default context's transportChannel + HttpJsonCallContext mergedWithNullChannel = context.merge(nullChannelContext); + Truth.assertThat(mergedWithNullChannel.getChannel()).isSameInstanceAs(channel1); + Truth.assertThat(mergedWithNullChannel.getTransportChannel()) + .isSameInstanceAs(transportChannel1); + } + + @Test + void testMergeClearsStaleTransportChannel() { + ManagedHttpJsonChannel channel1 = + mock(ManagedHttpJsonChannel.class, Mockito.withSettings().withoutAnnotations()); + ManagedHttpJsonChannel channel2 = + mock(ManagedHttpJsonChannel.class, Mockito.withSettings().withoutAnnotations()); + + HttpJsonTransportChannel transportChannel1 = + HttpJsonTransportChannel.newBuilder().setManagedChannel(channel1).build(); + + HttpJsonCallContext context1 = + HttpJsonCallContext.createDefault().withTransportChannel(transportChannel1); + HttpJsonCallContext context2 = HttpJsonCallContext.createDefault().withChannel(channel2); + + HttpJsonCallContext merged = context1.merge(context2); + Truth.assertThat(merged.getChannel()).isSameInstanceAs(channel2); + Truth.assertThat(merged.getTransportChannel()).isNull(); + } } diff --git a/sdk-platform-java/gax-java/gax-httpjson/src/test/java/com/google/api/gax/httpjson/InstantiatingHttpJsonChannelProviderTest.java b/sdk-platform-java/gax-java/gax-httpjson/src/test/java/com/google/api/gax/httpjson/InstantiatingHttpJsonChannelProviderTest.java index 8c95c1d2e1c4..f1a542a7c75d 100644 --- a/sdk-platform-java/gax-java/gax-httpjson/src/test/java/com/google/api/gax/httpjson/InstantiatingHttpJsonChannelProviderTest.java +++ b/sdk-platform-java/gax-java/gax-httpjson/src/test/java/com/google/api/gax/httpjson/InstantiatingHttpJsonChannelProviderTest.java @@ -31,8 +31,9 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.mock; -import com.google.api.client.http.javanet.NetHttpTransport; +import com.google.api.gax.rpc.HeaderProvider; import com.google.api.gax.rpc.TransportChannelProvider; import com.google.api.gax.rpc.mtls.AbstractMtlsTransportChannelTest; import com.google.api.gax.rpc.mtls.CertificateBasedAccess; @@ -46,6 +47,7 @@ import java.util.concurrent.ScheduledThreadPoolExecutor; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.Mockito; class InstantiatingHttpJsonChannelProviderTest extends AbstractMtlsTransportChannelTest { @@ -55,9 +57,10 @@ class InstantiatingHttpJsonChannelProviderTest extends AbstractMtlsTransportChan @BeforeEach public void setup() throws IOException { - certificateBasedAccess = - new CertificateBasedAccess( - name -> name.equals("GOOGLE_API_USE_MTLS_ENDPOINT") ? "never" : "false"); + certificateBasedAccess = org.mockito.Mockito.mock(CertificateBasedAccess.class); + org.mockito.Mockito.when(certificateBasedAccess.getMtlsEndpointUsagePolicy()) + .thenReturn(CertificateBasedAccess.MtlsEndpointUsagePolicy.NEVER); + org.mockito.Mockito.when(certificateBasedAccess.useMtlsClientCertificate()).thenReturn(false); } @Test @@ -179,39 +182,197 @@ void managedChannelUsesCustomExecutor() throws IOException { instantiatingHttpJsonChannelProvider.getTransportChannel().shutdownNow(); } - @Override - protected Object getMtlsObjectFromTransportChannel( - MtlsProvider provider, CertificateBasedAccess certificateBasedAccess) + @Test + void managedChannelDoesNotShutdownCustomHttpTransport() throws IOException { + com.google.api.client.http.HttpTransport mockHttpTransport = + org.mockito.Mockito.mock(com.google.api.client.http.HttpTransport.class); + + InstantiatingHttpJsonChannelProvider provider = + InstantiatingHttpJsonChannelProvider.newBuilder() + .setEndpoint(DEFAULT_ENDPOINT) + .setHttpTransport(mockHttpTransport) + .setCertificateBasedAccess(certificateBasedAccess) + .build(); + provider = (InstantiatingHttpJsonChannelProvider) provider.withHeaders(DEFAULT_HEADER_MAP); + + HttpJsonTransportChannel httpJsonTransportChannel = provider.getTransportChannel(); + + // Verify custom transport is injected (direct ManagedHttpJsonChannel when workloadCertPath is + // null) + ManagedHttpJsonInterceptorChannel interceptorChannel = + (ManagedHttpJsonInterceptorChannel) httpJsonTransportChannel.getManagedChannel(); + ManagedHttpJsonInterceptorChannel managedHttpJsonChannel = + (ManagedHttpJsonInterceptorChannel) interceptorChannel.getChannel(); + ManagedHttpJsonChannel channel = managedHttpJsonChannel.getChannel(); + + assertThat(channel.getHttpTransport()).isEqualTo(mockHttpTransport); + + // Perform a shutdown + provider.getTransportChannel().shutdownNow(); + + // Verify that shutdown() was NOT called on the custom HttpTransport + org.mockito.Mockito.verify(mockHttpTransport, org.mockito.Mockito.never()).shutdown(); + } + + @Test + void channelCreation_withWorkloadCertPath_wrapsWithRefreshingHttpJsonChannel() throws IOException, GeneralSecurityException { - InstantiatingHttpJsonChannelProvider channelProvider = + Mockito.when(certificateBasedAccess.useMtlsClientCertificate()).thenReturn(true); + Mockito.when(certificateBasedAccess.getWorkloadCertPath()).thenReturn("fake/cert/path.json"); + com.google.auth.mtls.MtlsProvider mtlsProvider = + new com.google.api.gax.rpc.testing.FakeMtlsProvider( + com.google.api.gax.rpc.testing.FakeMtlsProvider.createTestMtlsKeyStore(), "", false); + + InstantiatingHttpJsonChannelProvider provider = InstantiatingHttpJsonChannelProvider.newBuilder() - .setEndpoint("localhost:8080") - .setMtlsProvider(provider) + .setEndpoint(DEFAULT_ENDPOINT) + .setMtlsProvider(mtlsProvider) .setCertificateBasedAccess(certificateBasedAccess) - .setHeaderProvider(Collections::emptyMap) - .setExecutor(Runnable::run) .build(); - NetHttpTransport transport = (NetHttpTransport) channelProvider.createHttpTransport(); - return (transport != null && transport.isMtls()) ? transport : null; + provider = (InstantiatingHttpJsonChannelProvider) provider.withHeaders(DEFAULT_HEADER_MAP); + + HttpJsonTransportChannel httpJsonTransportChannel = provider.getTransportChannel(); + + ManagedHttpJsonInterceptorChannel interceptorChannel = + (ManagedHttpJsonInterceptorChannel) httpJsonTransportChannel.getManagedChannel(); + ManagedHttpJsonInterceptorChannel managedHttpJsonChannel = + (ManagedHttpJsonInterceptorChannel) interceptorChannel.getChannel(); + assertThat(managedHttpJsonChannel.getChannel()).isInstanceOf(RefreshingHttpJsonChannel.class); + + provider.getTransportChannel().shutdownNow(); } @Test - void testCreateHttpTransport_returnsValidTransport() throws Exception { + void channelCreation_withCustomHttpTransport_ignoresWorkloadCertPathAndDoesNotWrap() + throws IOException { + com.google.api.client.http.HttpTransport mockHttpTransport = + org.mockito.Mockito.mock(com.google.api.client.http.HttpTransport.class); + + InstantiatingHttpJsonChannelProvider provider = + InstantiatingHttpJsonChannelProvider.newBuilder() + .setEndpoint(DEFAULT_ENDPOINT) + .setHttpTransport(mockHttpTransport) + .setCertificateBasedAccess(certificateBasedAccess) + .build(); + provider = (InstantiatingHttpJsonChannelProvider) provider.withHeaders(DEFAULT_HEADER_MAP); + + HttpJsonTransportChannel httpJsonTransportChannel = provider.getTransportChannel(); + + ManagedHttpJsonInterceptorChannel interceptorChannel = + (ManagedHttpJsonInterceptorChannel) httpJsonTransportChannel.getManagedChannel(); + ManagedHttpJsonInterceptorChannel managedHttpJsonChannel = + (ManagedHttpJsonInterceptorChannel) interceptorChannel.getChannel(); + assertThat(managedHttpJsonChannel.getChannel()) + .isNotInstanceOf(RefreshingHttpJsonChannel.class); + Mockito.verify(certificateBasedAccess, Mockito.never()).getWorkloadCertPath(); + + httpJsonTransportChannel.shutdownNow(); + } + + @Test + void getTransportChannel_whenMtlsKeyStoreThrowsIOException_throwsCheckedIOException() + throws Exception { + Mockito.when(certificateBasedAccess.useMtlsClientCertificate()).thenReturn(true); + Mockito.when(certificateBasedAccess.getWorkloadCertPath()).thenReturn("fake/cert/path.json"); + MtlsProvider failingMtlsProvider = + Mockito.mock(MtlsProvider.class, Mockito.withSettings().withoutAnnotations()); + Mockito.when(failingMtlsProvider.getKeyStore()) + .thenThrow(new IOException("Simulated keystore read failure")); + + InstantiatingHttpJsonChannelProvider provider = + InstantiatingHttpJsonChannelProvider.newBuilder() + .setEndpoint(DEFAULT_ENDPOINT) + .setMtlsProvider(failingMtlsProvider) + .setCertificateBasedAccess(certificateBasedAccess) + .build(); + final InstantiatingHttpJsonChannelProvider finalProvider = + (InstantiatingHttpJsonChannelProvider) provider.withHeaders(DEFAULT_HEADER_MAP); + + // Must throw checked IOException directly (not wrapped in RuntimeException) + IOException thrown = + org.junit.jupiter.api.Assertions.assertThrows( + IOException.class, finalProvider::getTransportChannel); + assertThat(thrown).hasMessageThat().contains("Simulated keystore read failure"); + } + + @Test + void getTransportChannel_whenMtlsActiveAndKeyStoreNull_throwsIOException() { + Mockito.when(certificateBasedAccess.useMtlsClientCertificate()).thenReturn(true); + com.google.auth.mtls.MtlsProvider providerWithNullKeyStore = + new com.google.api.gax.rpc.testing.FakeMtlsProvider(null, "", false); + + InstantiatingHttpJsonChannelProvider provider = + InstantiatingHttpJsonChannelProvider.newBuilder() + .setEndpoint(DEFAULT_ENDPOINT) + .setMtlsProvider(providerWithNullKeyStore) + .setCertificateBasedAccess(certificateBasedAccess) + .build(); + InstantiatingHttpJsonChannelProvider finalProvider = + (InstantiatingHttpJsonChannelProvider) provider.withHeaders(DEFAULT_HEADER_MAP); + + IOException thrown = + org.junit.jupiter.api.Assertions.assertThrows( + IOException.class, finalProvider::getTransportChannel); + assertThat(thrown).hasMessageThat().contains("Failed to initialize mTLS HttpTransport"); + } + + @Test + void createHttpTransport_withMtlsAndConscrypt_configuresSecurityProvider() + throws IOException, GeneralSecurityException { + Mockito.when(certificateBasedAccess.useMtlsClientCertificate()).thenReturn(true); + com.google.auth.mtls.MtlsProvider provider = + new com.google.api.gax.rpc.testing.FakeMtlsProvider( + com.google.api.gax.rpc.testing.FakeMtlsProvider.createTestMtlsKeyStore(), "", false); + InstantiatingHttpJsonChannelProvider channelProvider = InstantiatingHttpJsonChannelProvider.newBuilder() - .setEndpoint("localhost:8080") - .setHeaderProvider(Collections::emptyMap) - .setExecutor(Runnable::run) + .setEndpoint(DEFAULT_ENDPOINT) + .setMtlsProvider(provider) + .setCertificateBasedAccess(certificateBasedAccess) .build(); - NetHttpTransport transport = (NetHttpTransport) channelProvider.createHttpTransport(); + + com.google.api.client.http.HttpTransport transport = channelProvider.createHttpTransport(); assertThat(transport).isNotNull(); + assertThat(transport).isInstanceOf(com.google.api.client.http.javanet.NetHttpTransport.class); } @Test - void testConfigureConscryptSecurityProvider_returnsConfiguredBuilder() { - NetHttpTransport.Builder builder = new NetHttpTransport.Builder(); - NetHttpTransport.Builder result = - HttpJsonConscryptUtils.configureConscryptSecurityProvider(builder); - assertThat(result).isSameInstanceAs(builder); + void createHttpTransport_whenMtlsProviderNullOrNotUsingClientCert_returnsNull() + throws IOException, GeneralSecurityException { + InstantiatingHttpJsonChannelProvider nullMtlsProviderChannelProvider = + InstantiatingHttpJsonChannelProvider.newBuilder() + .setEndpoint(DEFAULT_ENDPOINT) + .setMtlsProvider(null) + .setCertificateBasedAccess(certificateBasedAccess) + .build(); + assertThat(nullMtlsProviderChannelProvider.createHttpTransport()).isNull(); + + Mockito.when(certificateBasedAccess.useMtlsClientCertificate()).thenReturn(false); + com.google.auth.mtls.MtlsProvider provider = + new com.google.api.gax.rpc.testing.FakeMtlsProvider( + com.google.api.gax.rpc.testing.FakeMtlsProvider.createTestMtlsKeyStore(), "", false); + InstantiatingHttpJsonChannelProvider disabledMtlsChannelProvider = + InstantiatingHttpJsonChannelProvider.newBuilder() + .setEndpoint(DEFAULT_ENDPOINT) + .setMtlsProvider(provider) + .setCertificateBasedAccess(certificateBasedAccess) + .build(); + assertThat(disabledMtlsChannelProvider.createHttpTransport()).isNull(); + } + + @Override + protected Object getMtlsObjectFromTransportChannel( + MtlsProvider provider, CertificateBasedAccess certificateBasedAccess) + throws IOException, GeneralSecurityException { + InstantiatingHttpJsonChannelProvider channelProvider = + InstantiatingHttpJsonChannelProvider.newBuilder() + .setEndpoint("localhost:8080") + .setMtlsProvider(provider) + .setCertificateBasedAccess(certificateBasedAccess) + .setHeaderProvider( + mock(HeaderProvider.class, Mockito.withSettings().withoutAnnotations())) + .setExecutor(mock(Executor.class)) + .build(); + return channelProvider.createHttpTransport(); } } diff --git a/sdk-platform-java/gax-java/gax-httpjson/src/test/java/com/google/api/gax/httpjson/RefreshingHttpJsonChannelTest.java b/sdk-platform-java/gax-java/gax-httpjson/src/test/java/com/google/api/gax/httpjson/RefreshingHttpJsonChannelTest.java new file mode 100644 index 000000000000..aaaacd3c53db --- /dev/null +++ b/sdk-platform-java/gax-java/gax-httpjson/src/test/java/com/google/api/gax/httpjson/RefreshingHttpJsonChannelTest.java @@ -0,0 +1,621 @@ +/* + * Copyright 2026 Google LLC + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are + * met: + * + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above + * copyright notice, this list of conditions and the following disclaimer + * in the documentation and/or other materials provided with the + * distribution. + * * Neither the name of Google LLC nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS + * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT + * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR + * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT + * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, + * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT + * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, + * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ +package com.google.api.gax.httpjson; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Supplier; +import javax.annotation.Nullable; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class RefreshingHttpJsonChannelTest { + private static class FakeHttpJsonClientCall + extends HttpJsonClientCall { + protected Listener listener; + + @Override + public void start(Listener responseListener, HttpJsonMetadata requestHeaders) { + this.listener = responseListener; + } + + @Override + public void request(int numMessages) {} + + @Override + public void cancel(@Nullable String message, @Nullable Throwable cause) {} + + @Override + public void sendMessage(RequestT message) {} + + @Override + public void halfClose() {} + } + + private static class FakeManagedHttpJsonChannel extends ManagedHttpJsonChannel { + private volatile boolean isShutdown = false; + private volatile boolean isTerminated = false; + private HttpJsonClientCall nextCall = null; + + @Override + String getEndpoint() { + return "https://fake.endpoint:443"; + } + + @Override + public void shutdown() { + isShutdown = true; + } + + @Override + public void shutdownNow() { + isShutdown = true; + isTerminated = true; + } + + @Override + public boolean isShutdown() { + return isShutdown; + } + + @Override + public boolean isTerminated() { + return isTerminated; + } + + @Override + public boolean awaitTermination(long duration, TimeUnit unit) { + return isTerminated; + } + + @Override + @SuppressWarnings("unchecked") + public HttpJsonClientCall newCall( + ApiMethodDescriptor methodDescriptor, + HttpJsonCallOptions callOptions) { + if (nextCall != null) { + return (HttpJsonClientCall) nextCall; + } + return new FakeHttpJsonClientCall<>(); + } + } + + private AtomicInteger channelFactoryCount; + private FakeManagedHttpJsonChannel lastCreatedChannel; + private String testCertPath = "/fake/path"; + private String testFingerprint = "fingerprint1"; + private boolean shouldThrowOnFactory = false; + private List createdChannels; + + private Supplier channelFactory = + () -> { + if (shouldThrowOnFactory) { + throw new RuntimeException("Simulated factory failure"); + } + channelFactoryCount.incrementAndGet(); + lastCreatedChannel = new FakeManagedHttpJsonChannel(); + return lastCreatedChannel; + }; + + @BeforeEach + void setUp() { + channelFactoryCount = new AtomicInteger(0); + testCertPath = "/fake/path"; + testFingerprint = "fingerprint1"; + shouldThrowOnFactory = false; + createdChannels = new ArrayList<>(); + } + + @AfterEach + void tearDown() { + for (RefreshingHttpJsonChannel channel : createdChannels) { + channel.shutdownNow(); + } + } + + private RefreshingHttpJsonChannel createTestChannel() { + RefreshingHttpJsonChannel ch = + new RefreshingHttpJsonChannel(channelFactory, "fake/cert/path.json") { + @Override + String getWorkloadCertPath() { + return testCertPath; + } + + @Override + String getCertificateFingerprint(String certPath) { + return testFingerprint; + } + }; + createdChannels.add(ch); + return ch; + } + + @Test + void testShouldRefreshNullCertPath() { + testCertPath = null; + RefreshingHttpJsonChannel channel = createTestChannel(); + assertFalse(channel.shouldRefresh()); + } + + @Test + void testShouldRefreshFalseWhenUnchanged() throws InterruptedException { + RefreshingHttpJsonChannel channel = createTestChannel(); + + channel.invalidateDiskFingerprintCache(); // Invalidate 1-second cache + assertFalse(channel.shouldRefresh()); + } + + @Test + void testShouldRefreshTrueWhenChanged() throws InterruptedException { + RefreshingHttpJsonChannel channel = createTestChannel(); + + channel.invalidateDiskFingerprintCache(); // Invalidate 1-second cache + + // Simulate disk fingerprint changing + testFingerprint = "fingerprint2"; + + assertTrue(channel.shouldRefresh()); + } + + @Test + void shouldRefresh_doesNotCacheNegativeResultAndDetectsSubsequentRotationImmediately() { + RefreshingHttpJsonChannel channel = createTestChannel(); + + // First check returns false (unchanged fingerprint) + assertFalse(channel.shouldRefresh()); + + // Immediately change fingerprint WITHOUT invalidating cache + testFingerprint = "fingerprint2"; + + // Must immediately detect rotation because negative/unchanged checks are not cached for 1s + assertTrue(channel.shouldRefresh()); + + // Refresh updates activeCertFingerprint and clears cache + channel.refresh(); + assertFalse(channel.shouldRefresh()); + } + + @Test + void testRefreshSwapsChannel() throws InterruptedException { + RefreshingHttpJsonChannel channel = createTestChannel(); + FakeManagedHttpJsonChannel firstChannel = lastCreatedChannel; + assertEquals(1, channelFactoryCount.get()); + + channel.invalidateDiskFingerprintCache(); // Invalidate 1-second cache + + // Change fingerprint + testFingerprint = "fingerprint2"; + + // Act + channel.refresh(); + + // Verify a new channel was created and the old one retired + assertEquals(2, channelFactoryCount.get()); + FakeManagedHttpJsonChannel secondChannel = lastCreatedChannel; + + // The old channel should receive a shutdown request immediately since there are no active calls + assertTrue(firstChannel.isShutdown()); + assertFalse(secondChannel.isShutdown()); + } + + @Test + void testRefreshKeepsInFlightChannelsAlive() throws InterruptedException { + RefreshingHttpJsonChannel channel = createTestChannel(); + FakeManagedHttpJsonChannel firstChannel = lastCreatedChannel; + + // Simulate an in-flight API call + FakeHttpJsonClientCall fakeCall = new FakeHttpJsonClientCall<>(); + firstChannel.nextCall = fakeCall; + + HttpJsonClientCall activeCall = channel.newCall(null, null); + + channel.invalidateDiskFingerprintCache(); // Invalidate 1-second cache + + // Change fingerprint & refresh + testFingerprint = "fingerprint2"; + + channel.refresh(); + + // Verify a new channel was created + assertEquals(2, channelFactoryCount.get()); + + // IMPORTANT: The first channel should NOT be shut down yet because of the active call! + assertFalse(firstChannel.isShutdown()); + + // Now start the call + activeCall.start(new HttpJsonClientCall.Listener() {}, null); + + assertNotNull(fakeCall.listener); + + // Fire onClose + fakeCall.listener.onClose(0, null); + + // FIRST CHANNEL SHOULD BE SHUT DOWN NOW! + assertTrue(firstChannel.isShutdown()); + } + + @Test + void testCancelBeforeStartReleasesChannelEntry() { + RefreshingHttpJsonChannel channel = createTestChannel(); + FakeManagedHttpJsonChannel firstChannel = lastCreatedChannel; + + HttpJsonClientCall activeCall = channel.newCall(null, null); + + channel.invalidateDiskFingerprintCache(); + testFingerprint = "fingerprint2"; + channel.refresh(); + + // Because activeCall was created, the old channel should NOT be shut down yet + assertFalse(firstChannel.isShutdown()); + + // Cancel before start() is called + activeCall.cancel("Cancelled early", null); + + // Because cancel() safely released the entry, the old channel should now be shut down! + assertTrue(firstChannel.isShutdown()); + } + + @Test + void testRefreshDoesNotSpawnChannelWhenShutdown() throws InterruptedException { + RefreshingHttpJsonChannel channel = createTestChannel(); + FakeManagedHttpJsonChannel firstChannel = lastCreatedChannel; + assertEquals(1, channelFactoryCount.get()); + + // Simulate that the channel pool is shut down. + channel.shutdown(); + firstChannel.shutdown(); + + channel.invalidateDiskFingerprintCache(); // Invalidate 1-second cache + + // Change fingerprint + testFingerprint = "fingerprint2"; + + // Act + channel.refresh(); + + // Verify no new channel was spawned + assertEquals(1, channelFactoryCount.get()); + } + + @Test + void testRefreshFactoryExceptionDoesNotWedgeFingerprint() throws InterruptedException { + RefreshingHttpJsonChannel channel = createTestChannel(); + assertEquals(1, channelFactoryCount.get()); + + shouldThrowOnFactory = true; + channel.invalidateDiskFingerprintCache(); // Invalidate 1-second cache + testFingerprint = "fingerprint2"; + + assertThrows(RuntimeException.class, channel::refresh); + + // Because factory threw, activeCertFingerprint should NOT be updated to fingerprint2 + // Therefore shouldRefresh() should still return true + assertTrue(channel.shouldRefresh()); + + shouldThrowOnFactory = false; + channel.refresh(); + assertEquals(2, channelFactoryCount.get()); + assertFalse(channel.shouldRefresh()); + } + + @Test + void testShutdownNowSetsIsShutdown() { + RefreshingHttpJsonChannel channel = createTestChannel(); + assertFalse(channel.isShutdown()); + + channel.shutdownNow(); + + assertTrue(channel.isShutdown()); + } + + @Test + void testAwaitTerminationZeroTimeoutOnTerminatedChannelReturnsTrue() throws InterruptedException { + RefreshingHttpJsonChannel channel = createTestChannel(); + FakeManagedHttpJsonChannel firstChannel = lastCreatedChannel; + firstChannel.isTerminated = true; + + channel.shutdown(); + assertTrue(channel.awaitTermination(0, TimeUnit.MILLISECONDS)); + } + + @Test + void testChannelDelegationMethods() { + RefreshingHttpJsonChannel channel = createTestChannel(); + FakeManagedHttpJsonChannel firstChannel = lastCreatedChannel; + + assertEquals(firstChannel.getEndpoint(), channel.getEndpoint()); + assertEquals(firstChannel.getHttpTransport(), channel.getHttpTransport()); + assertEquals(firstChannel.getExecutor(), channel.getExecutor()); + } + + @Test + void testNewCallAfterShutdownNowThrowsIllegalStateException() { + RefreshingHttpJsonChannel channel = createTestChannel(); + channel.shutdownNow(); + + assertThrows( + IllegalStateException.class, + () -> channel.newCall(null, null), + "Channel has been shut down"); + } + + @Test + void testConcurrentNewCallDuringRefresh() throws InterruptedException { + RefreshingHttpJsonChannel channel = createTestChannel(); + int threadCount = 10; + java.util.concurrent.ExecutorService executorService = + java.util.concurrent.Executors.newFixedThreadPool(threadCount); + java.util.concurrent.CountDownLatch latch = + new java.util.concurrent.CountDownLatch(threadCount); + java.util.concurrent.atomic.AtomicInteger successCount = + new java.util.concurrent.atomic.AtomicInteger(0); + + for (int i = 0; i < threadCount; i++) { + executorService.submit( + () -> { + try { + channel.newCall(null, null); + successCount.incrementAndGet(); + } finally { + latch.countDown(); + } + }); + } + + channel.invalidateDiskFingerprintCache(); + testFingerprint = "fingerprint2"; + channel.refresh(); + + latch.await(5, TimeUnit.SECONDS); + executorService.shutdown(); + + assertEquals(threadCount, successCount.get()); + } + + @Test + void testGenerationIncrementAndLifecycleOnDelegatingWrapper() throws Exception { + RefreshingHttpJsonChannel channel = createTestChannel(); + assertEquals(0, channel.getGeneration()); + + channel.invalidateDiskFingerprintCache(); + testFingerprint = "fingerprint2"; + channel.refresh(); + + assertEquals(1, channel.getGeneration()); + + // Verify lifecycle methods on delegating wrapper do not throw NullPointerException + assertFalse(channel.isShutdown()); + assertFalse(channel.isTerminated()); + channel.shutdown(); + assertTrue(channel.isShutdown()); + channel.shutdownNow(); + assertTrue(channel.isTerminated()); + assertTrue(channel.awaitTermination(1, TimeUnit.SECONDS)); + } + + @Test + void testNewCall_whenDelegateThrowsError_releasesEntryAndShutsDownRetiredChannel() { + RefreshingHttpJsonChannel channel = createTestChannel(); + FakeManagedHttpJsonChannel firstChannel = lastCreatedChannel; + firstChannel.nextCall = null; + // Configure firstChannel to throw an Error on newCall + FakeManagedHttpJsonChannel throwingChannel = + new FakeManagedHttpJsonChannel() { + @Override + public HttpJsonClientCall newCall( + ApiMethodDescriptor methodDescriptor, + HttpJsonCallOptions callOptions) { + throw new LinkageError("Simulated native error in newCall"); + } + }; + channelFactory = + () -> { + channelFactoryCount.incrementAndGet(); + lastCreatedChannel = throwingChannel; + return throwingChannel; + }; + RefreshingHttpJsonChannel testChannel = createTestChannel(); + + assertThrows(LinkageError.class, () -> testChannel.newCall(null, null)); + + // Refresh should immediately shut down throwingChannel since ref count returned to 0 + testFingerprint = "fingerprint2"; + testChannel.refresh(); + assertTrue(throwingChannel.isShutdown()); + } + + @Test + void testStart_whenDelegateThrowsError_releasesEntryAndShutsDownRetiredChannel() { + RefreshingHttpJsonChannel channel = createTestChannel(); + FakeManagedHttpJsonChannel firstChannel = lastCreatedChannel; + firstChannel.nextCall = + new FakeHttpJsonClientCall() { + @Override + public void start(Listener responseListener, HttpJsonMetadata requestHeaders) { + throw new AssertionError("Simulated Error in start"); + } + }; + + HttpJsonClientCall call = channel.newCall(null, null); + testFingerprint = "fingerprint2"; + channel.refresh(); + assertFalse(firstChannel.isShutdown()); + + assertThrows( + AssertionError.class, () -> call.start(new HttpJsonClientCall.Listener() {}, null)); + assertTrue(firstChannel.isShutdown()); + } + + @Test + void testCancel_whenDelegateThrowsException_releasesEntryAndShutsDownRetiredChannel() { + RefreshingHttpJsonChannel channel = createTestChannel(); + FakeManagedHttpJsonChannel firstChannel = lastCreatedChannel; + firstChannel.nextCall = + new FakeHttpJsonClientCall() { + @Override + public void cancel(String message, Throwable cause) { + throw new RuntimeException("Simulated cancel failure"); + } + }; + + HttpJsonClientCall call = channel.newCall(null, null); + testFingerprint = "fingerprint2"; + channel.refresh(); + assertFalse(firstChannel.isShutdown()); + + assertThrows(RuntimeException.class, () -> call.cancel("cancel", null)); + assertTrue(firstChannel.isShutdown()); + } + + @Test + void testConcurrentStartAndCancel_neverLeaksOrDoubleReleasesEntry() throws Exception { + RefreshingHttpJsonChannel channel = createTestChannel(); + FakeManagedHttpJsonChannel firstChannel = lastCreatedChannel; + + int iterations = 100; + java.util.concurrent.ExecutorService executor = + java.util.concurrent.Executors.newFixedThreadPool(2); + try { + for (int i = 0; i < iterations; i++) { + firstChannel.nextCall = + new FakeHttpJsonClientCall() { + private boolean closed = false; + + @Override + public synchronized void start( + Listener responseListener, HttpJsonMetadata requestHeaders) { + if (closed) { + // Models HttpJsonClientCallImpl returning early when closed + return; + } + super.start(responseListener, requestHeaders); + } + + @Override + public synchronized void cancel(String message, Throwable cause) { + closed = true; + if (listener != null) { + listener.onClose(499, null); + } + } + }; + + HttpJsonClientCall call = channel.newCall(null, null); + java.util.concurrent.CyclicBarrier barrier = new java.util.concurrent.CyclicBarrier(2); + java.util.concurrent.Future f1 = + executor.submit( + () -> { + try { + barrier.await(); + call.start(new HttpJsonClientCall.Listener() {}, null); + } catch (Exception ignored) { + } + }); + java.util.concurrent.Future f2 = + executor.submit( + () -> { + try { + barrier.await(); + call.cancel("cancel", null); + } catch (Exception ignored) { + } + }); + f1.get(5, TimeUnit.SECONDS); + f2.get(5, TimeUnit.SECONDS); + } + } finally { + executor.shutdownNow(); + } + + testFingerprint = "fingerprint2"; + channel.refresh(); + assertTrue(firstChannel.isShutdown()); + } + + @Test + void cancel_whenStartedAndSuperCancelThrows_doesNotReleasePrematurelyUntilOnClose() { + RefreshingHttpJsonChannel channel = createTestChannel(); + FakeManagedHttpJsonChannel firstChannel = lastCreatedChannel; + FakeHttpJsonClientCall delegateCall = + new FakeHttpJsonClientCall() { + @Override + public void cancel(String message, Throwable cause) { + throw new RuntimeException("Simulated cancel failure"); + } + }; + firstChannel.nextCall = delegateCall; + + HttpJsonClientCall call = channel.newCall(null, null); + call.start(new HttpJsonClientCall.Listener() {}, null); + + assertThrows(RuntimeException.class, () -> call.cancel("abort", null)); + + // Rotate pool while call is still active (onClose hasn't fired yet): + // firstChannel must NOT be shut down yet because call is still active + testFingerprint = "fingerprint2"; + channel.refresh(); + assertFalse(firstChannel.isShutdown()); + + // Once onClose fires, entry is released and firstChannel shuts down + delegateCall.listener.onClose(200, null); + assertTrue(firstChannel.isShutdown()); + } + + @Test + void start_whenCalledTwice_throwsIllegalStateExceptionAndDoesNotReleaseFirstCallEntry() { + RefreshingHttpJsonChannel channel = createTestChannel(); + FakeManagedHttpJsonChannel firstChannel = lastCreatedChannel; + FakeHttpJsonClientCall delegateCall = new FakeHttpJsonClientCall<>(); + firstChannel.nextCall = delegateCall; + + HttpJsonClientCall call = channel.newCall(null, null); + call.start(new HttpJsonClientCall.Listener() {}, null); + + assertThrows( + IllegalStateException.class, + () -> call.start(new HttpJsonClientCall.Listener() {}, null)); + + testFingerprint = "fingerprint2"; + channel.refresh(); + assertFalse(firstChannel.isShutdown()); + + delegateCall.listener.onClose(200, null); + assertTrue(firstChannel.isShutdown()); + } +} diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/retrying/RetrySettings.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/retrying/RetrySettings.java index d69fd310c2c0..e6729459b775 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/retrying/RetrySettings.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/retrying/RetrySettings.java @@ -189,7 +189,7 @@ public final org.threeten.bp.Duration getInitialRpcTimeout() { * connection has been terminated). * *

{@link #getTotalTimeout()} caps how long the logic should keep trying the RPC until it gives - * up completely. If {@link #getTotalTimeout()} is set, initialRpcTimeout should be <= + * up completely. If {@link #getTotalTimeout()} is set, initialRpcTimeout should be <= * totalTimeout. * *

If there are no configurations, Retries have the default initial RPC timeout value of {@code @@ -356,7 +356,7 @@ public final Builder setInitialRpcTimeout(org.threeten.bp.Duration initialTimeou * the connection has been terminated). * *

{@link #getTotalTimeout()} caps how long the logic should keep trying the RPC until it - * gives up completely. If {@link #getTotalTimeout()} is set, initialRpcTimeout should be <= + * gives up completely. If {@link #getTotalTimeout()} is set, initialRpcTimeout should be <= * totalTimeout. * *

If there are no configurations, Retries have the default initial RPC timeout value of @@ -491,7 +491,7 @@ public final org.threeten.bp.Duration getInitialRpcTimeout() { * the connection has been terminated). * *

{@link #getTotalTimeout()} caps how long the logic should keep trying the RPC until it - * gives up completely. If {@link #getTotalTimeout()} is set, initialRpcTimeout should be <= + * gives up completely. If {@link #getTotalTimeout()} is set, initialRpcTimeout should be <= * totalTimeout. * *

If there are no configurations, Retries have the default initial RPC timeout value of diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ApiCallContext.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ApiCallContext.java index 67b5e5b285d1..36d6fd966e55 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ApiCallContext.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ApiCallContext.java @@ -42,7 +42,6 @@ import java.util.Map; import java.util.Set; import javax.annotation.Nonnull; -import org.jspecify.annotations.NullMarked; import org.jspecify.annotations.Nullable; /** @@ -55,7 +54,6 @@ * *

This is transport specific and each transport has an implementation with its own options. */ -@NullMarked @InternalExtensionOnly public interface ApiCallContext extends RetryingContext { @@ -65,6 +63,18 @@ public interface ApiCallContext extends RetryingContext { /** Returns a new ApiCallContext with the given channel set. */ ApiCallContext withTransportChannel(TransportChannel channel); + /** + * Returns the {@link TransportChannel} associated with this call context, or {@code null} if none + * is set. + * + *

Note: By default, this method returns {@code null}. If an implementation does not override + * this method, automatic mTLS certificate rotation and channel refreshing in retrying callables + * will be disabled. + */ + default TransportChannel getTransportChannel() { + return null; + } + /** Returns a new ApiCallContext with the given Endpoint Context. */ ApiCallContext withEndpointContext(EndpointContext endpointContext); diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ApiResultRetryAlgorithm.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ApiResultRetryAlgorithm.java index 7d04d38d2605..010e6462d09d 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ApiResultRetryAlgorithm.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ApiResultRetryAlgorithm.java @@ -30,16 +30,52 @@ package com.google.api.gax.rpc; import com.google.api.gax.retrying.BasicResultRetryAlgorithm; +import com.google.api.gax.retrying.RetrySettings; import com.google.api.gax.retrying.RetryingContext; +import com.google.api.gax.retrying.TimedAttemptSettings; import org.jspecify.annotations.NullMarked; +import org.jspecify.annotations.Nullable; /* Package-private for internal use. */ @NullMarked class ApiResultRetryAlgorithm extends BasicResultRetryAlgorithm { - /** Returns true if previousThrowable is an {@link ApiException} that is retryable. */ @Override - public boolean shouldRetry(Throwable previousThrowable, ResponseT previousResponse) { + public @Nullable TimedAttemptSettings createNextAttempt( + @Nullable Throwable previousThrowable, + @Nullable ResponseT previousResponse, + TimedAttemptSettings previousSettings) { + return createNextAttempt(null, previousThrowable, previousResponse, previousSettings); + } + + @Override + public @Nullable TimedAttemptSettings createNextAttempt( + @Nullable RetryingContext context, + @Nullable Throwable previousThrowable, + @Nullable ResponseT previousResponse, + TimedAttemptSettings previousSettings) { + if (previousThrowable instanceof UnauthenticatedException + && ((UnauthenticatedException) previousThrowable).isRetryable() + && previousSettings.getOverallAttemptCount() == previousSettings.getAttemptCount()) { + RetrySettings globalSettings = previousSettings.getGlobalSettings(); + if (globalSettings.getMaxAttempts() == 0 + && globalSettings.getTotalTimeoutDuration().isZero()) { + globalSettings = globalSettings.toBuilder().setMaxAttempts(1).build(); + } + return previousSettings.toBuilder() + .setGlobalSettings(globalSettings) + .setRetryDelayDuration(java.time.Duration.ZERO) + .setRandomizedRetryDelayDuration(java.time.Duration.ZERO) + .setAttemptCount(previousSettings.getAttemptCount()) + .setOverallAttemptCount(previousSettings.getOverallAttemptCount() + 1) + .build(); + } + return null; + } + + @Override + public boolean shouldRetry( + @Nullable Throwable previousThrowable, @Nullable ResponseT previousResponse) { return (previousThrowable instanceof ApiException) && ((ApiException) previousThrowable).isRetryable(); } @@ -52,7 +88,15 @@ public boolean shouldRetry(Throwable previousThrowable, ResponseT previousRespon */ @Override public boolean shouldRetry( - RetryingContext context, Throwable previousThrowable, ResponseT previousResponse) { + RetryingContext context, + @Nullable Throwable previousThrowable, + @Nullable ResponseT previousResponse) { + // Check UnauthenticatedException retryability first to ensure mTLS certificate + // rotation retries take precedence over static method retry codes. + if (previousThrowable instanceof UnauthenticatedException + && ((UnauthenticatedException) previousThrowable).isRetryable()) { + return true; + } if (context.getRetryableCodes() != null) { // Ignore the isRetryable() value of the throwable if the RetryingContext has a specific list // of codes that should be retried. diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/AttemptCallable.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/AttemptCallable.java index 897e1f2dae8e..a7211fb46d3b 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/AttemptCallable.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/AttemptCallable.java @@ -35,6 +35,8 @@ import com.google.api.gax.retrying.RetryingFuture; import com.google.common.base.Preconditions; import java.util.concurrent.Callable; +import java.util.logging.Level; +import java.util.logging.Logger; import org.jspecify.annotations.NullMarked; /** @@ -48,6 +50,7 @@ */ @NullMarked class AttemptCallable implements Callable { + private static final Logger LOG = Logger.getLogger(AttemptCallable.class.getName()); private final UnaryCallable callable; private final RequestT request; private final ApiCallContext originalCallContext; @@ -85,8 +88,51 @@ public ResponseT call() { .getTracer() .attemptStarted(request, externalFuture.getAttemptSettings().getOverallAttemptCount()); + TransportChannel transportChannel = callContext.getTransportChannel(); + final long attemptGeneration = + transportChannel != null ? transportChannel.getGeneration() : 0; + ApiFuture internalFuture = callable.futureCall(request, callContext); - externalFuture.setAttemptFuture(internalFuture); + final ApiCallContext finalContext = callContext; + ApiFuture mappedFuture = + ApiFutures.catching( + internalFuture, + UnauthenticatedException.class, + unauthenticatedException -> { + TransportChannel channel = finalContext.getTransportChannel(); + if (channel != null) { + if (channel.shouldRefresh()) { + try { + channel.refresh(); + } catch (Exception e) { + LOG.log( + Level.WARNING, + "Failed to refresh transport channel after authentication error", + e); + } + } + boolean shouldRetry = channel.getGeneration() > attemptGeneration; + + if (shouldRetry) { + UnauthenticatedException newEx = + new UnauthenticatedException( + unauthenticatedException.getMessage(), + unauthenticatedException.getCause(), + unauthenticatedException.getStatusCode(), + true, // isRetryable = true + unauthenticatedException.getErrorDetails()); + newEx.setStackTrace(unauthenticatedException.getStackTrace()); + for (Throwable suppressed : unauthenticatedException.getSuppressed()) { + newEx.addSuppressed(suppressed); + } + throw newEx; + } + } + throw unauthenticatedException; + }, + com.google.common.util.concurrent.MoreExecutors.directExecutor()); + + externalFuture.setAttemptFuture(mappedFuture); } catch (Throwable e) { externalFuture.setAttemptFuture(ApiFutures.immediateFailedFuture(e)); } diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/BidiStreamingCallable.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/BidiStreamingCallable.java index 97e22d6ee41f..c8fe6b2a59d3 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/BidiStreamingCallable.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/BidiStreamingCallable.java @@ -29,6 +29,8 @@ */ package com.google.api.gax.rpc; +import java.util.logging.Level; +import java.util.logging.Logger; import org.jspecify.annotations.NullMarked; import org.jspecify.annotations.Nullable; @@ -42,6 +44,7 @@ */ @NullMarked public abstract class BidiStreamingCallable { + private static final Logger LOG = Logger.getLogger(BidiStreamingCallable.class.getName()); protected BidiStreamingCallable() {} @@ -241,11 +244,47 @@ public BidiStreamingCallable withDefaultCallContext( return new BidiStreamingCallable() { @Override public ClientStream internalCall( - ResponseObserver responseObserver, + final ResponseObserver responseObserver, ClientStreamReadyObserver onReady, ApiCallContext thisCallContext) { - return BidiStreamingCallable.this.internalCall( - responseObserver, onReady, defaultCallContext.merge(thisCallContext)); + final ApiCallContext mergedContext = defaultCallContext.merge(thisCallContext); + ResponseObserver refreshingObserver = + new ResponseObserver() { + @Override + public void onStart(StreamController controller) { + responseObserver.onStart(controller); + } + + @Override + public void onResponse(ResponseT response) { + responseObserver.onResponse(response); + } + + @Override + public void onError(Throwable t) { + if (t instanceof UnauthenticatedException) { + TransportChannel transportChannel = mergedContext.getTransportChannel(); + if (transportChannel != null && transportChannel.shouldRefresh()) { + try { + transportChannel.refresh(); + } catch (Exception e) { + LOG.log( + Level.WARNING, + "Failed to refresh transport channel after authentication error", + e); + } + } + } + responseObserver.onError(t); + } + + @Override + public void onComplete() { + responseObserver.onComplete(); + } + }; + + return BidiStreamingCallable.this.internalCall(refreshingObserver, onReady, mergedContext); } }; } diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ClientStreamingCallable.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ClientStreamingCallable.java index 854ce3fd5de8..5c87e4c558b7 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ClientStreamingCallable.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ClientStreamingCallable.java @@ -29,6 +29,8 @@ */ package com.google.api.gax.rpc; +import java.util.logging.Level; +import java.util.logging.Logger; import org.jspecify.annotations.NullMarked; import org.jspecify.annotations.Nullable; @@ -42,6 +44,7 @@ */ @NullMarked public abstract class ClientStreamingCallable { + private static final Logger LOG = Logger.getLogger(ClientStreamingCallable.class.getName()); protected ClientStreamingCallable() {} @@ -77,9 +80,40 @@ public ClientStreamingCallable withDefaultCallContext( return new ClientStreamingCallable() { @Override public ApiStreamObserver clientStreamingCall( - ApiStreamObserver responseObserver, ApiCallContext thisCallContext) { - return ClientStreamingCallable.this.clientStreamingCall( - responseObserver, defaultCallContext.merge(thisCallContext)); + final ApiStreamObserver responseObserver, ApiCallContext thisCallContext) { + final ApiCallContext mergedContext = defaultCallContext.merge(thisCallContext); + ApiStreamObserver refreshingObserver = + new ApiStreamObserver() { + @Override + public void onNext(ResponseT response) { + responseObserver.onNext(response); + } + + @Override + public void onError(Throwable t) { + if (t instanceof UnauthenticatedException) { + TransportChannel transportChannel = mergedContext.getTransportChannel(); + if (transportChannel != null && transportChannel.shouldRefresh()) { + try { + transportChannel.refresh(); + } catch (Exception e) { + LOG.log( + Level.WARNING, + "Failed to refresh transport channel after authentication error", + e); + } + } + } + responseObserver.onError(t); + } + + @Override + public void onCompleted() { + responseObserver.onCompleted(); + } + }; + + return ClientStreamingCallable.this.clientStreamingCall(refreshingObserver, mergedContext); } }; } diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ServerStreamingAttemptCallable.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ServerStreamingAttemptCallable.java index af6e8014e38a..bf1fdef29c05 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ServerStreamingAttemptCallable.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ServerStreamingAttemptCallable.java @@ -37,6 +37,8 @@ import com.google.errorprone.annotations.concurrent.GuardedBy; import java.util.concurrent.Callable; import java.util.concurrent.CancellationException; +import java.util.logging.Level; +import java.util.logging.Logger; import org.jspecify.annotations.NullMarked; /** @@ -96,6 +98,8 @@ */ @NullMarked final class ServerStreamingAttemptCallable implements Callable { + private static final Logger LOG = + Logger.getLogger(ServerStreamingAttemptCallable.class.getName()); private final Object lock = new Object(); private final ServerStreamingCallable innerCallable; @@ -221,6 +225,10 @@ public Void call() { .getTracer() .attemptStarted(request, outerRetryingFuture.getAttemptSettings().getOverallAttemptCount()); + final ApiCallContext finalContext = attemptContext; + TransportChannel channelForAttempt = finalContext.getTransportChannel(); + final long attemptGeneration = + channelForAttempt != null ? channelForAttempt.getGeneration() : 0; innerCallable.call( request, new StateCheckingResponseObserver() { @@ -236,6 +244,40 @@ public void onResponseImpl(ResponseT response) { @Override public void onErrorImpl(Throwable t) { + Throwable cause = t; + if (cause instanceof ServerStreamingAttemptException) { + cause = cause.getCause(); + } + if (cause instanceof UnauthenticatedException) { + UnauthenticatedException unauthenticatedException = (UnauthenticatedException) cause; + TransportChannel transportChannel = finalContext.getTransportChannel(); + if (transportChannel != null) { + if (transportChannel.shouldRefresh()) { + try { + transportChannel.refresh(); + } catch (Exception e) { + LOG.log( + Level.WARNING, + "Failed to refresh transport channel after authentication error", + e); + } + } + boolean shouldRetry = transportChannel.getGeneration() > attemptGeneration; + if (shouldRetry) { + UnauthenticatedException newEx = + new UnauthenticatedException( + unauthenticatedException.getMessage(), + unauthenticatedException.getCause(), + unauthenticatedException.getStatusCode(), + true, + unauthenticatedException.getErrorDetails()); + for (Throwable suppressed : unauthenticatedException.getSuppressed()) { + newEx.addSuppressed(suppressed); + } + t = newEx; + } + } + } onAttemptError(t); } diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/TransportChannel.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/TransportChannel.java index 1866092e28f9..cad80d4f87b3 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/TransportChannel.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/TransportChannel.java @@ -31,10 +31,8 @@ import com.google.api.core.InternalExtensionOnly; import com.google.api.gax.core.BackgroundResource; -import org.jspecify.annotations.NullMarked; /** Class whose instances can issue RPCs on a particular transport. */ -@NullMarked @InternalExtensionOnly public interface TransportChannel extends BackgroundResource { @@ -49,4 +47,28 @@ public interface TransportChannel extends BackgroundResource { * Returns an empty {@link ApiCallContext} that is compatible with this {@code TransportChannel}. */ ApiCallContext getEmptyCallContext(); + + /** + * Refreshes or recreates the underlying network connections of this transport channel. + * + *

By default, this is a no-op for transports that do not require stateful connection lifecycle + * management. + */ + default void refresh() {} + + /** + * Returns true if a certificate rotation has been detected on disk and the transport channel + * should be refreshed, or false otherwise. + */ + default boolean shouldRefresh() { + return false; + } + + /** + * Returns a monotonic generation counter tracking the number of successful refreshes or channel + * rotations performed by this transport channel. + */ + default long getGeneration() { + return 0; + } } diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/mtls/CertificateBasedAccess.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/mtls/CertificateBasedAccess.java index 99f2cca9b53d..d77378ad9612 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/mtls/CertificateBasedAccess.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/mtls/CertificateBasedAccess.java @@ -32,21 +32,20 @@ import com.google.api.core.InternalApi; import com.google.api.gax.rpc.internal.EnvironmentProvider; -import org.jspecify.annotations.NullMarked; +import com.google.auth.mtls.MtlsUtils; +import com.google.auth.oauth2.PropertyProvider; /** * Utility class for handling certificate-based access configurations. * - *

This class handles the processing of GOOGLE_API_USE_CLIENT_CERTIFICATE and - * GOOGLE_API_USE_MTLS_ENDPOINT environment variables according to https://google.aip.dev/auth/4114 + *

This class handles the processing of GOOGLE_API_USE_CLIENT_CERTIFICATE, + * GOOGLE_API_CERTIFICATE_CONFIG, and GOOGLE_API_USE_MTLS_ENDPOINT configurations. */ -@NullMarked @InternalApi public class CertificateBasedAccess { private final EnvironmentProvider envProvider; - /** The EnvironmentProvider mechanism supports env var injection for unit tests. */ public CertificateBasedAccess(EnvironmentProvider envProvider) { this.envProvider = envProvider; } @@ -66,20 +65,35 @@ public enum MtlsEndpointUsagePolicy { ALWAYS; } + private com.google.auth.oauth2.EnvironmentProvider getAuthEnvProvider() { + return name -> envProvider.getenv(name); + } + + private PropertyProvider getAuthPropertyProvider() { + return System::getProperty; + } + /** Returns if mutual TLS client certificate should be used. */ public boolean useMtlsClientCertificate() { - String useClientCertificate = envProvider.getenv("GOOGLE_API_USE_CLIENT_CERTIFICATE"); - return "true".equals(useClientCertificate); + return MtlsUtils.useMtlsClientCertificate(getAuthEnvProvider(), getAuthPropertyProvider()); } /** Returns the current mutual TLS endpoint usage policy. */ public MtlsEndpointUsagePolicy getMtlsEndpointUsagePolicy() { String mtlsEndpointUsagePolicy = envProvider.getenv("GOOGLE_API_USE_MTLS_ENDPOINT"); - if ("never".equals(mtlsEndpointUsagePolicy)) { + if ("never".equalsIgnoreCase(mtlsEndpointUsagePolicy)) { return MtlsEndpointUsagePolicy.NEVER; - } else if ("always".equals(mtlsEndpointUsagePolicy)) { + } else if ("always".equalsIgnoreCase(mtlsEndpointUsagePolicy)) { return MtlsEndpointUsagePolicy.ALWAYS; } return MtlsEndpointUsagePolicy.AUTO; } + + /** + * Resolves and returns the path to the mutual TLS client certificate, or null if none should be + * used. + */ + public String getWorkloadCertPath() { + return MtlsUtils.getWorkloadCertPath(getAuthEnvProvider(), getAuthPropertyProvider()); + } } diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/mtls/CertificateRotationTracker.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/mtls/CertificateRotationTracker.java new file mode 100644 index 000000000000..fc58d67ab1d6 --- /dev/null +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/mtls/CertificateRotationTracker.java @@ -0,0 +1,197 @@ +/* + * Copyright 2026 Google LLC + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are + * met: + * + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above + * copyright notice, this list of conditions and the following disclaimer + * in the documentation and/or other materials provided with the + * distribution. + * * Neither the name of Google LLC nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS + * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT + * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR + * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT + * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, + * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT + * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, + * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ +package com.google.api.gax.rpc.mtls; + +import com.google.api.core.InternalApi; +import com.google.common.base.Strings; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Function; +import java.util.function.Supplier; +import org.jspecify.annotations.Nullable; + +/** + * Thread-safe helper that tracks workload certificate fingerprints on disk to detect mTLS + * certificate rotations for transport channels ({@code ChannelPool} and {@code + * RefreshingHttpJsonChannel}). + * + *

For internal use only. + */ +@InternalApi +public class CertificateRotationTracker { + + /** + * Duration (1 second) for which a detected certificate rotation (where the disk fingerprint + * differs from {@code activeCertFingerprint}) is cached in memory. + * + *

When a certificate rotates on disk, many concurrent in-flight RPCs may fail with {@code + * UNAUTHENTICATED} simultaneously and call {@link #shouldRefresh()}. Caching positive rotation + * detections for 1 second coalesces disk reads ({@code Files.readAllBytes}) and SHA-256 hashing + * across concurrent threads to prevent a thundering herd of file I/O, while bounding maximum + * staleness to 1 second. + */ + private static final long POSITIVE_ROTATION_CACHE_TTL_NANOS = TimeUnit.SECONDS.toNanos(1); + + private static final class DiskCheckResult { + final String fingerprint; + final long sequence; + final long timestampNanos; + + DiskCheckResult(String fingerprint, long sequence, long timestampNanos) { + this.fingerprint = fingerprint; + this.sequence = sequence; + this.timestampNanos = timestampNanos; + } + } + + private final Supplier certPathSupplier; + private final Function fingerprintReader; + private volatile String activeCertFingerprint; + private volatile DiskCheckResult lastDiskCheck = null; + private final ReentrantLock diskCheckLock = new ReentrantLock(); + private final AtomicLong diskCheckSequence = new AtomicLong(0); + + /** + * Creates a tracker for a fixed workload certificate path using {@link + * WorkloadCertificateUtils#getCertificateFingerprint(String)}. + */ + public CertificateRotationTracker(@Nullable String workloadCertPath) { + this(() -> workloadCertPath, WorkloadCertificateUtils::getCertificateFingerprint); + } + + /** + * Creates a tracker with custom certificate path and fingerprint suppliers (used by transports + * that expose package-private overrides for testing). + */ + public CertificateRotationTracker( + Supplier certPathSupplier, Function fingerprintReader) { + this.certPathSupplier = certPathSupplier; + this.fingerprintReader = fingerprintReader; + String initialCertPath = certPathSupplier.get(); + this.activeCertFingerprint = + initialCertPath != null + ? Strings.nullToEmpty(fingerprintReader.apply(initialCertPath)) + : ""; + } + + /** + * Returns {@code true} if a workload certificate path is configured, readable, and its current + * SHA-256 fingerprint on disk differs from the active fingerprint. + */ + public boolean shouldRefresh() { + String certPath = certPathSupplier.get(); + if (certPath == null) { + return false; + } + String currentDiskFingerprint = getOrUpdateDiskFingerprint(certPath); + if (currentDiskFingerprint.isEmpty()) { + return false; + } + return !currentDiskFingerprint.equalsIgnoreCase(activeCertFingerprint); + } + + private String getOrUpdateDiskFingerprint(String certPath) { + long seqBeforeLock = diskCheckSequence.get(); + long now = System.nanoTime(); + DiskCheckResult cached = lastDiskCheck; + if (cached != null + && !cached.fingerprint.isEmpty() + && !cached.fingerprint.equalsIgnoreCase(this.activeCertFingerprint) + && (now - cached.timestampNanos < POSITIVE_ROTATION_CACHE_TTL_NANOS)) { + return cached.fingerprint; + } + + diskCheckLock.lock(); + try { + now = System.nanoTime(); + cached = lastDiskCheck; + if (cached != null + && !cached.fingerprint.isEmpty() + && (cached.sequence > seqBeforeLock + || (!cached.fingerprint.equalsIgnoreCase(this.activeCertFingerprint) + && (now - cached.timestampNanos < POSITIVE_ROTATION_CACHE_TTL_NANOS)))) { + return cached.fingerprint; + } + long newSeq = diskCheckSequence.incrementAndGet(); + String fingerprint = Strings.nullToEmpty(fingerprintReader.apply(certPath)); + if (!fingerprint.isEmpty()) { + lastDiskCheck = new DiskCheckResult(fingerprint, newSeq, System.nanoTime()); + } else { + lastDiskCheck = null; + } + return fingerprint; + } finally { + diskCheckLock.unlock(); + } + } + + /** + * Reads the current certificate fingerprint directly from disk (bypassing the 1-second cache), + * returning {@code ""} if no workload certificate path is configured or if the file is currently + * unreadable/empty. + */ + public String readDiskFingerprint() { + String certPath = certPathSupplier.get(); + if (certPath == null) { + return ""; + } + return Strings.nullToEmpty(fingerprintReader.apply(certPath)); + } + + /** + * Returns {@code true} if {@code diskFingerprint} matches the currently active certificate + * fingerprint. + */ + public boolean isAlreadyActive(String diskFingerprint) { + return diskFingerprint != null && diskFingerprint.equalsIgnoreCase(this.activeCertFingerprint); + } + + /** + * Updates the active certificate fingerprint after a successful channel refresh and clears any + * cached disk check result. + */ + public void markRefreshed(String newFingerprint) { + if (newFingerprint != null && !newFingerprint.isEmpty()) { + this.activeCertFingerprint = newFingerprint; + this.lastDiskCheck = null; + } + } + + /** Returns the currently active certificate fingerprint (or {@code ""} if none). */ + public String getActiveCertFingerprint() { + return activeCertFingerprint; + } + + /** Invalidates the cached disk check result. Visible for testing. */ + public void invalidateCache() { + this.lastDiskCheck = null; + } +} diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/mtls/WorkloadCertificateUtils.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/mtls/WorkloadCertificateUtils.java new file mode 100644 index 000000000000..94fb36db1ac3 --- /dev/null +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/mtls/WorkloadCertificateUtils.java @@ -0,0 +1,64 @@ +/* + * Copyright 2026 Google LLC + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are + * met: + * + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above + * copyright notice, this list of conditions and the following disclaimer + * in the documentation and/or other materials provided with the + * distribution. + * * Neither the name of Google LLC nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS + * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT + * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR + * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT + * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, + * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT + * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, + * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ +package com.google.api.gax.rpc.mtls; + +import com.google.api.core.InternalApi; +import com.google.auth.mtls.MtlsUtils; +import java.io.File; + +/** Internal utility class for managing dynamic workload certificates. */ +@InternalApi +public class WorkloadCertificateUtils { + + private static final String EMPTY_FILE_SHA256 = + "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"; + + private WorkloadCertificateUtils() {} + + /** + * Computes the SHA-256 fingerprint of the certificate file at {@code certPath}, returning {@code + * ""} if the path is {@code null}, unreadable, empty (e.g., temporarily truncated to 0 bytes + * mid-write by an external certificate rotator), or hashes to the empty-byte digest. + * + *

Returning {@code ""} on unreadable or empty files ensures callers ({@code shouldRefresh()} + * and {@code refresh()}) safely skip refreshing during transient mid-write states rather than + * treating an empty digest as a certificate rotation mismatch. + */ + public static String getCertificateFingerprint(String certPath) { + if (certPath == null || new File(certPath).length() == 0) { + return ""; + } + String fingerprint = MtlsUtils.getCertificateFingerprint(certPath); + if (fingerprint == null || EMPTY_FILE_SHA256.equalsIgnoreCase(fingerprint)) { + return ""; + } + return fingerprint; + } +} diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ApiResultRetryAlgorithmTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ApiResultRetryAlgorithmTest.java index 300a0ad30130..41daa50745ea 100644 --- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ApiResultRetryAlgorithmTest.java +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ApiResultRetryAlgorithmTest.java @@ -29,14 +29,22 @@ */ package com.google.api.gax.rpc; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import com.google.api.core.NanoClock; +import com.google.api.gax.retrying.ExponentialRetryAlgorithm; +import com.google.api.gax.retrying.RetryAlgorithm; +import com.google.api.gax.retrying.RetrySettings; +import com.google.api.gax.retrying.TimedAttemptSettings; import com.google.api.gax.rpc.StatusCode.Code; import com.google.api.gax.rpc.testing.FakeStatusCode; import com.google.common.collect.Sets; +import java.time.Duration; import java.util.Collections; import org.junit.jupiter.api.Test; import org.mockito.Mockito; @@ -111,4 +119,133 @@ void testShouldRetryWithContextWithEmptyRetryableCodes() { ApiResultRetryAlgorithm algorithm = new ApiResultRetryAlgorithm<>(); assertFalse(algorithm.shouldRetry(context, unavailableException, null)); } + + @Test + void testRotationRetryWithNonRetryableSettings_maxAttemptsOne() { + ApiCallContext context = + mock(ApiCallContext.class, Mockito.withSettings().withoutAnnotations()); + when(context.getRetryableCodes()).thenReturn(Collections.emptySet()); + + RetrySettings settings = + RetrySettings.newBuilder() + .setMaxAttempts(1) + .setInitialRpcTimeoutDuration(Duration.ofSeconds(10)) + .setMaxRpcTimeoutDuration(Duration.ofSeconds(10)) + .setTotalTimeoutDuration(Duration.ofSeconds(10)) + .build(); + + ApiResultRetryAlgorithm resultAlgorithm = new ApiResultRetryAlgorithm<>(); + ExponentialRetryAlgorithm timedAlgorithm = + new ExponentialRetryAlgorithm(settings, NanoClock.getDefaultClock()); + RetryAlgorithm retryAlgorithm = new RetryAlgorithm<>(resultAlgorithm, timedAlgorithm); + + TimedAttemptSettings firstAttempt = retryAlgorithm.createFirstAttempt(context); + UnauthenticatedException rotationEx = + new UnauthenticatedException( + "Expired cert", null, new FakeStatusCode(Code.UNAUTHENTICATED), /* retryable= */ true); + + // First rotation failure: grants immediate free retry without incrementing attemptCount + TimedAttemptSettings nextAttempt = + retryAlgorithm.createNextAttempt(context, rotationEx, null, firstAttempt); + assertNotNull(nextAttempt); + assertEquals(Duration.ZERO, nextAttempt.getRetryDelayDuration()); + assertEquals(Duration.ZERO, nextAttempt.getRandomizedRetryDelayDuration()); + assertEquals(0, nextAttempt.getAttemptCount()); + assertEquals(1, nextAttempt.getOverallAttemptCount()); + assertTrue(retryAlgorithm.shouldRetry(context, rotationEx, null, nextAttempt)); + + // Second consecutive failure: overallAttemptCount (1) != attemptCount (0), so no free retry + TimedAttemptSettings thirdAttempt = + retryAlgorithm.createNextAttempt(context, rotationEx, null, nextAttempt); + assertNotNull(thirdAttempt); + assertEquals(1, thirdAttempt.getAttemptCount()); + assertEquals(2, thirdAttempt.getOverallAttemptCount()); + assertFalse(retryAlgorithm.shouldRetry(context, rotationEx, null, thirdAttempt)); + } + + @Test + void testRotationRetryWithNonRetryableSettings_zeroMaxAttemptsZeroTotalTimeout() { + ApiCallContext context = + mock(ApiCallContext.class, Mockito.withSettings().withoutAnnotations()); + when(context.getRetryableCodes()).thenReturn(Collections.emptySet()); + + RetrySettings settings = + RetrySettings.newBuilder().setMaxAttempts(0).setTotalTimeoutDuration(Duration.ZERO).build(); + + ApiResultRetryAlgorithm resultAlgorithm = new ApiResultRetryAlgorithm<>(); + ExponentialRetryAlgorithm timedAlgorithm = + new ExponentialRetryAlgorithm(settings, NanoClock.getDefaultClock()); + RetryAlgorithm retryAlgorithm = new RetryAlgorithm<>(resultAlgorithm, timedAlgorithm); + + TimedAttemptSettings firstAttempt = retryAlgorithm.createFirstAttempt(context); + UnauthenticatedException rotationEx = + new UnauthenticatedException( + "Expired cert", null, new FakeStatusCode(Code.UNAUTHENTICATED), /* retryable= */ true); + + TimedAttemptSettings nextAttempt = + retryAlgorithm.createNextAttempt(context, rotationEx, null, firstAttempt); + assertNotNull(nextAttempt); + assertEquals(1, nextAttempt.getGlobalSettings().getMaxAttempts()); + assertEquals(0, nextAttempt.getAttemptCount()); + assertEquals(1, nextAttempt.getOverallAttemptCount()); + assertTrue(retryAlgorithm.shouldRetry(context, rotationEx, null, nextAttempt)); + + // Subsequent failure is rejected + TimedAttemptSettings thirdAttempt = + retryAlgorithm.createNextAttempt(context, rotationEx, null, nextAttempt); + assertNotNull(thirdAttempt); + assertFalse(retryAlgorithm.shouldRetry(context, rotationEx, null, thirdAttempt)); + } + + @Test + void testRotationRetryAfterTransientErrorPreservesRemainingBudget() { + ApiCallContext context = + mock(ApiCallContext.class, Mockito.withSettings().withoutAnnotations()); + when(context.getRetryableCodes()).thenReturn(Sets.newHashSet(Code.UNAVAILABLE)); + + RetrySettings settings = + RetrySettings.newBuilder() + .setMaxAttempts(3) + .setInitialRetryDelayDuration(Duration.ofMillis(100)) + .setRetryDelayMultiplier(2.0) + .setMaxRetryDelayDuration(Duration.ofSeconds(1)) + .setTotalTimeoutDuration(Duration.ofSeconds(30)) + .build(); + + ApiResultRetryAlgorithm resultAlgorithm = new ApiResultRetryAlgorithm<>(); + ExponentialRetryAlgorithm timedAlgorithm = + new ExponentialRetryAlgorithm(settings, NanoClock.getDefaultClock()); + RetryAlgorithm retryAlgorithm = new RetryAlgorithm<>(resultAlgorithm, timedAlgorithm); + + TimedAttemptSettings attempt0 = retryAlgorithm.createFirstAttempt(context); + ApiException unavailableEx = + new ApiException(null, new FakeStatusCode(Code.UNAVAILABLE), /* retryable= */ true); + UnauthenticatedException rotationEx = + new UnauthenticatedException( + "Expired cert", null, new FakeStatusCode(Code.UNAUTHENTICATED), /* retryable= */ true); + + // Attempt 0 fails with UNAVAILABLE -> normal retry (attemptCount = 1, overallAttemptCount = 1) + TimedAttemptSettings attempt1 = + retryAlgorithm.createNextAttempt(context, unavailableEx, null, attempt0); + assertEquals(1, attempt1.getAttemptCount()); + assertEquals(1, attempt1.getOverallAttemptCount()); + assertTrue(retryAlgorithm.shouldRetry(context, unavailableEx, null, attempt1)); + + // Attempt 1 fails with rotation 401 -> free zero-delay retry (attemptCount = 1, + // overallAttemptCount = 2) + TimedAttemptSettings attempt2 = + retryAlgorithm.createNextAttempt(context, rotationEx, null, attempt1); + assertEquals(Duration.ZERO, attempt2.getRetryDelayDuration()); + assertEquals(1, attempt2.getAttemptCount()); + assertEquals(2, attempt2.getOverallAttemptCount()); + assertTrue(retryAlgorithm.shouldRetry(context, rotationEx, null, attempt2)); + + // Attempt 2 fails with UNAVAILABLE -> normal retry still allowed (attemptCount = 2 < + // maxAttempts = 3) + TimedAttemptSettings attempt3 = + retryAlgorithm.createNextAttempt(context, unavailableEx, null, attempt2); + assertEquals(2, attempt3.getAttemptCount()); + assertEquals(3, attempt3.getOverallAttemptCount()); + assertTrue(retryAlgorithm.shouldRetry(context, unavailableEx, null, attempt3)); + } } diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/AttemptCallableTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/AttemptCallableTest.java index 4b5da578862c..84c51da90707 100644 --- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/AttemptCallableTest.java +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/AttemptCallableTest.java @@ -34,11 +34,15 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import com.google.api.core.ApiFuture; import com.google.api.core.SettableApiFuture; import com.google.api.gax.retrying.RetrySettings; import com.google.api.gax.retrying.RetryingFuture; import com.google.api.gax.retrying.TimedAttemptSettings; import com.google.api.gax.rpc.testing.FakeCallContext; +import com.google.api.gax.rpc.testing.FakeChannel; +import com.google.api.gax.rpc.testing.FakeStatusCode; +import com.google.api.gax.rpc.testing.FakeTransportChannel; import com.google.api.gax.tracing.ApiTracer; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -133,4 +137,287 @@ void testRpcTimeoutIsNotErased() { assertThat(capturedCallContext.getValue().getTimeoutDuration()).isEqualTo(callerTimeout); } + + @Test + void testUnauthenticatedExceptionReThrowPreservesContext() { + FakeTransportChannel transportChannel = + FakeTransportChannel.create(new FakeChannel()).setShouldRefresh(true); + ApiCallContext callContext = + FakeCallContext.createDefault().withTransportChannel(transportChannel); + + UnauthenticatedException originalEx = + new UnauthenticatedException( + "Expired cert", + new IllegalStateException("Root cause"), + FakeStatusCode.of(StatusCode.Code.UNAUTHENTICATED), + false); + originalEx.setStackTrace( + new StackTraceElement[] {new StackTraceElement("foo", "bar", "Baz.java", 123)}); + originalEx.addSuppressed(new RuntimeException("Suppressed cause")); + + SettableApiFuture failedFuture = SettableApiFuture.create(); + failedFuture.setException(originalEx); + when(mockInnerCallable.futureCall(Mockito.anyString(), Mockito.any())).thenReturn(failedFuture); + + AttemptCallable callable = + new AttemptCallable<>(mockInnerCallable, "fake-request", callContext); + callable.setExternalFuture(mockExternalFuture); + + callable.call(); + + ArgumentCaptor futureCaptor = ArgumentCaptor.forClass(ApiFuture.class); + Mockito.verify(mockExternalFuture, Mockito.times(2)).setAttemptFuture(futureCaptor.capture()); + + Throwable thrown = null; + try { + futureCaptor.getValue().get(); + } catch (Exception e) { + thrown = e.getCause(); + } + + assertThat(thrown).isInstanceOf(UnauthenticatedException.class); + UnauthenticatedException rethrown = (UnauthenticatedException) thrown; + assertThat(rethrown.isRetryable()).isTrue(); + assertThat(rethrown.getCause()).isEqualTo(originalEx.getCause()); + assertThat(rethrown.getStackTrace()).isEqualTo(originalEx.getStackTrace()); + assertThat(rethrown.getSuppressed().length).isEqualTo(1); + assertThat(rethrown.getSuppressed()[0]).isInstanceOf(RuntimeException.class); + } + + @Test + void testSiblingInFlightRequest_channelRotatedInFlight_markedRetryableWithoutDuplicateRefresh() { + FakeChannel innerChannel = new FakeChannel(); + innerChannel.setGeneration(1); + // Initially shouldRefresh is false because sibling request already completed the refresh + innerChannel.setShouldRefresh(false); + FakeTransportChannel transportChannel = FakeTransportChannel.create(innerChannel); + + ApiCallContext callContext = + FakeCallContext.createDefault().withTransportChannel(transportChannel); + + UnauthenticatedException originalEx = + new UnauthenticatedException( + "Expired cert", null, FakeStatusCode.of(StatusCode.Code.UNAUTHENTICATED), false); + + SettableApiFuture failedFuture = SettableApiFuture.create(); + failedFuture.setException(originalEx); + when(mockInnerCallable.futureCall(Mockito.anyString(), Mockito.any())) + .thenAnswer( + invocation -> { + // While request was in flight, sibling finished rotation and bumped generation to 2 + innerChannel.setGeneration(2); + return failedFuture; + }); + + AttemptCallable callable = + new AttemptCallable<>(mockInnerCallable, "fake-request", callContext); + callable.setExternalFuture(mockExternalFuture); + + callable.call(); + + ArgumentCaptor futureCaptor = ArgumentCaptor.forClass(ApiFuture.class); + Mockito.verify(mockExternalFuture, Mockito.times(2)).setAttemptFuture(futureCaptor.capture()); + + Throwable thrown = null; + try { + futureCaptor.getValue().get(); + } catch (Exception e) { + thrown = e.getCause(); + } + + assertThat(thrown).isInstanceOf(UnauthenticatedException.class); + UnauthenticatedException rethrown = (UnauthenticatedException) thrown; + // Sibling request should be marked retryable to run on the new channel + assertThat(rethrown.isRetryable()).isTrue(); + // But should NOT have triggered a second refresh call + assertThat(innerChannel.getRefreshCount()).isEqualTo(0); + } + + @Test + void testPermanentUnauthenticatedFailure_sameGeneration_notMarkedRetryable() { + FakeChannel innerChannel = new FakeChannel(); + innerChannel.setGeneration(1); + innerChannel.setShouldRefresh(false); + FakeTransportChannel transportChannel = FakeTransportChannel.create(innerChannel); + + ApiCallContext callContext = + FakeCallContext.createDefault().withTransportChannel(transportChannel); + + UnauthenticatedException originalEx = + new UnauthenticatedException( + "Invalid credentials", null, FakeStatusCode.of(StatusCode.Code.UNAUTHENTICATED), false); + + SettableApiFuture failedFuture = SettableApiFuture.create(); + failedFuture.setException(originalEx); + when(mockInnerCallable.futureCall(Mockito.anyString(), Mockito.any())).thenReturn(failedFuture); + + AttemptCallable callable = + new AttemptCallable<>(mockInnerCallable, "fake-request", callContext); + callable.setExternalFuture(mockExternalFuture); + + callable.call(); + + ArgumentCaptor futureCaptor = ArgumentCaptor.forClass(ApiFuture.class); + Mockito.verify(mockExternalFuture, Mockito.times(2)).setAttemptFuture(futureCaptor.capture()); + + Throwable thrown = null; + try { + futureCaptor.getValue().get(); + } catch (Exception e) { + thrown = e.getCause(); + } + + assertThat(thrown).isInstanceOf(UnauthenticatedException.class); + UnauthenticatedException rethrown = (UnauthenticatedException) thrown; + // Genuine permanent error on same generation is NOT retryable + assertThat(rethrown.isRetryable()).isFalse(); + assertThat(innerChannel.getRefreshCount()).isEqualTo(0); + } + + @Test + void testRefreshThrowsException_notMarkedRetryableWhenGenerationUnchanged() { + FakeChannel fakeChannel = + new FakeChannel() { + @Override + public boolean shouldRefresh() { + return true; + } + + @Override + public void refresh() { + throw new RuntimeException("Refresh error"); + } + }; + FakeTransportChannel transportChannel = FakeTransportChannel.create(fakeChannel); + + ApiCallContext callContext = + FakeCallContext.createDefault().withTransportChannel(transportChannel); + + UnauthenticatedException originalEx = + new UnauthenticatedException( + "Expired cert", null, FakeStatusCode.of(StatusCode.Code.UNAUTHENTICATED), false); + + SettableApiFuture failedFuture = SettableApiFuture.create(); + failedFuture.setException(originalEx); + when(mockInnerCallable.futureCall(Mockito.anyString(), Mockito.any())).thenReturn(failedFuture); + + AttemptCallable callable = + new AttemptCallable<>(mockInnerCallable, "fake-request", callContext); + callable.setExternalFuture(mockExternalFuture); + + callable.call(); + + ArgumentCaptor futureCaptor = ArgumentCaptor.forClass(ApiFuture.class); + Mockito.verify(mockExternalFuture, Mockito.times(2)).setAttemptFuture(futureCaptor.capture()); + + Throwable thrown = null; + try { + futureCaptor.getValue().get(); + } catch (Exception e) { + thrown = e.getCause(); + } + + assertThat(thrown).isInstanceOf(UnauthenticatedException.class); + UnauthenticatedException rethrown = (UnauthenticatedException) thrown; + assertThat(rethrown.isRetryable()).isFalse(); + } + + @Test + void testRefreshReturnsWithoutAdvancingGeneration_notMarkedRetryable() { + FakeChannel fakeChannel = + new FakeChannel() { + @Override + public boolean shouldRefresh() { + return true; + } + + @Override + public void refresh() { + // Simulates refresh() returning early without rotating any channel (e.g. unreadable + // cert) + } + }; + FakeTransportChannel transportChannel = FakeTransportChannel.create(fakeChannel); + + ApiCallContext callContext = + FakeCallContext.createDefault().withTransportChannel(transportChannel); + + UnauthenticatedException originalEx = + new UnauthenticatedException( + "Expired cert", null, FakeStatusCode.of(StatusCode.Code.UNAUTHENTICATED), false); + + SettableApiFuture failedFuture = SettableApiFuture.create(); + failedFuture.setException(originalEx); + when(mockInnerCallable.futureCall(Mockito.anyString(), Mockito.any())).thenReturn(failedFuture); + + AttemptCallable callable = + new AttemptCallable<>(mockInnerCallable, "fake-request", callContext); + callable.setExternalFuture(mockExternalFuture); + + callable.call(); + + ArgumentCaptor futureCaptor = ArgumentCaptor.forClass(ApiFuture.class); + Mockito.verify(mockExternalFuture, Mockito.times(2)).setAttemptFuture(futureCaptor.capture()); + + Throwable thrown = null; + try { + futureCaptor.getValue().get(); + } catch (Exception e) { + thrown = e.getCause(); + } + + assertThat(thrown).isInstanceOf(UnauthenticatedException.class); + UnauthenticatedException rethrown = (UnauthenticatedException) thrown; + assertThat(rethrown.isRetryable()).isFalse(); + } + + @Test + void testRefreshThrowsException_markedRetryableIfConcurrentThreadAdvancedGeneration() { + FakeChannel fakeChannel = + new FakeChannel() { + @Override + public boolean shouldRefresh() { + return true; + } + + @Override + public void refresh() { + // Concurrent thread advanced generation before/during refresh failure + setGeneration(getGeneration() + 1); + throw new RuntimeException("Refresh error on this thread"); + } + }; + FakeTransportChannel transportChannel = FakeTransportChannel.create(fakeChannel); + + ApiCallContext callContext = + FakeCallContext.createDefault().withTransportChannel(transportChannel); + + UnauthenticatedException originalEx = + new UnauthenticatedException( + "Expired cert", null, FakeStatusCode.of(StatusCode.Code.UNAUTHENTICATED), false); + + SettableApiFuture failedFuture = SettableApiFuture.create(); + failedFuture.setException(originalEx); + when(mockInnerCallable.futureCall(Mockito.anyString(), Mockito.any())).thenReturn(failedFuture); + + AttemptCallable callable = + new AttemptCallable<>(mockInnerCallable, "fake-request", callContext); + callable.setExternalFuture(mockExternalFuture); + + callable.call(); + + ArgumentCaptor futureCaptor = ArgumentCaptor.forClass(ApiFuture.class); + Mockito.verify(mockExternalFuture, Mockito.times(2)).setAttemptFuture(futureCaptor.capture()); + + Throwable thrown = null; + try { + futureCaptor.getValue().get(); + } catch (Exception e) { + thrown = e.getCause(); + } + + assertThat(thrown).isInstanceOf(UnauthenticatedException.class); + UnauthenticatedException rethrown = (UnauthenticatedException) thrown; + assertThat(rethrown.isRetryable()).isTrue(); + } } diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/EndpointContextTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/EndpointContextTest.java index 4c1b5320de4c..c53ca1509430 100644 --- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/EndpointContextTest.java +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/EndpointContextTest.java @@ -77,8 +77,10 @@ void mtlsEndpointResolver_switchToMtlsAllowedIsFalse() throws IOException { FakeMtlsProvider.createTestMtlsKeyStore(), "", throwExceptionForGetKeyStore); boolean switchToMtlsEndpointAllowed = false; CertificateBasedAccess certificateBasedAccess = - new CertificateBasedAccess( - name -> name.equals("GOOGLE_API_USE_MTLS_ENDPOINT") ? "auto" : "true"); + org.mockito.Mockito.mock(CertificateBasedAccess.class); + org.mockito.Mockito.when(certificateBasedAccess.getMtlsEndpointUsagePolicy()) + .thenReturn(CertificateBasedAccess.MtlsEndpointUsagePolicy.AUTO); + org.mockito.Mockito.when(certificateBasedAccess.useMtlsClientCertificate()).thenReturn(true); Truth.assertThat( defaultEndpointContextBuilder.mtlsEndpointResolver( DEFAULT_ENDPOINT, @@ -97,8 +99,10 @@ void mtlsEndpointResolver_switchToMtlsAllowedIsTrue_mtlsUsageAuto() throws IOExc FakeMtlsProvider.createTestMtlsKeyStore(), "", throwExceptionForGetKeyStore); boolean switchToMtlsEndpointAllowed = true; CertificateBasedAccess certificateBasedAccess = - new CertificateBasedAccess( - name -> name.equals("GOOGLE_API_USE_MTLS_ENDPOINT") ? "auto" : "true"); + org.mockito.Mockito.mock(CertificateBasedAccess.class); + org.mockito.Mockito.when(certificateBasedAccess.getMtlsEndpointUsagePolicy()) + .thenReturn(CertificateBasedAccess.MtlsEndpointUsagePolicy.AUTO); + org.mockito.Mockito.when(certificateBasedAccess.useMtlsClientCertificate()).thenReturn(true); Truth.assertThat( defaultEndpointContextBuilder.mtlsEndpointResolver( DEFAULT_ENDPOINT, @@ -117,8 +121,10 @@ void mtlsEndpointResolver_switchToMtlsAllowedIsTrue_mtlsUsageAlways() throws IOE FakeMtlsProvider.createTestMtlsKeyStore(), "", throwExceptionForGetKeyStore); boolean switchToMtlsEndpointAllowed = true; CertificateBasedAccess certificateBasedAccess = - new CertificateBasedAccess( - name -> name.equals("GOOGLE_API_USE_MTLS_ENDPOINT") ? "always" : "true"); + org.mockito.Mockito.mock(CertificateBasedAccess.class); + org.mockito.Mockito.when(certificateBasedAccess.getMtlsEndpointUsagePolicy()) + .thenReturn(CertificateBasedAccess.MtlsEndpointUsagePolicy.ALWAYS); + org.mockito.Mockito.when(certificateBasedAccess.useMtlsClientCertificate()).thenReturn(true); Truth.assertThat( defaultEndpointContextBuilder.mtlsEndpointResolver( DEFAULT_ENDPOINT, @@ -137,8 +143,10 @@ void mtlsEndpointResolver_switchToMtlsAllowedIsTrue_mtlsUsageNever() throws IOEx FakeMtlsProvider.createTestMtlsKeyStore(), "", throwExceptionForGetKeyStore); boolean switchToMtlsEndpointAllowed = true; CertificateBasedAccess certificateBasedAccess = - new CertificateBasedAccess( - name -> name.equals("GOOGLE_API_USE_MTLS_ENDPOINT") ? "never" : "true"); + org.mockito.Mockito.mock(CertificateBasedAccess.class); + org.mockito.Mockito.when(certificateBasedAccess.getMtlsEndpointUsagePolicy()) + .thenReturn(CertificateBasedAccess.MtlsEndpointUsagePolicy.NEVER); + org.mockito.Mockito.when(certificateBasedAccess.useMtlsClientCertificate()).thenReturn(true); Truth.assertThat( defaultEndpointContextBuilder.mtlsEndpointResolver( DEFAULT_ENDPOINT, @@ -156,8 +164,10 @@ void mtlsEndpointResolver_switchToMtlsAllowedIsTrue_useCertificateIsFalse_nullMt MtlsProvider mtlsProvider = new FakeMtlsProvider(null, "", throwExceptionForGetKeyStore); boolean switchToMtlsEndpointAllowed = true; CertificateBasedAccess certificateBasedAccess = - new CertificateBasedAccess( - name -> name.equals("GOOGLE_API_USE_MTLS_ENDPOINT") ? "auto" : "false"); + org.mockito.Mockito.mock(CertificateBasedAccess.class); + org.mockito.Mockito.when(certificateBasedAccess.getMtlsEndpointUsagePolicy()) + .thenReturn(CertificateBasedAccess.MtlsEndpointUsagePolicy.AUTO); + org.mockito.Mockito.when(certificateBasedAccess.useMtlsClientCertificate()).thenReturn(false); Truth.assertThat( defaultEndpointContextBuilder.mtlsEndpointResolver( DEFAULT_ENDPOINT, @@ -174,8 +184,10 @@ void mtlsEndpointResolver_getKeyStore_throwsIOException() throws IOException { MtlsProvider mtlsProvider = new FakeMtlsProvider(null, "", throwExceptionForGetKeyStore); boolean switchToMtlsEndpointAllowed = true; CertificateBasedAccess certificateBasedAccess = - new CertificateBasedAccess( - name -> name.equals("GOOGLE_API_USE_MTLS_ENDPOINT") ? "auto" : "true"); + org.mockito.Mockito.mock(CertificateBasedAccess.class); + org.mockito.Mockito.when(certificateBasedAccess.getMtlsEndpointUsagePolicy()) + .thenReturn(CertificateBasedAccess.MtlsEndpointUsagePolicy.AUTO); + org.mockito.Mockito.when(certificateBasedAccess.useMtlsClientCertificate()).thenReturn(true); assertThrows( IOException.class, () -> @@ -272,8 +284,10 @@ void endpointContextBuild_mtlsConfigured_GDU() throws IOException { MtlsProvider mtlsProvider = new FakeMtlsProvider(FakeMtlsProvider.createTestMtlsKeyStore(), "", false); CertificateBasedAccess certificateBasedAccess = - new CertificateBasedAccess( - name -> name.equals("GOOGLE_API_USE_MTLS_ENDPOINT") ? "always" : "true"); + org.mockito.Mockito.mock(CertificateBasedAccess.class); + org.mockito.Mockito.when(certificateBasedAccess.getMtlsEndpointUsagePolicy()) + .thenReturn(CertificateBasedAccess.MtlsEndpointUsagePolicy.ALWAYS); + org.mockito.Mockito.when(certificateBasedAccess.useMtlsClientCertificate()).thenReturn(true); EndpointContext endpointContext = defaultEndpointContextBuilder .setClientSettingsEndpoint(null) @@ -293,8 +307,10 @@ void endpointContextBuild_mtlsConfigured_nonGDU_throwsIllegalArgumentException() MtlsProvider mtlsProvider = new FakeMtlsProvider(FakeMtlsProvider.createTestMtlsKeyStore(), "", false); CertificateBasedAccess certificateBasedAccess = - new CertificateBasedAccess( - name -> name.equals("GOOGLE_API_USE_MTLS_ENDPOINT") ? "always" : "true"); + org.mockito.Mockito.mock(CertificateBasedAccess.class); + org.mockito.Mockito.when(certificateBasedAccess.getMtlsEndpointUsagePolicy()) + .thenReturn(CertificateBasedAccess.MtlsEndpointUsagePolicy.ALWAYS); + org.mockito.Mockito.when(certificateBasedAccess.useMtlsClientCertificate()).thenReturn(true); EndpointContext.Builder endpointContextBuilder = defaultEndpointContextBuilder .setUniverseDomain("random.com") diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ServerStreamingAttemptCallableTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ServerStreamingAttemptCallableTest.java index 4cdd3cc018cc..740dc37b7ffe 100644 --- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ServerStreamingAttemptCallableTest.java +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ServerStreamingAttemptCallableTest.java @@ -29,6 +29,7 @@ */ package com.google.api.gax.rpc; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.Mockito.mock; import com.google.api.core.AbstractApiFuture; @@ -41,6 +42,9 @@ import com.google.api.gax.rpc.StatusCode.Code; import com.google.api.gax.rpc.testing.FakeApiException; import com.google.api.gax.rpc.testing.FakeCallContext; +import com.google.api.gax.rpc.testing.FakeChannel; +import com.google.api.gax.rpc.testing.FakeStatusCode; +import com.google.api.gax.rpc.testing.FakeTransportChannel; import com.google.api.gax.rpc.testing.MockStreamingApi.MockServerStreamingCall; import com.google.api.gax.rpc.testing.MockStreamingApi.MockServerStreamingCallable; import com.google.api.gax.tracing.BaseApiTracer; @@ -248,6 +252,191 @@ void testInitialRetry() { Truth.assertThat(call.getRequest()).isEqualTo("request > 0"); } + @Test + @SuppressWarnings("ConstantConditions") + void testUnauthenticatedRefresh() { + TransportChannel transportChannel = Mockito.mock(TransportChannel.class); + Mockito.when(transportChannel.shouldRefresh()).thenReturn(true); + + ApiCallContext context = Mockito.mock(ApiCallContext.class); + Mockito.when(context.getTransportChannel()).thenReturn(transportChannel); + Mockito.when(context.getTracer()).thenReturn(BaseApiTracer.getInstance()); + Mockito.when(context.getTimeoutDuration()).thenReturn(java.time.Duration.ofHours(5)); + + resumptionStrategy = new MyStreamResumptionStrategy(); + ServerStreamingAttemptCallable callable = createCallable(context); + callable.start(); + + MockServerStreamingCall call = innerCallable.popLastCall(); + + // Send initial error + UnauthenticatedException initialError = + new UnauthenticatedException( + "test", + null, + com.google.api.gax.rpc.testing.FakeStatusCode.of(Code.UNAUTHENTICATED), + false); + call.getController().getObserver().onError(initialError); + + // Should notify the outer future + ExecutionException ee = + assertThrows( + ExecutionException.class, + () -> fakeRetryingFuture.getAttemptResult().get(1, TimeUnit.SECONDS)); + Throwable outerError = ee.getCause(); + Mockito.verify(transportChannel).refresh(); + Truth.assertThat(outerError).isInstanceOf(ServerStreamingAttemptException.class); + Truth.assertThat(((ServerStreamingAttemptException) outerError).hasSeenResponses()).isFalse(); + Truth.assertThat(((ServerStreamingAttemptException) outerError).canResume()).isTrue(); + Truth.assertThat(outerError.getCause()).isInstanceOf(UnauthenticatedException.class); + Truth.assertThat(((UnauthenticatedException) outerError.getCause()).isRetryable()).isFalse(); + } + + @Test + @SuppressWarnings("ConstantConditions") + void testUnauthenticatedRefreshWithGenerationAdvanceRetries() { + TransportChannel transportChannel = Mockito.mock(TransportChannel.class); + java.util.concurrent.atomic.AtomicLong generation = + new java.util.concurrent.atomic.AtomicLong(0); + Mockito.when(transportChannel.getGeneration()).thenAnswer(inv -> generation.get()); + Mockito.when(transportChannel.shouldRefresh()).thenReturn(true); + Mockito.doAnswer( + inv -> { + generation.incrementAndGet(); + return null; + }) + .when(transportChannel) + .refresh(); + + ApiCallContext context = Mockito.mock(ApiCallContext.class); + Mockito.when(context.getTransportChannel()).thenReturn(transportChannel); + Mockito.when(context.getTracer()).thenReturn(BaseApiTracer.getInstance()); + Mockito.when(context.getTimeoutDuration()).thenReturn(java.time.Duration.ofHours(5)); + + resumptionStrategy = new MyStreamResumptionStrategy(); + ServerStreamingAttemptCallable callable = createCallable(context); + callable.start(); + + MockServerStreamingCall call = innerCallable.popLastCall(); + + UnauthenticatedException initialError = + new UnauthenticatedException( + "test", + null, + com.google.api.gax.rpc.testing.FakeStatusCode.of(Code.UNAUTHENTICATED), + false); + call.getController().getObserver().onError(initialError); + + ExecutionException ee = + assertThrows( + ExecutionException.class, + () -> fakeRetryingFuture.getAttemptResult().get(1, TimeUnit.SECONDS)); + Throwable outerError = ee.getCause(); + Mockito.verify(transportChannel).refresh(); + Truth.assertThat(outerError).isInstanceOf(ServerStreamingAttemptException.class); + Truth.assertThat(((ServerStreamingAttemptException) outerError).canResume()).isTrue(); + Truth.assertThat(outerError.getCause()).isInstanceOf(UnauthenticatedException.class); + Truth.assertThat(((UnauthenticatedException) outerError.getCause()).isRetryable()).isTrue(); + + // Verify retry call resumes stream + callable.call(); + call = innerCallable.popLastCall(); + Truth.assertThat(call.getRequest()).isEqualTo("request > 0"); + } + + @Test + @SuppressWarnings("ConstantConditions") + void testUnauthenticatedRefreshWithNonResumableStreamDoesNotRetry() { + TransportChannel transportChannel = Mockito.mock(TransportChannel.class); + java.util.concurrent.atomic.AtomicLong generation = + new java.util.concurrent.atomic.AtomicLong(0); + Mockito.when(transportChannel.getGeneration()).thenAnswer(inv -> generation.get()); + Mockito.when(transportChannel.shouldRefresh()).thenReturn(true); + Mockito.doAnswer( + inv -> { + generation.incrementAndGet(); + return null; + }) + .when(transportChannel) + .refresh(); + + ApiCallContext context = Mockito.mock(ApiCallContext.class); + Mockito.when(context.getTransportChannel()).thenReturn(transportChannel); + Mockito.when(context.getTracer()).thenReturn(BaseApiTracer.getInstance()); + Mockito.when(context.getTimeoutDuration()).thenReturn(java.time.Duration.ofHours(5)); + + // SimpleStreamResumptionStrategy cannot resume once a response has been received + resumptionStrategy = new com.google.api.gax.retrying.SimpleStreamResumptionStrategy<>(); + ServerStreamingAttemptCallable callable = createCallable(context); + callable.start(); + + MockServerStreamingCall call = innerCallable.popLastCall(); + call.getController().getObserver().onResponse("response1"); + + UnauthenticatedException initialError = + new UnauthenticatedException( + "test", + null, + com.google.api.gax.rpc.testing.FakeStatusCode.of(Code.UNAUTHENTICATED), + false); + call.getController().getObserver().onError(initialError); + + ExecutionException ee = + assertThrows( + ExecutionException.class, + () -> fakeRetryingFuture.getAttemptResult().get(1, TimeUnit.SECONDS)); + Throwable outerError = ee.getCause(); + Mockito.verify(transportChannel).refresh(); + Truth.assertThat(outerError).isInstanceOf(ServerStreamingAttemptException.class); + ServerStreamingAttemptException attemptEx = (ServerStreamingAttemptException) outerError; + Truth.assertThat(attemptEx.canResume()).isFalse(); + Truth.assertThat(((UnauthenticatedException) attemptEx.getCause()).isRetryable()).isTrue(); + Truth.assertThat( + new com.google.api.gax.retrying.StreamingRetryAlgorithm<>( + new ApiResultRetryAlgorithm<>(), + new com.google.api.gax.retrying.ExponentialRetryAlgorithm( + RetrySettings.newBuilder().build(), + com.google.api.core.NanoClock.getDefaultClock())) + .shouldRetry(attemptEx, null, fakeRetryingFuture.getAttemptSettings())) + .isFalse(); + } + + @Test + @SuppressWarnings("ConstantConditions") + void testRefreshThrowsException_originalErrorNotLost() { + TransportChannel transportChannel = Mockito.mock(TransportChannel.class); + Mockito.when(transportChannel.shouldRefresh()).thenReturn(true); + Mockito.doThrow(new RuntimeException("Refresh error")).when(transportChannel).refresh(); + + ApiCallContext context = Mockito.mock(ApiCallContext.class); + Mockito.when(context.getTransportChannel()).thenReturn(transportChannel); + Mockito.when(context.getTracer()).thenReturn(BaseApiTracer.getInstance()); + Mockito.when(context.getTimeoutDuration()).thenReturn(java.time.Duration.ofHours(5)); + + resumptionStrategy = new MyStreamResumptionStrategy(); + ServerStreamingAttemptCallable callable = createCallable(context); + callable.start(); + + MockServerStreamingCall call = innerCallable.popLastCall(); + + UnauthenticatedException initialError = + new UnauthenticatedException( + "test", + null, + com.google.api.gax.rpc.testing.FakeStatusCode.of(Code.UNAUTHENTICATED), + false); + call.getController().getObserver().onError(initialError); + + ExecutionException ee = + assertThrows( + ExecutionException.class, + () -> fakeRetryingFuture.getAttemptResult().get(1, TimeUnit.SECONDS)); + Throwable outerError = ee.getCause(); + Mockito.verify(transportChannel).refresh(); + Truth.assertThat(outerError).isInstanceOf(ServerStreamingAttemptException.class); + Truth.assertThat(outerError.getCause()).isEqualTo(initialError); + } + @Test @SuppressWarnings("ConstantConditions") void testMidRetry() { @@ -403,6 +592,67 @@ public String processResponse(String response) { .containsExactly("first+suffix", "second+suffix", "third+suffix"); } + @Test + void testUnauthenticatedException_whenChannelRefreshes_setsRetryableTrue() throws Exception { + FakeChannel fakeChannel = new FakeChannel(); + fakeChannel.setShouldRefresh(true); + ApiCallContext context = + FakeCallContext.createDefault() + .withTransportChannel(FakeTransportChannel.create(fakeChannel)); + + ServerStreamingAttemptCallable callable = createCallable(context); + callable.start(); + + MockServerStreamingCall call = innerCallable.popLastCall(); + UnauthenticatedException unauthEx = + new UnauthenticatedException( + "cert expired", null, new FakeStatusCode(Code.UNAUTHENTICATED), false); + call.getController().getObserver().onError(unauthEx); + + ExecutionException ex = + assertThrows( + ExecutionException.class, + () -> fakeRetryingFuture.getAttemptResult().get(1, TimeUnit.SECONDS)); + Truth.assertThat(ex.getCause()).isInstanceOf(ServerStreamingAttemptException.class); + Throwable cause = ex.getCause().getCause(); + Truth.assertThat(cause).isInstanceOf(UnauthenticatedException.class); + Truth.assertThat(((UnauthenticatedException) cause).isRetryable()).isTrue(); + Truth.assertThat(fakeChannel.getRefreshCount()).isEqualTo(1); + } + + @Test + void testUnauthenticatedException_whenChannelRefreshFails_remainsNonRetryable() throws Exception { + FakeChannel fakeChannel = + new FakeChannel() { + @Override + public void refresh() { + throw new RuntimeException("Refresh failed"); + } + }; + fakeChannel.setShouldRefresh(true); + ApiCallContext context = + FakeCallContext.createDefault() + .withTransportChannel(FakeTransportChannel.create(fakeChannel)); + + ServerStreamingAttemptCallable callable = createCallable(context); + callable.start(); + + MockServerStreamingCall call = innerCallable.popLastCall(); + UnauthenticatedException unauthEx = + new UnauthenticatedException( + "cert expired", null, new FakeStatusCode(Code.UNAUTHENTICATED), false); + call.getController().getObserver().onError(unauthEx); + + ExecutionException ex = + assertThrows( + ExecutionException.class, + () -> fakeRetryingFuture.getAttemptResult().get(1, TimeUnit.SECONDS)); + Truth.assertThat(ex.getCause()).isInstanceOf(ServerStreamingAttemptException.class); + Throwable cause = ex.getCause().getCause(); + Truth.assertThat(cause).isInstanceOf(UnauthenticatedException.class); + Truth.assertThat(((UnauthenticatedException) cause).isRetryable()).isFalse(); + } + static class MyStreamResumptionStrategy implements StreamResumptionStrategy { private int responseCount; diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/StreamingCallableTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/StreamingCallableTest.java index 6f9584826893..c46522114318 100644 --- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/StreamingCallableTest.java +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/StreamingCallableTest.java @@ -129,7 +129,7 @@ void testClientStreamingCall() { ClientStreamingCallable callable = stashCallable.withDefaultCallContext(defaultCallContext); callable.clientStreamingCall(observer); - assertSame(observer, stashCallable.getActualObserver()); + org.junit.jupiter.api.Assertions.assertNotNull(stashCallable.getActualObserver()); assertSame(defaultCallContext, stashCallable.getContext()); } @@ -158,7 +158,7 @@ void testClientStreamingCallWithContext() { ClientStreamingCallable callable = stashCallable.withDefaultCallContext(FakeCallContext.createDefault()); callable.clientStreamingCall(observer, context); - assertSame(observer, stashCallable.getActualObserver()); + org.junit.jupiter.api.Assertions.assertNotNull(stashCallable.getActualObserver()); FakeCallContext actualContext = (FakeCallContext) stashCallable.getContext(); assertSame(channel, actualContext.getChannel()); assertSame(credentials, actualContext.getCredentials()); diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/mtls/AbstractMtlsTransportChannelTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/mtls/AbstractMtlsTransportChannelTest.java index bea4674b765b..5ed563254965 100644 --- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/mtls/AbstractMtlsTransportChannelTest.java +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/mtls/AbstractMtlsTransportChannelTest.java @@ -64,8 +64,10 @@ void testNotUseClientCertificate() throws IOException, GeneralSecurityException @Test void testUseClientCertificate() throws IOException, GeneralSecurityException { CertificateBasedAccess certificateBasedAccess = - new CertificateBasedAccess( - name -> name.equals("GOOGLE_API_USE_MTLS_ENDPOINT") ? "auto" : "true"); + org.mockito.Mockito.mock(CertificateBasedAccess.class); + org.mockito.Mockito.when(certificateBasedAccess.getMtlsEndpointUsagePolicy()) + .thenReturn(CertificateBasedAccess.MtlsEndpointUsagePolicy.AUTO); + org.mockito.Mockito.when(certificateBasedAccess.useMtlsClientCertificate()).thenReturn(true); MtlsProvider provider = new FakeMtlsProvider(FakeMtlsProvider.createTestMtlsKeyStore(), "", false); assertNotNull(getMtlsObjectFromTransportChannel(provider, certificateBasedAccess)); @@ -74,8 +76,10 @@ void testUseClientCertificate() throws IOException, GeneralSecurityException { @Test void testNoClientCertificate() throws IOException, GeneralSecurityException { CertificateBasedAccess certificateBasedAccess = - new CertificateBasedAccess( - name -> name.equals("GOOGLE_API_USE_MTLS_ENDPOINT") ? "auto" : "true"); + org.mockito.Mockito.mock(CertificateBasedAccess.class); + org.mockito.Mockito.when(certificateBasedAccess.getMtlsEndpointUsagePolicy()) + .thenReturn(CertificateBasedAccess.MtlsEndpointUsagePolicy.AUTO); + org.mockito.Mockito.when(certificateBasedAccess.useMtlsClientCertificate()).thenReturn(true); MtlsProvider provider = new FakeMtlsProvider(null, "", false); assertNull(getMtlsObjectFromTransportChannel(provider, certificateBasedAccess)); } @@ -84,8 +88,10 @@ void testNoClientCertificate() throws IOException, GeneralSecurityException { void testGetKeyStoreThrows() throws GeneralSecurityException { // Test the case where provider.getKeyStore() throws. CertificateBasedAccess certificateBasedAccess = - new CertificateBasedAccess( - name -> name.equals("GOOGLE_API_USE_MTLS_ENDPOINT") ? "auto" : "true"); + org.mockito.Mockito.mock(CertificateBasedAccess.class); + org.mockito.Mockito.when(certificateBasedAccess.getMtlsEndpointUsagePolicy()) + .thenReturn(CertificateBasedAccess.MtlsEndpointUsagePolicy.AUTO); + org.mockito.Mockito.when(certificateBasedAccess.useMtlsClientCertificate()).thenReturn(true); MtlsProvider provider = new FakeMtlsProvider(null, "", true); IOException actual = assertThrows( diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/mtls/CertificateBasedAccessTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/mtls/CertificateBasedAccessTest.java index e328e0af4799..92c0dfb96f3a 100644 --- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/mtls/CertificateBasedAccessTest.java +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/mtls/CertificateBasedAccessTest.java @@ -32,52 +32,131 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import java.util.HashMap; +import java.util.Map; import org.junit.jupiter.api.Test; class CertificateBasedAccessTest { + private static class TestEnv { + private final Map env = new HashMap<>(); + + TestEnv() { + // Hermetically isolate tests from the host's ~/.config/gcloud/certificate_config.json + env.put("CLOUDSDK_CONFIG", "/nonexistent/test/gcloud"); + } + + void set(String key, String val) { + env.put(key, val); + } + + String get(String name) { + return env.get(name); + } + } + + private CertificateBasedAccess createCba(TestEnv env) { + return new CertificateBasedAccess(env::get); + } + @Test void testUseMtlsEndpointAlways() { - CertificateBasedAccess cba = - new CertificateBasedAccess( - name -> name.equals("GOOGLE_API_USE_MTLS_ENDPOINT") ? "always" : "false"); + TestEnv env = new TestEnv(); + env.set("GOOGLE_API_USE_MTLS_ENDPOINT", "always"); + CertificateBasedAccess cba = createCba(env); assertEquals( CertificateBasedAccess.MtlsEndpointUsagePolicy.ALWAYS, cba.getMtlsEndpointUsagePolicy()); } @Test void testUseMtlsEndpointAuto() { - CertificateBasedAccess cba = - new CertificateBasedAccess( - name -> name.equals("GOOGLE_API_USE_MTLS_ENDPOINT") ? "auto" : "false"); + TestEnv env = new TestEnv(); + env.set("GOOGLE_API_USE_MTLS_ENDPOINT", "auto"); + CertificateBasedAccess cba = createCba(env); assertEquals( CertificateBasedAccess.MtlsEndpointUsagePolicy.AUTO, cba.getMtlsEndpointUsagePolicy()); } @Test void testUseMtlsEndpointNever() { - CertificateBasedAccess cba = - new CertificateBasedAccess( - name -> name.equals("GOOGLE_API_USE_MTLS_ENDPOINT") ? "never" : "false"); + TestEnv env = new TestEnv(); + env.set("GOOGLE_API_USE_MTLS_ENDPOINT", "never"); + CertificateBasedAccess cba = createCba(env); + assertEquals( + CertificateBasedAccess.MtlsEndpointUsagePolicy.NEVER, cba.getMtlsEndpointUsagePolicy()); + } + + @Test + void testUseMtlsEndpointCaseInsensitive() { + TestEnv env = new TestEnv(); + env.set("GOOGLE_API_USE_MTLS_ENDPOINT", "ALWAYS"); + CertificateBasedAccess cba = createCba(env); + assertEquals( + CertificateBasedAccess.MtlsEndpointUsagePolicy.ALWAYS, cba.getMtlsEndpointUsagePolicy()); + + env.set("GOOGLE_API_USE_MTLS_ENDPOINT", "NEVER"); assertEquals( CertificateBasedAccess.MtlsEndpointUsagePolicy.NEVER, cba.getMtlsEndpointUsagePolicy()); } @Test - void testUseMtlsClientCertificateTrue() { - CertificateBasedAccess cba = - new CertificateBasedAccess( - name -> name.equals("GOOGLE_API_USE_CLIENT_CERTIFICATE") ? "true" : "auto"); + void testUseMtlsClientCertificateExplicitTrueNoCredentials() { + TestEnv env = new TestEnv(); + env.set("GOOGLE_API_USE_CLIENT_CERTIFICATE", "true"); + CertificateBasedAccess cba = createCba(env); + // Explicit 'true' enables mTLS client certificate usage (for ECP / custom MtlsProvider) even + // when no workload cert files are present, while getWorkloadCertPath returns null so file + // rotation polling is not active. assertTrue(cba.useMtlsClientCertificate()); + assertNull(cba.getWorkloadCertPath()); } @Test - void testUseMtlsClientCertificateFalse() { - CertificateBasedAccess cba = - new CertificateBasedAccess( - name -> name.equals("GOOGLE_API_USE_CLIENT_CERTIFICATE") ? "false" : "auto"); + void testUseMtlsClientCertificateExplicitFalse() { + TestEnv env = new TestEnv(); + env.set("GOOGLE_API_USE_CLIENT_CERTIFICATE", "false"); + + CertificateBasedAccess cba = createCba(env); + assertFalse(cba.useMtlsClientCertificate()); + assertNull(cba.getWorkloadCertPath()); + } + + @Test + void testUseMtlsClientCertificateUnsetNoFiles() { + TestEnv env = new TestEnv(); + CertificateBasedAccess cba = createCba(env); assertFalse(cba.useMtlsClientCertificate()); + assertNull(cba.getWorkloadCertPath()); + } + + @Test + void testUseMtlsClientCertificateConfigMissingConfigFile_throwsIllegalStateException() { + TestEnv env = new TestEnv(); + env.set("GOOGLE_API_CERTIFICATE_CONFIG", "/nonexistent/config.json"); + + CertificateBasedAccess cba = createCba(env); + + // Non-existent config file on disk specified via explicit env var throws IllegalStateException + // (Fail Closed) + assertThrows(IllegalStateException.class, () -> cba.useMtlsClientCertificate()); + assertThrows(IllegalStateException.class, () -> cba.getWorkloadCertPath()); + } + + @Test + void testWorkloadCertificateUtilsEmptyFileReturnsEmptyString() throws Exception { + java.io.File tempFile = java.io.File.createTempFile("test-cert-empty", ".pem"); + tempFile.deleteOnExit(); + // 0-byte truncated file mid-write should return empty string rather than SHA-256 of empty bytes + assertEquals( + "", WorkloadCertificateUtils.getCertificateFingerprint(tempFile.getAbsolutePath())); + + java.nio.file.Files.write( + tempFile.toPath(), "test-cert-content".getBytes(java.nio.charset.StandardCharsets.UTF_8)); + String fp = WorkloadCertificateUtils.getCertificateFingerprint(tempFile.getAbsolutePath()); + assertFalse(fp.isEmpty()); } } diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/testing/FakeCallContext.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/testing/FakeCallContext.java index 1cdefe435d55..e41dc041220d 100644 --- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/testing/FakeCallContext.java +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/testing/FakeCallContext.java @@ -241,6 +241,11 @@ public FakeChannel getChannel() { return channel; } + @Override + public TransportChannel getTransportChannel() { + return channel != null ? FakeTransportChannel.create(channel) : null; + } + @Override public java.time.Duration getTimeoutDuration() { return timeout; diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/testing/FakeChannel.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/testing/FakeChannel.java index b725da363b3b..9acfc703b9c3 100644 --- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/testing/FakeChannel.java +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/testing/FakeChannel.java @@ -32,4 +32,36 @@ import com.google.api.core.InternalApi; @InternalApi("for testing") -public class FakeChannel {} +public class FakeChannel { + private volatile boolean shouldRefresh = false; + private volatile int refreshCount = 0; + + public FakeChannel setShouldRefresh(boolean shouldRefresh) { + this.shouldRefresh = shouldRefresh; + return this; + } + + public boolean shouldRefresh() { + return shouldRefresh; + } + + public void refresh() { + refreshCount++; + generation++; + } + + public int getRefreshCount() { + return refreshCount; + } + + private volatile long generation = 0; + + public FakeChannel setGeneration(long generation) { + this.generation = generation; + return this; + } + + public long getGeneration() { + return generation; + } +} diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/testing/FakeTransportChannel.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/testing/FakeTransportChannel.java index 0d4abac8f1c6..7d92556b88d5 100644 --- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/testing/FakeTransportChannel.java +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/testing/FakeTransportChannel.java @@ -41,6 +41,48 @@ public class FakeTransportChannel implements TransportChannel { private volatile boolean isShutdown = false; private volatile Map headers; private volatile Executor executor; + private volatile boolean shouldRefresh = false; + private volatile int refreshCount = 0; + + public FakeTransportChannel setShouldRefresh(boolean shouldRefresh) { + if (channel != null) { + channel.setShouldRefresh(shouldRefresh); + } + this.shouldRefresh = shouldRefresh; + return this; + } + + @Override + public boolean shouldRefresh() { + return channel != null ? channel.shouldRefresh() : shouldRefresh; + } + + @Override + public void refresh() { + if (channel != null) { + channel.refresh(); + } + refreshCount++; + } + + public int getRefreshCount() { + return channel != null ? channel.getRefreshCount() : refreshCount; + } + + private volatile long generation = 0; + + public FakeTransportChannel setGeneration(long generation) { + if (channel != null) { + channel.setGeneration(generation); + } + this.generation = generation; + return this; + } + + @Override + public long getGeneration() { + return channel != null ? channel.getGeneration() : generation; + } private FakeTransportChannel(FakeChannel channel) { this.channel = channel;