<commit_before>// Copyright (c) 2016 GitHub, Inc.
// Copyright (c) 2011 The Chromium Authors. All rights reserved.
// Use of this source code is governed by the MIT license that can be
// found in the LICENSE file.

#include "atom/browser/net/atom_url_request.h"
#include <string>
#include "atom/browser/api/atom_api_url_request.h"
#include "atom/browser/atom_browser_context.h"
#include "base/callback.h"
#include "content/public/browser/browser_thread.h"
#include "net/base/elements_upload_data_stream.h"
#include "net/base/io_buffer.h"
#include "net/base/upload_bytes_element_reader.h"

namespace {
const int kBufferSize = 4096;
}  // namespace

namespace atom {

namespace internal {

class UploadOwnedIOBufferElementReader : public net::UploadBytesElementReader {
 public:
  explicit UploadOwnedIOBufferElementReader(
      scoped_refptr<const net::IOBufferWithSize> buffer)
      : net::UploadBytesElementReader(buffer->data(), buffer->size()),
        buffer_(buffer) {}

  ~UploadOwnedIOBufferElementReader() override {}

  static UploadOwnedIOBufferElementReader* CreateWithBuffer(
      scoped_refptr<const net::IOBufferWithSize> buffer) {
    return new UploadOwnedIOBufferElementReader(std::move(buffer));
  }

 private:
  scoped_refptr<const net::IOBuffer> buffer_;

  DISALLOW_COPY_AND_ASSIGN(UploadOwnedIOBufferElementReader);
};

}  // namespace internal

AtomURLRequest::AtomURLRequest(api::URLRequest* delegate)
    : delegate_(delegate),
      is_chunked_upload_(false),
      response_read_buffer_(new net::IOBuffer(kBufferSize)) {}

AtomURLRequest::~AtomURLRequest() {
  DCHECK(!request_context_getter_);
  DCHECK(!request_);
}

scoped_refptr<AtomURLRequest> AtomURLRequest::Create(
    AtomBrowserContext* browser_context,
    const std::string& method,
    const std::string& url,
    api::URLRequest* delegate) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);

  DCHECK(browser_context);
  DCHECK(!url.empty());
  DCHECK(delegate);
  if (!browser_context || url.empty() || !delegate) {
    return nullptr;
  }
  auto request_context_getter = browser_context->url_request_context_getter();
  DCHECK(request_context_getter);
  if (!request_context_getter) {
    return nullptr;
  }
  scoped_refptr<AtomURLRequest> atom_url_request(new AtomURLRequest(delegate));
  if (content::BrowserThread::PostTask(
          content::BrowserThread::IO, FROM_HERE,
          base::Bind(&AtomURLRequest::DoInitialize, atom_url_request,
                     request_context_getter, method, url))) {
    return atom_url_request;
  }
  return nullptr;
}

void AtomURLRequest::Terminate() {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  delegate_ = nullptr;
  content::BrowserThread::PostTask(
      content::BrowserThread::IO, FROM_HERE,
      base::Bind(&AtomURLRequest::DoTerminate, this));
}

void AtomURLRequest::DoInitialize(
    scoped_refptr<net::URLRequestContextGetter> request_context_getter,
    const std::string& method,
    const std::string& url) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  DCHECK(request_context_getter);

  request_context_getter_ = request_context_getter;
  request_context_getter_->AddObserver(this);
  auto context = request_context_getter_->GetURLRequestContext();
  if (!context) {
    // Called after shutdown.
    DoCancelWithError("Cannot start a request after shutdown.", true);
    return;
  }

  DCHECK(context);
  request_ = context->CreateRequest(
      GURL(url), net::RequestPriority::DEFAULT_PRIORITY, this);
  if (!request_) {
    DoCancelWithError("Failed to create a net::URLRequest.", true);
    return;
  }
  request_->set_method(method);
}

void AtomURLRequest::DoTerminate() {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  request_.reset();
  if (request_context_getter_) {
    request_context_getter_->RemoveObserver(this);
    request_context_getter_ = nullptr;
  }
}

bool AtomURLRequest::Write(scoped_refptr<const net::IOBufferWithSize> buffer,
                           bool is_last) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  return content::BrowserThread::PostTask(
      content::BrowserThread::IO, FROM_HERE,
      base::Bind(&AtomURLRequest::DoWriteBuffer, this, buffer, is_last));
}

void AtomURLRequest::SetChunkedUpload(bool is_chunked_upload) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);

  // The method can be called only before switching to multi-threaded mode,
  // i.e. before the first call to write.
  // So it is safe to change the object in the UI thread.
  is_chunked_upload_ = is_chunked_upload;
}

void AtomURLRequest::Cancel() {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  content::BrowserThread::PostTask(content::BrowserThread::IO, FROM_HERE,
                                   base::Bind(&AtomURLRequest::DoCancel, this));
}

void AtomURLRequest::SetExtraHeader(const std::string& name,
                                    const std::string& value) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  content::BrowserThread::PostTask(
      content::BrowserThread::IO, FROM_HERE,
      base::Bind(&AtomURLRequest::DoSetExtraHeader, this, name, value));
}

void AtomURLRequest::RemoveExtraHeader(const std::string& name) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  content::BrowserThread::PostTask(
      content::BrowserThread::IO, FROM_HERE,
      base::Bind(&AtomURLRequest::DoRemoveExtraHeader, this, name));
}

void AtomURLRequest::PassLoginInformation(
    const base::string16& username,
    const base::string16& password) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  if (username.empty() || password.empty()) {
    content::BrowserThread::PostTask(
        content::BrowserThread::IO, FROM_HERE,
        base::Bind(&AtomURLRequest::DoCancelAuth, this));
  } else {
    content::BrowserThread::PostTask(
        content::BrowserThread::IO, FROM_HERE,
        base::Bind(&AtomURLRequest::DoSetAuth, this, username, password));
  }
}

void AtomURLRequest::SetLoadFlags(int flags) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  content::BrowserThread::PostTask(
      content::BrowserThread::IO, FROM_HERE,
      base::Bind(&AtomURLRequest::DoSetLoadFlags, this, flags));
}

void AtomURLRequest::DoWriteBuffer(
    scoped_refptr<const net::IOBufferWithSize> buffer,
    bool is_last) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (!request_) {
    return;
  }

  if (is_chunked_upload_) {
    // Chunked encoding case.

    bool first_call = false;
    if (!chunked_stream_writer_) {
      std::unique_ptr<net::ChunkedUploadDataStream> chunked_stream(
          new net::ChunkedUploadDataStream(0));
      chunked_stream_writer_ = chunked_stream->CreateWriter();
      request_->set_upload(std::move(chunked_stream));
      first_call = true;
    }

    if (buffer)
      // Non-empty buffer.
      chunked_stream_writer_->AppendData(buffer->data(), buffer->size(),
                                         is_last);
    else if (is_last)
      // Empty buffer and last chunk, i.e. request.end().
      chunked_stream_writer_->AppendData(nullptr, 0, true);

    if (first_call) {
      request_->Start();
    }
  } else {
    if (buffer) {
      // Handling potential empty buffers.
      using internal::UploadOwnedIOBufferElementReader;
      auto element_reader =
          UploadOwnedIOBufferElementReader::CreateWithBuffer(std::move(buffer));
      upload_element_readers_.push_back(
          std::unique_ptr<net::UploadElementReader>(element_reader));
    }

    if (is_last) {
      auto elements_upload_data_stream = new net::ElementsUploadDataStream(
          std::move(upload_element_readers_), 0);
      request_->set_upload(
          std::unique_ptr<net::UploadDataStream>(elements_upload_data_stream));
      request_->Start();
    }
  }
}

void AtomURLRequest::DoCancel() {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (request_) {
    request_->Cancel();
  }
  DoTerminate();
}

void AtomURLRequest::DoSetExtraHeader(const std::string& name,
                                      const std::string& value) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (!request_) {
    return;
  }
  request_->SetExtraRequestHeaderByName(name, value, true);
}

void AtomURLRequest::DoRemoveExtraHeader(const std::string& name) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (!request_) {
    return;
  }
  request_->RemoveRequestHeaderByName(name);
}

void AtomURLRequest::DoSetAuth(const base::string16& username,
                               const base::string16& password) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (!request_) {
    return;
  }
  request_->SetAuth(net::AuthCredentials(username, password));
}

void AtomURLRequest::DoCancelAuth() const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (!request_) {
    return;
  }
  request_->CancelAuth();
}

void AtomURLRequest::DoCancelWithError(const std::string& error,
                                       bool isRequestError) {
  DoCancel();
  content::BrowserThread::PostTask(
      content::BrowserThread::UI, FROM_HERE,
      base::Bind(&AtomURLRequest::InformDelegateErrorOccured, this, error,
                 isRequestError));
}

void AtomURLRequest::DoSetLoadFlags(int flags) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (!request_) {
    return;
  }
  request_->SetLoadFlags(request_->load_flags() | flags);
}

void AtomURLRequest::OnAuthRequired(net::URLRequest* request,
                                    net::AuthChallengeInfo* auth_info) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);

  content::BrowserThread::PostTask(
      content::BrowserThread::UI, FROM_HERE,
      base::Bind(&AtomURLRequest::InformDelegateAuthenticationRequired, this,
                 scoped_refptr<net::AuthChallengeInfo>(auth_info)));
}

void AtomURLRequest::OnResponseStarted(net::URLRequest* request) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (!request_) {
    return;
  }
  DCHECK_EQ(request, request_.get());

  scoped_refptr<net::HttpResponseHeaders> response_headers =
      request->response_headers();
  const auto& status = request_->status();
  if (status.is_success()) {
    // Success or pending trigger a Read.
    content::BrowserThread::PostTask(
        content::BrowserThread::UI, FROM_HERE,
        base::Bind(&AtomURLRequest::InformDelegateResponseStarted, this,
                   response_headers));
    ReadResponse();
  } else if (status.status() == net::URLRequestStatus::Status::FAILED) {
    // Report error on Start.
    DoCancelWithError(net::ErrorToString(status.ToNetError()), true);
  }
  // We don't report an error is the request is canceled.
}

