src/stdio_redirection.cc

// Copyright 2023 Google Inc. All Rights Reserved.
//
// 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.

#include "stdio_redirection.h"

#include <thread>

#include "util.h"

StdioRedirector::StdioRedirector(IpcHandle& connection)
    : connection_(connection) {}

StdioRedirector::~StdioRedirector() {
  auto do_close = [](IpcHandle& old_handle, FILE* stream) {
    if (old_handle) {
      old_handle.CloneIntoStdio(stream);
      old_handle.Close();
    };
  };
  fflush(stderr);
  do_close(old_stderr_, stderr);

  fflush(stdout);
  do_close(old_stdout_, stdout);

  do_close(old_stdin_, stdin);
}

bool StdioRedirector::SendStandardDescriptors(std::string* err) {
  auto do_send = [this, err](FILE* stream) -> bool {
    return connection_.SendNativeHandle(IpcHandle::NativeForStdio(stream), err);
  };
  return do_send(stdin) && do_send(stdout) && do_send(stderr);
}

bool StdioRedirector::ReceiveStandardDescriptors(std::string* err) {
  auto do_receive = [this, err](FILE* stream, IpcHandle& old_handle) {
    IpcHandle new_handle;
    if (!connection_.ReceiveNativeHandle(&new_handle, err))
      return false;
    if (!new_handle) {
      *err = "Received invalid standard descriptor";
      return false;
    }
    old_handle = IpcHandle::CloneFromStdio(stream);
    if (!old_handle) {
      *err = "Could not save current standard descriptor";
      return false;
    }
    if (!new_handle.CloneIntoStdio(stream)) {
      *err = "Could not redirect standard descriptor";
      return false;
    }
    return true;
  };
  return do_receive(stdin, old_stdin_) && do_receive(stdout, old_stdout_) &&
         do_receive(stderr, old_stderr_);
}

StdioAsyncStringRedirector::StdioAsyncStringRedirector(AsyncLoop& async_loop,
                                                       FILE* stream)
    : stream_(stream), saved_handle_(IpcHandle::CloneFromStdio(stream)),
      async_loop_(async_loop) {
  fflush(stream_);
  IpcHandle read_handle, write_handle;
  std::string error;
  if (!IpcHandle::CreateAsyncPipe(&read_handle, &write_handle, &error))
    Fatal("ScopedBufferingStdio::CreatePipe: %s", error.c_str());

  // Ensure the read handle is never copied to remote processes.
  // IpcHandle::Clone() returns a new instance with CLOEXEC on Posix.
  read_handle.SetInheritable(false);

  read_handle_ = AsyncHandle::Create(std::move(read_handle), async_loop_,
      [this](AsyncError error, size_t size) {
        if (error || size == 0) {
          Close();
        } else {
          result_.append(buffer_, size);
          read_handle_.StartRead(buffer_, sizeof(buffer_));
        }
      });
  read_handle_.StartRead(buffer_, sizeof(buffer_));

#ifndef _WIN32
  // Ensure the write handle is in blocking mode since it will be
  // used as the new stdio descriptor.
  write_handle.SetNonBlocking(false);
#endif
  write_handle.CloneIntoStdio(stream_);
  write_handle.Close();
}

StdioAsyncStringRedirector::~StdioAsyncStringRedirector() {
  Close();
  Restore();
}

void StdioAsyncStringRedirector::Restore() {
  fflush(stream_);
  if (saved_handle_) {
    saved_handle_.CloneIntoStdio(stream_);
    saved_handle_.Close();
  }
}

bool StdioAsyncStringRedirector::IsActive() const {
  if (!read_handle_.IsRunning())
    return false;

  async_loop_.RunOnce(0);
  return read_handle_.IsRunning();
}

bool StdioAsyncStringRedirector::WaitForResult(std::string* result,
                                               int64_t timeout_ms) {
  Restore();
  auto status =
      async_loop_.RunUntil([this]() { return !read_handle_; }, timeout_ms);
  *result = result_;
  return status != AsyncLoop::ExitTimeout;
}

void StdioAsyncStringRedirector::Close() {
  read_handle_.Close();
}