Skip to content

Commit

Permalink
Merge pull request #7 from tayloraswift/bifurcated-tables
Browse files Browse the repository at this point in the history
encode separate ASN and country tables
  • Loading branch information
tayloraswift authored Nov 28, 2024
2 parents afd49fb + 4fd60d3 commit 0ef3a81
Show file tree
Hide file tree
Showing 15 changed files with 418 additions and 328 deletions.
2 changes: 1 addition & 1 deletion Package.swift
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ let package:Package = .init(
.library(name: "IPinfo", targets: ["IPinfo"]),
],
dependencies: [
.package(url: "https://github.com/tayloraswift/swift-bson", from: "0.3.0"),
.package(url: "https://github.com/tayloraswift/swift-bson", from: "0.3.1"),
.package(url: "https://github.com/tayloraswift/swift-json", from: "1.1.1"),

.package(url: "https://github.com/apple/swift-argument-parser", from: "1.5.0"),
Expand Down
5 changes: 3 additions & 2 deletions Scripts/Package
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ swift build -c release \
--target FirewallPrefabricator

# If JSON files are not present, download them from ipinfo.io
for file in country_asn ; do
for file in asn country ; do
if [ ! -f $file.json ]; then
if [ -z "$IPINFO_TOKEN" ]; then
echo "error: IPINFO_TOKEN is not set"
Expand Down Expand Up @@ -46,7 +46,8 @@ fi


swift run -c release FirewallPrefabricator \
--ipinfo country_asn.json \
--country country.json \
--asn asn.json \
--googlebot googlebot.json \
--bingbot bingbot.json \
--github github.json \
Expand Down
23 changes: 14 additions & 9 deletions Snippets/FirewallUsage.swift
Original file line number Diff line number Diff line change
Expand Up @@ -34,24 +34,29 @@ let data:Data = try .init(contentsOf: URL.init(fileURLWithPath: "firewall.bson")
let bson:BSON.Document = .init(bytes: [UInt8].init(data)[...])

let firewall:IP.Firewall = .load(from: try IP.Firewall.Image.init(bson: bson))
let (mapping, claimant):(IP.Mapping?, IP.Claimant?) = firewall.lookup(v6: ip)

guard
let mapping:IP.Mapping
let country:ISO.Country = firewall.country[v6: ip]
else
{
fatalError("No Country/ASN found for \(ip)")
fatalError("No Country found for \(ip)")
}

print("""
Address: \(ip)
Country: \(mapping.country)
Country: \(country)
""")

if let autonomousSystem:IP.AS = mapping.autonomousSystem
let (system, _):(IP.AS?, IP.Claimant?) = firewall.lookup(v6: ip)

guard
let system:IP.AS
else
{
print("""
ASN: \(autonomousSystem.number)
AS: \(autonomousSystem.domain) (\(autonomousSystem.name))
""")
fatalError("No ASN found for \(ip)")
}

print("""
ASN: \(system.number)
AS: \(system.domain) (\(system.name))
""")
56 changes: 35 additions & 21 deletions Sources/FirewallPrefabricator/IP.Firewall.Image (ext).swift
Original file line number Diff line number Diff line change
Expand Up @@ -4,55 +4,69 @@ import ISO

extension IP.Firewall.Image
{
static func build(from ranges:[IPinfo.AddressRange]) throws -> Self
static func build(from ranges:[IPinfo.ASNRange]) throws -> Self
{
var v4:[(IP.ASN, ISO.Country, ip:ClosedRange<IP.V4>)]
var v6:[(IP.ASN, ISO.Country, ip:ClosedRange<IP.V6>)]
var v4:[(IP.ASN, ip:ClosedRange<IP.V4>)]
var v6:[(IP.ASN, ip:ClosedRange<IP.V6>)]

let table:[IP.ASN: IP.AS]

(as: table, v4: v4, v6: v6) = ranges.reduce(into: ([:], [], []))
{
let asn:IP.ASN
$0.as[$1.asn] = .init(number: $1.asn, domain: $1.domain, name: $1.name)

if let autonomousSystem:IP.AS = $1.autonomousSystem
if let first:IP.V4 = $1.first.v4,
let last:IP.V4 = $1.last.v4
{
asn = autonomousSystem.number
$0.as[asn] = autonomousSystem
$0.v4.append(($1.asn, first ... last))
}
else
{
asn = 0
$0.v6.append(($1.asn, $1.first ... $1.last))
}
}

v4.sort { $0.ip.lowerBound < $1.ip.lowerBound }
v6.sort { $0.ip.lowerBound < $1.ip.lowerBound }

var image:Self = .init(autonomousSystems: table.values.sorted { $0.number < $1.number })
try image.colorByASN(v4: v4, v6: v6)

return image
}

mutating
func colorByCountry(from ranges:[IPinfo.CountryRange]) throws
{
var v4:[(ISO.Country, ip:ClosedRange<IP.V4>)]
var v6:[(ISO.Country, ip:ClosedRange<IP.V6>)]

(v4: v4, v6: v6) = ranges.reduce(into: ([], []))
{
if let first:IP.V4 = $1.first.v4,
let last:IP.V4 = $1.last.v4
{
$0.v4.append((asn, $1.country, first ... last))
$0.v4.append(($1.country, first ... last))
}
else
{
$0.v6.append((asn, $1.country, $1.first ... $1.last))
$0.v6.append(($1.country, $1.first ... $1.last))
}
}

v4.sort { $0.ip.lowerBound < $1.ip.lowerBound }
v6.sort { $0.ip.lowerBound < $1.ip.lowerBound }

var image:Self = .init(autonomousSystems: table.values.sorted { $0.number < $1.number })
try image.color(v4: v4, v6: v6)

return image
try self.colorByCountry(v4: v4, v6: v6)
}

mutating
func claim(_ claims:[IP.Claims]) throws
func colorByClaimant(_ claims:[IP.Claims]) throws
{
var (v4, v6):
(
v4:[(IP.Claimant, ip:ClosedRange<IP.V4>)],
v6:[(IP.Claimant, ip:ClosedRange<IP.V6>)]
) = claims.reduce(into: ([], []))
var v4:[(IP.Claimant, ip:ClosedRange<IP.V4>)]
var v6:[(IP.Claimant, ip:ClosedRange<IP.V6>)]

(v4: v4, v6: v6) = claims.reduce(into: ([], []))
{
for range:ClosedRange<IP.V4> in $1.v4
{
Expand All @@ -67,6 +81,6 @@ extension IP.Firewall.Image
v4.sort { $0.ip.lowerBound < $1.ip.lowerBound }
v6.sort { $0.ip.lowerBound < $1.ip.lowerBound }

try self.claim(v4: v4, v6: v6)
try self.colorByClaimant(v4: v4, v6: v6)
}
}
20 changes: 15 additions & 5 deletions Sources/FirewallPrefabricator/Main.swift
Original file line number Diff line number Diff line change
Expand Up @@ -33,9 +33,14 @@ struct Main:ParsableCommand
var github:String

@Option(
name: [.customLong("ipinfo")],
help: "IPinfo JSON file")
var ipinfo:String
name: [.customLong("country")],
help: "IPinfo Countries JSON file")
var country:String

@Option(
name: [.customLong("asn")],
help: "IPinfo Autonomous Systems JSON file")
var asn:String

@Option(
name: [.customLong("output"), .customShort("o")],
Expand All @@ -47,7 +52,12 @@ extension Main
func run() throws
{
var image:IP.Firewall.Image = try .build(from: try .splitting(try .init(
contentsOf: URL.init(fileURLWithPath: self.ipinfo),
contentsOf: URL.init(fileURLWithPath: self.asn),
encoding: .utf8),
where: \.isNewline))

try image.colorByCountry(from: try .splitting(try .init(
contentsOf: URL.init(fileURLWithPath: self.country),
encoding: .utf8),
where: \.isNewline))

Expand All @@ -69,7 +79,7 @@ extension Main
bingbot.add(to: &claims, as: .microsoft_bingbot)
github.add(to: &claims)

try image.claim(claims)
try image.colorByClaimant(claims)

let bson:BSON.Document = .init(encoding: image)
try Data.init(bson.bytes).write(to: URL.init(fileURLWithPath: self.output))
Expand Down
134 changes: 67 additions & 67 deletions Sources/FirewallTests/Bisection.swift
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ struct Bisection

#expect(throws: (any Error).self)
{
try image.color(v4: ranges.map { (0, .us, $0) }, v6: [])
try image.colorByCountry(v4: ranges.map { (.us, $0) }, v6: [])
}
}

Expand All @@ -31,49 +31,49 @@ struct Bisection
#expect(empty.lookup(v6: 0) == (nil, nil))

var image:IP.Firewall.Image = .init(autonomousSystems: [])
try image.color(
try image.colorByCountry(
v4: [
(0, .us, 1 ... 10),
(0, .gb, 11 ... 20),
(0, .il, 22 ... 30),
(.us, 1 ... 10),
(.gb, 11 ... 20),
(.il, 22 ... 30),
],
v6: [
(0, .ua, 11 ... 20),
(.ua, 11 ... 20),
])

let firewall:IP.Firewall = .load(from: image)

#expect(firewall.lookup(v6: 0).mapping?.country == nil)
#expect(firewall.lookup(v6: 1).mapping?.country == nil)
#expect(firewall.lookup(v6: 10).mapping?.country == nil)
#expect(firewall.lookup(v6: 11).mapping?.country == .ua)
#expect(firewall.lookup(v6: 19).mapping?.country == .ua)
#expect(firewall.lookup(v6: 20).mapping?.country == .ua)
#expect(firewall.lookup(v6: 21).mapping?.country == nil)

#expect(firewall.lookup(v4: 0).mapping?.country == nil)
#expect(firewall.lookup(v4: 1).mapping?.country == .us)
#expect(firewall.lookup(v4: 9).mapping?.country == .us)
#expect(firewall.lookup(v4: 10).mapping?.country == .us)
#expect(firewall.lookup(v4: 11).mapping?.country == .gb)
#expect(firewall.lookup(v4: 19).mapping?.country == .gb)
#expect(firewall.lookup(v4: 20).mapping?.country == .gb)
#expect(firewall.lookup(v4: 21).mapping?.country == nil)
#expect(firewall.lookup(v4: 22).mapping?.country == .il)
#expect(firewall.lookup(v4: 30).mapping?.country == .il)
#expect(firewall.lookup(v4: 31).mapping?.country == nil)

#expect(firewall.lookup(v6: .init(v4: 0)).mapping?.country == nil)
#expect(firewall.lookup(v6: .init(v4: 1)).mapping?.country == .us)
#expect(firewall.lookup(v6: .init(v4: 9)).mapping?.country == .us)
#expect(firewall.lookup(v6: .init(v4: 10)).mapping?.country == .us)
#expect(firewall.lookup(v6: .init(v4: 11)).mapping?.country == .gb)
#expect(firewall.lookup(v6: .init(v4: 19)).mapping?.country == .gb)
#expect(firewall.lookup(v6: .init(v4: 20)).mapping?.country == .gb)
#expect(firewall.lookup(v6: .init(v4: 21)).mapping?.country == nil)
#expect(firewall.lookup(v6: .init(v4: 22)).mapping?.country == .il)
#expect(firewall.lookup(v6: .init(v4: 30)).mapping?.country == .il)
#expect(firewall.lookup(v6: .init(v4: 31)).mapping?.country == nil)
#expect(firewall.country[v6: 0] == nil)
#expect(firewall.country[v6: 1] == nil)
#expect(firewall.country[v6: 10] == nil)
#expect(firewall.country[v6: 11] == .ua)
#expect(firewall.country[v6: 19] == .ua)
#expect(firewall.country[v6: 20] == .ua)
#expect(firewall.country[v6: 21] == nil)

#expect(firewall.country[v4: 0] == nil)
#expect(firewall.country[v4: 1] == .us)
#expect(firewall.country[v4: 9] == .us)
#expect(firewall.country[v4: 10] == .us)
#expect(firewall.country[v4: 11] == .gb)
#expect(firewall.country[v4: 19] == .gb)
#expect(firewall.country[v4: 20] == .gb)
#expect(firewall.country[v4: 21] == nil)
#expect(firewall.country[v4: 22] == .il)
#expect(firewall.country[v4: 30] == .il)
#expect(firewall.country[v4: 31] == nil)

#expect(firewall.country[v6: .init(v4: 0)] == nil)
#expect(firewall.country[v6: .init(v4: 1)] == .us)
#expect(firewall.country[v6: .init(v4: 9)] == .us)
#expect(firewall.country[v6: .init(v4: 10)] == .us)
#expect(firewall.country[v6: .init(v4: 11)] == .gb)
#expect(firewall.country[v6: .init(v4: 19)] == .gb)
#expect(firewall.country[v6: .init(v4: 20)] == .gb)
#expect(firewall.country[v6: .init(v4: 21)] == nil)
#expect(firewall.country[v6: .init(v4: 22)] == .il)
#expect(firewall.country[v6: .init(v4: 30)] == .il)
#expect(firewall.country[v6: .init(v4: 31)] == nil)
}

@Test
Expand All @@ -88,38 +88,38 @@ struct Bisection
]

var image:IP.Firewall.Image = .init(autonomousSystems: [])
try image.color(v4: [], v6: blocks.map { (0, .us, $0.range) })
try image.colorByCountry(v4: [], v6: blocks.map { (.us, $0.range) })

let firewall:IP.Firewall = .load(from: image)

#expect(firewall.lookup(v6: 0xAAAA_AAAA_AAAA_AAA9_FFFF_FFFF_FFFF_FFFF).mapping == nil)
#expect(firewall.lookup(v6: 0xAAAA_AAAA_AAAA_AAAA_0000_0000_0000_0000).mapping != nil)
#expect(firewall.lookup(v6: 0xAAAA_AAAA_AAAA_AAAA_0000_0000_0000_0001).mapping != nil)
#expect(firewall.lookup(v6: 0xAAAA_AAAA_AAAA_AAAA_FFFF_FFFF_FFFF_FFFF).mapping != nil)
#expect(firewall.lookup(v6: 0xAAAA_AAAA_AAAA_AAAB_0000_0000_0000_0000).mapping == nil)

#expect(firewall.lookup(v6: 0xBBBB_BBBB_BBBB_BBBB_AFFF_FFFF_FFFF_FFFF).mapping == nil)
#expect(firewall.lookup(v6: 0xBBBB_BBBB_BBBB_BBBB_B000_0000_0000_0000).mapping != nil)
#expect(firewall.lookup(v6: 0xBBBB_BBBB_BBBB_BBBB_B000_0000_0000_0001).mapping != nil)
#expect(firewall.lookup(v6: 0xBBBB_BBBB_BBBB_BBBB_BFFF_FFFF_FFFF_FFFF).mapping != nil)
#expect(firewall.lookup(v6: 0xBBBB_BBBB_BBBB_BBBB_C000_0000_0000_0000).mapping == nil)

#expect(firewall.lookup(v6: 0xCCCC_CCCC_CCCC_CCCC_CBFF_FFFF_FFFF_FFFF).mapping == nil)
#expect(firewall.lookup(v6: 0xCCCC_CCCC_CCCC_CCCC_CC00_0000_0000_0000).mapping != nil)
#expect(firewall.lookup(v6: 0xCCCC_CCCC_CCCC_CCCC_CC00_0000_0000_0001).mapping != nil)
#expect(firewall.lookup(v6: 0xCCCC_CCCC_CCCC_CCCC_CCFF_FFFF_FFFF_FFFF).mapping != nil)
#expect(firewall.lookup(v6: 0xCCCC_CCCC_CCCC_CCCC_CD00_0000_0000_0000).mapping == nil)

#expect(firewall.lookup(v6: 0xDDDD_DDDD_DDDD_DDDD_DDCF_FFFF_FFFF_FFFF).mapping == nil)
#expect(firewall.lookup(v6: 0xDDDD_DDDD_DDDD_DDDD_DDD0_0000_0000_0000).mapping != nil)
#expect(firewall.lookup(v6: 0xDDDD_DDDD_DDDD_DDDD_DDD0_0000_0000_0001).mapping != nil)
#expect(firewall.lookup(v6: 0xDDDD_DDDD_DDDD_DDDD_DDDF_FFFF_FFFF_FFFF).mapping != nil)
#expect(firewall.lookup(v6: 0xDDDD_DDDD_DDDD_DDDD_DDE0_0000_0000_0000).mapping == nil)

#expect(firewall.lookup(v6: 0xEEEE_EEEE_EEEE_EEEE_EEED_FFFF_FFFF_FFFF).mapping == nil)
#expect(firewall.lookup(v6: 0xEEEE_EEEE_EEEE_EEEE_EEEE_0000_0000_0000).mapping != nil)
#expect(firewall.lookup(v6: 0xEEEE_EEEE_EEEE_EEEE_EEEE_0000_0000_0001).mapping != nil)
#expect(firewall.lookup(v6: 0xEEEE_EEEE_EEEE_EEEE_EEEE_FFFF_FFFF_FFFF).mapping != nil)
#expect(firewall.lookup(v6: 0xEEEE_EEEE_EEEE_EEEE_EEEF_0000_0000_0000).mapping == nil)
#expect(firewall.country[v6: 0xAAAA_AAAA_AAAA_AAA9_FFFF_FFFF_FFFF_FFFF] == nil)
#expect(firewall.country[v6: 0xAAAA_AAAA_AAAA_AAAA_0000_0000_0000_0000] != nil)
#expect(firewall.country[v6: 0xAAAA_AAAA_AAAA_AAAA_0000_0000_0000_0001] != nil)
#expect(firewall.country[v6: 0xAAAA_AAAA_AAAA_AAAA_FFFF_FFFF_FFFF_FFFF] != nil)
#expect(firewall.country[v6: 0xAAAA_AAAA_AAAA_AAAB_0000_0000_0000_0000] == nil)

#expect(firewall.country[v6: 0xBBBB_BBBB_BBBB_BBBB_AFFF_FFFF_FFFF_FFFF] == nil)
#expect(firewall.country[v6: 0xBBBB_BBBB_BBBB_BBBB_B000_0000_0000_0000] != nil)
#expect(firewall.country[v6: 0xBBBB_BBBB_BBBB_BBBB_B000_0000_0000_0001] != nil)
#expect(firewall.country[v6: 0xBBBB_BBBB_BBBB_BBBB_BFFF_FFFF_FFFF_FFFF] != nil)
#expect(firewall.country[v6: 0xBBBB_BBBB_BBBB_BBBB_C000_0000_0000_0000] == nil)

#expect(firewall.country[v6: 0xCCCC_CCCC_CCCC_CCCC_CBFF_FFFF_FFFF_FFFF] == nil)
#expect(firewall.country[v6: 0xCCCC_CCCC_CCCC_CCCC_CC00_0000_0000_0000] != nil)
#expect(firewall.country[v6: 0xCCCC_CCCC_CCCC_CCCC_CC00_0000_0000_0001] != nil)
#expect(firewall.country[v6: 0xCCCC_CCCC_CCCC_CCCC_CCFF_FFFF_FFFF_FFFF] != nil)
#expect(firewall.country[v6: 0xCCCC_CCCC_CCCC_CCCC_CD00_0000_0000_0000] == nil)

#expect(firewall.country[v6: 0xDDDD_DDDD_DDDD_DDDD_DDCF_FFFF_FFFF_FFFF] == nil)
#expect(firewall.country[v6: 0xDDDD_DDDD_DDDD_DDDD_DDD0_0000_0000_0000] != nil)
#expect(firewall.country[v6: 0xDDDD_DDDD_DDDD_DDDD_DDD0_0000_0000_0001] != nil)
#expect(firewall.country[v6: 0xDDDD_DDDD_DDDD_DDDD_DDDF_FFFF_FFFF_FFFF] != nil)
#expect(firewall.country[v6: 0xDDDD_DDDD_DDDD_DDDD_DDE0_0000_0000_0000] == nil)

#expect(firewall.country[v6: 0xEEEE_EEEE_EEEE_EEEE_EEED_FFFF_FFFF_FFFF] == nil)
#expect(firewall.country[v6: 0xEEEE_EEEE_EEEE_EEEE_EEEE_0000_0000_0000] != nil)
#expect(firewall.country[v6: 0xEEEE_EEEE_EEEE_EEEE_EEEE_0000_0000_0001] != nil)
#expect(firewall.country[v6: 0xEEEE_EEEE_EEEE_EEEE_EEEE_FFFF_FFFF_FFFF] != nil)
#expect(firewall.country[v6: 0xEEEE_EEEE_EEEE_EEEE_EEEF_0000_0000_0000] == nil)
}
}
8 changes: 8 additions & 0 deletions Sources/Firewalls/IP.ASN.swift
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,14 @@ extension IP.ASN:LosslessStringConvertible
self.init(value: value)
}
}
extension IP.ASN:BSON.BinaryPackable
{
@inlinable public
static func get(_ storage:UInt32) -> Self { .init(value: .get(storage)) }

@inlinable public
consuming func set() -> UInt32 { self.value.set() }
}
extension IP.ASN:BSONEncodable
{
/// Encodes the ASN as a signed BSON ``Int32``, by mapping the range of the ``UInt32``
Expand Down
Loading

0 comments on commit 0ef3a81

Please sign in to comment.