Add support for datanode sasl negotation

Kimahriman · Kimahriman · commit 5e97c6c6aac6 · 2024-03-17T17:43:31.000-04:00
diff --git a/crates/hdfs-native/Cargo.toml b/crates/hdfs-native/Cargo.toml
@@ -22,7 +22,7 @@ libc = "0.2"
 libgssapi = { version = "0.7", default-features = false, optional = true }
 log = "0.4"
 num-traits = "0.2"
-once_cell = "1.19.0"
+once_cell = "1"
 prost = "0.12"
 prost-types = "0.12"
 roxmltree = "0.18"
diff --git a/crates/hdfs-native/minidfs/src/main/java/main/Main.java b/crates/hdfs-native/minidfs/src/main/java/main/Main.java
@@ -60,8 +60,10 @@ public static void main(String args[]) throws Exception {
             conf.set(DFSConfigKeys.DFS_DATANODE_KEYTAB_FILE_KEY, "target/test/hdfs.keytab");
             conf.set(DFS_DATANODE_KERBEROS_PRINCIPAL_KEY, "hdfs/localhost@" + kdc.getRealm());
             conf.set(DFSConfigKeys.DFS_BLOCK_ACCESS_TOKEN_ENABLE_KEY, "true");
-            // conf.set(DFS_DATA_TRANSFER_PROTECTION_KEY, "authentication");
             conf.set(DFSConfigKeys.IGNORE_SECURE_PORTS_FOR_TESTING_KEY, "true");
+            if (flags.contains("token")) {
+                conf.set(DFS_DATA_TRANSFER_PROTECTION_KEY, "authentication");
+            }
         }
 
         HdfsConfiguration hdfsConf = new HdfsConfiguration(conf);
