Skip to content

Commit

Permalink
feat(aws): fallback to IMDSv1 if session token cannot be retrieved
Browse files Browse the repository at this point in the history
  • Loading branch information
karol-kokoszka committed Feb 28, 2024
1 parent 503a9c0 commit b08287f
Showing 1 changed file with 38 additions and 30 deletions.
68 changes: 38 additions & 30 deletions pkg/rclone/aws.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,53 +9,34 @@ import (
"net/http"
"time"

"github.com/pkg/errors"
"github.com/rclone/rclone/fs"
)

// awsRegionFromMetadataAPI uses instance metadata API v2 to fetch region of the
// running instance see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
// Returns empty string if region can't be obtained for whatever reason.
// Fallbacks to IMDSv1 if the session token cannot be obtained.
func awsRegionFromMetadataAPI() string {
const (
tokenUrl = "http://169.254.169.254/latest/api/token"
docURL = "http://169.254.169.254/latest/dynamic/instance-identity/document"
)
const docURL = "http://169.254.169.254/latest/dynamic/instance-identity/document"

// Step 1: Request an IMDSv2 session token
reqToken, err := http.NewRequestWithContext(context.Background(), http.MethodPut, tokenUrl, nil)
token, err := awsAPIToken()
// fallback to IMDSv1 when error on token retrieval
if err != nil {
fs.Errorf(nil, "create token request: %+v", err)
return ""
}
reqToken.Header.Set("X-aws-ec2-metadata-token-ttl-seconds", "21600")
tokenClient := http.Client{
Timeout: 2 * time.Second,
}
resToken, err := tokenClient.Do(reqToken)
if err != nil {
fs.Errorf(nil, "IMDSv2 failed to fetch session token: %+v", err)
return ""
}
defer resToken.Body.Close()

if resToken.StatusCode != http.StatusOK {
fs.Errorf(nil, "failed to retrieve session token: %s", resToken.Status)
return ""
}

token, err := io.ReadAll(resToken.Body)
if err != nil {
fs.Errorf(nil, "Failed to read session token: %+v", err)
return ""
fs.Errorf(nil, "%+v", err)
token = ""
}

// Step 2: Use the session token to retrieve instance metadata
reqMetadata, err := http.NewRequestWithContext(context.Background(), http.MethodGet, docURL, nil)
reqMetadata, err := http.NewRequestWithContext(context.Background(), http.MethodGet, docURL, http.NoBody)
if err != nil {
fs.Errorf(nil, "create metadata request: %+v", err)
return ""
}
reqMetadata.Header.Set("X-aws-ec2-metadata-token", string(token))
if token != "" {
reqMetadata.Header.Set("X-aws-ec2-metadata-token", token)
}

metadataClient := http.Client{
Timeout: 2 * time.Second,
Expand All @@ -77,3 +58,30 @@ func awsRegionFromMetadataAPI() string {

return metadata.Region
}

func awsAPIToken() (string, error) {
const tokenURL = "http://169.254.169.254/latest/api/token"

reqToken, err := http.NewRequestWithContext(context.Background(), http.MethodPut, tokenURL, http.NoBody)
if err != nil {
return "", errors.Wrap(err, "create token request")
}
reqToken.Header.Set("X-aws-ec2-metadata-token-ttl-seconds", "21600")
tokenClient := http.Client{
Timeout: 2 * time.Second,
}
resToken, err := tokenClient.Do(reqToken)
if err != nil {
return "", errors.Wrap(err, "IMDSv2 failed to fetch session token")
}
defer resToken.Body.Close()

if resToken.StatusCode != http.StatusOK {
return "", errors.Wrap(err, "failed to retrieve session token")
}
token, err := io.ReadAll(resToken.Body)
if err != nil {
return "", errors.Wrap(err, "failed to read session token")
}
return string(token), nil
}

0 comments on commit b08287f

Please sign in to comment.