void AtomURLRequest::ReadResponse() {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);

  int bytes_read = -1;
  if (request_->Read(response_read_buffer_.get(), kBufferSize, &bytes_read)) {
    OnReadCompleted(request_.get(), bytes_read);
  }
}

void AtomURLRequest::OnReadCompleted(net::URLRequest* request, int bytes_read) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (!request_) {
    return;
  }
  DCHECK_EQ(request, request_.get());

  const auto status = request_->status();

  bool response_error = false;
  bool data_ended = false;
  bool data_transfer_error = false;
  do {
    if (!status.is_success()) {
      response_error = true;
      break;
    }
    if (bytes_read == 0) {
      data_ended = true;
      break;
    }
    if (bytes_read < 0 || !CopyAndPostBuffer(bytes_read)) {
      data_transfer_error = true;
      break;
    }
  } while (
      request_->Read(response_read_buffer_.get(), kBufferSize, &bytes_read));
  if (response_error) {
    DoCancelWithError(net::ErrorToString(status.ToNetError()), false);
  } else if (data_ended) {
    content::BrowserThread::PostTask(
        content::BrowserThread::UI, FROM_HERE,
        base::Bind(&AtomURLRequest::InformDelegateResponseCompleted, this));
    DoTerminate();
  } else if (data_transfer_error) {
    // We abort the request on corrupted data transfer.
    DoCancelWithError("Failed to transfer data from IO to UI thread.", false);
  }
}

void AtomURLRequest::OnContextShuttingDown() {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  DoCancel();
}

bool AtomURLRequest::CopyAndPostBuffer(int bytes_read) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);

  // data is only a wrapper for the asynchronous response_read_buffer_.
  // Make a deep copy of payload and transfer ownership to the UI thread.
  auto buffer_copy = new net::IOBufferWithSize(bytes_read);
  memcpy(buffer_copy->data(), response_read_buffer_->data(), bytes_read);

  return content::BrowserThread::PostTask(
      content::BrowserThread::UI, FROM_HERE,
      base::Bind(&AtomURLRequest::InformDelegateResponseData, this,
                 buffer_copy));
}

void AtomURLRequest::InformDelegateAuthenticationRequired(
    scoped_refptr<net::AuthChallengeInfo> auth_info) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  if (delegate_)
    delegate_->OnAuthenticationRequired(auth_info);
}

void AtomURLRequest::InformDelegateResponseStarted(
    scoped_refptr<net::HttpResponseHeaders> response_headers) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  if (delegate_)
    delegate_->OnResponseStarted(response_headers);
}

void AtomURLRequest::InformDelegateResponseData(
    scoped_refptr<net::IOBufferWithSize> data) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);

  // Transfer ownership of the data buffer, data will be released
  // by the delegate's OnResponseData.
  if (delegate_)
    delegate_->OnResponseData(data);
}

void AtomURLRequest::InformDelegateResponseCompleted() const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);

  if (delegate_)
    delegate_->OnResponseCompleted();
}

void AtomURLRequest::InformDelegateErrorOccured(const std::string& error,
                                                bool isRequestError) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);

  if (delegate_)
    delegate_->OnError(error, isRequestError);
}

}  // namespace atom
<commit_msg>net: disallow cookies from the cookie store<commit_after>// Copyright (c) 2016 GitHub, Inc.
// Copyright (c) 2011 The Chromium Authors. All rights reserved.
// Use of this source code is governed by the MIT license that can be
// found in the LICENSE file.

#include "atom/browser/net/atom_url_request.h"
#include <string>
#include "atom/browser/api/atom_api_url_request.h"
#include "atom/browser/atom_browser_context.h"
#include "base/callback.h"
#include "content/public/browser/browser_thread.h"
#include "net/base/elements_upload_data_stream.h"
#include "net/base/io_buffer.h"
#include "net/base/load_flags.h"
#include "net/base/upload_bytes_element_reader.h"

namespace {
const int kBufferSize = 4096;
}  // namespace

namespace atom {

namespace internal {

class UploadOwnedIOBufferElementReader : public net::UploadBytesElementReader {
 public:
  explicit UploadOwnedIOBufferElementReader(
      scoped_refptr<const net::IOBufferWithSize> buffer)
      : net::UploadBytesElementReader(buffer->data(), buffer->size()),
        buffer_(buffer) {}

  ~UploadOwnedIOBufferElementReader() override {}

  static UploadOwnedIOBufferElementReader* CreateWithBuffer(
      scoped_refptr<const net::IOBufferWithSize> buffer) {
    return new UploadOwnedIOBufferElementReader(std::move(buffer));
  }

 private:
  scoped_refptr<const net::IOBuffer> buffer_;

  DISALLOW_COPY_AND_ASSIGN(UploadOwnedIOBufferElementReader);
};

}  // namespace internal

AtomURLRequest::AtomURLRequest(api::URLRequest* delegate)
    : delegate_(delegate),
      is_chunked_upload_(false),
      response_read_buffer_(new net::IOBuffer(kBufferSize)) {}

AtomURLRequest::~AtomURLRequest() {
  DCHECK(!request_context_getter_);
  DCHECK(!request_);
}

scoped_refptr<AtomURLRequest> AtomURLRequest::Create(
    AtomBrowserContext* browser_context,
    const std::string& method,
    const std::string& url,
    api::URLRequest* delegate) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);

  DCHECK(browser_context);
  DCHECK(!url.empty());
  DCHECK(delegate);
  if (!browser_context || url.empty() || !delegate) {
    return nullptr;
  }
  auto request_context_getter = browser_context->url_request_context_getter();
  DCHECK(request_context_getter);
  if (!request_context_getter) {
    return nullptr;
  }
  scoped_refptr<AtomURLRequest> atom_url_request(new AtomURLRequest(delegate));
  if (content::BrowserThread::PostTask(
          content::BrowserThread::IO, FROM_HERE,
          base::Bind(&AtomURLRequest::DoInitialize, atom_url_request,
                     request_context_getter, method, url))) {
    return atom_url_request;
  }
  return nullptr;
}

void AtomURLRequest::Terminate() {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  delegate_ = nullptr;
  content::BrowserThread::PostTask(
      content::BrowserThread::IO, FROM_HERE,
      base::Bind(&AtomURLRequest::DoTerminate, this));
}

void AtomURLRequest::DoInitialize(
    scoped_refptr<net::URLRequestContextGetter> request_context_getter,
    const std::string& method,
    const std::string& url) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  DCHECK(request_context_getter);

  request_context_getter_ = request_context_getter;
  request_context_getter_->AddObserver(this);
  auto context = request_context_getter_->GetURLRequestContext();
  if (!context) {
    // Called after shutdown.
    DoCancelWithError("Cannot start a request after shutdown.", true);
    return;
  }

  DCHECK(context);
  request_ = context->CreateRequest(
      GURL(url), net::RequestPriority::DEFAULT_PRIORITY, this);
  if (!request_) {
    DoCancelWithError("Failed to create a net::URLRequest.", true);
    return;
  }
  request_->set_method(method);
  // Do not send cookies from the cookie store.
  DoSetLoadFlags(net::LOAD_DO_NOT_SEND_COOKIES);
}

void AtomURLRequest::DoTerminate() {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  request_.reset();
  if (request_context_getter_) {
    request_context_getter_->RemoveObserver(this);
    request_context_getter_ = nullptr;
  }
}

bool AtomURLRequest::Write(scoped_refptr<const net::IOBufferWithSize> buffer,
                           bool is_last) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  return content::BrowserThread::PostTask(
      content::BrowserThread::IO, FROM_HERE,
      base::Bind(&AtomURLRequest::DoWriteBuffer, this, buffer, is_last));
}

void AtomURLRequest::SetChunkedUpload(bool is_chunked_upload) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);

  // The method can be called only before switching to multi-threaded mode,
  // i.e. before the first call to write.
  // So it is safe to change the object in the UI thread.
  is_chunked_upload_ = is_chunked_upload;
}

void AtomURLRequest::Cancel() {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  content::BrowserThread::PostTask(content::BrowserThread::IO, FROM_HERE,
                                   base::Bind(&AtomURLRequest::DoCancel, this));
}

void AtomURLRequest::SetExtraHeader(const std::string& name,
                                    const std::string& value) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  content::BrowserThread::PostTask(
      content::BrowserThread::IO, FROM_HERE,
      base::Bind(&AtomURLRequest::DoSetExtraHeader, this, name, value));
}

void AtomURLRequest::RemoveExtraHeader(const std::string& name) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  content::BrowserThread::PostTask(
      content::BrowserThread::IO, FROM_HERE,
      base::Bind(&AtomURLRequest::DoRemoveExtraHeader, this, name));
}

void AtomURLRequest::PassLoginInformation(
    const base::string16& username,
    const base::string16& password) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  if (username.empty() || password.empty()) {
    content::BrowserThread::PostTask(
        content::BrowserThread::IO, FROM_HERE,
        base::Bind(&AtomURLRequest::DoCancelAuth, this));
  } else {
    content::BrowserThread::PostTask(
        content::BrowserThread::IO, FROM_HERE,
        base::Bind(&AtomURLRequest::DoSetAuth, this, username, password));
  }
}

void AtomURLRequest::SetLoadFlags(int flags) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  content::BrowserThread::PostTask(
      content::BrowserThread::IO, FROM_HERE,
      base::Bind(&AtomURLRequest::DoSetLoadFlags, this, flags));
}

