NIFI-10493 MiNiFi: Add C2 handler for Transfer/Debug operation

c2/c2-client-bundle/c2-client-api/src/main/java/org/apache/nifi/c2/client/api/C2Client.java

@@ -35,7 +35,7 @@ public interface C2Client {
     Optional<C2HeartbeatResponse> publishHeartbeat(C2Heartbeat heartbeat);
 
     /**
-     * Retrive the content of the new flow from the C2 Server
+     * Retrieve the content of the new flow from the C2 Server
      *
      * @param flowUpdateUrl url where the content should be downloaded from
      * @return the actual downloaded content. Will be empty if no content can be downloaded
@@ -48,4 +48,13 @@ public interface C2Client {
      * @param operationAck the acknowledge details to be sent
      */
     void acknowledgeOperation(C2OperationAck operationAck);
+
+    /**
+     * Uploads a binary bundle to C2 server
+     *
+     * @param callbackUrl url where the content should be uploaded to
+     * @param bundle      bundle content as byte array to be uploaded
+     * @return optional error message if any issues occurred
+     */
+    Optional<String> uploadBundle(String callbackUrl, byte[] bundle);
 }

c2/c2-client-bundle/c2-client-http/src/main/java/org/apache/nifi/c2/client/http/C2HttpClient.java

@@ -17,6 +17,9 @@
 
 package org.apache.nifi.c2.client.http;
 
+import static okhttp3.MultipartBody.FORM;
+import static okhttp3.RequestBody.create;
+
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.security.KeyStore;
@@ -31,9 +34,9 @@
 import javax.net.ssl.TrustManagerFactory;
 import javax.net.ssl.X509TrustManager;
 import okhttp3.MediaType;
+import okhttp3.MultipartBody;
 import okhttp3.OkHttpClient;
 import okhttp3.Request;
-import okhttp3.RequestBody;
 import okhttp3.Response;
 import okhttp3.ResponseBody;
 import okhttp3.logging.HttpLoggingInterceptor;
