From e89e5821ff6799c6e8fdd39aafe1653baf9dac5c Mon Sep 17 00:00:00 2001
From: jakubmatyszewski <jakubmatyszewski1@gmail.com>
Date: Fri, 15 Mar 2024 18:01:23 +0100
Subject: [PATCH] Add authenticationResponse context to OpaInput

---
 README.md                                            |  7 ++++++-
 example/druid.rego                                   |  4 ++--
 .../stackable/druid/opaauthorizer/OpaAuthorizer.java |  6 +-----
 .../druid/opaauthorizer/opatypes/OpaInput.java       | 12 +++++++++---
 .../druid/opaauthorizer/opatypes/OpaMessage.java     | 10 ++++++++--
 5 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/README.md b/README.md
index 09bc2cc..a5ce98a 100644
--- a/README.md
+++ b/README.md
@@ -36,7 +36,12 @@ Then the `myOpaAuth` authorizer needs to be referenced in your authenticator.
 The authorizer will send a request to the `uri` specified in the config. The input will be:
 
         {
-            user: <String: user name>
+            authenticationResult: {
+                identity: <String: user name>
+                authorizerName: <String>
+                authenticatedBy: <String>
+                context: Map<String, Object>
+            }
             action: <String: READ|WRITE>
             resource: {
                 name: <String>
diff --git a/example/druid.rego b/example/druid.rego
index c75151f..b3b52da 100644
--- a/example/druid.rego
+++ b/example/druid.rego
@@ -25,14 +25,14 @@ allow {
 # user_is_admin is true if...
 user_is_admin {
 	# "admin" is among the user's roles as per data.user_roles
-	"admin" in data.user_roles[input.user]
+	"admin" in data.user_roles[input.authenticationResult.identity]
 }
 
 # user_is_granted is a set of grants for the user identified in the request.
 # The `grant` will be contained if the set `user_is_granted` for every...
 user_is_granted[grant] {
 	# `role` assigned an element of the user_roles for this user...
-	some role in data.user_roles[input.user]
+	some role in data.user_roles[input.authenticationResult.identity]
 
 	# `grant` assigned a single grant from the grants list for 'role'...
 	some grant in data.role_grants[role]
diff --git a/src/main/java/tech/stackable/druid/opaauthorizer/OpaAuthorizer.java b/src/main/java/tech/stackable/druid/opaauthorizer/OpaAuthorizer.java
index 6448359..63d1a2f 100644
--- a/src/main/java/tech/stackable/druid/opaauthorizer/OpaAuthorizer.java
+++ b/src/main/java/tech/stackable/druid/opaauthorizer/OpaAuthorizer.java
@@ -45,11 +45,7 @@ public Access authorize(
         authenticationResult.getIdentity(), action.name(), resource.toString());
     LOG.trace("Creating OPA request JSON.");
     OpaMessage msg =
-        new OpaMessage(
-            authenticationResult.getIdentity(),
-            action.name(),
-            resource.getName(),
-            resource.getType());
+        new OpaMessage(authenticationResult, action.name(), resource.getName(), resource.getType());
     String msgJson;
     try {
       msgJson = objectMapper.writeValueAsString(msg);
diff --git a/src/main/java/tech/stackable/druid/opaauthorizer/opatypes/OpaInput.java b/src/main/java/tech/stackable/druid/opaauthorizer/opatypes/OpaInput.java
index b27c193..2b021e0 100644
--- a/src/main/java/tech/stackable/druid/opaauthorizer/opatypes/OpaInput.java
+++ b/src/main/java/tech/stackable/druid/opaauthorizer/opatypes/OpaInput.java
@@ -1,12 +1,18 @@
 package tech.stackable.druid.opaauthorizer.opatypes;
 
+import org.apache.druid.server.security.AuthenticationResult;
+
 public class OpaInput {
-  public String user;
+  public AuthenticationResult authenticationResult;
   public String action;
   public OpaResource resource;
 
-  public OpaInput(String user, String action, String resourceName, String resourceType) {
-    this.user = user;
+  public OpaInput(
+      AuthenticationResult authenticationResult,
+      String action,
+      String resourceName,
+      String resourceType) {
+    this.authenticationResult = authenticationResult;
     this.action = action;
     this.resource = new OpaResource(resourceName, resourceType);
   }
diff --git a/src/main/java/tech/stackable/druid/opaauthorizer/opatypes/OpaMessage.java b/src/main/java/tech/stackable/druid/opaauthorizer/opatypes/OpaMessage.java
index 98ec22a..bd1e1c7 100644
--- a/src/main/java/tech/stackable/druid/opaauthorizer/opatypes/OpaMessage.java
+++ b/src/main/java/tech/stackable/druid/opaauthorizer/opatypes/OpaMessage.java
@@ -1,9 +1,15 @@
 package tech.stackable.druid.opaauthorizer.opatypes;
 
+import org.apache.druid.server.security.AuthenticationResult;
+
 public class OpaMessage {
   public OpaInput input;
 
-  public OpaMessage(String user, String action, String resourceName, String resourceType) {
-    this.input = new OpaInput(user, action, resourceName, resourceType);
+  public OpaMessage(
+      AuthenticationResult authenticationResult,
+      String action,
+      String resourceName,
+      String resourceType) {
+    this.input = new OpaInput(authenticationResult, action, resourceName, resourceType);
   }
 }