void AtomURLRequest::DoWriteBuffer(
    scoped_refptr<const net::IOBufferWithSize> buffer,
    bool is_last) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (!request_) {
    return;
  }

  if (is_chunked_upload_) {
    // Chunked encoding case.

    bool first_call = false;
    if (!chunked_stream_writer_) {
      std::unique_ptr<net::ChunkedUploadDataStream> chunked_stream(
          new net::ChunkedUploadDataStream(0));
      chunked_stream_writer_ = chunked_stream->CreateWriter();
      request_->set_upload(std::move(chunked_stream));
      first_call = true;
    }

    if (buffer)
      // Non-empty buffer.
      chunked_stream_writer_->AppendData(buffer->data(), buffer->size(),
                                         is_last);
    else if (is_last)
      // Empty buffer and last chunk, i.e. request.end().
      chunked_stream_writer_->AppendData(nullptr, 0, true);

    if (first_call) {
      request_->Start();
    }
  } else {
    if (buffer) {
      // Handling potential empty buffers.
      using internal::UploadOwnedIOBufferElementReader;
      auto element_reader =
          UploadOwnedIOBufferElementReader::CreateWithBuffer(std::move(buffer));
      upload_element_readers_.push_back(
          std::unique_ptr<net::UploadElementReader>(element_reader));
    }

    if (is_last) {
      auto elements_upload_data_stream = new net::ElementsUploadDataStream(
          std::move(upload_element_readers_), 0);
      request_->set_upload(
          std::unique_ptr<net::UploadDataStream>(elements_upload_data_stream));
      request_->Start();
    }
  }
}

void AtomURLRequest::DoCancel() {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (request_) {
    request_->Cancel();
  }
  DoTerminate();
}

void AtomURLRequest::DoSetExtraHeader(const std::string& name,
                                      const std::string& value) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (!request_) {
    return;
  }
  request_->SetExtraRequestHeaderByName(name, value, true);
}

void AtomURLRequest::DoRemoveExtraHeader(const std::string& name) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (!request_) {
    return;
  }
  request_->RemoveRequestHeaderByName(name);
}

void AtomURLRequest::DoSetAuth(const base::string16& username,
                               const base::string16& password) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (!request_) {
    return;
  }
  request_->SetAuth(net::AuthCredentials(username, password));
}

void AtomURLRequest::DoCancelAuth() const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (!request_) {
    return;
  }
  request_->CancelAuth();
}

void AtomURLRequest::DoCancelWithError(const std::string& error,
                                       bool isRequestError) {
  DoCancel();
  content::BrowserThread::PostTask(
      content::BrowserThread::UI, FROM_HERE,
      base::Bind(&AtomURLRequest::InformDelegateErrorOccured, this, error,
                 isRequestError));
}

void AtomURLRequest::DoSetLoadFlags(int flags) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (!request_) {
    return;
  }
  request_->SetLoadFlags(request_->load_flags() | flags);
}

void AtomURLRequest::OnAuthRequired(net::URLRequest* request,
                                    net::AuthChallengeInfo* auth_info) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);

  content::BrowserThread::PostTask(
      content::BrowserThread::UI, FROM_HERE,
      base::Bind(&AtomURLRequest::InformDelegateAuthenticationRequired, this,
                 scoped_refptr<net::AuthChallengeInfo>(auth_info)));
}

void AtomURLRequest::OnResponseStarted(net::URLRequest* request) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (!request_) {
    return;
  }
  DCHECK_EQ(request, request_.get());

  scoped_refptr<net::HttpResponseHeaders> response_headers =
      request->response_headers();
  const auto& status = request_->status();
  if (status.is_success()) {
    // Success or pending trigger a Read.
    content::BrowserThread::PostTask(
        content::BrowserThread::UI, FROM_HERE,
        base::Bind(&AtomURLRequest::InformDelegateResponseStarted, this,
                   response_headers));
    ReadResponse();
  } else if (status.status() == net::URLRequestStatus::Status::FAILED) {
    // Report error on Start.
    DoCancelWithError(net::ErrorToString(status.ToNetError()), true);
  }
  // We don't report an error is the request is canceled.
}

void AtomURLRequest::ReadResponse() {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);

  int bytes_read = -1;
  if (request_->Read(response_read_buffer_.get(), kBufferSize, &bytes_read)) {
    OnReadCompleted(request_.get(), bytes_read);
  }
}

void AtomURLRequest::OnReadCompleted(net::URLRequest* request, int bytes_read) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  if (!request_) {
    return;
  }
  DCHECK_EQ(request, request_.get());

  const auto status = request_->status();

  bool response_error = false;
  bool data_ended = false;
  bool data_transfer_error = false;
  do {
    if (!status.is_success()) {
      response_error = true;
      break;
    }
    if (bytes_read == 0) {
      data_ended = true;
      break;
    }
    if (bytes_read < 0 || !CopyAndPostBuffer(bytes_read)) {
      data_transfer_error = true;
      break;
    }
  } while (
      request_->Read(response_read_buffer_.get(), kBufferSize, &bytes_read));
  if (response_error) {
    DoCancelWithError(net::ErrorToString(status.ToNetError()), false);
  } else if (data_ended) {
    content::BrowserThread::PostTask(
        content::BrowserThread::UI, FROM_HERE,
        base::Bind(&AtomURLRequest::InformDelegateResponseCompleted, this));
    DoTerminate();
  } else if (data_transfer_error) {
    // We abort the request on corrupted data transfer.
    DoCancelWithError("Failed to transfer data from IO to UI thread.", false);
  }
}

void AtomURLRequest::OnContextShuttingDown() {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
  DoCancel();
}

bool AtomURLRequest::CopyAndPostBuffer(int bytes_read) {
  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);

  // data is only a wrapper for the asynchronous response_read_buffer_.
  // Make a deep copy of payload and transfer ownership to the UI thread.
  auto buffer_copy = new net::IOBufferWithSize(bytes_read);
  memcpy(buffer_copy->data(), response_read_buffer_->data(), bytes_read);

  return content::BrowserThread::PostTask(
      content::BrowserThread::UI, FROM_HERE,
      base::Bind(&AtomURLRequest::InformDelegateResponseData, this,
                 buffer_copy));
}

void AtomURLRequest::InformDelegateAuthenticationRequired(
    scoped_refptr<net::AuthChallengeInfo> auth_info) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  if (delegate_)
    delegate_->OnAuthenticationRequired(auth_info);
}

void AtomURLRequest::InformDelegateResponseStarted(
    scoped_refptr<net::HttpResponseHeaders> response_headers) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
  if (delegate_)
    delegate_->OnResponseStarted(response_headers);
}

void AtomURLRequest::InformDelegateResponseData(
    scoped_refptr<net::IOBufferWithSize> data) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);

  // Transfer ownership of the data buffer, data will be released
  // by the delegate's OnResponseData.
  if (delegate_)
    delegate_->OnResponseData(data);
}

void AtomURLRequest::InformDelegateResponseCompleted() const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);

  if (delegate_)
    delegate_->OnResponseCompleted();
}

void AtomURLRequest::InformDelegateErrorOccured(const std::string& error,
                                                bool isRequestError) const {
  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);

  if (delegate_)
    delegate_->OnError(error, isRequestError);
}

}  // namespace atom
<|endoftext|>


<commit_before>/**
 * @file
 *
 * @brief Implementation of PluginDatabase(s)
 *
 * @copyright BSD License (see doc/COPYING or http://www.libelektra.org)
 *
 */

#include <plugindatabase.hpp>

#include <modules.hpp>

#include <algorithm>
#include <kdbconfig.h>

#ifdef HAVE_GLOB
#include <glob.h>
#endif

