1// Copyright (c) 2012 The Chromium Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5#include "chrome_frame/urlmon_upload_data_stream.h"
6
7#include "net/base/io_buffer.h"
8#include "net/base/net_errors.h"
9#include "net/base/upload_bytes_element_reader.h"
10#include "net/base/upload_file_element_reader.h"
11
12namespace {
13
14// Creates UploadDataStream from UploadData.
15net::UploadDataStream* CreateUploadDataStream(net::UploadData* upload_data) {
16  net::UploadDataStream* upload_data_stream = NULL;
17  const ScopedVector<net::UploadElement>& elements = upload_data->elements();
18
19  if (upload_data->is_chunked()) {
20    // Use AppendChunk when data is chunked.
21    upload_data_stream = new net::UploadDataStream(
22        net::UploadDataStream::CHUNKED, upload_data->identifier());
23
24    for (size_t i = 0; i < elements.size(); ++i) {
25      const net::UploadElement& element = *elements[i];
26      const bool is_last_chunk =
27          i == elements.size() - 1 && upload_data->last_chunk_appended();
28      DCHECK_EQ(net::UploadElement::TYPE_BYTES, element.type());
29      upload_data_stream->AppendChunk(element.bytes(), element.bytes_length(),
30                                      is_last_chunk);
31    }
32  } else {
33    // Not chunked.
34    ScopedVector<net::UploadElementReader> element_readers;
35    for (size_t i = 0; i < elements.size(); ++i) {
36      const net::UploadElement& element = *elements[i];
37      net::UploadElementReader* reader = NULL;
38      switch (element.type()) {
39        case net::UploadElement::TYPE_BYTES:
40          reader = new net::UploadBytesElementReader(element.bytes(),
41                                                     element.bytes_length());
42          break;
43        case net::UploadElement::TYPE_FILE:
44          reader = new net::UploadFileElementReaderSync(
45              element.file_path(),
46              element.file_range_offset(),
47              element.file_range_length(),
48              element.expected_file_modification_time());
49          break;
50      }
51      DCHECK(reader);
52      element_readers.push_back(reader);
53    }
54    upload_data_stream = new net::UploadDataStream(element_readers.Pass(),
55                                                   upload_data->identifier());
56  }
57  return upload_data_stream;
58}
59
60}  // namespace
61
62bool UrlmonUploadDataStream::Initialize(net::UploadData* upload_data) {
63  upload_data_ = upload_data;
64  request_body_stream_.reset(CreateUploadDataStream(upload_data));
65  return request_body_stream_->Init(net::CompletionCallback()) == net::OK;
66}
67
68STDMETHODIMP UrlmonUploadDataStream::Read(void* pv, ULONG cb, ULONG* read) {
69  if (pv == NULL) {
70    NOTREACHED();
71    return E_POINTER;
72  }
73
74  // Have we already read past the end of the stream?
75  if (request_body_stream_->IsEOF()) {
76    if (read) {
77      *read = 0;
78    }
79    return S_FALSE;
80  }
81
82  // The data in request_body_stream_ can be smaller than 'cb' so it's not
83  // guaranteed that we'll be able to read total_bytes_to_copy bytes.
84  uint64 total_bytes_to_copy = cb;
85
86  uint64 bytes_copied = 0;
87
88  char* write_pointer = reinterpret_cast<char*>(pv);
89  while (bytes_copied < total_bytes_to_copy) {
90    size_t bytes_to_copy_now = total_bytes_to_copy - bytes_copied;
91
92    scoped_refptr<net::IOBufferWithSize> buf(
93        new net::IOBufferWithSize(bytes_to_copy_now));
94    int bytes_read = request_body_stream_->Read(buf, buf->size(),
95                                                net::CompletionCallback());
96    DCHECK_NE(net::ERR_IO_PENDING, bytes_read);
97    if (bytes_read == 0)  // Reached the end of the stream.
98      break;
99
100    memcpy(write_pointer, buf->data(), bytes_read);
101
102    // Advance our copy tally
103    bytes_copied += bytes_read;
104
105    // Advance our write pointer
106    write_pointer += bytes_read;
107  }
108
109  DCHECK_LE(bytes_copied, total_bytes_to_copy);
110
111  if (read) {
112    *read = static_cast<ULONG>(bytes_copied);
113  }
114
115  return S_OK;
116}
117
118STDMETHODIMP UrlmonUploadDataStream::Seek(LARGE_INTEGER move, DWORD origin,
119                                          ULARGE_INTEGER* new_pos) {
120  // UploadDataStream is really not very seek-able, so for now allow
121  // STREAM_SEEK_SETs to work with a 0 offset, but fail on everything else.
122  if (origin == STREAM_SEEK_SET && move.QuadPart == 0) {
123    if (request_body_stream_->position() != 0) {
124      request_body_stream_.reset(CreateUploadDataStream(upload_data_));
125      const int result = request_body_stream_->Init(net::CompletionCallback());
126      DCHECK_EQ(net::OK, result);
127    }
128    if (new_pos) {
129      new_pos->QuadPart = 0;
130    }
131    return S_OK;
132  }
133
134  DCHECK(false) << __FUNCTION__;
135  return STG_E_INVALIDFUNCTION;
136}
137
138STDMETHODIMP UrlmonUploadDataStream::Stat(STATSTG *stat_stg,
139                                          DWORD grf_stat_flag) {
140  if (stat_stg == NULL)
141    return E_POINTER;
142
143  memset(stat_stg, 0, sizeof(STATSTG));
144  if (0 == (grf_stat_flag & STATFLAG_NONAME)) {
145    const wchar_t kStreamBuffer[] = L"PostStream";
146    stat_stg->pwcsName =
147        static_cast<wchar_t*>(::CoTaskMemAlloc(sizeof(kStreamBuffer)));
148    lstrcpy(stat_stg->pwcsName, kStreamBuffer);
149  }
150  stat_stg->type = STGTY_STREAM;
151  stat_stg->cbSize.QuadPart = request_body_stream_->size();
152  return S_OK;
153}
154