Add access control for Pinot server segment download api.

pinot-common/src/main/java/org/apache/pinot/common/utils/CommonConstants.java

@@ -256,6 +256,9 @@ public static class SegmentCompletionProtocol {
 
     public static final String DEFAULT_METRICS_PREFIX = "pinot.server.";
     public static final boolean DEFAULT_METRICS_GLOBAL_ENABLED = false;
+    public static final String ACCESS_CONTROL_FACTORY_CLASS = "pinot.server.admin.access.control.factory.class";
+    public static final String DEFAULT_ACCESS_CONTROL_FACTORY_CLASS =
+        "org.apache.pinot.server.api.access.AllowAllAccessFactory";
   }
 
   public static class Controller {

pinot-server/src/main/java/org/apache/pinot/server/api/access/AccessControl.java

@@ -0,0 +1,38 @@
+/**
+ * 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.server.api.access;
+
+import javax.ws.rs.core.HttpHeaders;
+import org.apache.pinot.spi.annotations.InterfaceAudience;
+import org.apache.pinot.spi.annotations.InterfaceStability;
+
+
+@InterfaceAudience.Public
+@InterfaceStability.Stable
+public interface AccessControl {
+
+  /**
+   * Return whether the client has data access to the given table.
+   *
+   * @param httpHeaders Http headers
+   * @param tableName Name of the table to be accessed
+   * @return Whether the client has data access to the table
+   */
+  boolean hasDataAccess(HttpHeaders httpHeaders, String tableName);
+}

pinot-server/src/main/java/org/apache/pinot/server/api/access/AccessControlFactory.java