namespace kdb
{

namespace tools
{

class ModulesPluginDatabase::Impl
{
public:
	Impl () {}
	~Impl () {}
	Modules modules;
};

ModulesPluginDatabase::ModulesPluginDatabase () :
	impl(new ModulesPluginDatabase::Impl())
{}

ModulesPluginDatabase::~ModulesPluginDatabase ()
{}

std::vector<std::string> ModulesPluginDatabase::listAllPlugins() const
{
	std::vector<std::string> ret;
#ifdef ELEKTRA_SHARED
#ifdef HAVE_GLOB
	glob_t pglob;
	if (glob(BUILTIN_PLUGIN_FOLDER "/libelektra-*", GLOB_NOSORT, NULL, &pglob) == 0)
	{
		for (size_t i=0; i<pglob.gl_pathc; ++i)
		{
			std::string fn (pglob.gl_pathv[i]);
			size_t start = fn.find_last_of('-');
			if (start == std::string::npos) continue; // ignore wrong file
			size_t end = fn.find_last_of('.');
			if (end == std::string::npos) continue; // ignore wrong file
			ret.push_back(fn.substr(start+1, end-start-1));
		}
	}
#endif
	if (!ret.empty())
	{
		std::sort (ret.begin(), ret.end());
		return ret;
	}
	// if we did not find plugins, return buildinPlugins
	// (even if they might be wrong for ELEKTRA_SHARED)
#endif
	std::string buildinPlugins = ELEKTRA_PLUGINS;
	std::istringstream ss(buildinPlugins);
	std::string plugin;
	while (getline(ss, plugin, ';'))
	{
		ret.push_back(plugin);
	}
	// remove duplicates:
	std::sort (ret.begin(), ret.end());
	ret.erase (std::unique (ret.begin(), ret.end()), ret.end());
	return ret;
}

namespace
{

bool hasProvides (PluginDatabase const & pd, std::string which)
{
	std::vector<std::string> allPlugins = pd.listAllPlugins();
	std::map<int, PluginSpec> foundPlugins;

	for (auto const & plugin : allPlugins)
	{
		std::istringstream ss (pd.lookupInfo (PluginSpec(plugin, KeySet(5, *Key("system/module",
			KEY_VALUE, "this plugin was loaded without a config", KEY_END), KS_END)),
			"provides"));
		std::string provide;
		while (ss >> provide)
		{
			if (provide == which)
			{
				return true;
			}
		}
	}
	return false;
}

}

// TODO: directly use data from CONTRACT.ini
const std::map <std::string, int> PluginDatabase::statusMap =
{
   {"default",      64000},
   {"recommended",  32000},
   {"productive",    8000},
   {"maintained",    4000},
   {"reviewed",      4000},
   {"conformant",    2000},
   {"compatible",    2000},
   {"coverage",      2000},
   {"specific",      1000},

   {"unittest",      1000},
   {"shelltest",     1000},
   {"tested",         500},
   {"nodep",          250},
   {"libc",           250},
   {"configurable",    50},
   {"final",           50},
   {"global",           1},
   {"preview",        -50},
   {"memleak",       -250},
   {"experimental",  -500},
   {"difficult",     -500},
   {"unfinished",   -1000},
   {"old",          -1000},
   {"nodoc",        -1000},
   {"concept",      -2000},
   {"orphan",       -4000},
   {"obsolete",     -4000},
   {"discouraged", -32000},

};

int PluginDatabase::calculateStatus (std::string statusString)
{
	int ret = 0;
	std::istringstream ss (statusString);
	std::string status;
	while (ss >> status)
	{
		auto it = statusMap.find(status);
		if (it != statusMap.end())
		{
			ret += it->second;
		}
		else
		{
			try {
				ret += stoi(status);
			} catch (std::invalid_argument)
			{}
		}
	}
	return ret;
}

PluginDatabase::Status ModulesPluginDatabase::status (PluginSpec const & spec) const
{
	PluginPtr plugin;
	try {
		plugin = impl->modules.load (spec.getName(), spec.getConfig());
		return real;
	}
	catch (...)
	{
		if (hasProvides (*this, spec.getName()))
		{
			return provides;
		}
		else
		{
			return missing;
		}
	}
}

std::string ModulesPluginDatabase::lookupInfo (PluginSpec const & spec, std::string const & which) const
{
	PluginPtr plugin;
	try {
		plugin = impl->modules.load (spec.getName(), spec.getConfig());
	}
	catch (...)
	{
		throw;
	}

	return plugin->lookupInfo (which);
}

PluginSpec ModulesPluginDatabase::lookupMetadata (std::string const & which) const
{
	std::vector<std::string> allPlugins = listAllPlugins();
	std::map<int, PluginSpec> foundPlugins;

	// collect possible plugins
	for (auto const & plugin : allPlugins)
	{
		try {
			// TODO remove /module hack
			std::istringstream ss (lookupInfo (PluginSpec(plugin, KeySet(5, *Key("system/module",
				KEY_VALUE, "this plugin was loaded without a config", KEY_END), KS_END)),
				"metadata"));
			std::string metadata;
			while (ss >> metadata)
			{
				if (metadata == which)
				{
					int s = calculateStatus (lookupInfo (PluginSpec(plugin, KeySet(5, *Key("system/module",
						KEY_VALUE, "this plugin was loaded without a config", KEY_END), KS_END)),
						"status"));
					foundPlugins.insert(std::make_pair(s, PluginSpec(plugin)));
					break;
				}
			}
		} catch (...) { } // assume not loaded
	}

	if (foundPlugins.empty())
	{
		throw NoPlugin ("Could not find plugin with metadata " + which);
	}

	// the largest element of the map contains the best-suited plugin:
	return foundPlugins.rbegin()->second;
}

PluginSpec ModulesPluginDatabase::lookupProvides (std::string const & which) const
{
	// check if plugin itself exists:
	if (status(PluginSpec(which)) == real)
	{
		return PluginSpec(which);
	}

	std::vector<std::string> allPlugins = listAllPlugins();
	std::map<int, PluginSpec> foundPlugins;
	for (auto const & plugin : allPlugins)
	{
		// TODO: make sure (non)-equal plugins (i.e. with same/different contract) are handled correctly
		try {
			// TODO: support for generic plugins with config
			std::istringstream ss (lookupInfo (PluginSpec(plugin, KeySet(5, *Key("system/module",
				KEY_VALUE, "this plugin was loaded without a config", KEY_END), KS_END)),
				"provides"));
			std::string provide;
			while (ss >> provide)
			{
				if (provide == which)
				{
					int s = calculateStatus (lookupInfo (PluginSpec(plugin, KeySet(5, *Key("system/module",
						KEY_VALUE, "this plugin was loaded without a config", KEY_END), KS_END)),
						"status"));
					foundPlugins.insert(std::make_pair(s, PluginSpec(plugin)));
				}
			}
		} catch (...) { } // assume not loaded
	}

	if (foundPlugins.empty())
	{
		throw NoPlugin("No plugin that provides " + which + " could be found");
	}

	// the largest element of the map contains the best-suited plugin:
	return foundPlugins.rbegin()->second;
}

std::vector<std::string> MockPluginDatabase::listAllPlugins() const
{
	std::vector<std::string> plugins;
	for (auto const & plugin : data)
	{
		plugins.push_back(plugin.first.getName());
	}
	return plugins;
}

PluginDatabase::Status MockPluginDatabase::status (PluginSpec const & spec) const
{
	auto it = data.find(spec);
	if (it != data.end())
	{
		return real;
	}

	if (hasProvides(*this, spec.getName()))
	{
		return provides;
	}

	return missing;
}

std::string MockPluginDatabase::lookupInfo(PluginSpec const & spec, std::string const & which) const
{
	auto it = data.find(spec);
	if (it != data.end())
	{
		return it->second[which];
	}

	return "";
}

}

}
<commit_msg>allow to load plugins with /module<commit_after>/**
 * @file
 *
 * @brief Implementation of PluginDatabase(s)
 *
 * @copyright BSD License (see doc/COPYING or http://www.libelektra.org)
 *
 */

#include <plugindatabase.hpp>

#include <modules.hpp>

#include <algorithm>
#include <kdbconfig.h>

#ifdef HAVE_GLOB
#include <glob.h>
#endif

namespace kdb
{

namespace tools
{

class ModulesPluginDatabase::Impl
{
public:
	Impl () {}
	~Impl () {}
	Modules modules;
};

ModulesPluginDatabase::ModulesPluginDatabase () :
	impl(new ModulesPluginDatabase::Impl())
{}

ModulesPluginDatabase::~ModulesPluginDatabase ()
{}

std::vector<std::string> ModulesPluginDatabase::listAllPlugins() const
{
	std::vector<std::string> ret;
#ifdef ELEKTRA_SHARED
#ifdef HAVE_GLOB
	glob_t pglob;
	if (glob(BUILTIN_PLUGIN_FOLDER "/libelektra-*", GLOB_NOSORT, NULL, &pglob) == 0)
	{
		for (size_t i=0; i<pglob.gl_pathc; ++i)
		{
			std::string fn (pglob.gl_pathv[i]);
			size_t start = fn.find_last_of('-');
			if (start == std::string::npos) continue; // ignore wrong file
			size_t end = fn.find_last_of('.');
			if (end == std::string::npos) continue; // ignore wrong file
			ret.push_back(fn.substr(start+1, end-start-1));
		}
	}
#endif
	if (!ret.empty())
	{
		std::sort (ret.begin(), ret.end());
		return ret;
	}
	// if we did not find plugins, return buildinPlugins
	// (even if they might be wrong for ELEKTRA_SHARED)
#endif
	std::string buildinPlugins = ELEKTRA_PLUGINS;
	std::istringstream ss(buildinPlugins);
	std::string plugin;
	while (getline(ss, plugin, ';'))
	{
		ret.push_back(plugin);
	}
	// remove duplicates:
	std::sort (ret.begin(), ret.end());
	ret.erase (std::unique (ret.begin(), ret.end()), ret.end());
	return ret;
}

namespace
{

bool hasProvides (PluginDatabase const & pd, std::string which)
{
	std::vector<std::string> allPlugins = pd.listAllPlugins();
	std::map<int, PluginSpec> foundPlugins;

	for (auto const & plugin : allPlugins)
	{
		std::istringstream ss (pd.lookupInfo (PluginSpec(plugin, KeySet(5, *Key("system/module",
			KEY_VALUE, "this plugin was loaded without a config", KEY_END), KS_END)),
			"provides"));
		std::string provide;
		while (ss >> provide)
		{
			if (provide == which)
			{
				return true;
			}
		}
	}
	return false;
}

}

// TODO: directly use data from CONTRACT.ini
const std::map <std::string, int> PluginDatabase::statusMap =
{
   {"default",      64000},
   {"recommended",  32000},
   {"productive",    8000},
   {"maintained",    4000},
   {"reviewed",      4000},
   {"conformant",    2000},
   {"compatible",    2000},
   {"coverage",      2000},
   {"specific",      1000},

   {"unittest",      1000},
   {"shelltest",     1000},
   {"tested",         500},
   {"nodep",          250},
   {"libc",           250},
   {"configurable",    50},
   {"final",           50},
   {"global",           1},
   {"preview",        -50},
   {"memleak",       -250},
   {"experimental",  -500},
   {"difficult",     -500},
   {"unfinished",   -1000},
   {"old",          -1000},
   {"nodoc",        -1000},
   {"concept",      -2000},
   {"orphan",       -4000},
   {"obsolete",     -4000},
   {"discouraged", -32000},

};

int PluginDatabase::calculateStatus (std::string statusString)
{
	int ret = 0;
	std::istringstream ss (statusString);
	std::string status;
	while (ss >> status)
	{
		auto it = statusMap.find(status);
		if (it != statusMap.end())
		{
			ret += it->second;
		}
		else
		{
			try {
				ret += stoi(status);
			} catch (std::invalid_argument)
			{}
		}
	}
	return ret;
}

PluginDatabase::Status ModulesPluginDatabase::status (PluginSpec const & spec) const
{
	PluginPtr plugin;
	try {
		KeySet conf = spec.getConfig();
		conf.append(Key("system/module",
				KEY_VALUE, "this plugin was loaded for the status", KEY_END));
		plugin = impl->modules.load (spec.getName(), conf);
		return real;
	}
	catch (...)
	{
		if (hasProvides (*this, spec.getName()))
		{
			return provides;
		}
		else
		{
			return missing;
		}
	}
}

std::string ModulesPluginDatabase::lookupInfo (PluginSpec const & spec, std::string const & which) const
{
	PluginPtr plugin;
	try {
		plugin = impl->modules.load (spec.getName(), spec.getConfig());
	}
	catch (...)
	{
		throw;
	}

	return plugin->lookupInfo (which);
}

PluginSpec ModulesPluginDatabase::lookupMetadata (std::string const & which) const
{
	std::vector<std::string> allPlugins = listAllPlugins();
	std::map<int, PluginSpec> foundPlugins;

	// collect possible plugins
	for (auto const & plugin : allPlugins)
	{
		try {
			// TODO remove /module hack
			std::istringstream ss (lookupInfo (PluginSpec(plugin, KeySet(5, *Key("system/module",
				KEY_VALUE, "this plugin was loaded without a config", KEY_END), KS_END)),
				"metadata"));
			std::string metadata;
			while (ss >> metadata)
			{
				if (metadata == which)
				{
					int s = calculateStatus (lookupInfo (PluginSpec(plugin, KeySet(5, *Key("system/module",
						KEY_VALUE, "this plugin was loaded without a config", KEY_END), KS_END)),
						"status"));
					foundPlugins.insert(std::make_pair(s, PluginSpec(plugin)));
					break;
				}
			}
		} catch (...) { } // assume not loaded
	}

	if (foundPlugins.empty())
	{
		throw NoPlugin ("Could not find plugin with metadata " + which);
	}

	// the largest element of the map contains the best-suited plugin:
	return foundPlugins.rbegin()->second;
}

PluginSpec ModulesPluginDatabase::lookupProvides (std::string const & which) const
{
	// check if plugin itself exists:
	if (status(PluginSpec(which)) == real)
	{
		return PluginSpec(which);
	}

	std::vector<std::string> allPlugins = listAllPlugins();
	std::map<int, PluginSpec> foundPlugins;
	for (auto const & plugin : allPlugins)
	{
		// TODO: make sure (non)-equal plugins (i.e. with same/different contract) are handled correctly
		try {
			// TODO: support for generic plugins with config
			std::istringstream ss (lookupInfo (PluginSpec(plugin, KeySet(5, *Key("system/module",
				KEY_VALUE, "this plugin was loaded without a config", KEY_END), KS_END)),
				"provides"));
			std::string provide;
			while (ss >> provide)
			{
				if (provide == which)
				{
					int s = calculateStatus (lookupInfo (PluginSpec(plugin, KeySet(5, *Key("system/module",
						KEY_VALUE, "this plugin was loaded without a config", KEY_END), KS_END)),
						"status"));
					foundPlugins.insert(std::make_pair(s, PluginSpec(plugin)));
				}
			}
		} catch (...) { } // assume not loaded
	}

	if (foundPlugins.empty())
	{
		throw NoPlugin("No plugin that provides " + which + " could be found");
	}

	// the largest element of the map contains the best-suited plugin:
	return foundPlugins.rbegin()->second;
}

std::vector<std::string> MockPluginDatabase::listAllPlugins() const
{
	std::vector<std::string> plugins;
	for (auto const & plugin : data)
	{
		plugins.push_back(plugin.first.getName());
	}
	return plugins;
}

PluginDatabase::Status MockPluginDatabase::status (PluginSpec const & spec) const
{
	auto it = data.find(spec);
	if (it != data.end())
	{
		return real;
	}

	if (hasProvides(*this, spec.getName()))
	{
		return provides;
	}

	return missing;
}

std::string MockPluginDatabase::lookupInfo(PluginSpec const & spec, std::string const & which) const
{
	auto it = data.find(spec);
	if (it != data.end())
	{
		return it->second[which];
	}

	return "";
}

}

}
<|endoftext|>


