@@ -4,17 +4,19 @@ use crate::{Cipher, Error, Result};
44use cipher:: KeyIvInit ;
55use core:: fmt:: { self , Debug } ;
66
7- #[ cfg( feature = "aes-ctr" ) ]
8- use crate :: { Ctr128BE , encryptor:: ctr_encrypt as ctr_decrypt} ;
97#[ cfg( any( feature = "aes-cbc" , feature = "aes-ctr" ) ) ]
108use aes:: { Aes128 , Aes192 , Aes256 } ;
119#[ cfg( any( feature = "aes-cbc" , feature = "tdes" ) ) ]
10+ use cipher:: SetIvState ;
11+ #[ cfg( any( feature = "aes-cbc" , feature = "tdes" ) ) ]
1212use cipher:: {
13- Block ,
13+ Block , IvState ,
1414 block:: { BlockCipherDecrypt , BlockModeDecrypt } ,
1515} ;
1616#[ cfg( feature = "tdes" ) ]
1717use des:: TdesEde3 ;
18+ #[ cfg( feature = "aes-ctr" ) ]
19+ use { crate :: encryptor:: ctr_encrypt as ctr_decrypt, cipher:: StreamCipherSeek , ctr:: Ctr128BE } ;
1820
1921/// Stateful decryptor object for unauthenticated SSH symmetric ciphers.
2022///
@@ -43,6 +45,16 @@ enum Inner {
4345 TDesCbc ( cbc:: Decryptor < TdesEde3 > ) ,
4446}
4547
48+ /// Current IV state or position within the cipher.
49+ enum State {
50+ #[ cfg( feature = "aes-cbc" ) ]
51+ AesCbc ( aes:: Block ) ,
52+ #[ cfg( feature = "aes-ctr" ) ]
53+ AesCtr ( u64 ) ,
54+ #[ cfg( feature = "tdes" ) ]
55+ TDesCbc ( Block < TdesEde3 > ) ,
56+ }
57+
4658impl Decryptor {
4759 /// Create a new decryptor object with the given [`Cipher`], `key`, and `iv` (i.e.
4860 /// initialization vector).
@@ -124,6 +136,81 @@ impl Decryptor {
124136
125137 Ok ( ( ) )
126138 }
139+
140+ /// Call the provided function with an ephemeral [`Decryptor`] state which will be reset upon
141+ /// completion, returning the result of the function.
142+ ///
143+ /// # Errors
144+ /// Returns errors propagated from `F`, or if an internal cryptographic error occurs.
145+ pub fn peek < T , F > ( & mut self , mut f : F ) -> Result < T >
146+ where
147+ F : FnMut ( & mut Self ) -> Result < T > ,
148+ {
149+ let state = self . state ( ) ;
150+ let ret = f ( self ) ;
151+ self . set_state ( state) ?;
152+ ret
153+ }
154+
155+ /// Get the current cipher state, i.e. IV or position within the stream cipher.
156+ fn state ( & self ) -> State {
157+ match & self . inner {
158+ #[ cfg( feature = "aes-cbc" ) ]
159+ Inner :: Aes128Cbc ( cipher) => State :: AesCbc ( cipher. iv_state ( ) ) ,
160+ #[ cfg( feature = "aes-cbc" ) ]
161+ Inner :: Aes192Cbc ( cipher) => State :: AesCbc ( cipher. iv_state ( ) ) ,
162+ #[ cfg( feature = "aes-cbc" ) ]
163+ Inner :: Aes256Cbc ( cipher) => State :: AesCbc ( cipher. iv_state ( ) ) ,
164+ #[ cfg( feature = "aes-ctr" ) ]
165+ Inner :: Aes128Ctr ( cipher) => State :: AesCtr ( cipher. current_pos ( ) ) ,
166+ #[ cfg( feature = "aes-ctr" ) ]
167+ Inner :: Aes192Ctr ( cipher) => State :: AesCtr ( cipher. current_pos ( ) ) ,
168+ #[ cfg( feature = "aes-ctr" ) ]
169+ Inner :: Aes256Ctr ( cipher) => State :: AesCtr ( cipher. current_pos ( ) ) ,
170+ #[ cfg( feature = "tdes" ) ]
171+ Inner :: TDesCbc ( cipher) => State :: TDesCbc ( cipher. iv_state ( ) ) ,
172+ }
173+ }
174+
175+ /// Set the current cipher state.
176+ fn set_state ( & mut self , state : State ) -> Result < ( ) > {
177+ match ( & mut self . inner , state) {
178+ #[ cfg( feature = "aes-cbc" ) ]
179+ ( Inner :: Aes128Cbc ( cipher) , State :: AesCbc ( iv) ) => {
180+ cipher. set_iv ( & iv) ;
181+ Ok ( ( ) )
182+ }
183+ #[ cfg( feature = "aes-cbc" ) ]
184+ ( Inner :: Aes192Cbc ( cipher) , State :: AesCbc ( iv) ) => {
185+ cipher. set_iv ( & iv) ;
186+ Ok ( ( ) )
187+ }
188+ #[ cfg( feature = "aes-cbc" ) ]
189+ ( Inner :: Aes256Cbc ( cipher) , State :: AesCbc ( iv) ) => {
190+ cipher. set_iv ( & iv) ;
191+ Ok ( ( ) )
192+ }
193+ #[ cfg( feature = "aes-ctr" ) ]
194+ ( Inner :: Aes128Ctr ( cipher) , State :: AesCtr ( pos) ) => {
195+ cipher. try_seek ( pos) . map_err ( |_| Error :: Crypto )
196+ }
197+ #[ cfg( feature = "aes-ctr" ) ]
198+ ( Inner :: Aes192Ctr ( cipher) , State :: AesCtr ( pos) ) => {
199+ cipher. try_seek ( pos) . map_err ( |_| Error :: Crypto )
200+ }
201+ #[ cfg( feature = "aes-ctr" ) ]
202+ ( Inner :: Aes256Ctr ( cipher) , State :: AesCtr ( pos) ) => {
203+ cipher. try_seek ( pos) . map_err ( |_| Error :: Crypto )
204+ }
205+ #[ cfg( feature = "tdes" ) ]
206+ ( Inner :: TDesCbc ( cipher) , State :: TDesCbc ( iv) ) => {
207+ cipher. set_iv ( & iv) ;
208+ Ok ( ( ) )
209+ }
210+ #[ allow( unreachable_patterns) ]
211+ _ => Err ( Error :: Crypto ) , // should be unreachable
212+ }
213+ }
127214}
128215
129216impl Debug for Decryptor {
0 commit comments