ensured data sources can be renamed and made it secure

This commit is contained in:
Paul Koudelka 2026-08-11 15:22:07 +02:00
parent 202e37410b
commit 92e986d9b9
13 changed files with 137 additions and 66 deletions

View File

@ -10030,6 +10030,9 @@ UI_TEXT_CONTENT["AISTUDIO::TOOLS::VALIDATION::DATASOURCEVALIDATION::T4001510395"
-- Please select a compliance level.
UI_TEXT_CONTENT["AISTUDIO::TOOLS::VALIDATION::DATASOURCEVALIDATION::T4066952091"] = "Please select a compliance level."
-- The name must not contain control characters.
UI_TEXT_CONTENT["AISTUDIO::TOOLS::VALIDATION::DATASOURCEVALIDATION::T4234589878"] = "The name must not contain control characters."
-- The selected embedding provider has confidence '{0}', but this data source requires provider confidence '{1}'. Select an embedding provider with equal or higher confidence or lower the compliance level.
UI_TEXT_CONTENT["AISTUDIO::TOOLS::VALIDATION::DATASOURCEVALIDATION::T476537662"] = "The selected embedding provider has confidence '{0}', but this data source requires provider confidence '{1}'. Select an embedding provider with equal or higher confidence or lower the compliance level."

View File

@ -168,9 +168,9 @@ public readonly record struct DataSourceERI_V1 : IERIDataSource
return false;
}
if (!table.TryGetValue("Name", out var nameValue) || !nameValue.TryRead<string>(out var name) || string.IsNullOrWhiteSpace(name))
if (!table.TryGetValue("Name", out var nameValue) || !nameValue.TryRead<string>(out var name) || string.IsNullOrWhiteSpace(name) || name.Length > 40 || name.Any(char.IsControl))
{
LOGGER.LogWarning($"The configured data source {idx} does not contain a valid name. (Plugin ID: {configPluginId})");
LOGGER.LogWarning($"The configured data source {idx} does not contain a valid name of at most 40 characters without control characters. (Plugin ID: {configPluginId})");
return false;
}
@ -406,4 +406,4 @@ public readonly record struct DataSourceERI_V1 : IERIDataSource
var cleanedHostname = hostname.Trim();
return cleanedHostname.EndsWith('/') ? cleanedHostname[..^1] : cleanedHostname;
}
}
}

View File

@ -22,7 +22,7 @@ public sealed class NoVectorStoreClient(string name, string? unavailableReason,
await Task.CompletedTask;
}
public override Task EnsureVectorStoreExists(string storeName, int vectorSize, CancellationToken token) =>
public override Task EnsureVectorStoreExists(string storeName, string dataSourceName, int vectorSize, CancellationToken token) =>
Task.FromException(this.CreateUnavailableException());
public override Task InsertEmbedding(string storeName, IReadOnlyList<VectorStoragePoint> points, CancellationToken token) =>

View File

@ -82,8 +82,8 @@ public sealed class QdrantEdgeClientImplementation(
yield return (TB("Number of vector stores"), displayStoresCount.ToString());
}
public override Task EnsureVectorStoreExists(string storeName, int vectorSize, CancellationToken token) =>
rustService.ExecuteDatabaseOperation(DATABASE_NAME, ENSURE_PATH, new EnsureVectorStoreRequest(storeName, vectorSize), token);
public override Task EnsureVectorStoreExists(string storeName, string dataSourceName, int vectorSize, CancellationToken token) =>
rustService.ExecuteDatabaseOperation(DATABASE_NAME, ENSURE_PATH, new EnsureVectorStoreRequest(storeName, dataSourceName, vectorSize), token);
public override Task InsertEmbedding(string storeName, IReadOnlyList<VectorStoragePoint> points, CancellationToken token) =>
rustService.ExecuteDatabaseOperation(DATABASE_NAME, INSERT_PATH, new InsertEmbeddingRequest(storeName, points), token);
@ -121,7 +121,7 @@ public sealed class QdrantEdgeClientImplementation(
}
// ReSharper disable NotAccessedPositionalProperty.Local
private sealed record EnsureVectorStoreRequest(string StoreName, int VectorSize);
private sealed record EnsureVectorStoreRequest(string StoreName, string DataSourceName, int VectorSize);
private sealed record InsertEmbeddingRequest(string StoreName, IReadOnlyList<VectorStoragePoint> Points);

View File

