diff --git a/pgvectorscale/Cargo.toml b/pgvectorscale/Cargo.toml index fd57cb20..c4a3f13c 100644 --- a/pgvectorscale/Cargo.toml +++ b/pgvectorscale/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "vectorscale" -version = "0.9.0" +version = "0.9.1" edition = "2021" [lib] diff --git a/pgvectorscale/sql/vectorscale--0.9.0--0.9.1.sql b/pgvectorscale/sql/vectorscale--0.9.0--0.9.1.sql new file mode 100644 index 00000000..9c4e3dec --- /dev/null +++ b/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; +$$; diff --git a/pgvectorscale/src/access_method/build.rs b/pgvectorscale/src/access_method/build.rs index 0aea049c..bf7e7cd4 100644 --- a/pgvectorscale/src/access_method/build.rs +++ b/pgvectorscale/src/access_method/build.rs @@ -250,7 +250,9 @@ fn get_meta_page( index_relation: &PgRelation, opt: PgBox, ) -> 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 diff --git a/pgvectorscale/src/access_method/meta_page.rs b/pgvectorscale/src/access_method/meta_page.rs index eecca125..fbb529c0 100644 --- a/pgvectorscale/src/access_method/meta_page.rs +++ b/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 } } diff --git a/pgvectorscale/src/access_method/mod.rs b/pgvectorscale/src/access_method/mod.rs index 42183629..a4b4c32d 100644 --- a/pgvectorscale/src/access_method/mod.rs +++ b/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, right: Array) -> 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::(&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 diff --git a/pgvectorscale/src/access_method/options.rs b/pgvectorscale/src/access_method/options.rs index 2855fd0f..9a23c392 100644 --- a/pgvectorscale/src/access_method/options.rs +++ b/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) diff --git a/pgvectorscale/src/access_method/pg_vector.rs b/pgvectorscale/src/access_method/pg_vector.rs index 5677357c..70e693f2 100644 --- a/pgvectorscale/src/access_method/pg_vector.rs +++ b/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::(), @@ -140,6 +147,12 @@ impl PgVector { assert!(is_copy, "Datum should be a copy"); let casted = detoasted.cast::(); + // 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() { @@ -147,8 +160,8 @@ impl PgVector { (*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); diff --git a/pgvectorscale/src/access_method/upgrade_test.rs b/pgvectorscale/src/access_method/upgrade_test.rs index 0c0b6a57..20e44775 100644 --- a/pgvectorscale/src/access_method/upgrade_test.rs +++ b/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"); + } } diff --git a/pgvectorscale/src/access_method/vector_type.rs b/pgvectorscale/src/access_method/vector_type.rs new file mode 100644 index 00000000..88e7bdfd --- /dev/null +++ b/pgvectorscale/src/access_method/vector_type.rs @@ -0,0 +1,148 @@ +use pgrx::pg_sys::{self, GETSTRUCT}; +use pgrx::varlena::varsize_any; +use pgrx::*; + +use super::build::MAX_DIMENSION; + +/// Return the pgvector-owned `vector` base type for `type_oid`. +/// +/// This verifies extension ownership and the type name without relying on +/// `search_path` or a backend-lifetime cache. +pub unsafe fn pgvector_vector_base_oid(type_oid: pg_sys::Oid) -> Option { + let base_oid = pg_sys::getBaseType(type_oid); + let vector_extension_oid = pg_sys::get_extension_oid(c"vector".as_ptr(), true); + if vector_extension_oid == pg_sys::InvalidOid + || pg_sys::getExtensionOfObject(pg_sys::TypeRelationId, base_oid) != vector_extension_oid + { + return None; + } + + let type_tuple = + pg_sys::SearchSysCache1(pg_sys::SysCacheIdentifier::TYPEOID as i32, base_oid.into()); + if type_tuple.is_null() { + return None; + } + + let type_form = GETSTRUCT(type_tuple) as pg_sys::Form_pg_type; + let is_vector = + core::ffi::CStr::from_ptr((*type_form).typname.data.as_ptr()).to_bytes() == b"vector"; + pg_sys::ReleaseSysCache(type_tuple); + + is_vector.then_some(base_oid) +} + +/// Validate a signed typmod as a vector dimension and return it as u32. +pub fn dimension_from_typmod(atttypmod: i32) -> Result { + if atttypmod < 1 || atttypmod as u32 > MAX_DIMENSION { + return Err(format!( + "diskann: indexed column has no valid vector dimension (atttypmod = {}); the column must be pgvector's vector(N)", + atttypmod + )); + } + Ok(atttypmod as u32) +} + +/// Validate persisted full and indexed dimensions. +pub fn ensure_valid_dimensions( + num_dimensions: u32, + num_dimensions_to_index: u32, +) -> Result<(), String> { + if num_dimensions < 1 || num_dimensions > MAX_DIMENSION { + return Err(format!( + "diskann: invalid full dimension {} (max is {})", + num_dimensions, MAX_DIMENSION + )); + } + if num_dimensions_to_index < 1 || num_dimensions_to_index > num_dimensions { + return Err(format!( + "diskann: num_dimensions={} must be between 1 and {}", + num_dimensions_to_index, num_dimensions + )); + } + Ok(()) +} + +pub fn ensure_datum_dimension(datum_dimension: usize, index_dimension: u32) -> Result<(), String> { + if datum_dimension != index_dimension as usize { + return Err(format!( + "diskann: vector datum dimension {} does not match index dimension {}", + datum_dimension, index_dimension + )); + } + Ok(()) +} + +fn validate_vector_layout(total: usize, dim: i32, reserved: i16) -> Result { + let header_bytes = std::mem::size_of::() + 2 * std::mem::size_of::(); + if total < header_bytes { + return Err("diskann: vector datum shorter than its header".to_string()); + } + if dim < 1 || dim as u32 > MAX_DIMENSION { + return Err(format!("diskann: vector dimension {dim} out of range")); + } + if reserved != 0 { + return Err("diskann: vector datum has non-zero reserved field".to_string()); + } + + let expected = header_bytes + .checked_add( + (dim as usize) + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| "diskann: vector datum size overflow".to_string())?, + ) + .ok_or_else(|| "diskann: vector datum size overflow".to_string())?; + if total != expected { + return Err(format!( + "diskann: vector datum is {total} bytes, dimension {dim} needs {expected}" + )); + } + + Ok(dim as usize) +} + +/// After detoast, require a well-formed pgvector layout before building any slice. +/// +/// Returns the signed dimension as usize. +pub unsafe fn checked_vector_dim( + ptr: *const super::pg_vector::PgVectorInternal, +) -> Result { + let total = varsize_any(ptr.cast()); + let dim = (*ptr).dim as i32; + validate_vector_layout(total, dim, (*ptr).reserved()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_invalid_typmods() { + assert!(dimension_from_typmod(-1).is_err()); + assert!(dimension_from_typmod(0).is_err()); + assert!(dimension_from_typmod((MAX_DIMENSION + 1) as i32).is_err()); + assert_eq!(dimension_from_typmod(3).unwrap(), 3); + } + + #[test] + fn rejects_invalid_metapage_dimensions() { + assert!(ensure_valid_dimensions(0, 0).is_err()); + assert!(ensure_valid_dimensions(MAX_DIMENSION + 1, 1).is_err()); + assert!(ensure_valid_dimensions(3, 0).is_err()); + assert!(ensure_valid_dimensions(3, 4).is_err()); + assert!(ensure_valid_dimensions(3, 3).is_ok()); + } + + #[test] + fn rejects_malformed_vector_layouts() { + assert!(validate_vector_layout(4, 1, 0).is_err()); + assert!(validate_vector_layout(16, 10, 0).is_err()); + assert!(validate_vector_layout(16, 2, 1).is_err()); + assert!(validate_vector_layout(16, 2, 0).is_ok()); + } + + #[test] + fn rejects_datum_dimension_mismatch() { + assert!(ensure_datum_dimension(2, 3).is_err()); + assert!(ensure_datum_dimension(3, 3).is_ok()); + } +} diff --git a/tests/test_security_hardening.py b/tests/test_security_hardening.py new file mode 100644 index 00000000..b57883f6 --- /dev/null +++ b/tests/test_security_hardening.py @@ -0,0 +1,110 @@ +"""Regression tests for DiskANN type and layout hardening.""" + +import os +import uuid + +import psycopg2 +import pytest + + +@pytest.mark.integration +def test_typmodless_vector_index_is_rejected(db_conn): + with db_conn.cursor() as cur: + cur.execute("DROP TABLE IF EXISTS test_typmodless_vector CASCADE") + cur.execute("CREATE TABLE test_typmodless_vector (embedding vector)") + + with pytest.raises( + psycopg2.Error, + match="indexed column has no valid vector dimension", + ): + cur.execute( + "CREATE INDEX ON test_typmodless_vector USING diskann (embedding)" + ) + + cur.execute("DROP TABLE test_typmodless_vector CASCADE") + + +@pytest.mark.integration +def test_schema_shadow_cannot_capture_vector_opclasses(db_connection_params): + database = f"vectorscale_security_{os.getpid()}_{uuid.uuid4().hex[:8]}" + admin = psycopg2.connect(**db_connection_params) + admin.autocommit = True + + try: + with admin.cursor() as cur: + cur.execute(f'CREATE DATABASE "{database}"') + + params = {**db_connection_params, "database": database} + conn = psycopg2.connect(**params) + conn.autocommit = True + try: + with conn.cursor() as cur: + cur.execute( + """ + CREATE SCHEMA real_vector; + CREATE EXTENSION vector WITH SCHEMA real_vector; + CREATE SCHEMA evil; + CREATE DOMAIN evil.vector AS pg_catalog.uuid; + CREATE FUNCTION evil.fake_dist(evil.vector, evil.vector) + RETURNS double precision + LANGUAGE sql IMMUTABLE STRICT PARALLEL SAFE + AS $$ SELECT 0::double precision $$; + CREATE OPERATOR evil.<=> ( + LEFTARG = evil.vector, + RIGHTARG = evil.vector, + FUNCTION = evil.fake_dist, + COMMUTATOR = OPERATOR(evil.<=>) + ); + CREATE OPERATOR evil.<-> ( + LEFTARG = evil.vector, + RIGHTARG = evil.vector, + FUNCTION = evil.fake_dist, + COMMUTATOR = OPERATOR(evil.<->) + ); + CREATE OPERATOR evil.<#> ( + LEFTARG = evil.vector, + RIGHTARG = evil.vector, + FUNCTION = evil.fake_dist, + COMMUTATOR = OPERATOR(evil.<#>) + ); + CREATE EXTENSION vectorscale WITH SCHEMA evil; + """ + ) + + cur.execute( + """ + WITH expected AS ( + SELECT t.oid + 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' + ) + SELECT pg_catalog.bool_and(c.opcintype = expected.oid) + FROM pg_catalog.pg_opclass c + JOIN pg_catalog.pg_am am ON am.oid = c.opcmethod + CROSS JOIN expected + WHERE am.amname = 'diskann' + AND c.opcname IN ( + 'vector_cosine_ops', + 'vector_l2_ops', + 'vector_ip_ops' + ) + """ + ) + assert cur.fetchone()[0] is True + finally: + conn.close() + finally: + with admin.cursor() as cur: + cur.execute( + """ + SELECT pg_catalog.pg_terminate_backend(pid) + FROM pg_catalog.pg_stat_activity + WHERE datname = %s + """, + (database,), + ) + cur.execute(f'DROP DATABASE IF EXISTS "{database}"') + admin.close()