mirror of
https://github.com/hanzoai/datastore.git
synced 2026-09-26 13:55:07 +00:00
3e3d61e115
Issue #81402 reported that availability-zone detection failed on EC2 instances configured for "IMDSv2 only (token required)", because `getAvailabilityZoneOrException` issued a plain unauthenticated `GET` against the instance metadata service, which such instances reject with `401 Unauthorized`. The behavior was fixed in master by #81323 (`kafka_aws_az_aware`), which makes `getAvailabilityZoneOrException` acquire an IMDSv2 session token via `PUT /latest/api/token` and send it with the availability-zone `GET`, falling back to the token-less IMDSv1 request when the metadata service does not implement the token endpoint. That change did not add a test covering the IMDSv2-only scenario in `test_placement_info`. This adds that missing coverage: a new mock metadata server `imdsv2_server.py` that issues tokens on `PUT /latest/api/token` and rejects `GET` requests without a valid `x-aws-ec2-metadata-token` header (`401`), plus a `node_imds_v2` instance and a `test_placement_info_from_imdsv2_only` case that asserts the availability zone is still loaded against such an endpoint. The test originates from #103549, which proposed an equivalent fix before #81323 was merged; that pull request is now superseded for the production change, so only its regression test is kept here. Co-Authored-By: ashrithb <[email protected]> Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
133 lines
3.9 KiB
Python
133 lines
3.9 KiB
Python
import os
|
|
|
|
import pytest
|
|
|
|
from helpers.cluster import ClickHouseCluster
|
|
from helpers.mock_servers import start_mock_servers
|
|
|
|
METADATA_SERVER_HOSTNAME = "node_imds"
|
|
METADATA_SERVER_PORT = 8080
|
|
METADATA_SERVER_V2_HOSTNAME = "node_imds_v2"
|
|
METADATA_SERVER_V2_PORT = 8081
|
|
|
|
cluster = ClickHouseCluster(__file__)
|
|
node_imds = cluster.add_instance(
|
|
"node_imds",
|
|
main_configs=["configs/imds_bootstrap.xml"],
|
|
env_variables={
|
|
"AWS_EC2_METADATA_SERVICE_ENDPOINT": f"http://{METADATA_SERVER_HOSTNAME}:{METADATA_SERVER_PORT}",
|
|
},
|
|
stay_alive=True,
|
|
)
|
|
node_imds_v2 = cluster.add_instance(
|
|
"node_imds_v2",
|
|
main_configs=["configs/imds_bootstrap.xml"],
|
|
env_variables={
|
|
"AWS_EC2_METADATA_SERVICE_ENDPOINT": f"http://{METADATA_SERVER_V2_HOSTNAME}:{METADATA_SERVER_V2_PORT}",
|
|
},
|
|
stay_alive=True,
|
|
)
|
|
node_config_value = cluster.add_instance(
|
|
"node_config_value",
|
|
main_configs=["configs/config_value.xml"],
|
|
)
|
|
node_file_value = cluster.add_instance(
|
|
"node_file_value",
|
|
main_configs=["configs/file_value.xml"],
|
|
stay_alive=True,
|
|
)
|
|
node_missing_value = cluster.add_instance(
|
|
"node_missing_value",
|
|
main_configs=["configs/missing_value.xml"],
|
|
)
|
|
|
|
|
|
def start_metadata_server(started_cluster):
|
|
script_dir = os.path.join(os.path.dirname(__file__), "metadata_servers")
|
|
start_mock_servers(
|
|
started_cluster,
|
|
script_dir,
|
|
[
|
|
(
|
|
"simple_server.py",
|
|
METADATA_SERVER_HOSTNAME,
|
|
METADATA_SERVER_PORT,
|
|
),
|
|
(
|
|
"imdsv2_server.py",
|
|
METADATA_SERVER_V2_HOSTNAME,
|
|
METADATA_SERVER_V2_PORT,
|
|
),
|
|
],
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="module", autouse=True)
|
|
def start_cluster():
|
|
try:
|
|
cluster.start()
|
|
start_metadata_server(cluster)
|
|
yield cluster
|
|
finally:
|
|
cluster.shutdown()
|
|
|
|
|
|
def test_placement_info_from_imds():
|
|
with open(os.path.join(os.path.dirname(__file__), "configs/imds.xml"), "r") as f:
|
|
node_imds.replace_config(
|
|
"/etc/clickhouse-server/config.d/imds_bootstrap.xml", f.read()
|
|
)
|
|
node_imds.stop_clickhouse(kill=True)
|
|
node_imds.start_clickhouse()
|
|
|
|
node_imds.query("SYSTEM FLUSH LOGS")
|
|
assert node_imds.contains_in_log(
|
|
"CloudPlacementInfo: Loaded info: availability_zone: ci-test-1a"
|
|
)
|
|
|
|
|
|
def test_placement_info_from_imdsv2_only():
|
|
# Regression test for #81402: an EC2 instance configured with
|
|
# "IMDSv2 only (token required)" rejects unauthenticated GETs
|
|
# against the metadata service. The mock server here only honors
|
|
# GETs that present a freshly-issued IMDSv2 session token.
|
|
with open(os.path.join(os.path.dirname(__file__), "configs/imds.xml"), "r") as f:
|
|
node_imds_v2.replace_config(
|
|
"/etc/clickhouse-server/config.d/imds_bootstrap.xml", f.read()
|
|
)
|
|
node_imds_v2.stop_clickhouse(kill=True)
|
|
node_imds_v2.start_clickhouse()
|
|
|
|
node_imds_v2.query("SYSTEM FLUSH LOGS")
|
|
assert node_imds_v2.contains_in_log(
|
|
"CloudPlacementInfo: Loaded info: availability_zone: ci-test-1a"
|
|
)
|
|
|
|
|
|
def test_placement_info_from_config():
|
|
node_config_value.query("SYSTEM FLUSH LOGS")
|
|
assert node_config_value.contains_in_log(
|
|
"CloudPlacementInfo: Loaded info: availability_zone: ci-test-1b"
|
|
)
|
|
|
|
|
|
def test_placement_info_from_file():
|
|
node_file_value.exec_in_container(
|
|
["bash", "-c", "echo ci-test-1c > /tmp/node-zone"]
|
|
)
|
|
|
|
node_file_value.stop_clickhouse(kill=True)
|
|
node_file_value.start_clickhouse()
|
|
|
|
node_file_value.query("SYSTEM FLUSH LOGS")
|
|
assert node_file_value.contains_in_log(
|
|
"CloudPlacementInfo: Loaded info: availability_zone: ci-test-1c"
|
|
)
|
|
|
|
|
|
def test_placement_info_missing_data():
|
|
node_missing_value.query("SYSTEM FLUSH LOGS")
|
|
assert node_missing_value.contains_in_log(
|
|
"CloudPlacementInfo: Availability zone info not found"
|
|
)
|