@@ -20,7 +20,7 @@ use crate::ln::msgs::DecodeError;
20
20
use crate :: ln:: onion_utils;
21
21
use crate :: onion_message:: messenger:: Destination ;
22
22
use crate :: crypto:: streams:: ChaChaPolyWriteAdapter ;
23
- use crate :: util:: ser:: { Readable , Writeable } ;
23
+ use crate :: util:: ser:: { Readable , Writeable , Writer } ;
24
24
25
25
use crate :: io;
26
26
@@ -135,17 +135,45 @@ fn encrypt_payload<P: Writeable>(payload: P, encrypted_tlvs_rho: [u8; 32]) -> Ve
135
135
write_adapter. encode ( )
136
136
}
137
137
138
- /// Blinded path encrypted payloads may be padded to ensure they are equal length.
139
- ///
140
- /// Reads padding to the end, ignoring what's read.
141
- pub ( crate ) struct Padding { }
138
+ /// Represents optional padding for encrypted payloads.
139
+ /// Padding is used to ensure payloads have a consistent length.
140
+ pub ( crate ) struct Padding {
141
+ length : usize ,
142
+ }
143
+
144
+ impl Padding {
145
+ /// Creates a new [`Padding`] instance with a specified size.
146
+ /// Use this method when defining the padding size before writing
147
+ /// an encrypted payload.
148
+ pub fn new ( length : usize ) -> Self {
149
+ Self { length }
150
+ }
151
+ }
152
+
142
153
impl Readable for Padding {
154
+ /// Reads and discards the padding data.
143
155
#[ inline]
144
156
fn read < R : io:: Read > ( reader : & mut R ) -> Result < Self , DecodeError > {
145
157
loop {
146
158
let mut buf = [ 0 ; 8192 ] ;
147
159
if reader. read ( & mut buf[ ..] ) ? == 0 { break ; }
148
160
}
149
- Ok ( Self { } )
161
+ Ok ( Self :: new ( 0 ) )
162
+ }
163
+ }
164
+
165
+ impl Writeable for Padding {
166
+ fn write < W : Writer > ( & self , writer : & mut W ) -> Result < ( ) , io:: Error > {
167
+ const BUFFER_SIZE : usize = 1024 ;
168
+ let buffer = [ 0u8 ; BUFFER_SIZE ] ;
169
+
170
+ let mut remaining = self . length ;
171
+ loop {
172
+ let to_write = core:: cmp:: min ( remaining, BUFFER_SIZE ) ;
173
+ writer. write_all ( & buffer[ ..to_write] ) ?;
174
+ remaining -= to_write;
175
+ if remaining == 0 { break ; }
176
+ }
177
+ Ok ( ( ) )
150
178
}
151
179
}
0 commit comments