Skip to content

Add more tests for column masking #10437

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
guyco33 wants to merge 3 commits into from
Closed

Add more tests for column masking #10437

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
65bfb53
Add more tests for column masking
guyco33 Dec 30, 2021
92a9702
Fix column masks plan
guyco33 Jan 4, 2022
110b6b8
Add tests
guyco33 Jan 4, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 9 additions & 12 deletions core/trino-main/src/main/java/io/trino/sql/planner/RelationPlanner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -296,26 +296,23 @@ private RelationPlan addColumnMasks(Table table, RelationPlan plan)
PlanBuilder planBuilder = newPlanBuilder(plan, analysis, lambdaDeclarationToSymbolMap)
.withScope(analysis.getAccessControlScope(table), plan.getFieldMappings()); // The fields in the access control scope has the same layout as those for the table scope

Map<Symbol, Expression> assignments = new LinkedHashMap<>();
for (Symbol symbol : planBuilder.getRoot().getOutputSymbols()) {
assignments.put(symbol, symbol.toSymbolReference());
}
for (int i = 0; i < plan.getDescriptor().getAllFieldCount(); i++) {
Field field = plan.getDescriptor().getFieldByIndex(i);

for (Expression mask : columnMasks.getOrDefault(field.getName().orElseThrow(), ImmutableList.of())) {
planBuilder = subqueryPlanner.handleSubqueries(planBuilder, mask, analysis.getSubqueries(mask));

Map<Symbol, Expression> assignments = new LinkedHashMap<>();
for (Symbol symbol : planBuilder.getRoot().getOutputSymbols()) {
assignments.put(symbol, symbol.toSymbolReference());
}
assignments.put(plan.getFieldMappings().get(i), coerceIfNecessary(analysis, mask, planBuilder.rewrite(mask)));

planBuilder = planBuilder
.withNewRoot(new ProjectNode(
idAllocator.getNextId(),
planBuilder.getRoot(),
Assignments.copyOf(assignments)));
}
}

planBuilder = planBuilder
.withNewRoot(new ProjectNode(
idAllocator.getNextId(),
planBuilder.getRoot(),
Assignments.copyOf(assignments)));
return new RelationPlan(planBuilder.getRoot(), plan.getScope(), plan.getFieldMappings(), outerContext);
}

Expand Down
123 changes: 113 additions & 10 deletions core/trino-main/src/test/java/io/trino/sql/query/TestColumnMask.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,7 @@ public void testSimpleMask()
}

@Test
public void testMultipleMasksOnSameColumn()
Copy link
Copy Markdown
Member Author

@guyco33 guyco33 Jan 4, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removing this test, since AC doesn't allow to have more than one masking per column.

trino/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/security/TableAccessControlRule.java

Line 72 in 35986d7

this.columnConstraints = Maps.uniqueIndex(requireNonNull(columns, "columns is null").orElse(ImmutableList.of()), ColumnConstraint::getName);

public void testMultipleMasksOnDifferentColumns()
{
accessControl.reset();
accessControl.columnMask(
Expand All @@ -191,31 +191,134 @@ public void testMultipleMasksOnSameColumn()

accessControl.columnMask(
new QualifiedObjectName(CATALOG, "tiny", "orders"),
"custkey",
"orderstatus",
USER,
new ViewExpression(USER, Optional.empty(), Optional.empty(), "custkey * 2"));
new ViewExpression(USER, Optional.empty(), Optional.empty(), "'X'"));

assertThat(assertions.query("SELECT custkey FROM orders WHERE orderkey = 1")).matches("VALUES BIGINT '-740'");
assertThat(assertions.query("SELECT custkey, orderstatus FROM orders WHERE orderkey = 1"))
.matches("VALUES (BIGINT '-370', 'X')");
}

