Add support for aggregator API keys

build.gradle.kts

@@ -45,12 +45,14 @@ dependencies {
     // Testing
     testImplementation(platform("org.junit:junit-bom:5.10.2"))
     testImplementation("org.junit.jupiter:junit-jupiter")
+    testRuntimeOnly("org.junit.platform:junit-platform-launcher")
     testImplementation("org.testcontainers:testcontainers:1.19.8")
     testImplementation("org.testcontainers:junit-jupiter:1.19.8")
     testImplementation("org.testcontainers:mongodb:1.19.8")
     testImplementation("org.awaitility:awaitility:4.2.0")
     testImplementation("org.slf4j:slf4j-simple:2.0.13")
     testImplementation("com.google.guava:guava:33.0.0-jre")
+    testImplementation("com.squareup.okhttp3:mockwebserver:4.12.0")
 
     // ✅ Cucumber for BDD
     testImplementation("io.cucumber:cucumber-java:7.27.2")

src/main/java/org/unicitylabs/sdk/api/AggregatorClient.java

@@ -9,9 +9,15 @@
 public class AggregatorClient implements IAggregatorClient {
 
   private final JsonRpcHttpTransport transport;
+  private final String apiKey;
 
   public AggregatorClient(String url) {
+    this(url, null);
+  }
+
+  public AggregatorClient(String url, String apiKey) {
     this.transport = new JsonRpcHttpTransport(url);
+    this.apiKey = apiKey;
   }
 
   public CompletableFuture<SubmitCommitmentResponse> submitCommitment(
@@ -21,7 +27,7 @@ public CompletableFuture<SubmitCommitmentResponse> submitCommitment(
 
     SubmitCommitmentRequest request = new SubmitCommitmentRequest(requestId, transactionHash,
         authenticator, false);
-    return this.transport.request("submit_commitment", request, SubmitCommitmentResponse.class);
+    return this.transport.request("submit_commitment", request, SubmitCommitmentResponse.class, this.apiKey);
   }
 
   public CompletableFuture<InclusionProof> getInclusionProof(RequestId requestId) {

src/main/java/org/unicitylabs/sdk/jsonrpc/JsonRpcHttpTransport.java

@@ -13,14 +13,18 @@
 import okhttp3.Response;
 import okhttp3.ResponseBody;
 
+import static java.net.HttpURLConnection.HTTP_UNAUTHORIZED;
+
 /**
  * JSON-RPC HTTP service.
  */
 public class JsonRpcHttpTransport {
 
-  private static final MediaType MEDIA_TYPE_JSON = MediaType.get("application/json; charset=utf-8");
+    private static final MediaType MEDIA_TYPE_JSON = MediaType.get("application/json; charset=utf-8");
+    private static final int HTTP_TOO_MANY_REQUESTS = 429;
+    private static final String HTTP_RETRY_AFTER = "Retry-After";
 
-  private final String url;
+    private final String url;
   private final OkHttpClient httpClient;
 
   /**
@@ -35,17 +39,29 @@ public JsonRpcHttpTransport(String url) {
    * Send a JSON-RPC request.
    */
   public <T> CompletableFuture<T> request(String method, Object params, Class<T> resultType) {
+    return request(method, params, resultType, null);
+  }
+
+  /**
+   * Send a JSON-RPC request with optional API key.
+   */
+  public <T> CompletableFuture<T> request(String method, Object params, Class<T> resultType, String apiKey) {
     CompletableFuture<T> future = new CompletableFuture<>();
 
     try {
-      Request request = new Request.Builder()
+      Request.Builder requestBuilder = new Request.Builder()
           .url(this.url)
           .post(
               RequestBody.create(
                   UnicityObjectMapper.JSON.writeValueAsString(new JsonRpcRequest(method, params)),
                   JsonRpcHttpTransport.MEDIA_TYPE_JSON)
-          )
-          .build();
+          );
+
+      if (apiKey != null) {
+        requestBuilder.header("Authorization", "Bearer " + apiKey);
+      }
+
+      Request request = requestBuilder.build();
 
       this.httpClient.newCall(request).enqueue(new Callback() {
         @Override
@@ -58,8 +74,21 @@ public void onResponse(Call call, Response response) throws IOException {
           try (ResponseBody body = response.body()) {
             if (!response.isSuccessful()) {
               String error = body != null ? body.string() : "";
-              future.completeExceptionally(new JsonRpcNetworkError(response.code(), error));
-              return;
+
+              if (response.code() == HTTP_UNAUTHORIZED) {
+                future.completeExceptionally(new UnauthorizedException(
+                    "Unauthorized: Invalid or missing API key"));
+                return;
+              } else if (response.code() == HTTP_TOO_MANY_REQUESTS) {
+                int retryAfterSeconds = extractRetryAfterSeconds(response);
+                future.completeExceptionally(new RateLimitExceededException(
+                    "Rate limit exceeded. Please retry after " + retryAfterSeconds + " seconds",
+                    retryAfterSeconds));
+                return;
+              } else {
+                  future.completeExceptionally(new JsonRpcNetworkError(response.code(), error));
+                  return;
+              }
             }
 
             JsonRpcResponse<T> data = UnicityObjectMapper.JSON.readValue(
@@ -85,4 +114,16 @@ public void onResponse(Call call, Response response) throws IOException {
 
     return future;
   }
+
+  private int extractRetryAfterSeconds(Response response) {
+    String retryAfterHeader = response.header(HTTP_RETRY_AFTER);
+    if (retryAfterHeader != null) {
+      try {
+        return Integer.parseInt(retryAfterHeader);
+      } catch (NumberFormatException ignored) {
+      }
+    }
+    // Default to 60 seconds if the HTTP header is missing, e.g. if the response is coming from a different component that is not using this header.
+    return 60;
+  }
 }

src/main/java/org/unicitylabs/sdk/jsonrpc/RateLimitExceededException.java

@@ -0,0 +1,20 @@
+package org.unicitylabs.sdk.jsonrpc;
+
+public class RateLimitExceededException extends RuntimeException {
+
+    private final int retryAfterSeconds;
+
+    public RateLimitExceededException(String message, int retryAfterSeconds) {
+        super(message);
+        this.retryAfterSeconds = retryAfterSeconds;
+    }
+
+    public RateLimitExceededException(String message, int retryAfterSeconds, Throwable cause) {
+        super(message, cause);
+        this.retryAfterSeconds = retryAfterSeconds;
+    }
+
+    public int getRetryAfterSeconds() {
+        return retryAfterSeconds;
+    }
+}

src/main/java/org/unicitylabs/sdk/jsonrpc/UnauthorizedException.java

@@ -0,0 +1,16 @@
+package org.unicitylabs.sdk.jsonrpc;
+
+/**
+ * Exception thrown when an API request is unauthorized (HTTP 401).
+ * This typically occurs when an API key is missing or invalid.
+ */
+public class UnauthorizedException extends RuntimeException {
+
+    public UnauthorizedException(String message) {
+        super(message);
+    }
+
+    public UnauthorizedException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}

src/test/java/org/unicitylabs/sdk/MockAggregatorServer.java

@@ -0,0 +1,152 @@
+package org.unicitylabs.sdk;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import okhttp3.mockwebserver.Dispatcher;
+import okhttp3.mockwebserver.MockResponse;
+import okhttp3.mockwebserver.MockWebServer;
+import okhttp3.mockwebserver.RecordedRequest;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.JsonNode;
+import org.jetbrains.annotations.Nullable;
+
+import java.io.IOException;
+import java.util.Set;
+import java.util.HashSet;
+import java.util.UUID;
+
+public class MockAggregatorServer {
+
+    private final MockWebServer server;
+    private final ObjectMapper objectMapper;
+    private final Set<String> protectedMethods;
+    private volatile boolean simulateRateLimit = false;
+    private volatile int rateLimitRetryAfter = 60;
+    private volatile String expectedApiKey = null;
+
+    public MockAggregatorServer() {
+        this.server = new MockWebServer();
+        this.objectMapper = new ObjectMapper();
+        this.protectedMethods = new HashSet<>();
+        this.protectedMethods.add("submit_commitment");
+
+        server.setDispatcher(new Dispatcher() {
+            @Override
+            public MockResponse dispatch(RecordedRequest request) {
+                return handleRequest(request);
+            }
+        });
+    }
+
+    public void start() throws IOException {
+        server.start();
+    }
+
+    public void shutdown() throws IOException {
+        server.shutdown();
+    }
+
+    public String getUrl() {
+        return server.url("/").toString();
+    }
+
+    public RecordedRequest takeRequest() throws InterruptedException {
+        return server.takeRequest();
+    }
+
+    public void simulateRateLimitForNextRequest(int retryAfterSeconds) {
+        this.simulateRateLimit = true;
+        this.rateLimitRetryAfter = retryAfterSeconds;
+    }
+
+    public void setExpectedApiKey(String apiKey) {
+        this.expectedApiKey = apiKey;
+    }
+
+    private MockResponse handleRequest(RecordedRequest request) {
+        try {
+            if (simulateRateLimit) {
+                simulateRateLimit = false; // Reset for next request
+                return new MockResponse()
+                    .setResponseCode(429)
+                    .setHeader("Retry-After", String.valueOf(rateLimitRetryAfter))
+                    .setBody("Too Many Requests");
+            }
+
+            String method = extractJsonRpcMethod(request);
+
+            if (protectedMethods.contains(method) && expectedApiKey != null && !hasValidApiKey(request)) {
+                return new MockResponse()
+                        .setResponseCode(401)
+                        .setHeader("WWW-Authenticate", "Bearer")
+                        .setBody("Unauthorized");
+            }
+
+            return generateSuccessResponse(method);
+
+        } catch (Exception e) {
+            return new MockResponse()
+                .setResponseCode(400)
+                .setBody("Bad Request");
+        }
+    }
+
+    private boolean hasValidApiKey(RecordedRequest request) {
+        String authHeader = request.getHeader("Authorization");
+        if (authHeader != null && authHeader.startsWith("Bearer ")) {
+            String providedKey = authHeader.substring(7);
+            return expectedApiKey.equals(providedKey);
+        }
+        return false;
+    }
+
+    private @Nullable String extractJsonRpcMethod(RecordedRequest request) throws JsonProcessingException {
+        if (!"POST".equals(request.getMethod())) {
+           return null;
+        }
+        JsonNode jsonRequest = objectMapper.readTree(request.getBody().readUtf8());
+        return jsonRequest.has("method") ? jsonRequest.get("method").asText() : null;
+    }
+
+    private MockResponse generateSuccessResponse(String method) {
+        String responseBody;
+        String id = UUID.randomUUID().toString();
+
+        switch (method != null ? method : "") {
+            case "submit_commitment":
+                responseBody = String.format(
+                    "{\n" +
+                    "    \"jsonrpc\": \"2.0\",\n" +
+                    "    \"result\": {\n" +
+                    "        \"status\": \"SUCCESS\"\n" +
+                    "    },\n" +
+                    "    \"id\": \"%s\"\n" +
+                    "}", id);
+                break;
+
+            case "get_block_height":
+                responseBody = String.format(
+                    "{\n" +
+                    "    \"jsonrpc\": \"2.0\",\n" +
+                    "    \"result\": {\n" +
+                    "        \"blockNumber\": \"67890\"\n" +
+                    "    },\n" +
+                    "    \"id\": \"%s\"\n" +
+                    "}", id);
+                break;
+
+            default:
+                responseBody = String.format(
+                    "{\n" +
+                    "    \"jsonrpc\": \"2.0\",\n" +
+                    "    \"result\": \"OK\",\n" +
+                    "    \"id\": \"%s\"\n" +
+                    "}", id);
+                break;
+        }
+
+        return new MockResponse()
+            .setResponseCode(200)
+            .setHeader("Content-Type", "application/json")
+            .setBody(responseBody);
+    }
+}