1- use std:: { collections:: BTreeMap , io:: IsTerminal } ;
1+ use std:: {
2+ collections:: { BTreeMap , HashMap } ,
3+ io:: IsTerminal ,
4+ } ;
25
36use secrecy:: { ExposeSecret , SecretString } ;
7+ use serde:: Serialize ;
48use snafu:: { ResultExt , Snafu } ;
59
610use crate :: {
7- cli:: { ImageCheckArguments , ImageListArguments , Pretty } ,
11+ cli:: { ImageCheckArguments , ImageListArguments , ImageSizeArguments , Pretty } ,
812 config:: Config ,
913 core:: bakefile:: { self , Targets , TargetsOptions } ,
10- models:: TagList ,
11- utils:: { format_image_index_manifest_tag, format_registry_token_env_var_name} ,
14+ models:: { Manifest , TagList } ,
15+ utils:: {
16+ format_image_index_manifest_tag, format_image_manifest_tag,
17+ format_registry_token_env_var_name,
18+ } ,
1219} ;
1320
1421#[ derive( Debug , Snafu ) ]
@@ -22,8 +29,8 @@ pub enum Error {
2229 #[ snafu( display( "failed to build request client" ) ) ]
2330 BuildClient { source : reqwest:: Error } ,
2431
25- #[ snafu( display( "failed to send request" ) ) ]
26- SendRequest { source : reqwest:: Error } ,
32+ #[ snafu( display( "failed to send request ({url}) " ) ) ]
33+ SendRequest { source : reqwest:: Error , url : String } ,
2734
2835 #[ snafu( display( "failed to deserialize response" ) ) ]
2936 DeserializeResponse { source : reqwest:: Error } ,
@@ -51,7 +58,7 @@ pub fn list_images(arguments: ImageListArguments) -> Result<(), Error> {
5158 . context ( BuildTargetsSnafu ) ?
5259 } ;
5360
54- let list = targets
61+ let list: BTreeMap < String , Vec < String > > = targets
5562 . into_iter ( )
5663 . map ( |( image_name, ( image_config, _) ) | {
5764 let versions: Vec < _ > = image_config
@@ -109,7 +116,7 @@ pub async fn check_images(arguments: ImageCheckArguments, config: Config) -> Res
109116 "https://{registry}/v2/{registry_namespace}/{image_name}/tags/list" ,
110117 registry_namespace = registry_options. namespace,
111118 ) ;
112- let request = client. get ( url) ;
119+ let request = client. get ( & url) ;
113120
114121 let request = match & registry_token {
115122 Some ( registry_token) => request. bearer_auth ( registry_token. expose_secret ( ) ) ,
@@ -119,7 +126,7 @@ pub async fn check_images(arguments: ImageCheckArguments, config: Config) -> Res
119126 let tag_list: TagList = request
120127 . send ( )
121128 . await
122- . context ( SendRequestSnafu ) ?
129+ . context ( SendRequestSnafu { url } ) ?
123130 . json ( )
124131 . await
125132 . context ( DeserializeResponseSnafu ) ?;
@@ -148,14 +155,108 @@ pub async fn check_images(arguments: ImageCheckArguments, config: Config) -> Res
148155 Ok ( ( ) )
149156}
150157
151- fn print_to_stdout ( list : BTreeMap < String , Vec < String > > , pretty : Pretty ) -> Result < ( ) , Error > {
158+ pub async fn calculate_size ( arguments : ImageSizeArguments , config : Config ) -> Result < ( ) , Error > {
159+ let targets = if arguments. image . is_empty ( ) {
160+ Targets :: all ( TargetsOptions {
161+ only_entry : true ,
162+ non_recursive : true ,
163+ } )
164+ . context ( BuildTargetsSnafu ) ?
165+ } else {
166+ Targets :: set (
167+ & arguments. image ,
168+ TargetsOptions {
169+ only_entry : true ,
170+ non_recursive : true ,
171+ } ,
172+ )
173+ . context ( BuildTargetsSnafu ) ?
174+ } ;
175+
176+ let mut registry_tokens = BTreeMap :: new ( ) ;
177+ let client = reqwest:: ClientBuilder :: new ( )
178+ . build ( )
179+ . context ( BuildClientSnafu ) ?;
180+
181+ #[ derive( Serialize ) ]
182+ struct SizeResult {
183+ images : HashMap < String , u64 > ,
184+ total : u64 ,
185+ }
186+
187+ let mut result = SizeResult {
188+ images : HashMap :: new ( ) ,
189+ total : 0 ,
190+ } ;
191+
192+ for ( image_name, ( image_config, _) ) in targets {
193+ for ( registry, registry_options) in image_config. metadata . registries {
194+ // Add tokens to a map so that we don't need construct the key and retrieve the value
195+ // over and over again.
196+ let registry_token = registry_tokens. entry ( registry. clone ( ) ) . or_insert_with ( || {
197+ let name = format_registry_token_env_var_name ( & registry) ;
198+ std:: env:: var ( name) . ok ( ) . map ( SecretString :: from)
199+ } ) ;
200+
201+ for ( image_version, _) in image_config. versions . iter ( ) {
202+ let image_index_manifest_tag = format_image_index_manifest_tag (
203+ image_version,
204+ & config. metadata . vendor_tag_prefix ,
205+ & arguments. image_version ,
206+ ) ;
207+
208+ let manifest_tag = format_image_manifest_tag (
209+ & image_index_manifest_tag,
210+ arguments. target_platform . architecture ( ) ,
211+ // Never strip the architecture, because we need to reference the exact manifest
212+ // to be able to calculate the sizes
213+ false ,
214+ ) ;
215+
216+ let url = format ! (
217+ "https://{registry}/v2/{registry_namespace}/{image_name}/manifests/{manifest_tag}" ,
218+ registry_namespace = registry_options. namespace,
219+ ) ;
220+
221+ let request = client. get ( & url) ;
222+
223+ let request = match & registry_token {
224+ Some ( registry_token) => request. bearer_auth ( registry_token. expose_secret ( ) ) ,
225+ None => request,
226+ } ;
227+
228+ let manifest: Manifest = request
229+ . send ( )
230+ . await
231+ . context ( SendRequestSnafu { url } ) ?
232+ . json ( )
233+ . await
234+ . context ( DeserializeResponseSnafu ) ?;
235+
236+ let layer_size = manifest. layers . iter ( ) . fold ( 0u64 , |acc, e| acc + e. size ) ;
237+
238+ let size = result. images . entry ( image_name. clone ( ) ) . or_default ( ) ;
239+ * size += layer_size;
240+
241+ result. total += layer_size;
242+ }
243+ }
244+ }
245+
246+ print_to_stdout ( result, Pretty :: Always )
247+ }
248+
249+ fn print_to_stdout < T > ( value : T , pretty : Pretty ) -> Result < ( ) , Error >
250+ where
251+ T : Serialize ,
252+ {
152253 let stdout = std:: io:: stdout ( ) ;
153254
154255 match pretty {
155256 Pretty :: Always | Pretty :: Auto if stdout. is_terminal ( ) => {
156- serde_json:: to_writer_pretty ( stdout, & list )
257+ serde_json:: to_writer_pretty ( stdout, & value )
157258 }
158- _ => serde_json:: to_writer ( stdout, & list ) ,
259+ _ => serde_json:: to_writer ( stdout, & value ) ,
159260 }
160261 . context ( SerializeListSnafu )
161262}
0 commit comments