<commit_before>#if ENABLE_S3

#include "s3.hh"
#include "s3-binary-cache-store.hh"
#include "nar-info.hh"
#include "nar-info-disk-cache.hh"
#include "globals.hh"
#include "compression.hh"
#include "download.hh"
#include "istringstream_nocopy.hh"

#include <aws/core/Aws.h>
#include <aws/core/client/ClientConfiguration.h>
#include <aws/core/client/DefaultRetryStrategy.h>
#include <aws/s3/S3Client.h>
#include <aws/s3/model/CreateBucketRequest.h>
#include <aws/s3/model/GetBucketLocationRequest.h>
#include <aws/s3/model/GetObjectRequest.h>
#include <aws/s3/model/HeadObjectRequest.h>
#include <aws/s3/model/ListObjectsRequest.h>
#include <aws/s3/model/PutObjectRequest.h>

namespace nix {

struct S3Error : public Error
{
    Aws::S3::S3Errors err;
    S3Error(Aws::S3::S3Errors err, const FormatOrString & fs)
        : Error(fs), err(err) { };
};

/* Helper: given an Outcome<R, E>, return R in case of success, or
   throw an exception in case of an error. */
template<typename R, typename E>
R && checkAws(const FormatOrString & fs, Aws::Utils::Outcome<R, E> && outcome)
{
    if (!outcome.IsSuccess())
        throw S3Error(
            outcome.GetError().GetErrorType(),
            fs.s + ": " + outcome.GetError().GetMessage());
    return outcome.GetResultWithOwnership();
}

static void initAWS()
{
    static std::once_flag flag;
    std::call_once(flag, []() {
        Aws::SDKOptions options;

        /* We install our own OpenSSL locking function (see
           shared.cc), so don't let aws-sdk-cpp override it. */
        options.cryptoOptions.initAndCleanupOpenSSL = false;

        Aws::InitAPI(options);
    });
}

S3Helper::S3Helper(const string & region)
    : config(makeConfig(region))
    , client(make_ref<Aws::S3::S3Client>(*config))
{
}

/* Log AWS retries. */
class RetryStrategy : public Aws::Client::DefaultRetryStrategy
{
    long CalculateDelayBeforeNextRetry(const Aws::Client::AWSError<Aws::Client::CoreErrors>& error, long attemptedRetries) const override
    {
        auto res = Aws::Client::DefaultRetryStrategy::CalculateDelayBeforeNextRetry(error, attemptedRetries);
        printError("AWS error '%s' (%s), will retry in %d ms",
            error.GetExceptionName(), error.GetMessage(), res);
        return res;
    }
};

ref<Aws::Client::ClientConfiguration> S3Helper::makeConfig(const string & region)
{
    initAWS();
    auto res = make_ref<Aws::Client::ClientConfiguration>();
    res->region = region;
    res->requestTimeoutMs = 600 * 1000;
    res->retryStrategy = std::make_shared<RetryStrategy>();
    res->caFile = settings.caFile;
    return res;
}

S3Helper::DownloadResult S3Helper::getObject(
    const std::string & bucketName, const std::string & key)
{
    debug("fetching ‘s3://%s/%s’...", bucketName, key);

    auto request =
        Aws::S3::Model::GetObjectRequest()
        .WithBucket(bucketName)
        .WithKey(key);

    request.SetResponseStreamFactory([&]() {
        return Aws::New<std::stringstream>("STRINGSTREAM");
    });

    DownloadResult res;

    auto now1 = std::chrono::steady_clock::now();

    try {

        auto result = checkAws(fmt("AWS error fetching ‘%s’", key),
            client->GetObject(request));

        res.data = decodeContent(
            result.GetContentEncoding(),
            make_ref<std::string>(
                dynamic_cast<std::stringstream &>(result.GetBody()).str()));

    } catch (S3Error & e) {
        if (e.err != Aws::S3::S3Errors::NO_SUCH_KEY) throw;
    }

    auto now2 = std::chrono::steady_clock::now();

    res.durationMs = std::chrono::duration_cast<std::chrono::milliseconds>(now2 - now1).count();

    return res;
}

struct S3BinaryCacheStoreImpl : public S3BinaryCacheStore
{
    const Setting<std::string> region{this, Aws::Region::US_EAST_1, "region", {"aws-region"}};
    const Setting<std::string> narinfoCompression{this, "", "narinfo-compression", "compression method for .narinfo files"};
    const Setting<std::string> lsCompression{this, "", "ls-compression", "compression method for .ls files"};
    const Setting<std::string> logCompression{this, "", "log-compression", "compression method for log/* files"};

    std::string bucketName;

    Stats stats;

    S3Helper s3Helper;

    S3BinaryCacheStoreImpl(
        const Params & params, const std::string & bucketName)
        : S3BinaryCacheStore(params)
        , bucketName(bucketName)
        , s3Helper(region)
    {
        diskCache = getNarInfoDiskCache();
    }

    std::string getUri() override
    {
        return "s3://" + bucketName;
    }

    void init() override
    {
        if (!diskCache->cacheExists(getUri(), wantMassQuery_, priority)) {

            /* Create the bucket if it doesn't already exists. */
            // FIXME: HeadBucket would be more appropriate, but doesn't return
            // an easily parsed 404 message.
            auto res = s3Helper.client->GetBucketLocation(
                Aws::S3::Model::GetBucketLocationRequest().WithBucket(bucketName));

            if (!res.IsSuccess()) {
                if (res.GetError().GetErrorType() != Aws::S3::S3Errors::NO_SUCH_BUCKET)
                    throw Error(format("AWS error checking bucket ‘%s’: %s") % bucketName % res.GetError().GetMessage());

                checkAws(format("AWS error creating bucket ‘%s’") % bucketName,
                    s3Helper.client->CreateBucket(
                        Aws::S3::Model::CreateBucketRequest()
                        .WithBucket(bucketName)
                        .WithCreateBucketConfiguration(
                            Aws::S3::Model::CreateBucketConfiguration()
                            /* .WithLocationConstraint(
                               Aws::S3::Model::BucketLocationConstraint::US) */ )));
            }

            BinaryCacheStore::init();

            diskCache->createCache(getUri(), storeDir, wantMassQuery_, priority);
        }
    }

