[CDAP-21211] Authenticate Preview Runner to Preview Manager communication done via CDAP Messaging Service

cdap-app-fabric/src/main/java/io/cdap/cdap/app/preview/PreviewDataPublisher.java

@@ -30,4 +30,16 @@ public interface PreviewDataPublisher {
    * @param previewMessage preview message
    */
   void publish(EntityId entityId, PreviewMessage previewMessage);
+
+  /**
+   * Sets the identifying information for the publisher of the preview data.
+   * This is typically called once to associate the publisher with a specific preview runner.
+   * The provided information should be a serialized {@code PreviewRequestPollerInfo} object,
+   * which contains the unique identity of the runner pod. This information is used by the
+   * receiving service to authenticate the origin of the messages.
+   *
+   * @param publisherInfo a byte array containing the serialized identifying information
+   *                      for the publisher (the preview runner).
+   */
+  void setPublisherInfo(byte[] publisherInfo);
 }

cdap-app-fabric/src/main/java/io/cdap/cdap/app/preview/PreviewMessage.java

@@ -23,14 +23,14 @@
 import io.cdap.cdap.proto.id.ProgramRunId;
 
 /**
- * A container for messages in the preview topic configured by {@link
- * Constants.Preview#MESSAGING_TOPIC}. It carries the message type and the payload as {@link
- * JsonElement}.
+ * A container for messages in the preview topic configured by
+ * {@link Constants.Preview#MESSAGING_TOPIC}. It carries the message type and the payload as
+ * {@link JsonElement}.
  */
 public final class PreviewMessage {
 
   /**
-   * The message type
+   * The message type.
    */
   public enum Type {
     DATA,
@@ -41,13 +41,14 @@ public enum Type {
   private final Type type;
   private final EntityId entityId;
   private final JsonElement payload;
+  private byte[] publisherInfo;
 
   /**
    * Create an instance of message.
    *
-   * @param type type of the message
+   * @param type     type of the message
    * @param entityId program run id associated with the message
-   * @param payload the payload
+   * @param payload  the payload
    */
   public PreviewMessage(Type type, EntityId entityId, JsonElement payload) {
     this.type = type;
@@ -72,15 +73,36 @@ public EntityId getEntityId() {
   /**
    * Returns the payload by decoding the json to the given type.
    *
-   * @param gson the {@link Gson} for decoding the json element
+   * @param gson    the {@link Gson} for decoding the json element
    * @param objType the resulting object type
-   * @param <T> the resulting object type
+   * @param <T>     the resulting object type
    * @return the decode object
    */
   public <T> T getPayload(Gson gson, java.lang.reflect.Type objType) {
     return gson.fromJson(payload, objType);
   }
 
+  /**
+   * Returns the serialized information that identifies and authenticates the publisher. This is
+   * expected to be a serialized {@link PreviewRequestPollerInfo} object.
+   *
+   * @return the publisher information byte array, or {@code null} if not set.
+   */
+  public byte[] getPublisherInfo() {
+    return publisherInfo;
+  }
+
+  /**
+   * Sets the identifying information for the publisher of this message. This information is used by
+   * the receiving service to authenticate the message's origin.
+   *
+   * @param publisherInfo a byte array containing the serialized {@link PreviewRequestPollerInfo},
+   *                      which includes the runner's identity and secret token for authentication.
+   */
+  public void setPublisherInfo(byte[] publisherInfo) {
+    this.publisherInfo = publisherInfo;
+  }
+
   @Override
   public String toString() {
     return "PreviewMessage{"

cdap-app-fabric/src/main/java/io/cdap/cdap/app/preview/PreviewRequest.java

@@ -32,8 +32,9 @@
   private final ProgramId program;
   private final AppRequest<?> appRequest;
   private final Principal principal;
+  private byte[] runnerInfo;
 
   public PreviewRequest(ProgramId program, AppRequest<?> appRequest,
       @Nullable Principal principal) {
     this.program = program;
     this.appRequest = appRequest;
@@ -58,6 +59,14 @@
     return principal;
   }
 
+  public byte[] getRunnerInfo() {
+    return runnerInfo;
+  }
+
+  public void setRunnerInfo(byte[] runnerInfo) {
+    this.runnerInfo = runnerInfo;
+  }
+
   private static ProgramId getProgramIdFromRequest(ApplicationId preview, AppRequest request) {
     PreviewConfig previewConfig = request.getPreview();
     if (previewConfig == null) {

cdap-app-fabric/src/main/java/io/cdap/cdap/gateway/handlers/preview/PreviewHttpHandlerInternal.java

@@ -23,6 +23,7 @@
 import io.cdap.cdap.app.preview.PreviewManager;
 import io.cdap.cdap.app.preview.PreviewRequest;
 import io.cdap.cdap.common.conf.Constants;
+import io.cdap.cdap.internal.app.runtime.k8s.PreviewRequestPollerInfo;
 import io.cdap.http.AbstractHttpHandler;
 import io.cdap.http.HttpHandler;
 import io.cdap.http.HttpResponder;
@@ -49,7 +50,7 @@
     this.previewManager = previewManager;
   }
 
   @POST
   @Path("/requests/pull")
   public void poll(FullHttpRequest request, HttpResponder responder) {
     byte[] pollerInfo = Bytes.toBytes(request.content().nioBuffer());
@@ -57,7 +58,7 @@
 
     if (previewRequest != null) {
       LOG.debug("Send preview request {} to poller {}", previewRequest.getProgram(),
-          Bytes.toString(pollerInfo));
+          GSON.fromJson(Bytes.toString(pollerInfo), PreviewRequestPollerInfo.class));
       responder.sendString(HttpResponseStatus.OK, GSON.toJson(previewRequest));
     } else {
       responder.sendStatus(HttpResponseStatus.OK);

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/preview/DefaultPreviewManager.java

@@ -133,13 +133,13 @@
   private final Configuration previewHConf;
   private final SConfiguration previewSConf;
   private final DiscoveryServiceClient discoveryServiceClient;
   private final DatasetFramework datasetFramework;
   private final TransactionSystemClient transactionSystemClient;
   private final LevelDBTableService previewLevelDBTableService;
   private final PreviewRequestQueue previewRequestQueue;
   private final PreviewStore previewStore;
   private final PreviewRunStopper previewRunStopper;
   private final MessagingService messagingService;
   private final PreviewDataCleanupService previewDataCleanupService;
   private final MetricsCollectionService metricsCollectionService;
   private Injector previewInjector;
@@ -159,10 +159,10 @@
       @Named(PreviewConfigModule.PREVIEW_HCONF) Configuration previewHConf,
       @Named(PreviewConfigModule.PREVIEW_SCONF) SConfiguration previewSConf,
       PreviewRequestQueue previewRequestQueue, PreviewStore previewStore,
       PreviewRunStopper previewRunStopper, MessagingService messagingService,
       PreviewDataCleanupService previewDataCleanupService,
       MetricsCollectionService metricsCollectionService) {
     this.authenticationContext = authenticationContext;
     this.previewCConf = previewCConf;
     this.previewHConf = previewHConf;
     this.previewSConf = previewSConf;
@@ -350,6 +350,9 @@
             bind(PreviewRequestQueue.class).toInstance(previewRequestQueue);
             expose(PreviewRequestQueue.class);
 
+            bind(PreviewRunStopper.class).toInstance(previewRunStopper);
+            expose(PreviewRunStopper.class);
+
             bind(Store.class).to(DefaultStore.class);
             bind(OwnerStore.class).to(DefaultOwnerStore.class);
             bind(OwnerAdmin.class).to(DefaultOwnerAdmin.class);

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/preview/DefaultPreviewRunner.java

@@ -112,7 +112,7 @@
   private final MetricsCollectionService metricsCollectionService;
   private final ProgramNotificationSubscriberService programNotificationSubscriberService;
   private final ProgramStopSubscriberService programStopSubscriberService;
   private final LevelDBTableService levelDBTableService;
   private final StructuredTableAdmin structuredTableAdmin;
   private final Path previewIdDirPath;
   private final PreferencesFetcher preferencesFetcher;
@@ -132,7 +132,7 @@
       MetricsCollectionService metricsCollectionService,
       ProgramNotificationSubscriberService programNotificationSubscriberService,
       ProgramStopSubscriberService programStopSubscriberService,
       LevelDBTableService levelDBTableService,
       StructuredTableAdmin structuredTableAdmin,
       CConfiguration cConf,
       PreferencesFetcher preferencesFetcher) {
@@ -158,7 +158,9 @@
   }
 
   @Override
-  public Future<PreviewRequest> startPreview(PreviewRequest previewRequest) throws Exception {
+  public Future<PreviewRequest> startPreview(PreviewRequest previewRequest)
+      throws Exception {
+    previewDataPublisher.setPublisherInfo(previewRequest.getRunnerInfo());
     ProgramId programId = previewRequest.getProgram();
     long submitTimeMillis = RunIds.getTime(programId.getApplication(), TimeUnit.MILLISECONDS);
     previewStarted(programId);
@@ -221,7 +223,7 @@
 
       @Override
       public void init(ProgramController.State currentState, @Nullable Throwable cause) {
         switch (currentState) {
           case STARTING:
           case ALIVE:
           case STOPPING:

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/preview/MessagingPreviewDataPublisher.java

@@ -28,9 +28,9 @@
 import io.cdap.cdap.common.service.Retries;
 import io.cdap.cdap.common.service.RetryStrategies;
 import io.cdap.cdap.common.service.RetryStrategy;
+import io.cdap.cdap.messaging.client.StoreRequestBuilder;
 import io.cdap.cdap.messaging.spi.MessagingService;
 import io.cdap.cdap.messaging.spi.StoreRequest;
-import io.cdap.cdap.messaging.client.StoreRequestBuilder;
 import io.cdap.cdap.proto.id.EntityId;
 import io.cdap.cdap.proto.id.NamespaceId;
 import io.cdap.cdap.proto.id.TopicId;
@@ -46,6 +46,7 @@ public class MessagingPreviewDataPublisher implements PreviewDataPublisher {
   private final TopicId topic;
   private final MessagingService messagingService;
   private final RetryStrategy retryStrategy;
+  private byte[] publisherInfo;
 
   @Inject
   MessagingPreviewDataPublisher(CConfiguration cConf,
@@ -57,6 +58,11 @@ public class MessagingPreviewDataPublisher implements PreviewDataPublisher {
 
   @Override
   public void publish(EntityId entityId, PreviewMessage previewMessage) {
+    if (publisherInfo != null) {
+      // The publisherInfo is null in case of non-distributed environments, because there is no separate
+      // runner container. The preview runner is a thread in the same process.
+      previewMessage.setPublisherInfo(publisherInfo);
+    }
     StoreRequest request = StoreRequestBuilder.of(topic).addPayload(GSON.toJson(previewMessage))
         .build();
     try {
@@ -68,4 +74,8 @@ public void publish(EntityId entityId, PreviewMessage previewMessage) {
           e);
     }
   }
+
+  public void setPublisherInfo(byte[] publisherInfo) {
+    this.publisherInfo = publisherInfo;
+  }
 }

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/preview/PreviewDataSubscriberService.java

@@ -35,16 +35,18 @@
 import io.cdap.cdap.common.utils.ImmutablePair;
 import io.cdap.cdap.internal.app.runtime.ProgramRunners;
 import io.cdap.cdap.internal.app.store.AppMetadataStore;
-import io.cdap.cdap.messaging.spi.MessagingService;
 import io.cdap.cdap.messaging.context.MultiThreadMessagingContext;
+import io.cdap.cdap.messaging.spi.MessagingService;
 import io.cdap.cdap.messaging.subscriber.AbstractMessagingSubscriberService;
+import io.cdap.cdap.proto.BasicThrowable;
 import io.cdap.cdap.proto.codec.EntityIdTypeAdapter;
 import io.cdap.cdap.proto.id.ApplicationId;
 import io.cdap.cdap.proto.id.EntityId;
 import io.cdap.cdap.proto.id.NamespaceId;
 import io.cdap.cdap.proto.id.ProgramRunId;
 import io.cdap.cdap.spi.data.StructuredTableContext;
 import io.cdap.cdap.spi.data.transaction.TransactionRunner;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
@@ -68,6 +70,7 @@ public class PreviewDataSubscriberService extends
   private final MultiThreadMessagingContext messagingContext;
   private final TransactionRunner transactionRunner;
   private final int maxRetriesOnError;
+  private final PreviewRunStopper previewRunStopper;
   private int errorCount;
   private String erroredMessageId;
   private MetricsCollectionService metricsCollectionService;
@@ -81,7 +84,8 @@ public class PreviewDataSubscriberService extends
       @Named(PreviewConfigModule.GLOBAL_METRICS) MetricsCollectionService
           metricsCollectionService,
       PreviewStore previewStore,
-      TransactionRunner transactionRunner) {
+      TransactionRunner transactionRunner,
+      PreviewRunStopper previewRunStopper) {
     super(
         NamespaceId.SYSTEM.topic(cConf.get(Constants.Preview.MESSAGING_TOPIC)),
         cConf.getInt(Constants.Metadata.MESSAGING_FETCH_SIZE),
@@ -101,6 +105,7 @@ public class PreviewDataSubscriberService extends
     this.transactionRunner = transactionRunner;
     this.maxRetriesOnError = cConf.getInt(Constants.Metadata.MESSAGING_RETRIES_ON_CONFLICT);
     this.metricsCollectionService = metricsCollectionService;
+    this.previewRunStopper = previewRunStopper;
   }
 
   @Override
@@ -131,6 +136,9 @@ protected void processMessages(StructuredTableContext structuredTableContext,
       ImmutablePair<String, PreviewMessage> next = messages.next();
       String messageId = next.getFirst();
       PreviewMessage message = next.getSecond();
+      if (!validateMessage(message)) {
+        continue;
+      }
 
       PreviewMessageProcessor processor = processors.computeIfAbsent(message.getType(), type -> {
         switch (type) {
@@ -201,13 +209,6 @@ private final class PreviewDataProcessor implements PreviewMessageProcessor {
 
     @Override
     public void processMessage(PreviewMessage message) {
-      if (!(message.getEntityId() instanceof ApplicationId)) {
-        LOG.warn(
-            "Missing application id from the preview data information. Ignoring the message {}",
-            message);
-        return;
-      }
-
       ApplicationId applicationId = (ApplicationId) message.getEntityId();
       PreviewDataPayload payload;
       try {
@@ -230,13 +231,6 @@ private final class PreviewStatusWriter implements PreviewMessageProcessor {
 
     @Override
     public void processMessage(PreviewMessage message) {
-      if (!(message.getEntityId() instanceof ApplicationId)) {
-        LOG.warn(
-            "Missing application id from the preview status information. Ignoring the message {}",
-            message);
-        return;
-      }
-
       ApplicationId applicationId = (ApplicationId) message.getEntityId();
       PreviewStatus payload;
       try {
@@ -268,12 +262,6 @@ private final class PreviewProgramRunIdWriter implements PreviewMessageProcessor
 
     @Override
     public void processMessage(PreviewMessage message) {
-      if (!(message.getEntityId() instanceof ApplicationId)) {
-        LOG.warn("Missing application id from the preview run information. Ignoring the message {}",
-            message);
-        return;
-      }
-
       ProgramRunId payload;
       try {
         payload = message.getPayload(GSON, ProgramRunId.class);
@@ -286,4 +274,76 @@ public void processMessage(PreviewMessage message) {
       previewStore.setProgramId(payload);
     }
   }
+
+  /**
+   * Validates an incoming {@link PreviewMessage} to enforce security policies. This method performs
+   * the following checks:
+   * <ul>
+   *   <li>Ensures the message is associated with a valid {@link ApplicationId}.</li>
+   *   <li>Allows messages that do not contain publisher information to pass (for non-authenticated flows).</li>
+   *   <li>If publisher information is present, it validates that the message's {@code ApplicationId}
+   *       matches the one registered for that publisher in the {@link PreviewStore}.</li>
+   * </ul>
+   * If a mismatch is found, it is treated as a violation, and an attempt is made to
+   * terminate the offending runner.
+   *
+   * @param message the message to be validated.
+   * @return {@code true} if the message is valid, {@code false} otherwise.
+   */
+  private boolean validateMessage(PreviewMessage message) {
+    if (!(message.getEntityId() instanceof ApplicationId)) {
+      LOG.warn("Missing application id from the preview run information. Ignoring message: {}",
+          message);
+      return false;
+    }
+
+    byte[] messageRunnerInfo = message.getPublisherInfo();
+    // If there's no publisher info, the message is considered valid but unauthenticated.
+    // This allows for backward compatibility or simpler, non-secure, non-distributed environments.
+    if (messageRunnerInfo == null) {
+      return true;
+    }
+
+    ApplicationId appIdFromMessage = (ApplicationId) message.getEntityId();
+    byte[] registeredRunnerInfo = previewStore.getPreviewRequestPollerInfo(appIdFromMessage);
+
+    // Happy Path: If no runner is registered for this app yet, or if the messageRunnerInfo from the message
+    // matches the registered messageRunnerInfo, the message is valid.
+    if (Arrays.equals(registeredRunnerInfo, messageRunnerInfo)) {
+      return true;
+    }
+
+    // Failure Path: The publisher information does not match.
+    LOG.warn(
+        "Authentication failure: A message for application '{}' was received with a non-registered "
+            + "runner. Terminating the runner.", appIdFromMessage);
+    terminateRunner(messageRunnerInfo);
+
+    return false;
+  }
+
+  /**
+   * This method finds the application the suspicious runner was registered to, marks its status as
+   * KILLED, and then stops the runner's container.
+   */
+  private void terminateRunner(byte[] runnerInfo) {
+    try {
+      ApplicationId appId = previewStore.getApplicationId(runnerInfo);
+      if (appId != null) {
+        PreviewStatus status = previewStore.getPreviewStatus(appId);
+        if (status != null && !status.getStatus().isEndState()) {
+          previewStore.setPreviewStatus(appId,
+              new PreviewStatus(PreviewStatus.Status.KILLED,
+                  status.getSubmitTime(), new BasicThrowable(new IllegalStateException(
+                  "Preview run stopped due to an authentication failure."
+                      + "Please try running preview again.")), null, null));
+        }
+      }
+      previewRunStopper.stop(runnerInfo);
+    } catch (Exception e) {
+      LOG.warn(
+          "Failed to stop runner after an authentication failure. The invalid message was still rejected.",
+          e);
+    }
+  }
 }