@ -2,7 +2,7 @@
public abstract class VectorStoreClient(string name, string path): DatabaseClient(name, path)
{
public abstract Task EnsureVectorStoreExists(string storeName, int vectorSize, CancellationToken token);
public abstract Task EnsureVectorStoreExists(string storeName, string dataSourceName, int vectorSize, CancellationToken token);
public abstract Task InsertEmbedding(string storeName, IReadOnlyList<VectorStoragePoint> points, CancellationToken token);

View File

@ -2,20 +2,11 @@ namespace AIStudio.Tools.Services;
internal static class DataSourceEmbeddingNames
{
public static string GetCollectionName(string dataSourceName, string dataSourceId)
public static string GetCollectionName(string dataSourceId)
{
var safeId = dataSourceId
.ToLowerInvariant()
.Replace("-", string.Empty, StringComparison.Ordinal);
if (!Guid.TryParse(dataSourceId, out var parsedDataSourceId))
throw new ArgumentException("Data source ID must be a valid GUID.", nameof(dataSourceId));
var safeName = new string(dataSourceName
.ToLowerInvariant()
.Where(c => c is >= 'a' and <= 'z' or >= '0' and <= '9')
.Take(32)
.ToArray());
safeName = string.IsNullOrWhiteSpace(safeName) ? "datasource" : safeName;
return $"rag_{safeName}_{safeId}";
return $"rag_{parsedDataSourceId:N}";
}
}

View File

@ -1158,8 +1158,8 @@ public sealed partial class DataSourceEmbeddingService
: null;
}
private string GetCollectionName(string dataSourceName, string dataSourceId) =>
DataSourceEmbeddingNames.GetCollectionName(dataSourceName, dataSourceId);
private string GetCollectionName(string dataSourceId) =>
DataSourceEmbeddingNames.GetCollectionName(dataSourceId);
private string CreatePointId(string dataSourceId, string fingerprint, int chunkIndex) =>
CreateStableGuid($"{dataSourceId}:chunk:{fingerprint}:{chunkIndex}");

View File

