// 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"

#ifndef _WIN32
#include <unistd.h>
#endif

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_);
}