    const Stats & getS3Stats() override
    {
        return stats;
    }

    /* This is a specialisation of isValidPath() that optimistically
       fetches the .narinfo file, rather than first checking for its
       existence via a HEAD request. Since .narinfos are small, doing
       a GET is unlikely to be slower than HEAD. */
    bool isValidPathUncached(const Path & storePath) override
    {
        try {
            queryPathInfo(storePath);
            return true;
        } catch (InvalidPath & e) {
            return false;
        }
    }

    bool fileExists(const std::string & path) override
    {
        stats.head++;

        auto res = s3Helper.client->HeadObject(
            Aws::S3::Model::HeadObjectRequest()
            .WithBucket(bucketName)
            .WithKey(path));

        if (!res.IsSuccess()) {
            auto & error = res.GetError();
            if (error.GetErrorType() == Aws::S3::S3Errors::UNKNOWN // FIXME
                && error.GetMessage().find("404") != std::string::npos)
                return false;
            throw Error(format("AWS error fetching ‘%s’: %s") % path % error.GetMessage());
        }

        return true;
    }

    void uploadFile(const std::string & path, const std::string & data,
        const std::string & mimeType,
        const std::string & contentEncoding)
    {
        auto request =
            Aws::S3::Model::PutObjectRequest()
            .WithBucket(bucketName)
            .WithKey(path);

        request.SetContentType(mimeType);

        if (contentEncoding != "")
            request.SetContentEncoding(contentEncoding);

        auto stream = std::make_shared<istringstream_nocopy>(data);

        request.SetBody(stream);

        stats.put++;
        stats.putBytes += data.size();

        auto now1 = std::chrono::steady_clock::now();

        auto result = checkAws(format("AWS error uploading ‘%s’") % path,
            s3Helper.client->PutObject(request));

        auto now2 = std::chrono::steady_clock::now();

        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(now2 - now1).count();

        printInfo(format("uploaded ‘s3://%1%/%2%’ (%3% bytes) in %4% ms")
            % bucketName % path % data.size() % duration);

        stats.putTimeMs += duration;
    }

    void upsertFile(const std::string & path, const std::string & data,
        const std::string & mimeType) override
    {
        if (narinfoCompression != "" && hasSuffix(path, ".narinfo"))
            uploadFile(path, *compress(narinfoCompression, data), mimeType, narinfoCompression);
        else if (lsCompression != "" && hasSuffix(path, ".ls"))
            uploadFile(path, *compress(lsCompression, data), mimeType, lsCompression);
        else if (logCompression != "" && hasPrefix(path, "log/"))
            uploadFile(path, *compress(logCompression, data), mimeType, logCompression);
        else
            uploadFile(path, data, mimeType, "");
    }

    void getFile(const std::string & path,
        std::function<void(std::shared_ptr<std::string>)> success,
        std::function<void(std::exception_ptr exc)> failure) override
    {
        sync2async<std::shared_ptr<std::string>>(success, failure, [&]() {
            stats.get++;

            auto res = s3Helper.getObject(bucketName, path);

            stats.getBytes += res.data ? res.data->size() : 0;
            stats.getTimeMs += res.durationMs;

            if (res.data)
                printTalkative("downloaded ‘s3://%s/%s’ (%d bytes) in %d ms",
                    bucketName, path, res.data->size(), res.durationMs);

            return res.data;
        });
    }

    PathSet queryAllValidPaths() override
    {
        PathSet paths;
        std::string marker;

        do {
            debug(format("listing bucket ‘s3://%s’ from key ‘%s’...") % bucketName % marker);

            auto res = checkAws(format("AWS error listing bucket ‘%s’") % bucketName,
                s3Helper.client->ListObjects(
                    Aws::S3::Model::ListObjectsRequest()
                    .WithBucket(bucketName)
                    .WithDelimiter("/")
                    .WithMarker(marker)));

            auto & contents = res.GetContents();

            debug(format("got %d keys, next marker ‘%s’")
                % contents.size() % res.GetNextMarker());

            for (auto object : contents) {
                auto & key = object.GetKey();
                if (key.size() != 40 || !hasSuffix(key, ".narinfo")) continue;
                paths.insert(storeDir + "/" + key.substr(0, key.size() - 8));
            }

            marker = res.GetNextMarker();
        } while (!marker.empty());

        return paths;
    }

};

static RegisterStoreImplementation regStore([](
    const std::string & uri, const Store::Params & params)
    -> std::shared_ptr<Store>
{
    if (std::string(uri, 0, 5) != "s3://") return 0;
    auto store = std::make_shared<S3BinaryCacheStoreImpl>(params, std::string(uri, 5));
    store->init();
    return store;
});

}

#endif
<commit_msg>Show aws-sdk-cpp log messages<commit_after>#if ENABLE_S3

#include "s3.hh"
#include "s3-binary-cache-store.hh"
#include "nar-info.hh"
#include "nar-info-disk-cache.hh"
#include "globals.hh"
#include "compression.hh"
#include "download.hh"
#include "istringstream_nocopy.hh"

#include <aws/core/Aws.h>
#include <aws/core/client/ClientConfiguration.h>
#include <aws/core/client/DefaultRetryStrategy.h>
#include <aws/core/utils/logging/FormattedLogSystem.h>
#include <aws/core/utils/logging/LogMacros.h>
#include <aws/s3/S3Client.h>
#include <aws/s3/model/CreateBucketRequest.h>
#include <aws/s3/model/GetBucketLocationRequest.h>
#include <aws/s3/model/GetObjectRequest.h>
#include <aws/s3/model/HeadObjectRequest.h>
#include <aws/s3/model/ListObjectsRequest.h>
#include <aws/s3/model/PutObjectRequest.h>

namespace nix {

struct S3Error : public Error
{
    Aws::S3::S3Errors err;
    S3Error(Aws::S3::S3Errors err, const FormatOrString & fs)
        : Error(fs), err(err) { };
};

/* Helper: given an Outcome<R, E>, return R in case of success, or
   throw an exception in case of an error. */
template<typename R, typename E>
R && checkAws(const FormatOrString & fs, Aws::Utils::Outcome<R, E> && outcome)
{
    if (!outcome.IsSuccess())
        throw S3Error(
            outcome.GetError().GetErrorType(),
            fs.s + ": " + outcome.GetError().GetMessage());
    return outcome.GetResultWithOwnership();
}

class AwsLogger : public Aws::Utils::Logging::FormattedLogSystem
{
    using Aws::Utils::Logging::FormattedLogSystem::FormattedLogSystem;

    void ProcessFormattedStatement(Aws::String && statement) override
    {
        debug("AWS: %s", chomp(statement));
    }
};

static void initAWS()
{
    static std::once_flag flag;
    std::call_once(flag, []() {
        Aws::SDKOptions options;

        /* We install our own OpenSSL locking function (see
           shared.cc), so don't let aws-sdk-cpp override it. */
        options.cryptoOptions.initAndCleanupOpenSSL = false;

        if (verbosity >= lvlDebug) {
            options.loggingOptions.logLevel =
                verbosity == lvlDebug
                ? Aws::Utils::Logging::LogLevel::Debug
                : Aws::Utils::Logging::LogLevel::Trace;
            options.loggingOptions.logger_create_fn = [options]() {
                return std::make_shared<AwsLogger>(options.loggingOptions.logLevel);
            };
        }

        Aws::InitAPI(options);
    });
}

S3Helper::S3Helper(const string & region)
    : config(makeConfig(region))
    , client(make_ref<Aws::S3::S3Client>(*config))
{
}

/* Log AWS retries. */
class RetryStrategy : public Aws::Client::DefaultRetryStrategy
{
    long CalculateDelayBeforeNextRetry(const Aws::Client::AWSError<Aws::Client::CoreErrors>& error, long attemptedRetries) const override
    {
        auto res = Aws::Client::DefaultRetryStrategy::CalculateDelayBeforeNextRetry(error, attemptedRetries);
        printError("AWS error '%s' (%s), will retry in %d ms",
            error.GetExceptionName(), error.GetMessage(), res);
        return res;
    }
};

ref<Aws::Client::ClientConfiguration> S3Helper::makeConfig(const string & region)
{
    initAWS();
    auto res = make_ref<Aws::Client::ClientConfiguration>();
    res->region = region;
    res->requestTimeoutMs = 600 * 1000;
    res->retryStrategy = std::make_shared<RetryStrategy>();
    res->caFile = settings.caFile;
    return res;
}

S3Helper::DownloadResult S3Helper::getObject(
    const std::string & bucketName, const std::string & key)
{
    debug("fetching ‘s3://%s/%s’...", bucketName, key);

    auto request =
        Aws::S3::Model::GetObjectRequest()
        .WithBucket(bucketName)
        .WithKey(key);

    request.SetResponseStreamFactory([&]() {
        return Aws::New<std::stringstream>("STRINGSTREAM");
    });

    DownloadResult res;

    auto now1 = std::chrono::steady_clock::now();

    try {

        auto result = checkAws(fmt("AWS error fetching ‘%s’", key),
            client->GetObject(request));

        res.data = decodeContent(
            result.GetContentEncoding(),
            make_ref<std::string>(
                dynamic_cast<std::stringstream &>(result.GetBody()).str()));

    } catch (S3Error & e) {
        if (e.err != Aws::S3::S3Errors::NO_SUCH_KEY) throw;
    }

    auto now2 = std::chrono::steady_clock::now();

    res.durationMs = std::chrono::duration_cast<std::chrono::milliseconds>(now2 - now1).count();

    return res;
}

struct S3BinaryCacheStoreImpl : public S3BinaryCacheStore
{
    const Setting<std::string> region{this, Aws::Region::US_EAST_1, "region", {"aws-region"}};
    const Setting<std::string> narinfoCompression{this, "", "narinfo-compression", "compression method for .narinfo files"};
    const Setting<std::string> lsCompression{this, "", "ls-compression", "compression method for .ls files"};
    const Setting<std::string> logCompression{this, "", "log-compression", "compression method for log/* files"};

