fix(GMS): Adding Rest.li Validation in GMS

gms/impl/src/main/java/com/linkedin/metadata/resources/ResourceUtils.java

@@ -0,0 +1,36 @@
+package com.linkedin.metadata.resources;
+
+import com.linkedin.common.urn.UrnValidator;
+import com.linkedin.data.schema.validation.ValidateDataAgainstSchema;
+import com.linkedin.data.schema.validation.ValidationOptions;
+import com.linkedin.data.schema.validation.ValidationResult;
+import com.linkedin.data.template.RecordTemplate;
+import com.linkedin.restli.common.HttpStatus;
+import com.linkedin.restli.server.RestLiServiceException;
+
+
+public class ResourceUtils {
+
+  private static final ValidationOptions DEFAULT_VALIDATION_OPTIONS = new ValidationOptions();
+  private static final UrnValidator URN_VALIDATOR = new UrnValidator();
+
+  /**
+   * Validates a {@link RecordTemplate} and throws {@link com.linkedin.restli.server.RestLiServiceException}
+   * if validation fails.
+   *
+   * @param record record to be validated.
+   * @param status the status code to return to the client on failure.
+   */
+  public static void validateRecord(RecordTemplate record, HttpStatus status) {
+    final ValidationResult result = ValidateDataAgainstSchema.validate(
+        record,
+        DEFAULT_VALIDATION_OPTIONS,
+        URN_VALIDATOR);
+    if (!result.isValid()) {
+      throw new RestLiServiceException(status, result.getMessages().toString());
+    }
+  }
+
+  private ResourceUtils() { }
+
+}

gms/impl/src/main/java/com/linkedin/metadata/resources/entity/AspectResource.java

@@ -5,6 +5,7 @@
 import com.linkedin.metadata.entity.EntityService;
 import com.linkedin.metadata.restli.RestliUtils;
 import com.linkedin.parseq.Task;
+import com.linkedin.restli.common.HttpStatus;
 import com.linkedin.restli.server.annotations.Optional;
 import com.linkedin.restli.server.annotations.QueryParam;
 import com.linkedin.restli.server.annotations.RestLiCollection;
@@ -18,6 +19,9 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import static com.linkedin.metadata.resources.ResourceUtils.*;
+
+
 /**
  * Single unified resource for fetching, updating, searching, & browsing DataHub entities
  */
@@ -46,9 +50,10 @@ public Task<VersionedAspect> get(
       final VersionedAspect aspect = _entityService.getVersionedAspect(urn, aspectName, version);
       if (aspect == null) {
         throw RestliUtils.resourceNotFoundException();
+      } else {
+        validateRecord(aspect, HttpStatus.S_500_INTERNAL_SERVER_ERROR);
       }
       return aspect;
     });
   }
-
 }

gms/impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java

@@ -16,12 +16,15 @@
 import com.linkedin.metadata.search.SearchService;
 import com.linkedin.metadata.search.utils.BrowsePathUtils;
 import com.linkedin.parseq.Task;
+import com.linkedin.restli.common.HttpStatus;
+import com.linkedin.restli.common.validation.RestLiDataValidator;
 import com.linkedin.restli.server.annotations.Action;
 import com.linkedin.restli.server.annotations.ActionParam;
 import com.linkedin.restli.server.annotations.Optional;
 import com.linkedin.restli.server.annotations.QueryParam;
 import com.linkedin.restli.server.annotations.RestLiCollection;
 import com.linkedin.restli.server.annotations.RestMethod;
+import com.linkedin.restli.server.annotations.ValidatorParam;
 import com.linkedin.restli.server.resources.CollectionResourceTaskTemplate;
 import java.net.URISyntaxException;
 import java.time.Clock;
@@ -39,6 +42,7 @@
 import org.slf4j.LoggerFactory;
 
 import static com.linkedin.metadata.PegasusUtils.urnToEntityName;
