Add logging to get information about refresh token version (#602)

* Add logging to get information about refresh token version * Fix unit tests * Record metrics using labels to separate versions * Remove unnecessary import * Remove global registry metrics in TokenEncodingTest * Remove unnecessary variables
IABTechLab · Jun 11, 2024 · 6017673 · 6017673
1 parent 3c2201e
commit 6017673
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 7 deletions.
diff --git a/src/main/java/com/uid2/operator/service/EncryptedTokenEncoder.java b/src/main/java/com/uid2/operator/service/EncryptedTokenEncoder.java
@@ -9,12 +9,13 @@
 import com.uid2.shared.model.KeysetKey;
 import com.uid2.shared.model.TokenVersion;
 import io.vertx.core.buffer.Buffer;
+import io.micrometer.core.instrument.Counter;
+import io.micrometer.core.instrument.Metrics;
 
 import java.time.Instant;
 import java.util.Base64;
 
 public class EncryptedTokenEncoder implements ITokenEncoder {
-
     private final KeyManager keyManager;
 
     public EncryptedTokenEncoder(KeyManager keyManager) {
@@ -259,13 +260,23 @@ public AdvertisingToken decodeAdvertisingTokenV3orV4(Buffer b, byte[] bytes, Tok
         );
     }
 
+    private void recordRefreshTokenVersionCount(String siteId, TokenVersion tokenVersion) {
+        Counter.builder("uid2_refresh_token_served_count")
+                .description(String.format("Counter for the amount of refresh token %s served", tokenVersion.toString().toLowerCase()))
+                .tags("site_id", String.valueOf(siteId))
+                .tags("refresh_token_version", tokenVersion.toString().toLowerCase())
+                .register(Metrics.globalRegistry).increment();
+    }
+
     public byte[] encode(RefreshToken t, Instant asOf) {
         final KeysetKey serviceKey = this.keyManager.getRefreshKey(asOf);
 
         switch (t.version) {
             case V2:
+                recordRefreshTokenVersionCount(String.valueOf(t.publisherIdentity.siteId), TokenVersion.V2);
                 return encodeV2(t, serviceKey);
             case V3:
+                recordRefreshTokenVersionCount(String.valueOf(t.publisherIdentity.siteId), TokenVersion.V3);
                 return encodeV3(t, serviceKey);
             default:
                 throw new ClientInputValidationException("RefreshToken version " + t.version + " not supported");

diff --git a/src/main/java/com/uid2/operator/service/UIDOperatorService.java b/src/main/java/com/uid2/operator/service/UIDOperatorService.java
@@ -2,7 +2,6 @@
 
 import com.uid2.operator.model.*;
 import com.uid2.operator.util.PrivacyBits;
-import com.uid2.operator.vertx.OperatorShutdownHandler;
 import com.uid2.shared.model.SaltEntry;
 import com.uid2.operator.store.IOptOutStore;
 import com.uid2.shared.store.ISaltProvider;

diff --git a/src/main/java/com/uid2/operator/vertx/UIDOperatorVerticle.java b/src/main/java/com/uid2/operator/vertx/UIDOperatorVerticle.java
@@ -114,7 +114,6 @@ public class UIDOperatorVerticle extends AbstractVerticle {
     public final static int MASTER_KEYSET_ID_FOR_SDKS = 9999999; //this is because SDKs have an issue where they assume keyset ids are always positive; that will be fixed.
     public final static long OPT_OUT_CHECK_CUTOFF_DATE = Instant.parse("2023-09-01T00:00:00.00Z").getEpochSecond();
     private final Handler<Boolean> saltRetrievalResponseHandler;
-
     private final int maxBidstreamLifetimeSeconds;
     private final int allowClockSkewSeconds;
     protected int maxSharingLifetimeSeconds;
@@ -1778,13 +1777,34 @@ private void recordOptOutStatusEndpointStats(RoutingContext rc, int inputCount,
         optOutDistSummary.record(optOutCount);
     }
 
+    public TokenVersion getRefreshTokenVersion(String s) {
+        if (s != null && !s.isEmpty()) {
+            final byte[] bytes = EncodingUtils.fromBase64(s);
+            final Buffer b = Buffer.buffer(bytes);
+            if (b.getByte(1) == TokenVersion.V3.rawVersion) {
+                return TokenVersion.V3;
+            } else if (b.getByte(0) == TokenVersion.V2.rawVersion) {
+                return TokenVersion.V2;
+            }
+        }
+        return null;
+    }
+
+    private void recordRefreshTokenVersionCount(String siteId, TokenVersion tokenVersion) {
+        Counter.builder("uid2_refresh_token_received_count")
+                .description(String.format("Counter for the amount of refresh token %s received", tokenVersion.toString().toLowerCase()))
+                .tags("site_id", siteId)
+                .tags("refresh_token_version", tokenVersion.toString().toLowerCase())
+                .register(Metrics.globalRegistry).increment();
+
+    }
+
     private RefreshResponse refreshIdentity(RoutingContext rc, String tokenStr) {
         final RefreshToken refreshToken;
         try {
             if (AuthMiddleware.isAuthenticated(rc)) {
                 rc.put(Const.RoutingContextData.SiteId, AuthMiddleware.getAuthClient(ClientKey.class, rc).getSiteId());
             }
-
             refreshToken = this.encoder.decodeRefreshToken(tokenStr);
         } catch (ClientInputValidationException cie) {
             return RefreshResponse.Invalid;
@@ -1795,6 +1815,7 @@ private RefreshResponse refreshIdentity(RoutingContext rc, String tokenStr) {
         if (!AuthMiddleware.isAuthenticated(rc)) {
             rc.put(Const.RoutingContextData.SiteId, refreshToken.publisherIdentity.siteId);
         }
+        recordRefreshTokenVersionCount(String.valueOf(rc.data().get(Const.RoutingContextData.SiteId)), this.getRefreshTokenVersion(tokenStr));
 
         return this.idService.refreshIdentity(refreshToken);
     }

diff --git a/src/test/java/com/uid2/operator/TokenEncodingTest.java b/src/test/java/com/uid2/operator/TokenEncodingTest.java
@@ -11,17 +11,16 @@
 import com.uid2.shared.store.reader.RotatingKeysetKeyStore;
 import com.uid2.shared.store.reader.RotatingKeysetProvider;
 import com.uid2.shared.store.scope.GlobalScope;
+import io.micrometer.core.instrument.Metrics;
 import io.vertx.core.buffer.Buffer;
 import io.vertx.core.json.JsonObject;
 import org.junit.Assert;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.EnumSource;
-import org.junit.jupiter.params.provider.ValueSource;
 
 import java.time.Instant;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.*;
 
 public class TokenEncodingTest {
 
@@ -80,6 +79,10 @@ public void testRefreshTokenEncoding(TokenVersion tokenVersion) {
         Buffer b = Buffer.buffer(encodedBytes);
         int keyId = b.getInt(tokenVersion == TokenVersion.V2 ? 25 : 2);
         assertEquals(Data.RefreshKeySiteId, keyManager.getSiteIdFromKeyId(keyId));
+
+        assertNotNull(Metrics.globalRegistry
+                .get("uid2_refresh_token_served_count")
+                .counter());
     }
 
     @ParameterizedTest

diff --git a/src/test/java/com/uid2/operator/UIDOperatorVerticleTest.java b/src/test/java/com/uid2/operator/UIDOperatorVerticleTest.java
@@ -1745,6 +1745,8 @@ void tokenRefreshExpiredTokenAuthenticated(String apiVersion, Vertx vertx, Vertx
 
             sendTokenRefresh(apiVersion, vertx, testContext, refreshToken, bodyJson.getString("refresh_response_key"), 400, refreshRespJson -> {
                 assertEquals("expired_token", refreshRespJson.getString("status"));
+                assertNotNull(Metrics.globalRegistry
+                        .get("uid2_refresh_token_received_count").counter());
                 testContext.completeNow();
             });
         });
@@ -1763,6 +1765,8 @@ void tokenRefreshExpiredTokenUnauthenticated(String apiVersion, Vertx vertx, Ver
 
             sendTokenRefresh(apiVersion, vertx, testContext, refreshToken, "", 400, refreshRespJson -> {
                 assertEquals("error", refreshRespJson.getString("status"));
+                assertNotNull(Metrics.globalRegistry
+                        .get("uid2_refresh_token_received_count").counter());
                 testContext.completeNow();
             });
         });