@Test
public void testMultipleMasksOnDifferentColumns()
public void testMultipleMasksUsingOtherMaskedColumns()
{
String query = "SELECT comment, orderstatus, clerk FROM orders WHERE orderkey = 1";
String expected = "VALUES (CAST('nstructions sleep furiously among ' as varchar(79)), 'O', 'Clerk#000000951')";

Copy link
Copy Markdown
Contributor

@findinpath findinpath Jan 3, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it would be worth for the readability of the test scenario to perform the query without any column masking

accessControl.reset();
assertThat(assertions.query(query)).matches(expected);

// mask "clerk" and "orderstatus" using "comment" ("comment" appears after both in table definition)
// columns will not be masked since rules are not apply to row values
accessControl.reset();
accessControl.columnMask(
new QualifiedObjectName(CATALOG, "tiny", "orders"),
"custkey",
"comment",
USER,
new ViewExpression(USER, Optional.empty(), Optional.empty(), "-custkey"));
new ViewExpression(USER, Optional.empty(), Optional.empty(), "cast(regexp_replace(comment,'(password: [^ ]+)','password: ****') as varchar(79))"));
Copy link
Copy Markdown
Contributor

@findinpath findinpath Jan 3, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't quite get the purpose of this masking.
The column comment has the value nstructions sleep furiously among .
Why do you check for password in this text?

Copy link
Copy Markdown
Member Author

@guyco33 guyco33 Jan 3, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I want to mask some values in comment column (in this case the password value) and also to use comment column as a reference for the decision about masking other columns (we called it "conditional masking")

In a real example there is a json payload column with some attributes that should be masked and a country attribute that used as a reference for masking other columns.

Copy link
Copy Markdown
Contributor

@findinpath findinpath Jan 3, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The column comment has the value nstructions sleep furiously among .

This value doesn't contain a password substring. So, I guess the masking is pointless in this case. Or am I missing something?

Copy link
Copy Markdown
Member Author

@guyco33 guyco33 Jan 3, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I found that issue happens in InlineProjections#inlineProjections
Maybe it relates on how column masks are added

Copy link
Copy Markdown
Member Author

@guyco33 guyco33 Jan 3, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The column comment has the value nstructions sleep furiously among .

This value doesn't contain a password substring. So, I guess the masking is pointless in this case. Or am I missing something?

You are right. That's why I added the last test (L293:L337)

Copy link
Copy Markdown
Contributor

@findinpath findinpath Jan 3, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I had a look again on the test and I am sorry to say that I still don't quite understand why isn't there any masking in the result of the queries.
Shouldn't the test demonstrate that the * is masking the content of the returned columns on which the masking should apply?

Copy link
Copy Markdown
Member Author

@guyco33 guyco33 Jan 3, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shouldn't the test demonstrate that the * is masking the content of the returned columns on which the masking should apply?

You are right. I fixed it now. Take a look on the tests with orderkey = 39 (L293:L337)

Copy link
Copy Markdown
Member Author

@guyco33 guyco33 Jan 3, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The column comment has the value nstructions sleep furiously among .

This value doesn't contain a password substring. So, I guess the masking is pointless in this case. Or am I missing something?

That's right. masking is pointless for this specific row, but I added it to demonstrate the issue when adding the masking on orderstatus while refering clerk (L291)

Copy link
Copy Markdown
Contributor

@findinpath findinpath Jan 4, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see now.

The purpose of the test is currently just to demonstrate that the masking doesn't work properly. Thank you for the patience in trying to explain the purpose of the tests.


accessControl.columnMask(
new QualifiedObjectName(CATALOG, "tiny", "orders"),
"orderstatus",
USER,
new ViewExpression(USER, Optional.empty(), Optional.empty(), "'X'"));
new ViewExpression(USER, Optional.empty(), Optional.empty(), "if(regexp_extract(comment,'(country: [^ ]+)') IN ('country: 1'), '*', orderstatus)"));

assertThat(assertions.query("SELECT custkey, orderstatus FROM orders WHERE orderkey = 1"))
.matches("VALUES (BIGINT '-370', 'X')");
accessControl.columnMask(
new QualifiedObjectName(CATALOG, "tiny", "orders"),
"clerk",
USER,
new ViewExpression(USER, Optional.empty(), Optional.empty(), "if(regexp_extract(comment,'(country: [^ ]+)') IN ('country: 1'), '***', clerk)"));

assertThat(assertions.query(query)).matches(expected);

// mask "comment" using "clerk" ("clerk" column appears before "comment" in table definition)
// columns will not be masked since rules are not apply to row values
accessControl.reset();
accessControl.columnMask(
new QualifiedObjectName(CATALOG, "tiny", "orders"),
"clerk",
USER,
new ViewExpression(USER, Optional.empty(), Optional.empty(), "cast(regexp_replace(clerk,'(password: [^ ]+)','password: ****') as varchar(15))"));