    std::string bucketName;

    Stats stats;

    S3Helper s3Helper;

    S3BinaryCacheStoreImpl(
        const Params & params, const std::string & bucketName)
        : S3BinaryCacheStore(params)
        , bucketName(bucketName)
        , s3Helper(region)
    {
        diskCache = getNarInfoDiskCache();
    }

    std::string getUri() override
    {
        return "s3://" + bucketName;
    }

    void init() override
    {
        if (!diskCache->cacheExists(getUri(), wantMassQuery_, priority)) {

            /* Create the bucket if it doesn't already exists. */
            // FIXME: HeadBucket would be more appropriate, but doesn't return
            // an easily parsed 404 message.
            auto res = s3Helper.client->GetBucketLocation(
                Aws::S3::Model::GetBucketLocationRequest().WithBucket(bucketName));

            if (!res.IsSuccess()) {
                if (res.GetError().GetErrorType() != Aws::S3::S3Errors::NO_SUCH_BUCKET)
                    throw Error(format("AWS error checking bucket ‘%s’: %s") % bucketName % res.GetError().GetMessage());

                checkAws(format("AWS error creating bucket ‘%s’") % bucketName,
                    s3Helper.client->CreateBucket(
                        Aws::S3::Model::CreateBucketRequest()
                        .WithBucket(bucketName)
                        .WithCreateBucketConfiguration(
                            Aws::S3::Model::CreateBucketConfiguration()
                            /* .WithLocationConstraint(
                               Aws::S3::Model::BucketLocationConstraint::US) */ )));
            }

            BinaryCacheStore::init();

            diskCache->createCache(getUri(), storeDir, wantMassQuery_, priority);
        }
    }

    const Stats & getS3Stats() override
    {
        return stats;
    }

    /* This is a specialisation of isValidPath() that optimistically
       fetches the .narinfo file, rather than first checking for its
       existence via a HEAD request. Since .narinfos are small, doing
       a GET is unlikely to be slower than HEAD. */
    bool isValidPathUncached(const Path & storePath) override
    {
        try {
            queryPathInfo(storePath);
            return true;
        } catch (InvalidPath & e) {
            return false;
        }
    }

    bool fileExists(const std::string & path) override
    {
        stats.head++;

        auto res = s3Helper.client->HeadObject(
            Aws::S3::Model::HeadObjectRequest()
            .WithBucket(bucketName)
            .WithKey(path));

        if (!res.IsSuccess()) {
            auto & error = res.GetError();
            if (error.GetErrorType() == Aws::S3::S3Errors::UNKNOWN // FIXME
                && error.GetMessage().find("404") != std::string::npos)
                return false;
            throw Error(format("AWS error fetching ‘%s’: %s") % path % error.GetMessage());
        }

        return true;
    }

    void uploadFile(const std::string & path, const std::string & data,
        const std::string & mimeType,
        const std::string & contentEncoding)
    {
        auto request =
            Aws::S3::Model::PutObjectRequest()
            .WithBucket(bucketName)
            .WithKey(path);

        request.SetContentType(mimeType);

        if (contentEncoding != "")
            request.SetContentEncoding(contentEncoding);

        auto stream = std::make_shared<istringstream_nocopy>(data);

        request.SetBody(stream);

        stats.put++;
        stats.putBytes += data.size();

        auto now1 = std::chrono::steady_clock::now();

        auto result = checkAws(format("AWS error uploading ‘%s’") % path,
            s3Helper.client->PutObject(request));

        auto now2 = std::chrono::steady_clock::now();

        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(now2 - now1).count();

        printInfo(format("uploaded ‘s3://%1%/%2%’ (%3% bytes) in %4% ms")
            % bucketName % path % data.size() % duration);

        stats.putTimeMs += duration;
    }

    void upsertFile(const std::string & path, const std::string & data,
        const std::string & mimeType) override
    {
        if (narinfoCompression != "" && hasSuffix(path, ".narinfo"))
            uploadFile(path, *compress(narinfoCompression, data), mimeType, narinfoCompression);
        else if (lsCompression != "" && hasSuffix(path, ".ls"))
            uploadFile(path, *compress(lsCompression, data), mimeType, lsCompression);
        else if (logCompression != "" && hasPrefix(path, "log/"))
            uploadFile(path, *compress(logCompression, data), mimeType, logCompression);
        else
            uploadFile(path, data, mimeType, "");
    }

    void getFile(const std::string & path,
        std::function<void(std::shared_ptr<std::string>)> success,
        std::function<void(std::exception_ptr exc)> failure) override
    {
        sync2async<std::shared_ptr<std::string>>(success, failure, [&]() {
            stats.get++;

            auto res = s3Helper.getObject(bucketName, path);

            stats.getBytes += res.data ? res.data->size() : 0;
            stats.getTimeMs += res.durationMs;

            if (res.data)
                printTalkative("downloaded ‘s3://%s/%s’ (%d bytes) in %d ms",
                    bucketName, path, res.data->size(), res.durationMs);

            return res.data;
        });
    }

    PathSet queryAllValidPaths() override
    {
        PathSet paths;
        std::string marker;

        do {
            debug(format("listing bucket ‘s3://%s’ from key ‘%s’...") % bucketName % marker);

            auto res = checkAws(format("AWS error listing bucket ‘%s’") % bucketName,
                s3Helper.client->ListObjects(
                    Aws::S3::Model::ListObjectsRequest()
                    .WithBucket(bucketName)
                    .WithDelimiter("/")
                    .WithMarker(marker)));

            auto & contents = res.GetContents();

            debug(format("got %d keys, next marker ‘%s’")
                % contents.size() % res.GetNextMarker());

            for (auto object : contents) {
                auto & key = object.GetKey();
                if (key.size() != 40 || !hasSuffix(key, ".narinfo")) continue;
                paths.insert(storeDir + "/" + key.substr(0, key.size() - 8));
            }

            marker = res.GetNextMarker();
        } while (!marker.empty());

        return paths;
    }

};

static RegisterStoreImplementation regStore([](
    const std::string & uri, const Store::Params & params)
    -> std::shared_ptr<Store>
{
    if (std::string(uri, 0, 5) != "s3://") return 0;
    auto store = std::make_shared<S3BinaryCacheStoreImpl>(params, std::string(uri, 5));
    store->init();
    return store;
});

}

#endif
<|endoftext|>


<commit_before>#line 2 "togo/core/lua/lua.cpp"
/**
@copyright MIT license; see @ref index or the accompanying LICENSE file.
*/

#include <togo/core/config.hpp>
#include <togo/core/error/assert.hpp>
#include <togo/core/utility/utility.hpp>
#include <togo/core/memory/memory.hpp>
#include <togo/core/string/string.hpp>
#include <togo/core/lua/lua.hpp>

#include <cstring>

