feat: Last leg acknowledgement for IAST scanning (#87)

lib/nr-security-agent/lib/core/commonUtils.js

@@ -301,8 +301,11 @@ function removeOlderlogfiles() {
             const secondStat = fs.statSync(`${basePath}${SLASH}${b}`);
             return new Date(secondStat.mtime).getTime() - new Date(firstStat.mtime).getTime();
         });
-
-        for (let i = 2; i < sorted.length; i++) {
+        let maxLogfile = process.env.NR_CSEC_DEBUG_LOGFILE_MAX_COUNT;
+        if(isNaN(maxLogfile)){
+            maxLogfile = 2;
+        }
+        for (let i = maxLogfile; i < sorted.length; i++) {
             const fileToDelete = basePath + sorted[i];
             try {
                 fs.unlink(fileToDelete, (err) => {

lib/nr-security-agent/lib/core/connections/websocket/response/IASTUtils.js

@@ -3,20 +3,53 @@
  * SPDX-License-Identifier: New Relic Pre-Release
  */
 const { Agent } = require('../../../agent');
-let completedRequestIds = new Set();
+let pendingRequestIds = new Set();
+let completedRequestsMap = new Map();
 const PolicyManager = require('../../../Policy');
 let batchSize;
 const batchthreshold = 160;
-function getCompletedRequestIds() {
-    return [...completedRequestIds];
+function getPendingRequestIds() {
+    return [...pendingRequestIds];
 }
 
-function addRequestId(id) {
-    completedRequestIds.add(id);
+function addPendingRequestId(id) {
+    pendingRequestIds.add(id);
 }
 
 function removeRequestId(id) {
-    completedRequestIds.delete(id);
+    completedRequestsMap.delete(id);
+}
+
+function removePendingRequestId(id) {
+    pendingRequestIds.delete(id);
+}
+
+// for last leg
+function completedRequestsMapInit(id) {
+    completedRequestsMap.set(id, []);
+}
+
+function addCompletedRequests(id, eventId) {
+    let data = completedRequestsMap.get(id);
+    data.push(eventId);
+    completedRequestsMap.set(id, data);
+}
+
+function getCompletedRequestsMap() {
+    return completedRequestsMap;
+}
+
+function clearCompletedRequestMap(){
+    completedRequestsMap.clear();
+}
+
+function clearPendingRequestIdSet(){
+    pendingRequestIds.clear();
+}
+
+function IASTCleanup(){
+    clearCompletedRequestMap();
+    clearPendingRequestIdSet();
 }
 
 function generateIASTDataRequest() {
@@ -27,24 +60,34 @@ function generateIASTDataRequest() {
             batchSize = policyInstance.data.vulnerabilityScan.iastScan.probing.batchSize;
         }
     }
-    if (getCompletedRequestIds().length >= batchSize || batchSize < batchthreshold) {
+    if (getCompletedRequestsMap().size >= batchSize || batchSize < batchthreshold || getPendingRequestIds().length < batchSize) {
         batchSize = (2 * batchSize);
     }
     let object = {};
     object['jsonName'] = 'iast-data-request';
     object['applicationUUID'] = Agent.getAgent().applicationInfo.applicationUUID;
     object["batchSize"] = batchSize ? batchSize : 300;
-    object['completedRequestIds'] = getCompletedRequestIds();
+    object['pendingRequestIds'] = getPendingRequestIds();
+    object['completedRequests'] = Object.fromEntries(getCompletedRequestsMap());
     return object;
-
 }
 
 
 module.exports = {
-    getCompletedRequestIds,
-    addRequestId,
-    removeRequestId,
+    getPendingRequestIds,
+    addPendingRequestId,
+    removePendingRequestId,
     generateIASTDataRequest,
+
+    // for last leg
+    addCompletedRequests,
+    completedRequestsMapInit,
+    getCompletedRequestsMap,
+    removeRequestId,
+    IASTCleanup
+
+
+
 }
 
 

lib/nr-security-agent/lib/core/connections/websocket/response/fuzz-request-handler.js

@@ -9,13 +9,12 @@ const stringify = require('fast-safe-stringify');
 const restClient = require('../../restclient');
 const { Agent } = require('../../../agent');
 const { FuzzFailEvent } = require('../../../FuzzFailEvent');
-const { CSEC_HOME_TMP, COLON, NR_CSEC_FUZZ_REQUEST_ID, SLASH } = require('../../../sec-agent-constants');
+const { CSEC_HOME_TMP, COLON, NR_CSEC_FUZZ_REQUEST_ID, SLASH, CSEC_SEP } = require('../../../sec-agent-constants');
 const LOCALHOST = 'localhost';
 const COLON_SLASH_SLASH = '://';
 const http = require('http');
 const https = require('https');
 const IASTUtil = require('./IASTUtils');
-const CSEC_SEP = ':IAST:';
 
 require('dns').setDefaultResultOrder('ipv4first')
 
@@ -46,10 +45,12 @@ function startIASTSchedular() {
     iastIntervalConst = setInterval(() => {
         let data = IASTUtil.generateIASTDataRequest();
         let currentTime = Date.now();
-        let completedListSize = IASTUtil.getCompletedRequestIds().length;
+        let completedListSize = IASTUtil.getCompletedRequestsMap().size;
+        let pendingListSize = IASTUtil.getPendingRequestIds().length;
         let timeDiffInSeconds = ((currentTime - lastFuzzEventTime) / 1000);
         logger.trace("Time difference since last fuzz request:", timeDiffInSeconds);
         logger.trace("Completed requests so far:", completedListSize);
+        logger.trace("Pending requests so far:", pendingListSize);
 
         if (timeDiffInSeconds > 5 && additionalCoolDownTime == 0) {
             Agent.getAgent().client.dispatcher(data);
@@ -83,10 +84,7 @@ function handler(json) {
     try {
         fuzzRequest = JSON.parse(rawFuzzRequest);
         fuzzRequest['id'] = json.id;
-        let completedList = IASTUtil.getCompletedRequestIds();
-        if (completedList.includes(json.id)) {
-            return;
-        }
+        IASTUtil.addPendingRequestId(json.id);
     } catch (error) {
         logger.error('Parsing exeception in fuzz request: ', error);
     }
@@ -130,6 +128,10 @@ function handleFuzzRequest(fuzzDetails) {
                 }
             }
 
+            config.headers['nr-csec-parent-id'] = fuzzRequest.id;
+
+            IASTUtil.completedRequestsMapInit(fuzzRequest.id);
+
             if(fuzzRequest.headers && fuzzRequest.headers[NR_CSEC_FUZZ_REQUEST_ID]){
                 logScannedApiId(fuzzRequest.headers[NR_CSEC_FUZZ_REQUEST_ID], fuzzRequest.requestURI)
             }
@@ -161,7 +163,7 @@ function handleFuzzRequest(fuzzDetails) {
 function handleFuzzResponse(response, fuzzDetails) {
     const { rawFuzzRequest, fuzzRequest } = fuzzDetails;
     if (response) {
-        IASTUtil.addRequestId(fuzzRequest.id);
+        IASTUtil.removePendingRequestId(fuzzRequest.id);
         response.then(() => {
             logger.info('Fuzz success: ' + rawFuzzRequest);
         }).catch(() => {

lib/nr-security-agent/lib/core/connections/websocket/websocket.js

@@ -9,16 +9,18 @@ const WebSocket = require('ws');
 const RingBuffer = require('ringbufferjs');
 const stringify = require('fast-safe-stringify');
 const { EventEmitter } = require('events');
-
+const API = require('../../../../../nr-security-api');
+const NRAgent = API.getNRAgent();
 const { Agent } = require('../../agent');
 const logs = require('../../logging');
 const hc = require('../../health-check');
 const fs = require('fs');
 const { promisify } = require('../../sec-util');
 const ResponseHandler = require('./response');
-const { EMPTY_APPLICATION_UUID, LOG_MESSAGES } = require('../../sec-agent-constants');
+const { EMPTY_APPLICATION_UUID, LOG_MESSAGES, CSEC_SEP, NR_CSEC_FUZZ_REQUEST_ID, EXITEVENT,JSON_NAME, RASP } = require('../../sec-agent-constants');
 const statusUtils = require('../../statusUtils');
 const commonUtils = require('../../commonUtils');
+const IASTUtil = require('../websocket/response/IASTUtils');
 
 const BUSY = 'busy';
 const IDLE = 'idle';
@@ -155,14 +157,15 @@ SecWebSocket.prototype.obeyReconnect = function obeyReconnect() {
     }
 
     initTimeOut = setTimeout(() => {
+        IASTUtil.IASTCleanup();
         logger.debug("Terminating ws instance and reconnecting");
         try {
             instance.terminate();
         } catch (err) {
             logger.debug("Error while terminating ws instance");
         }
         this.init();
-    }, 15000)
+    }, (5 + (Math.floor(Math.random() * 11))) * 1000)
 
 }
 
@@ -200,6 +203,18 @@ SecWebSocket.prototype.dispatch = async function dispatch(event) {
     if (event.jsonName === APPLICATION_INFO_JSON_NAME) {
         APP_INFO = event;
     }
+    if (event.parentId && event.httpRequest.headers[NR_CSEC_FUZZ_REQUEST_ID]) {
+        try {
+            let apiId = event.httpRequest.headers[NR_CSEC_FUZZ_REQUEST_ID].split(CSEC_SEP)[0]
+            if (apiId == event.apiId) {
+                IASTUtil.addCompletedRequests(event.parentId, event.id);
+            }
+        } catch (error) {
+            logger.debug("Error while mapping completedRequests:", error);
+        }
+
+    }
+
     RegExp.prototype.toJSON = RegExp.prototype.toString;
     const eventStr = stringify(event);
     if (!this.instance ||
@@ -212,6 +227,18 @@ SecWebSocket.prototype.dispatch = async function dispatch(event) {
     try {
         await promisify(this.instance, this.instance.send)(eventStr, { mask: true }, () => {
             logger.debug(LOG_MESSAGES.EVENT_SENT + eventStr);
+            if (event.jsonName == JSON_NAME.EVENT) {
+                if (NRAgent.config.security.mode == RASP) {
+                    hc.getInstance().raspEventStats.sent++;
+                }
+                else {
+                    hc.getInstance().iastEventStats.sent++;
+                }
+            }
+            if (event.jsonName == EXITEVENT) {
+                hc.getInstance().exitEventStats.sent++;
+            }
+
         });
     } catch (error) {
         logger.debug(LOG_MESSAGES.ERROR_WHILE_SEND_EVENT, eventStr, error);
@@ -291,6 +318,12 @@ const handleDispatchFailure = (self, event, eventStr) => {
     // so feeding an old applicationinfo JSON would be incorrect
     if (eventStr && event.jsonName !== APPLICATION_INFO_JSON_NAME) {
         hc.getInstance().registerEventDrop();
+        if (NRAgent.config.security.mode == "RASP") {
+            hc.getInstance().raspEventStats.errorCount++;
+        }
+        else {
+            hc.getInstance().iastEventStats.errorCount++;
+        }
     }
 };
 

lib/nr-security-agent/lib/core/event-stats.js

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2023 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: New Relic Pre-Release
+ */
+
+let eventStats = undefined;
+function EventStats() {
+    this.processed = 0;
+    this.sent = 0;
+    this.rejected = 0;
+    this.errorCount = 0;
+}
+EventStats.prototype.constructor = EventStats;
+
+/**
+ * Returns the current instanceof EventStats
+ * 
+ * @returns {EventStats} instance
+ */
+function getInstance() {
+        return new EventStats();
+}
+
+
+module.exports = {
+    getInstance
+};

lib/nr-security-agent/lib/core/event.js

@@ -99,6 +99,7 @@ function SecEvent(fileName, funcName, lineNumber, argument, eid, applicationInfo
     }
     if (httpRequest && httpRequest.headers[NR_CSEC_FUZZ_REQUEST_ID]) {
         this.isIASTRequest = true;
+        this.parentId = httpRequest.headers['nr-csec-parent-id'];
     }
     else {
         this.isIASTRequest = false;

lib/nr-security-agent/lib/core/health-check.js

@@ -8,6 +8,7 @@
 const { Agent } = require('./agent');
 const { JSON_NAME } = require('./sec-agent-constants');
 const { BasicInfo } = require('./event');
+const EventStats = require('./event-stats');
 
 const NODE_JS_LITERAL = 'Node.js:';
 
@@ -39,9 +40,10 @@ module.exports.getInstance = function () {
  *
  * @param {ApplicationInfo} info
  */
-function HC (info) {
+function HC(info) {
     BasicInfo.call(this);
 
+
     this.applicationUUID = info.applicationUUID;
     this.eventType = 'sec_health_check_lc';
     if (NRAgent && NRAgent.config) {
@@ -58,7 +60,9 @@ function HC (info) {
     this.buildNumber = info.buildNumber;
     this.jsonName = JSON_NAME.HC;
     this.pid = info.pid;
-
+    this.iastEventStats = EventStats.getInstance();
+    this.raspEventStats = EventStats.getInstance();
+    this.exitEventStats = EventStats.getInstance();
 
 
     /**
@@ -130,6 +134,45 @@ function HC (info) {
         this.httpRequestCount = 0;
     };
 
+    /**
+     * resets the IAST EventStats.
+     */
+    this.resetIASTEventStats = function () {
+        this.iastEventStats.processed = 0;
+        this.iastEventStats.sent = 0;
+        this.iastEventStats.rejected = 0;
+        this.iastEventStats.errorCount = 0;
+    };
+
+    /**
+    * resets the RASP EventStats.
+    */
+    this.resetRASPEventStats = function () {
+        this.raspEventStats.processed = 0;
+        this.raspEventStats.sent = 0;
+        this.raspEventStats.rejected = 0;
+        this.raspEventStats.errorCount = 0;
+    };
+
+    /**
+     * resets the EXIT EventStats.
+     */
+    this.resetEXITEventStats = function () {
+        this.exitEventStats.processed = 0;
+        this.exitEventStats.sent = 0;
+        this.exitEventStats.rejected = 0;
+        this.exitEventStats.errorCount = 0;
+    };
+
+    /**
+    * resets the EventStats.
+    */
+    this.resetEventStats = function () {
+        this.resetIASTEventStats();
+        this.resetRASPEventStats();
+        this.resetEXITEventStats();
+    };
+
 
     /**
      * sets gerate hc timestamp.

lib/nr-security-agent/lib/core/sec-agent-constants.js

@@ -105,6 +105,9 @@ module.exports = {
     DOUBLE_DOLLAR: '$$',
     LOADING_MODULE: 'Loading module ',
     EXITEVENT: 'exit-event',
+    CSEC_SEP: ':IAST:',
+    RASP: 'RASP',
+    IAST: 'IAST',
 
     LOG_MESSAGES: {
         LOADED_CSEC_ENVS_MSG: '[STEP-1][COMPLETE][env] Environment Information Gathering Done.',