From 507f9e3db338932f16cdb88c7fd52b450b181b51 Mon Sep 17 00:00:00 2001
From: SammyO <1315628+SammyO@users.noreply.github.com>
Date: Mon, 12 Feb 2024 12:19:10 +0000
Subject: [PATCH] Add support for DC and flight SDK request parameters, and
 attach to /authorize and /token requests

---
 .../identity/client/TokenParameters.java      | 35 +++++++++++++++++++
 .../internal/CommandParametersAdapter.java    |  6 ++++
 2 files changed, 41 insertions(+)

diff --git a/msal/src/main/java/com/microsoft/identity/client/TokenParameters.java b/msal/src/main/java/com/microsoft/identity/client/TokenParameters.java
index 3112532df..5a7b8b3ab 100644
--- a/msal/src/main/java/com/microsoft/identity/client/TokenParameters.java
+++ b/msal/src/main/java/com/microsoft/identity/client/TokenParameters.java
@@ -34,6 +34,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Locale;
+import java.util.Map;
 import java.util.UUID;
 
 /**
@@ -48,6 +49,8 @@ public abstract class TokenParameters {
     private AccountRecord mAccountRecord;
     private AuthenticationScheme mAuthenticationScheme;
     private String mCorrelationId;
+    private String mDc;
+    private Map<String, String> mFlightParams;
 
     protected TokenParameters(@NonNull final TokenParameters.Builder builder) {
         mAccount = builder.mAccount;
@@ -56,6 +59,8 @@ protected TokenParameters(@NonNull final TokenParameters.Builder builder) {
         mScopes = builder.mScopes;
         mAuthenticationScheme = builder.mAuthenticationScheme;
         mCorrelationId = builder.mCorrelationId;
+        mDc = builder.mDc;
+        mFlightParams = builder.mFlightParams;
     }
 
     /**
@@ -154,6 +159,24 @@ public String getCorrelationId() {
         return mCorrelationId;
     }
 
+    /**
+     * Optional DC parameter. If set, API requests will contain a DC query parameter with the
+     * respective value.
+     * @return
+     */
+    public String getDc() {
+        return mDc;
+    }
+
+    /**
+     * Optional flight parameters. If set, API requests will contain the keys and values from the
+     * respective map set as query parameters.
+     * @return
+     */
+    public Map<String, String> getFlightParameters() {
+        return mFlightParams;
+    }
+
     /**
      * TokenParameters builder
      *
@@ -167,6 +190,8 @@ public static abstract class Builder<B extends TokenParameters.Builder<B>> {
         private ClaimsRequest mClaimsRequest;
         private AuthenticationScheme mAuthenticationScheme;
         private String mCorrelationId;
+        private String mDc;
+        private Map<String, String> mFlightParams;
 
         public B withAuthenticationScheme(@NonNull final AuthenticationScheme scheme) {
             mAuthenticationScheme = scheme;
@@ -261,6 +286,16 @@ public B withCorrelationId(@NonNull final UUID correlationId) {
             return self();
         }
 
+        public B withDc(@NonNull final String dc) {
+            mDc = dc;
+            return self();
+        }
+
+        public B withFlightParameters(@NonNull final Map<String, String> flightParameters) {
+            mFlightParams = flightParameters;
+            return self();
+        }
+
         public abstract B self();
 
         public abstract TokenParameters build();
diff --git a/msal/src/main/java/com/microsoft/identity/client/internal/CommandParametersAdapter.java b/msal/src/main/java/com/microsoft/identity/client/internal/CommandParametersAdapter.java
index 62fdd5ce4..7d0b7abe4 100644
--- a/msal/src/main/java/com/microsoft/identity/client/internal/CommandParametersAdapter.java
+++ b/msal/src/main/java/com/microsoft/identity/client/internal/CommandParametersAdapter.java
@@ -199,6 +199,8 @@ public static InteractiveTokenCommandParameters createInteractiveTokenCommandPar
                 .handleNullTaskAffinity(configuration.isHandleNullTaskAffinityEnabled())
                 .powerOptCheckEnabled(configuration.isPowerOptCheckForEnabled())
                 .correlationId(parameters.getCorrelationId())
+                .dc(parameters.getDc())
+                .flightInformation(parameters.getFlightParameters())
                 .preferredAuthMethod(parameters.getPreferredAuthMethod())
                 .build();
 
@@ -249,6 +251,8 @@ public static SilentTokenCommandParameters createSilentTokenCommandParameters(
                 .scopes(new HashSet<>(parameters.getScopes()))
                 .powerOptCheckEnabled(configuration.isPowerOptCheckForEnabled())
                 .correlationId(parameters.getCorrelationId())
+                .dc(parameters.getDc())
+                .flightInformation(parameters.getFlightParameters())
                 .build();
 
         return commandParameters;
@@ -289,6 +293,8 @@ public static DeviceCodeFlowCommandParameters createDeviceCodeFlowWithClaimsComm
                 .authority(authority)
                 .claimsRequestJson(claimsRequestJson)
                 .correlationId(parameters.getCorrelationId())
+                .dc(parameters.getDc())
+                .flightInformation(parameters.getFlightParameters())
                 .build();
 
         return commandParameters;