Copilot commented on code in PR #12948:
URL: https://github.com/apache/trafficserver/pull/12948#discussion_r2927766211
##########
plugins/header_rewrite/conditions_geo_maxmind.cc:
##########
@@ -51,9 +76,23 @@ MMConditionGeo::initLibrary(const std::string &path)
if (MMDB_SUCCESS != status) {
Dbg(pi_dbg_ctl, "Cannot open %s - %s", path.c_str(),
MMDB_strerror(status));
delete gMaxMindDB;
+ gMaxMindDB = nullptr;
return;
}
- Dbg(pi_dbg_ctl, "Loaded %s", path.c_str());
+
+ // Probe the database schema at load time so we know which field paths to
+ // use for country lookups. Try a few well-known IPs until one hits.
+ for (auto *ip : probe_ips) {
+ int gai_error, mmdb_error;
+ MMDB_lookup_result_s result = MMDB_lookup_string(gMaxMindDB, ip,
&gai_error, &mmdb_error);
+ if (gai_error == 0 && MMDB_SUCCESS == mmdb_error && result.found_entry) {
+ gMmdbSchema = detect_schema(&result.entry);
+ Dbg(pi_dbg_ctl, "Loaded %s (schema: %s)", path.c_str(), gMmdbSchema ==
MmdbSchema::FLAT ? "flat" : "nested");
+ return;
+ }
+ }
+
+ Dbg(pi_dbg_ctl, "Loaded %s (schema: defaulting to nested, no probe IPs
matched)", path.c_str());
Review Comment:
Schema detection is based on probing a few hard-coded public IPs at load
time. If the database doesn’t contain any of these addresses (e.g.,
private/internal ranges), gMmdbSchema will remain the default (nested) and
flat-schema databases will never resolve country codes. Consider making country
lookup robust by attempting both paths per lookup (nested then flat, or
vice-versa), or lazily detecting the schema from the first successful lookup
result entry rather than relying on specific probe IPs.
##########
plugins/header_rewrite/header_rewrite_test.cc:
##########
@@ -538,12 +545,202 @@ test_tokenizer()
return errors;
}
+#if TS_USE_HRW_MAXMINDDB
+
+static bool
+file_exists(const char *path)
+{
+ struct stat st;
+ return stat(path, &st) == 0;
+}
+
+static int
+open_test_mmdb(MMDB_s &mmdb, const char *path)
+{
+ int status = MMDB_open(path, MMDB_MODE_MMAP, &mmdb);
+ if (MMDB_SUCCESS != status) {
+ std::cerr << "Cannot open " << path << ": " << MMDB_strerror(status) <<
std::endl;
+ return 1;
+ }
+ return 0;
+}
+
+static std::string
+lookup_country(MMDB_s &mmdb, const char *ip)
+{
+ int gai_error, mmdb_error;
+ MMDB_lookup_result_s result = MMDB_lookup_string(&mmdb, ip, &gai_error,
&mmdb_error);
+
+ if (gai_error != 0 || MMDB_SUCCESS != mmdb_error || !result.found_entry) {
+ return "(lookup_failed)";
+ }
+
+ MMDB_entry_data_s entry_data;
+
+ // Try nested path first (GeoLite2 style)
+ int status = MMDB_get_value(&result.entry, &entry_data, "country",
"iso_code", NULL);
+ if (MMDB_SUCCESS == status && entry_data.has_data && entry_data.type ==
MMDB_DATA_TYPE_UTF8_STRING) {
+ return std::string(entry_data.utf8_string, entry_data.data_size);
+ }
+
+ // Try flat path (vendor style)
+ status = MMDB_get_value(&result.entry, &entry_data, "country_code", NULL);
+ if (MMDB_SUCCESS == status && entry_data.has_data && entry_data.type ==
MMDB_DATA_TYPE_UTF8_STRING) {
+ return std::string(entry_data.utf8_string, entry_data.data_size);
+ }
+
+ return "(not_found)";
+}
+
+static uint32_t
+lookup_asn(MMDB_s &mmdb, const char *ip)
+{
+ int gai_error, mmdb_error;
+ MMDB_lookup_result_s result = MMDB_lookup_string(&mmdb, ip, &gai_error,
&mmdb_error);
+
+ if (gai_error != 0 || MMDB_SUCCESS != mmdb_error || !result.found_entry) {
+ return 0;
+ }
+
+ MMDB_entry_data_s entry_data;
+ int status = MMDB_get_value(&result.entry, &entry_data,
"autonomous_system_number", NULL);
+ if (MMDB_SUCCESS == status && entry_data.has_data && entry_data.type ==
MMDB_DATA_TYPE_UINT32) {
+ return entry_data.uint32;
+ }
+ return 0;
+}
+
+// Test that we can read country codes from a flat-schema MMDB (vendor
databases
+// where country_code is a top-level string field).
+int
+test_flat_schema()
+{
+ const char *path = getenv("MMDB_TEST_FLAT");
+ if (path == nullptr || !file_exists(path)) {
+ std::cout << "SKIP: flat-schema test mmdb not found (set MMDB_TEST_FLAT)"
<< std::endl;
+ return 0;
Review Comment:
This test code calls getenv(), but this file doesn’t include <cstdlib> (or
<stdlib.h>) directly. Relying on indirect declarations is non-portable and can
fail to compile on some toolchains; add the appropriate standard header.
##########
plugins/header_rewrite/CMakeLists.txt:
##########
@@ -60,7 +60,34 @@ if(BUILD_TESTING)
target_link_libraries(test_header_rewrite PRIVATE header_rewrite_parser
ts::inkevent ts::tscore)
if(maxminddb_FOUND)
+ target_compile_definitions(test_header_rewrite PRIVATE
TS_USE_HRW_MAXMINDDB=1)
target_link_libraries(test_header_rewrite PRIVATE maxminddb::maxminddb)
+
+ find_package(
+ Python3
+ COMPONENTS Interpreter
+ QUIET
+ )
+ if(Python3_FOUND)
+ set(_mmdb_test_dir "${CMAKE_CURRENT_BINARY_DIR}/test_mmdb")
+ add_custom_command(
+ OUTPUT "${_mmdb_test_dir}/test_flat_geo.mmdb"
"${_mmdb_test_dir}/test_nested_geo.mmdb"
+ COMMAND ${CMAKE_COMMAND} -E make_directory "${_mmdb_test_dir}"
+ COMMAND ${Python3_EXECUTABLE}
"${CMAKE_CURRENT_SOURCE_DIR}/generate_test_mmdb.py" "${_mmdb_test_dir}"
+ DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/generate_test_mmdb.py"
+ COMMENT "Generating test MMDB files for header_rewrite"
+ )
+ add_custom_target(
+ generate_test_mmdb DEPENDS "${_mmdb_test_dir}/test_flat_geo.mmdb"
"${_mmdb_test_dir}/test_nested_geo.mmdb"
+ )
+ add_dependencies(test_header_rewrite generate_test_mmdb)
+ set_tests_properties(
+ test_header_rewrite
+ PROPERTIES
+ ENVIRONMENT
+
"MMDB_TEST_FLAT=${_mmdb_test_dir}/test_flat_geo.mmdb;MMDB_TEST_NESTED=${_mmdb_test_dir}/test_nested_geo.mmdb"
+ )
Review Comment:
This custom command expects the Python script to generate the listed OUTPUT
files, but generate_test_mmdb.py currently exits 0 when dependencies are
missing, which can leave the outputs absent and break the build (e.g., Ninja
complains about missing outputs). Either (a) make CMake check for required
Python modules at configure time before adding this command/target, or (b) make
the script exit non-zero on missing deps so the failure is explicit.
```suggestion
execute_process(
COMMAND "${Python3_EXECUTABLE}" -c "import maxminddb"
RESULT_VARIABLE _maxminddb_python_result
OUTPUT_QUIET
ERROR_QUIET
)
if(_maxminddb_python_result EQUAL 0)
set(_have_python_maxminddb TRUE)
else()
set(_have_python_maxminddb FALSE)
message(STATUS "Python module 'maxminddb' not found; skipping
generation of test MMDB files for header_rewrite tests.")
endif()
if(_have_python_maxminddb)
set(_mmdb_test_dir "${CMAKE_CURRENT_BINARY_DIR}/test_mmdb")
add_custom_command(
OUTPUT "${_mmdb_test_dir}/test_flat_geo.mmdb"
"${_mmdb_test_dir}/test_nested_geo.mmdb"
COMMAND ${CMAKE_COMMAND} -E make_directory "${_mmdb_test_dir}"
COMMAND ${Python3_EXECUTABLE}
"${CMAKE_CURRENT_SOURCE_DIR}/generate_test_mmdb.py" "${_mmdb_test_dir}"
DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/generate_test_mmdb.py"
COMMENT "Generating test MMDB files for header_rewrite"
)
add_custom_target(
generate_test_mmdb DEPENDS "${_mmdb_test_dir}/test_flat_geo.mmdb"
"${_mmdb_test_dir}/test_nested_geo.mmdb"
)
add_dependencies(test_header_rewrite generate_test_mmdb)
set_tests_properties(
test_header_rewrite
PROPERTIES
ENVIRONMENT
"MMDB_TEST_FLAT=${_mmdb_test_dir}/test_flat_geo.mmdb;MMDB_TEST_NESTED=${_mmdb_test_dir}/test_nested_geo.mmdb"
)
endif()
```
##########
plugins/header_rewrite/header_rewrite_test.cc:
##########
@@ -24,9 +24,16 @@
#include <cstdarg>
#include <iostream>
#include <ostream>
+#include <sys/stat.h>
#include "parser.h"
+#if TS_USE_HRW_MAXMINDDB
+#include <maxminddb.h>
+#include <arpa/inet.h>
+#include <netdb.h>
Review Comment:
These headers appear unused in this file (no inet_*/getaddrinfo usage in the
added MaxMind test helpers). If they’re not needed, consider removing them to
keep the unit test’s dependencies minimal.
```suggestion
```
##########
plugins/header_rewrite/generate_test_mmdb.py:
##########
@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you 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
+#
+# http://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.
+"""
+Generate test MMDB files for header_rewrite geo lookup unit tests.
+
+Two schemas exist in the wild:
+
+ Nested (GeoLite2/GeoIP2/DBIP): country -> iso_code
+ Flat (vendor-specific): country_code (top-level)
+
+This script generates one MMDB file for each schema so the C++ test
+can verify that auto-detection works for both.
+
+Requires: pip install mmdb-writer
+"""
+
+import os
+import sys
+
+try:
+ from mmdb_writer import MMDBWriter, MmdbU32
+ import netaddr
+except ImportError:
+ print("SKIP: mmdb-writer or netaddr not installed (pip install
mmdb-writer)", file=sys.stderr)
+ sys.exit(0)
Review Comment:
The ImportError handler exits with status 0. When this script is used from
CMake as a generator for declared output files, exiting 0 without producing
outputs can break the build. Consider exiting non-zero (and making CMake skip
the generator when modules aren’t available) so the failure/skip behavior is
consistent.
```suggestion
sys.exit(1)
```
##########
plugins/header_rewrite/generate_test_mmdb.py:
##########
@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you 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
+#
+# http://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.
+"""
+Generate test MMDB files for header_rewrite geo lookup unit tests.
+
+Two schemas exist in the wild:
+
+ Nested (GeoLite2/GeoIP2/DBIP): country -> iso_code
+ Flat (vendor-specific): country_code (top-level)
+
+This script generates one MMDB file for each schema so the C++ test
+can verify that auto-detection works for both.
+
+Requires: pip install mmdb-writer
+"""
+
+import os
+import sys
+
+try:
+ from mmdb_writer import MMDBWriter, MmdbU32
+ import netaddr
+except ImportError:
+ print("SKIP: mmdb-writer or netaddr not installed (pip install
mmdb-writer)", file=sys.stderr)
+ sys.exit(0)
Review Comment:
The skip message suggests installing only mmdb-writer, but this script also
imports netaddr. Update the message to include all required packages so
developers can satisfy the dependency in one step.
##########
plugins/header_rewrite/generate_test_mmdb.py:
##########
@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you 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
+#
+# http://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.
+"""
+Generate test MMDB files for header_rewrite geo lookup unit tests.
+
+Two schemas exist in the wild:
+
+ Nested (GeoLite2/GeoIP2/DBIP): country -> iso_code
+ Flat (vendor-specific): country_code (top-level)
+
+This script generates one MMDB file for each schema so the C++ test
+can verify that auto-detection works for both.
+
+Requires: pip install mmdb-writer
+"""
+
+import os
+import sys
+
+try:
+ from mmdb_writer import MMDBWriter, MmdbU32
+ import netaddr
+except ImportError:
+ print("SKIP: mmdb-writer or netaddr not installed (pip install
mmdb-writer)", file=sys.stderr)
+ sys.exit(0)
+
+
+def net(cidr):
+ return netaddr.IPSet([netaddr.IPNetwork(cidr)])
+
+
+def generate_flat(path):
+ """Flat schema: country_code at top level (vendor databases)."""
+ w = MMDBWriter(ip_version=4, database_type="Test-Flat-GeoIP")
+ w.insert_network(
+ net("8.8.8.0/24"), {
+ "country_code": "US",
+ "autonomous_system_number": MmdbU32(15169),
+ "autonomous_system_organization": "GOOGLE",
+ })
+ w.insert_network(
+ net("1.2.3.0/24"), {
+ "country_code": "KR",
+ "autonomous_system_number": MmdbU32(9286),
+ "autonomous_system_organization": "KINX",
+ })
+ w.to_db_file(path)
+
+
+def generate_nested(path):
+ """Nested schema: country/iso_code (GeoLite2, GeoIP2, DBIP)."""
+ w = MMDBWriter(ip_version=4, database_type="Test-Nested-GeoIP2")
Review Comment:
Repo docs for HRW/HRW4U code style call for Python 3.11+ scripts to include
type hints on all functions (see .github/instructions/HRW.instructions.md).
Consider adding basic type annotations for
net()/generate_flat()/generate_nested() to align with the documented convention.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
