Skip to content

Commit

Permalink
feat: added bloomfilter changes
Browse files Browse the repository at this point in the history
  • Loading branch information
@TangoBeeAkto
TangoBeeAkto committed Sep 20, 2024
1 parent 6e3bd39 commit a118642
Show file tree
Hide file tree
Showing 16 changed files with 640 additions and 9 deletions.
45 changes: 39 additions & 6 deletions apps/api-runtime/src/main/java/com/akto/runtime/APICatalogSync.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,13 @@
import java.util.concurrent.ConcurrentHashMap;
import java.util.regex.Pattern;

import com.akto.DaoInit;
import com.akto.dao.*;
import com.akto.dao.context.Context;
import com.akto.dao.filter.MergedUrlsDao;
import com.akto.dto.*;
import com.akto.dto.billing.SyncLimit;
import com.akto.dto.dependency_flow.DependencyFlow;
import com.akto.dto.filter.MergedUrls;
import com.akto.dto.traffic.Key;
import com.akto.dto.traffic.SampleData;
import com.akto.dto.traffic.TrafficInfo;
Expand All @@ -21,6 +22,7 @@
import com.akto.dto.usage.MetricTypes;
import com.akto.log.LoggerMaker;
import com.akto.log.LoggerMaker.LogDb;
import com.akto.util.filter.DictionaryFilter;
import com.akto.runtime.merge.MergeOnHostOnly;
import com.akto.runtime.policies.AktoPolicyNew;
import com.akto.task.Cluster;
Expand All @@ -32,19 +34,16 @@
import com.google.common.hash.BloomFilter;
import com.google.common.hash.Funnels;
import com.mongodb.BasicDBObject;
import com.mongodb.ConnectionString;
import com.mongodb.bulk.BulkWriteResult;
import com.mongodb.client.model.*;
import com.mongodb.client.result.UpdateResult;
import org.apache.commons.lang3.math.NumberUtils;
import org.bson.Document;
import org.bson.conversions.Bson;
import org.bson.json.JsonParseException;
import org.bson.types.ObjectId;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.Map.Entry;
import java.util.concurrent.ConcurrentHashMap;
import java.util.regex.Pattern;
Expand All @@ -67,6 +66,8 @@ public class APICatalogSync {
public static boolean mergeAsyncOutside = true;
public BloomFilter<CharSequence> existingAPIsInDb = BloomFilter.create(Funnels.stringFunnel(Charsets.UTF_8), 1_000_000, 0.001 );

public static Set<MergedUrls> mergedUrls;

public APICatalogSync(String userIdentifier,int thresh, boolean fetchAllSTI) {
this(userIdentifier, thresh, fetchAllSTI, true);
}
Expand All @@ -79,6 +80,7 @@ public APICatalogSync(String userIdentifier, int thresh, boolean fetchAllSTI, bo
this.delta = new HashMap<>();
this.sensitiveParamInfoBooleanMap = new HashMap<>();
this.aktoPolicyNew = new AktoPolicyNew();
mergedUrls = new HashSet<>();
if (buildFromDb) {
buildFromDB(false, fetchAllSTI);
}
Expand Down Expand Up @@ -712,6 +714,7 @@ public static URLTemplate tryParamteresingUrl(URLStatic newUrl){
int start = newUrl.getUrl().startsWith("http") ? 3 : 0;
for(int i = start; i < tokens.length; i ++) {
String tempToken = tokens[i];
if(DictionaryFilter.isEnglishWord(tempToken)) continue;

if (NumberUtils.isParsable(tempToken)) {
newTypes[i] = isNumber(tempToken) ? SuperType.INTEGER : SuperType.FLOAT;
Expand Down Expand Up @@ -741,7 +744,21 @@ public static URLTemplate tryParamteresingUrl(URLStatic newUrl){
}

if (allNull) return null;
return new URLTemplate(tokens, newTypes, newUrl.getMethod());

URLTemplate urlTemplate = new URLTemplate(tokens, newTypes, newUrl.getMethod());

try {
for(MergedUrls mergedUrl : mergedUrls) {
if(mergedUrl.getUrl().equals(urlTemplate.getTemplateString()) &&
mergedUrl.getMethod().equals(urlTemplate.getMethod().name())) {
return null;
}
}
} catch(Exception e) {
loggerMaker.errorAndAddToDb("Error while creating a new URL object: " + e.getMessage(), LogDb.RUNTIME);
}

return urlTemplate;
}


Expand All @@ -763,6 +780,7 @@ public static URLTemplate tryMergeUrls(URLStatic dbUrl, URLStatic newUrl) {
for(int i = 0; i < newTokens.length; i ++) {
String tempToken = newTokens[i];
String dbToken = dbTokens[i];
if (DictionaryFilter.isEnglishWord(tempToken) || DictionaryFilter.isEnglishWord(dbToken)) continue;

int minCount = dbUrl.getUrl().startsWith("http") && newUrl.getUrl().startsWith("http") ? 3 : 0;
if (tempToken.equalsIgnoreCase(dbToken) || i < minCount) {
Expand Down Expand Up @@ -796,7 +814,20 @@ public static URLTemplate tryMergeUrls(URLStatic dbUrl, URLStatic newUrl) {
if (allNull) return null;

if (templatizedStrTokens <= 1) {
return new URLTemplate(newTokens, newTypes, newUrl.getMethod());
URLTemplate urlTemplate = new URLTemplate(newTokens, newTypes, newUrl.getMethod());

try {
for(MergedUrls mergedUrl : mergedUrls) {
if(mergedUrl.getUrl().equals(urlTemplate.getTemplateString()) &&
mergedUrl.getMethod().equals(urlTemplate.getMethod().name())) {
return null;
}
}
} catch(Exception e) {
loggerMaker.errorAndAddToDb("Error while creating a new URL object: " + e.getMessage(), LogDb.RUNTIME);
}

return urlTemplate;
}

return null;
Expand Down Expand Up @@ -1549,6 +1580,8 @@ public void buildFromDB(boolean calcDiff, boolean fetchAllSTI) {
loggerMaker.errorAndAddToDb("Error while filling urls in apiCollection: " + e.getMessage(), LogDb.RUNTIME);
}

mergedUrls = MergedUrlsDao.instance.getMergedUrls();

loggerMaker.infoAndAddToDb("Building from db completed", LogDb.RUNTIME);
aktoPolicyNew.buildFromDb(fetchAllSTI);
}
Expand Down
3 changes: 3 additions & 0 deletions apps/api-runtime/src/main/java/com/akto/runtime/Main.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import com.akto.log.LoggerMaker;
import com.akto.log.LoggerMaker.LogDb;
import com.akto.parsers.HttpCallParser;
import com.akto.util.filter.DictionaryFilter;
import com.akto.util.AccountTask;
import com.akto.util.DashboardMode;
import com.google.gson.Gson;
Expand Down Expand Up @@ -175,6 +176,8 @@ public static void main(String[] args) {
}
int maxPollRecordsConfig = Integer.parseInt(System.getenv("AKTO_KAFKA_MAX_POLL_RECORDS_CONFIG"));

DictionaryFilter.readDictionaryBinary();

if (topicName == null) topicName = "akto.api.logs";

DaoInit.init(new ConnectionString(mongoURI));
Expand Down
4 changes: 3 additions & 1 deletion apps/api-runtime/src/test/java/com/akto/parsers/TestDBSync.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import com.akto.runtime.APICatalogSync;
import com.akto.runtime.Main;
import com.akto.runtime.URLAggregator;
import com.akto.util.filter.DictionaryFilter;
import com.mongodb.BasicDBObject;
import com.mongodb.client.model.Filters;

Expand All @@ -41,6 +42,7 @@ public class TestDBSync extends MongoBasedTest {
public void changeAccountId() {
Context.accountId.set(currAccountId);
currAccountId += 1;
DictionaryFilter.readDictionaryBinary();
}

public void testInitializer(){
Expand Down Expand Up @@ -229,7 +231,7 @@ public void testInvalidMergeParameterizedURL() {
APICatalogSync.mergeUrlsAndSave(123, true, false, sync.existingAPIsInDb);
sync.buildFromDB(false, true);

assertEquals(0, sync.getDbState(123).getStrictURLToMethods().size());
assertEquals(1, sync.getDbState(123).getStrictURLToMethods().size());
assertEquals(1, sync.getDbState(123).getTemplateURLToMethods().size());


Expand Down
82 changes: 82 additions & 0 deletions apps/api-runtime/src/test/java/com/akto/parsers/TestMergingNew.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
import com.akto.dto.type.*;
import com.akto.runtime.APICatalogSync;
import com.akto.types.CappedSet;
import com.akto.util.filter.DictionaryFilter;
import com.akto.utils.RedactSampleData;
import com.google.api.client.util.Charsets;
import com.google.common.hash.BloomFilter;
Expand All @@ -24,6 +25,7 @@
import com.mongodb.client.model.Updates;
import org.bson.conversions.Bson;
import org.bson.types.ObjectId;
import org.junit.Before;
import org.junit.Test;

import java.util.*;
Expand All @@ -35,6 +37,11 @@

public class TestMergingNew extends MongoBasedTest {

@Before
public void initMain() {
DictionaryFilter.readDictionaryBinary();
}

public void testInitializer(){
Map<String, AktoDataType> aktoDataTypeMap = new HashMap<>();
aktoDataTypeMap.put("JWT", new AktoDataType(null, false, null, 0, new IgnoreData(new HashMap<>(), new HashSet<>()), false, true));
Expand Down Expand Up @@ -94,6 +101,81 @@ public void testMultipleIntegerMerging() {

}

@Test
public void testStringMerging() {
testInitializer();
SingleTypeInfoDao.instance.getMCollection().drop();
ApiCollectionsDao.instance.getMCollection().drop();
HttpCallParser parser = new HttpCallParser("userIdentifier", 1, 1, 1, true);

String baseUrl = "/api/";
List<HttpResponseParams> responseParams = new ArrayList<>();
List<String> urls = Arrays.asList(
baseUrl + "demo",
baseUrl + "cat",
baseUrl + "OSHE2CNS",
baseUrl + "2HOIWNJK",
baseUrl + "31a1a7c5-b4e3-47f5-8579-f7fc044c6a98",
baseUrl + "tree"
);

for (String c : urls) {
HttpResponseParams resp = createSampleParams("user1", c);
responseParams.add(resp);
}

parser.syncFunction(responseParams, false, true, null);
parser.apiCatalogSync.syncWithDB(false, true, SyncLimit.noLimit);
APICatalogSync.mergeUrlsAndSave(123, true, false, parser.apiCatalogSync.existingAPIsInDb);
parser.apiCatalogSync.buildFromDB(true, true);
Map<URLTemplate, RequestTemplate> urlTemplateMap = parser.apiCatalogSync.getDbState(123).getTemplateURLToMethods();
Map<URLStatic, RequestTemplate> strictUrlMap = parser.apiCatalogSync.getDbState(123).getStrictURLToMethods();


assertEquals(1, urlTemplateMap.size());
assertEquals(3, strictUrlMap.size());
}

@Test
public void testEnglishWordsUrlTestString() {
testInitializer();
SingleTypeInfoDao.instance.getMCollection().drop();
ApiCollectionsDao.instance.getMCollection().drop();
HttpCallParser parser = new HttpCallParser("userIdentifier", 1, 1, 1, true);
String url = "/link/";
List<HttpResponseParams> responseParams = new ArrayList<>();
List<String> urls = new ArrayList<>();
for (String x: Arrays.asList(
"apple", "banana", "cat", "dog", "elephant", "flower", "guitar", "house",
"island", "jungle", "kite", "lemon", "mountain", "night", "ocean", "piano",
"queen", "river", "sun", "tree", "umbrella", "village", "whale", "xylophone",
"yacht", "zebra", "bird", "clock", "desert", "engine", "forest", "garden",
"honey", "igloo", "jacket", "kangaroo", "lamp", "mirror", "notebook", "orange",
"pencil", "quilt", "rain", "star", "telephone", "uniform", "violin", "window",
"yellow", "zipper"
)) {
urls.add(url+x);
}
for (String c: urls) {
HttpResponseParams resp = createSampleParams("user1", c);
responseParams.add(resp);
}

parser.syncFunction(responseParams.subList(0,23), false, true, null);
parser.apiCatalogSync.syncWithDB(false, true, SyncLimit.noLimit);
assertEquals(23, getStaticURLsSize(parser));

parser.syncFunction(responseParams.subList(23,28), false, true, null);
parser.apiCatalogSync.syncWithDB(false, true, SyncLimit.noLimit);
APICatalogSync.mergeUrlsAndSave(123,true, false, parser.apiCatalogSync.existingAPIsInDb);
parser.apiCatalogSync.buildFromDB(false, true);
assertEquals(28, getStaticURLsSize(parser));

parser.syncFunction(responseParams.subList(28,33), false, true, null);
parser.apiCatalogSync.syncWithDB(false, true, SyncLimit.noLimit);
assertEquals(33, getStaticURLsSize(parser));
}

public int getStaticURLsSize(HttpCallParser parser) {
Map<URLStatic, RequestTemplate> urlStaticMap = parser.apiCatalogSync.getDbState(123).getStrictURLToMethods();

Expand Down
15 changes: 15 additions & 0 deletions apps/dashboard/src/main/java/com/akto/action/observe/InventoryAction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,14 @@
import com.akto.action.UserAction;
import com.akto.dao.*;
import com.akto.dao.context.Context;
import com.akto.dao.filter.MergedUrlsDao;
import com.akto.dao.testing_run_findings.TestingRunIssuesDao;
import com.akto.dto.*;
import com.akto.dto.ApiCollection.ENV_TYPE;
import com.akto.dto.ApiCollection.Type;
import com.akto.dto.ApiInfo.ApiInfoKey;
import com.akto.dto.CodeAnalysisApiInfo.CodeAnalysisApiInfoKey;
import com.akto.dto.filter.MergedUrls;
import com.akto.dto.traffic.SampleData;
import com.akto.dto.type.*;
import com.akto.dto.type.URLMethods.Method;
Expand Down Expand Up @@ -722,6 +724,19 @@ public String deMergeApi() {
return ERROR.toUpperCase();
}

try {
MergedUrlsDao.instance.updateOne(Filters.and(
Filters.eq(MergedUrls.URL, url),
Filters.eq(MergedUrls.METHOD, method),
Filters.eq(MergedUrls.API_COLLECTION_ID, apiCollectionId)
), Updates.combine(
Updates.set(MergedUrls.URL, url),
Updates.set(MergedUrls.METHOD, method),
Updates.set(MergedUrls.API_COLLECTION_ID, apiCollectionId)
));
} catch (Exception e) {
loggerMaker.errorAndAddToDb("Error while saving merged url in DB: " + e.getMessage(), LogDb.DASHBOARD);
}

SampleData sampleData = SampleDataDao.instance.fetchSampleDataForApi(apiCollectionId, url, urlMethod);
List<String> samples = sampleData.getSamples();
Expand Down
2 changes: 2 additions & 0 deletions apps/dashboard/src/main/java/com/akto/listener/InitializerListener.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@
import com.akto.util.UsageUtils;
import com.akto.util.enums.GlobalEnums.TestCategory;
import com.akto.util.enums.GlobalEnums.YamlTemplateSource;
import com.akto.util.filter.DictionaryFilter;
import com.akto.util.http_util.CoreHTTPClient;
import com.akto.util.tasks.OrganizationTask;
import com.akto.utils.*;
Expand Down Expand Up @@ -1814,6 +1815,7 @@ public static boolean isNotKubernetes() {
@Override
public void contextInitialized(javax.servlet.ServletContextEvent sce) {
setSubdomain();
DictionaryFilter.readDictionaryBinary();

String https = System.getenv("AKTO_HTTPS_FLAG");
if (Objects.equals(https, "true")) {
Expand Down
Loading

0 comments on commit a118642

Please sign in to comment.