1- use std:: path:: { Path , PathBuf } ;
2-
3- use entity_tag:: EntityTag ;
1+ use aws_sdk_s3:: primitives:: ByteStream ;
42use rocket:: {
53 data:: Capped ,
64 fairing:: { self , Fairing , Info , Kind } ,
75 fs:: TempFile ,
86 tokio:: io:: AsyncReadExt ,
9- Build , Rocket , State ,
7+ Build , Rocket ,
108} ;
9+ use serde:: Deserialize ;
1110use sha2:: { Digest , Sha256 } ;
1211use std:: fmt:: Write ;
1312
14- use crate :: { cached_file :: CachedFile , error:: AppError } ;
13+ use crate :: error:: AppError ;
1514
1615pub struct FilesInitializer ;
1716
17+ #[ derive( Deserialize ) ]
18+ struct S3Config {
19+ url : String ,
20+ bucket : String ,
21+ #[ serde( default ) ]
22+ use_mock : bool ,
23+ }
24+
1825pub struct Files {
19- #[ cfg( not( test) ) ]
20- upload_dir : PathBuf ,
21- #[ cfg( test) ]
22- upload_dir : tempfile:: TempDir ,
26+ s3_client : aws_sdk_s3:: Client ,
27+ s3_config : S3Config ,
2328}
2429
2530#[ rocket:: async_trait]
@@ -31,109 +36,88 @@ impl Fairing for FilesInitializer {
3136 }
3237 }
3338
34- #[ cfg( not( test) ) ]
3539 async fn on_ignite ( & self , rocket : Rocket < Build > ) -> fairing:: Result {
36- let upload_dir : PathBuf = match rocket. figment ( ) . extract_inner ( "upload_dir" ) {
40+ let s3_config : S3Config = match rocket. figment ( ) . focus ( "s3" ) . extract ( ) {
3741 Ok ( dir) => dir,
3842 Err ( e) => {
39- error ! ( "upload directory not specified : {}" , e) ;
43+ error ! ( "s3 configuration incomplete : {}" , e) ;
4044 return Err ( rocket) ;
4145 }
4246 } ;
4347
44- let upload_dir = match upload_dir. canonicalize ( ) {
45- Ok ( dir) => dir,
46- Err ( e) => {
47- error ! ( "invalid upload directory: {}" , e) ;
48- return Err ( rocket) ;
49- }
50- } ;
51-
52- if !upload_dir. is_dir ( ) {
53- error ! ( "given upload directory is not a directory" ) ;
54- return Err ( rocket) ;
55- }
56-
57- info ! ( "upload directory is set to {:?}" , upload_dir) ;
48+ let mut builder = aws_config:: load_defaults ( aws_config:: BehaviorVersion :: latest ( ) )
49+ . await
50+ . into_builder ( )
51+ . region ( aws_config:: Region :: new ( "eu-west-1" ) ) ;
52+ // For some stupid reason it isn't possible to conditionally set the endpoint url without a
53+ // mutable reference...
54+ builder. set_endpoint_url ( s3_config. use_mock . then ( || s3_config. url . clone ( ) ) ) ;
5855
59- let files = Files { upload_dir } ;
60- Ok ( rocket. manage ( files) . mount ( "/uploads" , routes ! [ uploads] ) )
61- }
56+ let config = builder. build ( ) ;
57+ let config = aws_sdk_s3:: config:: Builder :: from ( & config)
58+ . force_path_style ( s3_config. use_mock )
59+ . build ( ) ;
6260
63- /// Handle using a temp dir for tests
64- #[ cfg( test) ]
65- async fn on_ignite ( & self , rocket : Rocket < Build > ) -> fairing:: Result {
66- let dir = tempfile:: tempdir ( ) . unwrap ( ) ;
61+ let s3_client = aws_sdk_s3:: Client :: from_conf ( config) ;
6762
68- let files = Files { upload_dir : dir } ;
69- Ok ( rocket. manage ( files) . mount ( "/uploads" , routes ! [ uploads] ) )
63+ let files = Files {
64+ s3_client,
65+ s3_config,
66+ } ;
67+ Ok ( rocket. manage ( files) )
7068 }
7169}
7270
7371impl Files {
74- pub async fn upload_file ( & self , file : & mut Capped < TempFile < ' _ > > ) -> Result < PathBuf , AppError > {
72+ pub async fn upload_file ( & self , file : & mut Capped < TempFile < ' _ > > ) -> Result < String , AppError > {
7573 if !file. is_complete ( ) {
7674 return Err ( AppError :: FileTooBig ( file. len ( ) ) ) ;
7775 }
7876
79- let mut buf = file. open ( ) . await ?;
80- let mut hasher = Sha256 :: new ( ) ;
81- let mut buffer = [ 0u8 ; 1024 ] ;
77+ let hash = {
78+ let mut stream = file. open ( ) . await ?;
79+ let mut hasher = Sha256 :: new ( ) ;
80+ let mut buffer = [ 0u8 ; 1024 ] ;
8281
83- loop {
84- let count = buf . read ( & mut buffer ) . await ? ;
85- if count == 0 {
86- break ;
82+ while let count = stream . read ( & mut buffer ) . await ?
83+ && count != 0
84+ {
85+ hasher . update ( & buffer [ ..count ] ) ;
8786 }
88- hasher. update ( & buffer[ ..count] ) ;
89- }
90- std:: mem:: drop ( buf) ;
91-
92- let hash = hasher. finalize ( ) ;
93- let hash: String = hash. iter ( ) . fold ( "" . to_string ( ) , |mut s, b| {
94- write ! ( s, "{:02x}" , b) . unwrap ( ) ;
95- s
96- } ) ;
97- let extension = file. content_type ( ) . and_then ( |ct| ct. extension ( ) ) ;
98- let file_path = PathBuf :: from ( & hash[ ..2 ] ) ;
99- std:: fs:: create_dir_all ( self . get_path ( ) . join ( & file_path) ) ?;
100- let file_name = match extension {
101- Some ( ext) => format ! ( "{}.{}" , hash, ext) ,
102- None => hash,
103- } ;
104- let file_path = file_path. join ( file_name) ;
105-
106- let dest_path = self . get_path ( ) . join ( & file_path) ;
107- if !dest_path. try_exists ( ) ? {
108- // copying instead of `persist_to` because of cross-device limitations
109- file. move_copy_to ( dest_path) . await ?;
110- } else if !dest_path. is_file ( ) {
111- return Err ( AppError :: InternalError (
112- "destination path already exists, but it is not a file" ,
113- ) ) ;
114- }
11587
116- Ok ( file_path)
117- }
88+ hasher. finalize ( ) . iter ( ) . fold ( "" . to_string ( ) , |mut s, b| {
89+ write ! ( s, "{:02x}" , b) . unwrap ( ) ;
90+ s
91+ } )
92+ } ;
11893
119- #[ cfg( not( test) ) ]
120- fn get_path ( & self ) -> & Path {
121- self . upload_dir . as_path ( )
94+ let key = hash
95+ + "."
96+ + file
97+ . content_type ( )
98+ . and_then ( |ct| ct. extension ( ) )
99+ . map ( |string| string. as_str ( ) )
100+ . unwrap_or ( "" ) ;
101+
102+ let mut content = Vec :: new ( ) ;
103+ file. open ( ) . await ?. read_to_end ( & mut content) . await ?;
104+
105+ self . s3_client
106+ . put_object ( )
107+ . bucket ( & self . s3_config . bucket )
108+ . key ( & key)
109+ . body ( ByteStream :: from ( content) )
110+ . set_content_type (
111+ file. content_type ( )
112+ . map ( |content_type| content_type. to_string ( ) ) ,
113+ )
114+ . send ( )
115+ . await ?;
116+
117+ Ok ( key)
122118 }
123119
124- #[ cfg( test) ]
125- fn get_path ( & self ) -> & Path {
126- self . upload_dir . path ( )
120+ pub fn file_url ( & self , key : & str ) -> String {
121+ format ! ( "{}/{}/{}" , self . s3_config. url, self . s3_config. bucket, key)
127122 }
128123}
129-
130- #[ get( "/<path..>" ) ]
131- async fn uploads ( path : PathBuf , files_config : & State < Files > ) -> Option < CachedFile < ' static > > {
132- let path = files_config. get_path ( ) . join ( & path) ;
133- // Setting etag to filename as it is set to a hash of the contents on file upload.
134- let etag = EntityTag :: with_string ( false , path. file_stem ( ) ?. to_str ( ) ?. to_owned ( ) ) . ok ( ) ?;
135-
136- CachedFile :: open ( path, etag, chrono:: Duration :: weeks ( 1 ) )
137- . await
138- . ok ( )
139- }
0 commit comments