From 6a2abb5ebab8d475156e4fccb3784a6a661d8fb5 Mon Sep 17 00:00:00 2001
From: Derek Sonnenberg <derek@ardrive.io>
Date: Mon, 20 May 2024 14:15:33 -0500
Subject: [PATCH] refactor(ardrive): re-add public assertion arfsdao method
 PE-6155

---
 src/arfs/arfsdao.ts | 62 +++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 60 insertions(+), 2 deletions(-)

diff --git a/src/arfs/arfsdao.ts b/src/arfs/arfsdao.ts
index 87917321..b3685020 100644
--- a/src/arfs/arfsdao.ts
+++ b/src/arfs/arfsdao.ts
@@ -70,7 +70,7 @@ import {
 	ArFSPublicFolderCacheKey,
 	defaultArFSAnonymousCache
 } from './arfsdao_anonymous';
-import { deriveDriveKey, deriveFileKey } from '../utils/crypto';
+import { deriveDriveKey, deriveFileKey, driveDecrypt } from '../utils/crypto';
 import {
 	DEFAULT_APP_NAME,
 	DEFAULT_APP_VERSION,
@@ -175,7 +175,7 @@ import {
 } from './tx/arfs_tx_data_types';
 import { ArFSTagAssembler } from './tags/tag_assembler';
 import { assertDataRootsMatch, rePrepareV2Tx } from '../utils/arfsdao_utils';
-import { ArFSDataToUpload, ArFSFolderToUpload } from '../exports';
+import { ArFSDataToUpload, ArFSFolderToUpload, DrivePrivacy } from '../exports';
 import { Turbo, TurboCachesResponse } from './turbo';
 import { ArweaveSigner } from 'arbundles/src/signing';
 
@@ -1757,6 +1757,64 @@ export class ArFSDAO extends ArFSDAOAnonymous {
 		);
 	}
 
+	public async getOwnerAndAssertDrive(driveId: DriveID, driveKey?: DriveKey): Promise<ArweaveAddress> {
+		const cachedOwner = this.caches.ownerCache.get(driveId);
+		if (cachedOwner) {
+			return cachedOwner;
+		}
+
+		return this.caches.ownerCache.put(
+			driveId,
+			(async () => {
+				const gqlQuery = buildQuery({
+					tags: [
+						{ name: 'Entity-Type', value: 'drive' },
+						{ name: 'Drive-Id', value: `${driveId}` }
+					],
+					sort: ASCENDING_ORDER
+				});
+				const transactions = await this.gatewayApi.gqlRequest(gqlQuery);
+				const edges: GQLEdgeInterface[] = transactions.edges;
+
+				if (!edges.length) {
+					throw new Error(`Could not find a transaction with "Drive-Id": ${driveId}`);
+				}
+
+				const edgeOfFirstDrive = edges[0];
+				const driveOwnerAddress = edgeOfFirstDrive.node.owner.address;
+				const driveOwner = new ArweaveAddress(driveOwnerAddress);
+
+				const drivePrivacy: DrivePrivacy = driveKey ? 'private' : 'public';
+				const drivePrivacyFromTag = edgeOfFirstDrive.node.tags.find((t) => t.name === 'Drive-Privacy');
+
+				if (!drivePrivacyFromTag) {
+					throw new Error('Target drive has no "Drive-Privacy" tag!');
+				}
+
+				if (drivePrivacyFromTag.value !== drivePrivacy) {
+					throw new Error(`Target drive is not a ${drivePrivacy} drive!`);
+				}
+
+				if (driveKey) {
+					const cipherIVFromTag = edgeOfFirstDrive.node.tags.find((t) => t.name === 'Cipher-IV');
+					if (!cipherIVFromTag) {
+						throw new Error('Target private drive has no "Cipher-IV" tag!');
+					}
+
+					const driveDataBuffer = await this.gatewayApi.getTxData(TxID(edgeOfFirstDrive.node.id));
+					try {
+						// Attempt to decrypt drive to assert drive key is correct
+						await driveDecrypt(cipherIVFromTag.value, driveKey, driveDataBuffer);
+					} catch {
+						throw new Error('Provided drive key or password could not decrypt target private drive!');
+					}
+				}
+
+				return driveOwner;
+			})()
+		);
+	}
+
 	/**
 	 * Lists the children of certain private folder
 	 * @param {FolderID} folderId the folder ID to list children of