From c9413f65ec4ab2b354c093967e364d979f62f9a5 Mon Sep 17 00:00:00 2001
From: shikai liu <lsk569937453@gmail.com>
Date: Wed, 29 May 2024 10:39:18 +0800
Subject: [PATCH] Add the http2 and http2-prior-knowledge.

---
 README.md             |  3 +++
 src/cli/app_config.rs | 14 +++++++-------
 src/http/handler.rs   | 34 +++++++++++++++++++++++++---------
 3 files changed, 35 insertions(+), 16 deletions(-)

diff --git a/README.md b/README.md
index 67ba01b..334a656 100644
--- a/README.md
+++ b/README.md
@@ -24,10 +24,13 @@ Options:
   -e, --referer <URL>                                      Referrer URL
   -o, --output <file>                                      Write to file instead of stdout
   -T, --upload-file <file>                                 Transfer local FILE to destination
+  -Q, --quote <command>                                    Send command(s) to server before transfer
   -k, --insecure                                           Allow insecure server connections
   -I, --head                                               Show document info only
   -r, --range <range>                                      Retrieve only the bytes within RANGE
   -v, --verbose                                            Make the operation more talkative
+      --http2
+      --http2-prior-knowledge
   -h, --help                                               Print help
   -V, --version                                            Print version
 ```
diff --git a/src/cli/app_config.rs b/src/cli/app_config.rs
index 5539cff..d3203e9 100644
--- a/src/cli/app_config.rs
+++ b/src/cli/app_config.rs
@@ -28,12 +28,7 @@ pub struct Cli {
     #[arg(short = 'A', long = "user-agent", value_name = "name")]
     pub user_agent_option: Option<String>,
     /// The Cookie option.
-    #[arg(
-        short = 'b',
-        long = "cookie",
-        value_name = "data|filename",
-        hide_short_help = true
-    )]
+    #[arg(short = 'b', long = "cookie", value_name = "data|filename")]
     pub cookie_option: Option<String>,
     ///  Referrer URL
     #[arg(short = 'e', long = "referer", value_name = "URL")]
@@ -47,7 +42,6 @@ pub struct Cli {
         default_missing_value = "none"
     )]
     pub file_path_option: Option<String>,
-
     ///  Transfer local FILE to destination
     #[arg(long = "upload-file", short = 'T', value_name = "file")]
     pub uploadfile_option: Option<String>,
@@ -68,6 +62,10 @@ pub struct Cli {
     ///  Make the operation more talkative
     #[arg(short = 'v', long = "verbose")]
     pub debug: bool,
+    #[arg(long = "http2")]
+    pub http2: bool,
+    #[arg(long = "http2-prior-knowledge")]
+    pub http2_prior_knowledge: bool,
     // /// Print help
     // #[arg(short, long, action = clap::ArgAction::Help)]
     // pub help: bool,
@@ -95,6 +93,8 @@ impl Cli {
             skip_certificate_validate: false,
             header_option: false,
             range_option: None,
+            http2: false,
+            http2_prior_knowledge: false,
             debug: false,
             // help: false,
             // help_all: false,
diff --git a/src/http/handler.rs b/src/http/handler.rs
index 035aeaf..2c4f2f0 100644
--- a/src/http/handler.rs
+++ b/src/http/handler.rs
@@ -227,6 +227,9 @@ pub async fn http_request(
     let mut request_builder = Request::builder()
         .method(method.as_str())
         .uri(cli.url.clone());
+    if cli.http2_prior_knowledge {
+        request_builder = request_builder.version(hyper::Version::HTTP_2);
+    }
     let mut header_map = HeaderMap::new();
     if let Some(content_type) = content_type_option {
         header_map.insert(CONTENT_TYPE, HeaderValue::from_str(&content_type)?);
@@ -329,20 +332,33 @@ pub async fn http_request(
     let request_future = {
         trace!("Start request");
         let fut = if scheme == "https" {
-            let https = hyper_rustls::HttpsConnectorBuilder::new()
+            let connector_builder = hyper_rustls::HttpsConnectorBuilder::new()
                 .with_tls_config(tls_config)
-                .https_only()
-                .enable_all_versions()
-                .build();
+                .https_only();
+            let https_connector = if cli.http2 {
+                connector_builder.enable_all_versions().build()
+            } else if cli.http2_prior_knowledge {
+                connector_builder.enable_http2().build()
+            } else {
+                connector_builder.enable_http1().build()
+            };
+
             let https_clientt: Client<_, Full<Bytes>> =
-                Client::builder(TokioExecutor::new()).build(https);
+                Client::builder(TokioExecutor::new()).build(https_connector);
             https_clientt.request(request)
         } else {
-            let http_client = Client::builder(TokioExecutor::new())
+            let mut http_client_builder = Client::builder(TokioExecutor::new());
+            http_client_builder
                 .http1_title_case_headers(true)
-                .http1_preserve_header_case(true)
-                .build_http();
-            http_client.request(request)
+                .http1_preserve_header_case(true);
+            let https_connector = if cli.http2 {
+                http_client_builder.http2_only(false).build_http()
+            } else if cli.http2_prior_knowledge {
+                http_client_builder.http2_only(true).build_http()
+            } else {
+                http_client_builder.build_http()
+            };
+            https_connector.request(request)
         };
         fut
     };