// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#include "crypto/default/cms/cms_parser.h"

#include <stdint.h>

#include <string>
#include <vector>

#include "crypto/default/cms/ber_parser.h"
#include "crypto/default/cms/cms_error_code.h"
#include "crypto/default/cms/oids.h"
#include "openssl/base.h"
#include "openssl/bytestring.h"

namespace credentio_cms {

namespace {

constexpr unsigned ContextSpecificConstructed(int n) {
  return CBS_ASN1_CONSTRUCTED | CBS_ASN1_CONTEXT_SPECIFIC | n;
}
constexpr unsigned ContextSpecific(int n) {
  return CBS_ASN1_CONTEXT_SPECIFIC | n;
}

// Scoped class making it easier to keep track of the constructed tag by
// calling EndConstructed on destruction.
class ScopedConstructedElement {
 public:
  // Only calls EndConstructed when destroyed, use it when the constructed tag
  // was already read (usually with GetOptionalTag).
  explicit ScopedConstructedElement(BerParser* parser, const char* hint)
      : parser_(parser) {
    parser->SetElementName(hint);
  }
  // Reads expected_tag, provides a scope name for the error messages and calls
  // EndConstructed when destroyed.
  ScopedConstructedElement(unsigned expected_tag, const char* hint,
                           BerParser* parser)
      : parser_(parser) {
    parser->GetTag(expected_tag);
    parser->SetElementName(hint);
  }
  ~ScopedConstructedElement() { parser_->EndConstructed(); }