+import static com.linkedin.metadata.resources.ResourceUtils.*;
 import static com.linkedin.metadata.restli.RestliConstants.ACTION_AUTOCOMPLETE;
 import static com.linkedin.metadata.restli.RestliConstants.ACTION_BROWSE;
 import static com.linkedin.metadata.restli.RestliConstants.ACTION_GET_BROWSE_PATHS;
@@ -95,6 +99,8 @@ public Task<Entity> get(@Nonnull String urnStr, @QueryParam(PARAM_ASPECTS) @Opti
       final Entity entity = _entityService.getEntity(urn, projectedAspects);
       if (entity == null) {
         throw RestliUtils.resourceNotFoundException();
+      } else {
+        validateRecord(entity, HttpStatus.S_500_INTERNAL_SERVER_ERROR);
       }
       return entity;
     });
@@ -116,13 +122,17 @@ public Task<Map<String, Entity>> batchGet(
       return _entityService.getEntities(urns, projectedAspects)
           .entrySet()
           .stream()
+          .peek(entry -> validateRecord(entry.getValue(), HttpStatus.S_500_INTERNAL_SERVER_ERROR))
           .collect(Collectors.toMap(entry -> entry.getKey().toString(), Map.Entry::getValue));
     });
   }
 
   @Action(name = ACTION_INGEST)
   @Nonnull
-  public Task<Void> ingest(@ActionParam(PARAM_ENTITY) @Nonnull Entity entity) throws URISyntaxException {
+  public Task<Void> ingest(@ActionParam(PARAM_ENTITY) @Nonnull Entity entity, @ValidatorParam RestLiDataValidator validator) throws URISyntaxException {
+
+    validateRecord(entity, HttpStatus.S_422_UNPROCESSABLE_ENTITY);
+
     final Set<String> projectedAspects = new HashSet<>(Arrays.asList("browsePaths"));
     final RecordTemplate snapshotRecord = RecordUtils.getSelectedRecordTemplateFromUnion(entity.getValue());
     final Urn urn = com.linkedin.metadata.dao.utils.ModelUtils.getUrnFromSnapshot(snapshotRecord);
@@ -144,6 +154,10 @@ public Task<Void> ingest(@ActionParam(PARAM_ENTITY) @Nonnull Entity entity) thro
   @Action(name = ACTION_BATCH_INGEST)
   @Nonnull
   public Task<Void> batchIngest(@ActionParam(PARAM_ENTITIES) @Nonnull Entity[] entities) throws URISyntaxException {
+    for (Entity entity : entities) {
+      validateRecord(entity, HttpStatus.S_422_UNPROCESSABLE_ENTITY);
+    }
+
     final AuditStamp auditStamp =
         new AuditStamp().setTime(_clock.millis()).setActor(Urn.createFromString(DEFAULT_ACTOR));
     return RestliUtils.toTask(() -> {

li-utils/src/main/javaPegasus/com/linkedin/common/urn/UrnValidator.java

@@ -0,0 +1,30 @@
+package com.linkedin.common.urn;
+
+import com.linkedin.data.message.Message;
+import com.linkedin.data.schema.DataSchema;
+import com.linkedin.data.schema.NamedDataSchema;
+import com.linkedin.data.schema.validator.Validator;
+import com.linkedin.data.schema.validator.ValidatorContext;
+import java.net.URISyntaxException;
+
+
+/**
+ * Rest.li Validator responsible for ensuring that {@link Urn} objects are well-formed.
+ *
+ * Note that this validator does not validate the integrity of strongly typed urns,
+ * or validate Urn objects against their associated key aspect.
+ */
+public class UrnValidator implements Validator {
+  @Override
+  public void validate(ValidatorContext context) {
+    if (DataSchema.Type.TYPEREF.equals(context.dataElement().getSchema().getType())
+        && ((NamedDataSchema) context.dataElement().getSchema()).getName().endsWith("Urn")) {
+      try {
+        Urn.createFromString((String) context.dataElement().getValue());
+      } catch (URISyntaxException e) {
+        context.addResult(new Message(context.dataElement().path(), "\"Provided urn %s\" is invalid", context.dataElement().getValue()));
+        context.setHasFix(false);
+      }
+    }
+  }
+}