updating redirect implementation

adrian05-ms · adrian05-ms · commit a4970d201040 · 2026-02-27T13:47:26.000-06:00
diff --git a/components/http/okHttp/src/main/java/com/microsoft/kiota/http/middleware/RedirectHandler.java b/components/http/okHttp/src/main/java/com/microsoft/kiota/http/middleware/RedirectHandler.java
@@ -81,7 +81,11 @@ boolean isRedirected(
         return false;
     }
 
-    Request getRedirect(final Request request, final Response userResponse)
+    Request getRedirect(
+            final Request request,
+            final Response userResponse,
+            final RedirectHandlerOption redirectOption,
+            final Chain chain)
             throws ProtocolException {
         String location = userResponse.header("Location");
         if (location == null || location.length() == 0) return null;
@@ -105,18 +109,12 @@ Request getRedirect(final Request request, final Response userResponse)
         // Most redirects don't include a request body.
         Request.Builder requestBuilder = userResponse.request().newBuilder();
 
-        // When redirecting across hosts, drop all authentication headers. This
-        // is potentially annoying to the application layer since they have no
-        // way to retain them.
-        boolean sameScheme = locationUrl.scheme().equalsIgnoreCase(requestUrl.scheme());
-        boolean sameHost =
-                locationUrl.host().toString().equalsIgnoreCase(requestUrl.host().toString());
-        boolean samePort = locationUrl.port() == requestUrl.port();
-        if (!sameScheme || !sameHost || !samePort) {
-            requestBuilder.removeHeader("Authorization");
-            requestBuilder.removeHeader("Cookie");
-            requestBuilder.removeHeader("Proxy-Authorization");
-        }
+        // Scrub sensitive headers before following the redirect
+        java.util.function.Function<HttpUrl, java.net.Proxy> proxyResolver =
+                RedirectHandlerOption.getProxyResolver(
+                        chain.call().client().proxySelector());
+        redirectOption.scrubSensitiveHeaders().scrubHeaders(
+                requestBuilder, requestUrl, locationUrl, proxyResolver);
 
         // Response status code 303 See Other then POST changes to GET
         if (userResponse.code() == HTTP_SEE_OTHER) {
@@ -166,7 +164,8 @@ Request getRedirect(final Request request, final Response userResponse)
                         isRedirected(request, response, requestsCount, redirectOption)
                                 && redirectOption.shouldRedirect().shouldRedirect(response);
 
-                final Request followup = shouldRedirect ? getRedirect(request, response) : null;
+                final Request followup =
+                        shouldRedirect ? getRedirect(request, response, redirectOption, chain) : null;
                 if (followup != null) {
                     response.close();
                     request = followup;
diff --git a/components/http/okHttp/src/main/java/com/microsoft/kiota/http/middleware/options/RedirectHandlerOption.java b/components/http/okHttp/src/main/java/com/microsoft/kiota/http/middleware/options/RedirectHandlerOption.java
@@ -5,6 +5,17 @@
 import jakarta.annotation.Nonnull;
 import jakarta.annotation.Nullable;
 
+import okhttp3.HttpUrl;
+import okhttp3.Request;
+
+import java.net.InetSocketAddress;
+import java.net.Proxy;
+import java.net.ProxySelector;
+import java.net.URI;
+import java.util.List;
+import java.util.Objects;
+import java.util.function.Function;
+
 /**
  * Options to be passed to the redirect middleware.
  */
@@ -28,11 +39,59 @@ public class RedirectHandlerOption implements RequestOption {
      */
     @Nonnull public static final IShouldRedirect DEFAULT_SHOULD_REDIRECT = response -> true;
 
+    @Nonnull private final IScrubSensitiveHeaders scrubSensitiveHeaders;
+
+    /**
+     * Functional interface for scrubbing sensitive headers during redirects.
+     */
+    @FunctionalInterface
+    public interface IScrubSensitiveHeaders {
+        /**
+         * Scrubs sensitive headers from the request before following a redirect.
+         * @param requestBuilder The request builder to modify
+         * @param originalUrl The original request URL
+         * @param newUrl The new redirect URL
+         * @param proxyResolver A function that returns the proxy for a given destination, or null if no proxy applies
+         */
+        void scrubHeaders(
+                @Nonnull Request.Builder requestBuilder,
+                @Nonnull HttpUrl originalUrl,
+                @Nonnull HttpUrl newUrl,
+                @Nullable Function<HttpUrl, Proxy> proxyResolver);
+    }
+
+    /**
+     * The default implementation for scrubbing sensitive headers during redirects.
+     * This method removes Authorization and Cookie headers when the host or scheme changes,
+     * and removes Proxy-Authorization headers when no proxy is configured or the proxy is bypassed for the new URL.
+     */
+    @Nonnull public static final IScrubSensitiveHeaders DEFAULT_SCRUB_SENSITIVE_HEADERS =
+            (requestBuilder, originalUrl, newUrl, proxyResolver) -> {
+                Objects.requireNonNull(requestBuilder, "parameter requestBuilder cannot be null");
+                Objects.requireNonNull(originalUrl, "parameter originalUrl cannot be null");
+                Objects.requireNonNull(newUrl, "parameter newUrl cannot be null");
+
+                // Remove Authorization and Cookie headers if the request's scheme or host changes
+                boolean isDifferentHostOrScheme =
+                        !newUrl.host().equalsIgnoreCase(originalUrl.host())
+                                || !newUrl.scheme().equalsIgnoreCase(originalUrl.scheme());
+                if (isDifferentHostOrScheme) {
+                    requestBuilder.removeHeader("Authorization");
+                    requestBuilder.removeHeader("Cookie");
+                }
+
+                // Remove Proxy-Authorization if no proxy is configured or the URL is bypassed
+                boolean isProxyInactive = proxyResolver == null || proxyResolver.apply(newUrl) == null;
+                if (isProxyInactive) {
+                    requestBuilder.removeHeader("Proxy-Authorization");
+                }
+            };
+
     /**
      * Create default instance of redirect options, with default values of max redirects and should redirect
      */
     public RedirectHandlerOption() {
-        this(DEFAULT_MAX_REDIRECTS, DEFAULT_SHOULD_REDIRECT);
+        this(DEFAULT_MAX_REDIRECTS, DEFAULT_SHOULD_REDIRECT, DEFAULT_SCRUB_SENSITIVE_HEADERS);
     }
 
     /**
@@ -41,13 +100,30 @@ public RedirectHandlerOption() {
      * @param shouldRedirect Should redirect callback called before every redirect
      */
     public RedirectHandlerOption(int maxRedirects, @Nullable final IShouldRedirect shouldRedirect) {
+        this(maxRedirects, shouldRedirect, DEFAULT_SCRUB_SENSITIVE_HEADERS);
+    }
+
+    /**
+     * Create an instance with provided values
+     * @param maxRedirects Max redirects to occur
+     * @param shouldRedirect Should redirect callback called before every redirect
+     * @param scrubSensitiveHeaders Callback to scrub sensitive headers during redirects
+     */
+    public RedirectHandlerOption(
+            int maxRedirects,
+            @Nullable final IShouldRedirect shouldRedirect,
+            @Nullable final IScrubSensitiveHeaders scrubSensitiveHeaders) {
         if (maxRedirects < 0)
             throw new IllegalArgumentException("Max redirects cannot be negative");
         if (maxRedirects > MAX_REDIRECTS)
             throw new IllegalArgumentException("Max redirect cannot exceed " + MAX_REDIRECTS);
 
         this.maxRedirects = maxRedirects;
         this.shouldRedirect = shouldRedirect != null ? shouldRedirect : DEFAULT_SHOULD_REDIRECT;
+        this.scrubSensitiveHeaders =
+                scrubSensitiveHeaders != null
+                        ? scrubSensitiveHeaders
+                        : DEFAULT_SCRUB_SENSITIVE_HEADERS;
     }
 
     /**
@@ -66,6 +142,40 @@ public int maxRedirects() {
         return this.shouldRedirect;
     }
 
+    /**
+     * Gets the callback for scrubbing sensitive headers during redirects.
+     * @return scrub sensitive headers callback
+     */
+    @Nonnull public IScrubSensitiveHeaders scrubSensitiveHeaders() {
+        return this.scrubSensitiveHeaders;
+    }
+
+    /**
+     * Helper method to get a proxy resolver from a ProxySelector.
+     * @param proxySelector The ProxySelector to use, or null if no proxy is configured
+     * @return A function that resolves proxies for a given HttpUrl, or null if no proxy selector is provided
+     */
+    @Nullable public static Function<HttpUrl, Proxy> getProxyResolver(
+            @Nullable final ProxySelector proxySelector) {
+        if (proxySelector == null) {
+            return null;
+        }
+        return url -> {
+            try {
+                URI uri = new URI(url.scheme(), null, url.host(), url.port(), null, null, null);
+                List<Proxy> proxies = proxySelector.select(uri);
+                if (proxies != null && !proxies.isEmpty()) {
+                    Proxy proxy = proxies.get(0);
+                    // Return null for DIRECT proxies (no proxy)
+                    return proxy.type() == Proxy.Type.DIRECT ? null : proxy;
+                }
+                return null;
+            } catch (Exception e) {
+                return null;
+            }
+        };
+    }
+
     /** {@inheritDoc} */
     @SuppressWarnings("unchecked")
     @Override
diff --git a/components/http/okHttp/src/test/java/com/microsoft/kiota/http/middleware/RedirectHandlerTests.java b/components/http/okHttp/src/test/java/com/microsoft/kiota/http/middleware/RedirectHandlerTests.java
