JobS3LogLocationValidator implements AdmissionValidator<JobDescriptor> { @Override public Mono<Set<ValidationError>> validate(JobDescriptor jobDescriptor) { if (!enabledSupplier.get()) { metrics.incrementValidationSkipped(ValidatorMetrics.REASON_DISABLED); return Mono.just(Collections.emptySet()); } LogStorageInfos.S3Bucket s3BucketInfo = LogStorageInfos.findCustomS3Bucket(jobDescriptor).orElse(null); String customPrefix = LogStorageInfos.findCustomPathPrefix(jobDescriptor).orElse(null); if (s3BucketInfo == null && customPrefix == null) { metrics.incrementValidationSkipped(ValidatorMetrics.REASON_NOT_APPLICABLE); return Mono.just(Collections.emptySet()); } String bucketName = s3BucketInfo == null ? defaultBucket : s3BucketInfo.getBucketName(); String pathPrefix = LogStorageInfos.buildPathPrefix( customPrefix == null ? defaultPathPrefix : LogStorageInfos.buildPathPrefix(customPrefix, defaultPathPrefix), VALIDATOR_PATH ); String iamRole = jobDescriptor.getContainer().getSecurityProfile().getIamRole(); if (StringExt.isEmpty(iamRole)) { metrics.incrementValidationError(bucketName, REASON_ACCESS_DENIED); return Mono.just(Collections.singleton(new ValidationError("iamRole", "IAM role not set"))); } iamRole = iamRoleArnResolver.apply(iamRole); Mono<Set<ValidationError>> action = validationClient.validateS3BucketAccess( ComputeValidator.S3BucketAccessValidationRequest.newBuilder() .setBucket(bucketName) .setBucketPrefix(pathPrefix) .setIamRole(iamRole) .build() ).map(result -> { if (result.getResultCase() == ComputeValidator.S3BucketAccessValidationResponse.ResultCase.FAILURES) { List<ValidationFailure> failures = result.getFailures().getFailuresList(); if (!failures.isEmpty()) { metrics.incrementValidationError(bucketName, REASON_ACCESS_DENIED); return toValidationError(failures); } } metrics.incrementValidationSuccess(bucketName); return Collections.emptySet(); }); return action.retry(RETRY_COUNT) .onErrorMap(error -> { logger.warn("S3 validation failure: {}", error.getMessage()); logger.debug("Stack trace", error); metrics.incrementValidationError(bucketName, error.getClass().getSimpleName()); return new IllegalArgumentException(String.format("S3 bucket validation error: bucket=%s, pathPrefix=%s, error=%s", bucketName, pathPrefix, error.getMessage() ), error); }); } JobS3LogLocationValidator(ReactorValidationServiceClient validationClient,
                                     String defaultBucket,
                                     String defaultPathPrefix,
                                     Function<String, String> iamRoleArnResolver,
                                     Supplier<ValidationError.Type> validationErrorTypeProvider,
                                     Supplier<Boolean> enabledSupplier,
                                     TitusRuntime titusRuntime); @Override Mono<Set<ValidationError>> validate(JobDescriptor jobDescriptor); @Override ValidationError.Type getErrorType();  }
@Test public void testNoCustomBucket() { when(validationClient.validateS3BucketAccess(any())).thenReturn( Mono.just(S3BucketAccessValidationResponse.getDefaultInstance()) ); Set<ValidationError> result = validator.validate(JOB_WITH_DEFAULT_BUCKET).block(); assertThat(result).isEmpty(); verify(validationClient, times(0)).validateS3BucketAccess(any()); }
@Test public void testValidCustomBucket() { when(validationClient.validateS3BucketAccess(any())).thenReturn( Mono.just(S3BucketAccessValidationResponse.getDefaultInstance()) ); Set<ValidationError> result = validator.validate(JOB_WITH_CUSTOM_BUCKET).block(); assertThat(result).isEmpty(); verify(validationClient, times(1)).validateS3BucketAccess(any()); }
@Test public void testInValidCustomBucket() { when(validationClient.validateS3BucketAccess(any())).thenReturn( Mono.just(S3BucketAccessValidationResponse.newBuilder() .setFailures(ComputeValidator.ValidationFailures.newBuilder() .addFailures(ComputeValidator.ValidationFailure.getDefaultInstance()) .build() ) .build() ) ); Set<ValidationError> result = validator.validate(JOB_WITH_CUSTOM_BUCKET).block(); assertThat(result).hasSize(1); verify(validationClient, times(1)).validateS3BucketAccess(any()); }