Repository: Kimahriman/hdfs-native (language: python)

Issue:
Implement data transfer encryption
After #85, we can add encryption at the data transfer level which uses non-SASL mechanisms for stronger encryption algorithms


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 17d4a16..aa1de8b 100644
--- a/README.md
+++ b/README.md
@@ -26,7 +26,7 @@ Here is a list of currently supported and unsupported but possible future featur
 - [x] Token authentication (DIGEST-MD5 SASL support)
 - [x] NameNode SASL connection
 - [x] DataNode SASL connection
-- [ ] DataNode data transfer encryption
+- [x] DataNode data transfer encryption
 - [ ] Encryption at rest (KMS support)
 
 ## Supported HDFS Settings
diff --git a/crates/hdfs-native/minidfs/src/main/java/main/Main.java b/crates/hdfs-native/minidfs/src/main/java/main/Main.java
index 33c58cb..90823f7 100644
--- a/crates/hdfs-native/minidfs/src/main/java/main/Main.java
+++ b/crates/hdfs-native/minidfs/src/main/java/main/Main.java
@@ -56,10 +56,6 @@ public static void main(String args[]) throws Exception {
             if (flags.contains("privacy")) {
                 conf.set(HADOOP_RPC_PROTECTION, "privacy");
                 conf.set(DFS_DATA_TRANSFER_PROTECTION_KEY, "privacy");
-                if (flags.contains("data_transfer_encryption")) {
-                    // Force encryption for all connections
-                    conf.set(DFSConfigKeys.DFS_ENCRYPT_DATA_TRANSFER_KEY, "true");
-                }
                 if (flags.contains("aes")) {
                     conf.set(DFS_ENCRYPT_DATA_TRANSFER_CIPHER_SUITES_KEY, "AES/CTR/NoPadding");
                 }
@@ -70,6 +66,10 @@ public static void main(String args[]) throws Exception {
                 conf.set(HADOOP_RPC_PROTECTION, "authentication");
                 conf.set(DFS_DATA_TRANSFER_PROTECTION_KEY, "authentication");
             }
+            if (flags.contains("data_transfer_encryption")) {
+                // Force encryption for all connections, legacy method before SASL connections were a thing
+                conf.set(DFSConfigKeys.DFS_ENCRYPT_DATA_TRANSFER_KEY, "true");
+            }
             conf.set(DFSConfigKeys.DFS_NAMENODE_KEYTAB_FILE_KEY, "target/test/hdfs.keytab");
             conf.set(DFS_NAMENODE_KERBEROS_PRINCIPAL_KEY, "hdfs/localhost@" + kdc.getRealm());
             conf.set(DFSConfigKeys.DFS_DATANODE_KEYTAB_FILE_KEY, "target/test/hdfs.keytab");
diff --git a/crates/hdfs-native/src/client.rs b/crates/hdfs-native/src/client.rs
index 228b2fc..c78832e 100644
--- a/crates/hdfs-native/src/client.rs
+++ b/crates/hdfs-native/src/client.rs
@@ -286,7 +286,12 @@ impl Client {
                 }
 
                 if let Some(locations) = status.locations.take() {
-                    Ok(FileReader::new(status, locations, ec_schema))
+                    Ok(FileReader::new(
+                        Arc::clone(&link.protocol),
+                        status,
+                        locations,
+                        ec_schema,
+                    ))
                 } else {
                     Err(HdfsError::BlocksNotFound(path.to_string()))
                 }
@@ -305,14 +310,6 @@ impl Client {
         let write_options = write_options.as_ref();
 
         let (link, resolved_path) = self.mount_table.resolve(src);
-        let server_defaults = link.protocol.get_server_defaults().await?.server_defaults;
-
-        let block_size = write_options
-            .block_size
-            .unwrap_or(server_defaults.block_size);
-        let replication = write_options
-            .replication
-            .unwrap_or(server_defaults.replication);
 
         let create_response = link
             .protocol
@@ -321,8 +318,8 @@ impl Client {
                 write_options.permission,
                 write_options.overwrite,
                 write_options.create_parent,
-                replication,
-                block_size,
+                write_options.replication,
+                write_options.block_size,
             )
             .await?;
 
@@ -338,7 +335,6 @@ impl Client {
                     resolved_path,
                     status,
                     None,
-                    server_defaults,
                 ))
             }
             None => Err(HdfsError::FileNotFound(src.to_string())),
@@ -354,7 +350,6 @@ impl Client {
     /// coded, a new block will be created.
     pub async fn append(&self, src: &str) -> Result<FileWriter> {
         let (link, resolved_path) = self.mount_table.resolve(src);
-        let server_defaults = link.protocol.get_server_defaults().await?.server_defaults;
 
         // Assume the file is replicated and try to append to the current block. If the file is
         // erasure coded, then try again by appending to a new block.
@@ -380,7 +375,6 @@ impl Client {
                     resolved_path,
                     status,
                     append_response.block,
-                    server_defaults,
                 ))
             }
             None => Err(HdfsError::FileNotFound(src.to_string())),
diff --git a/crates/hdfs-native/src/file.rs b/crates/hdfs-native/src/file.rs
index e3ad9b1..772031d 100644
--- a/crates/hdfs-native/src/file.rs
+++ b/crates/hdfs-native/src/file.rs
@@ -17,6 +17,7 @@ const COMPLETE_RETRY_DELAY_MS: u64 = 500;
 const COMPLETE_RETRIES: u32 = 5;
 
 pub struct FileReader {
+    protocol: Arc<NamenodeProtocol>,
     status: hdfs::HdfsFileStatusProto,
     located_blocks: hdfs::LocatedBlocksProto,
     ec_schema: Option<EcSchema>,
@@ -25,11 +26,13 @@ pub struct FileReader {
 
 impl FileReader {
     pub(crate) fn new(
+        protocol: Arc<NamenodeProtocol>,
         status: hdfs::HdfsFileStatusProto,
         located_blocks: hdfs::LocatedBlocksProto,
         ec_schema: Option<EcSchema>,
     ) -> Self {
         Self {
+            protocol,
             status,
             located_blocks,
             ec_schema,
@@ -125,6 +128,7 @@ impl FileReader {
                     let block_start = offset - usize::min(offset, block_file_start);
                     let block_end = usize::min(offset + len, block_file_end) - block_file_start;
                     Some(get_block_stream(
+                        Arc::clone(&self.protocol),
                         block.clone(),
                         block_start,
                         block_end - block_start,
@@ -145,7 +149,6 @@ pub struct FileWriter {
     src: String,
     protocol: Arc<NamenodeProtocol>,
     status: hdfs::HdfsFileStatusProto,
-    server_defaults: hdfs::FsServerDefaultsProto,
     block_writer: Option<BlockWriter>,
     last_block: Option<hdfs::LocatedBlockProto>,
     closed: bool,
@@ -159,14 +162,12 @@ impl FileWriter {
         status: hdfs::HdfsFileStatusProto,
         // Some for append, None for create
         last_block: Option<hdfs::LocatedBlockProto>,
-        server_defaults: hdfs::FsServerDefaultsProto,
     ) -> Self {
         protocol.add_file_lease(status.file_id(), status.namespace.clone());
         Self {
             protocol,
             src,
             status,
-            server_defaults,
             block_writer: None,
             last_block,
             closed: false,
@@ -206,9 +207,10 @@ impl FileWriter {
         };
 
         let block_writer = BlockWriter::new(
+            Arc::clone(&self.protocol),
             new_block,
             self.status.blocksize() as usize,
-            self.server_defaults.clone(),
+            self.protocol.get_cached_server_defaults().await?,
             self.status
                 .ec_policy
                 .as_ref()
diff --git a/crates/hdfs-native/src/hdfs/block_reader.rs b/crates/hdfs-native/src/hdfs/block_reader.rs
index 041b47b..98cc2ee 100644
--- a/crates/hdfs-native/src/hdfs/block_reader.rs
+++ b/crates/hdfs-native/src/hdfs/block_reader.rs
@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, sync::Arc};
 
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 use futures::{
@@ -18,18 +18,21 @@ use crate::{
     HdfsError, Result,
 };
 
+use super::protocol::NamenodeProtocol;
+
 pub(crate) fn get_block_stream(
+    protocol: Arc<NamenodeProtocol>,
     block: hdfs::LocatedBlockProto,
     offset: usize,
     len: usize,
     ec_schema: Option<EcSchema>,
 ) -> BoxStream<'static, Result<Bytes>> {
     if let Some(ec_schema) = ec_schema {
-        StripedBlockStream::new(block, offset, len, ec_schema)
+        StripedBlockStream::new(protocol, block, offset, len, ec_schema)
             .into_stream()
             .boxed()
     } else {
-        ReplicatedBlockStream::new(block, offset, len)
+        ReplicatedBlockStream::new(protocol, block, offset, len)
             .into_stream()
             .boxed()
     }
@@ -37,6 +40,7 @@ pub(crate) fn get_block_stream(
 
 /// Connects to a DataNode to do a read, attempting to used cached connections.
 async fn connect_and_send(
+    protocol: &Arc<NamenodeProtocol>,
     datanode_id: &hdfs::DatanodeIdProto,
     block: &hdfs::ExtendedBlockProto,
     token: common::TokenProto,
@@ -68,7 +72,12 @@ async fn connect_and_send(
         }
         remaining_attempts -= 1;
     }
-    let mut conn = DatanodeConnection::connect(datanode_id, &token).await?;
+    let mut conn = DatanodeConnection::connect(
+        datanode_id,
+        &token,
+        protocol.get_cached_data_encryption_key().await?,
+    )
+    .await?;
 
     let message = hdfs::OpReadBlockProto {
         header: conn.build_header(block, Some(token)),
@@ -85,6 +94,7 @@ async fn connect_and_send(
 }
 
 struct ReplicatedBlockStream {
+    protocol: Arc<NamenodeProtocol>,
     block: hdfs::LocatedBlockProto,
     offset: usize,
     len: usize,
@@ -95,8 +105,14 @@ struct ReplicatedBlockStream {
 }
 
 impl ReplicatedBlockStream {
-    fn new(block: hdfs::LocatedBlockProto, offset: usize, len: usize) -> Self {
+    fn new(
+        protocol: Arc<NamenodeProtocol>,
+        block: hdfs::LocatedBlockProto,
+        offset: usize,
+        len: usize,
+    ) -> Self {
         Self {
+            protocol,
             block,
             offset,
             len,
@@ -119,6 +135,7 @@ impl ReplicatedBlockStream {
         let datanode = &self.block.locs[self.current_replica].id;
 
         let (connection, response) = connect_and_send(
+            &self.protocol,
             datanode,
             &self.block.b,
             self.block.block_token.clone(),
@@ -182,6 +199,7 @@ impl ReplicatedBlockStream {
 }
 
 struct StripedBlockStream {
+    protocol: Arc<NamenodeProtocol>,
     block: hdfs::LocatedBlockProto,
     offset: usize,
     len: usize,
@@ -189,8 +207,15 @@ struct StripedBlockStream {
 }
 
 impl StripedBlockStream {
-    fn new(block: hdfs::LocatedBlockProto, offset: usize, len: usize, ec_schema: EcSchema) -> Self {
+    fn new(
+        protocol: Arc<NamenodeProtocol>,
+        block: hdfs::LocatedBlockProto,
+        offset: usize,
+        len: usize,
+        ec_schema: EcSchema,
+    ) -> Self {
         Self {
+            protocol,
             block,
             offset,
             len,
@@ -399,8 +424,15 @@ impl StripedBlockStream {
             return Ok(());
         }
 
-        let (mut connection, response) =
-            connect_and_send(datanode, block, token.clone(), offset as u64, len as u64).await?;
+        let (mut connection, response) = connect_and_send(
+            &self.protocol,
+            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
index cabf6ca..85ce3ea 100644
--- a/crates/hdfs-native/src/hdfs/block_writer.rs
+++ b/crates/hdfs-native/src/hdfs/block_writer.rs
@@ -1,4 +1,4 @@
-use std::time::Duration;
+use std::{sync::Arc, time::Duration};
 
 use bytes::{BufMut, Bytes, BytesMut};
 use futures::future::join_all;
@@ -12,6 +12,8 @@ use crate::{
     HdfsError, Result,
 };
 
+use super::protocol::NamenodeProtocol;
+
 const HEART_BEAT_SEQNO: i64 = -1;
 const UNKNOWN_SEQNO: i64 = -2;
 
@@ -26,6 +28,7 @@ pub(crate) enum BlockWriter {
 
 impl BlockWriter {
     pub(crate) async fn new(
+        protocol: Arc<NamenodeProtocol>,
         block: hdfs::LocatedBlockProto,
         block_size: usize,
         server_defaults: hdfs::FsServerDefaultsProto,
@@ -33,13 +36,16 @@ impl BlockWriter {
     ) -> Result<Self> {
         let block_writer = if let Some(ec_schema) = ec_schema {
             Self::Striped(StripedBlockWriter::new(
+                protocol,
                 block,
                 ec_schema,
                 block_size,
                 server_defaults,
             ))
         } else {
-            Self::Replicated(ReplicatedBlockWriter::new(block, block_size, server_defaults).await?)
+            Self::Replicated(
+                ReplicatedBlockWriter::new(&protocol, block, block_size, server_defaults).await?,
+            )
         };
         Ok(block_writer)
     }
@@ -95,12 +101,18 @@ pub(crate) struct ReplicatedBlockWriter {
 
 impl ReplicatedBlockWriter {
     async fn new(
+        protocol: &Arc<NamenodeProtocol>,
         block: hdfs::LocatedBlockProto,
         block_size: usize,
         server_defaults: hdfs::FsServerDefaultsProto,
     ) -> Result<Self> {
         let datanode = &block.locs[0].id;
-        let mut connection = DatanodeConnection::connect(datanode, &block.block_token).await?;
+        let mut connection = DatanodeConnection::connect(
+            datanode,
+            &block.block_token,
+            protocol.get_cached_data_encryption_key().await?,
+        )
+        .await?;
 
         let checksum = hdfs::ChecksumProto {
             r#type: hdfs::ChecksumTypeProto::ChecksumCrc32c as i32,
@@ -462,6 +474,7 @@ impl CellBuffer {
 
 // Writer for erasure coded blocks.
 pub(crate) struct StripedBlockWriter {
+    protocol: Arc<NamenodeProtocol>,
     block: hdfs::LocatedBlockProto,
     server_defaults: hdfs::FsServerDefaultsProto,
     block_size: usize,
@@ -473,6 +486,7 @@ pub(crate) struct StripedBlockWriter {
 
 impl StripedBlockWriter {
     fn new(
+        protocol: Arc<NamenodeProtocol>,
         block: hdfs::LocatedBlockProto,
         ec_schema: &EcSchema,
         block_size: usize,
@@ -481,6 +495,7 @@ impl StripedBlockWriter {
         let block_writers = (0..block.block_indices().len()).map(|_| None).collect();
 
         Self {
+            protocol,
             block,
             block_size,
             server_defaults,
@@ -519,6 +534,7 @@ impl StripedBlockWriter {
 
                 *writer = Some(
                     ReplicatedBlockWriter::new(
+                        &self.protocol,
                         cloned,
                         self.block_size,
                         self.server_defaults.clone(),
diff --git a/crates/hdfs-native/src/hdfs/connection.rs b/crates/hdfs-native/src/hdfs/connection.rs
index 50c9849..da28105 100644
--- a/crates/hdfs-native/src/hdfs/connection.rs
+++ b/crates/hdfs-native/src/hdfs/connection.rs
@@ -22,7 +22,7 @@ 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::hdfs::{DataEncryptionKeyProto, DatanodeIdProto};
 use crate::proto::{common, hdfs};
 use crate::security::sasl::{SaslDatanodeConnection, SaslDatanodeReader, SaslDatanodeWriter};
 use crate::security::sasl::{SaslReader, SaslRpcClient, SaslWriter};
@@ -390,7 +390,6 @@ const CHECKSUM_BYTES: usize = 4;
 
 pub(crate) struct Packet {
     pub header: hdfs::PacketHeaderProto,
-    #[allow(dead_code)]
     checksum: BytesMut,
     data: BytesMut,
     bytes_per_checksum: usize,
@@ -519,13 +518,18 @@ pub(crate) struct DatanodeConnection {
 }
 
 impl DatanodeConnection {
-    #[allow(unused_variables)]
-    pub(crate) async fn connect(datanode_id: &DatanodeIdProto, token: &TokenProto) -> Result<Self> {
+    pub(crate) async fn connect(
+        datanode_id: &DatanodeIdProto,
+        token: &TokenProto,
+        encryption_key: Option<DataEncryptionKeyProto>,
+    ) -> Result<Self> {
         let url = format!("{}:{}", datanode_id.ip_addr, datanode_id.xfer_port);
         let stream = connect(&url).await?;
 
         let sasl_connection = SaslDatanodeConnection::create(stream);
-        let (reader, writer) = sasl_connection.negotiate(datanode_id, token).await?;
+        let (reader, writer) = sasl_connection
+            .negotiate(datanode_id, token, encryption_key.as_ref())
+            .await?;
 
         let conn = DatanodeConnection {
             client_name: Uuid::new_v4().to_string(),
diff --git a/crates/hdfs-native/src/hdfs/protocol.rs b/crates/hdfs-native/src/hdfs/protocol.rs
index f98dcde..5aff410 100644
--- a/crates/hdfs-native/src/hdfs/protocol.rs
+++ b/crates/hdfs-native/src/hdfs/protocol.rs
@@ -2,12 +2,15 @@ use std::collections::HashSet;
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, SystemTime};
 
+use chrono::Utc;
 use log::{debug, warn};
 use prost::Message;
 use tokio::task::JoinHandle;
 use uuid::Uuid;
 
-use crate::proto::hdfs;
+use crate::proto::hdfs::{
+    self, DataEncryptionKeyProto, FsServerDefaultsProto, GetDataEncryptionKeyResponseProto,
+};
 use crate::Result;
 
 use super::proxy::NameServiceProxy;
@@ -27,6 +30,9 @@ pub(crate) struct NamenodeProtocol {
     // Stores files currently opened for writing for lease renewal purposes
     open_files: Arc<Mutex<HashSet<LeasedFile>>>,
     lease_renewer: Mutex<Option<JoinHandle<()>>>,
+    // We need a sync mutex for these
+    server_defaults: tokio::sync::Mutex<Option<FsServerDefaultsProto>>,
+    encryption_key: tokio::sync::Mutex<Option<DataEncryptionKeyProto>>,
 }
 
 impl NamenodeProtocol {
@@ -37,6 +43,8 @@ impl NamenodeProtocol {
             client_name,
             open_files: Arc::new(Mutex::new(HashSet::new())),
             lease_renewer: Mutex::new(None),
+            server_defaults: tokio::sync::Mutex::new(None),
+            encryption_key: tokio::sync::Mutex::new(None),
         }
     }
 
@@ -119,14 +127,65 @@ impl NamenodeProtocol {
         Ok(decoded)
     }
 
+    pub(crate) async fn get_cached_server_defaults(&self) -> Result<FsServerDefaultsProto> {
+        let mut server_defaults = self.server_defaults.lock().await;
+        if let Some(defaults) = server_defaults.as_ref() {
+            Ok(defaults.clone())
+        } else {
+            let defaults = self.get_server_defaults().await?.server_defaults;
+            *server_defaults = Some(defaults.clone());
+            Ok(defaults)
+        }
+    }
+
+    pub(crate) async fn get_data_encryption_key(
+        &self,
+    ) -> Result<GetDataEncryptionKeyResponseProto> {
+        let message = hdfs::GetDataEncryptionKeyRequestProto::default();
+
+        let response = self
+            .proxy
+            .call(
+                "getDataEncryptionKey",
+                message.encode_length_delimited_to_vec(),
+            )
+            .await?;
+
+        let decoded = hdfs::GetDataEncryptionKeyResponseProto::decode_length_delimited(response)?;
+        debug!("get_data_encryption_key response: {:?}", &decoded);
+        Ok(decoded)
+    }
+
+    pub(crate) async fn get_cached_data_encryption_key(
+        &self,
+    ) -> Result<Option<DataEncryptionKeyProto>> {
+        let server_defaults = self.get_cached_server_defaults().await?;
+        if server_defaults.encrypt_data_transfer() {
+            let mut encryption_key = self.encryption_key.lock().await;
+            // Check if we have an encryption key and that it's not expired (with a 10 second buffer)
+            if encryption_key
+                .as_ref()
+                .is_some_and(|key| (key.expiry_date as i64) > Utc::now().timestamp_millis() + 10000)
+            {
+                Ok(encryption_key.clone())
+            } else {
+                let key = self.get_data_encryption_key().await?.data_encryption_key;
+                *encryption_key = key.clone();
+                Ok(key)
+            }
+        } else {
+            Ok(None)
+        }
+    }
+
     pub(crate) async fn create(
         &self,
         src: &str,
         permission: u32,
         overwrite: bool,
         create_parent: bool,
-        replication: u32,
-        block_size: u64,
+        replication: Option<u32>,
+        block_size: Option<u64>,
     ) -> Result<hdfs::CreateResponseProto> {
         let masked = hdfs::FsPermissionProto { perm: permission };
         let mut create_flag = hdfs::CreateFlagProto::Create as u32;
@@ -134,6 +193,10 @@ impl NamenodeProtocol {
             create_flag |= hdfs::CreateFlagProto::Overwrite as u32;
         }
 
+        let server_defaults = self.get_cached_server_defaults().await?;
+        let replication = replication.unwrap_or(server_defaults.replication);
+        let block_size = block_size.unwrap_or(server_defaults.block_size);
+
         let message = hdfs::CreateRequestProto {
             src: src.to_string(),
             masked,
diff --git a/crates/hdfs-native/src/minidfs.rs b/crates/hdfs-native/src/minidfs.rs
index 9edf6fe..05bf01b 100644
--- a/crates/hdfs-native/src/minidfs.rs
+++ b/crates/hdfs-native/src/minidfs.rs
@@ -14,6 +14,7 @@ pub enum DfsFeatures {
     Integrity,
     Privacy,
     AES,
+    DataTransferEncryption,
     HA,
     ViewFS,
     EC,
@@ -30,6 +31,7 @@ impl DfsFeatures {
             DfsFeatures::Security => "security",
             DfsFeatures::Integrity => "integrity",
             DfsFeatures::AES => "aes",
+            DfsFeatures::DataTransferEncryption => "data_transfer_encryption",
             DfsFeatures::Token => "token",
             DfsFeatures::RBF => "rbf",
         }
diff --git a/crates/hdfs-native/src/security/digest.rs b/crates/hdfs-native/src/security/digest.rs
index 1b61a0c..efe4a3f 100644
--- a/crates/hdfs-native/src/security/digest.rs
+++ b/crates/hdfs-native/src/security/digest.rs
@@ -12,7 +12,7 @@ use once_cell::sync::Lazy;
 use rand::Rng;
 use regex::Regex;
 
-use crate::{HdfsError, Result};
+use crate::{proto::hdfs::DataEncryptionKeyProto, HdfsError, Result};
 
 use super::{
     sasl::SaslSession,
@@ -268,7 +268,7 @@ pub(super) struct DigestSaslSession {
 }
 
 impl DigestSaslSession {
-    pub(super) fn new(service: String, hostname: String, token: &Token) -> Self {
+    pub(super) fn from_token(service: String, hostname: String, token: &Token) -> Self {
         Self {
             auth_id: general_purpose::STANDARD.encode(&token.identifier),
             password: general_purpose::STANDARD.encode(&token.password),
@@ -278,6 +278,25 @@ impl DigestSaslSession {
         }
     }
 
+    pub(super) fn from_encryption_key(
+        service: String,
+        hostname: String,
+        encryption_key: &DataEncryptionKeyProto,
+    ) -> Self {
+        Self {
+            auth_id: format!(
+                "{} {} {}",
+                encryption_key.key_id,
+                encryption_key.block_pool_id,
+                general_purpose::STANDARD.encode(&encryption_key.nonce)
+            ),
+            password: general_purpose::STANDARD.encode(&encryption_key.encryption_key),
+            service,
+            hostname,
+            state: DigestState::Pending,
+        }
+    }
+
     fn compute(&self, ctx: &DigestContext, initial: bool) -> String {
         let x = hex::encode(h(self.a1(ctx)));
         let y = format!(
@@ -605,7 +624,7 @@ mod test {
             service: "127.0.0.1:9000".to_string(),
         };
 
-        let session = DigestSaslSession::new("".to_string(), "default".to_string(), &token);
+        let session = DigestSaslSession::from_token("".to_string(), "default".to_string(), &token);
         let ctx = DigestContext {
             nonce: "A+DoU3+eajz9Ei11Ib0S9CUKLyLPh0qFJbwn1/OZ".to_string(),
             cnonce: "UGP4ejVb7M54KO4yqDwCsA==".to_string(),
@@ -619,7 +638,6 @@ mod test {
     }
 
     #[test]
-    #[ignore]
     fn test_digest_step_privacy() {
         let token = Token {
             alias: "127.0.0.1:9000".to_string(),
@@ -636,7 +654,7 @@ mod test {
             service: "127.0.0.1:9000".to_string(),
         };
 
-        let session = DigestSaslSession::new("".to_string(), "default".to_string(), &token);
+        let session = DigestSaslSession::from_token("".to_string(), "default".to_string(), &token);
         let ctx = DigestContext {
             nonce: "tm3kclm9F0JMECFNJi5xk/NaGgQ75ZOqb9/vCHt5".to_string(),
             cnonce: "kp49R9SjR4de6ynNgMwNwQ==".to_string(),
@@ -645,7 +663,7 @@ mod test {
         };
         assert_eq!(
             &session.compute(&ctx, true),
-            "621c0b2210a1d648983abd6bb43ff507"
+            "260f34cc6e484ede6277cb9c37672d07"
         );
     }
 }
diff --git a/crates/hdfs-native/src/security/sasl.rs b/crates/hdfs-native/src/security/sasl.rs
index 8e87876..cf3e1fb 100644
--- a/crates/hdfs-native/src/security/sasl.rs
+++ b/crates/hdfs-native/src/security/sasl.rs
@@ -12,7 +12,7 @@ use tokio::{
 };
 
 use super::user::BlockTokenIdentifier;
-use crate::proto::hdfs::{CipherOptionProto, CipherSuiteProto};
+use crate::proto::hdfs::{CipherOptionProto, CipherSuiteProto, DataEncryptionKeyProto};
 use crate::proto::{
     common::{
         rpc_response_header_proto::RpcStatusProto,
@@ -198,7 +198,7 @@ impl SaslRpcClient {
                     return Ok((auth.clone(), Some(Box::new(session))));
                 }
                 (Some(AuthMethod::Token), Some(token)) => {
-                    let session = DigestSaslSession::new(
+                    let session = DigestSaslSession::from_token(
                         auth.protocol().to_string(),
                         auth.server_id().to_string(),
                         token,
@@ -614,22 +614,38 @@ impl SaslDatanodeConnection {
         }
     }
 
+    /// There are a few different paths for negotiating a connection with a DataNode:
+    ///
+    /// 1. If `dfs.encrypt.data.transfer` is set on the NameNode, always encrypt the session
+    ///    and use an encryption key from the NameNode for the negotiation. This will happen
+    ///    if `encryption_key` is defined.
+    /// 2. If there is no block token or the DataNode transfer port is privileged (<= 1024), we
+    ///    skip the SASL handshake and assume it is trusted.
+    /// 3. Otherwise, we do a SAL handshake using the provided block token.
+    ///
+    /// For cases 1 and 3, we optionally negotiate a cipher to use for encryption instead of
+    /// SASL protection mechanisms.
     pub(crate) async fn negotiate(
         mut self,
         datanode_id: &DatanodeIdProto,
         token: &TokenProto,
+        encryption_key: Option<&DataEncryptionKeyProto>,
     ) -> Result<(SaslDatanodeReader, SaslDatanodeWriter)> {
-        // If there's no token identifier or it's a privileged port, don't do SASL negotation
-        if token.identifier.is_empty() || datanode_id.xfer_port <= 1024 {
+        let mut session = if let Some(key) = encryption_key {
+            DigestSaslSession::from_encryption_key("hdfs".to_string(), "0".to_string(), key)
+        } else if token.identifier.is_empty() || datanode_id.xfer_port <= 1024 {
             return self.split(None, None);
-        }
+        } else {
+            DigestSaslSession::from_token(
+                "hdfs".to_string(),
+                "0".to_string(),
+                &token.clone().into(),
+            )
+        };
 
         self.stream.write_i32(SASL_TRANSFER_MAGIC_NUMBER).await?;
         self.stream.flush().await?;
 
-        let mut session =
-            DigestSaslSession::new("hdfs".to_string(), "0".to_string(), &token.clone().into());
-
         let token_identifier = BlockTokenIdentifier::from_identifier(&token.identifier)?;
 
         let handshake_secret = if !token_identifier.handshake_secret.is_empty() {


Test patch:
diff --git a/crates/hdfs-native/tests/test_integration.rs b/crates/hdfs-native/tests/test_integration.rs
index 6778093..2021b14 100644
--- a/crates/hdfs-native/tests/test_integration.rs
+++ b/crates/hdfs-native/tests/test_integration.rs
@@ -91,6 +91,43 @@ mod test {
         .unwrap();
     }
 
+    #[tokio::test]
+    #[serial]
+    async fn test_forced_data_transfer_encryption() {
+        // DataTransferEncryption enabled but privacy isn't, still force encryption
+        test_with_features(&HashSet::from([
+            DfsFeatures::Security,
+            DfsFeatures::DataTransferEncryption,
+        ]))
+        .await
+        .unwrap();
+    }
+
+    #[tokio::test]
+    #[serial]
+    async fn test_data_transfer_encryption() {
+        test_with_features(&HashSet::from([
+            DfsFeatures::Security,
+            DfsFeatures::Privacy,
+            DfsFeatures::DataTransferEncryption,
+        ]))
+        .await
+        .unwrap();
+    }
+
+    #[tokio::test]
+    #[serial]
+    async fn test_data_transfer_encryption_aes() {
+        test_with_features(&HashSet::from([
+            DfsFeatures::Security,
+            DfsFeatures::Privacy,
+            DfsFeatures::DataTransferEncryption,
+            DfsFeatures::AES,
+        ]))
+        .await
+        .unwrap();
+    }
+
     #[tokio::test]
     #[serial]
     async fn test_basic_ha() {
diff --git a/crates/hdfs-native/tests/test_read.rs b/crates/hdfs-native/tests/test_read.rs
index fad48da..6469357 100644
--- a/crates/hdfs-native/tests/test_read.rs
+++ b/crates/hdfs-native/tests/test_read.rs
@@ -11,10 +11,43 @@ mod test {
 
     #[tokio::test]
     #[serial]
-    async fn test_read() -> Result<()> {
+    async fn test_read_simple() {
+        test_read(&HashSet::from([DfsFeatures::HA])).await.unwrap();
+    }
+
+    // These tests take a long time, so don't run by default
+    #[tokio::test]
+    #[serial]
+    #[ignore]
+    async fn test_read_sasl_encryption() {
+        test_read(&HashSet::from([
+            DfsFeatures::HA,
+            DfsFeatures::Security,
+            DfsFeatures::Privacy,
+        ]))
+        .await
+        .unwrap();
+    }
+
+    // These tests take a long time, so don't run by default
+    #[tokio::test]
+    #[serial]
+    #[ignore]
+    async fn test_read_cipher_encryption() {
+        test_read(&HashSet::from([
+            DfsFeatures::HA,
+            DfsFeatures::Security,
+            DfsFeatures::Privacy,
+            DfsFeatures::AES,
+        ]))
+        .await
+        .unwrap();
+    }
+
+    async fn test_read(features: &HashSet<DfsFeatures>) -> Result<()> {
         let _ = env_logger::builder().is_test(true).try_init();
 
-        let _dfs = setup(&HashSet::from([DfsFeatures::HA]));
+        let _dfs = setup(features);
         let client = Client::default();
 
         // Read the whole file
diff --git a/crates/hdfs-native/tests/test_write.rs b/crates/hdfs-native/tests/test_write.rs
index 93f1267..3ee43a6 100644
--- a/crates/hdfs-native/tests/test_write.rs
+++ b/crates/hdfs-native/tests/test_write.rs
@@ -11,14 +11,48 @@ mod test {
 
     #[tokio::test]
     #[serial]
-    async fn test_write() {
+    async fn test_write_simple() {
+        test_write(&HashSet::from([DfsFeatures::HA])).await.unwrap();
+    }
+
+    // These tests take a long time, so don't run by default
+    #[tokio::test]
+    #[serial]
+    #[ignore]
+    async fn test_write_sasl_encryption() {
+        test_write(&HashSet::from([
+            DfsFeatures::HA,
+            DfsFeatures::Security,
+            DfsFeatures::Privacy,
+        ]))
+        .await
+        .unwrap();
+    }
+
+    // These tests take a long time, so don't run by default
+    #[tokio::test]
+    #[serial]
+    #[ignore]
+    async fn test_write_cipher_encryption() {
+        test_write(&HashSet::from([
+            DfsFeatures::HA,
+            DfsFeatures::Security,
+            DfsFeatures::Privacy,
+            DfsFeatures::AES,
+        ]))
+        .await
+        .unwrap();
+    }
+
+    async fn test_write(features: &HashSet<DfsFeatures>) -> Result<()> {
         let _ = env_logger::builder().is_test(true).try_init();
 
-        let _dfs = setup(&HashSet::from([DfsFeatures::HA]));
+        let _dfs = setup(features);
         let client = Client::default();
 
-        test_create(&client).await.unwrap();
-        test_append(&client).await.unwrap();
+        test_create(&client).await?;
+        test_append(&client).await?;
+        Ok(())
     }
 
     async fn test_create(client: &Client) -> Result<()> {