@@ -3,14 +3,15 @@ use std::io;
33use std:: io:: ErrorKind ;
44#[ cfg( test) ]
55use std:: panic:: RefUnwindSafe ;
6+ use std:: sync:: Arc ;
67use std:: time:: Duration ;
78
89use crate :: io:: utils:: check_namespace_key_validity;
910use lightning:: util:: persist:: KVStore ;
1011use prost:: Message ;
1112use rand:: RngCore ;
1213use tokio:: runtime:: Runtime ;
13- use vss_client:: client:: VssClient ;
14+ use vss_client:: client:: { AuthMethod , VssClient } ;
1415use vss_client:: error:: VssError ;
1516use vss_client:: types:: {
1617 DeleteObjectRequest , GetObjectRequest , KeyValue , ListKeyVersionsRequest , PutObjectRequest ,
@@ -35,10 +36,14 @@ pub struct VssStore {
3536 store_id : String ,
3637 runtime : Runtime ,
3738 storable_builder : StorableBuilder < RandEntropySource > ,
39+ auth_custom : Arc < dyn AuthMethod > ,
3840}
3941
4042impl VssStore {
41- pub ( crate ) fn new ( base_url : String , store_id : String , data_encryption_key : [ u8 ; 32 ] ) -> Self {
43+ pub ( crate ) fn new (
44+ base_url : String , store_id : String , data_encryption_key : [ u8 ; 32 ] ,
45+ auth_custom : impl AuthMethod + ' static ,
46+ ) -> Self {
4247 let runtime = tokio:: runtime:: Builder :: new_multi_thread ( ) . enable_all ( ) . build ( ) . unwrap ( ) ;
4348 let storable_builder = StorableBuilder :: new ( data_encryption_key, RandEntropySource ) ;
4449 let retry_policy = ExponentialBackoffRetryPolicy :: new ( Duration :: from_millis ( 100 ) )
@@ -55,7 +60,7 @@ impl VssStore {
5560 } ) as _ ) ;
5661
5762 let client = VssClient :: new ( & base_url, retry_policy) ;
58- Self { client, store_id, runtime, storable_builder }
63+ Self { client, store_id, runtime, storable_builder, auth_custom : Arc :: new ( auth_custom ) }
5964 }
6065
6166 fn build_key (
@@ -118,18 +123,19 @@ impl KVStore for VssStore {
118123 key : self . build_key ( primary_namespace, secondary_namespace, key) ?,
119124 } ;
120125
121- let resp =
122- tokio:: task:: block_in_place ( || self . runtime . block_on ( self . client . get_object ( & request) ) )
123- . map_err ( |e| {
124- let msg = format ! (
125- "Failed to read from key {}/{}/{}: {}" ,
126- primary_namespace, secondary_namespace, key, e
127- ) ;
128- match e {
129- VssError :: NoSuchKeyError ( ..) => Error :: new ( ErrorKind :: NotFound , msg) ,
130- _ => Error :: new ( ErrorKind :: Other , msg) ,
131- }
132- } ) ?;
126+ let resp = tokio:: task:: block_in_place ( || {
127+ self . runtime . block_on ( self . client . get_object ( & request, self . auth_custom . clone ( ) ) )
128+ } )
129+ . map_err ( |e| {
130+ let msg = format ! (
131+ "Failed to read from key {}/{}/{}: {}" ,
132+ primary_namespace, secondary_namespace, key, e
133+ ) ;
134+ match e {
135+ VssError :: NoSuchKeyError ( ..) => Error :: new ( ErrorKind :: NotFound , msg) ,
136+ _ => Error :: new ( ErrorKind :: Other , msg) ,
137+ }
138+ } ) ?;
133139 // unwrap safety: resp.value must be always present for a non-erroneous VSS response, otherwise
134140 // it is an API-violation which is converted to [`VssError::InternalServerError`] in [`VssClient`]
135141 let storable = Storable :: decode ( & resp. value . unwrap ( ) . value [ ..] ) ?;
0 commit comments