Skip to content

Commit d56a3fe

Browse files
SwiddisSwiddis
committed
Apply PR feedback
Signed-off-by: Simeon Widdis <sawiddis@amazon.com>
1 parent d460bf1 commit d56a3fe

1 file changed

Lines changed: 83 additions & 18 deletions

File tree

‎integ-test/src/test/java/org/opensearch/sql/security/FGACIndexScanningIT.java‎

Lines changed: 83 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
import java.io.IOException;
1212
import java.util.Locale;
1313
import lombok.SneakyThrows;
14+
import org.apache.hc.core5.http.ParseException;
15+
import org.apache.hc.core5.http.io.entity.EntityUtils;
1416
import org.jetbrains.annotations.NotNull;
1517
import org.json.JSONObject;
1618
import org.junit.jupiter.api.BeforeAll;
@@ -81,15 +83,15 @@ private void configureEngine(boolean useCalcite) throws IOException {
8183
}
8284
}
8385

84-
private void setupTestIndices() throws IOException {
86+
private void setupTestIndices() throws IOException, ParseException {
8587
createPublicLogsIndex();
8688
createSensitiveLogsIndex();
8789
createEmployeeRecordsIndex();
8890
createSecureLogsIndex();
8991
}
9092

9193
/** Creates public_logs index with test documents. */
92-
private void createPublicLogsIndex() throws IOException {
94+
private void createPublicLogsIndex() throws IOException, ParseException {
9395
Request request = new Request("PUT", "/" + PUBLIC_LOGS);
9496
request.setJsonEntity(
9597
"""
@@ -113,7 +115,7 @@ private void createPublicLogsIndex() throws IOException {
113115
}
114116

115117
/** Creates sensitive_logs index with test documents. */
116-
private void createSensitiveLogsIndex() throws IOException {
118+
private void createSensitiveLogsIndex() throws IOException, ParseException {
117119
Request request = new Request("PUT", "/" + SENSITIVE_LOGS);
118120
request.setJsonEntity(
119121
"""
@@ -140,7 +142,7 @@ private void createSensitiveLogsIndex() throws IOException {
140142
* Creates employee_records index with sensitive fields for field-level security testing. Contains
141143
* fields: employee_id, name, department, salary, ssn
142144
*/
143-
private void createEmployeeRecordsIndex() throws IOException {
145+
private void createEmployeeRecordsIndex() throws IOException, ParseException {
144146
Request request = new Request("PUT", "/" + EMPLOYEE_RECORDS);
145147
request.setJsonEntity(
146148
"""
@@ -170,7 +172,7 @@ private void createEmployeeRecordsIndex() throws IOException {
170172
* Creates secure_logs index with mixed security levels. This index contains documents with
171173
* different security_level values to test row-level filtering.
172174
*/
173-
private void createSecureLogsIndex() throws IOException {
175+
private void createSecureLogsIndex() throws IOException, ParseException {
174176
Request request = new Request("PUT", "/" + SECURE_LOGS);
175177
request.setJsonEntity(
176178
"""
@@ -194,7 +196,7 @@ private void createSecureLogsIndex() throws IOException {
194196
}
195197

196198
/** Bulk inserts documents to trigger background scanning. */
197-
private void bulkInsertDocs(String indexName, String prefix) throws IOException {
199+
private void bulkInsertDocs(String indexName, String prefix) throws IOException, ParseException {
198200
StringBuilder bulk = new StringBuilder();
199201
for (int i = 0; i < FGACIndexScanningIT.LARGE_DATASET_SIZE; i++) {
200202
bulk.append(
@@ -219,10 +221,13 @@ private void bulkInsertDocs(String indexName, String prefix) throws IOException
219221

220222
Response response = client().performRequest(request);
221223
assertEquals(200, response.getStatusLine().getStatusCode());
224+
String body = EntityUtils.toString(response.getEntity());
225+
JSONObject json = new JSONObject(body);
226+
assertFalse("Bulk indexing reported errors: " + body, json.getBoolean("errors"));
222227
}
223228

224229
/** Bulk inserts employee records with sensitive fields for FLS testing. */
225-
private void bulkInsertEmployeeRecords() throws IOException {
230+
private void bulkInsertEmployeeRecords() throws IOException, ParseException {
226231
String bulk = getBulkEmployeeIndexRequest();
227232

228233
Request request = new Request("POST", "/_bulk");
@@ -235,6 +240,9 @@ private void bulkInsertEmployeeRecords() throws IOException {
235240

236241
Response response = client().performRequest(request);
237242
assertEquals(200, response.getStatusLine().getStatusCode());
243+
String body = EntityUtils.toString(response.getEntity());
244+
JSONObject json = new JSONObject(body);
245+
assertFalse("Bulk indexing reported errors: " + body, json.getBoolean("errors"));
238246
}
239247

240248
@NotNull
@@ -263,7 +271,7 @@ private static String getBulkEmployeeIndexRequest() {
263271
}
264272

265273
/** Bulk inserts documents with different security levels for row-level testing. */
266-
private void bulkInsertDocsWithSecurityLevel() throws IOException {
274+
private void bulkInsertDocsWithSecurityLevel() throws IOException, ParseException {
267275
StringBuilder bulk = new StringBuilder();
268276

269277
int publicCount = LARGE_DATASET_SIZE / 2;
@@ -307,6 +315,26 @@ private void bulkInsertDocsWithSecurityLevel() throws IOException {
307315
i));
308316
}
309317

318+
// Add document with null security_level to test DLS behavior with null values
319+
bulk.append(
320+
String.format(
321+
Locale.ROOT,
322+
"""
323+
{ "index": { "_index": "%s" } }
324+
{ "message": "null security level", "security_level": null, "timestamp": "2025-01-01T00:00:00Z" }
325+
""",
326+
FGACIndexScanningIT.SECURE_LOGS));
327+
328+
// Add document with missing security_level field to test DLS behavior with missing fields
329+
bulk.append(
330+
String.format(
331+
Locale.ROOT,
332+
"""
333+
{ "index": { "_index": "%s" } }
334+
{ "message": "missing security level", "timestamp": "2025-01-01T00:00:00Z" }
335+
""",
336+
FGACIndexScanningIT.SECURE_LOGS));
337+
310338
Request request = new Request("POST", "/_bulk");
311339
request.addParameter("refresh", "true");
312340
request.setJsonEntity(bulk.toString());
@@ -317,6 +345,9 @@ private void bulkInsertDocsWithSecurityLevel() throws IOException {
317345

318346
Response response = client().performRequest(request);
319347
assertEquals(200, response.getStatusLine().getStatusCode());
348+
String body = EntityUtils.toString(response.getEntity());
349+
JSONObject json = new JSONObject(body);
350+
assertFalse("Bulk indexing reported errors: " + body, json.getBoolean("errors"));
320351
}
321352

322353
/** Creates security roles and users for testing. */
@@ -505,25 +536,44 @@ public void testFieldLevelSecurityEnforcedWithLargeDataset(boolean useCalcite)
505536
throws IOException {
506537
configureEngine(useCalcite);
507538
// Verify with large result set that FLS is still enforced
539+
// Query all data (not just count) to actually exercise FLS at scale
508540
String queryLargeDataset =
509541
String.format(
510-
"search source=%s | fields name, salary, department | stats count()", EMPLOYEE_RECORDS);
542+
"search source=%s | fields name, salary, department, employee_id", EMPLOYEE_RECORDS);
511543
JSONObject managerLargeResult = executeQueryAsUser(queryLargeDataset, MANAGER_USER);
512544

513-
// Even with large dataset, manager should not see ssn
545+
// Verify the schema contains allowed fields but not restricted fields
514546
var largeSchema = managerLargeResult.getJSONArray("schema");
515-
boolean hasSSNInLarge = false;
547+
boolean hasName = false;
548+
boolean hasSalary = false;
549+
boolean hasDepartment = false;
550+
boolean hasEmployeeId = false;
551+
boolean hasSSN = false;
552+
516553
for (int i = 0; i < largeSchema.length(); i++) {
517-
if ("ssn".equals(largeSchema.getJSONObject(i).getString("name"))) {
518-
hasSSNInLarge = true;
519-
break;
520-
}
554+
String fieldName = largeSchema.getJSONObject(i).getString("name");
555+
if ("name".equals(fieldName)) hasName = true;
556+
if ("salary".equals(fieldName)) hasSalary = true;
557+
if ("department".equals(fieldName)) hasDepartment = true;
558+
if ("employee_id".equals(fieldName)) hasEmployeeId = true;
559+
if ("ssn".equals(fieldName)) hasSSN = true;
521560
}
522561

562+
// Verify allowed fields are present
563+
assertTrue("manager_user should see 'name' field in large dataset query", hasName);
564+
assertTrue("manager_user should see 'salary' field in large dataset query", hasSalary);
565+
assertTrue("manager_user should see 'department' field in large dataset query", hasDepartment);
566+
assertTrue("manager_user should see 'employee_id' field in large dataset query", hasEmployeeId);
567+
568+
// Verify restricted field is NOT present
523569
assertFalse(
524570
"SECURITY VIOLATION: manager_user should NOT see 'ssn' even with large dataset. "
525571
+ "Field-level security must be enforced.",
526-
hasSSNInLarge);
572+
hasSSN);
573+
574+
// Verify we actually got data back (not just an empty result)
575+
var datarows = managerLargeResult.getJSONArray("datarows");
576+
assertTrue("Expected to receive data from large dataset query", datarows.length() > 0);
527577
}
528578

529579
@ParameterizedTest
@@ -544,6 +594,21 @@ public void testRowLevelSecurity(boolean useCalcite) throws IOException {
544594
// Extract the datarows for validation
545595
var datarows = result.getJSONArray("datarows");
546596

597+
// Derive column indexes from schema to support both V2 and V3 engines
598+
var schema = result.getJSONArray("schema");
599+
int countIdx = -1;
600+
int levelIdx = -1;
601+
for (int i = 0; i < schema.length(); i++) {
602+
String name = schema.getJSONObject(i).getString("name");
603+
if ("security_level".equals(name)) {
604+
levelIdx = i;
605+
} else if ("count()".equalsIgnoreCase(name) || "count".equalsIgnoreCase(name)) {
606+
countIdx = i;
607+
}
608+
}
609+
assertTrue("Expected count() in schema", countIdx >= 0);
610+
assertTrue("Expected security_level in schema", levelIdx >= 0);
611+
547612
// Count total documents visible
548613
int totalDocs = 0;
549614
boolean sawConfidential = false;
@@ -552,8 +617,8 @@ public void testRowLevelSecurity(boolean useCalcite) throws IOException {
552617

553618
for (int i = 0; i < datarows.length(); i++) {
554619
var row = datarows.getJSONArray(i);
555-
int count = row.getInt(0);
556-
String securityLevel = row.getString(1);
620+
int count = row.getInt(countIdx);
621+
String securityLevel = row.getString(levelIdx);
557622
totalDocs += count;
558623

559624
if ("confidential".equals(securityLevel)) {

0 commit comments

Comments
 (0)