Skip to content

Commit c14e6ec

Browse files
craig[bot]spilchen
craig[bot]
and
spilchen
committed
Merge #140069
140069: sql/opt: Invalidate memo for queries using RLS when user changes r=spilchen a=spilchen When reusing a query from the statement cache, we generally allow reuse if the query is to be executed by a different user. However, this approach is problematic for queries that involve row-level security (RLS), as the executing user directly influences the injected expressions. To ensure correctness, this considers the memo as stale when the user changes if RLS was used in the query. Epic: CRDB-11724 Release note: None Closes #136717 Co-authored-by: Matt Spilchen <[email protected]>
2 parents 1eceee1 + e17bcc6 commit c14e6ec

File tree

10 files changed

+332
-9
lines changed

10 files changed

+332
-9
lines changed

‎pkg/sql/logictest/testdata/logic_test/row_level_security

Lines changed: 193 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -567,16 +567,15 @@ nationals NL
567567
statement ok
568568
set role buck
569569

570-
# TODO(136717): Reusing the statement cache prevents recompiling a new plan,
571-
# which leads to incorrect results in this case.
570+
# Retry the same query issue before to ensure memo is properly invalidated.
572571
query TT
573572
select team, league from bbteams where team != 'cardinals' order by league, team;
574573
----
575574
jays AL
576575
orioles AL
577576
tigers AL
578-
nationals NL
579577

578+
# Try another query we know isn't in the statement cache.
580579
query TT
581580
select team, league from bbteams where team != 'astros' order by league, team;
582581
----
@@ -595,14 +594,14 @@ statement ok
595594
set role buck;
596595

597596
# This is the same query as before, but since admin changed, we will see all of the rows.
598-
# TODO(136717): We are mistakenly reusing the statement plan here. So we see the rows as if
599-
# the policies were applied.
600597
query TT
601598
select team, league from bbteams where team != 'astros' order by league, team;
602599
----
603600
jays AL
604601
orioles AL
605602
tigers AL
603+
cardinals NL
604+
nationals NL
606605

607606
# Retry with a query never tried before so we avoid the statement cache.
608607
query TT
@@ -620,6 +619,20 @@ set role root
620619
statement ok
621620
REVOKE admin FROM buck;
622621

622+
statement ok
623+
set role buck
624+
625+
# Retry same query we ran before to ensure it's properly invalidated in the statement cache.
626+
query TT
627+
select team, league from bbteams where team != 'mariners' order by league, team;
628+
----
629+
jays AL
630+
orioles AL
631+
tigers AL
632+
633+
statement ok
634+
set role root
635+
623636
# Add policies that apply to other commands. Only SELECT will return rows.
624637
statement ok
625638
CREATE POLICY restrict_insert ON bbteams FOR INSERT TO buck USING (false);
@@ -767,4 +780,179 @@ use defaultdb;
767780
statement ok
768781
drop database db2 cascade;
769782

783+
# Ensure that functions defined with security behave as expected
784+
subtest function_security_definer
785+
786+
statement ok
787+
CREATE USER sensitive_user;
788+
789+
statement ok
790+
CREATE TABLE sensitive_data_table (C1 INT);
791+
792+
statement ok
793+
INSERT INTO sensitive_data_table VALUES (0),(1),(2);
794+
795+
statement ok
796+
ALTER TABLE sensitive_data_table ENABLE ROW LEVEL SECURITY;
797+
798+
statement ok
799+
GRANT ALL ON sensitive_data_table TO sensitive_user;
800+
801+
statement ok
802+
CREATE FUNCTION my_sec_definer_function() RETURNS TABLE(ID INT)
803+
LANGUAGE SQL AS
804+
$$
805+
SELECT * FROM sensitive_data_table
806+
$$ SECURITY DEFINER;
807+
808+
statement ok
809+
CREATE FUNCTION my_non_sec_definer_function() RETURNS TABLE(ID INT)
810+
LANGUAGE SQL AS
811+
$$
812+
SELECT * FROM sensitive_data_table
813+
$$;
814+
815+
statement ok
816+
SET ROLE sensitive_user;
817+
818+
query I rowsort
819+
SELECT * FROM sensitive_data_table;
820+
----
821+
822+
query I rowsort
823+
SELECT my_sec_definer_function();
824+
----
825+
0
826+
1
827+
2
828+
829+
query I rowsort
830+
SELECT my_non_sec_definer_function();
831+
-----
832+
833+
statement ok
834+
SET ROLE root
835+
836+
statement ok
837+
CREATE POLICY p1 ON sensitive_data_table FOR SELECT TO sensitive_user USING (C1 != 0);
838+
839+
statement ok
840+
SET ROLE sensitive_user;
841+
842+
query I rowsort
843+
SELECT my_sec_definer_function();
844+
----
845+
0
846+
1
847+
2
848+
849+
query I rowsort
850+
SELECT my_non_sec_definer_function()
851+
----
852+
1
853+
2
854+
855+
statement ok
856+
SET ROLE root
857+
858+
statement ok
859+
DROP FUNCTION my_sec_definer_function;
860+
861+
statement ok
862+
DROP FUNCTION my_non_sec_definer_function;
863+
864+
statement ok
865+
DROP TABLE sensitive_data_table;
866+
867+
subtest validate_statement_cache_after_rls_changes
868+
869+
statement ok
870+
CREATE TABLE rls_cache_test (c1 TEXT);
871+
872+
statement ok
873+
INSERT INTO rls_cache_test VALUES ('a'), ('b'), ('c');
874+
875+
statement ok
876+
CREATE USER rls_cache_user;
877+
878+
statement ok
879+
GRANT ALL ON rls_cache_test TO rls_cache_user;
880+
881+
statement ok
882+
SET ROLE rls_cache_user;
883+
884+
# Prime the cache
885+
query T
886+
SELECT * FROM rls_cache_test ORDER BY c1;
887+
----
888+
a
889+
b
890+
c
891+
892+
statement ok
893+
SET ROLE root
894+
895+
statement ok
896+
ALTER TABLE rls_cache_test ENABLE ROW LEVEL SECURITY;
897+
898+
statement ok
899+
SET ROLE rls_cache_user;
900+
901+
# The cache should be invalidated
902+
query T
903+
SELECT * FROM rls_cache_test ORDER BY c1;
904+
----
905+
906+
statement ok
907+
SET ROLE root
908+
909+
statement ok
910+
CREATE POLICY rls_cache_policy ON rls_cache_test FOR SELECT TO rls_cache_user USING (c1 != 'a');
911+
912+
statement ok
913+
SET ROLE rls_cache_user;
914+
915+
# The cache should be invalidated (again)
916+
query T
917+
SELECT * FROM rls_cache_test ORDER BY c1;
918+
----
919+
b
920+
c
921+
922+
statement ok
923+
SET ROLE root
924+
925+
statement ok
926+
ALTER TABLE rls_cache_test DISABLE ROW LEVEL SECURITY;
927+
928+
statement ok
929+
SET ROLE rls_cache_user;
930+
931+
# The cache should be invalidated (again)
932+
query T
933+
SELECT * FROM rls_cache_test ORDER BY c1;
934+
----
935+
a
936+
b
937+
c
938+
939+
statement ok
940+
SET ROLE root
941+
942+
# Ensure that the cache is invalidated when table is dropped
943+
statement ok
944+
DROP TABLE rls_cache_test;
945+
946+
statement ok
947+
SET ROLE rls_cache_user;
948+
949+
statement error pq: relation "rls_cache_test" does not exist
950+
SELECT * FROM rls_cache_test ORDER BY c1;
951+
952+
statement ok
953+
SET ROLE root
954+
955+
statement ok
956+
DROP ROLE rls_cache_user;
957+
770958
subtest end

