fix(optimize-analyzer) (#56)

Author: NikhilShahi

backend/src/analyze-traces.ts

@@ -9,10 +9,54 @@ import { RedisClient } from "utils/redis"
 import { TRACES_QUEUE } from "~/constants"
 import { QueryRunner, Raw } from "typeorm"
 import { QueuedApiTrace } from "@common/types"
-import { isSuspectedParamater, skipAutoGeneratedMatch } from "utils"
+import {
+  endpointAddNumberParams,
+  endpointUpdateDates,
+  isSuspectedParamater,
+  skipAutoGeneratedMatch,
+} from "utils"
 import { getPathTokens } from "@common/utils"
 import { AlertType } from "@common/enums"
 
+const GET_ENDPOINT_QUERY = `
+SELECT
+  endpoint. *,
+  CASE WHEN spec."isAutoGenerated" IS NULL THEN NULL ELSE json_build_object('isAutoGenerated', spec."isAutoGenerated") END as "openapiSpec"
+FROM
+  "api_endpoint" endpoint
+  LEFT JOIN "open_api_spec" spec ON endpoint."openapiSpecName" = spec.name
+WHERE
+  $1 ~ "pathRegex"
+  AND method = $2
+  AND host = $3
+GROUP BY
+  1,
+  spec."isAutoGenerated"
+ORDER BY
+  endpoint."numberParams" ASC
+LIMIT
+  1
+`
+
+const GET_DATA_FIELDS_QUERY = `
+SELECT
+  uuid,
+  "dataClasses"::text[],
+  "falsePositives"::text[],
+  "scannerIdentified"::text[],
+  "dataType",
+  "dataTag",
+  "dataSection",
+  "createdAt",
+  "updatedAt",
+  "dataPath",
+  "apiEndpointUuid"
+FROM
+  data_field
+WHERE
+  "apiEndpointUuid" = $1
+`
+
 const getQueuedApiTrace = async (): Promise<QueuedApiTrace> => {
   try {
     const traceString = await RedisClient.popValueFromRedisList(TRACES_QUEUE)
@@ -28,7 +72,7 @@ const analyze = async (
   queryRunner: QueryRunner,
   newEndpoint?: boolean,
 ) => {
-  apiEndpoint.updateDates(trace.createdAt)
+  endpointUpdateDates(trace.createdAt, apiEndpoint)
   const dataFields = DataFieldService.findAllDataFields(trace, apiEndpoint)
   let alerts = await SpecService.findOpenApiSpecDiff(
     trace,
@@ -48,6 +92,8 @@ const analyze = async (
       AlertType.NEW_ENDPOINT,
       apiEndpoint,
     )
+    newEndpointAlert.createdAt = trace.createdAt
+    newEndpointAlert.updatedAt = trace.createdAt
     alerts = alerts?.concat(newEndpointAlert)
   }
 
@@ -140,7 +186,7 @@ const generateEndpoint = async (
     apiEndpoint.pathRegex = pathRegex
     apiEndpoint.host = trace.host
     apiEndpoint.method = trace.method
-    apiEndpoint.addNumberParams()
+    endpointAddNumberParams(apiEndpoint)
     apiEndpoint.dataFields = []
 
     try {
@@ -178,7 +224,6 @@ const generateEndpoint = async (
         }
       } else {
         console.error(`Error generating new endpoint: ${err}`)
-        await queryRunner.rollbackTransaction()
       }
     }
   }
@@ -199,21 +244,24 @@ const analyzeTraces = async (): Promise<void> => {
       const trace = await getQueuedApiTrace()
       if (trace) {
         trace.createdAt = new Date(trace.createdAt)
-        const apiEndpoint = await queryRunner.manager.findOne(ApiEndpoint, {
-          where: {
-            pathRegex: Raw(alias => `:path ~ ${alias}`, { path: trace.path }),
-            method: trace.method,
-            host: trace.host,
-          },
-          relations: { openapiSpec: true, dataFields: true },
-          order: {
-            numberParams: "ASC",
-          },
-        })
+        const apiEndpoint: ApiEndpoint = (
+          await queryRunner.query(GET_ENDPOINT_QUERY, [
+            trace.path,
+            trace.method,
+            trace.host,
+          ])
+        )?.[0]
         if (apiEndpoint && !skipAutoGeneratedMatch(apiEndpoint, trace.path)) {
+          const dataFields: DataField[] = await queryRunner.query(
+            GET_DATA_FIELDS_QUERY,
+            [apiEndpoint.uuid],
+          )
+          apiEndpoint.dataFields = dataFields
           await analyze(trace, apiEndpoint, queryRunner)
         } else {
-          await generateEndpoint(trace, queryRunner)
+          if (trace.responseStatus !== 404) {
+            await generateEndpoint(trace, queryRunner)
+          }
         }
       }
     } catch (err) {

backend/src/analyzer.ts

@@ -6,9 +6,9 @@ const main = async () => {
   const options = {
     filename: path.resolve(__dirname, "analyze-traces.js"),
   }
-  const analyzers = Array.from({ length: parseInt(process.env.NUM_WORKERS || "1") }).map(() =>
-    pool.run({}, options),
-  )
+  const analyzers = Array.from({
+    length: parseInt(process.env.NUM_WORKERS || "1"),
+  }).map(() => pool.run({}, options))
   await Promise.all(analyzers)
 }
 main()

backend/src/data-source.ts

@@ -20,6 +20,7 @@ import { runMigration } from "utils"
 import { initMigration1665782029662 } from "migrations/1665782029662-init-migration"
 import { addUniqueConstraintApiEndpoint1666678487137 } from "migrations/1666678487137-add-unique-constraint-api-endpoint"
 import { dropAnalyzedColumnFromApiTrace1666752646836 } from "migrations/1666752646836-drop-analyzed-column-from-api-trace"
+import { addIndexForDataField1666941075032 } from "migrations/1666941075032-add-index-for-data-field"
 
 export const AppDataSource: DataSource = new DataSource({
   type: "postgres",
@@ -45,11 +46,12 @@ export const AppDataSource: DataSource = new DataSource({
     initMigration1665782029662,
     addUniqueConstraintApiEndpoint1666678487137,
     dropAnalyzedColumnFromApiTrace1666752646836,
+    addIndexForDataField1666941075032,
   ],
   migrationsRun: runMigration,
   logging: false,
   extra: {
     max: 100,
-    idleTimeoutMillis: process.env.IS_ANALYZER ? 0 : 10000,
+    idleTimeoutMillis: process.env.IS_ANALYZER ? 0 : 2500,
   },
 })

backend/src/migrations/1666941075032-add-index-for-data-field.ts

@@ -0,0 +1,13 @@
+import { MigrationInterface, QueryRunner } from "typeorm"
+
+export class addIndexForDataField1666941075032 implements MigrationInterface {
+  public async up(queryRunner: QueryRunner): Promise<void> {
+    await queryRunner.query(
+      `CREATE INDEX IF NOT EXISTS "apiEndpointUuid_data_field" ON "data_field" ("apiEndpointUuid")`,
+    )
+  }
+
+  public async down(queryRunner: QueryRunner): Promise<void> {
+    await queryRunner.query(`DROP INDEX IF EXISTS "apiEndpointUuid_data_field"`)
+  }
+}

backend/src/models/data-field.ts

@@ -7,6 +7,7 @@ import {
   CreateDateColumn,
   UpdateDateColumn,
   Unique,
+  Index,
 } from "typeorm"
 import { DataClass, DataTag, DataType, DataSection } from "@common/enums"
 import { ApiEndpoint } from "models/api-endpoint"
@@ -85,6 +86,7 @@ export class DataField extends BaseEntity {
   matches: Record<DataClass, string[]>
 
   @Column({ nullable: false })
+  @Index("apiEndpointUuid_data_field")
   apiEndpointUuid: string
 
   @ManyToOne(() => ApiEndpoint, apiEndpoint => apiEndpoint.dataFields)

backend/src/services/alert/index.ts

@@ -252,7 +252,6 @@ export class AlertService {
     context?: object,
     noDuplicate?: boolean,
   ): Promise<Alert> {
-    const alertRepository = AppDataSource.getRepository(Alert)
     let alertDescription = description
     if (!alertDescription) {
       switch (alertType) {
@@ -358,6 +357,8 @@ export class AlertService {
                     trace: apiTrace,
                   }
                   newAlert.description = basicAuthDescription
+                  newAlert.createdAt = apiTrace.createdAt
+                  newAlert.updatedAt = apiTrace.createdAt
                   alerts.push(newAlert)
                 }
               }
@@ -429,6 +430,8 @@ export class AlertService {
                 newAlert.apiEndpointUuid = apiEndpointUuid
                 newAlert.context = alert.context
                 newAlert.description = alert.description
+                newAlert.createdAt = apiTrace.createdAt
+                newAlert.updatedAt = apiTrace.createdAt
                 alerts.push(newAlert)
               }
             }
@@ -475,6 +478,8 @@ export class AlertService {
             pathPointer,
             trace: apiTrace,
           }
+          newAlert.createdAt = apiTrace.createdAt
+          newAlert.updatedAt = apiTrace.createdAt
           if (!openApiSpec.minimizedSpecContext[minimizedSpecKey]) {
             let lineNumber = null
             if (openApiSpec.extension === SpecExtension.JSON) {

backend/src/services/data-field/index.ts

@@ -6,6 +6,7 @@ import { getPathTokens } from "@common/utils"
 import { ScannerService } from "services/scanner/scan"
 import { AppDataSource } from "data-source"
 import Error404NotFound from "errors/error-404-not-found"
+import { addDataClass } from "./utils"
 
 export class DataFieldService {
   static dataFields: Record<string, DataField>
@@ -67,7 +68,7 @@ export class DataFieldService {
         dataField.createdAt = this.traceCreatedAt
         dataField.updatedAt = this.traceCreatedAt
         if (dataClass) {
-          dataField.addDataClass(dataClass)
+          addDataClass(dataField, dataClass)
           dataField.dataTag = DataTag.PII
         }
         this.dataFields[existingMatch] = dataField
@@ -77,7 +78,7 @@ export class DataFieldService {
     } else {
       const existingDataField = this.dataFields[existingMatch]
       let updated = false
-      updated = existingDataField.addDataClass(dataClass)
+      updated = addDataClass(existingDataField, dataClass)
       if (updated) {
         existingDataField.dataTag = DataTag.PII
       }

backend/src/services/data-field/utils.ts

@@ -0,0 +1,33 @@
+import { DataClass } from "@common/enums"
+import { DataField } from "models"
+
+export const addDataClass = (
+  dataField: DataField,
+  dataClass: DataClass,
+): boolean => {
+  if (dataField.dataClasses === null || dataField.dataClasses === undefined) {
+    dataField.dataClasses = Array<DataClass>()
+  }
+  if (
+    dataField.falsePositives === null ||
+    dataField.falsePositives === undefined
+  ) {
+    dataField.falsePositives = Array<DataClass>()
+  }
+  if (
+    dataField.scannerIdentified === null ||
+    dataField.scannerIdentified === undefined
+  ) {
+    dataField.scannerIdentified = Array<DataClass>()
+  }
+  if (
+    dataClass === null ||
+    dataField.dataClasses.includes(dataClass) ||
+    dataField.falsePositives.includes(dataClass)
+  ) {
+    return false
+  }
+  dataField.dataClasses.push(dataClass)
+  dataField.scannerIdentified.push(dataClass)
+  return true
+}

backend/src/services/jobs/monitor-endpoint-hsts.ts

@@ -19,7 +19,8 @@ const monitorEndpointForHSTS = async (): Promise<void> => {
         order: { createdAt: "DESC" },
       })
       if (
-        !latest_trace_for_endpoint.responseHeaders.find(v =>
+        latest_trace_for_endpoint &&
+        !latest_trace_for_endpoint?.responseHeaders.find(v =>
           v.name.includes("Strict-Transport-Security"),
         )
       ) {

backend/src/utils/index.ts

@@ -138,3 +138,36 @@ export const parsedJsonNonNull = (
 
 export const inSandboxMode =
   (process.env.SANDBOX_MODE || "false").toLowerCase() == "true"
+
+export const endpointUpdateDates = (
+  traceCreatedDate: Date,
+  apiEndpoint: ApiEndpoint,
+) => {
+  if (!apiEndpoint.firstDetected) {
+    apiEndpoint.firstDetected = traceCreatedDate
+  }
+  if (!apiEndpoint.lastActive) {
+    apiEndpoint.lastActive = traceCreatedDate
+  }
+
+  if (traceCreatedDate && traceCreatedDate < apiEndpoint.firstDetected) {
+    apiEndpoint.firstDetected = traceCreatedDate
+  }
+  if (traceCreatedDate && traceCreatedDate > apiEndpoint.lastActive) {
+    apiEndpoint.lastActive = traceCreatedDate
+  }
+}
+
+export const endpointAddNumberParams = (apiEndpoint: ApiEndpoint) => {
+  if (apiEndpoint.path) {
+    const pathTokens = getPathTokens(apiEndpoint.path)
+    let numParams = 0
+    for (let i = 0; i < pathTokens.length; i++) {
+      const token = pathTokens[i]
+      if (isParameter(token)) {
+        numParams += 1
+      }
+    }
+    apiEndpoint.numberParams = numParams
+  }
+}