Rewrote IMDS emulator in Go, added unit tests.

This commit is contained in:
slawek committed 2026-10-20 09:06:57 +02:00
1 parent cedac0a6ee
commit 630375c7ed
10 files changed
+940 -248

No files matched your search

+9 -8
View File
@@ -95,20 +95,21 @@ The agent reads the certificate and key from the mounted directory, so certifica
Role: **Server administrator** Role: **Server administrator**
The agent cannot use the X.509 credentials itself. Its launcher is a statically linked binary that cannot load the `cryptography` package, which certificate-based authentication requires, and the `--creds-file` option accepts only a service account key or user credentials. The agent therefore gets its tokens from the IMDS emulator, a Compute Engine metadata server emulator in a second container. The emulator authenticates with ADC and the X.509 credentials in `/etc/sts-agent`, and serves the access tokens of `sa-test` on the metadata token endpoint. Without `GOOGLE_APPLICATION_CREDENTIALS`, ADC in the agent falls back to the metadata server, which `GCE_METADATA_HOST` and `GCE_METADATA_IP` point to the emulator. The agent cannot use the X.509 credentials itself. Its launcher is a statically linked binary that cannot load the `cryptography` package, which certificate-based authentication requires, and the `--creds-file` option accepts only a service account key or user credentials. The agent therefore gets its tokens from the IMDS emulator, a Compute Engine metadata server emulator in a second container. The emulator is a Go program that authenticates with ADC and the X.509 credentials in `/etc/sts-agent`, and serves the access tokens of `sa-test` on the metadata token endpoint. Without `GOOGLE_APPLICATION_CREDENTIALS`, ADC in the agent falls back to the metadata server, which `GCE_METADATA_HOST` and `GCE_METADATA_IP` point to the emulator.
The containers run with the following settings: The containers run with the following settings:
| Container | Setting | Value | | Container | Setting | Value |
|---|---|---| |---|---|---|
| `imds-emulator` | Image | built from `sts-agent/Dockerfile` | | `imds-emulator` | Image | built from `sts-agent/Dockerfile`: a static Go binary on `gcr.io/distroless/static-debian13` |
| | Port | `8080`, set by `IMDS_PORT` |
| | Mount | `/etc/sts-agent`, read only | | | Mount | `/etc/sts-agent`, read only |
| | Environment | `GOOGLE_APPLICATION_CREDENTIALS=/etc/sts-agent/cred.json` | | | Environment | `GOOGLE_APPLICATION_CREDENTIALS=/etc/sts-agent/cred.json` |
| | Environment | `GOOGLE_API_CERTIFICATE_CONFIG=/etc/sts-agent/certificate_config.json` | | | Environment | `GOOGLE_API_CERTIFICATE_CONFIG=/etc/sts-agent/certificate_config.json` |
| | Environment | `PROJECT_ID`, `PROJECT_NUMBER` | | | Environment | `PROJECT_ID`, `PROJECT_NUMBER` |
| `sts-agent` | Image | `gcr.io/cloud-ingest/tsop-agent:latest` | | `sts-agent` | Image | `gcr.io/cloud-ingest/tsop-agent:latest` |
| | Mount | `/mnt/ext-storage`, read only | | | Mount | `/mnt/ext-storage`, read only |
| | Environment | `GCE_METADATA_HOST=imds-emulator`, `GCE_METADATA_IP=imds-emulator` | | | Environment | `GCE_METADATA_HOST=imds-emulator:8080`, `GCE_METADATA_IP=imds-emulator:8080` |
| | Locked memory limit | `64000000` bytes | | | Locked memory limit | `64000000` bytes |
| | Arguments | `--project-id=${PROJECT_ID} --hostname=$(hostname) --agent-pool=${STS_AGENT_POOL}` | | | Arguments | `--project-id=${PROJECT_ID} --hostname=$(hostname) --agent-pool=${STS_AGENT_POOL}` |
@@ -141,15 +142,15 @@ docker compose logs imds-emulator
Requests for paths that the emulator does not serve are logged with the status 404. Requests for paths that the emulator does not serve are logged with the status 404.
`imds_test.py` in the emulator container tests the emulator. It checks the metadata endpoints and the rejection of requests without the `Metadata-Flavor` header or with `X-Forwarded-For`. It then authenticates with the Compute Engine credentials of `google-auth` against the emulator, the same way as the agent, and checks the token with the Google token information endpoint. With `--bucket`, it also lists the objects of the bucket with the token: The emulator is a static binary without dependencies on files of the system: it embeds the root certificates and the time zone data, and it is the first process (PID 1) of its container. On `SIGTERM` or `SIGINT`, which `docker compose stop` sends, it stops accepting connections, lets the requests in progress finish for up to 5 seconds, and exits with status 0. A second signal makes it exit at once.
The unit tests of the emulator, `sts-agent/main_test.go`, cover the request checks, the metadata endpoints, the token response, the extraction of the service account email from the credentials file, the shutdown, and a run of the emulator as a process stopped with real signals. They run during the build of the image, so a failing test stops it. To run them on a workstation with Go:
```bash ```bash
source ../env.sh cd sts-agent
docker compose exec imds-emulator /venv/bin/python3 /app/imds_test.py --bucket "${BUCKET}" go test -race ./...
``` ```
It prints `PASS` or `FAIL` for each check and exits with status 1 when a check fails.
The source directory is mounted at the same path as on the host, because the transfer job refers to it by that path. The source directory is mounted at the same path as on the host, because the transfer job refers to it by that path.
The connected agents of the pool are listed on the Agent pools page of the Google Cloud console. An agent appears there only after it has authenticated. The connected agents of the pool are listed on the Agent pools page of the Google Cloud console. An agent appears there only after it has authenticated.
+1 -1
View File
@@ -151,7 +151,7 @@ sudo install -d -m 0750 /etc/vault-cert
} }
``` ```
The `caddy` instance sets the signing path of the `web-server` role and a hook. The `wif` instance uses the default signing path and has no hook. It keeps the certificate in `/etc/sts-agent`, the credential directory of the Storage Transfer agent that authenticates with it (`docs/storage-transfer.md`). The `caddy` instance sets the signing path of the `web-server` role and a hook. The `wif` instance uses the default signing path and has no hook. It keeps the certificate in `/etc/sts-agent`, the credential directory of that agent (`docs/storage-transfer.md`).
```bash ```bash
sudo /usr/local/sbin/vault-unwrap-secret-id /etc/vault-cert/caddy.json sudo /usr/local/sbin/vault-unwrap-secret-id /etc/vault-cert/caddy.json
+10 -9
View File
@@ -1,13 +1,14 @@
FROM python:3.13-slim-trixie AS build FROM golang:1.27-trixie AS build
# The interpreter path in the virtualenv must match the path in the distroless image. WORKDIR /src
RUN ln -s /usr/local/bin/python /usr/bin/python \ COPY go.mod go.sum ./
&& /usr/bin/python -m venv /venv \ RUN go mod download
&& /venv/bin/pip install --no-cache-dir "google-auth[requests,cryptography]>=2.39.0" COPY main.go main_test.go ./
RUN CGO_ENABLED=0 go test ./... \
&& CGO_ENABLED=0 go build -trimpath -ldflags="-s -w" -o /imds-emulator .
FROM gcr.io/distroless/python3-debian13 FROM gcr.io/distroless/static-debian13
COPY --from=build /venv /venv COPY --from=build /imds-emulator /imds-emulator
COPY imds_emulator.py imds_test.py /app/
ENTRYPOINT ["/venv/bin/python3", "/app/imds_emulator.py"] ENTRYPOINT ["/imds-emulator"]
+2 -2
View File
@@ -19,8 +19,8 @@ services:
volumes: volumes:
- /mnt/ext-storage:/mnt/ext-storage:ro - /mnt/ext-storage:/mnt/ext-storage:ro
environment: environment:
GCE_METADATA_HOST: imds-emulator GCE_METADATA_HOST: imds-emulator:8080
GCE_METADATA_IP: imds-emulator GCE_METADATA_IP: imds-emulator:8080
command: command:
- --project-id=${PROJECT_ID} - --project-id=${PROJECT_ID}
- --hostname=${HOSTNAME} - --hostname=${HOSTNAME}
+15
View File
@@ -0,0 +1,15 @@
module imds-emulator
go 1.27.1
require (
cloud.google.com/go/auth v0.24.1
golang.org/x/crypto/x509roots/fallback v0.0.0-20261016151342-0ae6809b7559
)
require (
cloud.google.com/go/compute/metadata v0.9.0 // indirect
github.com/googleapis/enterprise-certificate-proxy v0.3.17 // indirect
github.com/googleapis/gax-go/v2 v2.26.2 // indirect
golang.org/x/sys v0.47.0 // indirect
)
+49
View File
@@ -0,0 +1,49 @@
cloud.google.com/go/auth v0.24.1 h1:NTmd2iJb4DsGIav0KiqFMYI88nMshA2wnFMBatRmoCs=
cloud.google.com/go/auth v0.24.1/go.mod h1:1dmBXa9xfBvBr39ZnDUDkaaA91ccq9CcqENDgJA4d7c=
cloud.google.com/go/compute/metadata v0.9.0 h1:pDUj4QMoPejqq20dK0Pg2N4yG9zIkYGdBtwLoEkH9Zs=
cloud.google.com/go/compute/metadata v0.9.0/go.mod h1:E0bWwX5wTnLPedCKqk3pJmVgCBSM6qQI1yTBdEb3C10=
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg=
github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
github.com/go-logr/logr v1.4.4 h1:tG4xh9yMsRCAiodLVTxyrkzSZ9+o0L1Kg/+cPVcbP/8=
github.com/go-logr/logr v1.4.4/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
github.com/google/s2a-go v0.1.9 h1:LGD7gtMgezd8a/Xak7mEWL0PjoTQFvpRudN895yqKW0=
github.com/google/s2a-go v0.1.9/go.mod h1:YA0Ei2ZQL3acow2O62kdp9UlnvMmU7kA6Eutn0dXayM=
github.com/googleapis/enterprise-certificate-proxy v0.3.17 h1:73NfMHdiqo9JFU9+7a5ExpVa10/R29pXfZIaW559nrg=
github.com/googleapis/enterprise-certificate-proxy v0.3.17/go.mod h1:rSEsBUemEBZEexP2y6jPp16LUmUbjmSbcPMQizR0o4k=
github.com/googleapis/gax-go/v2 v2.26.2 h1:ydkmNXxj7bEmmeK5AihkKnWxyOyBR9TDebvp5L5izk8=
github.com/googleapis/gax-go/v2 v2.26.2/go.mod h1:sMKqnMesnKH+3wiRJROcttA+cJoZoGbZl1vDQ8XYtGk=
go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.67.0 h1:OyrsyzuttWTSur2qN/Lm0m2a8yqyIjUVBZcxFPuXq2o=
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.67.0/go.mod h1:C2NGBr+kAB4bk3xtMXfZ94gqFDtg/GkI7e9zqGh5Beg=
go.opentelemetry.io/otel v1.45.0 h1:pdrWmLHofpubmArBv1LgFSv1Z0Ie/ppdZzu+kUN5EeU=
go.opentelemetry.io/otel v1.45.0/go.mod h1:XZxIqPapzEYnhNSScF5DIqXhm/rYi0FzCe2XddAwZfQ=
go.opentelemetry.io/otel/metric v1.45.0 h1:7Eg1uH7CJ5cXv9is6tnBe1FI6rj1nwUdbFypRm3br/M=
go.opentelemetry.io/otel/metric v1.45.0/go.mod h1:HAPbm1nd3p1PmFH7v2dR+6BjXxw+Lq4a2+pndMAm08s=
go.opentelemetry.io/otel/trace v1.45.0 h1:l/mP6Uv7oNO7/TblbhpbgMidxhq1uO/rPsikOyVhxag=
go.opentelemetry.io/otel/trace v1.45.0/go.mod h1:qoJJA2xNMnxRrdISU/kLtfUH2wNeQbiv+jhs/CxI8bc=
golang.org/x/crypto v0.55.0 h1:+KWHjbgOaAQ66dh/YlkZKHlz9ZUlq61AFirAR9ntP8M=
golang.org/x/crypto v0.55.0/go.mod h1:uq0V9dE/fzQuJtbnL+2EhWOE63vo164FY8xqEnV9xis=
golang.org/x/crypto/x509roots/fallback v0.0.0-20261016151342-0ae6809b7559 h1:M68k3yCPiXhrSyh9Twe3Kv7nOsvUumZDN/3kdSwUqw0=
golang.org/x/crypto/x509roots/fallback v0.0.0-20261016151342-0ae6809b7559/go.mod h1:HPze8vhfG6fO06AM+VSvxRm4E3+5Yk375mgrJ5M2z1E=
golang.org/x/net v0.58.0 h1:ynWG7rqYi4ccpTEuPZ2QGWHktVEM9DMCj9yzDE0Q7To=
golang.org/x/net v0.58.0/go.mod h1:YwCddHnFlT7eLQqVprV19OnhLGtc5xOKgE0RyqgfWAU=
golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8=
golang.org/x/text v0.41.0/go.mod h1:jvf1O8ajNzZqhSrQBPbutR/EB83Cc0CFrezNQIwbb5M=
google.golang.org/api v0.288.0 h1:glhO/J88obKP5I269W3hB73dvBKrjU56ZfmNlNXpgTU=
google.golang.org/api v0.288.0/go.mod h1:lM2kYRzYUCBY91P9h6VF1PYmvhxii3O5hji37qRvIcY=
google.golang.org/genproto v0.0.0-20260715232425-e75dac1f907d h1:C9v1o0/4quuhOAfmRXA2j+we0PqZIp8traLdeogF3Ms=
google.golang.org/genproto/googleapis/rpc v0.0.0-20260715232425-e75dac1f907d h1:Jkpk39hlTZOIp3RbfvNX9R8Hv+Sw0X89nlU/xFOErsc=
google.golang.org/genproto/googleapis/rpc v0.0.0-20260715232425-e75dac1f907d/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8=
google.golang.org/grpc v1.83.2 h1:EManeRomTObA0BU7I8vXgg/78uE5MJ9M8B39EX2WscU=
google.golang.org/grpc v1.83.2/go.mod h1:YPI1hK3kDked6iHvgX3tR0y+nX/qpMFKhPgFsokw1S8=
google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE=
google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
-127
View File
@@ -1,127 +0,0 @@
#!/usr/bin/env python3
"""
Compute Engine metadata server (IMDS) emulator.
Serves access tokens of the credentials found by Application Default
Credentials (ADC) on the IMDS token endpoint, together with the metadata
that Google client libraries read when they authenticate on Compute Engine.
"""
import datetime
import http.server
import json
import os
import socket
import sys
import threading
import urllib.parse
import google.auth
import google.auth.credentials
import google.auth.transport.requests
DEFAULT_SCOPE = "https://www.googleapis.com/auth/cloud-platform"
ROOT = "/computeMetadata/v1/"
port = int(os.environ.get("IMDS_PORT", "80"))
hostname = os.environ.get("IMDS_HOSTNAME", socket.gethostname())
project_number = os.environ.get("PROJECT_NUMBER", "")
base_credentials, adc_project_id = google.auth.default(scopes=[DEFAULT_SCOPE])
project_id = os.environ.get("PROJECT_ID") or adc_project_id or ""
email = getattr(base_credentials, "service_account_email", None) or os.environ.get("SERVICE_ACCOUNT_EMAIL", "")
if not email:
sys.exit("The service account email is unknown, set SERVICE_ACCOUNT_EMAIL")
credentials_by_scopes = {}
lock = threading.Lock()
def token(scopes):
"""
Returns the access token and its remaining lifetime in seconds for the
scopes, refreshing the cached credentials of the scopes when needed.
"""
key = tuple(sorted(scopes))
with lock:
credentials = credentials_by_scopes.get(key)
if credentials is None:
credentials = google.auth.credentials.with_scopes_if_required(base_credentials, list(key))
credentials_by_scopes[key] = credentials
if not credentials.valid:
credentials.refresh(google.auth.transport.requests.Request())
expires_in = int((credentials.expiry - datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)).total_seconds())
return credentials.token, expires_in
class Handler(http.server.BaseHTTPRequestHandler):
"""
Answers the IMDS paths used by the Google client libraries, and 404 for
any other path.
"""
server_version = "Metadata Server for VM"
sys_version = ""
def reply(self, status, body, content_type="application/text"):
data = body.encode() if isinstance(body, str) else json.dumps(body).encode()
self.send_response(status)
self.send_header("Metadata-Flavor", "Google")
self.send_header("Content-Type", content_type)
self.send_header("Content-Length", str(len(data)))
self.end_headers()
self.wfile.write(data)
def do_GET(self):
url = urllib.parse.urlsplit(self.path)
query = urllib.parse.parse_qs(url.query)
# The real IMDS rejects requests without the flavor header and proxied requests.
if self.headers.get("Metadata-Flavor") != "Google" or self.headers.get("X-Forwarded-For"):
return self.reply(403, "Forbidden")
if url.path in ("/", ROOT):
return self.reply(200, "")
if not url.path.startswith(ROOT):
return self.reply(404, "Not Found")
path = url.path[len(ROOT):]
# Both "default" and the email address name the service account.
for alias in ("default", email):
prefix = f"instance/service-accounts/{alias}/"
if path.startswith(prefix):
path = "instance/service-accounts/default/" + path[len(prefix):]
break
account = {"aliases": ["default"], "email": email, "scopes": [DEFAULT_SCOPE]}
if path == "instance/service-accounts/default/token":
scopes = query.get("scopes", [DEFAULT_SCOPE])[0].split(",")
access_token, expires_in = token(scopes)
return self.reply(200, {"access_token": access_token, "expires_in": expires_in, "token_type": "Bearer"}, "application/json")
if path == "instance/service-accounts/default/email":
return self.reply(200, email)
if path == "instance/service-accounts/default/scopes":
return self.reply(200, DEFAULT_SCOPE + "\n")
if path == "instance/service-accounts/default/" and query.get("recursive") == ["true"]:
return self.reply(200, account, "application/json")
if path == "instance/service-accounts/" and query.get("recursive") == ["true"]:
return self.reply(200, {"default": account, email: account}, "application/json")
if path == "instance/service-accounts/":
return self.reply(200, f"default/\n{email}/\n")
if path == "project/project-id" and project_id:
return self.reply(200, project_id)
if path == "project/numeric-project-id" and project_number:
return self.reply(200, project_number)
if path == "instance/hostname":
return self.reply(200, hostname)
if path == "universe/universe-domain":
return self.reply(200, "googleapis.com")
return self.reply(404, "Not Found")
def log_request(self, code="-", size="-"):
print(f"{self.command} {urllib.parse.urlsplit(self.path).path} {code}", flush=True)
print(f"Serving tokens of {email} for project {project_id or '-'} on port {port}", flush=True)
http.server.ThreadingHTTPServer(("", port), Handler).serve_forever()
-101
View File
@@ -1,101 +0,0 @@
#!/usr/bin/env python3
"""
Tests the IMDS emulator: the metadata endpoints, the request checks, the
Compute Engine credentials of google-auth on top of the emulator, and the
validity of the issued access token at Google.
"""
import argparse
import json
import os
import sys
import urllib.error
import urllib.parse
import urllib.request
parser = argparse.ArgumentParser(description="Test the IMDS emulator.")
parser.add_argument("--host", default="localhost", help="host and optional port of the emulator")
parser.add_argument("--bucket", help="bucket to list with the token, optional")
args = parser.parse_args()
root = f"http://{args.host}/computeMetadata/v1/"
failures = 0
def check(name, condition, detail=""):
"""
Prints the result of a check and counts the failures.
"""
global failures
print(f"{'PASS' if condition else 'FAIL'} {name}{': ' + detail if detail else ''}")
if not condition:
failures += 1
def get(url, headers=None):
"""
Returns the status, the headers and the body of a GET request.
"""
request = urllib.request.Request(url, headers=headers or {})
try:
with urllib.request.urlopen(request, timeout=10) as response:
return response.status, response.headers, response.read().decode()
except urllib.error.HTTPError as error:
return error.code, error.headers, error.read().decode()
flavor = {"Metadata-Flavor": "Google"}
# Detection ping and request checks
status, headers, _ = get(f"http://{args.host}/", flavor)
check("ping", status == 200 and headers.get("Metadata-Flavor") == "Google", f"status {status}")
status, _, _ = get(root + "project/project-id")
check("request without Metadata-Flavor is rejected", status == 403, f"status {status}")
status, _, _ = get(root + "project/project-id", {**flavor, "X-Forwarded-For": "192.0.2.1"})
check("request with X-Forwarded-For is rejected", status == 403, f"status {status}")
# Metadata
status, _, project_id = get(root + "project/project-id", flavor)
check("project ID", status == 200 and project_id != "", project_id)
status, _, email = get(root + "instance/service-accounts/default/email", flavor)
check("service account email", status == 200 and "@" in email, email)
status, _, body = get(root + "instance/service-accounts/default/?recursive=true", flavor)
check("service account info", status == 200 and json.loads(body).get("email") == email, body)
status, _, body = get(root + f"instance/service-accounts/{email}/email", flavor)
check("service account by email", status == 200 and body == email, body)
# Token
status, _, body = get(root + "instance/service-accounts/default/token", flavor)
token = json.loads(body) if status == 200 else {}
check("token", status == 200 and token.get("access_token") and token.get("expires_in", 0) > 0, f"expires in {token.get('expires_in')} s")
# Compute Engine credentials of google-auth, the path used by the STS agent
for name in ("GOOGLE_APPLICATION_CREDENTIALS", "GOOGLE_API_CERTIFICATE_CONFIG"):
os.environ.pop(name, None)
os.environ["GCE_METADATA_HOST"] = args.host
os.environ["GCE_METADATA_IP"] = args.host
import google.auth
import google.auth.transport.requests
credentials, adc_project_id = google.auth.default(scopes=["https://www.googleapis.com/auth/cloud-platform"])
credentials.refresh(google.auth.transport.requests.Request())
check("google-auth uses Compute Engine credentials", type(credentials).__module__ == "google.auth.compute_engine.credentials", type(credentials).__module__)
check("google-auth project ID", adc_project_id == project_id, str(adc_project_id))
check("google-auth service account", credentials.service_account_email == email, credentials.service_account_email)
check("google-auth token", credentials.valid, f"expires {credentials.expiry} UTC")
# Validity of the token at Google
status, _, body = get("https://oauth2.googleapis.com/tokeninfo?" + urllib.parse.urlencode({"access_token": credentials.token}))
info = json.loads(body) if status == 200 else {}
check("token accepted by Google", status == 200, f"scope {info.get('scope')}, expires in {info.get('expires_in')} s")
# Use of the token with Cloud Storage
if args.bucket:
status, _, body = get(
f"https://storage.googleapis.com/storage/v1/b/{urllib.parse.quote(args.bucket)}/o?maxResults=1",
{"Authorization": f"Bearer {credentials.token}"},
)
check(f"list objects of {args.bucket}", status == 200, f"status {status}")
print(f"{failures} check(s) failed" if failures else "All checks passed")
sys.exit(1 if failures else 0)
+292
View File
@@ -0,0 +1,292 @@
// Command imds-emulator serves the access tokens of the Application Default
// Credentials on the Compute Engine metadata server (IMDS) endpoints.
package main
import (
"cmp"
"context"
"encoding/json"
"io"
"log"
"net"
"net/http"
"os"
"os/signal"
"regexp"
"slices"
"strings"
"sync"
"syscall"
"time"
_ "time/tzdata" // Time zones for the log timestamps without the files of the system.
"cloud.google.com/go/auth"
"cloud.google.com/go/auth/credentials"
_ "golang.org/x/crypto/x509roots/fallback" // Root certificates without the files of the system.
)
const (
defaultScope = "https://www.googleapis.com/auth/cloud-platform"
root = "/computeMetadata/v1/"
accounts = "instance/service-accounts/"
defaultAccount = accounts + "default/"
// shutdownGrace is shorter than the 10 seconds Docker waits after SIGTERM.
shutdownGrace = 5 * time.Second
)
// tokenFunc returns an access token for the scopes.
type tokenFunc func(ctx context.Context, scopes []string) (*auth.Token, error)
// serviceAccount is the metadata of a service account.
type serviceAccount struct {
Aliases []string `json:"aliases"`
Email string `json:"email"`
Scopes []string `json:"scopes"`
}
// metadataServer answers the IMDS requests of the Google client libraries.
type metadataServer struct {
email string
projectID string
projectNumber string
hostname string
token tokenFunc
now func() time.Time
}
// ServeHTTP answers a request and logs its method, path and status.
func (s *metadataServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Metadata-Flavor", "Google")
w.Header().Set("Server", "Metadata Server for VM")
status := s.handle(w, r)
log.Printf("%s %s %d", r.Method, r.URL.Path, status)
}
// handle writes the response of a request and returns its status.
func (s *metadataServer) handle(w http.ResponseWriter, r *http.Request) int {
if r.Method != http.MethodGet && r.Method != http.MethodHead {
return text(w, http.StatusMethodNotAllowed, "Method Not Allowed")
}
// The real IMDS rejects requests without the flavor header and proxied requests.
if r.Header.Get("Metadata-Flavor") != "Google" || r.Header.Get("X-Forwarded-For") != "" {
return text(w, http.StatusForbidden, "Forbidden")
}
if r.URL.Path == "/" || r.URL.Path == root {
return text(w, http.StatusOK, "")
}
path, found := strings.CutPrefix(r.URL.Path, root)
if !found {
return text(w, http.StatusNotFound, "Not Found")
}
// Both "default" and the email address name the service account.
for _, alias := range []string{"default", s.email} {
if rest, found := strings.CutPrefix(path, accounts+alias+"/"); found {
path = defaultAccount + rest
break
}
}
query := r.URL.Query()
recursive := query.Get("recursive") == "true"
account := serviceAccount{Aliases: []string{"default"}, Email: s.email, Scopes: []string{defaultScope}}
switch {
case path == defaultAccount+"token":
scopes := []string{defaultScope}
if value := query.Get("scopes"); value != "" {
scopes = strings.Split(value, ",")
}
token, err := s.token(r.Context(), scopes)
if err != nil {
log.Printf("token error: %v", err)
return text(w, http.StatusInternalServerError, "Internal Server Error")
}
expiresIn := 3600
if !token.Expiry.IsZero() {
expiresIn = max(int(token.Expiry.Sub(s.now()).Seconds()), 0)
}
return jsonBody(w, map[string]any{
"access_token": token.Value,
"expires_in": expiresIn,
"token_type": cmp.Or(token.Type, "Bearer"),
})
case path == defaultAccount+"email":
return text(w, http.StatusOK, s.email)
case path == defaultAccount+"scopes":
return text(w, http.StatusOK, defaultScope+"\n")
case path == defaultAccount && recursive:
return jsonBody(w, account)
case path == defaultAccount:
return text(w, http.StatusOK, "aliases\nemail\nscopes\ntoken\n")
case path == accounts && recursive:
return jsonBody(w, map[string]serviceAccount{"default": account, s.email: account})
case path == accounts:
return text(w, http.StatusOK, "default/\n"+s.email+"/\n")
case path == "project/project-id" && s.projectID != "":
return text(w, http.StatusOK, s.projectID)
case path == "project/numeric-project-id" && s.projectNumber != "":
return text(w, http.StatusOK, s.projectNumber)
case path == "instance/hostname":
return text(w, http.StatusOK, s.hostname)
case path == "universe/universe-domain":
return text(w, http.StatusOK, "googleapis.com")
}
return text(w, http.StatusNotFound, "Not Found")
}
// text writes a plain text response and returns its status.
func text(w http.ResponseWriter, status int, body string) int {
w.Header().Set("Content-Type", "application/text")
w.WriteHeader(status)
io.WriteString(w, body)
return status
}
// jsonBody writes a JSON response with the status 200 and returns it.
func jsonBody(w http.ResponseWriter, value any) int {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
json.NewEncoder(w).Encode(value)
return http.StatusOK
}
var impersonationURL = regexp.MustCompile(`/serviceAccounts/([^/:]+):generateAccessToken$`)
// serviceAccountEmail returns the service account of a credentials file, or an empty string.
func serviceAccountEmail(credentialsJSON []byte) string {
var file struct {
ClientEmail string `json:"client_email"`
ImpersonationURL string `json:"service_account_impersonation_url"`
}
if json.Unmarshal(credentialsJSON, &file) != nil {
return ""
}
if match := impersonationURL.FindStringSubmatch(file.ImpersonationURL); match != nil {
return match[1]
}
return file.ClientEmail
}
// providerCache keeps one token provider per set of scopes.
type providerCache struct {
mu sync.Mutex
detect func(scopes []string) (auth.TokenProvider, error)
providers map[string]auth.TokenProvider
}
// get returns the token provider of the scopes, creating it on first use.
func (c *providerCache) get(scopes []string) (auth.TokenProvider, error) {
c.mu.Lock()
defer c.mu.Unlock()
key := strings.Join(scopes, ",")
if provider, found := c.providers[key]; found {
return provider, nil
}
provider, err := c.detect(scopes)
if err != nil {
return nil, err
}
if c.providers == nil {
c.providers = map[string]auth.TokenProvider{}
}
c.providers[key] = provider
return provider, nil
}
// token returns a token for the scopes, independent of their order.
func (c *providerCache) token(ctx context.Context, scopes []string) (*auth.Token, error) {
sorted := slices.Sorted(slices.Values(scopes))
provider, err := c.get(sorted)
if err != nil {
return nil, err
}
return provider.Token(ctx)
}
// serve serves the listener until the context is canceled, then stops the
// server, waiting up to the grace period for the requests in progress.
func serve(ctx context.Context, server *http.Server, listener net.Listener, grace time.Duration) error {
failed := make(chan error, 1)
go func() { failed <- server.Serve(listener) }()
select {
case err := <-failed:
return err
case <-ctx.Done():
}
shutdownCtx, cancel := context.WithTimeout(context.Background(), grace)
defer cancel()
if err := server.Shutdown(shutdownCtx); err != nil {
server.Close()
return err
}
<-failed
return nil
}
func main() {
// The first signal stops the server gracefully, the second one exits at once.
signals := make(chan os.Signal, 2)
signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
go func() {
log.Printf("Received %v, shutting down", <-signals)
cancel()
log.Printf("Received %v again, exiting", <-signals)
os.Exit(1)
}()
detect := func(scopes []string) (auth.TokenProvider, error) {
creds, err := credentials.DetectDefault(&credentials.DetectOptions{Scopes: scopes})
if err != nil {
return nil, err
}
return creds, nil
}
defaults, err := credentials.DetectDefault(&credentials.DetectOptions{Scopes: []string{defaultScope}})
if err != nil {
log.Fatalf("Cannot find the Application Default Credentials: %v", err)
}
email := cmp.Or(os.Getenv("SERVICE_ACCOUNT_EMAIL"), serviceAccountEmail(defaults.JSON()))
if email == "" {
log.Fatal("The service account email is unknown, set SERVICE_ACCOUNT_EMAIL")
}
projectID := os.Getenv("PROJECT_ID")
if projectID == "" {
projectID, _ = defaults.ProjectID(ctx)
}
hostname, _ := os.Hostname()
server := &metadataServer{
email: email,
projectID: projectID,
projectNumber: os.Getenv("PROJECT_NUMBER"),
hostname: cmp.Or(os.Getenv("IMDS_HOSTNAME"), hostname),
token: (&providerCache{detect: detect}).token,
now: time.Now,
}
httpServer := &http.Server{
Handler: server,
ReadHeaderTimeout: 10 * time.Second,
WriteTimeout: 60 * time.Second,
IdleTimeout: 60 * time.Second,
}
listener, err := net.Listen("tcp", ":"+cmp.Or(os.Getenv("IMDS_PORT"), "8080"))
if err != nil {
log.Fatalf("Cannot listen: %v", err)
}
log.Printf("Serving tokens of %s for project %s on %s", email, cmp.Or(projectID, "-"), listener.Addr())
if err := serve(ctx, httpServer, listener, shutdownGrace); err != nil {
log.Fatalf("Server stopped: %v", err)
}
log.Print("Stopped")
}
+562
View File
@@ -0,0 +1,562 @@
package main
import (
"bytes"
"context"
"crypto/rand"
"crypto/rsa"
"crypto/x509"
"encoding/json"
"encoding/pem"
"errors"
"io"
"log"
"net"
"net/http"
"net/http/httptest"
"os"
"os/exec"
"path/filepath"
"runtime"
"slices"
"strings"
"sync"
"sync/atomic"
"syscall"
"testing"
"time"
"cloud.google.com/go/auth"
)
const (
testEmail = "sa-test@dom-lab.iam.gserviceaccount.com"
mainEnv = "IMDS_EMULATOR_RUN_MAIN"
)
var testNow = time.Date(2026, 10, 20, 12, 0, 0, 0, time.UTC)
func TestMain(m *testing.M) {
// The signal tests run the test binary as the emulator.
if os.Getenv(mainEnv) == "1" {
main()
return
}
log.SetOutput(io.Discard)
os.Exit(m.Run())
}
// tokenProviderFunc adapts a function to auth.TokenProvider.
type tokenProviderFunc func(ctx context.Context) (*auth.Token, error)
func (f tokenProviderFunc) Token(ctx context.Context) (*auth.Token, error) { return f(ctx) }
// newTestServer returns a server whose tokens are named after the scopes and valid for an hour.
func newTestServer(scopesSeen *[]string) *metadataServer {
return &metadataServer{
email: testEmail,
projectID: "dom-lab",
projectNumber: "267098479356",
hostname: "agent-host",
now: func() time.Time { return testNow },
token: func(ctx context.Context, scopes []string) (*auth.Token, error) {
if scopesSeen != nil {
*scopesSeen = scopes
}
return &auth.Token{Value: "token-for-" + strings.Join(scopes, ","), Type: "Bearer", Expiry: testNow.Add(time.Hour)}, nil
},
}
}
// get sends a request with the metadata flavor header to the server.
func get(server *metadataServer, target string) *httptest.ResponseRecorder {
request := httptest.NewRequest(http.MethodGet, target, nil)
request.Header.Set("Metadata-Flavor", "Google")
recorder := httptest.NewRecorder()
server.ServeHTTP(recorder, request)
return recorder
}
func TestRequestChecks(t *testing.T) {
tests := []struct {
name string
method string
headers map[string]string
want int
}{
{"valid", http.MethodGet, map[string]string{"Metadata-Flavor": "Google"}, http.StatusOK},
{"head", http.MethodHead, map[string]string{"Metadata-Flavor": "Google"}, http.StatusOK},
{"missing flavor", http.MethodGet, nil, http.StatusForbidden},
{"wrong flavor", http.MethodGet, map[string]string{"Metadata-Flavor": "Other"}, http.StatusForbidden},
{"forwarded", http.MethodGet, map[string]string{"Metadata-Flavor": "Google", "X-Forwarded-For": "192.0.2.1"}, http.StatusForbidden},
{"post", http.MethodPost, map[string]string{"Metadata-Flavor": "Google"}, http.StatusMethodNotAllowed},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
request := httptest.NewRequest(test.method, root+"project/project-id", nil)
for name, value := range test.headers {
request.Header.Set(name, value)
}
recorder := httptest.NewRecorder()
newTestServer(nil).ServeHTTP(recorder, request)
if recorder.Code != test.want {
t.Errorf("status = %d, want %d", recorder.Code, test.want)
}
if got := recorder.Header().Get("Metadata-Flavor"); got != "Google" {
t.Errorf("response Metadata-Flavor = %q, want Google", got)
}
})
}
}
func TestEndpoints(t *testing.T) {
tests := []struct {
name string
target string
status int
body string
}{
{"ping", "/", 200, ""},
{"root", root, 200, ""},
{"project ID", root + "project/project-id", 200, "dom-lab"},
{"project number", root + "project/numeric-project-id", 200, "267098479356"},
{"hostname", root + "instance/hostname", 200, "agent-host"},
{"universe domain", root + "universe/universe-domain", 200, "googleapis.com"},
{"email", root + "instance/service-accounts/default/email", 200, testEmail},
{"email by address", root + "instance/service-accounts/" + testEmail + "/email", 200, testEmail},
{"scopes", root + "instance/service-accounts/default/scopes", 200, defaultScope + "\n"},
{"account list", root + "instance/service-accounts/", 200, "default/\n" + testEmail + "/\n"},
{"account contents", root + "instance/service-accounts/default/", 200, "aliases\nemail\nscopes\ntoken\n"},
{"other account", root + "instance/service-accounts/other@dom-lab.iam.gserviceaccount.com/email", 404, "Not Found"},
{"unknown path", root + "instance/zone", 404, "Not Found"},
{"outside root", "/other", 404, "Not Found"},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
recorder := get(newTestServer(nil), test.target)
if recorder.Code != test.status || recorder.Body.String() != test.body {
t.Errorf("got %d %q, want %d %q", recorder.Code, recorder.Body.String(), test.status, test.body)
}
})
}
}
func TestMissingProject(t *testing.T) {
server := newTestServer(nil)
server.projectID, server.projectNumber = "", ""
for _, name := range []string{"project/project-id", "project/numeric-project-id"} {
if recorder := get(server, root+name); recorder.Code != http.StatusNotFound {
t.Errorf("%s status = %d, want 404", name, recorder.Code)
}
}
}
func TestAccountInfo(t *testing.T) {
want := serviceAccount{Aliases: []string{"default"}, Email: testEmail, Scopes: []string{defaultScope}}
recorder := get(newTestServer(nil), root+"instance/service-accounts/default/?recursive=true")
var account serviceAccount
if err := json.Unmarshal(recorder.Body.Bytes(), &account); err != nil {
t.Fatalf("invalid JSON %q: %v", recorder.Body.String(), err)
}
if recorder.Code != 200 || !slices.Equal(account.Aliases, want.Aliases) || account.Email != want.Email || !slices.Equal(account.Scopes, want.Scopes) {
t.Errorf("got %d %+v, want 200 %+v", recorder.Code, account, want)
}
if got := recorder.Header().Get("Content-Type"); got != "application/json" {
t.Errorf("Content-Type = %q, want application/json", got)
}
recorder = get(newTestServer(nil), root+"instance/service-accounts/?recursive=true")
var listed map[string]serviceAccount
if err := json.Unmarshal(recorder.Body.Bytes(), &listed); err != nil {
t.Fatalf("invalid JSON %q: %v", recorder.Body.String(), err)
}
if len(listed) != 2 || listed["default"].Email != testEmail || listed[testEmail].Email != testEmail {
t.Errorf("got %+v, want default and %s", listed, testEmail)
}
}
func TestToken(t *testing.T) {
tests := []struct {
name string
target string
wantScopes []string
wantToken string
}{
{"default scope", root + "instance/service-accounts/default/token", []string{defaultScope}, "token-for-" + defaultScope},
{"scopes", root + "instance/service-accounts/default/token?scopes=a,b", []string{"a", "b"}, "token-for-a,b"},
{"by address", root + "instance/service-accounts/" + testEmail + "/token", []string{defaultScope}, "token-for-" + defaultScope},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
var scopes []string
recorder := get(newTestServer(&scopes), test.target)
var response struct {
AccessToken string `json:"access_token"`
ExpiresIn int `json:"expires_in"`
TokenType string `json:"token_type"`
}
if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil {
t.Fatalf("invalid JSON %q: %v", recorder.Body.String(), err)
}
if recorder.Code != 200 || response.AccessToken != test.wantToken || response.ExpiresIn != 3600 || response.TokenType != "Bearer" {
t.Errorf("got %d %+v", recorder.Code, response)
}
if !slices.Equal(scopes, test.wantScopes) {
t.Errorf("scopes = %v, want %v", scopes, test.wantScopes)
}
})
}
}
func TestTokenExpiry(t *testing.T) {
tests := []struct {
name string
token auth.Token
wantType string
wantIn int
}{
{"remaining lifetime", auth.Token{Value: "t", Expiry: testNow.Add(90 * time.Second)}, "Bearer", 90},
{"no expiry", auth.Token{Value: "t"}, "Bearer", 3600},
{"expired", auth.Token{Value: "t", Expiry: testNow.Add(-time.Minute)}, "Bearer", 0},
{"token type", auth.Token{Value: "t", Type: "MAC", Expiry: testNow.Add(time.Minute)}, "MAC", 60},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
server := newTestServer(nil)
server.token = func(context.Context, []string) (*auth.Token, error) { return &test.token, nil }
var response struct {
ExpiresIn int `json:"expires_in"`
TokenType string `json:"token_type"`
}
recorder := get(server, root+"instance/service-accounts/default/token")
if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil {
t.Fatalf("invalid JSON %q: %v", recorder.Body.String(), err)
}
if response.ExpiresIn != test.wantIn || response.TokenType != test.wantType {
t.Errorf("got expires_in %d token_type %q, want %d %q", response.ExpiresIn, response.TokenType, test.wantIn, test.wantType)
}
})
}
}
func TestTokenError(t *testing.T) {
server := newTestServer(nil)
server.token = func(context.Context, []string) (*auth.Token, error) { return nil, errors.New("secret detail") }
recorder := get(server, root+"instance/service-accounts/default/token")
if recorder.Code != http.StatusInternalServerError {
t.Errorf("status = %d, want 500", recorder.Code)
}
if strings.Contains(recorder.Body.String(), "secret detail") {
t.Errorf("the response %q leaks the error", recorder.Body.String())
}
}
func TestServiceAccountEmail(t *testing.T) {
tests := []struct {
name string
json string
want string
}{
{"external account", `{"type":"external_account","service_account_impersonation_url":"https://iamcredentials.mtls.googleapis.com/v1/projects/-/serviceAccounts/sa-test@dom-lab.iam.gserviceaccount.com:generateAccessToken"}`, testEmail},
{"impersonated account", `{"type":"impersonated_service_account","service_account_impersonation_url":"https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/sa-test@dom-lab.iam.gserviceaccount.com:generateAccessToken"}`, testEmail},
{"service account key", `{"type":"service_account","client_email":"sa-test@dom-lab.iam.gserviceaccount.com"}`, testEmail},
{"impersonation wins", `{"client_email":"key@p.iam.gserviceaccount.com","service_account_impersonation_url":"https://x/serviceAccounts/sa-test@dom-lab.iam.gserviceaccount.com:generateAccessToken"}`, testEmail},
{"external account without impersonation", `{"type":"external_account"}`, ""},
{"unexpected impersonation URL", `{"service_account_impersonation_url":"https://x/other"}`, ""},
{"invalid JSON", `{`, ""},
{"empty", ``, ""},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
if got := serviceAccountEmail([]byte(test.json)); got != test.want {
t.Errorf("got %q, want %q", got, test.want)
}
})
}
}
func TestProviderCache(t *testing.T) {
var detected atomic.Int32
var failures atomic.Int32
cache := &providerCache{detect: func(scopes []string) (auth.TokenProvider, error) {
detected.Add(1)
if slices.Equal(scopes, []string{"fail"}) && failures.Add(1) == 1 {
return nil, errors.New("detection failed")
}
return tokenProviderFunc(func(context.Context) (*auth.Token, error) {
return &auth.Token{Value: strings.Join(scopes, ",")}, nil
}), nil
}}
ctx := context.Background()
for _, scopes := range [][]string{{"a", "b"}, {"b", "a"}, {"a", "b"}} {
token, err := cache.token(ctx, scopes)
if err != nil || token.Value != "a,b" {
t.Fatalf("token(%v) = %v, %v; want a,b", scopes, token, err)
}
}
if got := detected.Load(); got != 1 {
t.Errorf("detected %d times for the same scopes in any order, want 1", got)
}
if _, err := cache.token(ctx, []string{"c"}); err != nil {
t.Fatal(err)
}
if got := detected.Load(); got != 2 {
t.Errorf("detected %d times after other scopes, want 2", got)
}
if _, err := cache.token(ctx, []string{"fail"}); err == nil {
t.Error("the detection error was not returned")
}
if _, err := cache.token(ctx, []string{"fail"}); err != nil {
t.Errorf("the detection error was cached: %v", err)
}
}
func TestProviderCacheConcurrent(t *testing.T) {
var detected atomic.Int32
cache := &providerCache{detect: func([]string) (auth.TokenProvider, error) {
detected.Add(1)
return tokenProviderFunc(func(context.Context) (*auth.Token, error) { return &auth.Token{Value: "t"}, nil }), nil
}}
var wg sync.WaitGroup
for range 32 {
wg.Add(1)
go func() {
defer wg.Done()
if _, err := cache.token(context.Background(), []string{"a"}); err != nil {
t.Error(err)
}
}()
}
wg.Wait()
if got := detected.Load(); got != 1 {
t.Errorf("detected %d times, want 1", got)
}
}
// startServe serves the handler on a local listener and returns its address,
// the function that stops it and the channel with the result of serve.
func startServe(t *testing.T, handler http.Handler, grace time.Duration) (string, context.CancelFunc, <-chan error) {
t.Helper()
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
ctx, cancel := context.WithCancel(context.Background())
t.Cleanup(cancel)
result := make(chan error, 1)
go func() {
result <- serve(ctx, &http.Server{Handler: handler}, listener, grace)
}()
return listener.Addr().String(), cancel, result
}
// await returns the result of serve, failing the test when it takes too long.
func await(t *testing.T, result <-chan error) error {
t.Helper()
select {
case err := <-result:
return err
case <-time.After(5 * time.Second):
t.Fatal("serve did not return")
return nil
}
}
func TestServeStopsOnCancel(t *testing.T) {
address, cancel, result := startServe(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}), time.Second)
response, err := http.Get("http://" + address)
if err != nil {
t.Fatal(err)
}
response.Body.Close()
cancel()
if err := await(t, result); err != nil {
t.Errorf("serve = %v, want nil", err)
}
if connection, err := net.Dial("tcp", address); err == nil {
connection.Close()
t.Error("the listener still accepts connections")
}
}
func TestServeFinishesRequestsInProgress(t *testing.T) {
started := make(chan struct{})
release := make(chan struct{})
address, cancel, result := startServe(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
close(started)
<-release
io.WriteString(w, "done")
}), 5*time.Second)
body := make(chan string, 1)
go func() {
response, err := http.Get("http://" + address)
if err != nil {
body <- err.Error()
return
}
defer response.Body.Close()
data, _ := io.ReadAll(response.Body)
body <- string(data)
}()
<-started
cancel()
select {
case err := <-result:
t.Fatalf("serve returned %v while a request was in progress", err)
case <-time.After(200 * time.Millisecond):
}
close(release)
if got := <-body; got != "done" {
t.Errorf("response = %q, want done", got)
}
if err := await(t, result); err != nil {
t.Errorf("serve = %v, want nil", err)
}
}
func TestServeStopsAfterGracePeriod(t *testing.T) {
started := make(chan struct{})
address, cancel, result := startServe(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
close(started)
<-r.Context().Done()
}), 100*time.Millisecond)
failed := make(chan error, 1)
go func() {
response, err := http.Get("http://" + address)
if err == nil {
response.Body.Close()
}
failed <- err
}()
<-started
cancel()
if err := await(t, result); !errors.Is(err, context.DeadlineExceeded) {
t.Errorf("serve = %v, want %v", err, context.DeadlineExceeded)
}
select {
case err := <-failed:
if err == nil {
t.Error("the request in progress was not interrupted")
}
case <-time.After(5 * time.Second):
t.Error("the request in progress was not interrupted")
}
}
func TestServeReturnsListenerFailure(t *testing.T) {
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
listener.Close()
if err := serve(context.Background(), &http.Server{}, listener, time.Second); err == nil {
t.Error("serve = nil, want the listener error")
}
}
// writeServiceAccountKey writes a service account key with a generated private key.
func writeServiceAccountKey(t *testing.T) string {
t.Helper()
key, err := rsa.GenerateKey(rand.Reader, 2048)
if err != nil {
t.Fatal(err)
}
der, err := x509.MarshalPKCS8PrivateKey(key)
if err != nil {
t.Fatal(err)
}
content, err := json.Marshal(map[string]string{
"type": "service_account",
"project_id": "dom-lab",
"private_key_id": "key",
"private_key": string(pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: der})),
"client_email": testEmail,
"client_id": "1",
"token_uri": "https://oauth2.googleapis.com/token",
})
if err != nil {
t.Fatal(err)
}
path := filepath.Join(t.TempDir(), "sa.json")
if err := os.WriteFile(path, content, 0o600); err != nil {
t.Fatal(err)
}
return path
}
// TestSignals runs the emulator as a process and stops it with signals.
func TestSignals(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("POSIX signals")
}
for _, signal := range []syscall.Signal{syscall.SIGTERM, syscall.SIGINT} {
t.Run(signal.String(), func(t *testing.T) {
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
address := listener.Addr().String()
_, port, _ := net.SplitHostPort(address)
listener.Close()
var output bytes.Buffer
command := exec.Command(os.Args[0])
command.Env = append(os.Environ(), mainEnv+"=1", "GOOGLE_APPLICATION_CREDENTIALS="+writeServiceAccountKey(t), "IMDS_PORT="+port)
command.Stderr = &output
if err := command.Start(); err != nil {
t.Fatal(err)
}
exited := make(chan error, 1)
go func() { exited <- command.Wait() }()
t.Cleanup(func() { command.Process.Kill() })
// Wait until the emulator answers the ping.
ready := false
for deadline := time.Now().Add(10 * time.Second); time.Now().Before(deadline) && !ready; {
select {
case err := <-exited:
t.Fatalf("the emulator exited before it was ready: %v\n%s", err, output.String())
default:
}
request, _ := http.NewRequest(http.MethodGet, "http://"+address+"/", nil)
request.Header.Set("Metadata-Flavor", "Google")
if response, err := http.DefaultClient.Do(request); err == nil {
response.Body.Close()
ready = response.StatusCode == http.StatusOK
}
time.Sleep(50 * time.Millisecond)
}
if !ready {
t.Fatalf("the emulator did not answer\n%s", output.String())
}
if err := command.Process.Signal(signal); err != nil {
t.Fatal(err)
}
select {
case err := <-exited:
if err != nil {
t.Errorf("the emulator exited with %v, want status 0\n%s", err, output.String())
}
case <-time.After(10 * time.Second):
t.Fatalf("the emulator did not stop\n%s", output.String())
}
if !strings.Contains(output.String(), "shutting down") || !strings.Contains(output.String(), "Stopped") {
t.Errorf("unexpected log:\n%s", output.String())
}
})
}
}