Parse the PROXY protocol for Electrum RPC

Cargo.toml

@@ -37,6 +37,7 @@ socket2 = { version = "0.4", features = ["all"] }
 num_cpus = "1.12.0"
 page_size = "0.4.2"
 prometheus = "0.13"
+proxy-protocol = { version = "0.5", features = ["always_exhaustive"] }
 rayon = "1.5.0"
 rocksdb = "0.21.0"
 serde = "1.0.118"

src/electrum/server.rs

@@ -16,6 +16,7 @@ use std::thread;
 use bitcoin::hashes::sha256d::Hash as Sha256dHash;
 use error_chain::ChainedError;
 use hex;
+use proxy_protocol::{version1, version2, ProxyHeader};
 use serde_json::{from_str, Value};
 use sha2::{Digest, Sha256};
 
@@ -522,12 +523,6 @@ impl Connection {
                     match msg {
                         Message::Request(line) => {
                             let cmd: Value = from_str(&line).chain_err(|| "invalid JSON format")?;
-                            // These thread_locals hold the client IP address on the "peer" threads
-                            if get_rpc_addr().is_none() {
-                                if let Some(Value::String(proxy_ip)) = cmd.get("proxied-client-ip") {
-                                    set_rpc_addr(proxy_ip.trim().parse::<SocketAddr>().ok());
-                                }
-                            }
                             let reply = match (
                                 cmd.get("method"),
                                 cmd.get("params").unwrap_or(&empty_params),
@@ -565,26 +560,45 @@ impl Connection {
         tx: crossbeam_channel::Sender<Message>,
     ) -> Result<()> {
         loop {
-            let mut line = Vec::<u8>::new();
-            reader
-                .read_until(b'\n', &mut line)
-                .chain_err(|| "failed to read a request")?;
-            if line.is_empty() {
+            let mut recv_data = Vec::<u8>::new();
+            read_to_end(&mut reader, &mut recv_data).chain_err(|| "failed to read a request")?;
+            if recv_data.is_empty() {
                 tx.send(Message::Done).chain_err(|| "channel closed")?;
                 return Ok(());
             } else {
-                if line.starts_with(&[22, 3, 1]) {
+                if recv_data.starts_with(&[22, 3, 1]) {
                     // (very) naive SSL handshake detection
                     let _ = tx.send(Message::Done);
-                    bail!("invalid request - maybe SSL-encrypted data?: {:?}", line)
+                    bail!(
+                        "invalid request - maybe SSL-encrypted data?: {:?}",
+                        recv_data
+                    )
                 }
-                match String::from_utf8(line) {
-                    Ok(req) => tx
-                        .send(Message::Request(req))
-                        .chain_err(|| "channel closed")?,
-                    Err(err) => {
-                        let _ = tx.send(Message::Done);
-                        bail!("invalid UTF8: {}", err)
+                // Parse the PROXY protocol
+                // If it parses a header, then the slice will be moved to the start
+                // of the electrum data.
+                // If it doesn't detect a PROXY header it will return very quickly
+                // and not touch the slice position.
+                let mut slice = &recv_data[..];
+                if let Ok(p_header) = proxy_protocol::parse(&mut slice) {
+                    // These thread_locals hold the client IP address on the "peer" threads
+                    if get_rpc_addr().is_none() {
+                        let proxy_ip = proxy_header_to_source_socket_addr(p_header);
+                        if proxy_ip.is_some() {
+                            set_rpc_addr(proxy_ip);
+                        }
+                    }
+                };
+
+                for line in slice.lines() {
+                    match line {
+                        Ok(req) => tx
+                            .send(Message::Request(req))
+                            .chain_err(|| "channel closed")?,
+                        Err(err) => {
+                            let _ = tx.send(Message::Done);
+                            bail!("invalid UTF8: {}", err)
+                        }
                     }
                 }
             }
@@ -1053,3 +1067,46 @@ impl Read for ConnectionStream {
         }
     }
 }
+
+fn read_to_end<R: BufRead + ?Sized>(
+    r: &mut R,
+    buf: &mut Vec<u8>,
+) -> std::result::Result<usize, std::io::Error> {
+    let mut read = 0;
+    loop {
+        let count = {
+            let available = match r.fill_buf() {
+                Ok(n) => n,
+                Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
+                Err(e) => return Err(e),
+            };
+            buf.extend_from_slice(available);
+            available.len()
+        };
+        r.consume(count);
+        read += count;
+        if count == 0 {
+            return Ok(read);
+        }
+    }
+}
+
+fn proxy_header_to_source_socket_addr(p_header: ProxyHeader) -> Option<SocketAddr> {
+    match p_header {
+        ProxyHeader::Version1 {
+            addresses: version1::ProxyAddresses::Ipv4 { source, .. },
+        } => Some(SocketAddr::V4(source)),
+        ProxyHeader::Version1 {
+            addresses: version1::ProxyAddresses::Ipv6 { source, .. },
+        } => Some(SocketAddr::V6(source)),
+        ProxyHeader::Version2 {
+            addresses: version2::ProxyAddresses::Ipv4 { source, .. },
+            ..
+        } => Some(SocketAddr::V4(source)),
+        ProxyHeader::Version2 {
+            addresses: version2::ProxyAddresses::Ipv6 { source, .. },
+            ..
+        } => Some(SocketAddr::V6(source)),
+        _ => None,
+    }
+}

src/rest.rs

@@ -688,7 +688,12 @@ fn handle_request(
         None => HashMap::new(),
     };
 
-    info!("[client_ip: {:?}] handle {:?} {:?}", get_rest_addr(), method, uri);
+    info!(
+        "[client_ip: {:?}] handle {:?} {:?}",
+        get_rest_addr(),
+        method,
+        uri
+    );
     match (
         &method,
         path.first(),