You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by xi...@apache.org on 2026/08/25 22:50:03 UTC
(pinot) branch master updated: Validate segment upload destinations (#19231)
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 <xi...@apache.org>
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: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org