diff --git a/crates/hdfs-native/src/hdfs/block_reader.rs b/crates/hdfs-native/src/hdfs/block_reader.rs
@@ -37,15 +37,15 @@ pub(crate) fn get_block_stream(
 
 /// Connects to a DataNode to do a read, attempting to used cached connections.
 async fn connect_and_send(
-    url: &str,
+    datanode_id: &hdfs::DatanodeIdProto,
     block: &hdfs::ExtendedBlockProto,
     token: common::TokenProto,
     offset: u64,
     len: u64,
 ) -> Result<(DatanodeConnection, BlockOpResponseProto)> {
     let mut remaining_attempts = 2;
     while remaining_attempts > 0 {
-        if let Some(mut conn) = DATANODE_CACHE.get(url) {
+        if let Some(mut conn) = DATANODE_CACHE.get(datanode_id) {
             let message = hdfs::OpReadBlockProto {
                 header: conn.build_header(block, Some(token.clone())),
                 offset,
@@ -68,7 +68,7 @@ async fn connect_and_send(
         }
         remaining_attempts -= 1;
     }
-    let mut conn = DatanodeConnection::connect(url).await?;
+    let mut conn = DatanodeConnection::connect(datanode_id, &token).await?;
 
     let message = hdfs::OpReadBlockProto {
         header: conn.build_header(block, Some(token)),
@@ -117,10 +117,9 @@ impl ReplicatedBlockStream {
         }
 
         let datanode = &self.block.locs[self.current_replica].id;
-        let datanode_url = format!("{}:{}", datanode.ip_addr, datanode.xfer_port);
 
         let (connection, response) = connect_and_send(
-            &datanode_url,
+            datanode,
             &self.block.b,
             self.block.block_token.clone(),
             self.offset as u64,
@@ -389,15 +388,8 @@ impl StripedBlockStream {
             return Ok(());
         }
 
-        let datanode_url = format!("{}:{}", datanode.ip_addr, datanode.xfer_port);
-        let (mut connection, response) = connect_and_send(
-            &datanode_url,
-            block,
-            token.clone(),
-            offset as u64,
-            len as u64,
-        )
-        .await?;
+        let (mut connection, response) =
+            connect_and_send(datanode, block, token.clone(), offset as u64, len as u64).await?;
 
         if response.status() != hdfs::Status::Success {
             return Err(HdfsError::DataTransferError(response.message().to_string()));
diff --git a/crates/hdfs-native/src/hdfs/block_writer.rs b/crates/hdfs-native/src/hdfs/block_writer.rs
@@ -100,9 +100,7 @@ impl ReplicatedBlockWriter {
         server_defaults: hdfs::FsServerDefaultsProto,
     ) -> Result<Self> {
         let datanode = &block.locs[0].id;
-        let mut connection =
-            DatanodeConnection::connect(&format!("{}:{}", datanode.ip_addr, datanode.xfer_port))
-                .await?;
+        let mut connection = DatanodeConnection::connect(datanode, &block.block_token).await?;
 
         let checksum = hdfs::ChecksumProto {
             r#type: hdfs::ChecksumTypeProto::ChecksumCrc32c as i32,
diff --git a/crates/hdfs-native/src/hdfs/connection.rs b/crates/hdfs-native/src/hdfs/connection.rs
@@ -25,11 +25,16 @@ use tokio::{
 use uuid::Uuid;
 
 use crate::proto::common::rpc_response_header_proto::RpcStatusProto;
+use crate::proto::common::TokenProto;
+use crate::proto::hdfs::DatanodeIdProto;
 use crate::proto::{common, hdfs};
 use crate::security::sasl::{SaslReader, SaslRpcClient, SaslWriter};
 use crate::security::user::UserInfo;
 use crate::{HdfsError, Result};
 
+#[cfg(feature = "token")]
+use crate::security::sasl::SaslDatanodeConnection;
+
 const PROTOCOL: &str = "org.apache.hadoop.hdfs.protocol.ClientProtocol";
 const DATA_TRANSFER_VERSION: u16 = 28;
 const MAX_PACKET_HEADER_SIZE: usize = 33;
@@ -520,12 +525,24 @@ pub(crate) struct DatanodeConnection {
 }
 
 impl DatanodeConnection {
-    pub(crate) async fn connect(url: &str) -> Result<Self> {
-        let stream = BufStream::new(connect(url).await?);
+    #[allow(unused_variables)]
+    pub(crate) async fn connect(datanode_id: &DatanodeIdProto, token: &TokenProto) -> Result<Self> {
+        let url = format!("{}:{}", datanode_id.ip_addr, datanode_id.xfer_port);
+        let stream = connect(&url).await?;
+
+        // If the token has an identifier, we can do SASL negotation
+        #[cfg(feature = "token")]
+        let stream = if token.identifier.is_empty() {
+            stream
+        } else {
+            debug!("{:?}", token);
+            let sasl_connection = SaslDatanodeConnection::create(stream);
+            sasl_connection.negotiate(datanode_id, token).await?
+        };
 
         let conn = DatanodeConnection {
             client_name: Uuid::new_v4().to_string(),
-            stream,
+            stream: BufStream::new(stream),
             url: url.to_string(),
         };
         Ok(conn)
@@ -704,15 +721,16 @@ impl DatanodeConnectionCache {
         }
     }
 
-    pub(crate) fn get(&self, url: &str) -> Option<DatanodeConnection> {
+    pub(crate) fn get(&self, datanode_id: &hdfs::DatanodeIdProto) -> Option<DatanodeConnection> {
         // Keep things simply and just expire cache entries when checking the cache. We could
         // move this to its own task but that will add a little more complexity.
         self.remove_expired();
 
+        let url = format!("{}:{}", datanode_id.ip_addr, datanode_id.xfer_port);
         let mut cache = self.cache.lock().unwrap();
 
         cache
-            .get_mut(url)
+            .get_mut(&url)
             .iter_mut()
             .flat_map(|conns| conns.pop_front())
             .map(|(_, conn)| conn)
diff --git a/crates/hdfs-native/src/security/sasl.rs b/crates/hdfs-native/src/security/sasl.rs
@@ -15,23 +15,38 @@ use crate::proto::common::rpc_sasl_proto::{SaslAuth, SaslState};
 use crate::proto::common::{
     RpcKindProto, RpcRequestHeaderProto, RpcResponseHeaderProto, RpcSaslProto,
 };
+
 use crate::{HdfsError, Result};
 #[cfg(feature = "token")]
 use {
-    super::user::Token,
+    super::user::{BlockTokenIdentifier, Token},
     base64::{engine::general_purpose, Engine as _},
     gsasl_sys as gsasl,
     libc::{c_char, c_void, memcpy},
     std::ffi::CString,
     std::ptr,
     std::sync::atomic::AtomicPtr,
 };
+#[cfg(feature = "token")]
+use {
+    crate::proto::{
+        common::TokenProto,
+        hdfs::{
+            data_transfer_encryptor_message_proto::DataTransferEncryptorStatus, CipherOptionProto,
+            CipherSuiteProto, DataTransferEncryptorMessageProto, DatanodeIdProto,
+            HandshakeSecretProto,
+        },
+    },
+    tokio::io::{AsyncBufReadExt, BufStream},
+};
 
 #[cfg(feature = "kerberos")]
 use super::gssapi::GssapiSession;
 use super::user::{User, UserInfo};
 
 const SASL_CALL_ID: i32 = -33;
+#[cfg(feature = "token")]
+const SASL_TRANSFER_MAGIC_NUMBER: i32 = 0xDEADBEEFu32 as i32;
 const HDFS_DELEGATION_TOKEN: &str = "HDFS_DELEGATION_TOKEN";
 
 pub(crate) enum AuthMethod {
@@ -309,17 +324,6 @@ impl SaslReader {
     }
 }
 
-// TODO: Can we implement this?
-// impl AsyncRead for SaslReader {
-//     fn poll_read(
-//         self: Pin<&mut Self>,
-//         cx: &mut task::Context<'_>,
-//         buf: &mut ReadBuf<'_>,
-//     ) -> Poll<io::Result<()>> {
-//         todo!()
-//     }
-// }
-
 pub(crate) struct SaslWriter {
     stream: OwnedWriteHalf,
     session: Option<Arc<Mutex<Box<dyn SaslSession>>>>,
@@ -547,3 +551,107 @@ impl Drop for GSASLSession {
         }
     }
 }
+
+#[cfg(feature = "token")]
+pub(crate) struct SaslDatanodeConnection {
+    stream: BufStream<TcpStream>,
+}
+
+#[cfg(feature = "token")]
+impl SaslDatanodeConnection {
+    pub fn create(stream: TcpStream) -> Self {
+        Self {
+            stream: BufStream::new(stream),
+        }
+    }
+
+    pub(crate) async fn negotiate(
+        mut self,
+        datanode_id: &DatanodeIdProto,
+        token: &TokenProto,
+    ) -> Result<TcpStream> {
+        // If it's a privileged port, don't do SASL negotation
+        if datanode_id.xfer_port <= 1024 {
+            return Ok(self.stream.into_inner());
+        }
+
+        self.stream.write_i32(SASL_TRANSFER_MAGIC_NUMBER).await?;
+        self.stream.flush().await?;
+
+        let mut session = GSASLSession::new("hdfs", "0", &token.clone().into())?;
+
+        let token_identifier = BlockTokenIdentifier::from_identifier(&token.identifier)?;
+
+        let handshake_secret = if !token_identifier.handshake_secret.is_empty() {
+            Some(HandshakeSecretProto {
+                bpid: token_identifier.block_pool_id.clone(),
+                secret: token_identifier.handshake_secret.clone(),
+            })
+        } else {
+            None
+        };
+
+        let message = DataTransferEncryptorMessageProto {
+            handshake_secret,
+            status: DataTransferEncryptorStatus::Success as i32,
+            ..Default::default()
+        };
+
+        debug!("Sending data transfer encryptor message: {:?}", message);
+
+        self.stream
+            .write_all(&message.encode_length_delimited_to_vec())
+            .await?;
+        self.stream.flush().await?;
+
+        let response = self.read_sasl_response().await?;
+        debug!("Data transfer encryptor response: {:?}", response);
+
+        let (payload, finished) = session.step(response.payload.as_ref().map(|p| &p[..]))?;
+        assert!(!finished);
+
+        let message = DataTransferEncryptorMessageProto {
+            status: DataTransferEncryptorStatus::Success as i32,
+            payload: Some(payload),
+            cipher_option: vec![CipherOptionProto {
+                suite: CipherSuiteProto::AesCtrNopadding as i32,
+                ..Default::default()
+            }],
+            ..Default::default()
+        };
+
+        debug!("Sending data transfer encryptor message: {:?}", message);
+
+        self.stream
+            .write_all(&message.encode_length_delimited_to_vec())
+            .await?;
+        self.stream.flush().await?;
+
+        let response = self.read_sasl_response().await?;
+        debug!("Data transfer encryptor response: {:?}", response);
+
+        let (_, finished) = session.step(response.payload.as_ref().map(|p| &p[..]))?;
+
+        assert!(finished);
+
+        Ok(self.stream.into_inner())
+    }
+
+    async fn read_sasl_response(&mut self) -> Result<DataTransferEncryptorMessageProto> {
+        self.stream.fill_buf().await?;
+
+        let buf = self.stream.fill_buf().await?;
+        if buf.is_empty() {
+            return Err(std::io::Error::from(std::io::ErrorKind::UnexpectedEof))?;
+        }
+        let msg_length = prost::decode_length_delimiter(buf)?;
+        let total_size = msg_length + prost::length_delimiter_len(msg_length);
+
+        let mut response_buf = BytesMut::zeroed(total_size);
+        self.stream.read_exact(&mut response_buf).await?;
+
+        Ok(DataTransferEncryptorMessageProto::decode_length_delimited(
+            response_buf.freeze(),
+        )?)
+    }
+}
diff --git a/crates/hdfs-native/src/security/user.rs b/crates/hdfs-native/src/security/user.rs
