1 // Copyright 2018 The Chromium OS Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4 
5 use std::fmt::{self, Display};
6 use std::io;
7 
8 use libc::EINVAL;
9 
10 use crate::qcow::qcow_raw_file::QcowRawFile;
11 use crate::qcow::vec_cache::{CacheMap, Cacheable, VecCache};
12 
13 #[derive(Debug)]
14 pub enum Error {
15     /// `EvictingCache` - Error writing a refblock from the cache to disk.
16     EvictingRefCounts(io::Error),
17     /// `InvalidIndex` - Address requested isn't within the range of the disk.
18     InvalidIndex,
19     /// `NeedCluster` - Handle this error by reading the cluster and calling the function again.
20     NeedCluster(u64),
21     /// `NeedNewCluster` - Handle this error by allocating a cluster and calling the function again.
22     NeedNewCluster,
23     /// `ReadingRefCounts` - Error reading the file in to the refcount cache.
24     ReadingRefCounts(io::Error),
25 }
26 
27 pub type Result<T> = std::result::Result<T, Error>;
28 
29 impl Display for Error {
fmt(&self, f: &mut fmt::Formatter) -> fmt::Result30     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
31         use self::Error::*;
32 
33         match self {
34             EvictingRefCounts(e) => write!(
35                 f,
36                 "failed to write a refblock from the cache to disk: {}",
37                 e
38             ),
39             InvalidIndex => write!(f, "address requested is not within the range of the disk"),
40             NeedCluster(addr) => write!(f, "cluster with addr={} needs to be read", addr),
41             NeedNewCluster => write!(f, "new cluster needs to be allocated for refcounts"),
42             ReadingRefCounts(e) => {
43                 write!(f, "failed to read the file into the refcount cache: {}", e)
44             }
45         }
46     }
47 }
48 
49 /// Represents the refcount entries for an open qcow file.
50 #[derive(Debug)]
51 pub struct RefCount {
52     ref_table: VecCache<u64>,
53     refcount_table_offset: u64,
54     refblock_cache: CacheMap<VecCache<u16>>,
55     refcount_block_entries: u64, // number of refcounts in a cluster.
56     cluster_size: u64,
57     max_valid_cluster_offset: u64,
58 }
59 
60 impl RefCount {
61     /// Creates a `RefCount` from `file`, reading the refcount table from `refcount_table_offset`.
62     /// `refcount_table_entries` specifies the number of refcount blocks used by this image.
63     /// `refcount_block_entries` indicates the number of refcounts in each refcount block.
64     /// Each refcount table entry points to a refcount block.
new( raw_file: &mut QcowRawFile, refcount_table_offset: u64, refcount_table_entries: u64, refcount_block_entries: u64, cluster_size: u64, ) -> io::Result<RefCount>65     pub fn new(
66         raw_file: &mut QcowRawFile,
67         refcount_table_offset: u64,
68         refcount_table_entries: u64,
69         refcount_block_entries: u64,
70         cluster_size: u64,
71     ) -> io::Result<RefCount> {
72         let ref_table = VecCache::from_vec(raw_file.read_pointer_table(
73             refcount_table_offset,
74             refcount_table_entries,
75             None,
76         )?);
77         let max_valid_cluster_index = (ref_table.len() as u64) * refcount_block_entries - 1;
78         let max_valid_cluster_offset = max_valid_cluster_index * cluster_size;
79         Ok(RefCount {
80             ref_table,
81             refcount_table_offset,
82             refblock_cache: CacheMap::new(50),
83             refcount_block_entries,
84             cluster_size,
85             max_valid_cluster_offset,
86         })
87     }
88 
89     /// Returns the number of refcounts per block.
refcounts_per_block(&self) -> u6490     pub fn refcounts_per_block(&self) -> u64 {
91         self.refcount_block_entries
92     }
93 
94     /// Returns the maximum valid cluster offset in the raw file for this refcount table.
max_valid_cluster_offset(&self) -> u6495     pub fn max_valid_cluster_offset(&self) -> u64 {
96         self.max_valid_cluster_offset
97     }
98 
99     /// Returns `NeedNewCluster` if a new cluster needs to be allocated for refcounts. If an
100     /// existing cluster needs to be read, `NeedCluster(addr)` is returned. The Caller should
101     /// allocate a cluster or read the required one and call this function again with the cluster.
102     /// On success, an optional address of a dropped cluster is returned. The dropped cluster can
103     /// be reused for other purposes.
set_cluster_refcount( &mut self, raw_file: &mut QcowRawFile, cluster_address: u64, refcount: u16, mut new_cluster: Option<(u64, VecCache<u16>)>, ) -> Result<Option<u64>>104     pub fn set_cluster_refcount(
105         &mut self,
106         raw_file: &mut QcowRawFile,
107         cluster_address: u64,
108         refcount: u16,
109         mut new_cluster: Option<(u64, VecCache<u16>)>,
110     ) -> Result<Option<u64>> {
111         let (table_index, block_index) = self.get_refcount_index(cluster_address);
112 
113         let block_addr_disk = *self.ref_table.get(table_index).ok_or(Error::InvalidIndex)?;
114 
115         // Fill the cache if this block isn't yet there.
116         if !self.refblock_cache.contains_key(&table_index) {
117             // Need a new cluster
118             if let Some((addr, table)) = new_cluster.take() {
119                 self.ref_table[table_index] = addr;
120                 let ref_table = &self.ref_table;
121                 self.refblock_cache
122                     .insert(table_index, table, |index, evicted| {
123                         raw_file.write_refcount_block(ref_table[index], evicted.get_values())
124                     })
125                     .map_err(Error::EvictingRefCounts)?;
126             } else {
127                 if block_addr_disk == 0 {
128                     return Err(Error::NeedNewCluster);
129                 }
130                 return Err(Error::NeedCluster(block_addr_disk));
131             }
132         }
133 
134         // Unwrap is safe here as the entry was filled directly above.
135         let dropped_cluster = if !self.refblock_cache.get(&table_index).unwrap().dirty() {
136             // Free the previously used block and use a new one. Writing modified counts to new
137             // blocks keeps the on-disk state consistent even if it's out of date.
138             if let Some((addr, _)) = new_cluster.take() {
139                 self.ref_table[table_index] = addr;
140                 Some(block_addr_disk)
141             } else {
142                 return Err(Error::NeedNewCluster);
143             }
144         } else {
145             None
146         };
147 
148         self.refblock_cache.get_mut(&table_index).unwrap()[block_index] = refcount;
149         Ok(dropped_cluster)
150     }
151 
152     /// Flush the dirty refcount blocks. This must be done before flushing the table that points to
153     /// the blocks.
flush_blocks(&mut self, raw_file: &mut QcowRawFile) -> io::Result<()>154     pub fn flush_blocks(&mut self, raw_file: &mut QcowRawFile) -> io::Result<()> {
155         // Write out all dirty L2 tables.
156         for (table_index, block) in self.refblock_cache.iter_mut().filter(|(_k, v)| v.dirty()) {
157             let addr = self.ref_table[*table_index];
158             if addr != 0 {
159                 raw_file.write_refcount_block(addr, block.get_values())?;
160             } else {
161                 return Err(std::io::Error::from_raw_os_error(EINVAL));
162             }
163             block.mark_clean();
164         }
165         Ok(())
166     }
167 
168     /// Flush the refcount table that keeps the address of the refcounts blocks.
169     /// Returns true if the table changed since the previous `flush_table()` call.
flush_table(&mut self, raw_file: &mut QcowRawFile) -> io::Result<bool>170     pub fn flush_table(&mut self, raw_file: &mut QcowRawFile) -> io::Result<bool> {
171         if self.ref_table.dirty() {
172             raw_file.write_pointer_table(
173                 self.refcount_table_offset,
174                 &self.ref_table.get_values(),
175                 0,
176             )?;
177             self.ref_table.mark_clean();
178             Ok(true)
179         } else {
180             Ok(false)
181         }
182     }
183 
184     /// Gets the refcount for a cluster with the given address.
get_cluster_refcount( &mut self, raw_file: &mut QcowRawFile, address: u64, ) -> Result<u16>185     pub fn get_cluster_refcount(
186         &mut self,
187         raw_file: &mut QcowRawFile,
188         address: u64,
189     ) -> Result<u16> {
190         let (table_index, block_index) = self.get_refcount_index(address);
191         let block_addr_disk = *self.ref_table.get(table_index).ok_or(Error::InvalidIndex)?;
192         if block_addr_disk == 0 {
193             return Ok(0);
194         }
195         if !self.refblock_cache.contains_key(&table_index) {
196             let table = VecCache::from_vec(
197                 raw_file
198                     .read_refcount_block(block_addr_disk)
199                     .map_err(Error::ReadingRefCounts)?,
200             );
201             let ref_table = &self.ref_table;
202             self.refblock_cache
203                 .insert(table_index, table, |index, evicted| {
204                     raw_file.write_refcount_block(ref_table[index], evicted.get_values())
205                 })
206                 .map_err(Error::EvictingRefCounts)?;
207         }
208         Ok(self.refblock_cache.get(&table_index).unwrap()[block_index])
209     }
210 
211     /// Returns the refcount table for this file. This is only useful for debugging.
ref_table(&self) -> &[u64]212     pub fn ref_table(&self) -> &[u64] {
213         &self.ref_table.get_values()
214     }
215 
216     /// Returns the refcounts stored in the given block.
refcount_block( &mut self, raw_file: &mut QcowRawFile, table_index: usize, ) -> Result<Option<&[u16]>>217     pub fn refcount_block(
218         &mut self,
219         raw_file: &mut QcowRawFile,
220         table_index: usize,
221     ) -> Result<Option<&[u16]>> {
222         let block_addr_disk = *self.ref_table.get(table_index).ok_or(Error::InvalidIndex)?;
223         if block_addr_disk == 0 {
224             return Ok(None);
225         }
226         if !self.refblock_cache.contains_key(&table_index) {
227             let table = VecCache::from_vec(
228                 raw_file
229                     .read_refcount_block(block_addr_disk)
230                     .map_err(Error::ReadingRefCounts)?,
231             );
232             // TODO(dgreid) - closure needs to return an error.
233             let ref_table = &self.ref_table;
234             self.refblock_cache
235                 .insert(table_index, table, |index, evicted| {
236                     raw_file.write_refcount_block(ref_table[index], evicted.get_values())
237                 })
238                 .map_err(Error::EvictingRefCounts)?;
239         }
240         // The index must exist as it was just inserted if it didn't already.
241         Ok(Some(
242             self.refblock_cache.get(&table_index).unwrap().get_values(),
243         ))
244     }
245 
246     // Gets the address of the refcount block and the index into the block for the given address.
get_refcount_index(&self, address: u64) -> (usize, usize)247     fn get_refcount_index(&self, address: u64) -> (usize, usize) {
248         let block_index = (address / self.cluster_size) % self.refcount_block_entries;
249         let refcount_table_index = (address / self.cluster_size) / self.refcount_block_entries;
250         (refcount_table_index as usize, block_index as usize)
251     }
252 }
253