forked from fortuna/ss-example
-
Notifications
You must be signed in to change notification settings - Fork 192
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
refactor: extract IP info package and add ASN support (#154)
- Loading branch information
Showing
15 changed files
with
505 additions
and
167 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,89 @@ | ||
| // Copyright 2023 Jigsaw Operations LLC | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"); | ||
| // you may not use this file except in compliance with the License. | ||
| // You may obtain a copy of the License at | ||
| // | ||
| // https://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| package ipinfo | ||
|
|
||
| import ( | ||
| "errors" | ||
| "fmt" | ||
| "net" | ||
| ) | ||
|
|
||
| type IPInfoMap interface { | ||
| GetIPInfo(net.IP) (IPInfo, error) | ||
| } | ||
|
|
||
| type IPInfo struct { | ||
| CountryCode CountryCode | ||
| ASN int | ||
| } | ||
|
|
||
| type CountryCode string | ||
|
|
||
| func (cc CountryCode) String() string { | ||
| return string(cc) | ||
| } | ||
|
|
||
| const ( | ||
| // Codes in the X* range are reserved to be user-assigned. | ||
| // See https://en.wikipedia.org/wiki/ISO_3166-1_alpha-2#Decoding_table | ||
| errParseAddr CountryCode = "XA" | ||
| localLocation CountryCode = "XL" | ||
| errDbLookupError CountryCode = "XD" | ||
| // Code "ZZ" is also used by Unicode as unknown country. | ||
| unknownLocation CountryCode = "ZZ" | ||
| ) | ||
|
|
||
| // GetIPInfoFromAddr is a helper function to extract the IP address from the [net.Addr] | ||
| // and call [IPInfoMap].GetIPInfo. | ||
| // It uses special country codes to indicate errors: | ||
| // - "XA": failed to extract the IP from the address ("A" is for "Address"). | ||
| // - "XL": IP is not global ("L" is for "Local"). | ||
| // - "XD": database error looking up the country code ("D" is for "DB"). | ||
| // - "ZZ": lookup returned an empty country code (same as the Unicode unknown location). | ||
| func GetIPInfoFromAddr(ip2info IPInfoMap, addr net.Addr) (IPInfo, error) { | ||
| var info IPInfo | ||
| if ip2info == nil { | ||
| // Location is disabled. return empty info. | ||
| return info, nil | ||
| } | ||
|
|
||
| if addr == nil { | ||
| info.CountryCode = errParseAddr | ||
| return info, fmt.Errorf("address cannot be nil") | ||
| } | ||
| hostname, _, err := net.SplitHostPort(addr.String()) | ||
| if err != nil { | ||
| info.CountryCode = errParseAddr | ||
| return info, fmt.Errorf("failed to split hostname and port: %w", err) | ||
| } | ||
| ip := net.ParseIP(hostname) | ||
| if ip == nil { | ||
| info.CountryCode = errParseAddr | ||
| return info, errors.New("failed to parse address as IP") | ||
| } | ||
|
|
||
| if !ip.IsGlobalUnicast() { | ||
| info.CountryCode = localLocation | ||
| return info, nil | ||
| } | ||
| info, err = ip2info.GetIPInfo(ip) | ||
| if err != nil { | ||
| info.CountryCode = errDbLookupError | ||
| } | ||
| if info.CountryCode == "" { | ||
| info.CountryCode = unknownLocation | ||
| } | ||
| return info, err | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,109 @@ | ||
| // Copyright 2023 Jigsaw Operations LLC | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"); | ||
| // you may not use this file except in compliance with the License. | ||
| // You may obtain a copy of the License at | ||
| // | ||
| // https://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| package ipinfo | ||
|
|
||
| import ( | ||
| "errors" | ||
| "net" | ||
| "testing" | ||
|
|
||
| "github.com/stretchr/testify/require" | ||
| ) | ||
|
|
||
| type noopMap struct{} | ||
|
|
||
| func (*noopMap) GetIPInfo(ip net.IP) (IPInfo, error) { | ||
| return IPInfo{}, nil | ||
| } | ||
|
|
||
| type badMap struct{} | ||
|
|
||
| func (*badMap) GetIPInfo(ip net.IP) (IPInfo, error) { | ||
| return IPInfo{}, errors.New("bad map") | ||
| } | ||
|
|
||
| type badAddr struct { | ||
| address string | ||
| } | ||
|
|
||
| func (a *badAddr) String() string { | ||
| return a.address | ||
| } | ||
|
|
||
| func (a *badAddr) Network() string { | ||
| return "bad" | ||
| } | ||
|
|
||
| func TestGetIPInfoFromAddr(t *testing.T) { | ||
| var emptyInfo IPInfo | ||
| noInfoMap := &noopMap{} | ||
|
|
||
| // IP info disabled | ||
| info, err := GetIPInfoFromAddr(nil, nil) | ||
| require.Equal(t, emptyInfo, info) | ||
| require.NoError(t, err) | ||
|
|
||
| // Nil address | ||
| info, err = GetIPInfoFromAddr(noInfoMap, nil) | ||
| require.Error(t, err) | ||
| require.Equal(t, errParseAddr, info.CountryCode) | ||
|
|
||
| // Can't split host:port in address | ||
| info, err = GetIPInfoFromAddr(noInfoMap, &badAddr{"host-no-port"}) | ||
| require.Error(t, err) | ||
| require.Equal(t, errParseAddr, info.CountryCode) | ||
|
|
||
| // Host is not an IP | ||
| info, err = GetIPInfoFromAddr(noInfoMap, &badAddr{"host-is-not-ip:port"}) | ||
| require.Error(t, err) | ||
| require.Equal(t, errParseAddr, info.CountryCode) | ||
|
|
||
| // Localhost address | ||
| info, err = GetIPInfoFromAddr(noInfoMap, &badAddr{"127.0.0.1:port"}) | ||
| require.NoError(t, err) | ||
| require.Equal(t, localLocation, info.CountryCode) | ||
|
|
||
| // Local network address | ||
| info, err = GetIPInfoFromAddr(noInfoMap, &badAddr{"10.0.0.1:port"}) | ||
| require.NoError(t, err) | ||
| require.Equal(t, unknownLocation, info.CountryCode) | ||
|
|
||
| // No country found | ||
| info, err = GetIPInfoFromAddr(noInfoMap, &badAddr{"8.8.8.8:port"}) | ||
| require.NoError(t, err) | ||
| require.Equal(t, unknownLocation, info.CountryCode) | ||
|
|
||
| // Failed DB lookup | ||
| info, err = GetIPInfoFromAddr(&badMap{}, &badAddr{"8.8.8.8:port"}) | ||
| require.Error(t, err) | ||
| require.Equal(t, errDbLookupError, info.CountryCode) | ||
| } | ||
|
|
||
| func TestCountryCode(t *testing.T) { | ||
| require.Equal(t, "BR", CountryCode("BR").String()) | ||
| } | ||
|
|
||
| func BenchmarkGetIPInfoFromAddr(b *testing.B) { | ||
| ip2info := &noopMap{} | ||
| testAddr := &net.TCPAddr{IP: net.ParseIP("217.65.48.1"), Port: 12345} | ||
|
|
||
| b.ResetTimer() | ||
| // Repeatedly check the country for the same address. This is realistic, because | ||
| // servers call this method for each new connection, but typically many connections | ||
| // come from a single user in succession. | ||
| for i := 0; i < b.N; i++ { | ||
| GetIPInfoFromAddr(ip2info, testAddr) | ||
| } | ||
| } |
Oops, something went wrong.