‎pkg/sql/opt/BUILD.bazel

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ go_library(
1212
"operator.go",
1313
"ordering.go",
1414
"panic_injection.go",
15+
"row_level_security.go",
1516
"rule_name.go",
1617
"schema_dependencies.go",
1718
"table_meta.go",
@@ -25,6 +26,7 @@ go_library(
2526
visibility = ["//visibility:public"],
2627
deps = [
2728
"//pkg/clusterversion",
29+
"//pkg/security/username",
2830
"//pkg/server/telemetry",
2931
"//pkg/sql/catalog",
3032
"//pkg/sql/catalog/catpb",

‎pkg/sql/opt/cat/catalog.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,9 @@ type Catalog interface {
224224
// returns true. Returns an error if query on the `system.users` table failed
225225
HasAdminRole(ctx context.Context) (bool, error)
226226

227+
// UserHasAdminRole checks if the specified user has admin privileges.
228+
UserHasAdminRole(ctx context.Context, user username.SQLUsername) (bool, error)
229+
227230
// HasRoleOption converts the roleoption to its SQL column name and checks if
228231
// the user belongs to a role where the option has value true. Requires a
229232
// valid transaction to be open.

‎pkg/sql/opt/memo/BUILD.bazel

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ go_test(
8383
embed = [":memo"],
8484
deps = [
8585
"//pkg/kv/kvserver/concurrency/isolation",
86+
"//pkg/security/username",
8687
"//pkg/settings/cluster",
8788
"//pkg/sql/inverted",
8889
"//pkg/sql/opt",

‎pkg/sql/opt/memo/memo_test.go

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111
"testing"
1212

1313
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/isolation"
14+
"github.com/cockroachdb/cockroach/pkg/security/username"
1415
"github.com/cockroachdb/cockroach/pkg/settings/cluster"
1516
"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
1617
"github.com/cockroachdb/cockroach/pkg/sql/opt/norm"
@@ -612,6 +613,24 @@ func TestMemoIsStale(t *testing.T) {
612613
stale()
613614
catalog.Function("one").Version = 0
614615
notStale()
616+
617+
// User changes (without RLS)
618+
oldUser := evalCtx.SessionData().UserProto
619+
newUser := username.MakeSQLUsernameFromPreNormalizedString("newuser").EncodeProto()
620+
evalCtx.SessionData().UserProto = newUser
621+
notStale()
622+
evalCtx.SessionData().UserProto = oldUser
623+
notStale()
624+
625+
// User changes (with RLS)
626+
o.Memo().Metadata().SetRLSEnabled(evalCtx.SessionData().User(), true /* admin */)
627+
notStale()
628+
evalCtx.SessionData().UserProto = newUser
629+
stale()
630+
evalCtx.SessionData().UserProto = oldUser
631+
notStale()
632+
o.Memo().Metadata().ClearRLSEnabled()
633+
notStale()
615634
}
616635

617636
// TestStatsAvailable tests that the statisticsBuilder correctly identifies

0 commit comments

Comments
 (0)