Rewrote IMDS emulator in Go, added unit tests.
This commit is contained in:
slawek committed
1 parent
cedac0a6ee
commit
630375c7ed
10 files changed
+940
-248
No files matched your search
+10
-9
@@ -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.
|
||||
RUN ln -s /usr/local/bin/python /usr/bin/python \
|
||||
&& /usr/bin/python -m venv /venv \
|
||||
&& /venv/bin/pip install --no-cache-dir "google-auth[requests,cryptography]>=2.39.0"
|
||||
WORKDIR /src
|
||||
COPY go.mod go.sum ./
|
||||
RUN go mod download
|
||||
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 imds_emulator.py imds_test.py /app/
|
||||
COPY --from=build /imds-emulator /imds-emulator
|
||||
|
||||
ENTRYPOINT ["/venv/bin/python3", "/app/imds_emulator.py"]
|
||||
ENTRYPOINT ["/imds-emulator"]
|
||||
@@ -19,8 +19,8 @@ services:
|
||||
volumes:
|
||||
- /mnt/ext-storage:/mnt/ext-storage:ro
|
||||
environment:
|
||||
GCE_METADATA_HOST: imds-emulator
|
||||
GCE_METADATA_IP: imds-emulator
|
||||
GCE_METADATA_HOST: imds-emulator:8080
|
||||
GCE_METADATA_IP: imds-emulator:8080
|
||||
command:
|
||||
- --project-id=${PROJECT_ID}
|
||||
- --hostname=${HOSTNAME}
|
||||
|
||||
@@ -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
|
||||
)
|
||||
@@ -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=
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user