fix: harden DiskANN against type confusion and invalid vector datums

pgvectorscale/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "vectorscale"
-version = "0.9.0"
+version = "0.9.1"
 edition = "2021"
 
 [lib]

pgvectorscale/sql/vectorscale--0.9.0--0.9.1.sql

@@ -0,0 +1,68 @@
+-- Rebind the extension's C entry points to the versioned 0.9.1 shared library.
+CREATE OR REPLACE FUNCTION diskann_amhandler(internal)
+RETURNS index_am_handler
+PARALLEL SAFE IMMUTABLE STRICT COST 0.0001
+LANGUAGE c
+AS 'vectorscale-0.9.1', 'amhandler_wrapper';
+
+CREATE OR REPLACE FUNCTION distance_type_cosine()
+RETURNS smallint
+IMMUTABLE STRICT PARALLEL SAFE
+LANGUAGE c
+AS 'vectorscale-0.9.1', 'distance_type_cosine_wrapper';
+
+CREATE OR REPLACE FUNCTION distance_type_inner_product()
+RETURNS smallint
+IMMUTABLE STRICT PARALLEL SAFE
+LANGUAGE c
+AS 'vectorscale-0.9.1', 'distance_type_inner_product_wrapper';
+
+CREATE OR REPLACE FUNCTION distance_type_l2()
+RETURNS smallint
+IMMUTABLE STRICT PARALLEL SAFE
+LANGUAGE c
+AS 'vectorscale-0.9.1', 'distance_type_l2_wrapper';
+
+CREATE OR REPLACE FUNCTION smallint_array_overlap(
+    "left" smallint[],
+    "right" smallint[]
+)
+RETURNS bool
+IMMUTABLE STRICT PARALLEL SAFE
+LANGUAGE c
+AS 'vectorscale-0.9.1', 'smallint_array_overlap_wrapper';
+
+DO $$
+DECLARE
+    expected_vector_type oid;
+BEGIN
+    SELECT t.oid
+    INTO STRICT expected_vector_type
+    FROM pg_catalog.pg_extension e
+    JOIN pg_catalog.pg_type t
+      ON t.typnamespace = e.extnamespace
+     AND t.typname = 'vector'
+    WHERE e.extname = 'vector';
+
+    IF EXISTS (
+        SELECT 1
+        FROM pg_catalog.pg_opclass c
+        JOIN pg_catalog.pg_am am ON am.oid = c.opcmethod
+        WHERE am.amname = 'diskann'
+          AND c.opcnamespace = (
+              SELECT oid
+              FROM pg_catalog.pg_namespace
+              WHERE nspname = '@extschema@'
+          )
+          AND c.opcname IN (
+              'vector_cosine_ops',
+              'vector_l2_ops',
+              'vector_ip_ops'
+          )
+          AND c.opcintype IS DISTINCT FROM expected_vector_type
+    ) THEN
+        RAISE EXCEPTION
+            'diskann: a vector operator class is not bound to pgvector''s vector type; drop the affected operator class and recreate the extension objects';
+    END IF;
+END;
+$$;

pgvectorscale/src/access_method/build.rs

