This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 14bc3123b64 Validate segment upload destinations (#19231)
14bc3123b64 is described below
commit 14bc3123b64aa526bf45b84e8472687f3d248ab2
Author: Xiang Fu <[email protected]>
AuthorDate: Tue Aug 25 15:49:56 2026 -0700
Validate segment upload destinations (#19231)
---
.../api/access/AuthenticationFilter.java | 12 +-
.../PinotSegmentUploadDownloadRestletResource.java | 158 ++++++--
.../api/PinotSegmentUploadAuthorizationTest.java | 428 +++++++++++++++++++++
...otSegmentUploadDownloadRestletResourceTest.java | 60 +++
.../pinot/core/auth/FineGrainedAuthUtils.java | 13 +-
.../pinot/core/auth/FineGrainedAuthUtilsTest.java | 48 +++
.../minion/tasks/purge/PurgeTaskGenerator.java | 4 +-
.../RefreshSegmentTaskGenerator.java | 4 +-
.../RefreshSegmentTaskGeneratorTest.java | 98 +++++
9 files changed, 779 insertions(+), 46 deletions(-)
diff --git
a/pinot-controller/src/main/java/org/apache/pinot/controller/api/access/AuthenticationFilter.java
b/pinot-controller/src/main/java/org/apache/pinot/controller/api/access/AuthenticationFilter.java
index 015a97d59b3..f2bca693357 100644
---
a/pinot-controller/src/main/java/org/apache/pinot/controller/api/access/AuthenticationFilter.java
+++
b/pinot-controller/src/main/java/org/apache/pinot/controller/api/access/AuthenticationFilter.java
@@ -30,13 +30,16 @@ import javax.inject.Provider;
import javax.ws.rs.DELETE;
import javax.ws.rs.POST;
import javax.ws.rs.PUT;
+import javax.ws.rs.WebApplicationException;
import javax.ws.rs.container.ContainerRequestContext;
import javax.ws.rs.container.ContainerRequestFilter;
import javax.ws.rs.container.ResourceInfo;
import javax.ws.rs.core.Context;
import javax.ws.rs.core.HttpHeaders;
import javax.ws.rs.core.MultivaluedMap;
+import javax.ws.rs.core.Response;
import javax.ws.rs.core.UriInfo;
+import org.apache.commons.lang3.StringUtils;
import org.apache.pinot.common.auth.AuthProviderUtils;
import org.apache.pinot.common.utils.DatabaseUtils;
import org.apache.pinot.core.auth.Authorize;
@@ -103,8 +106,15 @@ public class AuthenticationFilter implements
ContainerRequestFilter {
// authorization. If table name is not available, it means the endpoint is
not a table-level endpoint.
String tableName = extractTableName(endpointMethod,
uriInfo.getPathParameters(), uriInfo.getQueryParameters());
if (tableName != null) {
+ if (StringUtils.isBlank(tableName)) {
+ throw new WebApplicationException("Table name must not be blank",
Response.Status.BAD_REQUEST);
+ }
// If table name is present, translate it to the fully qualified name
based on database header.
- tableName = DatabaseUtils.translateTableName(tableName, _httpHeaders);
+ try {
+ tableName = DatabaseUtils.translateTableName(tableName, _httpHeaders);
+ } catch (RuntimeException e) {
+ throw new WebApplicationException("Invalid table name: " +
e.getMessage(), e, Response.Status.BAD_REQUEST);
+ }
}
AccessType accessType = extractAccessType(endpointMethod);
AccessControlUtils.validatePermission(tableName, accessType, _httpHeaders,
endpointUrl, accessControl);
diff --git
a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java
b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java
index 599464a75af..5446967d9e3 100644
---
a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java
+++
b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java
@@ -260,20 +260,12 @@ public class PinotSegmentUploadDownloadRestletResource {
private SuccessResponse uploadSegment(@Nullable String tableName, TableType
tableType,
@Nullable FormDataMultiPart multiPart, boolean
copySegmentToFinalLocation, boolean enableParallelPushProtection,
- boolean allowRefresh, HttpHeaders headers, Request request) {
+ boolean allowRefresh, boolean requireMatchingMetadataTable, HttpHeaders
headers, Request request) {
long segmentUploadStartTimeMs = System.currentTimeMillis();
- if (StringUtils.isNotEmpty(tableName)) {
- TableType tableTypeFromTableName =
TableNameBuilder.getTableTypeFromTableName(tableName);
- if (tableTypeFromTableName != null && tableTypeFromTableName !=
tableType) {
- throw new ControllerApplicationException(LOGGER,
- String.format("Table name: %s does not match table type: %s",
tableName, tableType),
- Response.Status.BAD_REQUEST);
- }
- }
-
- // TODO: Consider validating the segment name and table name from the
header against the actual segment
+ // TODO: Consider validating the segment name from the header against the
actual segment
extractHttpHeader(headers,
CommonConstants.Controller.SEGMENT_NAME_HTTP_HEADER);
- extractHttpHeader(headers,
CommonConstants.Controller.TABLE_NAME_HTTP_HEADER);
+ String tableNameInHeader = extractHttpHeader(headers,
CommonConstants.Controller.TABLE_NAME_HTTP_HEADER);
+ String requestedTableName = resolveRequestedTableName(tableName,
tableNameInHeader, tableType, headers);
String uploadTypeStr = extractHttpHeader(headers,
FileUploadDownloadClient.CustomHeaders.UPLOAD_TYPE);
String sourceDownloadURIStr = extractHttpHeader(headers,
FileUploadDownloadClient.CustomHeaders.DOWNLOAD_URI);
@@ -318,7 +310,7 @@ public class PinotSegmentUploadDownloadRestletResource {
"Source download URI is required in header field
'DOWNLOAD_URI' for URI upload mode",
Response.Status.BAD_REQUEST);
}
- downloadSegmentFileFromURI(sourceDownloadURIStr, destFile,
tableName);
+ downloadSegmentFileFromURI(sourceDownloadURIStr, destFile,
requestedTableName);
segmentSizeInBytes = destFile.length();
break;
case METADATA:
@@ -367,19 +359,15 @@ public class PinotSegmentUploadDownloadRestletResource {
// Fetch segment name
String segmentName = segmentMetadata.getName();
- // Fetch table name. Try to derive the table name from the parameter and
then from segment metadata
- String rawTableName;
- if (StringUtils.isNotEmpty(tableName)) {
- rawTableName = TableNameBuilder.extractRawTableName(tableName);
- } else {
- // TODO: remove this when we completely deprecate the table name from
segment metadata
- rawTableName = segmentMetadata.getTableName();
- LOGGER.warn("Table name is not provided as request query parameter
when uploading segment: {} for table: {}",
- segmentName, rawTableName);
- }
- String tableNameWithType = tableType == TableType.OFFLINE
- ? TableNameBuilder.OFFLINE.tableNameWithType(rawTableName)
- : TableNameBuilder.REALTIME.tableNameWithType(rawTableName);
+ String rawTableName = resolveDestinationTableName(requestedTableName,
segmentMetadata.getTableName(), tableType,
+ headers, requireMatchingMetadataTable);
+ String tableNameWithType =
TableNameBuilder.forType(tableType).tableNameWithType(rawTableName);
+
+ // The v1 endpoints are authorized at cluster scope before segment
extraction because tableName is optional.
+ // Re-authorize all variants against the canonical destination so no
AccessControl implementation can permit a
+ // different table through cluster-level CREATE access or
database-header translation.
+ ResourceUtils.checkPermissionAndAccess(rawTableName, request, headers,
AccessType.CREATE,
+ Actions.Table.UPLOAD_SEGMENT, _accessControlFactory, LOGGER);
if
(UploadedRealtimeSegmentName.isUploadedRealtimeSegmentName(segmentName) &&
tableType != TableType.REALTIME) {
throw new ControllerApplicationException(LOGGER, "Cannot upload
segment: " + segmentName
@@ -574,9 +562,20 @@ public class PinotSegmentUploadDownloadRestletResource {
private SuccessResponse uploadSegments(String tableName, TableType
tableType, FormDataMultiPart multiPart,
boolean enableParallelPushProtection, boolean allowRefresh, HttpHeaders
headers, Request request) {
long segmentsUploadStartTimeMs = System.currentTimeMillis();
- String rawTableName = TableNameBuilder.extractRawTableName(tableName);
- String tableNameWithType = tableType == TableType.OFFLINE ?
TableNameBuilder.OFFLINE.tableNameWithType(rawTableName)
- : TableNameBuilder.REALTIME.tableNameWithType(rawTableName);
+ String rawTableName = normalizeTableName(tableName, tableType, headers,
"request tableName");
+ if (rawTableName == null) {
+ throw new ControllerApplicationException(LOGGER, "tableName is required
for batch segment upload",
+ Response.Status.BAD_REQUEST);
+ }
+ String tableNameInHeader = normalizeTableName(
+ extractHttpHeader(headers,
CommonConstants.Controller.TABLE_NAME_HTTP_HEADER), tableType, headers,
+ CommonConstants.Controller.TABLE_NAME_HTTP_HEADER + " header");
+ validateMatchingTableName(rawTableName, tableNameInHeader,
+ CommonConstants.Controller.TABLE_NAME_HTTP_HEADER + " header");
+ String tableNameWithType =
TableNameBuilder.forType(tableType).tableNameWithType(rawTableName);
+
+ ResourceUtils.checkPermissionAndAccess(rawTableName, request, headers,
AccessType.CREATE,
+ Actions.Table.UPLOAD_SEGMENT, _accessControlFactory, LOGGER);
TableConfig tableConfig =
_pinotHelixResourceManager.getTableConfig(tableNameWithType);
if (tableConfig == null) {
@@ -648,6 +647,9 @@ public class PinotSegmentUploadDownloadRestletResource {
String metadataProviderClass =
DefaultMetadataExtractor.class.getName();
SegmentMetadata segmentMetadata =
getSegmentMetadata(tempDecryptedFile, tempSegmentDir, metadataProviderClass);
+ String segmentMetadataTableName =
+ normalizeTableName(segmentMetadata.getTableName(), tableType,
headers, "segment metadata table name");
+ validateMatchingTableName(rawTableName, segmentMetadataTableName,
"segment metadata table name");
LOGGER.info("Processing upload request for segment: {} of table: {}
with upload type: {} from client: {}, "
+ "ingestion descriptor: {}", segmentName, tableNameWithType,
uploadType, clientAddress,
ingestionDescriptor);
@@ -735,6 +737,84 @@ public class PinotSegmentUploadDownloadRestletResource {
}
}
+ @VisibleForTesting
+ static String resolveDestinationTableName(@Nullable String requestTableName,
@Nullable String headerTableName,
+ @Nullable String metadataTableName, TableType tableType, HttpHeaders
headers,
+ boolean requireMatchingMetadataTable) {
+ String requestedTableName = resolveRequestedTableName(requestTableName,
headerTableName, tableType, headers);
+ return resolveDestinationTableName(requestedTableName, metadataTableName,
tableType, headers,
+ requireMatchingMetadataTable);
+ }
+
+ @Nullable
+ private static String resolveRequestedTableName(@Nullable String
requestTableName, @Nullable String headerTableName,
+ TableType tableType, HttpHeaders headers) {
+ String normalizedRequestTable = normalizeTableName(requestTableName,
tableType, headers, "request tableName");
+ String normalizedHeaderTable = normalizeTableName(headerTableName,
tableType, headers,
+ CommonConstants.Controller.TABLE_NAME_HTTP_HEADER + " header");
+ String destinationTable = normalizedRequestTable != null ?
normalizedRequestTable
+ : normalizedHeaderTable;
+ if (destinationTable != null) {
+ validateMatchingTableName(destinationTable, normalizedHeaderTable,
+ CommonConstants.Controller.TABLE_NAME_HTTP_HEADER + " header");
+ }
+ return destinationTable;
+ }
+
+ private static String resolveDestinationTableName(@Nullable String
requestedTableName,
+ @Nullable String metadataTableName, TableType tableType, HttpHeaders
headers,
+ boolean requireMatchingMetadataTable) {
+ String normalizedMetadataTable = null;
+ if (requireMatchingMetadataTable || requestedTableName == null) {
+ normalizedMetadataTable =
+ normalizeTableName(metadataTableName, tableType, headers, "segment
metadata table name");
+ }
+ String destinationTable = requestedTableName != null ? requestedTableName
: normalizedMetadataTable;
+ if (destinationTable == null) {
+ throw new ControllerApplicationException(LOGGER,
+ "Table name is required in the request, " +
CommonConstants.Controller.TABLE_NAME_HTTP_HEADER
+ + " header, or segment metadata",
+ Response.Status.BAD_REQUEST);
+ }
+
+ if (requireMatchingMetadataTable) {
+ validateMatchingTableName(destinationTable, normalizedMetadataTable,
"segment metadata table name");
+ }
+ return destinationTable;
+ }
+
+ @Nullable
+ private static String normalizeTableName(@Nullable String tableName,
TableType tableType, HttpHeaders headers,
+ String source) {
+ if (tableName == null) {
+ return null;
+ }
+ if (StringUtils.isBlank(tableName)) {
+ throw new ControllerApplicationException(LOGGER, "Invalid " + source +
": table name must not be blank",
+ Response.Status.BAD_REQUEST);
+ }
+ try {
+ TableType tableTypeFromName =
TableNameBuilder.getTableTypeFromTableName(tableName);
+ if (tableTypeFromName != null && tableTypeFromName != tableType) {
+ throw new IllegalArgumentException(
+ String.format("Table name: %s does not match table type: %s",
tableName, tableType));
+ }
+ return
DatabaseUtils.translateTableName(TableNameBuilder.extractRawTableName(tableName),
headers);
+ } catch (RuntimeException e) {
+ throw new ControllerApplicationException(LOGGER, "Invalid " + source +
": " + e.getMessage(),
+ Response.Status.BAD_REQUEST, e);
+ }
+ }
+
+ private static void validateMatchingTableName(String destinationTable,
@Nullable String suppliedTable,
+ String source) {
+ if (suppliedTable != null && !destinationTable.equals(suppliedTable)) {
+ throw new ControllerApplicationException(LOGGER,
+ String.format("%s '%s' does not match destination table '%s'",
source, suppliedTable, destinationTable),
+ Response.Status.BAD_REQUEST);
+ }
+ }
+
@Nullable
private String extractHttpHeader(HttpHeaders headers, String name) {
String value = headers.getHeaderString(name);
@@ -778,7 +858,8 @@ public class PinotSegmentUploadDownloadRestletResource {
return out;
}
- private void downloadSegmentFileFromURI(String currentSegmentLocationURI,
File destFile, String tableName)
+ private void downloadSegmentFileFromURI(String currentSegmentLocationURI,
File destFile,
+ @Nullable String tableName)
throws Exception {
if (currentSegmentLocationURI == null ||
currentSegmentLocationURI.isEmpty()) {
throw new ControllerApplicationException(LOGGER, "Failed to get
downloadURI, needed for URI upload",
@@ -830,8 +911,8 @@ public class PinotSegmentUploadDownloadRestletResource {
// request if a multipart object is not sent. This endpoint does not move
the segment to its final location;
// it keeps it at the downloadURI header that is set. We will not support
this endpoint going forward.
public void uploadSegmentAsJson(String segmentJsonStr,
- @ApiParam(value = "Name of the table to upload into. Overrides
segment.table.name in segment metadata when set "
- + "(allows promoting a segment built for another table). Falls back
to metadata when omitted.")
+ @ApiParam(value = "Name of the table to upload into. Must match
segment.table.name in segment metadata when both "
+ + "are set. Falls back to metadata when omitted.")
@QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_NAME)
String tableName,
@ApiParam(value = "Type of the table")
@QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_TYPE)
@@ -844,7 +925,7 @@ public class PinotSegmentUploadDownloadRestletResource {
@Context HttpHeaders headers, @Context Request request, @Suspended
AsyncResponse asyncResponse) {
try {
asyncResponse.resume(uploadSegment(tableName,
TableType.valueOf(tableType.toUpperCase()), null, false,
- enableParallelPushProtection, allowRefresh, headers, request));
+ enableParallelPushProtection, allowRefresh, true, headers, request));
} catch (Throwable t) {
asyncResponse.resume(t);
}
@@ -871,8 +952,8 @@ public class PinotSegmentUploadDownloadRestletResource {
@TrackedByGauge(gauge = ControllerGauge.SEGMENT_UPLOADS_IN_PROGRESS)
// For the multipart endpoint, we will always move segment to final location
regardless of the segment endpoint.
public void uploadSegmentAsMultiPart(FormDataMultiPart multiPart,
- @ApiParam(value = "Name of the table to upload into. Overrides
segment.table.name in segment metadata when set "
- + "(allows promoting a segment built for another table). Falls back
to metadata when omitted.")
+ @ApiParam(value = "Name of the table to upload into. Must match
segment.table.name in segment metadata when both "
+ + "are set. Falls back to metadata when omitted.")
@QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_NAME)
String tableName,
@ApiParam(value = "Type of the table")
@QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_TYPE)
@@ -885,7 +966,7 @@ public class PinotSegmentUploadDownloadRestletResource {
@Context HttpHeaders headers, @Context Request request, @Suspended
AsyncResponse asyncResponse) {
try {
asyncResponse.resume(uploadSegment(tableName,
TableType.valueOf(tableType.toUpperCase()), multiPart, true,
- enableParallelPushProtection, allowRefresh, headers, request));
+ enableParallelPushProtection, allowRefresh, true, headers, request));
} catch (Throwable t) {
asyncResponse.resume(t);
}
@@ -989,7 +1070,7 @@ public class PinotSegmentUploadDownloadRestletResource {
try {
asyncResponse.resume(
uploadSegment(tableName, TableType.valueOf(tableType.toUpperCase()),
null, true, enableParallelPushProtection,
- allowRefresh, headers, request));
+ allowRefresh, false, headers, request));
} catch (Throwable t) {
asyncResponse.resume(t);
}
@@ -1014,7 +1095,6 @@ public class PinotSegmentUploadDownloadRestletResource {
})
@TrackInflightRequestMetrics
@TrackedByGauge(gauge = ControllerGauge.SEGMENT_UPLOADS_IN_PROGRESS)
- // This behavior does not differ from v1 of the same endpoint.
public void uploadSegmentAsMultiPartV2(FormDataMultiPart multiPart,
@ApiParam(value = "Name of the table to upload into. Overrides
segment.table.name in segment metadata when set.")
@QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_NAME)
@@ -1029,7 +1109,7 @@ public class PinotSegmentUploadDownloadRestletResource {
@Context HttpHeaders headers, @Context Request request, @Suspended
AsyncResponse asyncResponse) {
try {
asyncResponse.resume(uploadSegment(tableName,
TableType.valueOf(tableType.toUpperCase()), multiPart, true,
- enableParallelPushProtection, allowRefresh, headers, request));
+ enableParallelPushProtection, allowRefresh, false, headers,
request));
} catch (Throwable t) {
asyncResponse.resume(t);
}
diff --git
a/pinot-controller/src/test/java/org/apache/pinot/controller/api/PinotSegmentUploadAuthorizationTest.java
b/pinot-controller/src/test/java/org/apache/pinot/controller/api/PinotSegmentUploadAuthorizationTest.java
new file mode 100644
index 00000000000..b078fe249bb
--- /dev/null
+++
b/pinot-controller/src/test/java/org/apache/pinot/controller/api/PinotSegmentUploadAuthorizationTest.java
@@ -0,0 +1,428 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.controller.api;
+
+import com.sun.net.httpserver.HttpServer;
+import java.io.File;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.net.InetSocketAddress;
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.attribute.FileTime;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
+import javax.ws.rs.core.HttpHeaders;
+import javax.ws.rs.core.Response;
+import org.apache.commons.io.FileUtils;
+import org.apache.hc.core5.http.Header;
+import org.apache.hc.core5.http.NameValuePair;
+import org.apache.hc.core5.http.message.BasicHeader;
+import org.apache.hc.core5.http.message.BasicNameValuePair;
+import org.apache.hc.core5.net.URIBuilder;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.metadata.segment.SegmentZKMetadata;
+import org.apache.pinot.common.utils.FileUploadDownloadClient;
+import org.apache.pinot.common.utils.SimpleHttpResponse;
+import org.apache.pinot.common.utils.TarCompressionUtils;
+import org.apache.pinot.common.utils.http.HttpClient;
+import org.apache.pinot.controller.ControllerConf;
+import org.apache.pinot.controller.api.access.AccessControl;
+import org.apache.pinot.controller.api.access.AccessControlFactory;
+import org.apache.pinot.controller.api.access.AccessType;
+import org.apache.pinot.controller.helix.ControllerTest;
+import org.apache.pinot.core.auth.TargetType;
+import org.apache.pinot.segment.local.constants.SegmentUploadConstants;
+import
org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
+import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader;
+import org.apache.pinot.segment.spi.V1Constants;
+import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.apache.pinot.spi.utils.builder.TableNameBuilder;
+import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+
+@Test(groups = "stateless")
+public class PinotSegmentUploadAuthorizationTest extends ControllerTest {
+ private static final String TABLE_A = "segmentAuthTableA";
+ private static final String TABLE_B = "segmentAuthTableB";
+ private static final String DATABASE = "segmentAuthDatabase";
+ private static final String DATABASE_TABLE_NAME = "segmentAuthDatabaseTable";
+ private static final String DATABASE_TABLE = DATABASE + "." +
DATABASE_TABLE_NAME;
+ private static final String TABLE_A_WITH_TYPE =
TableNameBuilder.OFFLINE.tableNameWithType(TABLE_A);
+ private static final String TABLE_B_WITH_TYPE =
TableNameBuilder.OFFLINE.tableNameWithType(TABLE_B);
+ private static final String DATABASE_TABLE_WITH_TYPE =
+ TableNameBuilder.OFFLINE.tableNameWithType(DATABASE_TABLE);
+ private static final String SEGMENT_A = "segmentA";
+ private static final String SEGMENT_B = "segmentB";
+ private static final String SEGMENT_DATABASE = "segmentDatabase";
+
+ private File _testDir;
+ private File _segmentA;
+ private File _segmentB;
+ private File _segmentBRefresh;
+ private File _segmentDatabase;
+
+ @BeforeClass
+ public void setUp()
+ throws Exception {
+ startZk();
+ Map<String, Object> controllerConfig = getDefaultControllerConfiguration();
+ controllerConfig.put(ControllerConf.ACCESS_CONTROL_FACTORY_CLASS,
+ DestinationAccessControlFactory.class.getName());
+ startController(controllerConfig);
+ addFakeBrokerInstancesToAutoJoinHelixCluster(1, true);
+ addFakeServerInstancesToAutoJoinHelixCluster(1, true);
+
+ _testDir = new File(FileUtils.getTempDirectory(),
getClass().getSimpleName());
+ FileUtils.deleteQuietly(_testDir);
+ FileUtils.forceMkdir(_testDir);
+
+ addTable(TABLE_A);
+ addTable(TABLE_B);
+ addTable(DATABASE_TABLE);
+ _segmentA = buildSegment(TABLE_A, SEGMENT_A, SEGMENT_A, 1);
+ _segmentB = buildSegment(TABLE_B, SEGMENT_B, SEGMENT_B, 1);
+ _segmentBRefresh = buildSegment(TABLE_B, SEGMENT_B, SEGMENT_B + "Refresh",
2);
+ _segmentDatabase = buildSegment(DATABASE_TABLE, SEGMENT_DATABASE,
SEGMENT_DATABASE, 1);
+ DestinationAccessControlFactory.setAllowedTables(Set.of(TABLE_A));
+ }
+
+ @Test
+ public void testDestinationAuthorizationForSegmentUploads()
+ throws Exception {
+ URI uploadUri = URI.create(getControllerBaseApiUrl() + "/segments");
+ URI v2UploadUri = URI.create(getControllerBaseApiUrl() + "/v2/segments");
+ URI batchUploadUri = URI.create(getControllerBaseApiUrl() +
"/segments/batchUpload");
+ AtomicInteger segmentDownloadCount = new AtomicInteger();
+ HttpServer segmentServer = startSegmentServer(
+ Map.of(SEGMENT_A, _segmentA, SEGMENT_B, _segmentB, SEGMENT_DATABASE,
_segmentDatabase),
+ segmentDownloadCount);
+ File batchMetadataTar = null;
+ try (FileUploadDownloadClient client = new FileUploadDownloadClient()) {
+ batchMetadataTar = createBatchMetadataTar(_segmentDatabase,
SEGMENT_DATABASE,
+ segmentUrl(segmentServer, SEGMENT_DATABASE));
+ File batchMetadataTarForRequest = batchMetadataTar;
+ SimpleHttpResponse uploadResponse = uploadMultipart(client, uploadUri,
_segmentA, TABLE_A, List.of());
+ Assert.assertEquals(uploadResponse.getStatusCode(),
Response.Status.OK.getStatusCode());
+
+ SegmentZKMetadata segmentAMetadata =
+ _helixResourceManager.getSegmentZKMetadata(TABLE_A_WITH_TYPE,
SEGMENT_A);
+ Assert.assertNotNull(segmentAMetadata);
+ Assert.assertEquals(segmentAMetadata.getRefreshTime(), Long.MIN_VALUE);
+
+ String segmentAUri = segmentUrl(segmentServer, SEGMENT_A);
+ int downloadsBeforeInvalidRequest = segmentDownloadCount.get();
+ HttpErrorStatusException earlyHeaderMismatch =
assertHttpStatus(Response.Status.BAD_REQUEST,
+ () -> uploadJson(client, uploadUri, segmentAUri, TABLE_A,
+ List.of(new
BasicHeader(CommonConstants.Controller.TABLE_NAME_HTTP_HEADER, TABLE_B))));
+ Assert.assertTrue(
+
earlyHeaderMismatch.getMessage().contains(CommonConstants.Controller.TABLE_NAME_HTTP_HEADER),
+ earlyHeaderMismatch.getMessage());
+ Assert.assertEquals(segmentDownloadCount.get(),
downloadsBeforeInvalidRequest);
+
+ SimpleHttpResponse refreshResponse = uploadJson(client, uploadUri,
segmentAUri, TABLE_A, List.of());
+ Assert.assertEquals(refreshResponse.getStatusCode(),
Response.Status.OK.getStatusCode());
+ segmentAMetadata =
_helixResourceManager.getSegmentZKMetadata(TABLE_A_WITH_TYPE, SEGMENT_A);
+ Assert.assertTrue(segmentAMetadata.getRefreshTime() > 0);
+
+ assertHttpStatus(Response.Status.BAD_REQUEST,
+ () -> uploadMultipart(client, v2UploadUri, _segmentA, null,
List.of()));
+ assertHttpStatus(Response.Status.BAD_REQUEST,
+ () -> uploadMultipart(client, v2UploadUri, _segmentA, " ",
List.of()));
+ assertHttpStatus(Response.Status.BAD_REQUEST,
+ () -> uploadMultipart(client, v2UploadUri, _segmentA,
"otherDatabase." + TABLE_A,
+ List.of(new BasicHeader(CommonConstants.DATABASE, DATABASE))));
+
+ // V2 remains request-bound and keeps the supported promotion behavior
when segment metadata names another table.
+ Assert.assertEquals(uploadMultipart(client, v2UploadUri, _segmentB,
TABLE_A, List.of()).getStatusCode(),
+ Response.Status.OK.getStatusCode());
+
Assert.assertNotNull(_helixResourceManager.getSegmentMetadataZnRecord(TABLE_A_WITH_TYPE,
SEGMENT_B));
+ assertNoSegmentState(TABLE_B, TABLE_B_WITH_TYPE, SEGMENT_B);
+
+ assertHttpStatus(Response.Status.FORBIDDEN,
+ () -> uploadMultipart(client, v2UploadUri, _segmentA, TABLE_B,
List.of()));
+ assertNoSegmentState(TABLE_B, TABLE_B_WITH_TYPE, SEGMENT_A);
+
+ assertHttpStatus(Response.Status.BAD_REQUEST,
+ () -> uploadMultipart(client, batchUploadUri, _segmentDatabase,
TABLE_A,
+ List.of(new
BasicHeader(CommonConstants.Controller.TABLE_NAME_HTTP_HEADER, TABLE_B))));
+ assertNoSegmentState(TABLE_A, TABLE_A_WITH_TYPE, SEGMENT_DATABASE);
+
+ List<Header> batchUploadHeaders = List.of(new
BasicHeader(FileUploadDownloadClient.CustomHeaders.UPLOAD_TYPE,
+ FileUploadDownloadClient.FileUploadType.METADATA.toString()));
+ HttpErrorStatusException batchMetadataMismatch =
assertHttpStatus(Response.Status.BAD_REQUEST,
+ () -> client.uploadSegmentMetadata(batchUploadUri,
batchMetadataTarForRequest.getName(),
+ batchMetadataTarForRequest, batchUploadHeaders,
uploadParameters(TABLE_A),
+ HttpClient.DEFAULT_SOCKET_TIMEOUT_MS));
+ Assert.assertTrue(batchMetadataMismatch.getMessage().contains("segment
metadata table name"),
+ batchMetadataMismatch.getMessage());
+ assertNoSegmentState(TABLE_A, TABLE_A_WITH_TYPE, SEGMENT_DATABASE);
+
+ List<Header> blankTableHeader = List.of(
+ new BasicHeader(FileUploadDownloadClient.CustomHeaders.UPLOAD_TYPE,
+ FileUploadDownloadClient.FileUploadType.METADATA.toString()),
+ new BasicHeader(CommonConstants.Controller.TABLE_NAME_HTTP_HEADER, "
"));
+ assertHttpStatus(Response.Status.BAD_REQUEST,
+ () -> client.uploadSegmentMetadata(batchUploadUri,
batchMetadataTarForRequest.getName(),
+ batchMetadataTarForRequest, blankTableHeader,
uploadParameters(TABLE_A),
+ HttpClient.DEFAULT_SOCKET_TIMEOUT_MS));
+ assertNoSegmentState(TABLE_A, TABLE_A_WITH_TYPE, SEGMENT_DATABASE);
+
+ DestinationAccessControlFactory.setAllowedTables(Set.of(TABLE_A,
DATABASE_TABLE));
+ List<Header> databaseHeaders = List.of(
+ new BasicHeader(FileUploadDownloadClient.CustomHeaders.UPLOAD_TYPE,
+ FileUploadDownloadClient.FileUploadType.METADATA.toString()),
+ new BasicHeader(CommonConstants.DATABASE, DATABASE));
+ Assert.assertEquals(client.uploadSegmentMetadata(batchUploadUri,
batchMetadataTar.getName(), batchMetadataTar,
+ databaseHeaders, uploadParameters(DATABASE_TABLE_NAME),
HttpClient.DEFAULT_SOCKET_TIMEOUT_MS).getStatusCode(),
+ Response.Status.OK.getStatusCode());
+ Assert.assertNotNull(
+
_helixResourceManager.getSegmentMetadataZnRecord(DATABASE_TABLE_WITH_TYPE,
SEGMENT_DATABASE));
+ Assert.assertNull(_helixResourceManager.getSegmentMetadataZnRecord(
+ TableNameBuilder.OFFLINE.tableNameWithType(DATABASE_TABLE_NAME),
SEGMENT_DATABASE));
+ DestinationAccessControlFactory.setAllowedTables(Set.of(TABLE_A));
+
+ assertHttpStatus(Response.Status.FORBIDDEN, () -> uploadMultipart(
+ client, uploadUri, _segmentB, null, List.of()));
+ assertNoSegmentState(TABLE_B, TABLE_B_WITH_TYPE, SEGMENT_B);
+
+ assertHttpStatus(Response.Status.FORBIDDEN,
+ () -> uploadJson(client, uploadUri, segmentUrl(segmentServer,
SEGMENT_B), null, List.of()));
+ assertNoSegmentState(TABLE_B, TABLE_B_WITH_TYPE, SEGMENT_B);
+
+ HttpErrorStatusException mismatch =
assertHttpStatus(Response.Status.BAD_REQUEST,
+ () -> uploadMultipart(client, uploadUri, _segmentB, TABLE_A,
List.of()));
+ Assert.assertTrue(mismatch.getMessage().contains("segment metadata table
name"), mismatch.getMessage());
+ assertNoSegmentState(TABLE_B, TABLE_B_WITH_TYPE, SEGMENT_B);
+
+ long segmentARefreshTime = segmentAMetadata.getRefreshTime();
+ HttpErrorStatusException headerMismatch =
assertHttpStatus(Response.Status.BAD_REQUEST,
+ () -> uploadMultipart(client, uploadUri, _segmentA, TABLE_A,
+ List.of(new
BasicHeader(CommonConstants.Controller.TABLE_NAME_HTTP_HEADER, TABLE_B))));
+
Assert.assertTrue(headerMismatch.getMessage().contains(CommonConstants.Controller.TABLE_NAME_HTTP_HEADER),
+ headerMismatch.getMessage());
+
Assert.assertEquals(_helixResourceManager.getSegmentZKMetadata(TABLE_A_WITH_TYPE,
SEGMENT_A).getRefreshTime(),
+ segmentARefreshTime);
+
+ DestinationAccessControlFactory.setAllowedTables(Set.of(TABLE_A,
TABLE_B));
+ Assert.assertEquals(uploadMultipart(client, uploadUri, _segmentB,
TABLE_B, List.of()).getStatusCode(),
+ Response.Status.OK.getStatusCode());
+ DestinationAccessControlFactory.setAllowedTables(Set.of(TABLE_A));
+
+ ZNRecord segmentBZkBefore =
_helixResourceManager.getSegmentMetadataZnRecord(TABLE_B_WITH_TYPE, SEGMENT_B);
+ Assert.assertNotNull(segmentBZkBefore);
+ File segmentBDeepStoreFile = new File(new
File(_controllerConfig.getDataDir(), TABLE_B), SEGMENT_B);
+ Assert.assertTrue(segmentBDeepStoreFile.isFile());
+ byte[] segmentBBytesBefore =
Files.readAllBytes(segmentBDeepStoreFile.toPath());
+ FileTime segmentBModifiedBefore =
Files.getLastModifiedTime(segmentBDeepStoreFile.toPath());
+
Assert.assertFalse(Arrays.equals(Files.readAllBytes(_segmentBRefresh.toPath()),
segmentBBytesBefore));
+
+ // Omit tableName so the request passes cluster CREATE authorization and
reaches the final-table gate.
+ assertHttpStatus(Response.Status.FORBIDDEN,
+ () -> uploadMultipart(client, uploadUri, _segmentBRefresh, null,
List.of()));
+
+ ZNRecord segmentBZkAfter =
_helixResourceManager.getSegmentMetadataZnRecord(TABLE_B_WITH_TYPE, SEGMENT_B);
+ Assert.assertEquals(segmentBZkAfter.getVersion(),
segmentBZkBefore.getVersion());
+ Assert.assertEquals(segmentBZkAfter.getSimpleFields(),
segmentBZkBefore.getSimpleFields());
+ Assert.assertEquals(segmentBZkAfter.getListFields(),
segmentBZkBefore.getListFields());
+ Assert.assertEquals(segmentBZkAfter.getMapFields(),
segmentBZkBefore.getMapFields());
+
Assert.assertTrue(Arrays.equals(Files.readAllBytes(segmentBDeepStoreFile.toPath()),
segmentBBytesBefore));
+
Assert.assertEquals(Files.getLastModifiedTime(segmentBDeepStoreFile.toPath()),
segmentBModifiedBefore);
+ } finally {
+ DestinationAccessControlFactory.setAllowedTables(Set.of(TABLE_A));
+ FileUtils.deleteQuietly(batchMetadataTar);
+ segmentServer.stop(0);
+ }
+ }
+
+ private void addTable(String tableName)
+ throws Exception {
+ Schema schema = new Schema.SchemaBuilder().setSchemaName(tableName)
+ .addSingleValueDimension("value", DataType.INT).build();
+ _helixResourceManager.addSchema(schema, false, false);
+ _helixResourceManager.addTable(
+ new
TableConfigBuilder(TableType.OFFLINE).setTableName(tableName).setNumReplicas(1).build());
+ }
+
+ private File buildSegment(String tableName, String segmentName, String
artifactName, int value)
+ throws Exception {
+ Schema schema = new Schema.SchemaBuilder().setSchemaName(tableName)
+ .addSingleValueDimension("value", DataType.INT).build();
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.OFFLINE).setTableName(tableName).build();
+ SegmentGeneratorConfig generatorConfig = new
SegmentGeneratorConfig(tableConfig, schema);
+ File outputDir = new File(_testDir, artifactName + "-output");
+ generatorConfig.setOutDir(outputDir.getAbsolutePath());
+ generatorConfig.setSegmentName(segmentName);
+
+ GenericRow row = new GenericRow();
+ row.putValue("value", value);
+ SegmentIndexCreationDriverImpl driver = new
SegmentIndexCreationDriverImpl();
+ driver.init(generatorConfig, new GenericRowRecordReader(List.of(row)));
+ driver.build();
+
+ File segmentTar = new File(_testDir, artifactName +
TarCompressionUtils.TAR_GZ_FILE_EXTENSION);
+ TarCompressionUtils.createCompressedTarFile(new File(outputDir,
segmentName), segmentTar);
+ return segmentTar;
+ }
+
+ private File createBatchMetadataTar(File segmentTar, String segmentName,
String segmentUri)
+ throws IOException {
+ File metadataDir = new File(_testDir, segmentName + "-batch-metadata");
+ FileUtils.forceMkdir(metadataDir);
+ TarCompressionUtils.untarOneFile(segmentTar,
V1Constants.SEGMENT_CREATION_META,
+ new File(metadataDir, segmentName + "." +
V1Constants.SEGMENT_CREATION_META));
+ TarCompressionUtils.untarOneFile(segmentTar,
V1Constants.MetadataKeys.METADATA_FILE_NAME,
+ new File(metadataDir, segmentName + "." +
V1Constants.MetadataKeys.METADATA_FILE_NAME));
+ Files.writeString(new File(metadataDir,
SegmentUploadConstants.ALL_SEGMENTS_METADATA_FILENAME).toPath(),
+ segmentName + System.lineSeparator() + segmentUri +
System.lineSeparator(), StandardCharsets.UTF_8);
+ File metadataTar = new File(_testDir,
+ SegmentUploadConstants.ALL_SEGMENTS_METADATA_TAR_FILE_PREFIX +
segmentName
+ + TarCompressionUtils.TAR_GZ_FILE_EXTENSION);
+ TarCompressionUtils.createCompressedTarFile(metadataDir, metadataTar);
+ return metadataTar;
+ }
+
+ private static SimpleHttpResponse uploadMultipart(FileUploadDownloadClient
client, URI uploadUri, File segment,
+ String tableName, List<Header> headers)
+ throws Exception {
+ return client.uploadSegment(uploadUri, segment.getName(), segment, headers,
+ uploadParameters(tableName), HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
+ }
+
+ private static SimpleHttpResponse uploadJson(FileUploadDownloadClient
client, URI uploadUri, String segmentUri,
+ String tableName, List<Header> headers)
+ throws Exception {
+ URI requestUri = new
URIBuilder(uploadUri).addParameters(uploadParameters(tableName)).build();
+ return client.sendSegmentUri(requestUri, segmentUri, headers, null,
+ HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
+ }
+
+ private static List<NameValuePair> uploadParameters(String tableName) {
+ NameValuePair tableType =
+ new
BasicNameValuePair(FileUploadDownloadClient.QueryParameters.TABLE_TYPE,
TableType.OFFLINE.name());
+ return tableName == null ? List.of(tableType)
+ : List.of(new
BasicNameValuePair(FileUploadDownloadClient.QueryParameters.TABLE_NAME,
tableName), tableType);
+ }
+
+ private static HttpServer startSegmentServer(Map<String, File> segments,
AtomicInteger segmentDownloadCount)
+ throws IOException {
+ HttpServer server = HttpServer.create(new InetSocketAddress(0), 0);
+ for (Map.Entry<String, File> entry : segments.entrySet()) {
+ server.createContext("/" + entry.getKey(), exchange -> {
+ segmentDownloadCount.incrementAndGet();
+ long length = entry.getValue().length();
+ exchange.sendResponseHeaders(Response.Status.OK.getStatusCode(),
length);
+ try (OutputStream outputStream = exchange.getResponseBody()) {
+ Files.copy(entry.getValue().toPath(), outputStream);
+ }
+ });
+ }
+ server.start();
+ return server;
+ }
+
+ private static String segmentUrl(HttpServer server, String segmentName) {
+ return "http://localhost:"; + server.getAddress().getPort() + "/" +
segmentName;
+ }
+
+ private void assertNoSegmentState(String rawTableName, String
tableNameWithType, String segmentName) {
+
Assert.assertNull(_helixResourceManager.getSegmentMetadataZnRecord(tableNameWithType,
segmentName));
+ Assert.assertFalse(new File(new File(_controllerConfig.getDataDir(),
rawTableName), segmentName).exists());
+ }
+
+ private static HttpErrorStatusException assertHttpStatus(Response.Status
expectedStatus, ThrowingRunnable request)
+ throws Exception {
+ try {
+ request.run();
+ Assert.fail("Expected HTTP status " + expectedStatus);
+ return null;
+ } catch (HttpErrorStatusException e) {
+ Assert.assertEquals(e.getStatusCode(), expectedStatus.getStatusCode(),
e.getMessage());
+ return e;
+ }
+ }
+
+ @AfterClass
+ public void tearDown() {
+ FileUtils.deleteQuietly(_testDir);
+ stopFakeInstances();
+ stopController();
+ stopZk();
+ }
+
+ @FunctionalInterface
+ private interface ThrowingRunnable {
+ void run()
+ throws Exception;
+ }
+
+ public static class DestinationAccessControlFactory implements
AccessControlFactory {
+ private static volatile Set<String> _allowedTables = Set.of(TABLE_A);
+
+ static void setAllowedTables(Set<String> allowedTables) {
+ _allowedTables = allowedTables;
+ }
+
+ @Override
+ public AccessControl create() {
+ return new AccessControl() {
+ @Override
+ public boolean protectAnnotatedOnly() {
+ return false;
+ }
+
+ @Override
+ public boolean hasAccess(String tableName, AccessType accessType,
HttpHeaders httpHeaders,
+ String endpointUrl) {
+ return tableName == null || _allowedTables.contains(tableName);
+ }
+
+ @Override
+ public boolean hasAccess(AccessType accessType, HttpHeaders
httpHeaders, String endpointUrl) {
+ return true;
+ }
+
+ @Override
+ public boolean hasAccess(HttpHeaders httpHeaders, TargetType
targetType, String targetId, String action) {
+ return targetType == TargetType.CLUSTER ||
_allowedTables.contains(targetId);
+ }
+
+ @Override
+ public boolean hasAccess(HttpHeaders httpHeaders, TargetType
targetType) {
+ return true;
+ }
+ };
+ }
+ }
+}
diff --git
a/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResourceTest.java
b/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResourceTest.java
index 60aaa346d46..cd6107e581c 100644
---
a/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResourceTest.java
+++
b/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResourceTest.java
@@ -30,15 +30,19 @@ import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
+import javax.ws.rs.core.HttpHeaders;
+import javax.ws.rs.core.Response;
import org.apache.commons.io.FileUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.pinot.common.utils.TarCompressionUtils;
import org.apache.pinot.controller.ControllerConf;
import
org.apache.pinot.controller.api.exception.ControllerApplicationException;
import org.apache.pinot.controller.api.upload.SegmentMetadataInfo;
+import org.apache.pinot.spi.config.table.TableType;
import org.apache.pinot.spi.crypt.NoOpPinotCrypter;
import org.apache.pinot.spi.crypt.PinotCrypterFactory;
import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.utils.CommonConstants;
import org.glassfish.jersey.media.multipart.BodyPart;
import org.glassfish.jersey.media.multipart.FormDataBodyPart;
import org.glassfish.jersey.media.multipart.FormDataMultiPart;
@@ -260,6 +264,62 @@ public class PinotSegmentUploadDownloadRestletResourceTest
{
PinotSegmentUploadDownloadRestletResource.validateMultiPartForBatchSegmentUpload(bodyParts);
}
+ @Test
+ public void testResolveDestinationTableName() {
+ HttpHeaders headers = mock(HttpHeaders.class);
+
+
assertEquals(PinotSegmentUploadDownloadRestletResource.resolveDestinationTableName(
+ TABLE_NAME, TABLE_NAME + "_OFFLINE", TABLE_NAME, TableType.OFFLINE,
headers, true), TABLE_NAME);
+
assertEquals(PinotSegmentUploadDownloadRestletResource.resolveDestinationTableName(
+ null, null, TABLE_NAME, TableType.OFFLINE, headers, true), TABLE_NAME);
+
+ // V2 keeps its request-table override behavior because the request table
is already the authorized destination.
+
assertEquals(PinotSegmentUploadDownloadRestletResource.resolveDestinationTableName(
+ TABLE_NAME, null, "source_table", TableType.OFFLINE, headers, false),
TABLE_NAME);
+
+
when(headers.getHeaderString(CommonConstants.DATABASE)).thenReturn("testDatabase");
+
assertEquals(PinotSegmentUploadDownloadRestletResource.resolveDestinationTableName(
+ TABLE_NAME, TABLE_NAME, TABLE_NAME, TableType.OFFLINE, headers, true),
"testDatabase." + TABLE_NAME);
+
assertEquals(PinotSegmentUploadDownloadRestletResource.resolveDestinationTableName(
+ TABLE_NAME, null, "sourceDatabase." + TABLE_NAME, TableType.OFFLINE,
headers, false),
+ "testDatabase." + TABLE_NAME);
+ }
+
+ @Test
+ public void testRejectMissingOrMismatchedDestinationTableName() {
+ HttpHeaders headers = mock(HttpHeaders.class);
+
+ assertBadRequest(() ->
PinotSegmentUploadDownloadRestletResource.resolveDestinationTableName(
+ null, null, null, TableType.OFFLINE, headers, true), "Table name is
required");
+ assertBadRequest(() ->
PinotSegmentUploadDownloadRestletResource.resolveDestinationTableName(
+ " ", null, TABLE_NAME, TableType.OFFLINE, headers, true), "Invalid
request tableName");
+ assertBadRequest(() ->
PinotSegmentUploadDownloadRestletResource.resolveDestinationTableName(
+ TABLE_NAME, "\t", TABLE_NAME, TableType.OFFLINE, headers, true),
+ "Invalid " + CommonConstants.Controller.TABLE_NAME_HTTP_HEADER + "
header");
+ assertBadRequest(() ->
PinotSegmentUploadDownloadRestletResource.resolveDestinationTableName(
+ TABLE_NAME, null, " ", TableType.OFFLINE, headers, true), "Invalid
segment metadata table name");
+ assertBadRequest(() ->
PinotSegmentUploadDownloadRestletResource.resolveDestinationTableName(
+ TABLE_NAME, "other_table", TABLE_NAME, TableType.OFFLINE, headers,
true),
+ CommonConstants.Controller.TABLE_NAME_HTTP_HEADER + " header");
+ assertBadRequest(() ->
PinotSegmentUploadDownloadRestletResource.resolveDestinationTableName(
+ TABLE_NAME, null, "other_table", TableType.OFFLINE, headers, true),
"segment metadata table name");
+ assertBadRequest(() ->
PinotSegmentUploadDownloadRestletResource.resolveDestinationTableName(
+ TABLE_NAME + "_REALTIME", null, TABLE_NAME, TableType.OFFLINE,
headers, true),
+ "does not match table type");
+
+
when(headers.getHeaderString(CommonConstants.DATABASE)).thenReturn("databaseA");
+ assertBadRequest(() ->
PinotSegmentUploadDownloadRestletResource.resolveDestinationTableName(
+ "databaseB." + TABLE_NAME, null, "databaseB." + TABLE_NAME,
TableType.OFFLINE, headers, true),
+ "does not match database name");
+ }
+
+ private static void assertBadRequest(Runnable runnable, String
expectedMessage) {
+ ControllerApplicationException exception =
+ Assert.expectThrows(ControllerApplicationException.class, () ->
runnable.run());
+ assertEquals(exception.getResponse().getStatus(),
Response.Status.BAD_REQUEST.getStatusCode());
+ Assert.assertTrue(exception.getMessage().contains(expectedMessage),
exception.getMessage());
+ }
+
@Test
public void testCreateSegmentFileFromMultipart()
throws NoSuchMethodException, InvalidControllerConfigException,
IOException {
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/auth/FineGrainedAuthUtils.java
b/pinot-core/src/main/java/org/apache/pinot/core/auth/FineGrainedAuthUtils.java
index aca98e01fce..384bf439dca 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/auth/FineGrainedAuthUtils.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/auth/FineGrainedAuthUtils.java
@@ -100,14 +100,19 @@ public class FineGrainedAuthUtils {
// find the paramName in the path or query params
targetId = findRawTargetId(auth, uriInfo.getPathParameters(),
uriInfo.getQueryParameters());
- if (StringUtils.isEmpty(targetId)) {
+ if (StringUtils.isBlank(targetId)) {
throw new WebApplicationException(
- "Could not find paramName " + auth.paramName() + " in path or
query params of the API: "
- + uriInfo.getRequestUri(),
Response.Status.INTERNAL_SERVER_ERROR);
+ "Missing required table parameter '" + auth.paramName() + "' for
API: " + uriInfo.getRequestUri(),
+ Response.Status.BAD_REQUEST);
}
// Table name may contain type, hence get raw table name for checking
access
- targetId =
DatabaseUtils.translateTableName(TableNameBuilder.extractRawTableName(targetId),
httpHeaders);
+ try {
+ targetId =
DatabaseUtils.translateTableName(TableNameBuilder.extractRawTableName(targetId),
httpHeaders);
+ } catch (RuntimeException e) {
+ throw new WebApplicationException("Invalid table parameter '" +
auth.paramName() + "': " + e.getMessage(),
+ e, Response.Status.BAD_REQUEST);
+ }
accessDeniedMsg = "Access denied to " + auth.action() + " for table: "
+ targetId;
} else if (auth.targetType() == TargetType.CLUSTER) {
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/auth/FineGrainedAuthUtilsTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/auth/FineGrainedAuthUtilsTest.java
index 1328065b7e6..4d9ec48c527 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/auth/FineGrainedAuthUtilsTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/auth/FineGrainedAuthUtilsTest.java
@@ -19,12 +19,14 @@
package org.apache.pinot.core.auth;
import java.lang.reflect.Method;
+import java.net.URI;
import javax.ws.rs.WebApplicationException;
import javax.ws.rs.core.HttpHeaders;
import javax.ws.rs.core.MultivaluedHashMap;
import javax.ws.rs.core.MultivaluedMap;
import javax.ws.rs.core.Response;
import javax.ws.rs.core.UriInfo;
+import org.apache.pinot.spi.utils.CommonConstants;
import org.mockito.Mockito;
import org.testng.Assert;
import org.testng.annotations.Test;
@@ -102,6 +104,40 @@ public class FineGrainedAuthUtilsTest {
}
}
+ @Test
+ public void testMissingTableParameterIsBadRequest() {
+ UriInfo mockUriInfo = Mockito.mock(UriInfo.class);
+ Mockito.when(mockUriInfo.getPathParameters()).thenReturn(new
MultivaluedHashMap<>());
+ Mockito.when(mockUriInfo.getQueryParameters()).thenReturn(new
MultivaluedHashMap<>());
+
Mockito.when(mockUriInfo.getRequestUri()).thenReturn(URI.create("http://localhost/v2/segments";));
+
+ WebApplicationException exception =
Assert.expectThrows(WebApplicationException.class,
+ () ->
FineGrainedAuthUtils.validateFineGrainedAuth(getTableAnnotatedMethod(),
mockUriInfo,
+ Mockito.mock(HttpHeaders.class),
Mockito.mock(FineGrainedAccessControl.class)));
+
+ Assert.assertEquals(exception.getResponse().getStatus(),
Response.Status.BAD_REQUEST.getStatusCode());
+ Assert.assertTrue(exception.getMessage().contains("Missing required table
parameter 'tableName'"));
+ }
+
+ @Test
+ public void testInvalidTableParameterIsBadRequest() {
+ UriInfo mockUriInfo = Mockito.mock(UriInfo.class);
+ MultivaluedHashMap<String, String> queryParameters = new
MultivaluedHashMap<>();
+ queryParameters.putSingle("tableName", "databaseB.testTable");
+ Mockito.when(mockUriInfo.getPathParameters()).thenReturn(new
MultivaluedHashMap<>());
+ Mockito.when(mockUriInfo.getQueryParameters()).thenReturn(queryParameters);
+
Mockito.when(mockUriInfo.getRequestUri()).thenReturn(URI.create("http://localhost/v2/segments";));
+ HttpHeaders headers = Mockito.mock(HttpHeaders.class);
+
Mockito.when(headers.getHeaderString(CommonConstants.DATABASE)).thenReturn("databaseA");
+
+ WebApplicationException exception =
Assert.expectThrows(WebApplicationException.class,
+ () ->
FineGrainedAuthUtils.validateFineGrainedAuth(getTableAnnotatedMethod(),
mockUriInfo, headers,
+ Mockito.mock(FineGrainedAccessControl.class)));
+
+ Assert.assertEquals(exception.getResponse().getStatus(),
Response.Status.BAD_REQUEST.getStatusCode());
+ Assert.assertTrue(exception.getMessage().contains("Invalid table parameter
'tableName'"));
+ }
+
static class TestResource {
@Authorize(targetType = TargetType.CLUSTER, action = "getCluster")
void getCluster() {
@@ -110,6 +146,10 @@ public class FineGrainedAuthUtilsTest {
@Authorize(targetType = TargetType.TABLE, paramName = "tableName", action
= "getTable")
void getTable() {
}
+
+ @Authorize(targetType = TargetType.TABLE, paramName = "tableName", action
= "uploadSegment")
+ void uploadSegment() {
+ }
}
private Method getAnnotatedMethod() {
@@ -119,4 +159,12 @@ public class FineGrainedAuthUtilsTest {
throw new RuntimeException(e);
}
}
+
+ private Method getTableAnnotatedMethod() {
+ try {
+ return TestResource.class.getDeclaredMethod("uploadSegment");
+ } catch (NoSuchMethodException e) {
+ throw new RuntimeException(e);
+ }
+ }
}
diff --git
a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskGenerator.java
b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskGenerator.java
index 16506881cef..31c5359027b 100644
---
a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskGenerator.java
+++
b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskGenerator.java
@@ -144,8 +144,10 @@ public class PurgeTaskGenerator extends BaseTaskGenerator {
break;
}
configs.put(MinionConstants.DOWNLOAD_URL_KEY,
segmentZKMetadata.getDownloadUrl());
+ // Purge can reuse the original index directory, whose metadata may
name the source table. Other conversion
+ // tasks use v1 because they regenerate destination-bound metadata;
any future artifact-reuse path must use v2.
configs.put(MinionConstants.UPLOAD_URL_KEY,
- _clusterInfoAccessor.getVipUrlForLeadController(tableName) +
"/segments");
+ _clusterInfoAccessor.getVipUrlForLeadController(tableName) +
"/v2/segments");
configs.put(MinionConstants.ORIGINAL_SEGMENT_CRC_KEY,
String.valueOf(segmentZKMetadata.getCrc()));
pinotTaskConfigs.add(new PinotTaskConfig(taskType, configs));
tableNumTasks++;
diff --git
a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/refreshsegment/RefreshSegmentTaskGenerator.java
b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/refreshsegment/RefreshSegmentTaskGenerator.java
index fd4a5a72ee0..b65d8332fc6 100644
---
a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/refreshsegment/RefreshSegmentTaskGenerator.java
+++
b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/refreshsegment/RefreshSegmentTaskGenerator.java
@@ -129,8 +129,10 @@ public class RefreshSegmentTaskGenerator extends
BaseTaskGenerator {
Map<String, String> configs = new
HashMap<>(getBaseTaskConfigs(tableConfig, List.of(segmentName)));
configs.putAll(MinionTaskUtils.getPushTaskConfig(tableNameWithType,
taskConfigs, _clusterInfoAccessor));
configs.put(MinionConstants.DOWNLOAD_URL_KEY,
segmentZKMetadata.getDownloadUrl());
+ // Refresh can reuse the original index directory, whose metadata may
name the source table. Other conversion
+ // tasks use v1 because they regenerate destination-bound metadata; any
future artifact-reuse path must use v2.
configs.put(MinionConstants.UPLOAD_URL_KEY,
- _clusterInfoAccessor.getVipUrlForLeadController(tableNameWithType) +
"/segments");
+ _clusterInfoAccessor.getVipUrlForLeadController(tableNameWithType) +
"/v2/segments");
configs.put(MinionConstants.ORIGINAL_SEGMENT_CRC_KEY,
String.valueOf(segmentZKMetadata.getCrc()));
pinotTaskConfigs.add(new PinotTaskConfig(taskType, configs));
tableNumTasks++;
diff --git
a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/refreshsegment/RefreshSegmentTaskGeneratorTest.java
b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/refreshsegment/RefreshSegmentTaskGeneratorTest.java
new file mode 100644
index 00000000000..9e8be717f0f
--- /dev/null
+++
b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/refreshsegment/RefreshSegmentTaskGeneratorTest.java
@@ -0,0 +1,98 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.plugin.minion.tasks.refreshsegment;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.model.IdealState;
+import org.apache.pinot.common.metadata.segment.SegmentZKMetadata;
+import org.apache.pinot.controller.helix.core.PinotHelixResourceManager;
+import org.apache.pinot.controller.helix.core.minion.ClusterInfoAccessor;
+import org.apache.pinot.core.common.MinionConstants;
+import org.apache.pinot.core.minion.PinotTaskConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.apache.pinot.spi.utils.builder.TableNameBuilder;
+import org.apache.zookeeper.data.Stat;
+import org.testng.annotations.Test;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+
+
+/// Tests task configuration generated by [RefreshSegmentTaskGenerator].
+public class RefreshSegmentTaskGeneratorTest {
+ @Test
+ public void testGeneratesV2SegmentUploadUrl()
+ throws Exception {
+ String rawTableName = "testTable";
+ String tableNameWithType =
TableNameBuilder.OFFLINE.tableNameWithType(rawTableName);
+ String segmentName = "testSegment";
+ String controllerUrl = "http://controller:9000";;
+
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.OFFLINE).setTableName(rawTableName).build();
+ Schema schema = new Schema.SchemaBuilder()
+ .setSchemaName(rawTableName)
+ .addSingleValueDimension("value", FieldSpec.DataType.INT)
+ .build();
+
+ Stat tableStat = new Stat();
+ tableStat.setMtime(1L);
+ Stat schemaStat = new Stat();
+ schemaStat.setMtime(1L);
+
+ SegmentZKMetadata segmentMetadata = new SegmentZKMetadata(segmentName);
+ segmentMetadata.setDownloadUrl("http://deep-store/testSegment.tar.gz";);
+ segmentMetadata.setCrc(123L);
+
+ IdealState idealState = new IdealState(tableNameWithType);
+ idealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED);
+ idealState.setPartitionState(segmentName, "Server_0", "ONLINE");
+
+ PinotHelixResourceManager resourceManager =
mock(PinotHelixResourceManager.class);
+
when(resourceManager.getTableStat(tableNameWithType)).thenReturn(tableStat);
+ when(resourceManager.getTableSchema(tableNameWithType)).thenReturn(schema);
+ when(resourceManager.getSchemaStat(rawTableName)).thenReturn(schemaStat);
+
+ ClusterInfoAccessor accessor = mock(ClusterInfoAccessor.class);
+ when(accessor.getPinotHelixResourceManager()).thenReturn(resourceManager);
+
when(accessor.getTaskStates(MinionConstants.RefreshSegmentTask.TASK_TYPE)).thenReturn(Map.of());
+ when(accessor.getIdealState(tableNameWithType)).thenReturn(idealState);
+
when(accessor.getSegmentsZKMetadata(tableNameWithType)).thenReturn(List.of(segmentMetadata));
+ when(accessor.getDataDir()).thenReturn("file:///tmp");
+
when(accessor.getVipUrlForLeadController(tableNameWithType)).thenReturn(controllerUrl);
+
+ RefreshSegmentTaskGenerator generator = new RefreshSegmentTaskGenerator();
+ generator.init(accessor);
+
+ List<PinotTaskConfig> tasks = generator.generateTasks(tableConfig, new
HashMap<>());
+
+ assertEquals(tasks.size(), 1);
+ Map<String, String> configs = tasks.get(0).getConfigs();
+ assertEquals(configs.get(MinionConstants.UPLOAD_URL_KEY), controllerUrl +
"/v2/segments");
+ assertEquals(configs.get(MinionConstants.TABLE_NAME_KEY),
tableNameWithType);
+ assertEquals(configs.get(MinionConstants.SEGMENT_NAME_KEY), segmentName);
+ assertEquals(configs.get(MinionConstants.ORIGINAL_SEGMENT_CRC_KEY), "123");
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