namespace togo {

static void* alloc_func(void* ud, void* ptr, size_t osize, size_t nsize) {
	auto a = static_cast<Allocator*>(ud);
	if (nsize == 0) {
		a->deallocate(ptr);
		return nullptr;
	}

	void* new_ptr = a->allocate(nsize);
	// TODO: does Lua do resize blocks a lot? maybe this should be a nop
	// (allocators don't do reallocs)
	if (osize > 0) {
		std::memcpy(new_ptr, ptr, nsize);
	}
	a->deallocate(ptr);
	return new_ptr;
}

/// Create a new state with the given allocator.
///
/// allocator must exist as long as the state.
lua_State* lua::new_state(Allocator& allocator) {
	return lua_newstate(alloc_func, &allocator);
}

/// Create a new state with Lua's allocator.
lua_State* lua::new_state() {
	return luaL_newstate();
}

/// lua_pcall() error message handler.
///
/// Returns a traceback as a string.
signed lua::pcall_error_message_handler(lua_State* L) {
	if (lua_isstring(L, 1)) {
		luaL_traceback(L, L, lua_tostring(L, 1), 1);
	}
	return 1;
}

namespace {

static signed module_loader(lua_State* L) {
	StringRef specified_chunk_name = lua::get_string(L, lua_upvalueindex(1));
	StringRef source = lua::get_string(L, lua_upvalueindex(2));
	auto funcs = array_cref(
		unsigned_cast(lua::get_integer(L, lua_upvalueindex(3))),
		lua::get_lightuserdata_typed<LuaModuleFunction const>(L, lua_upvalueindex(4), false)
	);
	StringRef modname = lua::get_string(L, 1);
	FixedArray<char, 128> chunk_name{};

	string::append(chunk_name, "@");
	if (specified_chunk_name.any()) {
		string::append(chunk_name, specified_chunk_name);
	} else {
		string::append(chunk_name, modname);
	}

	if (luaL_loadbufferx(L, source.data, source.size, begin(chunk_name), "t")) {
		return lua_error(L);
	}
	// -> chunk
	lua::push_value(L, modname);
	lua_call(L, 1, 1);
	// -> result

	if (lua_type(L, -1) != LUA_TTABLE) {
		luaL_error(L, "wanted a table return value from module, but got a %s", luaL_typename(L, -1));
	}
	if (begin(funcs)) {
		for (auto const& mf : funcs) {
			lua::table_set(L, mf.name, mf.func);
		}
	}

	lua::table_get(L, "__module_init__");
	if (lua_type(L, -1) == LUA_TFUNCTION) {
		lua_pushvalue(L, -2);
		lua_call(L, 1, 0);
	} else {
		lua_pop(L, 1);
	}

	// remove preloader
	lua::table_get_raw(L, LUA_REGISTRYINDEX, "_PRELOAD");
	lua::table_set_raw(L, modname, null_tag{});
	lua_pop(L, 1);

	return 1; // result
}

static void push_module_loader(lua_State* L, LuaModuleRef const& module) {
	lua::push_value(L, module.chunk_name);
	lua::push_value(L, module.source);
	lua::push_value(L, module.funcs.size());
	lua::push_lightuserdata(L, const_cast<LuaModuleFunction*>(begin(module.funcs)));
	lua_pushcclosure(L, module_loader, 4);
}

static void call_module_loader(lua_State* L, StringRef name, bool return_module) {
	lua::push_value(L, name);
	if (lua_pcall(L, 1, 1, -3)) {
		auto err = lua::get_string(L, -1);
		TOGO_ASSERTF(false, "error loading module: %.*s", err.size, err.data);
	}
	lua_remove(L, -2); // pcall_message_handler

	lua::table_get_raw(L, LUA_REGISTRYINDEX, "_LOADED");
	lua::table_set_copy_raw(L, name, -2);
	lua_pop(L, 1); // _LOADED

	if (!return_module) {
		lua_pop(L, 1);
	}
}

} // anonymous namespace

/// Put module in package.preload[modname].
///
/// source must be valid as long as the state is alive.
void lua::preload_module(lua_State* L, LuaModuleRef const& module) {
	lua::table_get_raw(L, LUA_REGISTRYINDEX, "_PRELOAD");
	lua::push_value(L, module.name);
	push_module_loader(L, module);
	lua_rawset(L, -3); // _PRELOAD[module.name] = closure
	lua_pop(L, 1); // _PRELOAD
}

/// Load module.
///
/// If return_module == true, this will return the module table at the top of
/// the stack.
void lua::load_module(
	lua_State* L,
	LuaModuleRef const& module,
	bool return_module IGEN_DEFAULT(false)
) {
	lua::push_value(L, lua::pcall_error_message_handler);
	push_module_loader(L, module);
	call_module_loader(L, module.name, return_module);
}

/// Load module by name.
void lua::load_module(
	lua_State* L,
	StringRef name,
	bool return_module IGEN_DEFAULT(false)
) {
	lua::table_get_raw(L, LUA_REGISTRYINDEX, "_LOADED");
	lua::table_get_raw(L, name);
	lua_remove(L, -2);
	if (lua_type(L, -1) != LUA_TNIL) {
		if (!return_module) {
			lua_pop(L, 1);
		}
		return;
	}
	lua_pop(L, 1);

	lua::push_value(L, lua::pcall_error_message_handler);
	lua::table_get_raw(L, LUA_REGISTRYINDEX, "_PRELOAD");
	lua::table_get_raw(L, name);
	lua_remove(L, -2);

	TOGO_ASSERTF(
		lua_type(L, -1) == LUA_TFUNCTION,
		"module not found in _LOADED or _PRELOAD: %.*s",
		name.size, name.data
	);
	call_module_loader(L, name, return_module);
}

/// Register basic interfaces.
void lua::register_core(lua_State* L) {
	lua_createtable(L, 0, 32);
	lua::table_set_copy_raw(L, LUA_REGISTRYINDEX, "togo_class", -1);
	lua_pop(L, 1);

	utility::register_lua_interface(L);
}

} // namespace togo
<commit_msg>lib/core/lua: call_module_loader(): handle possible return value of non-togo-ian modules.<commit_after>#line 2 "togo/core/lua/lua.cpp"
/**
@copyright MIT license; see @ref index or the accompanying LICENSE file.
*/

#include <togo/core/config.hpp>
#include <togo/core/error/assert.hpp>
#include <togo/core/utility/utility.hpp>
#include <togo/core/memory/memory.hpp>
#include <togo/core/string/string.hpp>
#include <togo/core/lua/lua.hpp>

#include <cstring>

namespace togo {

static void* alloc_func(void* ud, void* ptr, size_t osize, size_t nsize) {
	auto a = static_cast<Allocator*>(ud);
	if (nsize == 0) {
		a->deallocate(ptr);
		return nullptr;
	}

	void* new_ptr = a->allocate(nsize);
	// TODO: does Lua do resize blocks a lot? maybe this should be a nop
	// (allocators don't do reallocs)
	if (osize > 0) {
		std::memcpy(new_ptr, ptr, nsize);
	}
	a->deallocate(ptr);
	return new_ptr;
}

/// Create a new state with the given allocator.
///
/// allocator must exist as long as the state.
lua_State* lua::new_state(Allocator& allocator) {
	return lua_newstate(alloc_func, &allocator);
}

/// Create a new state with Lua's allocator.
lua_State* lua::new_state() {
	return luaL_newstate();
}

/// lua_pcall() error message handler.
///
/// Returns a traceback as a string.
signed lua::pcall_error_message_handler(lua_State* L) {
	if (lua_isstring(L, 1)) {
		luaL_traceback(L, L, lua_tostring(L, 1), 1);
	}
	return 1;
}

namespace {

static signed module_loader(lua_State* L) {
	StringRef specified_chunk_name = lua::get_string(L, lua_upvalueindex(1));
	StringRef source = lua::get_string(L, lua_upvalueindex(2));
	auto funcs = array_cref(
		unsigned_cast(lua::get_integer(L, lua_upvalueindex(3))),
		lua::get_lightuserdata_typed<LuaModuleFunction const>(L, lua_upvalueindex(4), false)
	);
	StringRef modname = lua::get_string(L, 1);
	FixedArray<char, 128> chunk_name{};

	string::append(chunk_name, "@");
	if (specified_chunk_name.any()) {
		string::append(chunk_name, specified_chunk_name);
	} else {
		string::append(chunk_name, modname);
	}

	if (luaL_loadbufferx(L, source.data, source.size, begin(chunk_name), "t")) {
		return lua_error(L);
	}
	// -> chunk
	lua::push_value(L, modname);
	lua_call(L, 1, 1);
	// -> result

	if (lua_type(L, -1) != LUA_TTABLE) {
		luaL_error(L, "wanted a table return value from module, but got a %s", luaL_typename(L, -1));
	}
	if (begin(funcs)) {
		for (auto const& mf : funcs) {
			lua::table_set(L, mf.name, mf.func);
		}
	}

	lua::table_get(L, "__module_init__");
	if (lua_type(L, -1) == LUA_TFUNCTION) {
		lua_pushvalue(L, -2);
		lua_call(L, 1, 0);
	} else {
		lua_pop(L, 1);
	}

	// remove preloader
	lua::table_get_raw(L, LUA_REGISTRYINDEX, "_PRELOAD");
	lua::table_set_raw(L, modname, null_tag{});
	lua_pop(L, 1);

	return 1; // result
}

static void push_module_loader(lua_State* L, LuaModuleRef const& module) {
	lua::push_value(L, module.chunk_name);
	lua::push_value(L, module.source);
	lua::push_value(L, module.funcs.size());
	lua::push_lightuserdata(L, const_cast<LuaModuleFunction*>(begin(module.funcs)));
	lua_pushcclosure(L, module_loader, 4);
}

static void call_module_loader(lua_State* L, StringRef name, bool return_module) {
	lua::push_value(L, name);
	if (lua_pcall(L, 1, 1, -3)) {
		auto err = lua::get_string(L, -1);
		TOGO_ASSERTF(false, "error loading module: %.*s", err.size, err.data);
	}
	lua_remove(L, -2); // pcall_error_message_handler

	lua::table_get_raw(L, LUA_REGISTRYINDEX, "_LOADED");
	{auto rtype = lua_type(L, -2);
	if (rtype == LUA_TNIL || (rtype == LUA_TBOOLEAN && !lua::get_boolean(L, -2))) {
		lua::table_set_raw(L, name, true);
	} else {
		lua::table_set_copy_raw(L, name, -2);
	}}
	lua_pop(L, 1);

	if (!return_module) {
		lua_pop(L, 1);
	}
}

} // anonymous namespace

/// Put module in package.preload[modname].
///
/// source must be valid as long as the state is alive.
void lua::preload_module(lua_State* L, LuaModuleRef const& module) {
	lua::table_get_raw(L, LUA_REGISTRYINDEX, "_PRELOAD");
	lua::push_value(L, module.name);
	push_module_loader(L, module);
	lua_rawset(L, -3); // _PRELOAD[module.name] = closure
	lua_pop(L, 1); // _PRELOAD
}

/// Load module.
///
/// If return_module == true, this will return the module table at the top of
/// the stack.
void lua::load_module(
	lua_State* L,
	LuaModuleRef const& module,
	bool return_module IGEN_DEFAULT(false)
) {
	lua::push_value(L, lua::pcall_error_message_handler);
	push_module_loader(L, module);
	call_module_loader(L, module.name, return_module);
}

/// Load module by name.
void lua::load_module(
	lua_State* L,
	StringRef name,
	bool return_module IGEN_DEFAULT(false)
) {
	lua::table_get_raw(L, LUA_REGISTRYINDEX, "_LOADED");
	lua::table_get_raw(L, name);
	lua_remove(L, -2);
	if (lua_type(L, -1) != LUA_TNIL) {
		if (!return_module) {
			lua_pop(L, 1);
		}
		return;
	}
	lua_pop(L, 1);

	lua::push_value(L, lua::pcall_error_message_handler);
	lua::table_get_raw(L, LUA_REGISTRYINDEX, "_PRELOAD");
	lua::table_get_raw(L, name);
	lua_remove(L, -2);

	TOGO_ASSERTF(
		lua_type(L, -1) == LUA_TFUNCTION,
		"module not found in _LOADED or _PRELOAD: %.*s",
		name.size, name.data
	);
	call_module_loader(L, name, return_module);
}

/// Register basic interfaces.
void lua::register_core(lua_State* L) {
	lua_createtable(L, 0, 32);
	lua::table_set_copy_raw(L, LUA_REGISTRYINDEX, "togo_class", -1);
	lua_pop(L, 1);

	utility::register_lua_interface(L);
}

} // namespace togo
<|endoftext|>