diff --git a/lib/ec2macosinit/imds.go b/lib/ec2macosinit/imds.go index be40c14..b431060 100644 --- a/lib/ec2macosinit/imds.go +++ b/lib/ec2macosinit/imds.go @@ -2,73 +2,150 @@ package ec2macosinit import ( "fmt" + "io" + "net" "net/http" + "net/url" + "os" "strconv" + "strings" + "time" ) const ( - imdsBase = "http://169.254.169.254/latest/" + imdsProbeTimeout = 200 * time.Millisecond imdsTokenTTL = 21600 - tokenEndpoint = "api/token" + tokenEndpoint = "latest/api/token" tokenRequestTTLHeader = "X-aws-ec2-metadata-token-ttl-seconds" tokenHeader = "X-aws-ec2-metadata-token" + imdsEndpointModeEnv = "EC2_METADATA_SERVICE_ENDPOINT_MODE" ) -// IMDS config contains the current instance ID and a place for the IMDSv2 token to be stored. +// baseURL returns an IMDS base URL for the given host. +func baseURL(host string) *url.URL { + return &url.URL{ + Scheme: "http", + Host: host, + } +} + +var ( + imdsIPv4Base = baseURL("169.254.169.254") + imdsIPv6Base = baseURL("[fd00:ec2::254]") +) + +// IMDSConfig contains the current instance ID and a place for the IMDSv2 token to be stored. // Using IMDSv2: -// https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html#instance-metadata-v2-how-it-works +// https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html type IMDSConfig struct { token string + imdsBase *url.URL InstanceID string } +// getIMDSBase returns the IMDS base URL, probing for the correct endpoint +// if not yet determined. If neither endpoint is reachable (e.g. during early +// boot before the network interface is configured), it returns the IPv4 +// default without caching, so the existing init retry loop in setup.go will +// re-probe on the next attempt. +// +// The probe timeout is kept short (200ms) to avoid significantly increasing +// the per-retry cost in setup.go's 1-second retry loop (up to 600 attempts). +// IMDS is link-local and responds in sub-millisecond when reachable, so +// 200ms is more than sufficient to detect availability. +func (i *IMDSConfig) getIMDSBase() *url.URL { + if i.imdsBase != nil { + return i.imdsBase + } + + // Honor explicit override if set (matches SDK convention) + if mode := os.Getenv(imdsEndpointModeEnv); mode != "" { + switch strings.ToLower(mode) { + case "ipv6": + i.imdsBase = imdsIPv6Base + return i.imdsBase + case "ipv4": + i.imdsBase = imdsIPv4Base + return i.imdsBase + } + } + + // Auto-detect: try IPv4 first, fall back to IPv6 + for _, candidate := range []struct { + addr string + base *url.URL + }{ + {"169.254.169.254:80", imdsIPv4Base}, + {"[fd00:ec2::254]:80", imdsIPv6Base}, + } { + conn, err := net.DialTimeout("tcp", candidate.addr, imdsProbeTimeout) + if err == nil { + conn.Close() + i.imdsBase = candidate.base + return i.imdsBase + } + } + + // Neither endpoint confirmed reachable yet — return IPv4 default + // without caching so we re-probe on the next call. + return imdsIPv4Base +} + // getIMDSProperty gets a given endpoint property from IMDS. func (i *IMDSConfig) getIMDSProperty(endpoint string) (value string, httpResponseCode int, err error) { // Check that an IMDSv2 token exists - get one if it doesn't if i.token == "" { err = i.getNewToken() if err != nil { - return "", 0, fmt.Errorf("ec2macosinit: error while getting new IMDS token: %s\n", err) + return "", 0, fmt.Errorf("ec2macosinit: error while getting new IMDS token: %w\n", err) } } // Create request + imdsURL := i.getIMDSBase().JoinPath("latest", endpoint) client := &http.Client{} - req, err := http.NewRequest(http.MethodGet, imdsBase+endpoint, nil) + req, err := http.NewRequest(http.MethodGet, imdsURL.String(), nil) if err != nil { - return "", 0, fmt.Errorf("ec2macosinit: error while creating new HTTP request: %s\n", err) + return "", 0, fmt.Errorf("ec2macosinit: error while creating new HTTP request: %w\n", err) } - req.Header.Set(tokenHeader, i.token) // set IMDSv2 token + req.Header.Set(tokenHeader, i.token) // Make request resp, err := client.Do(req) if err != nil { - return "", 0, fmt.Errorf("ec2macosinit: error while requesting IMDS property: %s\n", err) + return "", 0, fmt.Errorf("ec2macosinit: error while requesting IMDS property: %w\n", err) + } + if resp != nil && resp.Body != nil { + defer resp.Body.Close() } - // Convert returned io.ReadCloser to string - value, err = ioReadCloserToString(resp.Body) + // Read response body + data, err := io.ReadAll(resp.Body) if err != nil { - return "", 0, fmt.Errorf("ec2macosinit: error reading response body: %s\n", err) + return "", 0, fmt.Errorf("ec2macosinit: error reading response body: %w\n", err) } - return value, resp.StatusCode, nil + return string(data), resp.StatusCode, nil } // getNewToken gets a new IMDSv2 token from the IMDS API. func (i *IMDSConfig) getNewToken() (err error) { // Create request + tokenURL := i.getIMDSBase().JoinPath(tokenEndpoint) client := &http.Client{} - req, err := http.NewRequest(http.MethodPut, imdsBase+tokenEndpoint, nil) + req, err := http.NewRequest(http.MethodPut, tokenURL.String(), nil) if err != nil { - return fmt.Errorf("ec2macosinit: error while creating new HTTP request: %s\n", err) + return fmt.Errorf("ec2macosinit: error while creating new HTTP request: %w\n", err) } req.Header.Set(tokenRequestTTLHeader, strconv.FormatInt(int64(imdsTokenTTL), 10)) // Make request resp, err := client.Do(req) if err != nil { - return fmt.Errorf("ec2macosinit: error while requesting new token: %s\n", err) + return fmt.Errorf("ec2macosinit: error while requesting new token: %w\n", err) + } + if resp != nil && resp.Body != nil { + defer resp.Body.Close() } // Validate response code @@ -79,29 +156,27 @@ func (i *IMDSConfig) getNewToken() (err error) { ) } - // Set returned value - i.token, err = ioReadCloserToString(resp.Body) + // Read returned value + data, err := io.ReadAll(resp.Body) if err != nil { - return fmt.Errorf("ec2macosinit: error reading response body: %s\n", err) + return fmt.Errorf("ec2macosinit: error reading response body: %w\n", err) } + i.token = string(data) return nil } -// UpdateInstanceID is a wrapper for getIMDSProperty that gets the current instance ID for the attached config. +// UpdateInstanceID gets the current instance ID from IMDS. func (i *IMDSConfig) UpdateInstanceID() (err error) { - // If instance ID is already set, this doesn't need to be run if i.InstanceID != "" { return nil } - // Get IMDS property "meta-data/instance-id" i.InstanceID, _, err = i.getIMDSProperty("meta-data/instance-id") if err != nil { - return fmt.Errorf("ec2macosinit: error getting instance ID from IMDS: %s\n", err) + return fmt.Errorf("ec2macosinit: error getting instance ID from IMDS: %w\n", err) } - // Validate that an ID was returned if i.InstanceID == "" { return fmt.Errorf("ec2macosinit: an empty instance ID was returned from IMDS\n") }