Skip to main content

embed_impl/
embed_impl.rs

1#![recursion_limit = "1024"]
2#![forbid(unsafe_code)]
3#[macro_use]
4extern crate quote;
5extern crate proc_macro;
6
7use std::{
8    collections::BTreeMap,
9    env,
10    iter::FromIterator,
11    path::{Path, PathBuf},
12};
13
14use embed_utils::PathMatcher;
15use proc_macro::TokenStream;
16use proc_macro2::TokenStream as TokenStream2;
17use syn::{Data, DeriveInput, Expr, ExprLit, Fields, Lit, Meta, MetaNameValue, parse_macro_input};
18
19fn embedded(
20    ident: &syn::Ident,
21    relative_folder_path: Option<&str>,
22    absolute_folder_path: String,
23    prefix: Option<&str>,
24    includes: &[String],
25    excludes: &[String],
26    metadata_only: bool,
27    crate_path: &syn::Path,
28) -> syn::Result<TokenStream2> {
29    extern crate embed_utils;
30
31    let mut match_values = BTreeMap::new();
32    let mut list_values = Vec::<String>::new();
33
34    let includes: Vec<&str> = includes.iter().map(AsRef::as_ref).collect();
35    let excludes: Vec<&str> = excludes.iter().map(AsRef::as_ref).collect();
36    let matcher = PathMatcher::new(&includes, &excludes);
37    for embed_utils::FileEntry {
38        rel_path,
39        full_canonical_path,
40    } in embed_utils::get_files(absolute_folder_path.clone(), matcher)
41    {
42        match_values.insert(
43            rel_path.clone(),
44            embed_file(
45                relative_folder_path,
46                ident,
47                &rel_path,
48                &full_canonical_path,
49                metadata_only,
50                crate_path,
51            )?,
52        );
53
54        list_values.push(if let Some(prefix) = prefix {
55            format!("{}{}", prefix, rel_path)
56        } else {
57            rel_path
58        });
59    }
60
61    let array_len = list_values.len();
62
63    // If debug-embed is on, unconditionally include the code below. Otherwise,
64    // make it conditional on cfg(not(debug_assertions)).
65    let not_debug_attr = if cfg!(feature = "debug-embed") {
66        quote! {}
67    } else {
68        quote! { #[cfg(not(debug_assertions))]}
69    };
70
71    let handle_prefix = if let Some(prefix) = prefix {
72        quote! {
73          let file_path = file_path.strip_prefix(#prefix)?;
74        }
75    } else {
76        TokenStream2::new()
77    };
78    let match_values = match_values.into_iter().map(|(path, bytes)| {
79        quote! {
80            (#path, #bytes),
81        }
82    });
83    let value_type = if cfg!(feature = "compression") {
84        quote! { fn() -> #crate_path::EmbeddedFile }
85    } else {
86        quote! { #crate_path::EmbeddedFile }
87    };
88    let get_value = if cfg!(feature = "compression") {
89        quote! {|idx| (ENTRIES[idx].1)()}
90    } else {
91        quote! {|idx| ENTRIES[idx].1.clone()}
92    };
93    Ok(quote! {
94        #not_debug_attr
95        impl #ident {
96            /// Get an embedded file and its metadata.
97            pub fn get(file_path: &str) -> ::std::option::Option<#crate_path::EmbeddedFile> {
98              #handle_prefix
99              let key = file_path.replace("\\", "/");
100              const ENTRIES: &'static [(&'static str, #value_type)] = &[
101                  #(#match_values)*];
102              let position = ENTRIES.binary_search_by_key(&key.as_str(), |entry| entry.0);
103              position.ok().map(#get_value)
104
105            }
106
107            fn names() -> ::std::slice::Iter<'static, &'static str> {
108                const ITEMS: [&str; #array_len] = [#(#list_values),*];
109                ITEMS.iter()
110            }
111
112            /// Iterates over the file paths in the folder.
113            pub fn iter() -> impl ::std::iter::Iterator<Item = ::std::borrow::Cow<'static, str>> {
114                Self::names().map(|x| ::std::borrow::Cow::from(*x))
115            }
116        }
117
118        #not_debug_attr
119        impl #crate_path::RustEmbed for #ident {
120          fn get(file_path: &str) -> ::std::option::Option<#crate_path::EmbeddedFile> {
121            #ident::get(file_path)
122          }
123          fn iter() -> #crate_path::Filenames {
124            #crate_path::Filenames::Embedded(#ident::names())
125          }
126        }
127    })
128}
129
130fn dynamic(
131    ident: &syn::Ident,
132    folder_path: String,
133    prefix: Option<&str>,
134    includes: &[String],
135    excludes: &[String],
136    metadata_only: bool,
137    crate_path: &syn::Path,
138) -> TokenStream2 {
139    let (handle_prefix, map_iter) = if let ::std::option::Option::Some(prefix) = prefix {
140        (
141            quote! { let file_path = file_path.strip_prefix(#prefix)?; },
142            quote! { ::std::borrow::Cow::Owned(format!("{}{}", #prefix, e.rel_path)) },
143        )
144    } else {
145        (TokenStream2::new(), quote! { ::std::borrow::Cow::from(e.rel_path) })
146    };
147
148    let declare_includes = quote! {
149      const INCLUDES: &[&str] = &[#(#includes),*];
150    };
151
152    let declare_excludes = quote! {
153      const EXCLUDES: &[&str] = &[#(#excludes),*];
154    };
155
156    // In metadata_only mode, we still need to read file contents to generate the
157    // file hash, but then we drop the file data.
158    let strip_contents = metadata_only.then_some(quote! {
159        .map(|mut file| { file.data = ::std::default::Default::default(); file })
160    });
161
162    let canonical_folder_path = Path::new(&folder_path)
163        .canonicalize()
164        .expect("folder path must resolve to an absolute path");
165    let canonical_folder_path = canonical_folder_path
166        .to_str()
167        .expect("absolute folder path must be valid unicode");
168
169    quote! {
170        #[cfg(debug_assertions)]
171        impl #ident {
172
173
174          fn matcher() -> #crate_path::utils::PathMatcher {
175              #declare_includes
176              #declare_excludes
177              static PATH_MATCHER: ::std::sync::OnceLock<#crate_path::utils::PathMatcher> = ::std::sync::OnceLock::new();
178              PATH_MATCHER.get_or_init(|| #crate_path::utils::PathMatcher::new(INCLUDES, EXCLUDES)).clone()
179          }
180            /// Get an embedded file and its metadata.
181            pub fn get(file_path: &str) -> ::std::option::Option<#crate_path::EmbeddedFile> {
182                #handle_prefix
183
184                let rel_file_path = file_path.replace("\\", "/");
185                let file_path = ::std::path::Path::new(#folder_path).join(&rel_file_path);
186
187                // Make sure the path requested does not escape the folder path
188                let canonical_file_path = file_path.canonicalize().ok()?;
189                if !canonical_file_path.starts_with(#canonical_folder_path) {
190                    // Tried to request a path that is not in the embedded folder
191
192                    // TODO: Currently it allows "path_traversal_attack" for the symlink files
193                    // For it to be working properly we need to get absolute path first
194                    // and check that instead if it starts with `canonical_folder_path`
195                    // https://doc.rust-lang.org/std/path/fn.absolute.html (currently nightly)
196                    // Should be allowed only if it was a symlink
197                    let metadata = ::std::fs::symlink_metadata(&file_path).ok()?;
198                    if !metadata.is_symlink() {
199                      return ::std::option::Option::None;
200                    }
201                }
202                let path_matcher = Self::matcher();
203                if path_matcher.is_path_included(&rel_file_path) {
204                  #crate_path::utils::read_file_from_fs(&canonical_file_path).ok() #strip_contents
205                } else {
206                  ::std::option::Option::None
207                }
208            }
209
210            /// Iterates over the file paths in the folder.
211            pub fn iter() -> impl ::std::iter::Iterator<Item = ::std::borrow::Cow<'static, str>> {
212                use ::std::path::Path;
213
214
215                #crate_path::utils::get_files(::std::string::String::from(#folder_path), Self::matcher())
216                    .map(|e| #map_iter)
217            }
218        }
219
220        #[cfg(debug_assertions)]
221        impl #crate_path::RustEmbed for #ident {
222          fn get(file_path: &str) -> ::std::option::Option<#crate_path::EmbeddedFile> {
223            #ident::get(file_path)
224          }
225          fn iter() -> #crate_path::Filenames {
226            // the return type of iter() is unnamable, so we have to box it
227            #crate_path::Filenames::Dynamic(::std::boxed::Box::new(#ident::iter()))
228          }
229        }
230    }
231}
232
233fn generate_assets(
234    ident: &syn::Ident,
235    relative_folder_path: Option<&str>,
236    absolute_folder_path: String,
237    prefix: Option<String>,
238    includes: Vec<String>,
239    excludes: Vec<String>,
240    metadata_only: bool,
241    crate_path: &syn::Path,
242) -> syn::Result<TokenStream2> {
243    let embedded_impl = embedded(
244        ident,
245        relative_folder_path,
246        absolute_folder_path.clone(),
247        prefix.as_deref(),
248        &includes,
249        &excludes,
250        metadata_only,
251        crate_path,
252    );
253    if cfg!(feature = "debug-embed") {
254        return embedded_impl;
255    }
256    let embedded_impl = embedded_impl?;
257    let dynamic_impl = dynamic(
258        ident,
259        absolute_folder_path,
260        prefix.as_deref(),
261        &includes,
262        &excludes,
263        metadata_only,
264        crate_path,
265    );
266
267    Ok(quote! {
268        #embedded_impl
269        #dynamic_impl
270    })
271}
272
273fn embed_file(
274    folder_path: Option<&str>,
275    ident: &syn::Ident,
276    rel_path: &str,
277    full_canonical_path: &str,
278    metadata_only: bool,
279    crate_path: &syn::Path,
280) -> syn::Result<TokenStream2> {
281    let file = embed_utils::read_file_from_fs(Path::new(full_canonical_path)).expect("File should be readable");
282    let hash = file.metadata.hash();
283    let last_modified = match file.metadata.last_modified() {
284        Some(last_modified) => quote! { ::std::option::Option::Some(#last_modified) },
285        None => quote! { ::std::option::Option::None },
286    };
287    let created = match file.metadata.created() {
288        Some(created) => quote! { ::std::option::Option::Some(#created) },
289        None => quote! { ::std::option::Option::None },
290    };
291    #[cfg(feature = "mime-guess")]
292    let mimetype_tokens = {
293        let mt = file.metadata.mimetype();
294        quote! { , #mt }
295    };
296    #[cfg(not(feature = "mime-guess"))]
297    let mimetype_tokens = TokenStream2::new();
298
299    let embedding_code = if metadata_only {
300        quote! {
301            const BYTES: &'static [u8] = &[];
302        }
303    } else if cfg!(feature = "compression") {
304        let folder_path = folder_path.ok_or(syn::Error::new(
305            ident.span(),
306            "`folder` must be provided under `compression` feature.",
307        ))?;
308        // Print some debugging information
309        let full_relative_path = PathBuf::from_iter([folder_path, rel_path]);
310        let full_relative_path = full_relative_path.to_string_lossy();
311        quote! {
312          #crate_path::flate!(static BYTES: [u8] from #full_relative_path);
313        }
314    } else {
315        quote! {
316          const BYTES: &'static [u8] = include_bytes!(#full_canonical_path);
317        }
318    };
319    let closure_args = if cfg!(feature = "compression") {
320        quote! { || }
321    } else {
322        quote! {}
323    };
324    Ok(quote! {
325         #closure_args {
326          #embedding_code
327
328          #crate_path::EmbeddedFile {
329              data: ::std::borrow::Cow::Borrowed(&BYTES),
330              metadata: #crate_path::Metadata::__embed_new([#(#hash),*], #last_modified, #created #mimetype_tokens)
331          }
332        }
333    })
334}
335
336/// Find all pairs of the `name = "value"` attribute from the derive input
337fn find_attribute_values(ast: &syn::DeriveInput, attr_name: &str) -> Vec<String> {
338    ast.attrs
339        .iter()
340        .filter(|value| value.path().is_ident(attr_name))
341        .filter_map(|attr| match &attr.meta {
342            Meta::NameValue(MetaNameValue {
343                value: Expr::Lit(ExprLit {
344                    lit: Lit::Str(val), ..
345                }),
346                ..
347            }) => Some(val.value()),
348            _ => None,
349        })
350        .collect()
351}
352
353fn find_bool_attribute(ast: &syn::DeriveInput, attr_name: &str) -> Option<bool> {
354    ast.attrs
355        .iter()
356        .find(|value| value.path().is_ident(attr_name))
357        .and_then(|attr| match &attr.meta {
358            Meta::NameValue(MetaNameValue {
359                value: Expr::Lit(ExprLit {
360                    lit: Lit::Bool(val), ..
361                }),
362                ..
363            }) => Some(val.value()),
364            _ => None,
365        })
366}
367
368fn impl_embed(ast: &syn::DeriveInput) -> syn::Result<TokenStream2> {
369    match ast.data {
370        Data::Struct(ref data) => match data.fields {
371            Fields::Unit => {}
372            _ => return Err(syn::Error::new_spanned(ast, "RustEmbed can only be derived for unit structs")),
373        },
374        _ => return Err(syn::Error::new_spanned(ast, "RustEmbed can only be derived for unit structs")),
375    };
376
377    let crate_path: syn::Path = find_attribute_values(ast, "crate_path")
378        .last()
379        .map(|v| syn::parse_str(&v).unwrap())
380        .unwrap_or_else(|| syn::parse_str("embed").unwrap());
381
382    let mut folder_paths = find_attribute_values(ast, "folder");
383    if folder_paths.len() != 1 {
384        return Err(syn::Error::new_spanned(
385            ast,
386            "#[derive(RustEmbed)] must contain one attribute like this #[folder = \"examples/public/\"]",
387        ));
388    }
389    let folder_path = folder_paths.remove(0);
390
391    let prefix = find_attribute_values(ast, "prefix").into_iter().next();
392    let includes = find_attribute_values(ast, "include");
393    let excludes = find_attribute_values(ast, "exclude");
394    let metadata_only = find_bool_attribute(ast, "metadata_only").unwrap_or(false);
395
396    #[cfg(not(feature = "include-exclude"))]
397    if !includes.is_empty() || !excludes.is_empty() {
398        return Err(syn::Error::new_spanned(
399            ast,
400            "Please turn on the `include-exclude` feature to use the `include` and `exclude` attributes",
401        ));
402    }
403
404    #[cfg(feature = "interpolate-folder-path")]
405    let folder_path = shellexpand::full(&folder_path)
406        .map_err(|v| syn::Error::new_spanned(ast, v.to_string()))?
407        .to_string();
408
409    // Base relative paths on the Cargo.toml location
410    let (relative_path, absolute_folder_path) = if Path::new(&folder_path).is_relative() {
411        let absolute_path = Path::new(&env::var("CARGO_MANIFEST_DIR").unwrap())
412            .join(&folder_path)
413            .to_str()
414            .unwrap()
415            .to_owned();
416        (Some(folder_path.clone()), absolute_path)
417    } else {
418        if cfg!(feature = "compression") {
419            return Err(syn::Error::new_spanned(
420                ast,
421                "`folder` must be a relative path under `compression` feature.",
422            ));
423        }
424        (None, folder_path)
425    };
426
427    if !Path::new(&absolute_folder_path).exists() {
428        let mut message = format!(
429            "#[derive(RustEmbed)] folder '{}' does not exist. cwd: '{}'",
430            absolute_folder_path,
431            std::env::current_dir().unwrap().to_str().unwrap()
432        );
433
434        // Add a message about the interpolate-folder-path feature if the path may
435        // include a variable
436        if absolute_folder_path.contains('$') && cfg!(not(feature = "interpolate-folder-path")) {
437            message += "\nA variable has been detected. RustEmbed can expand variables \
438                  when the `interpolate-folder-path` feature is enabled.";
439        }
440
441        return Err(syn::Error::new_spanned(ast, message));
442    };
443
444    generate_assets(
445        &ast.ident,
446        relative_path.as_deref(),
447        absolute_folder_path,
448        prefix,
449        includes,
450        excludes,
451        metadata_only,
452        &crate_path,
453    )
454}
455
456#[proc_macro_derive(RustEmbed, attributes(folder, prefix, include, exclude, metadata_only, crate_path))]
457pub fn derive_input_object(input: TokenStream) -> TokenStream {
458    let ast = parse_macro_input!(input as DeriveInput);
459    match impl_embed(&ast) {
460        Ok(ok) => ok.into(),
461        Err(e) => e.to_compile_error().into(),
462    }
463}