accessControl.columnMask(
new QualifiedObjectName(CATALOG, "tiny", "orders"),
"comment",
USER,
new ViewExpression(USER, Optional.empty(), Optional.empty(), "if(regexp_extract(clerk,'(country: [^ ]+)') IN ('country: 1'), '***', comment)"));

assertThat(assertions.query(query)).matches(expected);

// now add mask for "orderstatus" using "clerk"
// columns will not be masked since rules are not apply to row values
accessControl.reset();
accessControl.columnMask(
new QualifiedObjectName(CATALOG, "tiny", "orders"),
"clerk",
USER,
new ViewExpression(USER, Optional.empty(), Optional.empty(), "cast(regexp_replace(clerk,'(password: [^ ]+)','password: ****') as varchar(15))"));

accessControl.columnMask(
new QualifiedObjectName(CATALOG, "tiny", "orders"),
"orderstatus",
USER,
new ViewExpression(USER, Optional.empty(), Optional.empty(), "if(regexp_extract(clerk,'(country: [^ ]+)') IN ('country: 1'), '*', orderstatus)"));

accessControl.columnMask(
new QualifiedObjectName(CATALOG, "tiny", "orders"),
"comment",
USER,
new ViewExpression(USER, Optional.empty(), Optional.empty(), "if(regexp_extract(clerk,'(country: [^ ]+)') IN ('country: 1'), '***', comment)"));

assertThat(assertions.query(query)).matches(expected);
Copy link
Copy Markdown
Member Author

@guyco33 guyco33 Jan 3, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fails with Symbol clerk already has assignment


query = "SELECT comment, orderstatus, clerk FROM orders WHERE orderkey = 39";

accessControl.reset();
assertThat(assertions.query(query))
.matches("VALUES (CAST('ole express, ironic requests: ir' as varchar(79)), 'O', 'Clerk#000000659')");

// mask "comment" and "orderstatus" using "clerk" ("clerk" appears before "comment" table definition)
accessControl.reset();
accessControl.columnMask(
new QualifiedObjectName(CATALOG, "tiny", "orders"),
"clerk",
USER,
new ViewExpression(USER, Optional.empty(), Optional.empty(), "cast(regexp_replace(clerk,'(Clerk#)','***#') as varchar(15))"));

accessControl.columnMask(
new QualifiedObjectName(CATALOG, "tiny", "orders"),
"comment",
USER,
new ViewExpression(USER, Optional.empty(), Optional.empty(), "if(regexp_extract(clerk,'([1-9]+)') IN ('659'), '***', comment)"));

assertThat(assertions.query(query))
.matches("VALUES (CAST('***' as varchar(79)), 'O', CAST('***#000000659' as varchar(15)))");

assertThat(assertions.query("SELECT comment, orderstatus, clerk, length(clerk) FROM orders WHERE orderkey = 39"))
.matches("VALUES (CAST('***' as varchar(79)), 'O', CAST('***#000000659' as varchar(15)), bigint'13')");

// mask "comment" and "orderstatus" using "clerk" ("clerk" appears before "comment" table definition)
accessControl.reset();
accessControl.columnMask(
new QualifiedObjectName(CATALOG, "tiny", "orders"),
"clerk",
USER,
new ViewExpression(USER, Optional.empty(), Optional.empty(), "cast(regexp_replace(clerk,'(Clerk#)','***#') as varchar(15))"));

accessControl.columnMask(
new QualifiedObjectName(CATALOG, "tiny", "orders"),
"orderstatus",
USER,
new ViewExpression(USER, Optional.empty(), Optional.empty(), "if(regexp_extract(clerk,'([1-9]+)') IN ('659'), '*', orderstatus)"));

accessControl.columnMask(
new QualifiedObjectName(CATALOG, "tiny", "orders"),
"comment",
USER,
new ViewExpression(USER, Optional.empty(), Optional.empty(), "if(regexp_extract(clerk,'([1-9]+)') IN ('659'), '***', comment)"));

assertThat(assertions.query(query))
.matches("VALUES (CAST('***' as varchar(79)), '*', CAST('***#000000659' as varchar(15)))");
}

@Test
Expand Down