io.camunda.document.store.aws.AwsDocumentStore Maven / Gradle / Ivy

Go to download
/*
 * Copyright Camunda Services GmbH and/or licensed to Camunda Services GmbH under
 * one or more contributor license agreements. See the NOTICE file distributed
 * with this work for additional information regarding copyright ownership.
 * Licensed under the Camunda License 1.0. You may not use this file
 * except in compliance with the Camunda License 1.0.
 */
package io.camunda.document.store.aws;

import io.camunda.document.api.DocumentContent;
import io.camunda.document.api.DocumentCreationRequest;
import io.camunda.document.api.DocumentError;
import io.camunda.document.api.DocumentError.DocumentAlreadyExists;
import io.camunda.document.api.DocumentError.DocumentNotFound;
import io.camunda.document.api.DocumentError.InvalidInput;
import io.camunda.document.api.DocumentError.UnknownDocumentError;
import io.camunda.document.api.DocumentLink;
import io.camunda.document.api.DocumentMetadataModel;
import io.camunda.document.api.DocumentReference;
import io.camunda.document.api.DocumentStore;
import io.camunda.zeebe.util.Either;
import java.time.Duration;
import java.time.Instant;
import java.time.OffsetDateTime;
import java.time.ZoneId;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import software.amazon.awssdk.core.ResponseInputStream;
import software.amazon.awssdk.core.sync.RequestBody;
import software.amazon.awssdk.http.HttpStatusCode;
import software.amazon.awssdk.services.s3.S3Client;
import software.amazon.awssdk.services.s3.model.DeleteObjectRequest;
import software.amazon.awssdk.services.s3.model.GetObjectRequest;
import software.amazon.awssdk.services.s3.model.GetObjectResponse;
import software.amazon.awssdk.services.s3.model.HeadObjectRequest;
import software.amazon.awssdk.services.s3.model.HeadObjectResponse;
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
import software.amazon.awssdk.services.s3.model.S3Exception;
import software.amazon.awssdk.services.s3.model.Tag;
import software.amazon.awssdk.services.s3.model.Tagging;
import software.amazon.awssdk.services.s3.presigner.S3Presigner;
import software.amazon.awssdk.services.s3.presigner.model.GetObjectPresignRequest;
import software.amazon.awssdk.services.s3.presigner.model.PresignedGetObjectRequest;

public class AwsDocumentStore implements DocumentStore {

  private static final Tag NO_AUTO_DELETE_TAG =
      Tag.builder().key("NoAutoDelete").value("true").build();

  private static final String METADATA_PROCESS_DEFINITION_ID = "camunda.processDefinitionId";
  private static final String METADATA_PROCESS_INSTANCE_KEY = "camunda.processInstanceKey";

  private final String bucketName;
  private final S3Client client;
  private final ExecutorService executor;
  private final S3Presigner preSigner;
  private final Long defaultTTL;
  private final String bucketPath;

  public AwsDocumentStore(
      final String bucketName,
      final Long defaultTTL,
      final String bucketPath,
      final ExecutorService executor) {
    this(bucketName, defaultTTL, bucketPath, S3Client.create(), executor, S3Presigner.create());
  }

  public AwsDocumentStore(
      final String bucketName,
      final Long defaultTTL,
      final String bucketPath,
      final S3Client client,
      final ExecutorService executor,
      final S3Presigner preSigner) {
    this.bucketName = bucketName;
    this.defaultTTL = defaultTTL;
    this.bucketPath = bucketPath;
    this.client = client;
    this.executor = executor;
    this.preSigner = preSigner;
  }

  @Override
  public CompletableFuture> createDocument(
      final DocumentCreationRequest request) {
    return CompletableFuture.supplyAsync(() -> createDocumentInternal(request), executor);
  }

  @Override
  public CompletableFuture> getDocument(
      final String documentId) {
    return CompletableFuture.supplyAsync(() -> getDocumentInternal(documentId), executor);
  }

  @Override
  public CompletableFuture> deleteDocument(final String documentId) {
    return CompletableFuture.supplyAsync(() -> deleteDocumentInternal(documentId), executor);
  }

  @Override
  public CompletableFuture> createLink(
      final String documentId, final long durationInMillis) {
    return CompletableFuture.supplyAsync(
        () -> linkDocumentInternal(documentId, durationInMillis), executor);
  }

  private Either createDocumentInternal(
      final DocumentCreationRequest request) {
    try {
      final String documentId =
          Objects.requireNonNullElse(request.documentId(), UUID.randomUUID().toString());

      final HeadObjectResponse documentInfo = getDocumentInfo(documentId);
      if (documentInfo != null) {
        return Either.left(new DocumentAlreadyExists(documentId));
      }

      return uploadDocument(request, documentId);
    } catch (final Exception e) {
      return Either.left(new UnknownDocumentError(e));
    }
  }

  private Either getDocumentInternal(final String documentId) {
    try {
      final GetObjectRequest getObjectRequest =
          GetObjectRequest.builder().key(resolveKey(documentId)).bucket(bucketName).build();

      final HeadObjectResponse documentInfo = getDocumentInfo(documentId);
      if (documentInfo != null && isDocumentExpired(documentInfo.metadata(), documentId)) {
        return Either.left(new DocumentNotFound(documentId));
      }

      final ResponseInputStream responseResponseInputStream =
          client.getObject(getObjectRequest);

      final String contentType =
          Optional.ofNullable(documentInfo).map(HeadObjectResponse::contentType).orElse(null);

      return Either.right(new DocumentContent(responseResponseInputStream, contentType));
    } catch (final Exception e) {
      return Either.left(getDocumentError(documentId, e));
    }
  }

