From 19b6a2f5f58164d965b687d4d0ef7d8a851fc3fe Mon Sep 17 00:00:00 2001
From: Aditya Thebe <contact@adityathebe.com>
Date: Thu, 7 Nov 2024 11:57:33 +0545
Subject: [PATCH] feat: use the agent id in casbin rule

[skip ci]
---
 functions/postgrest.sql   |  9 +++------
 models/permission.go      | 21 ++++++++++++++++++++-
 models/permission_test.go |  8 ++++++++
 views/034_rls_enable.sql  |  1 +
 4 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/functions/postgrest.sql b/functions/postgrest.sql
index 9e6df7c3..f4606a26 100644
--- a/functions/postgrest.sql
+++ b/functions/postgrest.sql
@@ -17,12 +17,9 @@ END $$;
 
 DO $$
 BEGIN
-    IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = 'postgrest_api') THEN
-      -- CREATE a ROLE that will own all views where we need to enforce RLS.
-      CREATE ROLE api_views_owner NOSUPERUSER NOBYPASSRLS;
-
-      GRANT SELECT ON ALL TABLES IN SCHEMA public TO api_views_owner;
-    END IF ;
+  -- CREATE a ROLE that will own all views where we need to enforce RLS.
+  CREATE ROLE api_views_owner NOSUPERUSER NOBYPASSRLS;
+  GRANT SELECT ON ALL TABLES IN SCHEMA public TO api_views_owner;
 END
 $$;
 
diff --git a/models/permission.go b/models/permission.go
index 8cfad08f..7c5fb48c 100644
--- a/models/permission.go
+++ b/models/permission.go
@@ -31,7 +31,7 @@ type Permission struct {
 	UpdatedBy    *uuid.UUID `json:"updated_by"`
 
 	// List of agent ids whose configs/components are accessible to a person when RLS is enabled
-	Agents pq.StringArray `json:"agents,omitempty"`
+	Agents pq.StringArray `json:"agents,omitempty" gorm:"type:[]text"`
 
 	// List of config/component tags a person is allowed access to when RLS is enabled
 	Tags types.JSONStringMap `json:"tags,omitempty"`
@@ -68,6 +68,25 @@ func (t *Permission) Condition() string {
 		rule = append(rule, fmt.Sprintf("r.obj.playbook != undefined && r.obj.playbook.id == %q", t.PlaybookID.String()))
 	}
 
+	if len(t.Agents) > 0 {
+		var agents []string
+		for _, agentID := range t.Agents {
+			agents = append(agents, fmt.Sprintf("'%s'", agentID))
+		}
+
+		rule = append(rule, fmt.Sprintf("r.obj.config != undefined && r.obj.config.agent_id in (%s)", strings.Join(agents, ",")))
+		rule = append(rule, fmt.Sprintf("r.obj.component != undefined && r.obj.component.agent_id in (%s)", strings.Join(agents, ",")))
+		rule = append(rule, fmt.Sprintf("r.obj.canary != undefined && r.obj.canary.agent_id in (%s)", strings.Join(agents, ",")))
+	}
+
+	// if len(t.Tags) > 0 {
+	// 	var tagsClause []string
+	// 	for _, agentID := range t.Tags {
+	// 	}
+	//
+	// 	rule = append(rule, strings.Join(tagsClause, " || "))
+	// }
+
 	return strings.Join(rule, " && ")
 }
 
diff --git a/models/permission_test.go b/models/permission_test.go
index 1ddd9525..c5ca8df1 100644
--- a/models/permission_test.go
+++ b/models/permission_test.go
@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"github.com/google/uuid"
+	"github.com/lib/pq"
 	"github.com/samber/lo"
 )
 
@@ -33,6 +34,13 @@ func TestPermission_Condition(t *testing.T) {
 			perm:     Permission{},
 			expected: "",
 		},
+		{
+			name: "agents",
+			perm: Permission{
+				Agents: pq.StringArray([]string{"aws", "azure"}),
+			},
+			expected: "r.obj.config != undefined && r.obj.config.agent_id in ('aws','azure') && r.obj.component != undefined && r.obj.component.agent_id in ('aws','azure') && r.obj.canary != undefined && r.obj.canary.agent_id in ('aws','azure')",
+		},
 	}
 
 	for _, tt := range tests {
diff --git a/views/034_rls_enable.sql b/views/034_rls_enable.sql
index e185b85f..473440cd 100644
--- a/views/034_rls_enable.sql
+++ b/views/034_rls_enable.sql
@@ -5,6 +5,7 @@ ALTER TABLE components ENABLE ROW LEVEL SECURITY;
 -- Policy config items
 DROP POLICY IF EXISTS config_items_auth ON config_items;
 
+-- TODO:: Don't re-add policy if it exists
 CREATE POLICY config_items_auth ON config_items
   FOR ALL TO postgrest_api, postgrest_anon
     USING (tags::jsonb @> (current_setting('request.jwt.claims', TRUE)::json ->> 'tags')::jsonb