Store additional display claims returned by XSTS step

Closes #26
RaphiMC · Jun 19, 2024 · 99bed8b · 99bed8b
1 parent eefd2d2
commit 99bed8b
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 1 deletion.
diff --git a/src/main/java/net/raphimc/minecraftauth/step/xbl/StepXblSisuAuthentication.java b/src/main/java/net/raphimc/minecraftauth/step/xbl/StepXblSisuAuthentication.java
@@ -32,6 +32,7 @@
 
 import java.time.Instant;
 import java.time.ZoneId;
+import java.util.Map;
 
 public class StepXblSisuAuthentication extends AbstractStep<StepInitialXblSession.InitialXblSession, StepXblSisuAuthentication.XblSisuTokens> {
 
@@ -134,6 +135,11 @@ public String getUserHash() {
             return this.xstsToken.getUserHash();
         }
 
+        @Override
+        public Map<String, String> getDisplayClaims() {
+            return this.xstsToken.getDisplayClaims();
+        }
+
         @Override
         public StepFullXblSession.FullXblSession getFullXblSession() {
             return this.xstsToken.getFullXblSession();

diff --git a/src/main/java/net/raphimc/minecraftauth/step/xbl/StepXblXstsToken.java b/src/main/java/net/raphimc/minecraftauth/step/xbl/StepXblXstsToken.java
@@ -31,10 +31,14 @@
 import net.raphimc.minecraftauth.step.xbl.session.StepInitialXblSession;
 import net.raphimc.minecraftauth.util.CryptUtil;
 import net.raphimc.minecraftauth.util.JsonContent;
+import net.raphimc.minecraftauth.util.JsonUtil;
 import org.jetbrains.annotations.ApiStatus;
 
 import java.time.Instant;
 import java.time.ZoneId;
+import java.util.Collections;
+import java.util.Map;
+import java.util.stream.Collectors;
 
 public class StepXblXstsToken extends AbstractStep<StepFullXblSession.FullXblSession, StepXblXstsToken.XblXstsToken> {
 
@@ -103,24 +107,30 @@ public static class XblXstsToken extends XblXsts<StepFullXblSession.FullXblSessi
         long expireTimeMs;
         String token;
         String userHash;
+        Map<String, String> displayClaims;
         StepFullXblSession.FullXblSession fullXblSession;
 
         @ApiStatus.Internal
         public static XblXstsToken fromMicrosoftJson(final JsonObject obj, final StepFullXblSession.FullXblSession fullXblSession) {
+            final JsonObject displayClaims = obj.getAsJsonObject("DisplayClaims").getAsJsonArray("xui").get(0).getAsJsonObject();
+            final Map<String, String> displayClaimsMap = displayClaims.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().getAsString()));
             return new XblXstsToken(
                     Instant.parse(obj.get("NotAfter").getAsString()).toEpochMilli(),
                     obj.get("Token").getAsString(),
-                    obj.getAsJsonObject("DisplayClaims").getAsJsonArray("xui").get(0).getAsJsonObject().get("uhs").getAsString(),
+                    displayClaims.get("uhs").getAsString(),
+                    Collections.unmodifiableMap(displayClaimsMap),
                     fullXblSession
             );
         }
 
         @ApiStatus.Internal
         public static XblXstsToken fromJson(final JsonObject json, final StepFullXblSession.FullXblSession fullXblSession) {
+            final JsonObject displayClaims = JsonUtil.getObjectNonNull(json, "displayClaims");
             return new XblXstsToken(
                     json.get("expireTimeMs").getAsLong(),
                     json.get("token").getAsString(),
                     json.get("userHash").getAsString(),
+                    Collections.unmodifiableMap(displayClaims.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().getAsString()))),
                     fullXblSession
             );
         }
@@ -131,6 +141,7 @@ public static JsonObject toJson(final XblXstsToken xblXstsToken) {
             json.addProperty("expireTimeMs", xblXstsToken.expireTimeMs);
             json.addProperty("token", xblXstsToken.token);
             json.addProperty("userHash", xblXstsToken.userHash);
+            json.add("displayClaims", JsonUtil.GSON.toJsonTree(xblXstsToken.displayClaims));
             return json;
         }
 
@@ -154,6 +165,8 @@ public abstract static class XblXsts<P extends AbstractStep.StepResult<?>> exten
 
         public abstract String getUserHash();
 
+        public abstract Map<String, String> getDisplayClaims();
+
         public String getServiceToken() {
             return this.getUserHash() + ';' + this.getToken();
         }

diff --git a/src/main/java/net/raphimc/minecraftauth/util/JsonUtil.java b/src/main/java/net/raphimc/minecraftauth/util/JsonUtil.java
@@ -65,4 +65,13 @@ public static long getLongOr(final JsonObject obj, final String key, final long
         }
     }
 
+    public static JsonObject getObjectNonNull(final JsonObject obj, final String key) {
+        final JsonElement element = obj.get(key);
+        if (element != null && !element.isJsonNull()) {
+            return element.getAsJsonObject();
+        } else {
+            return new JsonObject();
+        }
+    }
+
 }