// Copyright 2023 Northern.tech AS
//
//    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 <mender-update/update_module/v3/update_module.hpp>

#include <mender-update/progress_reader/progress_reader.hpp>

#include <common/events.hpp>
#include <common/events_io.hpp>
#include <common/log.hpp>
#include <common/path.hpp>
#include <common/processes.hpp>

namespace mender {
namespace update {
namespace update_module {
namespace v3 {

namespace log = mender::common::log;
namespace path = mender::common::path;
namespace processes = mender::common::processes;
namespace progress = mender::update::progress;

/* A note on how streams tree works in the download process:

According to Update Module protocol documentation, Update Modules need to do one full read of the
stream-next pipe, then a full read of a file that is mentioned in this pipe (payload pipe), and so
on, until they get a zero-length read from stream-next. This raises the question: is is possible
that stream-next pipe or the payload pipe are not  ready or even not exist when an Update Module
tries to read from it? How do we design against that?

The answer is: with current implementation it is not possible for such race conditions to occur.
The full download flow is as follows:

----------------------------------------------------------------------------------------------------
1. Before calling an Update Module with Download command, mender (meaning mender-orchestrator) first
creates a named pipe "stream-next", but doesn't open it for writing,
(UpdateModule::PrepareStreamNextPipe).

2. Now we have two possibilities:

2a. Mender is faster

mender opens the pipe for writing (UpdateModule::OpenStreamNextPipe), but there's no reader yet
(Update Module did not connect to the pipe): mender blocks until it can write to the pipe. As soon
as the Update Module connects, mender writes data to the pipe.

2b. Update Module is faster

Update Module tries to read from the pipe, but mender did not connect to the pipe for writing yet.
Because named pipes are blocking by default, Update Module will block, waiting for data to come from
the pipe. As soon as mender connects to the pipe, it will immediately start writing to it
(UpdateModule::OpenStreamNextPipe) and Update Module will begin reading from it.

There's no possibility of an Update Module getting an error saying that the pipe does not exist
instead, as it was created before the Update Module was called with a Download command. The Update
Module will always block on a pipe that exists but does not have a writer yet.

3. Update Module has read stream-next, receiving the name of the payload pipe.

When mender has written everything to stream-next, it does not close the pipe for writing yet. This
causes the Update Module to wait for additional data from the pipe indefinitely. Instead, mender
first calls UpdateModule::PrepareAndOpenStreamPipe where it creates the payload pipe and connects to
it.

Only after the payload pipe exists, mender closes the stream-next pipe, causing the Update Module to
unblock. Notice that thanks to this, there is no possibility of the Update Module being so fast that
it tries to read from the payload pipe before mender has a chance to create it - there's no
possibility of "file does not exist" happening. Mender is now free to start writing to the payload
pipe anytime
- it doesn't matter if the Update Module tries to read from it before mender is ready, as it will
just block on the pipe, waiting for mender to connect to it.

4. Update Module finished reading from the payload pipe. Tries to read again from stream-next.

Mender closes the payload pipe.
Because mender only closed stream-next pipe, but did not delete it, the Update Module won't get
"file does not exist" under any circumstances. Instead, it blocks on stream-next until mender is
ready to write to it. Mender is now free to open the stream-next for writing anytime, and the Update
Module will wait for it. As soon as mender connects to the pipe for writing, the Update Module
ublocks and starts reading from it.

5. Steps 2 - 4 repeat as many times as there are payloads and thus payload pipes.
Mender marks the download as completed internally (UpdateModule::StreamNextOpenHandler).

6. As soon as the Update Module tries to read from the stream-next pipe again, mender does not write
anything to it, and instead just connects to the pipe and immediately closes its connection. This
causes the Update Module to unblock and get a zero-length read from stream-next.

Update Module knows that there are no more payloads. Mender closes the stream-next pipe
(UpdateModule::StreamNextOpenHandler) and is free to continue to the next command.
----------------------------------------------------------------------------------------------------
Summary:
There are no places where it is possible for a race condition to occur. If the Update Module is
faster, it will wait for mender. If mender is faster, it will wait for the Update Module. As the
named pipes are created beforehand, there's no possibility of the Update Module getting "file does
not exist" error. This all assumes that a given Update Module implements the API correctly. If one
were to use non-blocking reads from the pipes in an Update Module, then it could break. API
explicitly says that "full reads" are required so if an Update Module doesn't block while there's a
writer on the other end of the pipe - it does not implement the API correctly.

Note:
What happens if an Update Module is broken, for example it does not read from stream-next the last
time, where it would receive a zero-length read? As soon as the Update Module returns after being
called with Download command, mender calls UpdateModule::ProcessEndedHandler where, if it sees that
not all reads were completed, it marks the download process as failed and closes all pipes
unblocking itself.
*/

void UpdateModule::StartDownloadProcess() {
	string download_command = "Download";
	if (download_->downloading_with_sizes_) {
		download_command = "DownloadWithFileSizes";
	}
	log::Debug(
		"Calling Update Module with command `" + update_module_path_ + " " + download_command + " "
		+ update_module_workdir_ + "`.");
	download_->proc_ = make_shared<procs::Process>(
		vector<string> {update_module_path_, download_command, update_module_workdir_});

	download_->proc_->SetWorkDir(update_module_workdir_);

	auto err = PrepareStreamNextPipe();
	if (err != error::NoError) {
		DownloadErrorHandler(err);
		return;
	}

	processes::OutputHandler stdout_handler {"Update Module output (stdout): "};
	processes::OutputHandler stderr_handler {"Update Module output (stderr): "};

	err = download_->proc_->Start(stdout_handler, stderr_handler);
	if (err != error::NoError) {
		DownloadErrorHandler(GetProcessError(err));
		return;
	}

	err = download_->proc_->AsyncWait(
		download_->event_loop_,
		[this](error::Error err) {
			if (err.code == make_error_condition(errc::timed_out)) {
				DownloadTimeoutHandler();
			} else {
				ProcessEndedHandler(err);
			}
		},
		chrono::seconds(ctx_.GetConfig().module_timeout_seconds));
	if (err != error::NoError) {
		DownloadErrorHandler(err);
		return;
	}

	DownloadErrorHandler(OpenStreamNextPipe(
		[this](io::ExpectedAsyncWriterPtr writer) { StreamNextOpenHandler(writer); }));
}

void UpdateModule::StreamNextOpenHandler(io::ExpectedAsyncWriterPtr writer) {
	if (!writer) {
		DownloadErrorHandler(writer.error());
		return;
	}
	download_->stream_next_writer_ = writer.value();

	download_->module_has_started_download_ = true;

	auto reader = download_->payload_.Next();
	if (!reader) {
		if (reader.error().code
			== artifact::parser_error::MakeError(
				   artifact::parser_error::NoMorePayloadFilesError, "")
				   .code) {
			download_->module_has_finished_download_ = true;
			log::Debug("Update Module finished all downloads");
			EndStreamNext();
		} else {
			DownloadErrorHandler(reader.error());
		}
		return;
	}
	auto payload_reader = make_shared<artifact::Reader>(std::move(reader.value()));

	auto progress_reader = make_shared<progress::Reader>(payload_reader, payload_reader->Size());

	download_->current_payload_reader_ =
		make_shared<events::io::AsyncReaderFromReader>(download_->event_loop_, progress_reader);
	download_->current_payload_name_ = payload_reader->Name();
	download_->current_payload_size_ = payload_reader->Size();

	auto stream_path =
		path::Join(update_module_workdir_, string("streams"), download_->current_payload_name_);
	auto exp_path_is_safe =
		path::IsWithinOrEqual(stream_path, path::Join(update_module_workdir_, string("streams")));
	if (!exp_path_is_safe.has_value()) {
		DownloadErrorHandler(exp_path_is_safe.error().WithContext(
			"Error checking if path is equal to or within directory"));
		return;
	}
	if (!exp_path_is_safe.value()) {
		DownloadErrorHandler(error::Error(
			make_error_condition(errc::invalid_argument),
			"Error downloading payload: Provided payload file (" + download_->current_payload_name_
				+ ") would point outside work directory when extracted."));
		return;
	}

	auto err = PrepareAndOpenStreamPipe(
		stream_path, [this](io::ExpectedAsyncWriterPtr writer) { StreamOpenHandler(writer); });
	if (err != error::NoError) {
		DownloadErrorHandler(err);
		return;
	}

	string stream_next_string;
	if (download_->downloading_with_sizes_) {
		stream_next_string = path::Join("streams", download_->current_payload_name_) + " "
							 + to_string(download_->current_payload_size_);
	} else {
		stream_next_string = path::Join("streams", download_->current_payload_name_);
	}
	size_t entry_size = stream_next_string.size() + 1;
	if (entry_size > download_->buffer_.size()) {
		DownloadErrorHandler(error::Error(
			make_error_condition(errc::no_buffer_space), "Payload name is too large for buffer"));
		return;
	}
	copy(stream_next_string.begin(), stream_next_string.end(), download_->buffer_.begin());
	download_->buffer_[entry_size - 1] = '\n';
	DownloadErrorHandler(download_->stream_next_writer_->AsyncWrite(
		download_->buffer_.begin(),
		download_->buffer_.begin() + entry_size,
		[this, entry_size](io::ExpectedSize result) {
			StreamNextWriteHandler(entry_size, result);
		}));
}

void UpdateModule::StreamOpenHandler(io::ExpectedAsyncWriterPtr writer) {
	if (!writer) {
		DownloadErrorHandler(writer.error());
		return;
	}
	download_->current_stream_writer_ = writer.value();

	DownloadErrorHandler(download_->current_payload_reader_->AsyncRead(
		download_->buffer_.begin(), download_->buffer_.end(), [this](io::ExpectedSize result) {
			PayloadReadHandler(result);
		}));
}

void UpdateModule::StreamNextWriteHandler(size_t expected_n, io::ExpectedSize result) {
	// Close stream-next writer.
	download_->stream_next_writer_.reset();
	if (!result) {
		DownloadErrorHandler(result.error());
	} else if (expected_n != result.value()) {
		DownloadErrorHandler(error::Error(
			make_error_condition(errc::io_error),
			"Unexpected number of written bytes to stream-next"));
	}
}

void UpdateModule::PayloadReadHandler(io::ExpectedSize result) {
	if (!result) {
		// Close streams.
		download_->current_stream_writer_.reset();
		download_->current_payload_reader_.reset();
		DownloadErrorHandler(result.error());
	} else if (result.value() > 0) {
		DownloadErrorHandler(download_->current_stream_writer_->AsyncWrite(
			download_->buffer_.begin(),
			download_->buffer_.begin() + result.value(),
			[this, result](io::ExpectedSize write_result) {
				StreamWriteHandler(0, result.value(), write_result);
			}));
	} else {
		// Close streams.
		download_->current_stream_writer_.reset();
		download_->current_payload_reader_.reset();

		if (download_->downloading_to_files_) {
			StartDownloadToFile();
		} else {
			DownloadErrorHandler(OpenStreamNextPipe(
				[this](io::ExpectedAsyncWriterPtr writer) { StreamNextOpenHandler(writer); }));
		}
	}
}

void UpdateModule::StreamWriteHandler(size_t offset, size_t expected_n, io::ExpectedSize result) {
	if (!result) {
		DownloadErrorHandler(result.error());
	} else if (result.value() == 0 || result.value() > expected_n) {
		DownloadErrorHandler(error::Error(
			make_error_condition(errc::io_error),
			"Unexpected number of written bytes to download stream"));
	} else if (result.value() < expected_n) {
		auto new_offset = offset + result.value();
		auto new_expected = expected_n - result.value();
		DownloadErrorHandler(download_->current_stream_writer_->AsyncWrite(
			download_->buffer_.begin() + new_offset,
			download_->buffer_.begin() + new_offset + new_expected,
			[this, new_offset, new_expected](io::ExpectedSize write_result) {
				StreamWriteHandler(new_offset, new_expected, write_result);
			}));
	} else {
		download_->written_ += result.value();
		log::Trace("Wrote " + to_string(download_->written_) + " bytes to Update Module");
		DownloadErrorHandler(download_->current_payload_reader_->AsyncRead(
			download_->buffer_.begin(), download_->buffer_.end(), [this](io::ExpectedSize result) {
				PayloadReadHandler(result);
			}));
	}
}

void UpdateModule::EndStreamNext() {
	// Empty write.
	DownloadErrorHandler(download_->stream_next_writer_->AsyncWrite(
		download_->buffer_.begin(), download_->buffer_.begin(), [this](io::ExpectedSize result) {
			if (!result) {
				DownloadErrorHandler(result.error());
			} else {
				DownloadErrorHandler(error::NoError);
			}
			// Close writer.
			download_->stream_next_writer_.reset();
			// No further action necessary. Now we just need to wait for the process to finish.
		}));
}

void UpdateModule::DownloadErrorHandler(const error::Error &err) {
	if (err != error::NoError) {
		EndDownloadLoop(err);
	}
}

void UpdateModule::EndDownloadLoop(const error::Error &err) {
	download_->download_finished_handler_(err);
}

void UpdateModule::DownloadTimeoutHandler() {
	download_->proc_->EnsureTerminated();
	EndDownloadLoop(error::Error(
		make_error_condition(errc::timed_out), "Update Module Download process timed out"));
}

void UpdateModule::ProcessEndedHandler(error::Error err) {
	if (err != error::NoError) {
		err = GetProcessError(err);
		DownloadErrorHandler(error::Error(
			err.code, "Download: Update Module returned non-zero status: " + err.message));
	} else if (download_->module_has_finished_download_) {
		EndDownloadLoop(error::NoError);
	} else if (download_->module_has_started_download_) {
		DownloadErrorHandler(error::Error(
			make_error_condition(errc::broken_pipe),
			"Update Module started downloading, but did not finish"));
	} else {
		download_->downloading_to_files_ = true;
		download_->stream_next_opener_.reset();
		download_->current_stream_opener_.reset();
		err = DeleteStreamsFiles();
		if (err != error::NoError) {
			DownloadErrorHandler(err);
		} else {
			StartDownloadToFile();
		}
	}
}

void UpdateModule::StartDownloadToFile() {
	auto reader = download_->payload_.Next();
	if (!reader) {
		if (reader.error().code
			== artifact::parser_error::MakeError(
				   artifact::parser_error::NoMorePayloadFilesError, "")
				   .code) {
			log::Debug("Downloaded all files to `files` directory.");
			EndDownloadLoop(error::NoError);
		} else {
			DownloadErrorHandler(reader.error());
		}
		return;
	}
	auto payload_reader = make_shared<artifact::Reader>(std::move(reader.value()));
	download_->current_payload_reader_ =
		make_shared<events::io::AsyncReaderFromReader>(download_->event_loop_, payload_reader);
	download_->current_payload_name_ = payload_reader->Name();

	auto stream_path = path::Join(update_module_workdir_, string("files"));
	auto err = PrepareDownloadDirectory(stream_path);
	if (err != error::NoError) {
		DownloadErrorHandler(err);
		return;
	}

	stream_path = path::Join(stream_path, download_->current_payload_name_);
	auto exp_path_is_safe =
		path::IsWithinOrEqual(stream_path, path::Join(update_module_workdir_, string("files")));
	if (!exp_path_is_safe.has_value()) {
		DownloadErrorHandler(exp_path_is_safe.error().WithContext(
			"Error checking if path is equal to or within directory"));
		return;
	}
	if (!exp_path_is_safe.value()) {
		DownloadErrorHandler(error::Error(
			make_error_condition(errc::invalid_argument),
			"Error downloading payload: Provided payload file (" + download_->current_payload_name_
				+ ") would point outside work directory when extracted."));
		return;
	}

	auto current_stream_writer =
		make_shared<events::io::AsyncFileDescriptorWriter>(download_->event_loop_);
	err = current_stream_writer->Open(stream_path);
	if (err != error::NoError) {
		DownloadErrorHandler(err);
		return;
	}
	download_->current_stream_writer_ = current_stream_writer;

	DownloadErrorHandler(download_->current_payload_reader_->AsyncRead(
		download_->buffer_.begin(), download_->buffer_.end(), [this](io::ExpectedSize result) {
			PayloadReadHandler(result);
		}));
}

} // namespace v3
} // namespace update_module
} // namespace update
} // namespace mender
