From c07bac4a5f81b8b2e5cd89d2e1f805d6694282c7 Mon Sep 17 00:00:00 2001
From: "Evgeniy A. Dushistov" <dushistov@mail.ru>
Date: Wed, 20 May 2020 12:50:21 +0300
Subject: [PATCH] feature: make possible to set port number in Uri::builder

---
 src/uri/builder.rs | 37 ++++++++++++++++++++++++++++++++++---
 src/uri/tests.rs   | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/src/uri/builder.rs b/src/uri/builder.rs
index 825c0faf..85dc619f 100644
--- a/src/uri/builder.rs
+++ b/src/uri/builder.rs
@@ -1,7 +1,10 @@
-use std::convert::{TryFrom, TryInto};
+use std::{
+    convert::{TryFrom, TryInto},
+    fmt::Write,
+};
 
-use super::{Authority, Parts, PathAndQuery, Scheme};
-use crate::Uri;
+use super::{Authority, ErrorKind, InvalidUriParts, Parts, PathAndQuery, Scheme};
+use crate::{byte_str::ByteStr, Uri};
 
 /// A builder for `Uri`s.
 ///
@@ -78,6 +81,34 @@ impl Builder {
         })
     }
 
+    /// Set the port number for URI, will be part of `Authority`
+    pub fn port<P>(self, port: P) -> Self
+    where
+        P: Into<u16>,
+    {
+        let port: u16 = port.into();
+        self.map(move |mut parts| {
+            let prev_auth = match parts.authority.as_ref() {
+                Some(auth) => {
+                    if auth.port().is_some() {
+                        return Err(InvalidUriParts::from(ErrorKind::InvalidPort).into());
+                    }
+
+                    auth.as_str()
+                }
+                None => "",
+            };
+            // 1 for ':', 5 for port number digists
+            let mut auth = String::with_capacity(prev_auth.len() + 6);
+            auth.push_str(prev_auth);
+            auth.push(':');
+            write!(&mut auth, "{}", port).expect("write to String failed");
+            let data: ByteStr = auth.into();
+            parts.authority = Some(Authority { data });
+            Ok(parts)
+        })
+    }
+
     /// Set the `PathAndQuery` for this URI.
     ///
     /// # Examples
diff --git a/src/uri/tests.rs b/src/uri/tests.rs
index 719cb94e..344e4565 100644
--- a/src/uri/tests.rs
+++ b/src/uri/tests.rs
@@ -517,3 +517,37 @@ fn test_partial_eq_path_with_terminating_questionmark() {
 
     assert_eq!(uri, a);
 }
+
+#[test]
+fn test_uri_builder() {
+    assert_eq!(
+        "ws://localhost:8000/demo/",
+        Uri::builder()
+            .scheme("ws")
+            .authority("localhost")
+            .port(8000u16)
+            .path_and_query("/demo/")
+            .build()
+            .unwrap()
+    );
+
+    let uri = Uri::from_str("192.168.1.30:33").unwrap();
+    assert_eq!(33, uri.port().unwrap());
+    assert_eq!(
+        "ws://[2001:db8:1f70::999:de8:7648:6e8]:33/demo/",
+        Uri::builder()
+            .scheme("ws")
+            .authority("[2001:db8:1f70::999:de8:7648:6e8]")
+            .port(uri.port().unwrap())
+            .path_and_query("/demo/")
+            .build()
+            .unwrap()
+    );
+    Uri::builder()
+        .scheme("ws")
+        .authority("localhost:8000")
+        .port(8000u16)
+        .path_and_query("/demo/")
+        .build()
+        .unwrap_err();
+}