Tag DynamoDB and S3 client spans with the owning AWS account

dd-java-agent/instrumentation/aws-java/aws-java-common/src/main/java/datadog/trace/instrumentation/aws/AwsAccountIdentity.java

@@ -0,0 +1,182 @@
+package datadog.trace.instrumentation.aws;
+
+/**
+ * Derives AWS account identity for resources addressed by AWS SDK requests.
+ *
+ * <p>Two sources are trustworthy and used here:
+ *
+ * <ul>
+ *   <li>An ARN carried by the request itself (for example a DynamoDB {@code TableName} given as an
+ *       ARN, or an SNS {@code TopicArn}), parsed once with {@link AwsArn}.
+ *   <li>The account that owns the credentials signing the request. Several AWS services resolve a
+ *       bare resource name in the requestor's own account (DynamoDB documents this explicitly: "If
+ *       you only provide the table name parameter instead of a complete ARN, the API operation will
+ *       be performed on the table in the account to which the requestor belongs"), so for those
+ *       services the caller account is the resource owner.
+ * </ul>
+ *
+ * <p>The caller account is taken from the credentials object when the SDK exposes it ({@code
+ * AwsCredentialsIdentity.accountId()} in AWS SDK for Java v2 2.26+, populated by the STS, SSO,
+ * profile, process and container credential providers). As an opt-in fallback for older SDKs the
+ * account can be decoded from the access key ID, which encodes it in its trailing characters.
+ */
+public final class AwsAccountIdentity {
+
+  private AwsAccountIdentity() {}
+
+  /** {@code true} for exactly twelve ASCII digits. */
+  public static boolean isAccountId(final String value) {
+    if (value == null || value.length() != 12) {
+      return false;
+    }
+    for (int i = 0; i < 12; i++) {
+      char c = value.charAt(i);
+      if (c < '0' || c > '9') {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  /** Maps a Region code to its partition. Unknown prefixes map to the commercial partition. */
+  public static String partitionForRegion(final String region) {
+    if (region == null) {
+      return "aws";
+    }
+    if (region.startsWith("cn-")) {
+      return "aws-cn";
+    }
+    if (region.startsWith("us-gov-")) {
+      return "aws-us-gov";
+    }
+    if (region.startsWith("us-isob-")) {
+      return "aws-iso-b";
+    }
+    if (region.startsWith("us-isof-")) {
+      return "aws-iso-f";
+    }
+    if (region.startsWith("us-iso-")) {
+      return "aws-iso";
+    }
+    if (region.startsWith("eu-isoe-")) {
+      return "aws-iso-e";
+    }
+    if (region.startsWith("eusc-")) {
+      return "aws-eusc";
+    }
+    return "aws";
+  }
+
+  /** Builds a DynamoDB table ARN, or returns {@code null} when the Region or account is unknown. */
+  public static String dynamoDbTableArn(
+      final String region, final String account, final String tableName) {
+    if (region == null || !isAccountId(account) || tableName == null || tableName.isEmpty()) {
+      return null;
+    }
+    return "arn:"
+        + partitionForRegion(region)
+        + ":dynamodb:"
+        + region
+        + ':'
+        + account
+        + ":table/"
+        + tableName;
+  }
+
+  /**
+   * Decodes the owning account from an AWS access key ID.
+   *
+   * <p>Access key IDs are a 4 character prefix ({@code AKIA} for long-term keys, {@code ASIA} for
+   * temporary ones) followed by 16 base32 characters that encode 10 bytes. For keys issued since
+   * late March 2019 the account ID is held in the first 48 bits of those bytes, shifted left by 7,
+   * and the top bit of those 48 is always set (the fifth character of the key is {@code Q} or later
+   * in the base32 alphabet). Older keys carry no account at all and that bit is clear, so they are
+   * rejected rather than decoded into an arbitrary value. This encoding is not part of the
+   * documented AWS API surface, which is why callers only use it when explicitly enabled.
+   *
+   * @return the 12-digit account, or {@code null} when the key does not have the expected shape or
+   *     predates the account-encoding format.
+   */
+  public static String accountFromAccessKeyId(final String accessKeyId) {
+    if (accessKeyId == null || accessKeyId.length() != 20) {
+      return null;
+    }
+    if (!(accessKeyId.startsWith("AKIA") || accessKeyId.startsWith("ASIA"))) {
+      return null;
+    }
+    byte[] decoded = decodeBase32(accessKeyId, 4, 20);
+    return decoded == null ? null : accountFromEncodedBytes(decoded);
+  }
+
+  /**
+   * Decodes RFC 4648 base32 (no padding, case-insensitive) from {@code value[from, to)}.
+   *
+   * @return the decoded bytes, or {@code null} when a character is outside the alphabet.
+   */
+  static byte[] decodeBase32(final CharSequence value, final int from, final int to) {
+    byte[] out = new byte[(to - from) * 5 / 8];
+    int buffer = 0;
+    int bits = 0;
+    int index = 0;
+    for (int i = from; i < to; i++) {
+      int digit = base32Value(value.charAt(i));
+      if (digit < 0) {
+        return null;
+      }
+      buffer = (buffer << 5) | digit;
+      bits += 5;
+      if (bits >= 8) {
+        bits -= 8;
+        out[index++] = (byte) ((buffer >> bits) & 0xff);
+      }
+    }
+    return out;
+  }
+
+  /** Set on the first 48 bits of every access key issued in the account-encoding format. */
+  static final long ACCOUNT_FORMAT_MARKER = 0x800000000000L;
+
+  /**
+   * Extracts the account ID from the decoded bytes of an access key ID: the first 6 bytes hold
+   * {@code account << 7} with the top bit set as the format marker.
+   *
+   * @return the zero-padded 12-digit account, or {@code null} when fewer than 6 bytes are given,
+   *     the format marker is clear (a legacy key with no encoded account), or the value does not
+   *     fit in 12 digits.
+   */
+  static String accountFromEncodedBytes(final byte[] decoded) {
+    if (decoded.length < 6) {
+      return null;
+    }
+    long firstSixBytes = 0;
+    for (int i = 0; i < 6; i++) {
+      firstSixBytes = (firstSixBytes << 8) | (decoded[i] & 0xff);
+    }
+    if ((firstSixBytes & ACCOUNT_FORMAT_MARKER) == 0) {
+      return null;
+    }
+    long account = (firstSixBytes & 0x7fffffffff80L) >>> 7;
+    String text = Long.toString(account);
+    if (text.length() > 12) {
+      return null;
+    }
+    StringBuilder padded = new StringBuilder(12);
+    for (int i = text.length(); i < 12; i++) {
+      padded.append('0');
+    }
+    return padded.append(text).toString();
+  }
+
+  private static int base32Value(final char c) {
+    if (c >= 'A' && c <= 'Z') {
+      return c - 'A';
+    }
+    if (c >= 'a' && c <= 'z') {
+      return c - 'a';
+    }
+    if (c >= '2' && c <= '7') {
+      return c - '2' + 26;
+    }
+    return -1;
+  }
+}

dd-java-agent/instrumentation/aws-java/aws-java-common/src/main/java/datadog/trace/instrumentation/aws/AwsArn.java

@@ -0,0 +1,125 @@
+package datadog.trace.instrumentation.aws;
+
+/**
+ * A parsed Amazon Resource Name: {@code arn:<partition>:<service>:<region>:<account>:<resource>}.
+ *
+ * <p>Parsed once so callers that need several fields do not re-scan the string. The AWS SDK ships
+ * its own parser ({@code software.amazon.awssdk.arns.Arn}) but it lives in a module that neither
+ * the SDK v2 2.2.0 floor these instrumentations compile against nor SDK v1 provide, so a minimal
+ * one is kept here.
+ */
+public final class AwsArn {
+
+  private static final String TABLE_PREFIX = "table/";
+
+  private final String raw;
+  private final String partition;
+  private final String service;
+  private final String region;
+  private final String account;
+  private final String resource;
+
+  private AwsArn(
+      final String raw,
+      final String partition,
+      final String service,
+      final String region,
+      final String account,
+      final String resource) {
+    this.raw = raw;
+    this.partition = partition;
+    this.service = service;
+    this.region = region;
+    this.account = account;
+    this.resource = resource;
+  }
+
+  /**
+   * Parses an ARN.
+   *
+   * @return the parsed ARN, or {@code null} when the value is not of the form {@code
+   *     arn:partition:service:region:account:resource} (partition and service non-empty).
+   */
+  public static AwsArn parse(final String value) {
+    if (value == null || !value.startsWith("arn:")) {
+      return null;
+    }
+    int c1 = value.indexOf(':', 4);
+    if (c1 < 0) {
+      return null;
+    }
+    int c2 = value.indexOf(':', c1 + 1);
+    if (c2 < 0) {
+      return null;
+    }
+    int c3 = value.indexOf(':', c2 + 1);
+    if (c3 < 0) {
+      return null;
+    }
+    int c4 = value.indexOf(':', c3 + 1);
+    if (c4 < 0) {
+      return null;
+    }
+    if (c1 == 4 || c2 == c1 + 1 || c4 == value.length() - 1) {
+      // empty partition, empty service or empty resource
+      return null;
+    }
+    return new AwsArn(
+        value,
+        value.substring(4, c1),
+        value.substring(c1 + 1, c2),
+        emptyToNull(value.substring(c2 + 1, c3)),
+        emptyToNull(value.substring(c3 + 1, c4)),
+        value.substring(c4 + 1));
+  }
+
+  /** The original string. */
+  public String raw() {
+    return raw;
+  }
+
+  /** {@code aws}, {@code aws-cn}, {@code aws-us-gov}, ... */
+  public String partition() {
+    return partition;
+  }
+
+  /** {@code dynamodb}, {@code sns}, {@code s3}, ... */
+  public String service() {
+    return service;
+  }
+
+  /** The Region, or {@code null} for global services (S3, IAM). */
+  public String region() {
+    return region;
+  }
+
+  /** The 12-digit account, or {@code null} when absent (S3 ARNs) or not exactly twelve digits. */
+  public String account() {
+    return AwsAccountIdentity.isAccountId(account) ? account : null;
+  }
+
+  /** Everything after the fifth colon. Never empty. */
+  public String resource() {
+    return resource;
+  }
+
+  /**
+   * The table name of a DynamoDB table ARN, with any sub-resource ({@code /index/<name>}, {@code
+   * /stream/<label>}, ...) removed.
+   *
+   * @return the bare table name, or {@code null} when the resource is not a table.
+   */
+  public String dynamoDbTableName() {
+    if (!resource.startsWith(TABLE_PREFIX)) {
+      return null;
+    }
+    int start = TABLE_PREFIX.length();
+    int slash = resource.indexOf('/', start);
+    String name = slash < 0 ? resource.substring(start) : resource.substring(start, slash);
+    return name.isEmpty() ? null : name;
+  }
+
+  private static String emptyToNull(final String value) {
+    return value.isEmpty() ? null : value;
+  }
+}