Keep only base/extractor.[cc|h].

base/apply_tag.cc

-// Copyright 2006-2009 Google Inc.
-//
-// 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
-//
-//      http://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.
-// ========================================================================
-//
-// Applies a tag to a signed file.
-
-#include "omaha/base/apply_tag.h"
-#include <atlrx.h>
-#include <vector>
-#include "base/scoped_ptr.h"
-#include "omaha/base/utils.h"
-#include "omaha/base/extractor.h"
-
-namespace omaha {
-
-const char kMagicBytes[] = "Gact";
-const uint32 kPEHeaderOffset = 60;
-
-ApplyTag::ApplyTag()
-    : prev_tag_string_length_(0),
-      prev_cert_length_(0),
-      append_(0) {}
-
-bool ApplyTag::IsValidTagString(const char* tag_string) {
-  ASSERT1(tag_string);
-
-  CAtlRegExp<CAtlRECharTraitsA> regex;
-  REParseError error = regex.Parse(kValidTagStringRegEx);
-  if (error != REPARSE_ERROR_OK) {
-    return false;
-  }
-
-  CAtlREMatchContext<CAtlRECharTraitsA> context;
-  return !!regex.Match(tag_string, &context);
-}
-
-HRESULT ApplyTag::Init(const TCHAR* signed_exe_file,
-                       const char* tag_string,
-                       int tag_string_length,
-                       const TCHAR* tagged_file,
-                       bool append) {
-  ASSERT1(signed_exe_file);
-  ASSERT1(tag_string);
-  ASSERT1(tagged_file);
-
-  signed_exe_file_ = signed_exe_file;
-  tagged_file_ = tagged_file;
-  append_ = append;
-
-  // Check the tag_string for invalid characters.
-  if (!IsValidTagString(tag_string)) {
-    return E_INVALIDARG;
-  }
-
-  for (int i = 0; i < tag_string_length; ++i) {
-    tag_string_.push_back(tag_string[i]);
-  }
-
-  return S_OK;
-}
-
-HRESULT ApplyTag::EmbedTagString() {
-  std::vector<byte> input_file_buffer;
-  HRESULT hr = ReadEntireFile(signed_exe_file_, 0, &input_file_buffer);
-  if (FAILED(hr)) {
-    return hr;
-  }
-
-  ASSERT1(!input_file_buffer.empty());
-  VERIFY1(ReadExistingTag(&input_file_buffer));
-  if (!append_ && prev_tag_string_length_) {
-    // If there is a previous tag and the append flag is not set, then
-    // we should error out.
-    return APPLYTAG_E_ALREADY_TAGGED;
-  }
-
-  if (!CreateBufferToWrite()) {
-    return E_FAIL;
-  }
-
-  // The input_file_buffer might contain the previously read tag, in which
-  // case the buffer_data_ is larger than the actual output buffer length.
-  // The real output buffer length is returned by the ApplyTagToBuffer
-  // method.
-  buffer_data_.resize(input_file_buffer.size() + tag_buffer_.size());
-
-  copy(input_file_buffer.begin(),
-       input_file_buffer.end(),
-       buffer_data_.begin());
-
-  int output_length = 0;
-  if (!ApplyTagToBuffer(&output_length))
-    return E_FAIL;
-
-  std::vector<byte> output_buffer(output_length);
-  ASSERT1(static_cast<size_t>(output_length) <= buffer_data_.size());
-  copy(buffer_data_.begin(),
-       buffer_data_.begin() + output_length,
-       output_buffer.begin());
-  return WriteEntireFile(tagged_file_, output_buffer);
-}
-
-uint32 ApplyTag::GetUint32(const void* p) {
-  ASSERT1(p);
-
-  const uint32* pu = reinterpret_cast<const uint32*>(p);
-  return *pu;
-}
-
-void ApplyTag::PutUint32(uint32 i, void* p) {
-  ASSERT1(p);
-
-  uint32* pu = reinterpret_cast<uint32*>(p);
-  *pu = i;
-}
-
-bool ApplyTag::ReadExistingTag(std::vector<byte>* binary) {
-  ASSERT1(binary);
-
-  int len = 0;
-  TagExtractor tag;
-  char* bin = reinterpret_cast<char*>(&binary->front());
-  ASSERT1(bin);
-  if (tag.ExtractTag(bin, binary->size(), NULL, &len)) {
-    prev_tag_string_.resize(len);
-    if (tag.ExtractTag(bin, binary->size(), &prev_tag_string_.front(), &len)) {
-      // The extractor returns the actual length
-      // of the string + 1 for the terminating null.
-      prev_tag_string_length_ = len - 1;
-    }
-  }
-
-  // Set the existing certificate length even if previous
-  // tag does not exist.
-  prev_cert_length_ = tag.cert_length();
-  return true;
-}
-
-bool ApplyTag::CreateBufferToWrite() {
-  ASSERT1(!append_ && !prev_tag_string_length_ || append_);
-  ASSERT1(!tag_string_.empty());
-  ASSERT1(!prev_tag_string_.size() ||
-          prev_tag_string_.size() ==
-          static_cast<size_t>(prev_tag_string_length_ + 1));
-
-  // Build the tag buffer.
-  // The format of the tag buffer is:
-  // 000000-000003: 4-byte magic (big-endian)
-  // 000004-000005: unsigned 16-bit int string length (big-endian)
-  // 000006-??????: ASCII string
-  int tag_string_len = tag_string_.size() + prev_tag_string_length_;
-  int kMagicBytesLen = ::lstrlenA(kMagicBytes);
-  int tag_header_len = kMagicBytesLen + 2;
-  int unpadded_tag_buffer_len = tag_string_len + tag_header_len;
-  // The tag buffer should be padded to multiples of 8, otherwise it will
-  // break the signature of the executable file.
-  int padded_tag_buffer_length = (unpadded_tag_buffer_len + 15) & (-8);
-
-  tag_buffer_.clear();
-  tag_buffer_.resize(padded_tag_buffer_length, 0);
-  memcpy(&tag_buffer_.front(), kMagicBytes, kMagicBytesLen);
-  tag_buffer_[kMagicBytesLen] =
-      static_cast<char>((tag_string_len & 0xff00) >> 8);
-  tag_buffer_[kMagicBytesLen+1] = static_cast<char>(tag_string_len & 0xff);
-
-  if (prev_tag_string_length_ > 0) {
-    copy(prev_tag_string_.begin(),
-         prev_tag_string_.end(),
-         tag_buffer_.begin() + tag_header_len);
-  }
-
-  copy(tag_string_.begin(),
-       tag_string_.end(),
-       tag_buffer_.begin() + tag_header_len + prev_tag_string_length_);
-  ASSERT1(static_cast<int>(tag_buffer_.size()) == padded_tag_buffer_length);
-
-  return true;
-}
-
-bool ApplyTag::ApplyTagToBuffer(int* output_len) {
-  ASSERT1(output_len);
-
-  uint32 original_data_len = buffer_data_.size() - tag_buffer_.size();
-  uint32 peheader = GetUint32(&buffer_data_.front() + kPEHeaderOffset);
-  uint32 kCertDirAddressOffset = 152;
-  uint32 kCertDirInfoSize = 4 + 4;
-
-  ASSERT1(peheader + kCertDirAddressOffset + kCertDirInfoSize <=
-          original_data_len);
-
-  // Read certificate directory info.
-  uint32 cert_dir_offset = GetUint32(&buffer_data_.front() + peheader +
-                                     kCertDirAddressOffset);
-  if (cert_dir_offset == 0)
-    return false;
-  uint32 cert_dir_len = GetUint32(&buffer_data_.front() + peheader +
-                                  kCertDirAddressOffset + 4);
-  ASSERT1(cert_dir_offset + cert_dir_len <= original_data_len);
-
-  // Calculate the new output length.
-  int prev_pad_length = cert_dir_len - prev_cert_length_ -
-                        prev_tag_string_length_;
-  ASSERT1(prev_pad_length >= 0);
-  int orig_dir_len = cert_dir_len - prev_tag_string_length_ -
-                     prev_pad_length;
-  ASSERT1(orig_dir_len == prev_cert_length_);
-  int output_length = original_data_len - prev_tag_string_length_ -
-                      prev_pad_length + tag_buffer_.size();
-  *output_len = output_length;
-  ASSERT1(static_cast<size_t>(output_length) <= buffer_data_.size());
-  ASSERT1(output_length >= orig_dir_len);
-
-  // Increase the size of certificate directory.
-  int new_cert_len = prev_cert_length_ + tag_buffer_.size();
-  PutUint32(new_cert_len,
-            &buffer_data_.front() + peheader + kCertDirAddressOffset + 4);
-
-  // Read certificate struct info.
-  uint32 cert_struct_len = GetUint32(&buffer_data_.front() + cert_dir_offset);
-  ASSERT1(!(cert_struct_len > cert_dir_len ||
-            cert_struct_len < cert_dir_len - 8));
-
-  // Increase the certificate struct size.
-  PutUint32(new_cert_len, &buffer_data_.front() + cert_dir_offset);
-
-  // Copy the tag buffer.
-  copy(tag_buffer_.begin(), tag_buffer_.end(),
-       buffer_data_.begin() + cert_dir_offset + prev_cert_length_);
-
-  return true;
-}
-
-}  // namespace omaha