pro.cpp 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327
  1. /*
  2. * Copyright (c) 2018-2020, Andreas Kling <kling@serenityos.org>
  3. *
  4. * SPDX-License-Identifier: BSD-2-Clause
  5. */
  6. #include <AK/FileStream.h>
  7. #include <AK/GenericLexer.h>
  8. #include <AK/LexicalPath.h>
  9. #include <AK/NumberFormat.h>
  10. #include <AK/URL.h>
  11. #include <LibCore/ArgsParser.h>
  12. #include <LibCore/EventLoop.h>
  13. #include <LibCore/File.h>
  14. #include <LibCore/System.h>
  15. #include <LibMain/Main.h>
  16. #include <LibProtocol/Request.h>
  17. #include <LibProtocol/RequestClient.h>
  18. #include <ctype.h>
  19. #include <stdio.h>
  20. // FIXME: Move this somewhere else when it's needed (e.g. in the Browser)
  21. class ContentDispositionParser {
  22. public:
  23. ContentDispositionParser(StringView value)
  24. {
  25. GenericLexer lexer(value);
  26. lexer.ignore_while(isspace);
  27. if (lexer.consume_specific("inline")) {
  28. m_kind = Kind::Inline;
  29. if (!lexer.is_eof())
  30. m_might_be_wrong = true;
  31. return;
  32. }
  33. if (lexer.consume_specific("attachment")) {
  34. m_kind = Kind::Attachment;
  35. if (lexer.consume_specific(";")) {
  36. lexer.ignore_while(isspace);
  37. if (lexer.consume_specific("filename=")) {
  38. // RFC 2183: "A short (length <= 78 characters)
  39. // parameter value containing only non-`tspecials' characters SHOULD be
  40. // represented as a single `token'."
  41. // Some people seem to take this as generic advice of "if it doesn't have special characters,
  42. // it's safe to specify as a single token"
  43. // So let's just be as lenient as possible.
  44. if (lexer.next_is('"'))
  45. m_filename = lexer.consume_quoted_string();
  46. else
  47. m_filename = lexer.consume_until(is_any_of("()<>@,;:\\\"/[]?= "));
  48. } else {
  49. m_might_be_wrong = true;
  50. }
  51. }
  52. return;
  53. }
  54. if (lexer.consume_specific("form-data")) {
  55. m_kind = Kind::FormData;
  56. while (lexer.consume_specific(";")) {
  57. lexer.ignore_while(isspace);
  58. if (lexer.consume_specific("name=")) {
  59. m_name = lexer.consume_quoted_string();
  60. } else if (lexer.consume_specific("filename=")) {
  61. if (lexer.next_is('"'))
  62. m_filename = lexer.consume_quoted_string();
  63. else
  64. m_filename = lexer.consume_until(is_any_of("()<>@,;:\\\"/[]?= "));
  65. } else {
  66. m_might_be_wrong = true;
  67. }
  68. }
  69. return;
  70. }
  71. // FIXME: Support 'filename*'
  72. m_might_be_wrong = true;
  73. }
  74. enum class Kind {
  75. Inline,
  76. Attachment,
  77. FormData,
  78. };
  79. StringView filename() const { return m_filename; }
  80. StringView name() const { return m_name; }
  81. Kind kind() const { return m_kind; }
  82. bool might_be_wrong() const { return m_might_be_wrong; }
  83. private:
  84. StringView m_filename;
  85. StringView m_name;
  86. Kind m_kind { Kind::Inline };
  87. bool m_might_be_wrong { false };
  88. };
  89. template<typename ConditionT>
  90. class ConditionalOutputFileStream final : public OutputFileStream {
  91. public:
  92. template<typename... Args>
  93. ConditionalOutputFileStream(ConditionT&& condition, Args... args)
  94. : OutputFileStream(args...)
  95. , m_condition(condition)
  96. {
  97. }
  98. ~ConditionalOutputFileStream()
  99. {
  100. if (!m_condition())
  101. return;
  102. if (!m_buffer.is_empty()) {
  103. OutputFileStream::write(m_buffer);
  104. m_buffer.clear();
  105. }
  106. }
  107. private:
  108. size_t write(ReadonlyBytes bytes) override
  109. {
  110. if (!m_condition()) {
  111. write_to_buffer:;
  112. // FIXME: Propagate errors.
  113. if (m_buffer.try_append(bytes.data(), bytes.size()).is_error())
  114. return 0;
  115. return bytes.size();
  116. }
  117. if (!m_buffer.is_empty()) {
  118. auto size = OutputFileStream::write(m_buffer);
  119. m_buffer = m_buffer.slice(size, m_buffer.size() - size);
  120. }
  121. if (!m_buffer.is_empty())
  122. goto write_to_buffer;
  123. return OutputFileStream::write(bytes);
  124. }
  125. ConditionT m_condition;
  126. ByteBuffer m_buffer;
  127. };
  128. ErrorOr<int> serenity_main(Main::Arguments arguments)
  129. {
  130. char const* url_str = nullptr;
  131. bool save_at_provided_name = false;
  132. bool should_follow_url = false;
  133. char const* data = nullptr;
  134. StringView proxy_spec;
  135. String method = "GET";
  136. HashMap<String, String, CaseInsensitiveStringTraits> request_headers;
  137. Core::ArgsParser args_parser;
  138. args_parser.set_general_help(
  139. "Request a file from an arbitrary URL. This command uses RequestServer, "
  140. "and thus supports at least http, https, and gemini.");
  141. args_parser.add_option(save_at_provided_name, "Write to a file named as the remote file", nullptr, 'O');
  142. args_parser.add_option(data, "(HTTP only) Send the provided data via an HTTP POST request", "data", 'd', "data");
  143. args_parser.add_option(should_follow_url, "(HTTP only) Follow the Location header if a 3xx status is encountered", "follow", 'l');
  144. args_parser.add_option(Core::ArgsParser::Option {
  145. .requires_argument = true,
  146. .help_string = "Add a header entry to the request",
  147. .long_name = "header",
  148. .short_name = 'H',
  149. .value_name = "header-value",
  150. .accept_value = [&](auto* s) {
  151. StringView header { s };
  152. auto split = header.find(':');
  153. if (!split.has_value())
  154. return false;
  155. request_headers.set(header.substring_view(0, split.value()), header.substring_view(split.value() + 1));
  156. return true;
  157. } });
  158. args_parser.add_option(proxy_spec, "Specify a proxy server to use for this request (proto://ip:port)", "proxy", 'p', "proxy");
  159. args_parser.add_positional_argument(url_str, "URL to download from", "url");
  160. args_parser.parse(arguments);
  161. if (data) {
  162. method = "POST";
  163. // FIXME: Content-Type?
  164. }
  165. URL url(url_str);
  166. if (!url.is_valid()) {
  167. warnln("'{}' is not a valid URL", url_str);
  168. return 1;
  169. }
  170. Core::ProxyData proxy_data {};
  171. if (!proxy_spec.is_empty())
  172. proxy_data = TRY(Core::ProxyData::parse_url(proxy_spec));
  173. Core::EventLoop loop;
  174. bool received_actual_headers = false;
  175. bool should_save_stream_data = false;
  176. bool following_url = false;
  177. u32 previous_downloaded_size = 0;
  178. u32 const report_time_in_ms = 100;
  179. u32 const speed_update_time_in_ms = 4000;
  180. timeval previous_time, current_time, time_diff;
  181. gettimeofday(&previous_time, nullptr);
  182. RefPtr<Protocol::Request> request;
  183. auto protocol_client = TRY(Protocol::RequestClient::try_create());
  184. auto output_stream = ConditionalOutputFileStream { [&] { return should_save_stream_data; }, stdout };
  185. Function<void()> setup_request = [&] {
  186. if (!request) {
  187. warnln("Failed to start request for '{}'", url_str);
  188. exit(1);
  189. }
  190. request->on_progress = [&](Optional<u32> maybe_total_size, u32 downloaded_size) {
  191. gettimeofday(&current_time, nullptr);
  192. timersub(&current_time, &previous_time, &time_diff);
  193. auto time_diff_ms = time_diff.tv_sec * 1000 + time_diff.tv_usec / 1000;
  194. if (time_diff_ms < report_time_in_ms)
  195. return;
  196. warn("\r\033[2K");
  197. if (maybe_total_size.has_value()) {
  198. warn("\033]9;{};{};\033\\", downloaded_size, maybe_total_size.value());
  199. warn("Download progress: {} / {}", human_readable_size(downloaded_size), human_readable_size(maybe_total_size.value()));
  200. } else {
  201. warn("Download progress: {} / ???", human_readable_size(downloaded_size));
  202. }
  203. auto size_diff = downloaded_size - previous_downloaded_size;
  204. if (time_diff_ms > speed_update_time_in_ms) {
  205. previous_time = current_time;
  206. previous_downloaded_size = downloaded_size;
  207. }
  208. warn(" at {}/s", human_readable_size(((float)size_diff / (float)time_diff_ms) * 1000));
  209. };
  210. request->on_headers_received = [&](auto& response_headers, auto status_code) {
  211. if (received_actual_headers)
  212. return;
  213. dbgln("Received headers! response code = {}", status_code.value_or(0));
  214. received_actual_headers = true; // And not trailers!
  215. should_save_stream_data = true;
  216. if (!following_url && save_at_provided_name) {
  217. String output_name;
  218. if (auto content_disposition = response_headers.get("Content-Disposition"); content_disposition.has_value()) {
  219. auto& value = content_disposition.value();
  220. ContentDispositionParser parser(value);
  221. output_name = parser.filename();
  222. }
  223. if (output_name.is_empty())
  224. output_name = url.path();
  225. LexicalPath path { output_name };
  226. output_name = path.basename();
  227. // The URL didn't have a name component, e.g. 'serenityos.org'
  228. if (output_name.is_empty() || output_name == "/") {
  229. int i = -1;
  230. do {
  231. output_name = url.host();
  232. if (i > -1)
  233. output_name = String::formatted("{}.{}", output_name, i);
  234. ++i;
  235. } while (Core::File::exists(output_name));
  236. }
  237. if (freopen(output_name.characters(), "w", stdout) == nullptr) {
  238. perror("freopen");
  239. loop.quit(1);
  240. return;
  241. }
  242. }
  243. auto status_code_value = status_code.value_or(0);
  244. if (should_follow_url && status_code_value >= 300 && status_code_value < 400) {
  245. if (auto location = response_headers.get("Location"); location.has_value()) {
  246. auto was_following_url = following_url;
  247. following_url = true;
  248. received_actual_headers = false;
  249. should_save_stream_data = false;
  250. request->on_finish = nullptr;
  251. request->on_headers_received = nullptr;
  252. request->on_progress = nullptr;
  253. request->stop();
  254. Core::deferred_invoke([&, was_following_url, url = location.value()] {
  255. warnln("{}Following to {}", was_following_url ? "" : "\n", url);
  256. request = protocol_client->start_request(method, url, request_headers, ReadonlyBytes {}, proxy_data);
  257. setup_request();
  258. });
  259. }
  260. } else {
  261. following_url = false;
  262. }
  263. };
  264. request->on_finish = [&](bool success, auto) {
  265. if (following_url)
  266. return;
  267. warn("\033]9;-1;\033\\");
  268. warnln();
  269. if (!success)
  270. warnln("Request failed :(");
  271. loop.quit(0);
  272. };
  273. request->stream_into(output_stream);
  274. };
  275. request = protocol_client->start_request(method, url, request_headers, data ? StringView { data }.bytes() : ReadonlyBytes {}, proxy_data);
  276. setup_request();
  277. dbgln("started request with id {}", request->id());
  278. auto rc = loop.exec();
  279. // FIXME: This shouldn't be needed.
  280. fclose(stdout);
  281. return rc;
  282. }