@@ -250,7 +250,9 @@ fn get_meta_page(
     index_relation: &PgRelation,
     opt: PgBox<TSVIndexOptions>,
 ) -> MetaPage {
-    let dimensions = index_relation.tuple_desc().get(0).unwrap().atttypmod;
+    let typmod = index_relation.tuple_desc().get(0).unwrap().atttypmod;
+    let dimensions = crate::access_method::vector_type::dimension_from_typmod(typmod)
+        .unwrap_or_else(|message| error!("{}", message));
 
     let distance_type = unsafe {
         let fmgr_info = index_getprocinfo(indexrel, 1, DISKANN_DISTANCE_TYPE_PROC);
@@ -265,16 +267,7 @@ fn get_meta_page(
         error!("Inner product distance type is not supported with plain storage");
     }
 
-    let meta_page =
-        unsafe { MetaPage::create(index_relation, dimensions as _, distance_type, opt) };
-
-    if meta_page.get_num_dimensions_to_index() == 0 {
-        error!("No dimensions to index");
-    }
-
-    if meta_page.get_num_dimensions_to_index() > MAX_DIMENSION {
-        error!("Too many dimensions to index (max is {})", MAX_DIMENSION);
-    }
+    let meta_page = unsafe { MetaPage::create(index_relation, dimensions, distance_type, opt) };
 
     if meta_page.get_num_dimensions_to_index() > MAX_DIMENSION_NO_SBQ
         && meta_page.get_storage_type() == StorageType::Plain

pgvectorscale/src/access_method/meta_page.rs

@@ -309,6 +309,12 @@ impl MetaPage {
             opt.num_dimensions
         };
 
+        crate::access_method::vector_type::ensure_valid_dimensions(
+            num_dimensions,
+            num_dimensions_to_index,
+        )
+        .unwrap_or_else(|message| pgrx::error!("{}", message));
+
         let bq_num_bits_per_dimension =
             if opt.bq_num_bits_per_dimension == SBQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL {
                 if (*opt).get_storage_type() == StorageType::SbqCompression
@@ -400,7 +406,7 @@ impl MetaPage {
         unsafe {
             let page = page::ReadablePage::read(index, META_BLOCK_NUMBER);
             let page_type = page.get_type();
-            match page_type {
+            let meta = match page_type {
                 PageType::MetaV1 => {
                     let old_meta = MetaPageV1::page_get_meta(*page, *(*(page.get_buffer())));
                     let new_meta: MetaPage = (&*old_meta).into();
@@ -414,7 +420,13 @@ impl MetaPage {
                 PageType::MetaV2 => MetaPageV2::from_page(page).into(),
                 PageType::Meta => Self::load(index),
                 _ => pgrx::error!("Meta page is not of type Meta"),
-            }
+            };
+            crate::access_method::vector_type::ensure_valid_dimensions(
+                meta.get_num_dimensions(),
+                meta.get_num_dimensions_to_index(),
+            )
+            .unwrap_or_else(|message| pgrx::error!("{}", message));
+            meta
         }
     }
 

pgvectorscale/src/access_method/mod.rs

@@ -20,6 +20,7 @@ mod storage;
 mod storage_common;
 mod upgrade_test;
 mod vacuum;
+mod vector_type;
 
 /// Access method support function numbers
 pub const DISKANN_DISTANCE_TYPE_PROC: u16 = 1;
@@ -171,6 +172,7 @@ DECLARE
   have_l2_ops int;
   have_ip_ops int;
   have_label_ops int;
+  vector_schema text;
 BEGIN
     -- Has cosine operator class been installed previously?
     SELECT count(*)
@@ -204,11 +206,22 @@ BEGIN
     AND c.opcmethod = (SELECT oid FROM pg_catalog.pg_am am WHERE am.amname = 'diskann')
     AND c.opcnamespace = (SELECT oid FROM pg_catalog.pg_namespace where nspname='@extschema@');
 
+    SELECT n.nspname
+    INTO STRICT vector_schema
+    FROM pg_catalog.pg_extension e
+    JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace
+    WHERE e.extname = 'vector';
+
     IF have_cos_ops = 0 THEN
-        CREATE OPERATOR CLASS vector_cosine_ops DEFAULT
-        FOR TYPE vector USING diskann AS
-	        OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops,
-            FUNCTION 1 distance_type_cosine();
+        EXECUTE pg_catalog.format(
+            $opclass$
+            CREATE OPERATOR CLASS vector_cosine_ops DEFAULT
+            FOR TYPE %1$I.vector USING diskann AS
+                OPERATOR 1 %1$I.<=> (%1$I.vector, %1$I.vector) FOR ORDER BY pg_catalog.float_ops,
+                FUNCTION 1 distance_type_cosine()
+            $opclass$,
+            vector_schema
+        );
     ELSIF have_l2_ops = 0 THEN
         -- Upgrade from 0.4.0 to 0.5.0.  Update cosine opclass to include
         -- the distance_type_cosine function.
@@ -219,17 +232,27 @@ BEGIN
     END IF;
 
     IF have_l2_ops = 0 THEN
-        CREATE OPERATOR CLASS vector_l2_ops
-        FOR TYPE vector USING diskann AS
-            OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops,
-            FUNCTION 1 distance_type_l2();
+        EXECUTE pg_catalog.format(
+            $opclass$
+            CREATE OPERATOR CLASS vector_l2_ops
+            FOR TYPE %1$I.vector USING diskann AS
+                OPERATOR 1 %1$I.<-> (%1$I.vector, %1$I.vector) FOR ORDER BY pg_catalog.float_ops,
+                FUNCTION 1 distance_type_l2()
+            $opclass$,
+            vector_schema
+        );
     END IF;
 
     IF have_ip_ops = 0 THEN
-        CREATE OPERATOR CLASS vector_ip_ops
-        FOR TYPE vector USING diskann AS
-            OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops,
-            FUNCTION 1 distance_type_inner_product();
+        EXECUTE pg_catalog.format(
+            $opclass$
+            CREATE OPERATOR CLASS vector_ip_ops
+            FOR TYPE %1$I.vector USING diskann AS
+                OPERATOR 1 %1$I.<#> (%1$I.vector, %1$I.vector) FOR ORDER BY pg_catalog.float_ops,
+                FUNCTION 1 distance_type_inner_product()
+            $opclass$,
+            vector_schema
+        );
     END IF;
 
     -- First, check if the && operator exists for smallint[]
@@ -275,8 +298,29 @@ $$;
 );
 
 #[pg_guard]
-pub extern "C-unwind" fn amvalidate(_opclassoid: pg_sys::Oid) -> bool {
-    true
+pub extern "C-unwind" fn amvalidate(opclassoid: pg_sys::Oid) -> bool {
+    unsafe {
+        let tup =
+            pg_sys::SearchSysCache1(pg_sys::SysCacheIdentifier::CLAOID as i32, opclassoid.into());
+        if tup.is_null() {
+            return false;
+        }
+        let form = pg_sys::GETSTRUCT(tup) as pg_sys::Form_pg_opclass;
+        let opcname = core::ffi::CStr::from_ptr((*form).opcname.data.as_ptr())
+            .to_bytes()
+            .to_owned();
+        let opcintype = (*form).opcintype;
+        pg_sys::ReleaseSysCache(tup);
+
+        if !matches!(
+            opcname.as_slice(),
+            b"vector_cosine_ops" | b"vector_l2_ops" | b"vector_ip_ops"
+        ) {
+            return true;
+        }
+
+        vector_type::pgvector_vector_base_oid(opcintype) == Some(opcintype)
+    }
 }
 
 /// Implementation of the array overlap operator (&&) for smallint arrays
@@ -321,6 +365,22 @@ pub fn smallint_array_overlap(left: Array<i16>, right: Array<i16>) -> bool {
 mod tests {
     use super::*;
 
+    #[pg_test]
+    fn test_vector_opclasses_validate() -> spi::Result<()> {
+        for opclass in ["vector_cosine_ops", "vector_l2_ops", "vector_ip_ops"] {
+            let oid = Spi::get_one::<pg_sys::Oid>(&format!(
+                "SELECT c.oid
+                 FROM pg_catalog.pg_opclass c
+                 JOIN pg_catalog.pg_am am ON am.oid = c.opcmethod
+                 WHERE am.amname = 'diskann'
+                   AND c.opcname = '{opclass}'"
+            ))?
+            .expect("operator class was not installed");
+            assert!(amvalidate(oid), "{opclass} failed validation");
+        }
+        Ok(())
+    }
+
     #[pg_test]
     fn test_empty_overlap() -> spi::Result<()> {
         // Test overlap with arrays containing only NULL values

pgvectorscale/src/access_method/options.rs

@@ -344,7 +344,7 @@ mod tests {
 
     #[pg_test]
     unsafe fn test_index_options_custom() -> spi::Result<()> {
-        Spi::run("CREATE TABLE test(encoding vector(3));
+        Spi::run("CREATE TABLE test(encoding vector(30));
         CREATE INDEX idxtest
                   ON test
                USING diskann(encoding)
@@ -368,7 +368,7 @@ mod tests {
 
     #[pg_test]
     unsafe fn test_index_options_custom_mem_optimized() -> spi::Result<()> {
-        Spi::run("CREATE TABLE test(encoding vector(3));
+        Spi::run("CREATE TABLE test(encoding vector(30));
         CREATE INDEX idxtest
                   ON test
                USING diskann(encoding)

pgvectorscale/src/access_method/pg_vector.rs

@@ -21,6 +21,10 @@ impl PgVectorInternal {
         let raw_slice = unsafe { self.x.as_slice(dim as _) };
         raw_slice
     }
+
+    pub(crate) unsafe fn reserved(&self) -> i16 {
+        self.unused.assume_init()
+    }
 }
 
 #[derive(Debug)]
@@ -130,6 +134,9 @@ impl PgVector {
         //TODO: we are using a copy here to avoid lifetime issues and because in some cases we have to
         //modify the datum in preprocess_cosine. We should find a way to avoid the copy if the vector is
         //normalized and preprocess_cosine is a noop;
+        //
+        // Callers must have verified the index attribute is pgvector's vector type before reaching
+        // here. Detoast still assumes a well-formed varlena; layout is checked immediately after.
         let detoasted = pg_sys::pg_detoast_datum_copy(datum.cast_mut_ptr());
         let is_copy = !std::ptr::eq(
             detoasted.cast::<PgVectorInternal>(),
@@ -140,15 +147,21 @@ impl PgVector {
         assert!(is_copy, "Datum should be a copy");
         let casted = detoasted.cast::<PgVectorInternal>();
 
+        // Validate the on-disk/layout dimension before any slice is constructed.
+        let dim = super::vector_type::checked_vector_dim(casted)
+            .unwrap_or_else(|message| error!("{}", message));
+        super::vector_type::ensure_datum_dimension(dim, meta_page.get_num_dimensions())
+            .unwrap_or_else(|message| error!("{}", message));
+
         if is_index_distance
             && meta_page.get_num_dimensions() != meta_page.get_num_dimensions_to_index()
         {
             assert!((*casted).dim > meta_page.get_num_dimensions_to_index() as _);
             (*casted).dim = meta_page.get_num_dimensions_to_index() as _;
         }
 
-        let dim = (*casted).dim;
-        let raw_slice = unsafe { (*casted).x.as_mut_slice(dim as _) };
+        let dim = (*casted).dim as usize;
+        let raw_slice = unsafe { (*casted).x.as_mut_slice(dim) };
 
         if meta_page.get_distance_type() == DistanceType::Cosine {
             preprocess_cosine(raw_slice);

pgvectorscale/src/access_method/upgrade_test.rs

@@ -334,4 +334,11 @@ pub mod tests {
     fn test_upgrade_from_0_8_0() {
         test_upgrade_base("0.8.0", "0.12.9", "pgvectorscale", "vectorscale", "diskann");
     }
+
+    #[ignore]
+    #[serial]
+    #[test]
+    fn test_upgrade_from_0_9_0() {
+        test_upgrade_base("0.9.0", "0.16.1", "pgvectorscale", "vectorscale", "diskann");
+    }
 }