fix: suggested changes

Author: AgraVator

api/src/main/java/io/grpc/ManagedChannelProvider.java

@@ -98,7 +98,7 @@ protected NewChannelBuilderResult newChannelBuilder(String target, ChannelCreden
    *     search the registry
    * @return a {@link NewChannelBuilderResult} containing either the builder or an
    *     error description
-   * @since 1.79.0
+   * @since 1.81.0
    */
   protected NewChannelBuilderResult newChannelBuilder(String target, ChannelCredentials creds,
                                                       NameResolverRegistry nameResolverRegistry,

api/src/main/java/io/grpc/NameResolverProvider.java

@@ -65,7 +65,7 @@ public abstract class NameResolverProvider extends NameResolver.Factory {
    *
    * @since 1.40.0
    * */
-  protected String getScheme() {
+  public String getScheme() {
     return getDefaultScheme();
   }
 

core/src/main/java/io/grpc/internal/ManagedChannelImplBuilder.java

@@ -762,7 +762,7 @@ public ManagedChannel build() {
         clientTransportFactoryBuilder.buildClientTransportFactory();
     ResolvedNameResolver resolvedResolver =
         InternalFeatureFlags.getRfc3986UrisEnabled()
-            ? getNameResolverProviderRfc3986(target, nameResolverRegistry)
+            ? getNameResolverProviderRfc3986(target, nameResolverRegistry, nameResolverProvider)
             : getNameResolverProvider(target, nameResolverRegistry, nameResolverProvider);
     resolvedResolver.checkAddressTypes(clientTransportFactory.getSupportedSocketAddressTypes());
     return new ManagedChannelOrphanWrapper(new ManagedChannelImpl(
@@ -902,7 +902,7 @@ static ResolvedNameResolver getNameResolverProvider(
       // parsed as the scheme. Will hit the next case and try "dns:///localhost:8080".
       // Use the explicit provider if its scheme matches the target URI.
       if (nameResolverProvider != null
-          && targetUri.getScheme().equals(nameResolverProvider.getDefaultScheme())) {
+          && targetUri.getScheme().equals(nameResolverProvider.getScheme())) {
         provider = nameResolverProvider;
       } else {
         provider = nameResolverRegistry.getProviderForScheme(targetUri.getScheme());
@@ -914,7 +914,7 @@ static ResolvedNameResolver getNameResolverProvider(
       // the default scheme from the registry (if provider is not specified) or
       // the provider's default scheme (if provider is specified).
       String scheme = nameResolverProvider != null
-          ? nameResolverProvider.getDefaultScheme()
+          ? nameResolverProvider.getScheme()
           : nameResolverRegistry.getDefaultScheme();
       try {
         targetUri = new URI(scheme, "", "/" + target, null);
@@ -940,7 +940,8 @@ static ResolvedNameResolver getNameResolverProvider(
 
   @VisibleForTesting
   static ResolvedNameResolver getNameResolverProviderRfc3986(
-      String target, NameResolverRegistry nameResolverRegistry) {
+      String target, NameResolverRegistry nameResolverRegistry,
+      NameResolverProvider nameResolverProvider) {
     // Finding a NameResolver. Try using the target string as the URI. If that fails, try prepending
     // "dns:///".
     NameResolverProvider provider = null;
@@ -955,15 +956,25 @@ static ResolvedNameResolver getNameResolverProviderRfc3986(
     if (targetUri != null) {
       // For "localhost:8080" this would likely cause provider to be null, because "localhost" is
       // parsed as the scheme. Will hit the next case and try "dns:///localhost:8080".
-      provider = nameResolverRegistry.getProviderForScheme(targetUri.getScheme());
+      // Use the explicit provider if its scheme matches the target URI.
+      if (nameResolverProvider != null
+          && targetUri.getScheme().equals(nameResolverProvider.getScheme())) {
+        provider = nameResolverProvider;
+      } else {
+        provider = nameResolverRegistry.getProviderForScheme(targetUri.getScheme());
+      }
     }
 
     if (provider == null && !URI_PATTERN.matcher(target).matches()) {
-      // It doesn't look like a URI target. Maybe it's an authority string. Try with the default
-      // scheme from the registry.
+      // It doesn't look like a URI target. Maybe it's an authority string. Try with
+      // the default scheme from the registry (if provider is not specified) or
+      // the provider's default scheme (if provider is specified).
+      String scheme = nameResolverProvider != null
+          ? nameResolverProvider.getScheme()
+          : nameResolverRegistry.getDefaultScheme();
       targetUri =
           Uri.newBuilder()
-              .setScheme(nameResolverRegistry.getDefaultScheme())
+              .setScheme(scheme)
               .setHost("")
               .setPath("/" + target)
               .build();

core/src/test/java/io/grpc/internal/ManagedChannelImplBuilderTest.java

@@ -808,6 +808,7 @@ public void getNameResolverProvider_explicitProviderWithIpTarget() {
     String target = "127.0.0.1:8080";
     NameResolverRegistry registry = new NameResolverRegistry();
     NameResolverProvider explicitProvider = mock(NameResolverProvider.class);
+    when(explicitProvider.getScheme()).thenReturn("dns");
     when(explicitProvider.getDefaultScheme()).thenReturn("dns");
 
     ManagedChannelImplBuilder.ResolvedNameResolver resolved;
@@ -823,6 +824,7 @@ public void getNameResolverProvider_explicitProviderWithInvalidUri() {
     String target = "::1";
     NameResolverRegistry registry = new NameResolverRegistry();
     NameResolverProvider explicitProvider = mock(NameResolverProvider.class);
+    when(explicitProvider.getScheme()).thenReturn("dns");
     when(explicitProvider.getDefaultScheme()).thenReturn("dns");
 
     ManagedChannelImplBuilder.ResolvedNameResolver resolved;
@@ -907,4 +909,43 @@ protected int priority() {
 
     assertThat(resolved.provider).isSameInstanceAs(registryProvider);
   }
+
+  @Test
+  public void getNameResolverProviderRfc3986_explicitProviderWithAuthorityTarget() {
+    String target = "authority-target";
+    NameResolverRegistry registry = new NameResolverRegistry();
+    NameResolverProvider explicitProvider = new NameResolverProvider() {
+      @Override
+      public NameResolver newNameResolver(URI targetUri, NameResolver.Args args) {
+        return null;
+      }
+
+      @Override
+      public String getScheme() {
+        return "customscheme";
+      }
+
+      @Override
+      public String getDefaultScheme() {
+        return "customscheme";
+      }
+
+      @Override
+      protected boolean isAvailable() {
+        return true;
+      }
+
+      @Override
+      protected int priority() {
+        return 5;
+      }
+    };
+    registry.register(explicitProvider);
+
+    ManagedChannelImplBuilder.ResolvedNameResolver resolved = ManagedChannelImplBuilder
+        .getNameResolverProviderRfc3986(target, registry, explicitProvider);
+
+    assertThat(resolved.provider).isSameInstanceAs(explicitProvider);
+    assertThat(resolved.targetUri.toString()).isEqualTo("customscheme:///authority-target");
+  }
 }

core/src/test/java/io/grpc/internal/ManagedChannelImplGetNameResolverRfc3986Test.java

@@ -136,7 +136,7 @@ public void validTargetNoProvider() {
     NameResolverRegistry nameResolverRegistry = new NameResolverRegistry();
     try {
       ManagedChannelImplBuilder.getNameResolverProviderRfc3986(
-          "foo.googleapis.com:8080", nameResolverRegistry);
+          "foo.googleapis.com:8080", nameResolverRegistry, null);
       fail("Should fail");
     } catch (IllegalArgumentException e) {
       // expected
@@ -148,7 +148,7 @@ public void validTargetProviderAddrTypesNotSupported() {
     NameResolverRegistry nameResolverRegistry = getTestRegistry("testscheme");
     try {
       ManagedChannelImplBuilder.getNameResolverProviderRfc3986(
-              "testscheme:///foo.googleapis.com:8080", nameResolverRegistry)
+              "testscheme:///foo.googleapis.com:8080", nameResolverRegistry, null)
           .checkAddressTypes(Collections.singleton(CustomSocketAddress.class));
       fail("Should fail");
     } catch (IllegalArgumentException e) {
@@ -163,7 +163,8 @@ public void validTargetProviderAddrTypesNotSupported() {
   private void testValidTarget(String target, String expectedUriString, Uri expectedUri) {
     NameResolverRegistry nameResolverRegistry = getTestRegistry(expectedUri.getScheme());
     ManagedChannelImplBuilder.ResolvedNameResolver resolved =
-        ManagedChannelImplBuilder.getNameResolverProviderRfc3986(target, nameResolverRegistry);
+        ManagedChannelImplBuilder.getNameResolverProviderRfc3986(target, nameResolverRegistry,
+            null);
     assertThat(resolved.provider).isInstanceOf(FakeNameResolverProvider.class);
     assertThat(resolved.targetUri).isEqualTo(wrap(expectedUri));
     assertThat(resolved.targetUri.toString()).isEqualTo(expectedUriString);
@@ -174,7 +175,8 @@ private void testInvalidTarget(String target) {
 
     try {
       ManagedChannelImplBuilder.ResolvedNameResolver resolved =
-          ManagedChannelImplBuilder.getNameResolverProviderRfc3986(target, nameResolverRegistry);
+          ManagedChannelImplBuilder.getNameResolverProviderRfc3986(target, nameResolverRegistry,
+              null);
       FakeNameResolverProvider nameResolverProvider = (FakeNameResolverProvider) resolved.provider;
       fail("Should have failed, but got resolver provider " + nameResolverProvider);
     } catch (IllegalArgumentException e) {