  private Either deleteDocumentInternal(final String documentId) {
    try {
      client.deleteObject(
          DeleteObjectRequest.builder().bucket(bucketName).key(resolveKey(documentId)).build());

      return Either.right(null);
    } catch (final Exception e) {
      return Either.left(getDocumentError(documentId, e));
    }
  }

  private Either linkDocumentInternal(
      final String documentId, final long durationInMillis) {
    try {
      if (durationInMillis <= 0) {
        return Either.left(new InvalidInput("Duration must be greater than 0"));
      }

      final HeadObjectResponse documentInfo = getDocumentInfo(documentId);
      if (documentInfo == null || isDocumentExpired(documentInfo.metadata(), documentId)) {
        return Either.left(new DocumentNotFound(documentId));
      }

      final GetObjectRequest objectRequest =
          GetObjectRequest.builder().bucket(bucketName).key(resolveKey(documentId)).build();

      final GetObjectPresignRequest preSignRequest =
          GetObjectPresignRequest.builder()
              .signatureDuration(Duration.ofMillis(durationInMillis))
              .getObjectRequest(objectRequest)
              .build();

      final PresignedGetObjectRequest preSignedRequest = preSigner.presignGetObject(preSignRequest);
      final Instant expiration = Instant.now().plusMillis(durationInMillis);

      return Either.right(
          new DocumentLink(
              preSignedRequest.url().toString(),
              OffsetDateTime.ofInstant(expiration, ZoneId.systemDefault())));
    } catch (final Exception e) {
      return Either.left(getDocumentError(documentId, e));
    }
  }

  private HeadObjectResponse getDocumentInfo(final String documentId) {
    try {
      final HeadObjectRequest headObjectRequest =
          HeadObjectRequest.builder().bucket(bucketName).key(resolveKey(documentId)).build();

      return client.headObject(headObjectRequest);
    } catch (final S3Exception e) {
      if (e.statusCode() == HttpStatusCode.NOT_FOUND) {
        return null;
      }
      throw e;
    }
  }

  private boolean isDocumentExpired(final Map metadata, final String documentId) {
    if (metadata != null) {
      final String expiresAt = metadata.get("expires-at");

      if (expiresAt != null && OffsetDateTime.parse(expiresAt).isBefore(OffsetDateTime.now())) {
        deleteDocumentInternal(documentId);
        return true;
      }
    }
    return false;
  }

  private Either uploadDocument(
      final DocumentCreationRequest request, final String documentId) {
    final String fileName = resolveFileName(request.metadata(), documentId);
    final PutObjectRequest putObjectRequest =
        PutObjectRequest.builder()
            .key(resolveKey(documentId))
            .bucket(bucketName)
            .metadata(toS3MetaData(request.metadata(), fileName))
            .tagging(generateExpiryTag(request.metadata().expiresAt()))
            .build();

    client.putObject(
        putObjectRequest,
        RequestBody.fromInputStream(request.contentInputStream(), request.metadata().size()));

    final var updatedMetadata =
        new DocumentMetadataModel(
            request.metadata().contentType(),
            resolveFileName(request.metadata(), documentId),
            request.metadata().expiresAt(),
            request.metadata().size(),
            request.metadata().processDefinitionId(),
            request.metadata().processInstanceKey(),
            request.metadata().customProperties());
    return Either.right(new DocumentReference(documentId, updatedMetadata));
  }

  private Map toS3MetaData(
      final DocumentMetadataModel metadata, final String fileName) {
    if (metadata == null) {
      return Collections.emptyMap();
    }

    final Map metadataMap = new HashMap<>();

    putIfPresent("content-type", metadata.contentType(), metadataMap);
    putIfPresent("size", metadata.size(), metadataMap);
    putIfPresent("filename", fileName, metadataMap);
    putIfPresent("expires-at", metadata.expiresAt(), metadataMap);

    if (metadata.customProperties() != null) {
      metadata
          .customProperties()
          .forEach((key, value) -> metadataMap.put(key, String.valueOf(value)));
    }

    putIfPresent(METADATA_PROCESS_DEFINITION_ID, metadata.processDefinitionId(), metadataMap);
    putIfPresent(METADATA_PROCESS_INSTANCE_KEY, metadata.processInstanceKey(), metadataMap);

    return metadataMap;
  }

  private  void putIfPresent(
      final String key, final T value, final Map metadataMap) {
    if (value != null) {
      metadataMap.put(key, value.toString());
    }
  }

  private static DocumentError getDocumentError(final String documentId, final Exception e) {
    if (e instanceof final S3Exception s3Exception
        && s3Exception.statusCode() == HttpStatusCode.NOT_FOUND) {
      return new DocumentNotFound(documentId);
    }
    return new UnknownDocumentError(e);
  }

  private Tagging generateExpiryTag(final OffsetDateTime expiryDate) {
    final boolean isExpiryDateBeyondBucketTTL =
        expiryDate != null
            && defaultTTL != null
            && expiryDate.isAfter(OffsetDateTime.now().plus(Duration.ofDays(defaultTTL)));

    return Tagging.builder()
        .tagSet(
            isExpiryDateBeyondBucketTTL
                ? Collections.singletonList(NO_AUTO_DELETE_TAG)
                : Collections.emptyList())
        .build();
  }

  private String resolveKey(final String documentId) {
    return bucketPath + documentId;
  }

  private String resolveFileName(
      final DocumentMetadataModel documentMetadata, final String documentId) {
    return documentMetadata.fileName() != null ? documentMetadata.fileName() : documentId;
  }
}