@@ -0,0 +1,30 @@
+/**
+ * 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.server.api.access;
+
+import org.apache.pinot.spi.annotations.InterfaceAudience;
+import org.apache.pinot.spi.annotations.InterfaceStability;
+
+
+@InterfaceAudience.Public
+@InterfaceStability.Stable
+public interface AccessControlFactory {
+
+  AccessControl create();
+}

pinot-server/src/main/java/org/apache/pinot/server/api/access/AllowAllAccessFactory.java

@@ -0,0 +1,37 @@
+/**
+ * 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.server.api.access;
+
+import javax.ws.rs.core.HttpHeaders;
+
+
+public class AllowAllAccessFactory implements AccessControlFactory {
+  private static final AccessControl ALLOW_ALL_ACCESS = new AccessControl() {
+    @Override
+    public boolean hasDataAccess(HttpHeaders httpHeaders, String tableName) {
+      return true;
+    }
+  };
+
+  @Override
+  public AccessControl create() {
+    return ALLOW_ALL_ACCESS;
+  }
+}
+

pinot-server/src/main/java/org/apache/pinot/server/api/resources/TablesResource.java

@@ -55,6 +55,8 @@
 import org.apache.pinot.core.data.manager.SegmentDataManager;
 import org.apache.pinot.core.data.manager.TableDataManager;
 import org.apache.pinot.core.segment.index.metadata.SegmentMetadataImpl;
+import org.apache.pinot.server.api.access.AccessControl;
+import org.apache.pinot.server.api.access.AccessControlFactory;
 import org.apache.pinot.server.starter.ServerInstance;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -69,6 +71,9 @@ public class TablesResource {
   @Inject
   ServerInstance serverInstance;
 
+  @Inject
+  private AccessControlFactory _accessControlFactory;
+
   @GET
   @Path("/tables")
   @Produces(MediaType.APPLICATION_JSON)
@@ -197,6 +202,19 @@ public Response downloadSegment(
       @Context HttpHeaders httpHeaders)
       throws Exception {
     LOGGER.info("Received a request to download segment {} for table {}", segmentName, tableNameWithType);
+    // Validate data access
+    boolean hasDataAccess;
+    try {
+      AccessControl accessControl = _accessControlFactory.create();
+      hasDataAccess = accessControl.hasDataAccess(httpHeaders, tableNameWithType);
+    } catch (Exception e) {
+      throw new WebApplicationException(
+          "Caught exception while validating access to table: " + tableNameWithType, Response.Status.INTERNAL_SERVER_ERROR);
+    }
+    if (!hasDataAccess) {
+      throw new WebApplicationException("No data access to table: " + tableNameWithType, Response.Status.FORBIDDEN);
+    }
+
     TableDataManager tableDataManager = checkGetTableDataManager(tableNameWithType);
     SegmentDataManager segmentDataManager = tableDataManager.acquireSegment(segmentName);
     if (segmentDataManager == null) {

pinot-server/src/main/java/org/apache/pinot/server/starter/helix/AdminApiApplication.java

@@ -26,6 +26,7 @@
 import javax.ws.rs.container.ContainerRequestContext;
 import javax.ws.rs.container.ContainerResponseContext;
 import javax.ws.rs.container.ContainerResponseFilter;
+import org.apache.pinot.server.api.access.AccessControlFactory;
 import org.apache.pinot.server.starter.ServerInstance;
 import org.glassfish.grizzly.http.server.CLStaticHttpHandler;
 import org.glassfish.grizzly.http.server.HttpServer;
@@ -41,18 +42,21 @@ public class AdminApiApplication extends ResourceConfig {
   private static final Logger LOGGER = LoggerFactory.getLogger(AdminApiApplication.class);
 
   private final ServerInstance serverInstance;
+  private final AccessControlFactory accessControlFactory;
   private URI baseUri;
   private boolean started = false;
   private HttpServer httpServer;
   public static final String RESOURCE_PACKAGE = "org.apache.pinot.server.api.resources";
 
-  public AdminApiApplication(ServerInstance instance) {
+  public AdminApiApplication(ServerInstance instance, AccessControlFactory accessControlFactory) {
     this.serverInstance = instance;
+    this.accessControlFactory = accessControlFactory;
     packages(RESOURCE_PACKAGE);
     register(new AbstractBinder() {
       @Override
       protected void configure() {
         bind(serverInstance).to(ServerInstance.class);
+        bind(accessControlFactory).to(AccessControlFactory.class);
       }
     });
 

pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixServerStarter.java

@@ -59,11 +59,13 @@
 import org.apache.pinot.core.data.manager.InstanceDataManager;
 import org.apache.pinot.core.realtime.impl.invertedindex.RealtimeLuceneIndexRefreshState;
 import org.apache.pinot.core.segment.memory.PinotDataBuffer;
+import org.apache.pinot.server.api.access.AccessControlFactory;
 import org.apache.pinot.server.conf.ServerConf;
 import org.apache.pinot.server.realtime.ControllerLeaderLocator;
 import org.apache.pinot.server.realtime.ServerSegmentCompletionProtocolHandler;
 import org.apache.pinot.server.starter.ServerInstance;
 import org.apache.pinot.spi.filesystem.PinotFSFactory;
+import org.apache.pinot.spi.plugin.PluginManager;
 import org.apache.pinot.spi.utils.builder.TableNameBuilder;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -153,8 +155,20 @@ public HelixServerStarter(String helixClusterName, String zkAddress, Configurati
     updateInstanceConfigIfNeeded(host, port);
 
     // Start restlet server for admin API endpoint
+    String accessControlFactoryClass =
+        _serverConf.getString(ACCESS_CONTROL_FACTORY_CLASS, DEFAULT_ACCESS_CONTROL_FACTORY_CLASS);
+    LOGGER.info("Using class: {} as the AccessControlFactory", accessControlFactoryClass);
+    final AccessControlFactory accessControlFactory;
+    try {
+      accessControlFactory = PluginManager.get().createInstance(accessControlFactoryClass);
+    } catch (Exception e) {
+      throw new RuntimeException(
+          "Caught exception while creating new AccessControlFactory instance in class " + this.getClass()
+              .getSimpleName(), e);
+    }
+
     int adminApiPort = _serverConf.getInt(CONFIG_OF_ADMIN_API_PORT, DEFAULT_ADMIN_API_PORT);
-    _adminApiApplication = new AdminApiApplication(_serverInstance);
+    _adminApiApplication = new AdminApiApplication(_serverInstance, accessControlFactory);
     _adminApiApplication.start(adminApiPort);
     setAdminApiPort(adminApiPort);
 

pinot-server/src/test/java/org/apache/pinot/server/api/AccessControlTest.java

@@ -0,0 +1,87 @@
+/**
+ * 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.server.api;
+
+import java.io.File;
+import java.util.HashMap;
+import java.util.Map;
+import javax.ws.rs.client.ClientBuilder;
+import javax.ws.rs.client.WebTarget;
+import javax.ws.rs.core.Response;
+import org.apache.commons.io.FileUtils;
+import org.apache.pinot.common.utils.CommonConstants;
+import org.apache.pinot.core.data.manager.TableDataManager;
+import org.apache.pinot.server.api.access.AccessControl;
+import org.apache.pinot.server.api.access.AccessControlFactory;
+import org.apache.pinot.server.starter.ServerInstance;
+import org.apache.pinot.server.starter.helix.AdminApiApplication;
+import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import static org.mockito.Mockito.mock;
+
+
+public class AccessControlTest {
+  private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "AccessControlTest");
+  protected static final String TABLE_NAME = "testTable";
+
+  private final Map<String, TableDataManager> _tableDataManagerMap = new HashMap<>();
+  private AdminApiApplication _adminApiApplication;
+  protected WebTarget _webTarget;
+
+  @BeforeClass
+  public void setUp()
+      throws Exception {
+    FileUtils.deleteQuietly(INDEX_DIR);
+    Assert.assertTrue(INDEX_DIR.mkdirs());
+
+    // Mock the instance data manager
+    // Mock the server instance
+    ServerInstance serverInstance = mock(ServerInstance.class);
+
+    _adminApiApplication = new AdminApiApplication(serverInstance, new DenyAllAccessFactory());
+    _adminApiApplication.start(CommonConstants.Server.DEFAULT_ADMIN_API_PORT);
+    _webTarget = ClientBuilder.newClient().target(_adminApiApplication.getBaseUri());
+  }
+
+  @AfterClass
+  public void tearDown() {
+    _adminApiApplication.stop();
+    FileUtils.deleteQuietly(INDEX_DIR);
+  }
+
+  @Test
+  public void testAccessDenied() {
+    String segmentPath = "/segments/testTable_REALTIME/segment_name";
+    // Download any segment will be denied.
+    Response response = _webTarget.path(segmentPath).request().get(Response.class);
+    Assert.assertNotEquals(response.getStatus(), Response.Status.FORBIDDEN);
+  }
+
+  public static class DenyAllAccessFactory implements AccessControlFactory {
+    private static final AccessControl DENY_ALL_ACCESS = (httpHeaders, tableName) -> false;
+
+    @Override
+    public AccessControl create() {
+      return DENY_ALL_ACCESS;
+    }
+  }
+}

pinot-server/src/test/java/org/apache/pinot/server/api/BaseResourceTest.java

@@ -26,7 +26,6 @@
 import java.util.Map;
 import javax.ws.rs.client.ClientBuilder;
 import javax.ws.rs.client.WebTarget;
-import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.io.FileUtils;
 import org.apache.helix.store.zk.ZkHelixPropertyStore;
 import org.apache.pinot.common.metrics.ServerMetrics;
@@ -42,6 +41,7 @@
 import org.apache.pinot.core.segment.creator.SegmentIndexCreationDriver;
 import org.apache.pinot.core.segment.creator.impl.SegmentIndexCreationDriverImpl;
 import org.apache.pinot.segments.v1.creator.SegmentTestUtils;
+import org.apache.pinot.server.api.access.AllowAllAccessFactory;
 import org.apache.pinot.server.starter.ServerInstance;
 import org.apache.pinot.server.starter.helix.AdminApiApplication;
 import org.apache.pinot.spi.utils.builder.TableNameBuilder;
@@ -86,7 +86,6 @@ public void setUp()
     ServerInstance serverInstance = mock(ServerInstance.class);
     when(serverInstance.getInstanceDataManager()).thenReturn(instanceDataManager);
     when(serverInstance.getInstanceDataManager().getSegmentFileDirectory()).thenReturn(FileUtils.getTempDirectoryPath());
-
     // Add the default tables and segments.
     String realtimeTableName = TableNameBuilder.REALTIME.tableNameWithType(TABLE_NAME);
     String offlineTableName = TableNameBuilder.OFFLINE.tableNameWithType(TABLE_NAME);
@@ -96,7 +95,7 @@ public void setUp()
     setUpSegment(realtimeTableName, "default", _realtimeIndexSegments);
     setUpSegment(offlineTableName, "default", _offlineIndexSegments);
 
-    _adminApiApplication = new AdminApiApplication(serverInstance);
+    _adminApiApplication = new AdminApiApplication(serverInstance, new AllowAllAccessFactory());
     _adminApiApplication.start(CommonConstants.Server.DEFAULT_ADMIN_API_PORT);
     _webTarget = ClientBuilder.newClient().target(_adminApiApplication.getBaseUri());
   }