Skip to content

Commit

Permalink
Improved geo-blocking trie logic
Browse files Browse the repository at this point in the history
  • Loading branch information
@tudor-timcu
tudor-timcu committed Dec 11, 2024
1 parent 4152a7e commit 698cbae
Show file tree
Hide file tree
Showing 10 changed files with 204 additions and 22 deletions.
13 changes: 7 additions & 6 deletions lib/request-processor/aikido_types/config.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,10 +37,11 @@ type EndpointKey struct {
}

type CloudConfigData struct {
ConfigUpdatedAt int64
Endpoints map[EndpointKey]EndpointData
BlockedUserIds map[string]bool
BypassedIps map[string]bool
GeoBlockedIpsTrie *ipaddr.AddressTrie
Block int
ConfigUpdatedAt int64
Endpoints map[EndpointKey]EndpointData
BlockedUserIds map[string]bool
BypassedIps map[string]bool
GeoBlockedIpsTrieV4 *ipaddr.IPv4AddressTrie
GeoBlockedIpsTrieV6 *ipaddr.IPv6AddressTrie
Block int
}
22 changes: 19 additions & 3 deletions lib/request-processor/grpc/config.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
. "main/aikido_types"
"main/globals"
"main/ipc/protos"
"main/log"
"time"

"github.com/seancfoley/ipaddress-go/ipaddr"
Expand All @@ -16,14 +17,29 @@ var (

func buildGeoBlockedIpsTrie(geoBlockedIps []string) {
if len(geoBlockedIps) == 0 {
globals.CloudConfig.GeoBlockedIpsTrie = nil
globals.CloudConfig.GeoBlockedIpsTrieV4 = nil
globals.CloudConfig.GeoBlockedIpsTrieV6 = nil
return
}

globals.CloudConfig.GeoBlockedIpsTrie = &ipaddr.AddressTrie{}
globals.CloudConfig.GeoBlockedIpsTrieV4 = &ipaddr.IPv4AddressTrie{}
globals.CloudConfig.GeoBlockedIpsTrieV6 = &ipaddr.IPv6AddressTrie{}
for _, ip := range geoBlockedIps {
globals.CloudConfig.GeoBlockedIpsTrie.Add(ipaddr.NewIPAddressString(ip).GetAddress().ToAddressBase())
ipAddress, err := ipaddr.NewIPAddressString(ip).ToAddress()
if err != nil {
log.Infof("Invalid geoip address: %s\n", ip)
continue
}

if ipAddress.IsIPv4() {
globals.CloudConfig.GeoBlockedIpsTrieV4.Add(ipAddress.ToIPv4())
} else if ipAddress.IsIPv6() {
globals.CloudConfig.GeoBlockedIpsTrieV6.Add(ipAddress.ToIPv6())
}
}

log.Debugf("Built GeoBlockedIpsTrieV4: %d\n%v", globals.CloudConfig.GeoBlockedIpsTrieV4.Size(), globals.CloudConfig.GeoBlockedIpsTrieV4)
log.Debugf("Built GeoBlockedIpsTrieV6: %d\n%v", globals.CloudConfig.GeoBlockedIpsTrieV6.Size(), globals.CloudConfig.GeoBlockedIpsTrieV6)
}

func setCloudConfig(cloudConfigFromAgent *protos.CloudConfig) {
Expand Down
4 changes: 2 additions & 2 deletions lib/request-processor/handle_blocking_request.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ func OnGetBlockingStatus() string {

ip := context.GetIp()
if utils.IsIpGeoBlocked(ip) {
return GetStoreAction("blocked", "geo", ip)
return GetStoreAction("blocked", "geoip", ip)
}

method := context.GetMethod()
Expand All @@ -46,7 +46,7 @@ func OnGetBlockingStatus() string {

endpointData, err := utils.GetEndpointConfig(method, route)
if err != nil {
log.Debugf("Method+route in not configured in endpoints! Skipping checks...")
log.Debugf("Method+route is not configured in endpoints! Skipping checks...")
return ""
}

Expand Down
16 changes: 14 additions & 2 deletions lib/request-processor/utils/utils.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"encoding/json"
"fmt"
"main/globals"
"main/log"
"net"
"strings"

Expand Down Expand Up @@ -183,10 +184,21 @@ func IsIpGeoBlocked(ip string) bool {
globals.CloudConfigMutex.Lock()
defer globals.CloudConfigMutex.Unlock()

if globals.CloudConfig.GeoBlockedIpsTrie == nil {
ipAddress, err := ipaddr.NewIPAddressString(ip).ToAddress()
if err != nil {
log.Infof("Invalid ip address: %s\n", ip)
return false
}
return globals.CloudConfig.GeoBlockedIpsTrie.Contains(ipaddr.NewIPAddressString(ip).GetAddress().ToAddressBase())

if ipAddress.IsIPv4() && globals.CloudConfig.GeoBlockedIpsTrieV4 != nil {
log.Infof("Checking ipv4 address for geo-blocking: %v\n", ipAddress.ToIPv4())
return globals.CloudConfig.GeoBlockedIpsTrieV4.ElementContains(ipAddress.ToIPv4())
} else if ipAddress.IsIPv6() && globals.CloudConfig.GeoBlockedIpsTrieV6 != nil {
log.Infof("Checking ipv6 address for geo-blocking: %v\n", ipAddress.ToIPv6())
return globals.CloudConfig.GeoBlockedIpsTrieV6.ElementContains(ipAddress.ToIPv6())
}

return false
}

type DatabaseType int
Expand Down
10 changes: 10 additions & 0 deletions tests/server/test_geo_blocked_ip/change_config_remove_geo_blocked_ips.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
{
"success": true,
"serviceId": 1,
"heartbeatIntervalInMS": 600000,
"endpoints": [],
"blockedUserIds": [],
"allowedIPAddresses": [],
"blockedIPAddresses": [],
"receivedAnyStats": true
}
5 changes: 5 additions & 0 deletions tests/server/test_geo_blocked_ip/env.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
{
"AIKIDO_BLOCK": "1",
"AIKIDO_LOCALHOST_ALLOWED_BY_DEFAULT": "0",
"AIKIDO_FEATURE_COLLECT_API_SCHEMA": "1"
}
17 changes: 17 additions & 0 deletions tests/server/test_geo_blocked_ip/index.php
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
<?php

$_SERVER['REMOTE_ADDR'] = '5.8.19.22';

if (extension_loaded('aikido')) {
$decision = \aikido\should_block_request();

if ($decision->block && $decision->type == "blocked" && $decision->trigger == "geoip") {
http_response_code(403);
echo "Your IP address is blocked due to geo restrictions. (Your IP: {$decision->ip})";
exit();
}
}

echo "Something";

?>
79 changes: 79 additions & 0 deletions tests/server/test_geo_blocked_ip/start_config.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
{
"success": true,
"serviceId": 1,
"heartbeatIntervalInMS": 600000,
"endpoints": [],
"blockedUserIds": [],
"allowedIPAddresses": [],
"receivedAnyStats": true,
"blockedIPAddresses": [
{
"source": "geoip",
"description": "geo restrictions",
"ips": [
"64:ff9b:1::/48",
"1.2.9.0/24",
"1.4.9.0/24",
"1.173.94.92/30",
"2.16.20.0/23",
"2.16.53.0/24",
"2.16.103.0/24",
"2.16.154.0/24",
"2.17.144.0/23",
"2.17.146.0/24",
"2.19.181.0/24",
"2.19.204.0/23",
"2.20.254.0/23",
"2.22.237.0/24",
"2.23.167.0/24",
"2.56.26.0/23",
"2.56.88.0/23",
"2.56.138.0/24",
"2.56.178.0/24",
"2.56.180.0/22",
"2.56.228.0/22",
"2.56.240.0/23",
"2.56.242.0/24",
"2.57.0.0/24",
"2.57.37.0/24",
"2.57.38.0/24",
"2.57.52.0/22",
"2.57.80.0/22",
"2.57.112.0/22",
"2.57.184.0/22",
"2.58.68.0/22",
"2.58.96.0/23",
"2.58.98.0/24",
"2.58.124.0/22",
"2.59.40.0/22",
"2.59.48.0/23",
"2.59.50.0/24",
"2.59.51.0/24",
"2.59.76.0/22",
"2.59.80.0/22",
"2.59.160.0/23",
"2.59.176.0/22",
"2.59.213.0/24",
"2.59.214.0/23",
"2.59.216.0/22",
"2.59.240.0/22",
"2.60.0.0/15",
"2.62.0.0/16",
"2.63.0.0/17",
"2.63.128.0/20",
"2.63.160.0/20",
"2.63.192.0/19",
"2.63.240.0/20",
"2.92.0.0/14",
"5.1.41.0/24",
"5.1.47.0/24",
"5.1.48.0/21",
"5.2.32.0/19",
"5.3.0.0/16",
"5.8.8.0/21",
"5.8.16.0/23",
"5.8.19.0/24"
]
}
]
}
35 changes: 35 additions & 0 deletions tests/server/test_geo_blocked_ip/test.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
import requests
import time
import sys
from testlib import *

'''
1. Sets up the allowed IP address config for route '/test'. Checks that requests are blocked.
2. Changes the config to remote the allowed IP address. Checks that requests are passing.
3. Changes the config again to enable allowed IP address. Checks that requests are blocked.
'''


def run_test():
response = php_server_get("/test")
assert_response_code_is(response, 403)
assert_response_header_contains(response, "Content-Type", "text")
assert_response_body_contains(response, "Your IP address is blocked due to geo restrictions. (Your IP: ")

apply_config("change_config_remove_geo_blocked_ips.json")

response = php_server_get("/test")
assert_response_code_is(response, 200)
assert_response_body_contains(response, "Something")

apply_config("start_config.json")

response = php_server_get("/test")
assert_response_code_is(response, 403)
assert_response_header_contains(response, "Content-Type", "text")
assert_response_body_contains(response, "Your IP address is blocked due to geo restrictions. (Your IP: ")


if __name__ == "__main__":
load_test_args()
run_test()
25 changes: 16 additions & 9 deletions tools/mock_aikido_core.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,23 @@
responses = {
"config": {},
"configUpdatedAt": {},
"lists": {}
}

events = []
server_down = False

excluded_routes = ['mock_get_events', 'mock_tests_simple', 'mock_down', 'mock_up']

def load_config(j):
configUpdatedAt = int(time.time())
responses["lists"] = { "success": True, "serviceId": j["serviceId"], "blockedIPAddresses": j["blockedIPAddresses"] }
del j["blockedIPAddresses"]
responses["config"] = j
responses["config"]["configUpdatedAt"] = configUpdatedAt
responses["configUpdatedAt"] = { "serviceId": 1, "configUpdatedAt": configUpdatedAt }
print(f"Loaded new runtime config!")

@app.before_request
def check_server_status():
global server_down
Expand All @@ -33,6 +43,10 @@ def get_config():
def get_runtime_config():
return jsonify(responses["config"])

@app.route('/api/runtime/firewall/lists', methods=['GET'])
def get_lists_config():
return jsonify(responses["lists"])


@app.route('/api/runtime/events', methods=['POST'])
def post_events():
Expand All @@ -44,10 +58,7 @@ def post_events():

@app.route('/mock/config', methods=['POST'])
def mock_set_config():
configUpdatedAt = int(time.time())
responses["config"] = request.get_json()
responses["config"]["configUpdatedAt"] = configUpdatedAt
responses["configUpdatedAt"] = { "serviceId": 1, "configUpdatedAt": configUpdatedAt }
load_config(request.get_json())
return jsonify({})

@app.route('/mock/down', methods=['POST'])
Expand Down Expand Up @@ -83,11 +94,7 @@ def mock_tests_simple():
if os.path.exists(config_file):
try:
with open(config_file, 'r') as file:
configUpdatedAt = int(time.time())
responses["config"] = json.load(file)
responses["config"]["configUpdatedAt"] = configUpdatedAt
responses["configUpdatedAt"] = { "serviceId": 1, "configUpdatedAt": configUpdatedAt }
print(f"Loaded runtime config from {config_file}")
load_config(json.load(file))
except json.JSONDecodeError:
print(f"Error: Could not decode JSON from {config_file}")
sys.exit(1)
Expand Down

0 comments on commit 698cbae

Please sign in to comment.