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

@@ -1,16 +1,27 @@
 package org.unicitylabs.sdk.api;
 
 import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import org.unicitylabs.sdk.hash.DataHash;
 import org.unicitylabs.sdk.jsonrpc.JsonRpcHttpTransport;
 
+import static com.google.common.net.HttpHeaders.AUTHORIZATION;
+
 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(
@@ -20,7 +31,11 @@ public CompletableFuture<SubmitCommitmentResponse> submitCommitment(
 
     SubmitCommitmentRequest request = new SubmitCommitmentRequest(requestId, transactionHash,
         authenticator, false);
-    return this.transport.request("submit_commitment", request, SubmitCommitmentResponse.class);
+    Map<String, List<String>> headers = new LinkedHashMap<>();
+    if (apiKey != null) {
+        headers.put(AUTHORIZATION, Collections.singletonList("Bearer " + apiKey));
+    }
+    return this.transport.request("submit_commitment", request, SubmitCommitmentResponse.class, headers);
   }
 
   public CompletableFuture<InclusionProofResponse> getInclusionProof(RequestId requestId) {

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

@@ -2,6 +2,8 @@
 package org.unicitylabs.sdk.jsonrpc;
 
 import java.io.IOException;
+import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import okhttp3.Call;
 import okhttp3.Callback;
@@ -18,9 +20,9 @@
  */
 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 final String url;
+    private final String url;
   private final OkHttpClient httpClient;
 
   /**
@@ -35,17 +37,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, Map.of());
+  }
+
+  /**
+   * Send a JSON-RPC request with optional API key.
+   */
+  public <T> CompletableFuture<T> request(String method, Object params, Class<T> resultType, Map<String, List<String>> headers) {
     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();
+          );
+
+      headers.forEach((header, values) ->
+        values.forEach(value ->
+          requestBuilder.addHeader(header, value)));
+
+      Request request = requestBuilder.build();
 
       this.httpClient.newCall(request).enqueue(new Callback() {
         @Override

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

@@ -0,0 +1,157 @@
+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 = 0;
+    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) {
+                try {
+                    return new MockResponse()
+                            .setResponseCode(429)
+                            .setHeader("Retry-After", String.valueOf(rateLimitRetryAfter))
+                            .setBody("Too Many Requests");
+                } finally {
+                    // Reset for next request
+                    simulateRateLimit = false;
+                    rateLimitRetryAfter = 0;
+                }
+            }
+
+            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);
+    }
+}