@ -6,13 +6,12 @@ namespace AIStudio.Tools.Services;
public sealed partial class DataSourceEmbeddingService
{
private async Task ResetPersistedStateAsync(
string dataSourceName,
string dataSourceId,
VectorStoreClient? vectorStore,
EmbeddingStateClient? embeddingState,
CancellationToken token)
{
await this.DeleteCollectionAsync(this.GetCollectionName(dataSourceName, dataSourceId), vectorStore, token);
await this.DeleteCollectionAsync(this.GetCollectionName(dataSourceId), vectorStore, token);
embeddingState ??= await databaseClientProvider.GetEmbeddingStateAsync(token);
if (!embeddingState.IsAvailable)

View File

@ -285,7 +285,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
}
this.statuses.TryRemove(dataSource.Id, out _);
await this.ResetPersistedStateAsync(dataSource.Name, dataSource.Id, null, null, CancellationToken.None);
await this.ResetPersistedStateAsync(dataSource.Id, null, null, CancellationToken.None);
this.statuses.TryRemove(dataSource.Id, out _);
this.PublishStatusChanged();
}
@ -460,13 +460,26 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
dataSource.Name,
dataSource.Id);
var collectionName = this.GetCollectionName(dataSource.Name, dataSource.Id);
var collectionName = this.GetCollectionName(dataSource.Id);
var manifest = await this.EnsureCompatibleManifestAsync(dataSource, embeddingProvider, collectionName, vectorStore, embeddingState, token);
token.ThrowIfCancellationRequested();
if (manifest.VectorSize > 0)
await this.EnsureCollectionExistsAsync(vectorStore, collectionName, dataSource.Name, manifest.VectorSize, token);
var inputFiles = this.GetInputFiles(dataSource);
var indexedFiles = inputFiles.Files;
var totalFiles = indexedFiles.Count + inputFiles.FailedFiles;
foreach (var failure in inputFiles.Failures)
{
logger.LogWarning(
"Cannot index data source input '{FilePath}' for data source '{DataSourceName}' ({DataSourceId}). Reason='{Reason}'.",
failure.FilePath,
dataSource.Name,
dataSource.Id,
failure.Reason);
}
logger.LogInformation(
"Prepared data source '{DataSourceName}' ({DataSourceId}) for embedding. AccessibleFiles={AccessibleFiles}, FailedFiles={FailedFiles}, Collection='{CollectionName}'.",
dataSource.Name,
@ -659,7 +672,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
VectorStoreOptimizationTracker optimizationTracker,
CancellationToken token)
{
var collectionName = this.GetCollectionName(dataSource.Name, dataSource.Id);
var collectionName = this.GetCollectionName(dataSource.Id);
logger.LogDebug(
"Resetting stored embeddings for file '{FilePath}' in collection '{CollectionName}' before re-indexing.",
file.FullName,
@ -757,7 +770,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
if (manifest.VectorSize == 0)
{
token.ThrowIfCancellationRequested();
await this.EnsureCollectionExistsAsync(vectorStore, collectionName, vectorSize, token);
await this.EnsureCollectionExistsAsync(vectorStore, collectionName, dataSource.Name, vectorSize, token);
await embeddingState.UpdateVectorSizeAsync(dataSource.Id, vectorSize, token);
manifest.VectorSize = vectorSize;
logger.LogInformation(
@ -806,9 +819,9 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
batch.Clear();
}
private async Task EnsureCollectionExistsAsync(VectorStoreClient vectorStore, string collectionName, int vectorSize, CancellationToken token)
private async Task EnsureCollectionExistsAsync(VectorStoreClient vectorStore, string collectionName, string dataSourceName, int vectorSize, CancellationToken token)
{
await vectorStore.EnsureVectorStoreExists(collectionName, vectorSize, token);
await vectorStore.EnsureVectorStoreExists(collectionName, dataSourceName, vectorSize, token);
}
private async Task UpsertPointsAsync(
@ -1061,7 +1074,7 @@ public sealed partial class DataSourceEmbeddingService(SettingsManager settingsM
dataSource.Id,
manifest.EmbeddingSignature,
embeddingSignature);
await this.ResetPersistedStateAsync(dataSource.Name, dataSource.Id, vectorStore, embeddingState, token);
await this.ResetPersistedStateAsync(dataSource.Id, vectorStore, embeddingState, token);
manifest = await embeddingState.GetManifestAsync(dataSource.Id, token);
}

View File

@ -59,7 +59,7 @@ public sealed class DataSourceLocalRetrievalService(
if (maxMatches == 0)
return [];
var collectionName = DataSourceEmbeddingNames.GetCollectionName(dataSource.Name, dataSource.Id);
var collectionName = DataSourceEmbeddingNames.GetCollectionName(dataSource.Id);
var vectorTask = this.SearchVectorAsync(dataSource, query, maxMatches, collectionName, token);
var bm25Task = this.SearchBm25Async(dataSource, query, maxMatches, token);

View File

@ -121,6 +121,9 @@ public sealed class DataSourceValidation
if (dataSourceName.Length > 40)
return TB("The name must not exceed 40 characters.");
if (dataSourceName.Any(char.IsControl))
return TB("The name must not contain control characters.");
var lowerName = dataSourceName.ToLowerInvariant();
if(lowerName != this.GetPreviousDataSourceName() && this.GetUsedDataSourceNames().Contains(lowerName))

View File

@ -15,7 +15,6 @@ use qdrant_edge::{
UpdateOperation, ValueVariants, VectorInternal, Vectors, WithPayloadInterface, WithVector,
};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use tauri::Manager;
use crate::api_token::APIToken;
@ -30,6 +29,8 @@ const HNSW_MAX_INDEXING_THREADS: usize = 0;
const VECTOR_INDEXING_THRESHOLD_KB: usize = 10_000;
const STORE_INITIALIZATION_MARKER: &str = "store_name.txt";
const STORE_INITIALIZATION_MARKER_TEMP: &str = "store_name.tmp";
const STORE_DISPLAY_NAME_MARKER: &str = "data_source_name.txt";
const STORE_DISPLAY_NAME_MARKER_TEMP: &str = "data_source_name.tmp";
type QdrantEdgeResult<T> = Result<T, Box<dyn std::error::Error + Send + Sync>>;
@ -92,6 +93,7 @@ pub struct QdrantEdgeStoragePoint {
#[derive(Deserialize)]
pub struct EnsureQdrantEdgeStoreRequest {
pub store_name: String,
pub data_source_name: String,
pub vector_size: usize,
}
@ -287,9 +289,12 @@ impl QdrantEdgeDatabase {
})
}
fn ensure_store_exists(&mut self, store_name: &str, vector_size: usize) -> QdrantEdgeResult<()> {
fn ensure_store_exists(&mut self, store_name: &str, data_source_name: &str, vector_size: usize) -> QdrantEdgeResult<()> {
validate_vector_size(vector_size)?;
validate_data_source_name(data_source_name)?;
let store_path = self.store_path(store_name)?;
self.get_or_create_store(store_name, vector_size)?;
write_store_display_name(&store_path, data_source_name)?;
Ok(())
}
@ -304,11 +309,19 @@ impl QdrantEdgeDatabase {
return Err("All vectors in one insert request must have the same size.".into());
}
let data_source_name = first_point.data_source_name.clone();
validate_data_source_name(&data_source_name)?;
if points.iter().any(|point| point.data_source_name != data_source_name) {
return Err("All points in one insert request must belong to the same data source name.".into());
}
let store_path = self.store_path(store_name)?;
let shard = self.get_or_create_store(store_name, vector_size)?;
write_store_display_name(&store_path, &data_source_name)?;
let points = points
.into_iter()
.map(to_qdrant_edge_point)
.collect::<Vec<_>>();
.collect::<QdrantEdgeResult<Vec<_>>>()?;
shard.update(UpdateOperation::PointOperation(
PointOperations::UpsertPoints(PointInsertOperations::PointsList(points)),
@ -389,13 +402,8 @@ impl QdrantEdgeDatabase {
}
fn store_directory_name(store_name: &str) -> String {
// Qdrant creates deeply nested files, so keep the physical path short on Windows.
let digest = Sha256::digest(store_name.as_bytes());
let short_hash = digest[..12]
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>();
format!("store_{short_hash}")
let stable_id = store_name.strip_prefix("rag_").unwrap_or(store_name);
format!("store_{stable_id}")
}
fn qdrant_edge_base_path() -> QdrantEdgeResult<PathBuf> {
@ -433,7 +441,7 @@ pub async fn qdrant_edge_info(_token: APIToken) -> Json<QdrantEdgeServiceInfo> {
pub async fn ensure_qdrant_edge_store(_token: APIToken, Json(request): Json<EnsureQdrantEdgeStoreRequest>) -> Json<QdrantEdgeOperationResponse> {
execute_qdrant_edge_operation(|database| {
database.ensure_store_exists(&request.store_name, request.vector_size)
database.ensure_store_exists(&request.store_name, &request.data_source_name, request.vector_size)
})
}
@ -702,6 +710,20 @@ fn write_store_initialization_marker(path: &Path, store_name: &str) -> std::io::
fs::rename(temporary_marker_path, marker_path)
}
fn write_store_display_name(path: &Path, data_source_name: &str) -> std::io::Result<()> {
let marker_path = path.join(STORE_DISPLAY_NAME_MARKER);
if fs::read_to_string(&marker_path).is_ok_and(|current_name| current_name == data_source_name) {
return Ok(());
}
let temporary_marker_path = path.join(STORE_DISPLAY_NAME_MARKER_TEMP);
fs::write(&temporary_marker_path, data_source_name)?;
if marker_path.exists() {
fs::remove_file(&marker_path)?;
}
fs::rename(temporary_marker_path, marker_path)
}
fn remove_partial_store(path: &Path) -> String {
match fs::remove_dir_all(path) {
Ok(()) => String::new(),
@ -717,6 +739,24 @@ fn validate_vector_size(vector_size: usize) -> QdrantEdgeResult<()> {
Ok(())
}
fn validate_data_source_name(data_source_name: &str) -> QdrantEdgeResult<()> {
const MAX_DATA_SOURCE_NAME_LENGTH: usize = 40;
if data_source_name.trim().is_empty() {
return Err("Data source name cannot be empty.".into());
}
if data_source_name.chars().count() > MAX_DATA_SOURCE_NAME_LENGTH {
return Err(format!("Data source name exceeds the maximum length of {MAX_DATA_SOURCE_NAME_LENGTH} characters.").into());
}
if data_source_name.chars().any(|c| c.is_control()) {
return Err("Data source name contains unsupported control characters.".into());
}
Ok(())
}
fn vector_store_version() -> QdrantEdgeResult<String> {
let metadata = META_DATA
.lock()
@ -728,9 +768,9 @@ fn vector_store_version() -> QdrantEdgeResult<String> {
Ok(metadata.vector_store_version.clone())
}
fn to_qdrant_edge_point(point: QdrantEdgeStoragePoint) -> qdrant_edge::PointStructPersisted {
PointStruct::new(
to_point_id(&point.point_id),
fn to_qdrant_edge_point(point: QdrantEdgeStoragePoint) -> QdrantEdgeResult<qdrant_edge::PointStructPersisted> {
Ok(PointStruct::new(
to_point_id(&point.point_id)?,
Vectors::new_named([(VECTOR_NAME, point.vector)]),
json!({
"data_source_id": point.data_source_id,
@ -754,7 +794,7 @@ fn to_qdrant_edge_point(point: QdrantEdgeStoragePoint) -> qdrant_edge::PointStru
"compliance_level_rank": point.compliance_level_rank,
}),
)
.into()
.into())
}
fn to_qdrant_edge_search_result(point: ScoredPoint) -> QdrantEdgeSearchResult {
@ -784,10 +824,10 @@ fn to_qdrant_edge_search_result(point: ScoredPoint) -> QdrantEdgeSearchResult {
}
}
fn to_point_id(point_id: &str) -> PointId {
fn to_point_id(point_id: &str) -> QdrantEdgeResult<PointId> {
Uuid::parse_str(point_id)
.map(PointId::Uuid)
.unwrap_or_else(|_| PointId::NumId(stable_u64(point_id)))
.map_err(|_| "Vector point ID must be a valid UUID.".into())
}
fn point_id_to_string(point_id: PointId) -> String {
@ -814,16 +854,6 @@ fn payload_i32(payload: &Payload, key: &str) -> Option<i32> {
.and_then(|value| i32::try_from(value).ok())
}
fn stable_u64(value: &str) -> u64 {
let mut hash = 0xcbf29ce484222325_u64;
for byte in value.as_bytes() {
hash ^= u64::from(*byte);
hash = hash.wrapping_mul(0x100000001b3);
}
hash
}
fn match_keyword_filter(field_name: &str, value: &str) -> QdrantEdgeResult<Filter> {
Ok(Filter {
should: None,
@ -841,17 +871,19 @@ fn match_keyword_filter(field_name: &str, value: &str) -> QdrantEdgeResult<Filte
}
fn validate_store_name(store_name: &str) -> QdrantEdgeResult<()> {
const MAX_STORE_NAME_LENGTH: usize = 128;
if store_name.is_empty() {
return Err("Vector store name cannot be empty.".into());
}
if matches!(store_name, "." | "..") {
return Err(format!("Vector store name '{store_name}' is not supported.").into());
if store_name.len() > MAX_STORE_NAME_LENGTH {
return Err(format!("Vector store name exceeds the maximum length of {MAX_STORE_NAME_LENGTH} bytes.").into());
}
if store_name
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '.')
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
{
return Ok(());
}
@ -865,12 +897,41 @@ mod tests {
#[test]
fn validate_store_name_allows_safe_store_names() {
assert!(validate_store_name("rag_1234-abcd.ef").is_ok());
assert!(validate_store_name("rag_1234-abcd").is_ok());
}
#[test]
fn validate_store_name_rejects_path_traversal_names() {
fn validate_store_name_rejects_path_syntax() {
assert!(validate_store_name(".").is_err());
assert!(validate_store_name("..").is_err());
assert!(validate_store_name("../store").is_err());
assert!(validate_store_name("store\\name").is_err());
}
#[test]
fn validate_store_name_rejects_oversized_names() {
assert!(validate_store_name(&"a".repeat(129)).is_err());
}
#[test]
fn store_directory_name_contains_the_stable_data_source_id() {
assert_eq!(
store_directory_name("rag_6cc665a82b1e4d42bc748015b7b391ec"),
"store_6cc665a82b1e4d42bc748015b7b391ec"
);
}
#[test]
fn validate_data_source_name_allows_display_names_but_rejects_invalid_values() {
assert!(validate_data_source_name("Mäßig Confidence C#").is_ok());
assert!(validate_data_source_name(" ").is_err());
assert!(validate_data_source_name("invalid\nname").is_err());
assert!(validate_data_source_name(&"a".repeat(41)).is_err());
}
#[test]
fn point_ids_must_be_valid_uuids() {
assert!(to_point_id("6cc665a8-2b1e-4d42-bc74-8015b7b391ec").is_ok());
assert!(to_point_id("deliberate-collision-input").is_err());
}
}

View File

@ -28,7 +28,8 @@
],
"resources": [
"resources/libraries/*",
"resources/notices/*"
"resources/notices/*",
"resources/tokenizers/*"
],
"macOS": {
"exceptionDomain": "localhost"