enhancement: Ability to set headers for requests

src/main/java/dev/cerbos/sdk/CerbosBlockingAdminClient.java

@@ -6,25 +6,57 @@
 import dev.cerbos.api.v1.schema.SchemaOuterClass;
 import dev.cerbos.api.v1.svc.CerbosAdminServiceGrpc;
 import io.grpc.Channel;
+import io.grpc.Metadata;
 import io.grpc.StatusRuntimeException;
+import io.grpc.stub.MetadataUtils;
 
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
 public class CerbosBlockingAdminClient {
     private final CerbosAdminServiceGrpc.CerbosAdminServiceBlockingStub stub;
     private final long timeoutMillis;
+    private final Optional<Metadata> headerMetadata;
 
     CerbosBlockingAdminClient(Channel channel, long timeoutMillis, AdminApiCredentials adminCredentials) {
         CerbosAdminServiceGrpc.CerbosAdminServiceBlockingStub c = CerbosAdminServiceGrpc.newBlockingStub(channel);
         this.stub = c.withCallCredentials(adminCredentials);
         this.timeoutMillis = timeoutMillis;
+        this.headerMetadata = Optional.empty();
+    }
+
+    CerbosBlockingAdminClient(CerbosAdminServiceGrpc.CerbosAdminServiceBlockingStub stub, long timeoutMillis, Optional<Metadata> headerMetadata) {
+        this.stub = stub;
+        this.timeoutMillis = timeoutMillis;
+        this.headerMetadata = headerMetadata;
     }
 
     private CerbosAdminServiceGrpc.CerbosAdminServiceBlockingStub withClient() {
-        return stub.withDeadlineAfter(timeoutMillis, TimeUnit.MILLISECONDS);
+        CerbosAdminServiceGrpc.CerbosAdminServiceBlockingStub s = headerMetadata.map(md -> stub.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(md))).orElse(stub);
+        return s.withDeadlineAfter(timeoutMillis, TimeUnit.MILLISECONDS);
+    }
+
+    /**
+     * Add header metadata to Cerbos requests
+     * @param md {@link Metadata}
+     * @return CerbosBlockingAdminClient configured to attach headers to each request.
+     */
+    public CerbosBlockingAdminClient withHeaders(Metadata md) {
+        return new CerbosBlockingAdminClient(stub, timeoutMillis, Optional.ofNullable(md));
+    }
+
+    /**
+     * Attach the given headers to the Cerbos request.
+     * @param headers Map of key-value pairs
+     * @return new CerbosBlockingAdminClient configured to attach the given headers to the requests.
+     */
+    public CerbosBlockingAdminClient withHeaders(Map<String, String> headers) {
+        Metadata md = new Metadata();
+        headers.forEach((k, v) -> md.put(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER), v));
+        return withHeaders(md);
     }
 
     /**

src/main/java/dev/cerbos/sdk/CerbosBlockingClient.java

@@ -12,9 +12,12 @@
 import dev.cerbos.sdk.builders.Principal;
 import dev.cerbos.sdk.builders.Resource;
 import io.grpc.Channel;
+import io.grpc.Metadata;
 import io.grpc.StatusRuntimeException;
+import io.grpc.stub.MetadataUtils;
 
 import java.util.Arrays;
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 
@@ -26,6 +29,8 @@ public class CerbosBlockingClient {
   private final CerbosServiceGrpc.CerbosServiceBlockingStub cerbosStub;
   private final long timeoutMillis;
   private final Optional<AuxData> auxData;
+  private final Optional<Metadata> headerMetadata;
+
 
   CerbosBlockingClient(
       Channel channel, long timeoutMillis, PlaygroundInstanceCredentials playgroundCredentials) {
@@ -37,17 +42,20 @@ public class CerbosBlockingClient {
     }
     this.timeoutMillis = timeoutMillis;
     this.auxData = Optional.empty();
+    this.headerMetadata = Optional.empty();
   }
 
   CerbosBlockingClient(
-      CerbosServiceGrpc.CerbosServiceBlockingStub cerbosStub, long timeoutMillis, AuxData auxData) {
+      CerbosServiceGrpc.CerbosServiceBlockingStub cerbosStub, long timeoutMillis, Optional<AuxData> auxData, Optional<Metadata> headerMetadata) {
     this.cerbosStub = cerbosStub;
     this.timeoutMillis = timeoutMillis;
-    this.auxData = Optional.ofNullable(auxData);
+    this.auxData = auxData;
+    this.headerMetadata = headerMetadata;
   }
 
   private CerbosServiceGrpc.CerbosServiceBlockingStub withClient() {
-    return cerbosStub.withDeadlineAfter(timeoutMillis, TimeUnit.MILLISECONDS);
+    CerbosServiceGrpc.CerbosServiceBlockingStub stub = this.headerMetadata.map(md -> cerbosStub.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(md))).orElse(cerbosStub);
+    return stub.withDeadlineAfter(timeoutMillis, TimeUnit.MILLISECONDS);
   }
 
   /**
@@ -57,7 +65,27 @@ private CerbosServiceGrpc.CerbosServiceBlockingStub withClient() {
    * @return new CerbosBlockingClient configured to attach the auxiliary data to requests.
    */
   public CerbosBlockingClient with(AuxData auxData) {
-    return new CerbosBlockingClient(cerbosStub, timeoutMillis, auxData);
+    return new CerbosBlockingClient(cerbosStub, timeoutMillis, Optional.ofNullable(auxData), headerMetadata);
+  }
+
+  /**
+   * Attach the given header metadata to the Cerbos request
+   * @param md {@link Metadata}
+   * @return new CerbosBlockingClient configured to attach given headers to the requests.
+   */
+  public CerbosBlockingClient withHeaders(Metadata md) {
+    return new CerbosBlockingClient(cerbosStub, timeoutMillis, auxData, Optional.ofNullable(md));
+  }
+
+  /**
+   * Attach the given headers to the Cerbos request.
+   * @param headers Map of key-value pairs
+   * @return new CerbosBlockingClient configured to attach the given headers to the requests.
+   */
+  public CerbosBlockingClient withHeaders(Map<String, String> headers) {
+    Metadata md = new Metadata();
+    headers.forEach((k, v) -> md.put(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER), v));
+    return withHeaders(md);
   }
 
   /**

src/test/java/dev/cerbos/sdk/CerbosBlockingAdminClientTest.java

@@ -20,6 +20,7 @@
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 
 @Testcontainers
@@ -39,7 +40,7 @@ class CerbosBlockingAdminClientTest extends CerbosClientTests {
     @BeforeAll
     public void initClient() throws CerbosClientBuilder.InvalidClientConfigurationException, URISyntaxException {
         String target = cerbosContainer.getTarget();
-        this.adminClient = new CerbosClientBuilder(target).withPlaintext().buildBlockingAdminClient("cerbos", "cerbosAdmin");
+        this.adminClient = new CerbosClientBuilder(target).withPlaintext().buildBlockingAdminClient("cerbos", "cerbosAdmin").withHeaders(Map.of("wibble", "wobble"));
         this.client = new CerbosClientBuilder(target).withPlaintext().buildBlockingClient();
         loadSchemas();
         loadPolicies();

src/test/java/dev/cerbos/sdk/CerbosBlockingClientTest.java

@@ -14,6 +14,8 @@
 import org.testcontainers.junit.jupiter.Container;
 import org.testcontainers.junit.jupiter.Testcontainers;
 
+import java.util.Map;
+
 @Testcontainers
 @TestInstance(TestInstance.Lifecycle.PER_CLASS)
 class CerbosBlockingClientTest extends CerbosClientTests {
@@ -29,6 +31,6 @@ class CerbosBlockingClientTest extends CerbosClientTests {
     @BeforeAll
     public void initClient() throws CerbosClientBuilder.InvalidClientConfigurationException {
         String target = cerbosContainer.getTarget();
-        this.client = new CerbosClientBuilder(target).withPlaintext().buildBlockingClient();
+        this.client = new CerbosClientBuilder(target).withPlaintext().buildBlockingClient().withHeaders(Map.of("wibble", "wobble"));
     }
 }