Add configurable thread pool size for MySQL connections and update schematics loading logic

🔧 Add invalid_nbt flag.

schemsearch-cli/src/main.rs

@@ -259,7 +259,7 @@ fn main() {
         if let Some(x) = matches.get_many::<String>("sql-filter-name") {
             filter = filter.name(x.collect());
         }
-        for schem in block_on(load_all_schematics(filter)) {
+        for schem in block_on(load_all_schematics(filter, *matches.get_one::<usize>("threads").expect("Could not get Threads") as u32)) {
             schematics.push(SchematicSupplierType::SQL(SqlSchematicSupplier {
                 node: schem
             }))

schemsearch-files/src/lib.rs

@@ -15,10 +15,10 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
+use nbt::{CompoundTag, Tag};
 use std::collections::hash_map::HashMap;
 use std::io::Read;
 use std::path::PathBuf;
-use nbt::{CompoundTag, Tag};
 
 #[derive(Clone, Debug)]
 pub struct SpongeSchematic {
@@ -62,8 +62,12 @@ pub struct Entity {
 }
 
 impl SpongeSchematic {
-    pub fn load_data<R>(data: &mut R) -> Result<SpongeSchematic, String> where R: Read {
-        let nbt: CompoundTag = nbt::decode::read_gzip_compound_tag(data).map_err(|e| e.to_string())?;
+    pub fn load_data<R>(data: &mut R) -> Result<SpongeSchematic, String>
+    where
+        R: Read,
+    {
+        let nbt: CompoundTag =
+            nbt::decode::read_gzip_compound_tag(data).map_err(|e| e.to_string())?;
         let version = nbt.get_i32("Version").unwrap_or_else(|_| {
             return if nbt.contains_key("Blocks") {
                 3
@@ -79,7 +83,10 @@ impl SpongeSchematic {
         match version {
             1 => SpongeSchematic::from_nbt_1(nbt),
             2 => SpongeSchematic::from_nbt_2(nbt),
-            3 => SpongeSchematic::from_nbt_3(nbt),
+            3 => SpongeSchematic::from_nbt_3(
+                nbt.get_compound_tag("Schematic")
+                    .map_err(|e| e.to_string())?,
+            ),
             _ => Err("Invalid schematic: Unknown Version".to_string()),
         }
     }
@@ -92,7 +99,10 @@ impl SpongeSchematic {
     pub fn from_nbt_1(nbt: CompoundTag) -> Result<Self, String> {
         Ok(Self {
             data_version: 0,
-            metadata: nbt.get_compound_tag("Metadata").map_err(|e| e.to_string())?.clone(),
+            metadata: nbt
+                .get_compound_tag("Metadata")
+                .map(|v| v.clone())
+                .unwrap_or_else(|_| CompoundTag::new()),
             width: nbt.get_i16("Width").map_err(|e| e.to_string())? as u16,
             height: nbt.get_i16("Height").map_err(|e| e.to_string())? as u16,
             length: nbt.get_i16("Length").map_err(|e| e.to_string())? as u16,
@@ -100,15 +110,21 @@ impl SpongeSchematic {
             palette_max: nbt.get_i32("PaletteMax").map_err(|e| e.to_string())?,
             palette: read_palette(nbt.get_compound_tag("Palette").map_err(|e| e.to_string())?),
             block_data: read_blocks(nbt.get_i8_vec("BlockData").map_err(|e| e.to_string())?),
-            block_entities: read_tile_entities(nbt.get_compound_tag_vec("TileEntities").unwrap_or_else(|_| vec![]))?,
+            block_entities: read_tile_entities(
+                nbt.get_compound_tag_vec("TileEntities")
+                    .unwrap_or_else(|_| vec![]),
+            )?,
             entities: None,
         })
     }
 
     pub fn from_nbt_2(nbt: CompoundTag) -> Result<Self, String> {
-        Ok(Self{
+        Ok(Self {
             data_version: nbt.get_i32("DataVersion").map_err(|e| e.to_string())?,
-            metadata: nbt.get_compound_tag("Metadata").map_err(|e| e.to_string())?.clone(),
+            metadata: nbt
+                .get_compound_tag("Metadata")
+                .map(|v| v.clone())
+                .unwrap_or_else(|_| CompoundTag::new()),
             width: nbt.get_i16("Width").map_err(|e| e.to_string())? as u16,
             height: nbt.get_i16("Height").map_err(|e| e.to_string())? as u16,
             length: nbt.get_i16("Length").map_err(|e| e.to_string())? as u16,
@@ -116,32 +132,37 @@ impl SpongeSchematic {
             palette_max: nbt.get_i32("PaletteMax").map_err(|e| e.to_string())?,
             palette: read_palette(nbt.get_compound_tag("Palette").map_err(|e| e.to_string())?),
             block_data: read_blocks(nbt.get_i8_vec("BlockData").map_err(|e| e.to_string())?),
-            block_entities: read_tile_entities(nbt.get_compound_tag_vec("BlockEntities").unwrap_or_else(|_| vec![]))?,
+            block_entities: read_tile_entities(
+                nbt.get_compound_tag_vec("BlockEntities")
+                    .unwrap_or_else(|_| vec![]),
+            )?,
             entities: None,
         })
     }
 
-    pub fn from_nbt_3(nbt: CompoundTag) -> Result<Self, String> {
+    pub fn from_nbt_3(nbt: &CompoundTag) -> Result<Self, String> {
         let blocks = nbt.get_compound_tag("Blocks").map_err(|e| e.to_string())?;
-        Ok(Self{
+        Ok(Self {
             data_version: nbt.get_i32("DataVersion").map_err(|e| e.to_string())?,
-            metadata: nbt.get_compound_tag("Metadata").map_err(|e| e.to_string())?.clone(),
+            metadata: nbt
+                .get_compound_tag("Metadata")
+                .map(|v| v.clone())
+                .unwrap_or_else(|_| CompoundTag::new()),
             width: nbt.get_i16("Width").map_err(|e| e.to_string())? as u16,
             height: nbt.get_i16("Height").map_err(|e| e.to_string())? as u16,
             length: nbt.get_i16("Length").map_err(|e| e.to_string())? as u16,
             offset: read_offset(nbt.get_i32_vec("Offset").map_err(|e| e.to_string())?)?,
-            palette_max: compute_palette_max(blocks.get_compound_tag("Palette").map_err(|e| e.to_string())?),
+            palette_max: compute_palette_max(blocks.get_compound_tag("Palette").map_err(|e| e.to_string())?) + 1,
             palette: read_palette(blocks.get_compound_tag("Palette").map_err(|e| e.to_string())?),
-            block_data: read_blocks(blocks.get_i8_vec("BlockData").map_err(|e| e.to_string())?),
+            block_data: read_blocks(blocks.get_i8_vec("Data").map_err(|e| e.to_string())?),
             block_entities: read_tile_entities(blocks.get_compound_tag_vec("BlockEntities").unwrap_or_else(|_| vec![]))?,
             entities: None,
         })
     }
-
 }
 
 fn read_tile_entities(tag: Vec<&CompoundTag>) -> Result<Vec<BlockEntity>, String> {
-let mut tile_entities = Vec::new();
+    let mut tile_entities = Vec::new();
     for t in tag {
         tile_entities.push(BlockEntity {
             id: t.get_str("Id").map_err(|e| e.to_string())?.to_string(),
@@ -164,8 +185,10 @@ fn read_palette(p: &CompoundTag) -> HashMap<String, i32> {
     let mut palette = HashMap::new();
     for (key, value) in p.iter() {
         match value {
-            Tag::Int(n) => { palette.insert(key.clone(), *n); },
-            _ => {},
+            Tag::Int(n) => {
+                palette.insert(key.clone(), *n);
+            }
+            _ => {}
         };
     }
     palette
@@ -173,10 +196,15 @@ fn read_palette(p: &CompoundTag) -> HashMap<String, i32> {
 
 #[inline]
 fn compute_palette_max(palette: &CompoundTag) -> i32 {
-    palette.iter().map(|(_, v)| v).filter_map(|v| match v {
-        Tag::Int(n) => Some(*n),
-        _ => None,
-    }).max().unwrap_or(0)
+    palette
+        .iter()
+        .map(|(_, v)| v)
+        .filter_map(|v| match v {
+            Tag::Int(n) => Some(*n),
+            _ => None,
+        })
+        .max()
+        .unwrap_or(0)
 }
 
 #[inline]
@@ -193,13 +221,16 @@ pub fn read_varint_array(read: &Vec<i8>) -> Vec<i32> {
     let mut cursor = 0;
     loop {
         match read.get(cursor) {
-            Some(byte) => { current_byte = *byte as u8; cursor += 1; },
+            Some(byte) => {
+                current_byte = *byte as u8;
+                cursor += 1;
+            }
             None => break,
         };
 
         value |= (((current_byte & 0x7F) as u32) << position) as i32;
 
-        if(current_byte & 0x80) == 0 {
+        if (current_byte & 0x80) == 0 {
             data.push(value);
             value = 0;
             position = 0;

schemsearch-lib/src/nbt_search.rs

@@ -7,7 +7,9 @@ use schemsearch_files::SpongeSchematic;
 const NBT_BLOCKS: &str = include_str!("blocks.txt");
 
 lazy_static! {
-    static ref NBT_BLOCKS_SET: HashSet<String> = NBT_BLOCKS.lines().map(ToOwned::to_owned).collect();
+    static ref NBT_BLOCKS_SET: HashSet<String> = {
+        NBT_BLOCKS.lines().map(|x| format!("minecraft:{}", x)).collect()
+    };
 }
 
 pub fn has_invalid_nbt(schem: SpongeSchematic) -> bool {
@@ -15,13 +17,94 @@ pub fn has_invalid_nbt(schem: SpongeSchematic) -> bool {
         return true;
     }
 
-    let nbt_blocks = schem.palette.iter().filter(|(k, _)| NBT_BLOCKS_SET.contains(*k)).map(|(_, v)| *v).collect::<HashSet<i32>>();
+    let nbt_blocks = schem.palette.iter().filter(|(k, _)| NBT_BLOCKS_SET.contains(k.to_owned())).map(|(_, v)| *v).collect::<HashSet<i32>>();
 
-    for block_entity in schem.block_data.iter() {
-        if nbt_blocks.contains(block_entity) {
-            return true;
+    for (i, block_entity) in schem.block_data.iter().enumerate() {
+        if nbt_blocks.contains(&*block_entity) {
+            // i = x + z * Width + y * Width * Length
+            let x = i % schem.width as usize;
+            let z = (i / schem.width as usize) % schem.length as usize;
+            let y = i / (schem.width as usize * schem.length as usize);
+            if schem.block_entities.iter().any(|e| !e.pos.eq(&[x as i32, y as i32, z as i32])) {
+                return true;
+            }
         }
     }
 
     return false;
+}
+
+#[allow(unused_imports)]
+#[cfg(test)]
+mod tests {
+    use nbt::CompoundTag;
+    use schemsearch_files::{BlockEntity, SpongeSchematic};
+    use super::*;
+
+    #[test]
+    fn test_has_invalid_nbt() {
+        let schem = SpongeSchematic {
+            data_version: 1,
+            metadata: CompoundTag::new(),
+            width: 0,
+            height: 0,
+            length: 0,
+            offset: [0, 0, 0],
+            palette_max: 1,
+            palette: vec![("minecraft:chest".to_owned(), 1)].into_iter().collect(),
+            block_data: vec![1],
+            block_entities: vec![],
+            entities: None,
+        };
+
+        assert_eq!(has_invalid_nbt(schem), true);
+    }
+
+    #[test]
+    fn test_has_invalid_nbt_2() {
+        let schem = SpongeSchematic {
+            data_version: 1,
+            metadata: CompoundTag::new(),
+            width: 1,
+            height: 1,
+            length: 1,
+            offset: [0, 0, 0],
+            palette_max: 1,
+            palette: vec![("minecraft:chest".to_owned(), 1)].into_iter().collect(),
+            block_data: vec![1],
+            block_entities: vec![
+                BlockEntity {
+                    id: "minecraft:chest".to_owned(),
+                    pos: [0, 0, 0],
+                }
+            ],
+            entities: None,
+        };
+
+        assert_eq!(has_invalid_nbt(schem), false);
+    }
+
+    #[test]
+    fn test_has_invalid_nbt_3() {
+        let schem = SpongeSchematic {
+            data_version: 1,
+            metadata: CompoundTag::new(),
+            width: 2,
+            height: 1,
+            length: 1,
+            offset: [0, 0, 0],
+            palette_max: 1,
+            palette: vec![("minecraft:chest".to_owned(), 1), ("minecraft:stone".to_owned(), 2)].into_iter().collect(),
+            block_data: vec![1, 2],
+            block_entities: vec![
+                BlockEntity {
+                    id: "minecraft:chest".to_owned(),
+                    pos: [1, 0, 0],
+                }
+            ],
+            entities: None,
+        };
+
+        assert_eq!(has_invalid_nbt(schem), true);
+    }
 }

schemsearch-sql/src/lib.rs

@@ -30,12 +30,12 @@ pub struct SchematicNode {
     pub name: String
 }
 
-pub async unsafe fn get_connection() {
+pub async unsafe fn get_connection(connections: u32) {
     let mut conn = CONN.lock().unwrap();
     if conn.is_none() {
         let properties = properties::load_mysql_properties();
         let _ = conn.insert(MySqlPoolOptions::new()
-            .max_connections(5)
+            .max_connections(connections)
             .connect_with(MySqlConnectOptions::new()
                 .host(properties.host.as_str())
                 .port(3306)
@@ -46,10 +46,10 @@ pub async unsafe fn get_connection() {
     }
 }
 
-pub async fn load_all_schematics(filter: SchematicFilter) -> Vec<SchematicNode> {
-    unsafe { get_connection().await; }
+pub async fn load_all_schematics(filter: SchematicFilter, connections: u32) -> Vec<SchematicNode> {
+    unsafe { get_connection(connections).await; }
     let mut schematics = Vec::new();
-    let rows = unsafe { &CONN }.lock().unwrap().as_mut().unwrap().fetch_all(format!("SELECT SN.NodeId, SN.NodeName FROM NodeData ND INNER JOIN SchematicNode SN ON SN.NodeId = ND.NodeId WHERE NodeFormat = true {}", filter.build()).as_str()).await.expect("Failed to fetch schematics");
+    let rows = unsafe { &CONN }.lock().unwrap().as_mut().unwrap().fetch_all(format!("SELECT SN.NodeId, SN.NodeName FROM NodeData ND INNER JOIN SchematicNode SN ON SN.NodeId = ND.NodeId WHERE NodeFormat >= 2 {}", filter.build()).as_str()).await.expect("Failed to fetch schematics");
     for row in rows {
         schematics.push(SchematicNode {
             id: row.get(0),
@@ -60,7 +60,7 @@ pub async fn load_all_schematics(filter: SchematicFilter) -> Vec<SchematicNode>
 }
 
 pub async fn load_schemdata(id: i32) -> Vec<u8> {
-    unsafe { get_connection().await; }
+    unsafe { get_connection(5).await; }
     let rows = unsafe { &CONN }.lock().unwrap().as_mut().unwrap().fetch_one(format!("SELECT SchemData FROM NodeData WHERE NodeId = {}", id).as_str()).await.expect("Failed to fetch schematics");
     rows.get(0)
 }