Multiple Uids Cookies Support (#3668)

Author: AntoxaAntoxic

src/main/java/org/prebid/server/cookie/PrioritizedCoopSyncProvider.java

@@ -74,8 +74,4 @@ public boolean isPrioritizedFamily(String cookieFamilyName) {
         final String bidder = prioritizedCookieFamilyNameToBidderName.get(cookieFamilyName);
         return prioritizedBidders.contains(bidder);
     }
-
-    public boolean hasPrioritizedBidders() {
-        return !prioritizedBidders.isEmpty();
-    }
 }

src/main/java/org/prebid/server/cookie/UidsCookie.java

@@ -102,7 +102,7 @@ public UidsCookie updateOptout(boolean optout) {
     /**
      * Converts {@link Uids} to JSON string.
      */
-    String toJson() {
+    public String toJson() {
         return mapper.encodeToString(uids);
     }
 

src/main/java/org/prebid/server/cookie/UidsCookieService.java

@@ -6,20 +6,23 @@
 import io.vertx.ext.web.RoutingContext;
 import org.apache.commons.lang3.StringUtils;
 import org.prebid.server.cookie.model.UidWithExpiry;
-import org.prebid.server.cookie.model.UidsCookieUpdateResult;
 import org.prebid.server.cookie.proto.Uids;
 import org.prebid.server.json.DecodeException;
 import org.prebid.server.json.JacksonMapper;
 import org.prebid.server.log.Logger;
 import org.prebid.server.log.LoggerFactory;
 import org.prebid.server.metric.Metrics;
 import org.prebid.server.model.HttpRequestContext;
+import org.prebid.server.model.UpdateResult;
 import org.prebid.server.util.HttpUtil;
 
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Base64;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
@@ -34,15 +37,20 @@ public class UidsCookieService {
     private static final Logger logger = LoggerFactory.getLogger(UidsCookieService.class);
 
     private static final String COOKIE_NAME = "uids";
+    private static final String COOKIE_NAME_FORMAT = "uids%d";
     private static final int MIN_COOKIE_SIZE_BYTES = 500;
+    private static final int MIN_NUMBER_OF_UID_COOKIES = 1;
+    private static final int MAX_NUMBER_OF_UID_COOKIES = 30;
 
     private final String optOutCookieName;
     private final String optOutCookieValue;
     private final String hostCookieFamily;
     private final String hostCookieName;
     private final String hostCookieDomain;
     private final long ttlSeconds;
+
     private final int maxCookieSizeBytes;
+    private final int numberOfUidCookies;
 
     private final PrioritizedCoopSyncProvider prioritizedCoopSyncProvider;
     private final Metrics metrics;
@@ -55,6 +63,7 @@ public UidsCookieService(String optOutCookieName,
                              String hostCookieDomain,
                              int ttlDays,
                              int maxCookieSizeBytes,
+                             int numberOfUidCookies,
                              PrioritizedCoopSyncProvider prioritizedCoopSyncProvider,
                              Metrics metrics,
                              JacksonMapper mapper) {
@@ -64,13 +73,20 @@ public UidsCookieService(String optOutCookieName,
                     "Configured cookie size is less than allowed minimum size of " + MIN_COOKIE_SIZE_BYTES);
         }
 
+        if (numberOfUidCookies < MIN_NUMBER_OF_UID_COOKIES || numberOfUidCookies > MAX_NUMBER_OF_UID_COOKIES) {
+            throw new IllegalArgumentException(
+                    "Configured number of uid cookies should be in the range from %d to %d"
+                            .formatted(MIN_NUMBER_OF_UID_COOKIES, MAX_NUMBER_OF_UID_COOKIES));
+        }
+
         this.optOutCookieName = optOutCookieName;
         this.optOutCookieValue = optOutCookieValue;
         this.hostCookieFamily = hostCookieFamily;
         this.hostCookieName = hostCookieName;
         this.hostCookieDomain = StringUtils.isNotBlank(hostCookieDomain) ? hostCookieDomain : null;
         this.ttlSeconds = Duration.ofDays(ttlDays).getSeconds();
         this.maxCookieSizeBytes = maxCookieSizeBytes;
+        this.numberOfUidCookies = numberOfUidCookies;
         this.prioritizedCoopSyncProvider = Objects.requireNonNull(prioritizedCoopSyncProvider);
         this.metrics = Objects.requireNonNull(metrics);
         this.mapper = Objects.requireNonNull(mapper);
@@ -105,57 +121,66 @@ public UidsCookie parseFromRequest(HttpRequestContext httpRequest) {
      */
     UidsCookie parseFromCookies(Map<String, String> cookies) {
         final Uids parsedUids = parseUids(cookies);
+        final boolean isOptedOut = isOptedOut(cookies);
 
-        final Boolean optout;
-        final Map<String, UidWithExpiry> uidsMap;
-
-        if (isOptedOut(cookies)) {
-            optout = true;
-            uidsMap = Collections.emptyMap();
-        } else {
-            optout = parsedUids != null ? parsedUids.getOptout() : null;
-            uidsMap = enrichAndSanitizeUids(parsedUids, cookies);
-        }
-
-        final Uids uids = Uids.builder().uids(uidsMap).optout(optout).build();
+        final Uids uids = Uids.builder()
+                .uids(isOptedOut ? Collections.emptyMap() : enrichAndSanitizeUids(parsedUids, cookies))
+                .optout(isOptedOut)
+                .build();
 
         return new UidsCookie(uids, mapper);
     }
 
     /**
      * Parses cookies {@link Map} and composes {@link Uids} model.
      */
-    public Uids parseUids(Map<String, String> cookies) {
-        if (cookies.containsKey(COOKIE_NAME)) {
-            final String cookieValue = cookies.get(COOKIE_NAME);
+    private Uids parseUids(Map<String, String> cookies) {
+        final Map<String, UidWithExpiry> uids = new HashMap<>();
+
+        for (Map.Entry<String, String> cookie : cookies.entrySet()) {
+            final String cookieKey = cookie.getKey();
+            if (!cookieKey.startsWith(COOKIE_NAME)) {
+                continue;
+            }
+
             try {
-                return mapper.decodeValue(Buffer.buffer(Base64.getUrlDecoder().decode(cookieValue)), Uids.class);
+                final Uids parsedUids = mapper.decodeValue(
+                        Buffer.buffer(Base64.getUrlDecoder().decode(cookie.getValue())), Uids.class);
+                if (parsedUids != null && parsedUids.getUids() != null) {
+                    parsedUids.getUids().forEach((key, value) -> uids.merge(key, value, (newValue, oldValue) ->
+                            newValue.getExpires().compareTo(oldValue.getExpires()) > 0 ? newValue : oldValue));
+                }
             } catch (IllegalArgumentException | DecodeException e) {
-                logger.debug("Could not decode or parse {} cookie value {}", e, COOKIE_NAME, cookieValue);
+                logger.debug("Could not decode or parse {} cookie value {}", e, COOKIE_NAME, cookie.getValue());
             }
         }
-        return null;
+
+        return Uids.builder().uids(uids).build();
     }
 
     /**
      * Creates a {@link Cookie} with 'uids' as a name and encoded JSON string representing supplied {@link UidsCookie}
      * as a value.
      */
-    public Cookie toCookie(UidsCookie uidsCookie) {
-        return makeCookie(uidsCookie);
+    public Cookie aliveCookie(String cookieName, UidsCookie uidsCookie) {
+        final String value = Base64.getUrlEncoder().encodeToString(uidsCookie.toJson().getBytes());
+        return makeCookie(cookieName, value, ttlSeconds);
+    }
+
+    public Cookie aliveCookie(UidsCookie uidsCookie) {
+        return aliveCookie(COOKIE_NAME, uidsCookie);
     }
 
-    private int cookieBytesLength(UidsCookie uidsCookie) {
-        return makeCookie(uidsCookie).encode().getBytes().length;
+    public Cookie expiredCookie(String cookieName) {
+        return makeCookie(cookieName, StringUtils.EMPTY, 0);
     }
 
-    private Cookie makeCookie(UidsCookie uidsCookie) {
-        return Cookie
-                .cookie(COOKIE_NAME, Base64.getUrlEncoder().encodeToString(uidsCookie.toJson().getBytes()))
+    private Cookie makeCookie(String cookieName, String value, long maxAge) {
+        return Cookie.cookie(cookieName, value)
                 .setPath("/")
                 .setSameSite(CookieSameSite.NONE)
                 .setSecure(true)
-                .setMaxAge(ttlSeconds)
+                .setMaxAge(maxAge)
                 .setDomain(hostCookieDomain);
     }
 
@@ -221,20 +246,18 @@ private static boolean facebookSentinelOrEmpty(Map.Entry<String, UidWithExpiry>
 
     /***
      * Removes expired {@link Uids}, updates {@link UidsCookie} with new uid for family name according to priority
-     * and trims it to the limit
      */
-    public UidsCookieUpdateResult updateUidsCookie(UidsCookie uidsCookie, String familyName, String uid) {
-        final UidsCookie initialCookie = trimToLimit(removeExpiredUids(uidsCookie)); // if already exceeded limit
-
-        if (StringUtils.isBlank(uid)) {
-            return UidsCookieUpdateResult.unaltered(initialCookie.deleteUid(familyName));
-        } else if (UidsCookie.isFacebookSentinel(familyName, uid)) {
-            // At the moment, Facebook calls /setuid with a UID of 0 if the user isn't logged into Facebook.
-            // They shouldn't be sending us a sentinel value... but since they are, we're refusing to save that ID.
-            return UidsCookieUpdateResult.unaltered(initialCookie);
+    public UpdateResult<UidsCookie> updateUidsCookie(UidsCookie uidsCookie, String familyName, String uid) {
+        final UidsCookie initialCookie = removeExpiredUids(uidsCookie);
+
+        // At the moment, Facebook calls /setuid with a UID of 0 if the user isn't logged into Facebook.
+        // They shouldn't be sending us a sentinel value... but since they are, we're refusing to save that ID.
+        if (StringUtils.isBlank(uid) || UidsCookie.isFacebookSentinel(familyName, uid)) {
+            return UpdateResult.unaltered(initialCookie);
         }
 
-        return updateUidsCookieByPriority(initialCookie, familyName, uid);
+        final UidsCookie updatedCookie = initialCookie.updateUid(familyName, uid);
+        return UpdateResult.updated(updatedCookie);
     }
 
     private static UidsCookie removeExpiredUids(UidsCookie uidsCookie) {
@@ -250,47 +273,58 @@ private static UidsCookie removeExpiredUids(UidsCookie uidsCookie) {
         return updatedCookie;
     }
 
-    private UidsCookieUpdateResult updateUidsCookieByPriority(UidsCookie uidsCookie, String familyName, String uid) {
-        final UidsCookie updatedCookie = uidsCookie.updateUid(familyName, uid);
-        if (!cookieExceededMaxLength(updatedCookie)) {
-            return UidsCookieUpdateResult.updated(updatedCookie);
-        }
+    public List<Cookie> splitUidsIntoCookies(UidsCookie uidsCookie) {
+        final Uids cookieUids = uidsCookie.getCookieUids();
+        final Map<String, UidWithExpiry> uids = cookieUids.getUids();
+        final boolean hasOptout = !uidsCookie.allowsSync();
 
-        if (!prioritizedCoopSyncProvider.hasPrioritizedBidders()
-                || prioritizedCoopSyncProvider.isPrioritizedFamily(familyName)) {
-            return UidsCookieUpdateResult.updated(trimToLimit(updatedCookie));
-        } else {
-            metrics.updateUserSyncSizeBlockedMetric(familyName);
-            return UidsCookieUpdateResult.unaltered(uidsCookie);
-        }
-    }
+        final Iterator<String> cookieFamilies = cookieFamilyNamesByDescPriorityAndExpiration(uidsCookie);
+        final List<Cookie> splitCookies = new ArrayList<>();
 
-    private boolean cookieExceededMaxLength(UidsCookie uidsCookie) {
-        return maxCookieSizeBytes > 0 && cookieBytesLength(uidsCookie) > maxCookieSizeBytes;
-    }
+        final int cookieSchemaSize = UidsCookieSize.schemaSize(makeCookie(COOKIE_NAME, StringUtils.EMPTY, ttlSeconds));
+        String nextCookieFamily = null;
+        for (int i = 0; i < numberOfUidCookies; i++) {
+            final int digits = i < 10 ? Integer.signum(i) : 2;
+            final UidsCookieSize uidsCookieSize = new UidsCookieSize(cookieSchemaSize + digits, maxCookieSizeBytes);
 
-    private UidsCookie trimToLimit(UidsCookie uidsCookie) {
-        if (!cookieExceededMaxLength(uidsCookie)) {
-            return uidsCookie;
-        }
+            final Map<String, UidWithExpiry> tempUids = new HashMap<>();
+            while (nextCookieFamily != null || cookieFamilies.hasNext()) {
+                nextCookieFamily = nextCookieFamily == null ? cookieFamilies.next() : nextCookieFamily;
+                final UidWithExpiry uidWithExpiry = uids.get(nextCookieFamily);
 
-        UidsCookie trimmedUids = uidsCookie;
-        final Iterator<String> familyToRemoveIterator = cookieFamilyNamesByAscendingPriority(uidsCookie);
+                uidsCookieSize.addUid(nextCookieFamily, uidWithExpiry.getUid());
+                if (!uidsCookieSize.isValid()) {
+                    break;
+                }
+
+                tempUids.put(nextCookieFamily, uidWithExpiry);
+                nextCookieFamily = null;
+            }
+
+            final String uidsName = i == 0 ? COOKIE_NAME : COOKIE_NAME_FORMAT.formatted(i + 1);
+
+            if (tempUids.isEmpty()) {
+                splitCookies.add(expiredCookie(uidsName));
+            } else {
+                splitCookies.add(aliveCookie(
+                        uidsName,
+                        new UidsCookie(Uids.builder().uids(tempUids).optout(hasOptout).build(), mapper)));
+            }
+        }
 
-        while (familyToRemoveIterator.hasNext() && cookieExceededMaxLength(trimmedUids)) {
-            final String familyToRemove = familyToRemoveIterator.next();
-            metrics.updateUserSyncSizedOutMetric(familyToRemove);
-            trimmedUids = trimmedUids.deleteUid(familyToRemove);
+        if (nextCookieFamily != null) {
+            updateSyncSizeMetrics(nextCookieFamily);
         }
 
-        return trimmedUids;
+        cookieFamilies.forEachRemaining(this::updateSyncSizeMetrics);
+
+        return splitCookies;
     }
 
-    private Iterator<String> cookieFamilyNamesByAscendingPriority(UidsCookie uidsCookie) {
+    private Iterator<String> cookieFamilyNamesByDescPriorityAndExpiration(UidsCookie uidsCookie) {
         return uidsCookie.getCookieUids().getUids().entrySet().stream()
                 .sorted(this::compareCookieFamilyNames)
                 .map(Map.Entry::getKey)
-                .toList()
                 .iterator();
     }
 
@@ -303,9 +337,17 @@ private int compareCookieFamilyNames(Map.Entry<String, UidWithExpiry> left,
         if ((leftPrioritized && rightPrioritized) || (!leftPrioritized && !rightPrioritized)) {
             return left.getValue().getExpires().compareTo(right.getValue().getExpires());
         } else if (leftPrioritized) {
-            return 1;
-        } else { // right is prioritized
             return -1;
+        } else { // right is prioritized
+            return 1;
+        }
+    }
+
+    private void updateSyncSizeMetrics(String nextCookieFamily) {
+        if (prioritizedCoopSyncProvider.isPrioritizedFamily(nextCookieFamily)) {
+            metrics.updateUserSyncSizedOutMetric(nextCookieFamily);
+        } else {
+            metrics.updateUserSyncSizeBlockedMetric(nextCookieFamily);
         }
     }
 

src/main/java/org/prebid/server/cookie/UidsCookieSize.java

@@ -0,0 +1,73 @@
+package org.prebid.server.cookie;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import io.vertx.core.http.Cookie;
+import org.apache.commons.lang3.StringUtils;
+import org.prebid.server.json.ObjectMapperProvider;
+
+import java.time.Instant;
+import java.time.ZoneId;
+import java.time.ZonedDateTime;
+
+public class UidsCookieSize {
+
+    // {"tempUIDs":{},"optout":false}
+    private static final int TEMP_UIDS_BASE64_BYTES = "eyJ0ZW1wVUlEcyI6e30sIm9wdG91dCI6ZmFsc2V9".length();
+    private static final int UID_TEMPLATE_BYTES;
+
+    static {
+        try {
+            UID_TEMPLATE_BYTES = "\"\":{\"uid\":\"\",\"expires\":\"%s\"},"
+                    .formatted(ObjectMapperProvider.mapper().writeValueAsString(
+                            ZonedDateTime.ofInstant(Instant.ofEpochSecond(0, 1), ZoneId.of("UTC"))))
+                    .length();
+        } catch (JsonProcessingException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private final int cookieSchemaSize;
+    private final int maxSize;
+    private int encodedUidsSize;
+
+    public UidsCookieSize(int cookieSchemaSize, int maxSize) {
+        this.cookieSchemaSize = cookieSchemaSize;
+        this.maxSize = maxSize;
+
+        encodedUidsSize = 0;
+    }
+
+    public static int schemaSize(Cookie cookieSchema) {
+        return cookieSchema.setValue(StringUtils.EMPTY).encode().length();
+    }
+
+    public boolean isValid() {
+        return maxSize <= 0 || totalSize() <= maxSize;
+    }
+
+    public int totalSize() {
+        return cookieSchemaSize
+                + TEMP_UIDS_BASE64_BYTES
+                + Base64Size.base64Size(encodedUidsSize);
+    }
+
+    public void addUid(String cookieFamily, String uid) {
+        final int uidSize = UID_TEMPLATE_BYTES + cookieFamily.length() + uid.length();
+        encodedUidsSize = Base64Size.encodeSize(Base64Size.decodeSize(encodedUidsSize) + uidSize);
+    }
+
+    private static class Base64Size {
+
+        public static int encodeSize(int size) {
+            return size / 3 * 4 + size % 3;
+        }
+
+        public static int decodeSize(int encodedSize) {
+            return encodedSize / 4 * 3 + encodedSize % 4;
+        }
+
+        private static int base64Size(int encodedSize) {
+            return (encodedSize & -4) + 4 * Integer.signum(encodedSize % 4);
+        }
+    }
+}