 private:
  BerParser* parser_;  // Not owned.
};

// Some applications encode OCTET STRING as a constructed element composed
// of octet strings.
void ParseBerOctetString(int tag, BerParser* parser,
                         std::vector<ByteString>* data) {
  data->clear();
  if (parser->Peek(tag)) {
    data->push_back(ByteString());
    parser->GetElementData(tag, data->back().cbs_ptr());
    return;
  }
  ScopedConstructedElement constructed_string(tag | CBS_ASN1_CONSTRUCTED,
                                              "BerOctetString", parser);
  for (; !parser->IsNullOrEmpty();) {
    data->push_back(ByteString());
    parser->GetElementData(CBS_ASN1_OCTETSTRING, data->back().cbs_ptr());
  }
}

// Expect a sequence with a single OID in it.
void ParseAlgorithmIdentifier(const char* hint, BerParser* parser,
                              AlgorithmIdentifier* id) {
  parser->GetAnyElementAndSkipChildren(id->raw_value.cbs_ptr());
  // Extract the algorithm identifier parts.
  // https://tools.ietf.org/html/rfc5280#section-4.1.1.2
  BerParser id_parser(CBS_data(id->raw_value.cbs_ptr()),
                      CBS_len(id->raw_value.cbs_ptr()));
  id_parser.GetTag(CBS_ASN1_SEQUENCE);
  id_parser.GetElementData(CBS_ASN1_OBJECT, id->algorithm_oid.cbs_ptr());
  if (!id_parser.IsNullOrEmpty()) {
    id_parser.GetAnyElementAndSkipChildren(id->parameter.cbs_ptr());
  }
  id_parser.EndConstructed();
  if (!id_parser.ok()) {
    parser->SetError(id_parser.status(), hint);
  }
}

template <typename T>
void GetIssuerAndSerialNumber(BerParser* parser, T* ids) {
  ScopedConstructedElement issuer_and_serial_number(
      CBS_ASN1_SEQUENCE, "Issuer and serial number", parser);
  parser->GetAnyElementAndSkipChildren(ids->issuer_name.cbs_ptr());
  parser->GetElementData(CBS_ASN1_INTEGER, ids->serial_number.cbs_ptr());
}

// Sn = Serial Number
// Ski = Subject Key Identifier
template <typename T>
void GetIssuerAndSnOrSki(BerParser* parser, T* ids) {
  if (parser->Peek(CBS_ASN1_SEQUENCE)) {
    GetIssuerAndSerialNumber(parser, ids);
  } else if (parser->Peek(ContextSpecific(0))) {
    parser->GetElementData(ContextSpecific(0),
                           ids->subject_key_identifier.cbs_ptr());
  }
}

// http://tools.ietf.org/html/rfc5652#section-6.2.1
void ParseKeyTransRecipientInfo(BerParser* parser, RecipientInfo* recipient) {
  ScopedConstructedElement key_agree(parser, "KeyTransRecipientInfo");
  parser->SetElementName("RecipientInfo");
  recipient->version = parser->GetInt();
  if (recipient->version != 0 && recipient->version != 2) {
    parser->SetError(static_cast<ErrorCode>(ErrorCode::UNSUPPORTED),
                     "Unsupported recipient version");
    return;
  }
  recipient->encrypted_keys.push_back(RecipientEncryptedKey());
  GetIssuerAndSnOrSki(parser, &recipient->encrypted_keys.back());
  ParseAlgorithmIdentifier("Key Encryption Algorithm", parser,
                           &recipient->key_encryption_algorithm);
  parser->GetElementData(
      CBS_ASN1_OCTETSTRING,
      recipient->encrypted_keys.back().encrypted_key.cbs_ptr());
}

void ParseRecipientEncryptedKey(BerParser* parser, RecipientEncryptedKey* key) {
  ScopedConstructedElement recipient_key(CBS_ASN1_SEQUENCE,
                                         "RecicientEncryptedKey", parser);
  if (parser->Peek(CBS_ASN1_SEQUENCE)) {
    GetIssuerAndSerialNumber(parser, key);
  } else {
    ScopedConstructedElement rid(ContextSpecificConstructed(0),
                                 "RecipientKeyIdentifier", parser);
    GetIssuerAndSerialNumber(parser, key);
    parser->GetOptionalElementData(CBS_ASN1_GENERALIZEDTIME,
                                   key->date.cbs_ptr());
    if (!parser->IsNullOrEmpty()) {
      parser->GetAnyElementAndSkipChildren(key->other.cbs_ptr());
    }
  }
  parser->GetElementData(CBS_ASN1_OCTETSTRING, key->encrypted_key.cbs_ptr());
}

// http://tools.ietf.org/html/rfc5652#section-6.2.2
void ParseOriginatorPublicKey(BerParser* parser,
                              OriginatorIdentifierOrKey* oik) {
  ScopedConstructedElement opk(parser, "OriginatorPublicKey");
  ParseAlgorithmIdentifier("Public Key Algorithm", parser,
                           &oik->public_key_algorithm);
  parser->GetElementData(CBS_ASN1_BITSTRING, oik->public_key_value.cbs_ptr());
}

// http://tools.ietf.org/html/rfc5652#section-6.2.2
void ParseOriginatorIdentifierOrKey(BerParser* parser,
                                    OriginatorIdentifierOrKey* oik) {
  ScopedConstructedElement originator(ContextSpecificConstructed(0),
                                      "OriginatorIdentifierOrKey", parser);
  if (parser->Peek(CBS_ASN1_SEQUENCE)) {
    GetIssuerAndSerialNumber(parser, oik);
  } else if (parser->GetOptionalTag(ContextSpecificConstructed(0))) {
    parser->GetElementData(CBS_ASN1_OCTETSTRING,
                           oik->subject_key_identifier.cbs_ptr());
    parser->EndConstructed();
  } else if (parser->GetOptionalTag(ContextSpecificConstructed(1))) {
    ParseOriginatorPublicKey(parser, oik);
  } else {
    parser->SetError(static_cast<ErrorCode>(ErrorCode::UNSUPPORTED),
                     "Unsupported originator tag");
  }
}

// http://tools.ietf.org/html/rfc5652#section-6.2.2
void ParseKeyAgreeRecipientInfo(BerParser* parser, RecipientInfo* recipient) {
  ScopedConstructedElement key_agree(parser, "KeyAgreeRecipientInfo");
  recipient->version = parser->GetInt();
  if (recipient->version != 3) {
    parser->SetError(static_cast<ErrorCode>(ErrorCode::UNSUPPORTED),
                     "Unsupported key agree recipient version");
    return;
  }
  ParseOriginatorIdentifierOrKey(parser,
                                 &recipient->originator_identifier_or_key);
  // User Key Material (UKM)
  if (parser->GetOptionalTag(ContextSpecificConstructed(1))) {
    parser->GetElementData(CBS_ASN1_OCTETSTRING,
                           recipient->user_keying_material.cbs_ptr());
    parser->EndConstructed();
  }
  ParseAlgorithmIdentifier("Key Encryption Algorithm", parser,
                           &recipient->key_encryption_algorithm);

  ScopedConstructedElement recipient_keys(CBS_ASN1_SEQUENCE,
                                          "RecipientEncryptedKeys", parser);
  for (; !parser->IsNullOrEmpty();) {
    RecipientEncryptedKey key;
    ParseRecipientEncryptedKey(parser, &key);
    recipient->encrypted_keys.push_back(key);
  }
}

// https://tools.ietf.org/html/rfc5751#section-2.5.3
bool ExtractSmimeEncryptionKeyPreference(SignerInfo* signer,
                                         ByteString key_preference) {
  SmimeEncryptionKeyPreference& preference = signer->encryption_key_preference;
  BerParser parser(CBS_data(key_preference.cbs_ptr()),
                   CBS_len(key_preference.cbs_ptr()));
  if (parser.Peek(ContextSpecificConstructed(0))) {
    ScopedConstructedElement issuer_sn(ContextSpecificConstructed(0),
                                       "IssuerAndSerialNumber", &parser);
    parser.GetAnyElementAndSkipChildren(preference.issuer_name.cbs_ptr());
    parser.GetElementData(CBS_ASN1_INTEGER, preference.serial_number.cbs_ptr());
  } else if (parser.Peek(ContextSpecificConstructed(1))) {
    ScopedConstructedElement issuer_sn(ContextSpecificConstructed(1),
                                       "RecipientKeyIdentifier", &parser);
    parser.GetElementData(CBS_ASN1_OCTETSTRING,
                          preference.subject_key_identifier.cbs_ptr());
    parser.GetOptionalElementData(CBS_ASN1_GENERALIZEDTIME,
                                  preference.date.cbs_ptr());
    parser.GetAnyElementAndSkipChildren(preference.other.cbs_ptr());
  } else if (parser.Peek(ContextSpecificConstructed(2))) {
    parser.GetElementData(ContextSpecific(2),
                          preference.subject_key_identifier.cbs_ptr());
  } else {
    return false;
  }
  return parser.ok();
}

bool ExtractSingleValue(const Attribute& attribute, unsigned expected_tag,
                        CBS* out) {
  if (attribute.values.size() != 1) {
    return false;
  }
  CBS copy;
  CBS_init(&copy, CBS_data(attribute.values[0].cbs_ptr()),
           CBS_len(attribute.values[0].cbs_ptr()));
  return CBS_get_asn1(&copy, out, expected_tag) == 1;
}

bool ExtractCommonAttributes(SignerInfo* signer) {
  bool status = true;
  for (const auto& attribute : signer->signed_attributes) {
    if (CompareOid(attribute.type.cbs(), kContentAttributeType,
                   sizeof(kContentAttributeType))) {
      status &= ExtractSingleValue(attribute, CBS_ASN1_OBJECT,
                                   signer->content_type_signed.cbs_ptr());
    } else if (CompareOid(attribute.type.cbs(), kContentAttributeMessageDigest,
                          sizeof(kContentAttributeMessageDigest))) {
      status &= ExtractSingleValue(attribute, CBS_ASN1_OCTETSTRING,
                                   signer->message_digest.cbs_ptr());
    } else if (CompareOid(attribute.type.cbs(),
                          kSmimeEncryptionKeyPreferenceOid,
                          sizeof(kSmimeEncryptionKeyPreferenceOid))) {
      status &= (attribute.values.size() == 1);
      if (status) {
        status &=
            ExtractSmimeEncryptionKeyPreference(signer, attribute.values[0]);
      }
    }
  }
  return status;
}

// Extracts the message digest and content type from the signed attributes.
// http://tools.ietf.org/html/rfc5652#section-5.3
void ParseSignedAttributes(BerParser* parser, SignerInfo* signer) {
  // Keep the raw signed attributes to make the signature check easier.
  parser->GetAnyElementAndSkipChildren(signer->raw_signed_attributes.cbs_ptr());
  // Extract the content type and content digest.
  BerParser att_parser(CBS_data(signer->raw_signed_attributes.cbs_ptr()),
                       CBS_len(signer->raw_signed_attributes.cbs_ptr()));
  att_parser.GetTag(ContextSpecificConstructed(0));
  for (; !att_parser.IsNullOrEmpty();) {
    ScopedConstructedElement seq(CBS_ASN1_SEQUENCE, "Authenticated Attributes",
                                 &att_parser);
    signer->signed_attributes.push_back(Attribute());
    Attribute& attribute = signer->signed_attributes.back();
    att_parser.GetElementData(CBS_ASN1_OBJECT, attribute.type.cbs_ptr());
    ScopedConstructedElement s(CBS_ASN1_SET, "Signed attribute values",
                               &att_parser);
    for (; !att_parser.IsNullOrEmpty();) {
      attribute.values.push_back(ByteString());
      att_parser.GetAnyElementAndSkipChildren(
          attribute.values.back().cbs_ptr());
    }
  }
  att_parser.EndConstructed();
  if (!att_parser.ok()) {
    parser->SetError(att_parser.status(),
                     "Failed to parse the signed attributes");
    return;
  }
  if (!ExtractCommonAttributes(signer)) {
    parser->SetError(att_parser.status(),
                     "Invalid content type or digest attributes");
  }
}

// Extract the unsigned attributes for anybody who needs to look at them.
// http://tools.ietf.org/html/rfc5652#section-5.3
void ParseUnsignedAttributes(BerParser* parser, SignerInfo* signer) {
  parser->GetAnyElementAndSkipChildren(
      signer->raw_unsigned_attributes.cbs_ptr());
  BerParser att_parser(CBS_data(signer->raw_unsigned_attributes.cbs_ptr()),
                       CBS_len(signer->raw_unsigned_attributes.cbs_ptr()));
  att_parser.GetTag(ContextSpecificConstructed(1));
  for (; !att_parser.IsNullOrEmpty();) {
    ScopedConstructedElement seq(CBS_ASN1_SEQUENCE, "Unsigned Attributes",
                                 &att_parser);
    signer->unsigned_attributes.push_back(Attribute());
    Attribute& attribute = signer->unsigned_attributes.back();
    att_parser.GetElementData(CBS_ASN1_OBJECT, attribute.type.cbs_ptr());
    ScopedConstructedElement s(CBS_ASN1_SET, "Unsigned attribute values",
                               &att_parser);
    for (; !att_parser.IsNullOrEmpty();) {
      attribute.values.push_back(ByteString());
      att_parser.GetAnyElementAndSkipChildren(
          attribute.values.back().cbs_ptr());
    }
  }
  att_parser.EndConstructed();
  if (!att_parser.ok()) {
    parser->SetError(att_parser.status(),
                     "Failed to parse the unsigned attributes");
  }
}

// http://tools.ietf.org/html/rfc5652#section-5.3
static void ParseSignerInfoInternal(BerParser* parser, SignerInfo* signer) {
  ScopedConstructedElement signer_info(CBS_ASN1_SEQUENCE, "SignerInfo", parser);
  signer->version = parser->GetInt();
  if (signer->version != 1 && signer->version != 3) {
    parser->SetError(ErrorCode::UNSUPPORTED, "Unsupported signer version");
    return;
  }
  GetIssuerAndSnOrSki(parser, signer);
  ParseAlgorithmIdentifier("Digest Algorithm", parser,
                           &signer->digest_algorithm);
  // Authenticated attributes
  if (parser->Peek(ContextSpecificConstructed(0))) {
    ParseSignedAttributes(parser, signer);
  }
  ParseAlgorithmIdentifier("Signature Algorithm", parser,
                           &signer->signature_algorithm);
  parser->GetElementData(CBS_ASN1_OCTETSTRING,
                         signer->signature_value.cbs_ptr());
  // Implicit tag 1: Unauthenticated attributes
  if (parser->Peek(ContextSpecificConstructed(1))) {
    ParseUnsignedAttributes(parser, signer);
  }
}

void ParseCertificates(BerParser* parser,
                       std::vector<ByteString>* certificates) {
  ScopedConstructedElement s(ContextSpecificConstructed(0), "Certificates",
                             parser);
  for (; !parser->IsNullOrEmpty();) {
    ByteString out;
    parser->GetAnyElementAndSkipChildren(out.cbs_ptr());
    certificates->push_back(out);
  }
}

// Certificates Revocation Lists are ignored.
void ParseCrls(BerParser* parser) {
  ScopedConstructedElement s(ContextSpecificConstructed(1), "CRL", parser);
  for (; !parser->IsNullOrEmpty();) {
    ByteString out;
    parser->GetAnyElementAndSkipChildren(out.cbs_ptr());
  }
}

// https://tools.ietf.org/html/rfc5652#section-5.1
void ParseSignedData(BerParser* parser, Content* content) {
  content->type = ContentType::SIGNED_DATA;
  ScopedConstructedElement signed_data(CBS_ASN1_SEQUENCE, "SignedData", parser);
  content->version = parser->GetInt();
  {
    ScopedConstructedElement algorigthm_set(CBS_ASN1_SET, "AlgorithmSet",
                                            parser);
    for (; !parser->IsNullOrEmpty();) {
      content->digest_algorithms.push_back(AlgorithmIdentifier());
      ParseAlgorithmIdentifier("DigestAlgorithm", parser,
                               &content->digest_algorithms.back());
    }
  }
  {
    ScopedConstructedElement contentInfo(CBS_ASN1_SEQUENCE, "Content", parser);
    parser->GetElementData(CBS_ASN1_OBJECT, content->content_type.cbs_ptr());
    content->content_tag = 0;
    if (parser->GetOptionalTag(ContextSpecificConstructed(0))) {
      if (parser->Peek(CBS_ASN1_OCTETSTRING) ||
          parser->Peek(CBS_ASN1_OCTETSTRING | CBS_ASN1_CONSTRUCTED)) {
        content->content_offsets.first = parser->GetCurrentOffset();
        ParseBerOctetString(CBS_ASN1_OCTETSTRING, parser, &content->content);
        content->content_offsets.second = parser->GetCurrentOffset();
        content->content_tag = CBS_ASN1_OCTETSTRING;
      } else if (parser->Peek(CBS_ASN1_SEQUENCE)) {
        // PKCS#7 compatibility.
        // See https://tools.ietf.org/html/rfc5652#section-5.2.1
        content->content_offsets.first = parser->GetCurrentOffset();
        ByteString bytes;
        parser->GetAnyElementAndSkipChildren(bytes.cbs_ptr());
        content->content_offsets.second = parser->GetCurrentOffset();
        content->content.push_back(bytes);
        content->content_tag = CBS_ASN1_SEQUENCE;
      } else {
        parser->SetError(ErrorCode::UNSUPPORTED, "Unexpected type");
      }
      parser->EndConstructed();  // Content
    }
  }
  // Certificates
  if (parser->Peek(ContextSpecificConstructed(0))) {
    ParseCertificates(parser, &content->certificates);
  }
  // Certificate Revocation List are ignored.
  if (parser->Peek(ContextSpecificConstructed(1))) {
    ParseCrls(parser);
  }
  // Signers
  {
    ScopedConstructedElement signers(CBS_ASN1_SET, "Signers", parser);
    for (; !parser->IsNullOrEmpty();) {
      content->signers.push_back(SignerInfo());
      ParseSignerInfoInternal(parser, &content->signers.back());
    }
  }
}

// http://tools.ietf.org/html/rfc5652#section-6.1
void ParseEncryptedContentInfo(BerParser* parser, Content* content) {
  ScopedConstructedElement s(CBS_ASN1_SEQUENCE, "EncryptedContentInfo", parser);
  CBS oid;
  parser->GetElementData(CBS_ASN1_OBJECT, &oid);
  if (!CompareOid(oid, kDataOid, sizeof(kDataOid))) {
    parser->SetError(ErrorCode::UNSUPPORTED, "Unexpected content type.");
    return;
  }
  ParseAlgorithmIdentifier("content encryption algorithm", parser,
                           &content->encryption_algorithm);
  if (parser->Peek(CBS_ASN1_CONTEXT_SPECIFIC | CBS_ASN1_CONSTRUCTED) ||
      parser->Peek(CBS_ASN1_CONTEXT_SPECIFIC)) {
    content->content_offsets.first = parser->GetCurrentOffset();
    ParseBerOctetString(CBS_ASN1_CONTEXT_SPECIFIC, parser, &content->content);
    content->content_offsets.second = parser->GetCurrentOffset();
  } else {
    content->content.clear();
  }
}

// http://tools.ietf.org/html/rfc5652#section-6.2
void ParseRecipientsInfo(BerParser* parser, Content* content) {
  ScopedConstructedElement recipients(CBS_ASN1_SET, "Recipients", parser);
  for (; !parser->IsNullOrEmpty();) {
    content->recipients.push_back(RecipientInfo());
    if (parser->GetOptionalTag(CBS_ASN1_SEQUENCE)) {
      // http://tools.ietf.org/html/rfc2630#section-6.2.1
      ParseKeyTransRecipientInfo(parser, &content->recipients.back());
    } else if (parser->GetOptionalTag(ContextSpecificConstructed(1))) {
      // http://tools.ietf.org/html/rfc2630#section-6.2.2
      ParseKeyAgreeRecipientInfo(parser, &content->recipients.back());
    } else {
      // Unsupported:
      // KEKRecipientInfo, PasswordRecipientInfo, OtherRecipientInfo
      parser->SetError(ErrorCode::UNSUPPORTED, "Unsupported recipient type.");
      return;
    }
  }
}

// https://tools.ietf.org/html/rfc5652#section-6.1
void ParseEnvelopedData(BerParser* parser, Content* content) {
  content->type = ContentType::ENVELOPED_DATA;
  ScopedConstructedElement enveloped_data(CBS_ASN1_SEQUENCE, "EnvelopedData",
                                          parser);
  content->version = parser->GetInt();
  if (content->version != 0 && content->version != 2) {
    parser->SetError(ErrorCode::UNSUPPORTED, "Unsupported envelope version.");
    return;
  }
  if (parser->Peek(ContextSpecificConstructed(0))) {
    ScopedConstructedElement s(ContextSpecificConstructed(0), "OriginatorInfo",
                               parser);
    if (parser->Peek(ContextSpecificConstructed(0))) {
      ParseCertificates(parser, &content->certificates);
    }
    if (parser->Peek(ContextSpecificConstructed(1))) {
      ParseCrls(parser);
    }
  }
  ParseRecipientsInfo(parser, content);
  ParseEncryptedContentInfo(parser, content);
  // UnprotectedAttributes are not used in this implementation.
  if (parser->Peek(ContextSpecificConstructed(1))) {
    CBS dummy;
    parser->GetAnyElementAndSkipChildren(&dummy);
  }
}

}  // namespace

// RFC 5652 (partial support), RFC 2633 SMIME and RFC 2315 PKCS #7
ErrorCode ParseCms(const char* data, size_t size, Content* contents,
                   std::string* error_message) {
  BerParser parser(reinterpret_cast<const uint8_t*>(data), size);
  parser.SetElementName("ContentInfo");
  parser.GetTag(CBS_ASN1_SEQUENCE);
  CBS oid;
  parser.GetElementData(CBS_ASN1_OBJECT, &oid);
  parser.GetTag(ContextSpecificConstructed(0));
  if (CompareOid(oid, kSignedDataOid, sizeof(kSignedDataOid))) {
    ParseSignedData(&parser, contents);
  } else if (CompareOid(oid, kEnvelopedDataOid, sizeof(kEnvelopedDataOid))) {
    ParseEnvelopedData(&parser, contents);
  } else {
    parser.SetError(ErrorCode::UNSUPPORTED, "Unsupported data type.");
  }
  parser.EndConstructed();
  *error_message = parser.error_message();
  return parser.status();
}

}  // namespace credentio_cms