Use `DisplayBuffer` for socket addresses.

7 files changed, 34 insertions, 46 deletions

diff --git a/library/std/src/net/ip/display_buffer.rs b/library/std/src/net/addr/display_buffer.rs
index bd852d5da8e..7aadf06e92f 100644
--- a/library/std/src/net/ip/display_buffer.rs
+++ b/library/std/src/net/addr/display_buffer.rs
@@ -3,12 +3,12 @@ use crate::mem::MaybeUninit;
 use crate::str;
 
 /// Used for slow path in `Display` implementations when alignment is required.
-pub struct IpDisplayBuffer<const SIZE: usize> {
+pub struct DisplayBuffer<const SIZE: usize> {
     buf: [MaybeUninit<u8>; SIZE],
     len: usize,
 }
 
-impl<const SIZE: usize> IpDisplayBuffer<SIZE> {
+impl<const SIZE: usize> DisplayBuffer<SIZE> {
     #[inline]
     pub const fn new() -> Self {
         Self { buf: MaybeUninit::uninit_array(), len: 0 }
@@ -25,7 +25,7 @@ impl<const SIZE: usize> IpDisplayBuffer<SIZE> {
     }
 }
 
-impl<const SIZE: usize> fmt::Write for IpDisplayBuffer<SIZE> {
+impl<const SIZE: usize> fmt::Write for DisplayBuffer<SIZE> {
     fn write_str(&mut self, s: &str) -> fmt::Result {
         let bytes = s.as_bytes();
 
diff --git a/library/std/src/net/ip.rs b/library/std/src/net/addr/ip/mod.rs
index 6004810655e..a670f716833 100644
--- a/library/std/src/net/ip.rs
+++ b/library/std/src/net/addr/ip/mod.rs
@@ -8,8 +8,7 @@ use crate::mem::transmute;
 use crate::sys::net::netc as c;
 use crate::sys_common::{FromInner, IntoInner};
 
-mod display_buffer;
-use display_buffer::IpDisplayBuffer;
+use super::display_buffer::DisplayBuffer;
 
 /// An IP address, either IPv4 or IPv6.
 ///
@@ -997,7 +996,7 @@ impl fmt::Display for Ipv4Addr {
         } else {
             const LONGEST_IPV4_ADDR: &str = "255.255.255.255";
 
-            let mut buf = IpDisplayBuffer::<{ LONGEST_IPV4_ADDR.len() }>::new();
+            let mut buf = DisplayBuffer::<{ LONGEST_IPV4_ADDR.len() }>::new();
             // Buffer is long enough for the longest possible IPv4 address, so this should never fail.
             write!(buf, "{}.{}.{}.{}", octets[0], octets[1], octets[2], octets[3]).unwrap();
 
@@ -1844,7 +1843,7 @@ impl fmt::Display for Ipv6Addr {
         } else {
             const LONGEST_IPV6_ADDR: &str = "ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff";
 
-            let mut buf = IpDisplayBuffer::<{ LONGEST_IPV6_ADDR.len() }>::new();
+            let mut buf = DisplayBuffer::<{ LONGEST_IPV6_ADDR.len() }>::new();
             // Buffer is long enough for the longest possible IPv6 address, so this should never fail.
             write!(buf, "{}", self).unwrap();
 
diff --git a/library/std/src/net/ip/tests.rs b/library/std/src/net/addr/ip/tests.rs
index 7c3430b2b21..7c3430b2b21 100644
--- a/library/std/src/net/ip/tests.rs
+++ b/library/std/src/net/addr/ip/tests.rs
diff --git a/library/std/src/net/addr/mod.rs b/library/std/src/net/addr/mod.rs
new file mode 100644
index 00000000000..afecab30498
--- /dev/null
+++ b/library/std/src/net/addr/mod.rs
@@ -0,0 +1,4 @@
+mod display_buffer;
+
+pub mod ip;
+pub mod socket;
diff --git a/library/std/src/net/addr.rs b/library/std/src/net/addr/socket/mod.rs
index 53fee952a7a..33b0dfa03e0 100644
--- a/library/std/src/net/addr.rs
+++ b/library/std/src/net/addr/socket/mod.rs
@@ -2,9 +2,9 @@
 mod tests;
 
 use crate::cmp::Ordering;
-use crate::fmt;
+use crate::fmt::{self, Write};
 use crate::hash;
-use crate::io::{self, Write};
+use crate::io;
 use crate::iter;
 use crate::mem;
 use crate::net::{IpAddr, Ipv4Addr, Ipv6Addr};
@@ -15,6 +15,8 @@ use crate::sys_common::net::LookupHost;
 use crate::sys_common::{FromInner, IntoInner};
 use crate::vec;
 
+use super::display_buffer::DisplayBuffer;
+
 /// An internet socket address, either IPv4 or IPv6.
 ///
 /// Internet socket addresses consist of an [IP address], a 16-bit port number, as well
@@ -616,25 +618,18 @@ impl fmt::Debug for SocketAddr {
 #[stable(feature = "rust1", since = "1.0.0")]
 impl fmt::Display for SocketAddrV4 {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        // Fast path: if there's no alignment stuff, write to the output buffer
-        // directly
+        // If there are no alignment requirements, write the socket address directly to `f`.
+        // Otherwise, write it to a local buffer and then use `f.pad`.
         if f.precision().is_none() && f.width().is_none() {
             write!(f, "{}:{}", self.ip(), self.port())
         } else {
-            const IPV4_SOCKET_BUF_LEN: usize = (3 * 4)  // the segments
-                + 3  // the separators
-                + 1 + 5; // the port
-            let mut buf = [0; IPV4_SOCKET_BUF_LEN];
-            let mut buf_slice = &mut buf[..];
-
-            // Unwrap is fine because writing to a sufficiently-sized
-            // buffer is infallible
-            write!(buf_slice, "{}:{}", self.ip(), self.port()).unwrap();
-            let len = IPV4_SOCKET_BUF_LEN - buf_slice.len();
-
-            // This unsafe is OK because we know what is being written to the buffer
-            let buf = unsafe { crate::str::from_utf8_unchecked(&buf[..len]) };
-            f.pad(buf)
+            const LONGEST_IPV4_SOCKET_ADDR: &str = "255.255.255.255:65536";
+
+            let mut buf = DisplayBuffer::<{ LONGEST_IPV4_SOCKET_ADDR.len() }>::new();
+            // Buffer is long enough for the longest possible IPv4 socket address, so this should never fail.
+            write!(buf, "{}:{}", self.ip(), self.port()).unwrap();
+
+            f.pad(buf.as_str())
         }
     }
 }
@@ -649,35 +644,26 @@ impl fmt::Debug for SocketAddrV4 {
 #[stable(feature = "rust1", since = "1.0.0")]
 impl fmt::Display for SocketAddrV6 {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        // Fast path: if there's no alignment stuff, write to the output
-        // buffer directly
+        // If there are no alignment requirements, write the socket address directly to `f`.
+        // Otherwise, write it to a local buffer and then use `f.pad`.
         if f.precision().is_none() && f.width().is_none() {
             match self.scope_id() {
                 0 => write!(f, "[{}]:{}", self.ip(), self.port()),
                 scope_id => write!(f, "[{}%{}]:{}", self.ip(), scope_id, self.port()),
             }
         } else {
-            const IPV6_SOCKET_BUF_LEN: usize = (4 * 8)  // The address
-            + 7  // The colon separators
-            + 2  // The brackets
-            + 1 + 10 // The scope id
-            + 1 + 5; // The port
-
-            let mut buf = [0; IPV6_SOCKET_BUF_LEN];
-            let mut buf_slice = &mut buf[..];
+            const LONGEST_IPV6_SOCKET_ADDR: &str =
+                "[ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff%4294967296]:65536";
 
+            let mut buf = DisplayBuffer::<{ LONGEST_IPV6_SOCKET_ADDR.len() }>::new();
             match self.scope_id() {
-                0 => write!(buf_slice, "[{}]:{}", self.ip(), self.port()),
-                scope_id => write!(buf_slice, "[{}%{}]:{}", self.ip(), scope_id, self.port()),
+                0 => write!(buf, "[{}]:{}", self.ip(), self.port()),
+                scope_id => write!(buf, "[{}%{}]:{}", self.ip(), scope_id, self.port()),
             }
-            // Unwrap is fine because writing to a sufficiently-sized
-            // buffer is infallible
+            // Buffer is long enough for the longest possible IPv6 socket address, so this should never fail.
             .unwrap();
-            let len = IPV6_SOCKET_BUF_LEN - buf_slice.len();
 
-            // This unsafe is OK because we know what is being written to the buffer
-            let buf = unsafe { crate::str::from_utf8_unchecked(&buf[..len]) };
-            f.pad(buf)
+            f.pad(buf.as_str())
         }
     }
 }
diff --git a/library/std/src/net/addr/tests.rs b/library/std/src/net/addr/socket/tests.rs
index 585a17451a0..585a17451a0 100644
--- a/library/std/src/net/addr/tests.rs
+++ b/library/std/src/net/addr/socket/tests.rs
diff --git a/library/std/src/net/mod.rs b/library/std/src/net/mod.rs
index 6f9743f3a0e..8dd0936133d 100644
--- a/library/std/src/net/mod.rs
+++ b/library/std/src/net/mod.rs
@@ -24,9 +24,9 @@
 use crate::io::{self, ErrorKind};
 
 #[stable(feature = "rust1", since = "1.0.0")]
-pub use self::addr::{SocketAddr, SocketAddrV4, SocketAddrV6, ToSocketAddrs};
+pub use self::addr::ip::{IpAddr, Ipv4Addr, Ipv6Addr, Ipv6MulticastScope};
 #[stable(feature = "rust1", since = "1.0.0")]
-pub use self::ip::{IpAddr, Ipv4Addr, Ipv6Addr, Ipv6MulticastScope};
+pub use self::addr::socket::{SocketAddr, SocketAddrV4, SocketAddrV6, ToSocketAddrs};
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::parser::AddrParseError;
 #[unstable(feature = "tcplistener_into_incoming", issue = "88339")]
@@ -37,7 +37,6 @@ pub use self::tcp::{Incoming, TcpListener, TcpStream};
 pub use self::udp::UdpSocket;
 
 mod addr;
-mod ip;
 mod parser;
 mod tcp;
 #[cfg(test)]