@@ -51,6 +54,9 @@ public class C2HttpClient implements C2Client {
 
     static final MediaType MEDIA_TYPE_APPLICATION_JSON = MediaType.parse("application/json");
     private static final Logger logger = LoggerFactory.getLogger(C2HttpClient.class);
+    private static final String MULTIPART_FORM_FILE_FIELD_NAME = "file";
+    private static final String BUNDLE_FILE_NAME = "debug.tar.gz";
+    private static final MediaType BUNDLE_MIME_TYPE = MediaType.parse("application/gzip");
 
     private final AtomicReference<OkHttpClient> httpClientReference = new AtomicReference<>();
     private final C2ClientConfig clientConfig;
@@ -123,18 +129,41 @@ public Optional<byte[]> retrieveUpdateContent(String flowUpdateUrl) {
 
     @Override
     public void acknowledgeOperation(C2OperationAck operationAck) {
-        logger.info("Acknowledging Operation [{}] C2 URL [{}]", operationAck.getOperationId(), clientConfig.getC2AckUrl());
+        logger.info("Acknowledging Operation {} to C2 server {}", operationAck.getOperationId(), clientConfig.getC2AckUrl());
         serializer.serialize(operationAck)
-            .map(operationAckBody -> RequestBody.create(operationAckBody, MEDIA_TYPE_APPLICATION_JSON))
+            .map(operationAckBody -> create(operationAckBody, MEDIA_TYPE_APPLICATION_JSON))
             .map(requestBody -> new Request.Builder().post(requestBody).url(clientConfig.getC2AckUrl()).build())
             .map(C2RequestCompression.forType(clientConfig.getC2RequestCompression())::compress)
             .ifPresent(this::sendAck);
     }
 
+    @Override
+    public Optional<String> uploadBundle(String callbackUrl, byte[] bundle) {
+        Request request = new Request.Builder()
+            .url(callbackUrl)
+            .post(new MultipartBody.Builder()
+                .setType(FORM)
+                .addFormDataPart(MULTIPART_FORM_FILE_FIELD_NAME, BUNDLE_FILE_NAME, create(bundle, BUNDLE_MIME_TYPE))
+                .build())
+            .build();
+
+        logger.info("Uploading bundle to C2 server {} with size {}", callbackUrl, bundle.length);
+        try (Response response = httpClientReference.get().newCall(request).execute()) {
+            if (!response.isSuccessful()) {
+                logger.error("Upload bundle failed to C2 server {} with status code {}", callbackUrl, response.code());
+                return Optional.of("Upload bundle failed to C2 server");
+            }
+        } catch (IOException e) {
+            logger.error("Could not upload bundle to C2 server {}", callbackUrl, e);
+            return Optional.of("Could not upload bundle to C2 server");
+        }
+        return Optional.empty();
+    }
+
     private Optional<C2HeartbeatResponse> sendHeartbeat(String heartbeat) {
         Optional<C2HeartbeatResponse> c2HeartbeatResponse = Optional.empty();
         Request request = new Request.Builder()
-            .post(RequestBody.create(heartbeat, MEDIA_TYPE_APPLICATION_JSON))
+            .post(create(heartbeat, MEDIA_TYPE_APPLICATION_JSON))
             .url(clientConfig.getC2Url())
             .build();
 
@@ -143,7 +172,7 @@ private Optional<C2HeartbeatResponse> sendHeartbeat(String heartbeat) {
         try (Response heartbeatResponse = httpClientReference.get().newCall(decoratedRequest).execute()) {
             c2HeartbeatResponse = getResponseBody(heartbeatResponse).flatMap(response -> serializer.deserialize(response, C2HeartbeatResponse.class));
         } catch (IOException ce) {
-            logger.error("Send Heartbeat failed [{}]", clientConfig.getC2Url(), ce);
+            logger.error("Send Heartbeat failed to C2 server {}", clientConfig.getC2Url(), ce);
         }
 
         return c2HeartbeatResponse;
@@ -243,10 +272,10 @@ private void assertTruststorePropertiesSet(String location, String password, Str
     private void sendAck(Request request) {
         try (Response heartbeatResponse = httpClientReference.get().newCall(request).execute()) {
             if (!heartbeatResponse.isSuccessful()) {
-                logger.warn("Acknowledgement was not successful with c2 server [{}] with status code {}", clientConfig.getC2AckUrl(), heartbeatResponse.code());
+                logger.warn("Acknowledgement was not successful with C2 server {} with status code {}", clientConfig.getC2AckUrl(), heartbeatResponse.code());
             }
         } catch (IOException e) {
-            logger.error("Could not transmit ack to c2 server [{}]", clientConfig.getC2AckUrl(), e);
+            logger.error("Could not transmit ack to C2 server {}", clientConfig.getC2AckUrl(), e);
         }
     }
 }

c2/c2-client-bundle/c2-client-service/pom.xml

@@ -43,5 +43,10 @@ limitations under the License.
             <artifactId>c2-client-http</artifactId>
             <version>1.18.0-SNAPSHOT</version>
         </dependency>
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-compress</artifactId>
+            <version>1.21</version>
+        </dependency>
     </dependencies>
 </project>

c2/c2-client-bundle/c2-client-service/src/main/java/org/apache/nifi/c2/client/service/operation/DebugOperationHandler.java

@@ -0,0 +1,193 @@
+/*
+ * 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.nifi.c2.client.service.operation;
+
+import static java.nio.file.Files.copy;
+import static java.nio.file.Files.createTempDirectory;
+import static java.nio.file.Files.lines;
+import static java.nio.file.Files.walk;
+import static java.util.Optional.empty;
+import static java.util.Optional.ofNullable;
+import static org.apache.commons.compress.utils.IOUtils.closeQuietly;
+import static org.apache.commons.lang3.StringUtils.EMPTY;
+import static org.apache.nifi.c2.protocol.api.C2OperationState.OperationState.FULLY_APPLIED;
+import static org.apache.nifi.c2.protocol.api.C2OperationState.OperationState.NOT_APPLIED;
+import static org.apache.nifi.c2.protocol.api.OperandType.DEBUG;
+import static org.apache.nifi.c2.protocol.api.OperationType.TRANSFER;
+
+import java.io.ByteArrayOutputStream;
+import java.io.File;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import java.util.function.Predicate;
+import java.util.stream.Stream;
+import org.apache.commons.compress.archivers.tar.TarArchiveEntry;
+import org.apache.commons.compress.archivers.tar.TarArchiveOutputStream;
+import org.apache.commons.compress.compressors.gzip.GzipCompressorOutputStream;
+import org.apache.nifi.c2.client.api.C2Client;
+import org.apache.nifi.c2.protocol.api.C2Operation;
+import org.apache.nifi.c2.protocol.api.C2OperationAck;
+import org.apache.nifi.c2.protocol.api.C2OperationState;
+import org.apache.nifi.c2.protocol.api.C2OperationState.OperationState;
+import org.apache.nifi.c2.protocol.api.OperandType;
+import org.apache.nifi.c2.protocol.api.OperationType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class DebugOperationHandler implements C2OperationHandler {
+
+    private static final Logger LOG = LoggerFactory.getLogger(DebugOperationHandler.class);
+
+    private static final String C2_CALLBACK_URL_NOT_FOUND = "C2 Server callback URL was not found in request";
+    private static final String SUCCESSFUL_UPLOAD = "Debug bundle was uploaded successfully";
+    private static final String UNABLE_TO_CREATE_BUNDLE = "Unable to create debug bundle";
+
+    static final String TARGET_ARG = "target";
+    static final String NEW_LINE = "\n";
+
+    private final C2Client c2Client;
+    private final List<Path> bundleFilePaths;
+    private final Predicate<String> contentFilter;
+
+    private DebugOperationHandler(C2Client c2Client, List<Path> bundleFilePaths, Predicate<String> contentFilter) {
+        this.c2Client = c2Client;
+        this.bundleFilePaths = bundleFilePaths;
+        this.contentFilter = contentFilter;
+    }
+
+    public static DebugOperationHandler create(C2Client c2Client, List<Path> bundleFilePaths, Predicate<String> contentFilter) {
+        if (c2Client == null) {
+            throw new IllegalArgumentException("C2Client should not be null");
+        }
+        if (bundleFilePaths == null || bundleFilePaths.isEmpty()) {
+            throw new IllegalArgumentException("bundleFilePaths should not be not null or empty");
+        }
+        if (contentFilter == null) {
+            throw new IllegalArgumentException("Content filter should not be null");
+        }
+
+        return new DebugOperationHandler(c2Client, bundleFilePaths, contentFilter);
+    }
+
+    @Override
+    public OperationType getOperationType() {
+        return TRANSFER;
+    }
+
+    @Override
+    public OperandType getOperandType() {
+        return DEBUG;
+    }
+
+    @Override
+    public C2OperationAck handle(C2Operation operation) {
+        String debugCallbackUrl = operation.getArgs().get(TARGET_ARG);
+        if (debugCallbackUrl == null) {
+            LOG.error("Callback URL was not found in C2 request.");
+            return operationAck(operation, operationState(NOT_APPLIED, C2_CALLBACK_URL_NOT_FOUND));
+        }
+
+        List<Path> contentFilteredFilePaths = null;
+        C2OperationState operationState;
+        try {
+            contentFilteredFilePaths = filterContent(operation.getIdentifier(), bundleFilePaths);
+            operationState = createDebugBundle(contentFilteredFilePaths)
+                .map(bundle -> c2Client.uploadBundle(debugCallbackUrl, bundle)
+                    .map(errorMessage -> operationState(NOT_APPLIED, errorMessage))
+                    .orElseGet(() -> operationState(FULLY_APPLIED, SUCCESSFUL_UPLOAD)))
+                .orElseGet(() -> operationState(NOT_APPLIED, UNABLE_TO_CREATE_BUNDLE));
+        } catch (Exception e) {
+            LOG.error("Unexpected error happened", e);
+            operationState = operationState(NOT_APPLIED, UNABLE_TO_CREATE_BUNDLE);
+        } finally {
+            ofNullable(contentFilteredFilePaths).ifPresent(this::cleanup);
+        }
+
+        LOG.debug("Returning operation ack for operation {} with state {} and details {}", operation.getIdentifier(), operationState.getState(), operationState.getDetails());
+        return operationAck(operation, operationState);
+    }
+
+    private C2OperationAck operationAck(C2Operation operation, C2OperationState state) {
+        C2OperationAck operationAck = new C2OperationAck();
+        operationAck.setOperationId(ofNullable(operation.getIdentifier()).orElse(EMPTY));
+        operationAck.setOperationState(state);
+        return operationAck;
+    }
+
+    private C2OperationState operationState(OperationState operationState, String details) {
+        C2OperationState state = new C2OperationState();
+        state.setState(operationState);
+        state.setDetails(details);
+        return state;
+    }
+
+    private List<Path> filterContent(String operationId, List<Path> bundleFilePaths) {
+        List<Path> contentFilteredFilePaths = new ArrayList<>();
+        for (Path path : bundleFilePaths) {
+            String fileName = path.getFileName().toString();
+            try (Stream<String> fileStream = lines(path)) {
+                Path tempDirectory = createTempDirectory(operationId);
+                Path tempFile = Paths.get(tempDirectory.toAbsolutePath().toString(), fileName);
+                Files.write(tempFile, (Iterable<String>) fileStream.filter(contentFilter)::iterator);
+                contentFilteredFilePaths.add(tempFile);
+            } catch (IOException e) {
+                throw new UncheckedIOException(e);
+            }
+        }
+        return contentFilteredFilePaths;
+    }
+
+    private Optional<byte[]> createDebugBundle(List<Path> filePaths) {
+        ByteArrayOutputStream byteOutputStream = new ByteArrayOutputStream();
+        try (GzipCompressorOutputStream gzipCompressorOutputStream = new GzipCompressorOutputStream(byteOutputStream);
+             TarArchiveOutputStream tarOutputStream = new TarArchiveOutputStream(gzipCompressorOutputStream)) {
+            for (Path filePath : filePaths) {
+                TarArchiveEntry tarArchiveEntry = new TarArchiveEntry(filePath.toFile(), filePath.getFileName().toString());
+                tarOutputStream.putArchiveEntry(tarArchiveEntry);
+                copy(filePath, tarOutputStream);
+                tarOutputStream.closeArchiveEntry();
+            }
+            tarOutputStream.finish();
+        } catch (Exception e) {
+            LOG.error("Error during create compressed bundle", e);
+            return empty();
+        } finally {
+            closeQuietly(byteOutputStream);
+        }
+        return Optional.of(byteOutputStream).map(ByteArrayOutputStream::toByteArray);
+    }
+
+    private void cleanup(List<Path> paths) {
+        paths.stream()
+            .findFirst()
+            .map(Path::getParent)
+            .ifPresent(basePath -> {
+                try (Stream<Path> walk = walk(basePath)) {
+                    walk.map(Path::toFile).forEach(File::delete);
+                } catch (IOException e) {
+                    LOG.warn("Unable to clean up temporary directory {}", basePath, e);
+                }
+            });
+    }
+}