mirror of
https://github.com/logancyang/obsidian-copilot.git
synced 2026-07-22 07:50:24 +00:00
Fix cross platform encryption (#1034)
* Update encryption for backward compatibility * Support backward compatibility * Update all encryption touchpoints to be async
This commit is contained in:
Logan Yang
GitHub
parent
ddcbb23c74
commit
4a0eae53ef
11 changed files with 258 additions and 103 deletions
|
|
@ -94,7 +94,7 @@ export class BrevilabsClient {
|
|||
method,
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
Authorization: `Bearer ${getDecryptedKey(getSettings().plusLicenseKey)}`,
|
||||
Authorization: `Bearer ${await getDecryptedKey(getSettings().plusLicenseKey)}`,
|
||||
},
|
||||
...(method === "POST" && { body: JSON.stringify(body) }),
|
||||
});
|
||||
|
|
|
|||
|
|
@ -53,8 +53,12 @@ export default class ChainManager {
|
|||
this.memoryManager = MemoryManager.getInstance();
|
||||
this.chatModelManager = ChatModelManager.getInstance();
|
||||
this.promptManager = PromptManager.getInstance();
|
||||
this.createChainWithNewModel();
|
||||
subscribeToModelKeyChange(() => this.createChainWithNewModel());
|
||||
|
||||
// Initialize async operations
|
||||
this.initialize();
|
||||
|
||||
// Set up subscriptions
|
||||
subscribeToModelKeyChange(async () => await this.createChainWithNewModel());
|
||||
subscribeToChainTypeChange(() =>
|
||||
this.setChain(getChainType(), {
|
||||
refreshIndex:
|
||||
|
|
@ -63,7 +67,11 @@ export default class ChainManager {
|
|||
getChainType() === ChainType.COPILOT_PLUS_CHAIN),
|
||||
})
|
||||
);
|
||||
subscribeToSettingsChange(() => this.createChainWithNewModel());
|
||||
subscribeToSettingsChange(async () => await this.createChainWithNewModel());
|
||||
}
|
||||
|
||||
private async initialize() {
|
||||
await this.createChainWithNewModel();
|
||||
}
|
||||
|
||||
static getChain(): RunnableSequence {
|
||||
|
|
@ -102,7 +110,7 @@ export default class ChainManager {
|
|||
* Update the active model and create a new chain with the specified model
|
||||
* name.
|
||||
*/
|
||||
createChainWithNewModel(): void {
|
||||
async createChainWithNewModel(): Promise<void> {
|
||||
let newModelKey = getModelKey();
|
||||
try {
|
||||
let customModel = findCustomModel(newModelKey, getSettings().activeModels);
|
||||
|
|
@ -112,7 +120,7 @@ export default class ChainManager {
|
|||
customModel = BUILTIN_CHAT_MODELS[0];
|
||||
newModelKey = customModel.name + "|" + customModel.provider;
|
||||
}
|
||||
this.chatModelManager.setChatModel(customModel);
|
||||
await this.chatModelManager.setChatModel(customModel);
|
||||
// Must update the chatModel for chain because ChainFactory always
|
||||
// retrieves the old chain without the chatModel change if it exists!
|
||||
// Create a new chain with the new chatModel
|
||||
|
|
|
|||
|
|
@ -70,7 +70,7 @@ export default class ChatModelManager {
|
|||
return ChatModelManager.instance;
|
||||
}
|
||||
|
||||
private getModelConfig(customModel: CustomModel): ModelConfig {
|
||||
private async getModelConfig(customModel: CustomModel): Promise<ModelConfig> {
|
||||
const settings = getSettings();
|
||||
|
||||
// Check if the model starts with "o1"
|
||||
|
|
@ -90,7 +90,7 @@ export default class ChatModelManager {
|
|||
} = {
|
||||
[ChatModelProviders.OPENAI]: {
|
||||
modelName: modelName,
|
||||
openAIApiKey: getDecryptedKey(customModel.apiKey || settings.openAIApiKey),
|
||||
openAIApiKey: await getDecryptedKey(customModel.apiKey || settings.openAIApiKey),
|
||||
configuration: {
|
||||
baseURL: customModel.baseUrl,
|
||||
fetch: customModel.enableCors ? safeFetch : undefined,
|
||||
|
|
@ -100,7 +100,7 @@ export default class ChatModelManager {
|
|||
...this.handleOpenAIExtraArgs(isO1Model, settings.maxTokens, settings.temperature),
|
||||
},
|
||||
[ChatModelProviders.ANTHROPIC]: {
|
||||
anthropicApiKey: getDecryptedKey(customModel.apiKey || settings.anthropicApiKey),
|
||||
anthropicApiKey: await getDecryptedKey(customModel.apiKey || settings.anthropicApiKey),
|
||||
modelName: modelName,
|
||||
anthropicApiUrl: customModel.baseUrl,
|
||||
clientOptions: {
|
||||
|
|
@ -110,7 +110,7 @@ export default class ChatModelManager {
|
|||
},
|
||||
},
|
||||
[ChatModelProviders.AZURE_OPENAI]: {
|
||||
azureOpenAIApiKey: getDecryptedKey(customModel.apiKey || settings.azureOpenAIApiKey),
|
||||
azureOpenAIApiKey: await getDecryptedKey(customModel.apiKey || settings.azureOpenAIApiKey),
|
||||
azureOpenAIApiInstanceName: settings.azureOpenAIApiInstanceName,
|
||||
azureOpenAIApiDeploymentName: settings.azureOpenAIApiDeploymentName,
|
||||
azureOpenAIApiVersion: settings.azureOpenAIApiVersion,
|
||||
|
|
@ -121,11 +121,11 @@ export default class ChatModelManager {
|
|||
...this.handleOpenAIExtraArgs(isO1Model, settings.maxTokens, settings.temperature),
|
||||
},
|
||||
[ChatModelProviders.COHEREAI]: {
|
||||
apiKey: getDecryptedKey(customModel.apiKey || settings.cohereApiKey),
|
||||
apiKey: await getDecryptedKey(customModel.apiKey || settings.cohereApiKey),
|
||||
model: modelName,
|
||||
},
|
||||
[ChatModelProviders.GOOGLE]: {
|
||||
apiKey: getDecryptedKey(customModel.apiKey || settings.googleApiKey),
|
||||
apiKey: await getDecryptedKey(customModel.apiKey || settings.googleApiKey),
|
||||
modelName: modelName,
|
||||
safetySettings: [
|
||||
{
|
||||
|
|
@ -149,14 +149,14 @@ export default class ChatModelManager {
|
|||
},
|
||||
[ChatModelProviders.OPENROUTERAI]: {
|
||||
modelName: modelName,
|
||||
openAIApiKey: getDecryptedKey(customModel.apiKey || settings.openRouterAiApiKey),
|
||||
openAIApiKey: await getDecryptedKey(customModel.apiKey || settings.openRouterAiApiKey),
|
||||
configuration: {
|
||||
baseURL: customModel.baseUrl || "https://openrouter.ai/api/v1",
|
||||
fetch: customModel.enableCors ? safeFetch : undefined,
|
||||
},
|
||||
},
|
||||
[ChatModelProviders.GROQ]: {
|
||||
apiKey: getDecryptedKey(customModel.apiKey || settings.groqApiKey),
|
||||
apiKey: await getDecryptedKey(customModel.apiKey || settings.groqApiKey),
|
||||
modelName: modelName,
|
||||
},
|
||||
[ChatModelProviders.OLLAMA]: {
|
||||
|
|
@ -177,7 +177,7 @@ export default class ChatModelManager {
|
|||
},
|
||||
[ChatModelProviders.OPENAI_FORMAT]: {
|
||||
modelName: modelName,
|
||||
openAIApiKey: getDecryptedKey(customModel.apiKey || settings.openAIApiKey),
|
||||
openAIApiKey: await getDecryptedKey(customModel.apiKey || settings.openAIApiKey),
|
||||
configuration: {
|
||||
baseURL: customModel.baseUrl,
|
||||
fetch: customModel.enableCors ? safeFetch : undefined,
|
||||
|
|
@ -251,7 +251,7 @@ export default class ChatModelManager {
|
|||
return ChatModelManager.chatModel;
|
||||
}
|
||||
|
||||
setChatModel(model: CustomModel): void {
|
||||
async setChatModel(model: CustomModel): Promise<void> {
|
||||
const modelKey = `${model.name}|${model.provider}`;
|
||||
if (!ChatModelManager.modelMap.hasOwnProperty(modelKey)) {
|
||||
throw new Error(`No model found for: ${modelKey}`);
|
||||
|
|
@ -266,7 +266,7 @@ export default class ChatModelManager {
|
|||
throw new Error(errorMessage);
|
||||
}
|
||||
|
||||
const modelConfig = this.getModelConfig(model);
|
||||
const modelConfig = await this.getModelConfig(model);
|
||||
|
||||
setModelKey(`${model.name}|${model.provider}`);
|
||||
try {
|
||||
|
|
@ -312,7 +312,7 @@ export default class ChatModelManager {
|
|||
async ping(model: CustomModel): Promise<boolean> {
|
||||
const tryPing = async (enableCors: boolean) => {
|
||||
const modelToTest = { ...model, enableCors };
|
||||
const modelConfig = this.getModelConfig(modelToTest);
|
||||
const modelConfig = await this.getModelConfig(modelToTest);
|
||||
// eslint-disable-next-line @typescript-eslint/no-unused-vars
|
||||
const { streaming, temperature, ...pingConfig } = modelConfig;
|
||||
pingConfig.maxTokens = 10;
|
||||
|
|
|
|||
|
|
@ -131,7 +131,7 @@ export default class EmbeddingManager {
|
|||
})[0];
|
||||
}
|
||||
|
||||
getEmbeddingsAPI(): Embeddings {
|
||||
async getEmbeddingsAPI(): Promise<Embeddings> {
|
||||
const { embeddingModelKey } = getSettings();
|
||||
|
||||
if (!EmbeddingManager.modelMap.hasOwnProperty(embeddingModelKey)) {
|
||||
|
|
@ -146,7 +146,7 @@ export default class EmbeddingManager {
|
|||
}
|
||||
|
||||
const customModel = this.getCustomModel(embeddingModelKey);
|
||||
const config = this.getEmbeddingConfig(customModel);
|
||||
const config = await this.getEmbeddingConfig(customModel);
|
||||
|
||||
try {
|
||||
EmbeddingManager.embeddingModel = new selectedModel.EmbeddingConstructor(config);
|
||||
|
|
@ -158,7 +158,7 @@ export default class EmbeddingManager {
|
|||
}
|
||||
}
|
||||
|
||||
private getEmbeddingConfig(customModel: CustomModel): any {
|
||||
private async getEmbeddingConfig(customModel: CustomModel): Promise<any> {
|
||||
const settings = getSettings();
|
||||
const modelName = customModel.name;
|
||||
|
||||
|
|
@ -187,7 +187,7 @@ export default class EmbeddingManager {
|
|||
} = {
|
||||
[EmbeddingModelProviders.COPILOT_PLUS]: {
|
||||
modelName,
|
||||
apiKey: getDecryptedKey(settings.plusLicenseKey),
|
||||
apiKey: await getDecryptedKey(settings.plusLicenseKey),
|
||||
timeout: 10000,
|
||||
configuration: {
|
||||
baseURL: BREVILABS_API_BASE_URL,
|
||||
|
|
@ -196,7 +196,7 @@ export default class EmbeddingManager {
|
|||
},
|
||||
[EmbeddingModelProviders.COPILOT_PLUS_JINA]: {
|
||||
model: modelName,
|
||||
apiKey: getDecryptedKey(settings.plusLicenseKey),
|
||||
apiKey: await getDecryptedKey(settings.plusLicenseKey),
|
||||
timeout: 10000,
|
||||
batchSize: 128,
|
||||
dimensions: 512,
|
||||
|
|
@ -207,7 +207,7 @@ export default class EmbeddingManager {
|
|||
},
|
||||
[EmbeddingModelProviders.OPENAI]: {
|
||||
modelName,
|
||||
openAIApiKey: getDecryptedKey(customModel.apiKey || settings.openAIApiKey),
|
||||
openAIApiKey: await getDecryptedKey(customModel.apiKey || settings.openAIApiKey),
|
||||
timeout: 10000,
|
||||
configuration: {
|
||||
baseURL: customModel.baseUrl,
|
||||
|
|
@ -216,14 +216,14 @@ export default class EmbeddingManager {
|
|||
},
|
||||
[EmbeddingModelProviders.COHEREAI]: {
|
||||
model: modelName,
|
||||
apiKey: getDecryptedKey(customModel.apiKey || settings.cohereApiKey),
|
||||
apiKey: await getDecryptedKey(customModel.apiKey || settings.cohereApiKey),
|
||||
},
|
||||
[EmbeddingModelProviders.GOOGLE]: {
|
||||
modelName: modelName,
|
||||
apiKey: getDecryptedKey(settings.googleApiKey),
|
||||
apiKey: await getDecryptedKey(settings.googleApiKey),
|
||||
},
|
||||
[EmbeddingModelProviders.AZURE_OPENAI]: {
|
||||
azureOpenAIApiKey: getDecryptedKey(customModel.apiKey || settings.azureOpenAIApiKey),
|
||||
azureOpenAIApiKey: await getDecryptedKey(customModel.apiKey || settings.azureOpenAIApiKey),
|
||||
azureOpenAIApiInstanceName: settings.azureOpenAIApiInstanceName,
|
||||
azureOpenAIApiDeploymentName: settings.azureOpenAIApiEmbeddingDeploymentName,
|
||||
azureOpenAIApiVersion: settings.azureOpenAIApiVersion,
|
||||
|
|
@ -239,7 +239,7 @@ export default class EmbeddingManager {
|
|||
},
|
||||
[EmbeddingModelProviders.LM_STUDIO]: {
|
||||
modelName,
|
||||
openAIApiKey: getDecryptedKey(customModel.apiKey || "default-key"),
|
||||
openAIApiKey: await getDecryptedKey(customModel.apiKey || "default-key"),
|
||||
configuration: {
|
||||
baseURL: customModel.baseUrl || "http://localhost:1234/v1",
|
||||
fetch: customModel.enableCors ? safeFetch : undefined,
|
||||
|
|
@ -247,7 +247,7 @@ export default class EmbeddingManager {
|
|||
},
|
||||
[EmbeddingModelProviders.OPENAI_FORMAT]: {
|
||||
modelName,
|
||||
openAIApiKey: getDecryptedKey(customModel.apiKey || ""),
|
||||
openAIApiKey: await getDecryptedKey(customModel.apiKey || ""),
|
||||
configuration: {
|
||||
baseURL: customModel.baseUrl,
|
||||
fetch: customModel.enableCors ? safeFetch : undefined,
|
||||
|
|
@ -265,7 +265,7 @@ export default class EmbeddingManager {
|
|||
async ping(model: CustomModel): Promise<boolean> {
|
||||
const tryPing = async (enableCors: boolean) => {
|
||||
const modelToTest = { ...model, enableCors };
|
||||
const config = this.getEmbeddingConfig(modelToTest);
|
||||
const config = await this.getEmbeddingConfig(modelToTest);
|
||||
const testModel = new (this.getProviderConstructor(modelToTest))(config);
|
||||
await testModel.embedQuery("test");
|
||||
};
|
||||
|
|
|
|||
|
|
@ -14,10 +14,27 @@ function getSafeStorage() {
|
|||
return safeStorageInternal;
|
||||
}
|
||||
|
||||
// Add new prefixes to distinguish encryption methods
|
||||
const DESKTOP_PREFIX = "enc_desk_";
|
||||
const WEBCRYPTO_PREFIX = "enc_web_";
|
||||
// Keep old prefix for backward compatibility
|
||||
const ENCRYPTION_PREFIX = "enc_";
|
||||
const DECRYPTION_PREFIX = "dec_";
|
||||
|
||||
export function encryptAllKeys(settings: Readonly<CopilotSettings>): Readonly<CopilotSettings> {
|
||||
// Add these constants for the Web Crypto implementation
|
||||
const ENCRYPTION_KEY = new TextEncoder().encode("obsidian-copilot-v1");
|
||||
const ALGORITHM = { name: "AES-GCM", iv: new Uint8Array(12) };
|
||||
|
||||
async function getEncryptionKey(): Promise<CryptoKey> {
|
||||
return await crypto.subtle.importKey("raw", ENCRYPTION_KEY, ALGORITHM.name, false, [
|
||||
"encrypt",
|
||||
"decrypt",
|
||||
]);
|
||||
}
|
||||
|
||||
export async function encryptAllKeys(
|
||||
settings: Readonly<CopilotSettings>
|
||||
): Promise<Readonly<CopilotSettings>> {
|
||||
if (!settings.enableEncryption) {
|
||||
return settings;
|
||||
}
|
||||
|
|
@ -28,20 +45,22 @@ export function encryptAllKeys(settings: Readonly<CopilotSettings>): Readonly<Co
|
|||
|
||||
for (const key of keysToEncrypt) {
|
||||
const apiKey = settings[key as keyof CopilotSettings] as string;
|
||||
(newSettings[key as keyof CopilotSettings] as any) = getEncryptedKey(apiKey);
|
||||
(newSettings[key as keyof CopilotSettings] as any) = await getEncryptedKey(apiKey);
|
||||
}
|
||||
|
||||
if (Array.isArray(settings.activeModels)) {
|
||||
newSettings.activeModels = settings.activeModels.map((model) => ({
|
||||
...model,
|
||||
apiKey: getEncryptedKey(model.apiKey || ""),
|
||||
}));
|
||||
newSettings.activeModels = await Promise.all(
|
||||
settings.activeModels.map(async (model) => ({
|
||||
...model,
|
||||
apiKey: await getEncryptedKey(model.apiKey || ""),
|
||||
}))
|
||||
);
|
||||
}
|
||||
|
||||
return newSettings;
|
||||
}
|
||||
|
||||
export function getEncryptedKey(apiKey: string): string {
|
||||
export async function getEncryptedKey(apiKey: string): Promise<string> {
|
||||
if (!apiKey || apiKey.startsWith(ENCRYPTION_PREFIX)) {
|
||||
return apiKey;
|
||||
}
|
||||
|
|
@ -50,20 +69,25 @@ export function getEncryptedKey(apiKey: string): string {
|
|||
apiKey = apiKey.replace(DECRYPTION_PREFIX, "");
|
||||
}
|
||||
|
||||
if (getSafeStorage() && getSafeStorage().isEncryptionAvailable()) {
|
||||
// Convert the encrypted buffer to a Base64 string and prepend the prefix
|
||||
const encryptedBuffer = getSafeStorage().encryptString(apiKey) as Buffer;
|
||||
// Convert the encrypted buffer to a Base64 string and prepend the prefix
|
||||
return ENCRYPTION_PREFIX + encryptedBuffer.toString("base64");
|
||||
} else {
|
||||
// Simple fallback for mobile (just for demonstration)
|
||||
const encoder = new TextEncoder();
|
||||
const data = encoder.encode(apiKey);
|
||||
return ENCRYPTION_PREFIX + arrayBufferToBase64(data);
|
||||
try {
|
||||
// Try desktop encryption first
|
||||
if (getSafeStorage()?.isEncryptionAvailable()) {
|
||||
const encryptedBuffer = getSafeStorage().encryptString(apiKey) as Buffer;
|
||||
return DESKTOP_PREFIX + encryptedBuffer.toString("base64");
|
||||
}
|
||||
|
||||
// Fallback to Web Crypto API
|
||||
const key = await getEncryptionKey();
|
||||
const encodedData = new TextEncoder().encode(apiKey);
|
||||
const encryptedData = await crypto.subtle.encrypt(ALGORITHM, key, encodedData);
|
||||
return WEBCRYPTO_PREFIX + arrayBufferToBase64(encryptedData);
|
||||
} catch (error) {
|
||||
console.error("Encryption failed:", error);
|
||||
return apiKey;
|
||||
}
|
||||
}
|
||||
|
||||
export function getDecryptedKey(apiKey: string): string {
|
||||
export async function getDecryptedKey(apiKey: string): Promise<string> {
|
||||
if (!apiKey || isPlainText(apiKey)) {
|
||||
return apiKey;
|
||||
}
|
||||
|
|
@ -71,17 +95,41 @@ export function getDecryptedKey(apiKey: string): string {
|
|||
return apiKey.replace(DECRYPTION_PREFIX, "");
|
||||
}
|
||||
|
||||
// Handle different encryption methods
|
||||
if (apiKey.startsWith(DESKTOP_PREFIX)) {
|
||||
const base64Data = apiKey.replace(DESKTOP_PREFIX, "");
|
||||
const buffer = Buffer.from(base64Data, "base64");
|
||||
return getSafeStorage().decryptString(buffer) as string;
|
||||
}
|
||||
|
||||
if (apiKey.startsWith(WEBCRYPTO_PREFIX)) {
|
||||
const base64Data = apiKey.replace(WEBCRYPTO_PREFIX, "");
|
||||
const key = await getEncryptionKey();
|
||||
const encryptedData = base64ToArrayBuffer(base64Data);
|
||||
const decryptedData = await crypto.subtle.decrypt(ALGORITHM, key, encryptedData);
|
||||
return new TextDecoder().decode(decryptedData);
|
||||
}
|
||||
|
||||
// Legacy support for old enc_ prefix
|
||||
const base64Data = apiKey.replace(ENCRYPTION_PREFIX, "");
|
||||
try {
|
||||
if (getSafeStorage() && getSafeStorage().isEncryptionAvailable()) {
|
||||
const buffer = Buffer.from(base64Data, "base64");
|
||||
return getSafeStorage().decryptString(buffer) as string;
|
||||
} else {
|
||||
// Simple fallback for mobile (just for demonstration)
|
||||
const data = base64ToArrayBuffer(base64Data);
|
||||
const decoder = new TextDecoder();
|
||||
return decoder.decode(data);
|
||||
// Try desktop decryption first
|
||||
if (getSafeStorage()?.isEncryptionAvailable()) {
|
||||
try {
|
||||
const buffer = Buffer.from(base64Data, "base64");
|
||||
return getSafeStorage().decryptString(buffer) as string;
|
||||
} catch {
|
||||
// Silent catch is intentional - if desktop decryption fails,
|
||||
// it means this key was likely encrypted with Web Crypto.
|
||||
// We'll fall through to the Web Crypto decryption below.
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to Web Crypto API
|
||||
const key = await getEncryptionKey();
|
||||
const encryptedData = base64ToArrayBuffer(base64Data);
|
||||
const decryptedData = await crypto.subtle.decrypt(ALGORITHM, key, encryptedData);
|
||||
return new TextDecoder().decode(decryptedData);
|
||||
} catch (err) {
|
||||
console.error("Decryption failed:", err);
|
||||
return "Copilot failed to decrypt API keys!";
|
||||
|
|
|
|||
|
|
@ -54,12 +54,12 @@ export default class CopilotPlugin extends Plugin {
|
|||
|
||||
async onload(): Promise<void> {
|
||||
await this.loadSettings();
|
||||
this.settingsUnsubscriber = subscribeToSettingsChange(() => {
|
||||
this.settingsUnsubscriber = subscribeToSettingsChange(async () => {
|
||||
const settings = getSettings();
|
||||
if (settings.enableEncryption) {
|
||||
this.saveData(encryptAllKeys(settings));
|
||||
await this.saveData(await encryptAllKeys(settings));
|
||||
} else {
|
||||
this.saveData(settings);
|
||||
await this.saveData(settings);
|
||||
}
|
||||
registerBuiltInCommands(this);
|
||||
});
|
||||
|
|
|
|||
|
|
@ -45,7 +45,7 @@ export class DBOperations {
|
|||
this.oramaDb = undefined;
|
||||
} else if (Platform.isMobile && !settings.disableIndexOnMobile && !this.oramaDb) {
|
||||
// Re-initialize DB if mobile setting is enabled
|
||||
await this.initializeDB(EmbeddingsManager.getInstance().getEmbeddingsAPI());
|
||||
await this.initializeDB(await EmbeddingsManager.getInstance().getEmbeddingsAPI());
|
||||
}
|
||||
|
||||
// Handle index sync setting change
|
||||
|
|
@ -55,7 +55,7 @@ export class DBOperations {
|
|||
console.log("Path change detected, reinitializing database...");
|
||||
this.dbPath = newPath;
|
||||
await this.initializeChunkedStorage();
|
||||
await this.initializeDB(EmbeddingsManager.getInstance().getEmbeddingsAPI());
|
||||
await this.initializeDB(await EmbeddingsManager.getInstance().getEmbeddingsAPI());
|
||||
console.log("Database reinitialized with new path:", newPath);
|
||||
}
|
||||
});
|
||||
|
|
|
|||
|
|
@ -346,7 +346,8 @@ export class HybridRetriever extends BaseRetriever {
|
|||
}
|
||||
|
||||
private async convertQueryToVector(query: string): Promise<number[]> {
|
||||
const vector = await EmbeddingManager.getInstance().getEmbeddingsAPI().embedQuery(query);
|
||||
const embeddingsAPI = await EmbeddingManager.getInstance().getEmbeddingsAPI();
|
||||
const vector = await embeddingsAPI.embedQuery(query);
|
||||
if (vector.length === 0) {
|
||||
throw new Error("Query embedding returned an empty vector");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -52,7 +52,7 @@ export class IndexOperations {
|
|||
let rateLimitNoticeShown = false;
|
||||
|
||||
try {
|
||||
const embeddingInstance = this.embeddingsManager.getEmbeddingsAPI();
|
||||
const embeddingInstance = await this.embeddingsManager.getEmbeddingsAPI();
|
||||
if (!embeddingInstance) {
|
||||
console.error("Embedding instance not found.");
|
||||
return 0;
|
||||
|
|
@ -169,7 +169,7 @@ export class IndexOperations {
|
|||
fileInfo: any;
|
||||
}>
|
||||
> {
|
||||
const embeddingInstance = this.embeddingsManager.getEmbeddingsAPI();
|
||||
const embeddingInstance = await this.embeddingsManager.getEmbeddingsAPI();
|
||||
if (!embeddingInstance) {
|
||||
console.error("Embedding instance not found.");
|
||||
return [];
|
||||
|
|
@ -444,7 +444,7 @@ export class IndexOperations {
|
|||
|
||||
public async reindexFile(file: TFile): Promise<void> {
|
||||
try {
|
||||
const embeddingInstance = this.embeddingsManager.getEmbeddingsAPI();
|
||||
const embeddingInstance = await this.embeddingsManager.getEmbeddingsAPI();
|
||||
if (!embeddingInstance) {
|
||||
return;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -48,7 +48,7 @@ export default class VectorStoreManager {
|
|||
const oldPath = this.dbOps.getCurrentDbPath();
|
||||
|
||||
if (oldPath !== newPath) {
|
||||
await this.dbOps.initializeDB(this.embeddingsManager.getEmbeddingsAPI());
|
||||
await this.dbOps.initializeDB(await this.embeddingsManager.getEmbeddingsAPI());
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -71,7 +71,7 @@ export default class VectorStoreManager {
|
|||
let retries = 3;
|
||||
while (retries > 0) {
|
||||
try {
|
||||
await this.dbOps.initializeDB(this.embeddingsManager.getEmbeddingsAPI());
|
||||
await this.dbOps.initializeDB(await this.embeddingsManager.getEmbeddingsAPI());
|
||||
break;
|
||||
} catch (error) {
|
||||
if (
|
||||
|
|
@ -84,6 +84,10 @@ export default class VectorStoreManager {
|
|||
continue;
|
||||
}
|
||||
}
|
||||
new Notice(
|
||||
"Failed to initialize vector store. Please make sure you have a valid API key " +
|
||||
"for your embedding model and restart the plugin."
|
||||
);
|
||||
console.error("Failed to initialize vector store:", error);
|
||||
break;
|
||||
}
|
||||
|
|
@ -109,7 +113,7 @@ export default class VectorStoreManager {
|
|||
|
||||
public async clearIndex(): Promise<void> {
|
||||
await this.waitForInitialization();
|
||||
await this.dbOps.clearIndex(this.embeddingsManager.getEmbeddingsAPI());
|
||||
await this.dbOps.clearIndex(await this.embeddingsManager.getEmbeddingsAPI());
|
||||
}
|
||||
|
||||
public async garbageCollectVectorStore(): Promise<number> {
|
||||
|
|
|
|||
|
|
@ -1,20 +1,50 @@
|
|||
import { getDecryptedKey, getEncryptedKey, encryptAllKeys } from "@/encryptionService";
|
||||
import { Platform } from "obsidian";
|
||||
import { type CopilotSettings } from "@/settings/model";
|
||||
import { TextDecoder, TextEncoder } from "util";
|
||||
|
||||
// Mocking Electron's safeStorage
|
||||
jest.mock("electron", () => {
|
||||
return {
|
||||
remote: {
|
||||
safeStorage: {
|
||||
encryptString: jest.fn().mockImplementation((text) => `encrypted_${text}`),
|
||||
decryptString: jest
|
||||
.fn()
|
||||
.mockImplementation((buffer) => buffer.toString().replace("encrypted_", "")),
|
||||
isEncryptionAvailable: jest.fn().mockReturnValue(true),
|
||||
},
|
||||
// Mock electron module with proper types
|
||||
const mockElectron = {
|
||||
remote: {
|
||||
safeStorage: {
|
||||
encryptString: jest.fn().mockImplementation((text) => Buffer.from(`${text}_encrypted`)),
|
||||
decryptString: jest
|
||||
.fn()
|
||||
.mockImplementation((buffer) => buffer.toString().replace("_encrypted", "")),
|
||||
isEncryptionAvailable: jest.fn().mockReturnValue(true),
|
||||
},
|
||||
};
|
||||
},
|
||||
};
|
||||
|
||||
jest.mock("electron", () => mockElectron);
|
||||
|
||||
global.TextEncoder = TextEncoder as any;
|
||||
global.TextDecoder = TextDecoder as any;
|
||||
|
||||
// Now we can import our modules
|
||||
import { encryptAllKeys, getDecryptedKey, getEncryptedKey } from "@/encryptionService";
|
||||
import { type CopilotSettings } from "@/settings/model";
|
||||
import { Platform } from "obsidian";
|
||||
|
||||
// Mock window.btoa and window.atob for base64 encoding/decoding
|
||||
global.btoa = jest.fn().mockImplementation((str) => Buffer.from(str).toString("base64"));
|
||||
global.atob = jest.fn().mockImplementation((str) => Buffer.from(str, "base64").toString());
|
||||
|
||||
const mockSubtle = {
|
||||
importKey: jest.fn().mockResolvedValue("mockCryptoKey"),
|
||||
encrypt: jest.fn().mockImplementation((algorithm, key, data) => {
|
||||
const originalText = new TextDecoder().decode(data);
|
||||
const encryptedText = `${originalText}_encrypted`;
|
||||
return Promise.resolve(new TextEncoder().encode(encryptedText).buffer);
|
||||
}),
|
||||
decrypt: jest.fn().mockImplementation((algorithm, key, data) => {
|
||||
const encryptedText = new TextDecoder().decode(new Uint8Array(data));
|
||||
const originalText = encryptedText.replace("_encrypted", "");
|
||||
return Promise.resolve(new TextEncoder().encode(originalText).buffer);
|
||||
}),
|
||||
};
|
||||
|
||||
// Mock crypto.subtle instead of the entire crypto object
|
||||
Object.defineProperty(global.crypto, "subtle", {
|
||||
value: mockSubtle,
|
||||
configurable: true,
|
||||
});
|
||||
|
||||
describe("Platform-specific Tests", () => {
|
||||
|
|
@ -39,52 +69,61 @@ describe("EncryptionService", () => {
|
|||
});
|
||||
|
||||
describe("getEncryptedKey", () => {
|
||||
it("should encrypt an API key", () => {
|
||||
it("should encrypt an API key", async () => {
|
||||
const apiKey = "testApiKey";
|
||||
const encryptedKey = getEncryptedKey(apiKey);
|
||||
expect(encryptedKey).toBe(`enc_encrypted_${apiKey}`);
|
||||
const encryptedKey = await getEncryptedKey(apiKey);
|
||||
// The key is base64 encoded, so we should expect that format
|
||||
expect(encryptedKey).toMatch(/^enc_(desk|web)_[A-Za-z0-9+/=]+$/);
|
||||
// Verify we can decrypt it back
|
||||
const decryptedKey = await getDecryptedKey(encryptedKey);
|
||||
expect(decryptedKey).toBe(apiKey);
|
||||
});
|
||||
|
||||
it("should return the original key if already encrypted", () => {
|
||||
it("should return the original key if already encrypted", async () => {
|
||||
const apiKey = "enc_testApiKey";
|
||||
const encryptedKey = getEncryptedKey(apiKey);
|
||||
const encryptedKey = await getEncryptedKey(apiKey);
|
||||
expect(encryptedKey).toBe(apiKey);
|
||||
});
|
||||
});
|
||||
|
||||
describe("getDecryptedKey", () => {
|
||||
it("should decrypt an encrypted API key", () => {
|
||||
it("should decrypt an encrypted API key", async () => {
|
||||
const apiKey = "testApiKey";
|
||||
const mockEncryptedKey = `encrypted_${apiKey}`;
|
||||
const base64Encoded = Buffer.from(mockEncryptedKey).toString("base64");
|
||||
const encryptedKey = `enc_${base64Encoded}`;
|
||||
|
||||
const decryptedKey = getDecryptedKey(encryptedKey);
|
||||
const encryptedKey = await getEncryptedKey(apiKey);
|
||||
const decryptedKey = await getDecryptedKey(encryptedKey);
|
||||
expect(decryptedKey).toBe(apiKey);
|
||||
});
|
||||
|
||||
it("should return the original key if it is in plain text", () => {
|
||||
it("should return the original key if it is in plain text", async () => {
|
||||
const apiKey = "testApiKey";
|
||||
const decryptedKey = getDecryptedKey(apiKey);
|
||||
const decryptedKey = await getDecryptedKey(apiKey);
|
||||
expect(decryptedKey).toBe(apiKey);
|
||||
});
|
||||
});
|
||||
|
||||
describe("encryptAllKeys", () => {
|
||||
it('should encrypt all keys containing "apikey"', () => {
|
||||
const newSettings = encryptAllKeys({
|
||||
it("should encrypt all keys containing 'apikey'", async () => {
|
||||
const settings = {
|
||||
enableEncryption: true,
|
||||
openAIApiKey: "testApiKey",
|
||||
cohereApiKey: "anotherTestApiKey",
|
||||
userSystemPrompt: "shouldBeIgnored",
|
||||
} as unknown as CopilotSettings);
|
||||
expect(newSettings.openAIApiKey).toBe("enc_encrypted_testApiKey");
|
||||
expect(newSettings.cohereApiKey).toBe("enc_encrypted_anotherTestApiKey");
|
||||
} as unknown as CopilotSettings;
|
||||
|
||||
const newSettings = await encryptAllKeys(settings);
|
||||
expect(newSettings.openAIApiKey).toMatch(/^enc_(desk|web)_[A-Za-z0-9+/=]+$/);
|
||||
expect(newSettings.cohereApiKey).toMatch(/^enc_(desk|web)_[A-Za-z0-9+/=]+$/);
|
||||
expect(newSettings.userSystemPrompt).toBe("shouldBeIgnored");
|
||||
|
||||
// Verify we can decrypt the keys back
|
||||
const decryptedOpenAI = await getDecryptedKey(newSettings.openAIApiKey);
|
||||
const decryptedCohere = await getDecryptedKey(newSettings.cohereApiKey);
|
||||
expect(decryptedOpenAI).toBe("testApiKey");
|
||||
expect(decryptedCohere).toBe("anotherTestApiKey");
|
||||
});
|
||||
|
||||
it("should not encrypt keys when encryption is not enabled", () => {
|
||||
const newSettings = encryptAllKeys({
|
||||
it("should not encrypt keys when encryption is not enabled", async () => {
|
||||
const newSettings = await encryptAllKeys({
|
||||
enableEncryption: false,
|
||||
openAIApiKey: "testApiKey",
|
||||
cohereApiKey: "anotherTestApiKey",
|
||||
|
|
@ -96,3 +135,58 @@ describe("EncryptionService", () => {
|
|||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("Cross-platform compatibility", () => {
|
||||
let originalConsoleError: typeof console.error;
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
// Save original console.error
|
||||
originalConsoleError = console.error;
|
||||
// Mock console.error to suppress expected encryption fallback messages
|
||||
console.error = jest.fn();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
// Restore original console.error
|
||||
console.error = originalConsoleError;
|
||||
});
|
||||
|
||||
it("should encrypt and decrypt consistently on mobile", async () => {
|
||||
// Mock as mobile by making safeStorage unavailable
|
||||
mockElectron.remote.safeStorage.isEncryptionAvailable.mockReturnValue(false);
|
||||
|
||||
const originalKey = "testApiKey";
|
||||
const encryptedKey = await getEncryptedKey(originalKey);
|
||||
expect(encryptedKey).toMatch(/^enc_(desk|web)_[A-Za-z0-9+/=]+$/);
|
||||
|
||||
// Reset the mock counts before decryption
|
||||
mockSubtle.encrypt.mockClear();
|
||||
mockSubtle.decrypt.mockClear();
|
||||
|
||||
const decryptedKey = await getDecryptedKey(encryptedKey);
|
||||
expect(decryptedKey).toBe(originalKey);
|
||||
|
||||
// On mobile, we should use Web Crypto API for decryption
|
||||
expect(mockSubtle.decrypt).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("should be able to decrypt mobile-encrypted keys on desktop", async () => {
|
||||
// First encrypt on mobile
|
||||
mockElectron.remote.safeStorage.isEncryptionAvailable.mockReturnValue(false);
|
||||
|
||||
const originalKey = "testApiKey";
|
||||
const mobileEncryptedKey = await getEncryptedKey(originalKey);
|
||||
expect(mobileEncryptedKey).toMatch(/^enc_(desk|web)_[A-Za-z0-9+/=]+$/);
|
||||
expect(mockSubtle.encrypt).toHaveBeenCalled();
|
||||
|
||||
// Reset the mock counts before desktop decryption
|
||||
mockSubtle.encrypt.mockClear();
|
||||
mockSubtle.decrypt.mockClear();
|
||||
|
||||
// Then decrypt on desktop
|
||||
mockElectron.remote.safeStorage.isEncryptionAvailable.mockReturnValue(true);
|
||||
const decryptedKey = await getDecryptedKey(mobileEncryptedKey);
|
||||
expect(decryptedKey).toBe(originalKey);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
Loading…
Reference in a new issue