feat: encryption and decryption schema typing

FedericoAmura · FedericoAmura · commit d156ab6366d3 · 2024-10-16T16:59:01.000+02:00
diff --git a/packages/lit-node-client-nodejs/src/lib/lit-node-client-nodejs.ts b/packages/lit-node-client-nodejs/src/lib/lit-node-client-nodejs.ts
@@ -64,11 +64,11 @@ import {
 } from '@lit-protocol/misc-browser';
 import { nacl } from '@lit-protocol/nacl';
 import { LitNodeClientConfigSchema } from '@lit-protocol/schemas';
+import { ILitResource, ISessionCapabilityObject } from '@lit-protocol/types';
 import {
   uint8arrayFromString,
   uint8arrayToString,
 } from '@lit-protocol/uint8arrays';
-import { ILitResource, ISessionCapabilityObject } from '@lit-protocol/types';
 
 import { encodeCode } from './helpers/encode-code';
 import { getBlsSignatures } from './helpers/get-bls-signatures';
@@ -1441,8 +1441,7 @@ export class LitNodeClientNodeJs
    *
    */
   decrypt = async (params: DecryptRequest): Promise<DecryptResponse> => {
-    const { sessionSigs, authSig, chain, ciphertext, dataToEncryptHash } =
-      params;
+    const { chain, ciphertext, dataToEncryptHash } = params;
 
     // ========== Validate Params ==========
     // -- validate if it's ready
@@ -1527,7 +1526,8 @@ export class LitNodeClientNodeJs
     const requestId = this._getNewRequestId();
     const nodePromises = this.getNodePromises((url: string) => {
       // -- if session key is available, use it
-      const authSigToSend = sessionSigs ? sessionSigs[url] : authSig;
+      const authSigToSend =
+        'sessionSigs' in params ? params.sessionSigs[url] : params.authSig;
 
       if (!authSigToSend) {
         throw new InvalidArgumentException(
@@ -1794,7 +1794,7 @@ export class LitNodeClientNodeJs
 
     // ========== Extract shares from response data ==========
     // -- 1. combine signed data as a list, and get the signatures from it
-    let curveType = responseData[0]?.curveType;
+    const curveType = responseData[0]?.curveType;
 
     if (curveType === 'ECDSA') {
       throw new Error(
diff --git a/packages/misc/src/lib/params-validators.ts b/packages/misc/src/lib/params-validators.ts
@@ -30,9 +30,9 @@ import {
   EvmContractConditions,
   GetSignedTokenRequest,
   JsonExecutionSdkParams,
-  SessionSigsOrAuthSig,
   SolRpcConditions,
   UnifiedAccessControlConditions,
+  ChainedSessionSigsOrAuthSig,
 } from '@lit-protocol/types';
 
 import { checkIfAuthSigRequiresChainParam, checkType, is, log } from './misc';
@@ -456,9 +456,7 @@ class FileValidator implements ParamsValidator {
   }
 }
 
-export interface AuthMaterialValidatorProps extends SessionSigsOrAuthSig {
-  chain?: string;
-}
+export type AuthMaterialValidatorProps = ChainedSessionSigsOrAuthSig;
 
 class AuthMaterialValidator implements ParamsValidator {
   private readonly fnName: string;
@@ -476,15 +474,16 @@ class AuthMaterialValidator implements ParamsValidator {
   }
 
   validate(): IEither<void> {
-    const { authSig, sessionSigs } = this.authMaterial;
-
-    if (authSig && !is(authSig, 'Object', 'authSig', this.fnName))
+    if (
+      'authSig' in this.authMaterial &&
+      !is(this.authMaterial.authSig, 'Object', 'authSig', this.fnName)
+    )
       return ELeft(
         new InvalidParamType(
           {
             info: {
               param: 'authSig',
-              value: authSig,
+              value: this.authMaterial.authSig,
               functionName: this.fnName,
             },
           },
@@ -507,9 +506,9 @@ class AuthMaterialValidator implements ParamsValidator {
         );
 
       if (
-        authSig &&
+        'authSig' in this.authMaterial &&
         !checkIfAuthSigRequiresChainParam(
-          authSig,
+          this.authMaterial.authSig,
           this.authMaterial.chain,
           this.fnName
         )
@@ -519,7 +518,7 @@ class AuthMaterialValidator implements ParamsValidator {
             {
               info: {
                 param: 'authSig',
-                value: authSig,
+                value: this.authMaterial.authSig,
                 functionName: this.fnName,
               },
             },
@@ -528,43 +527,47 @@ class AuthMaterialValidator implements ParamsValidator {
         );
     }
 
-    if (sessionSigs && !is(sessionSigs, 'Object', 'sessionSigs', this.fnName))
+    if (
+      'sessionSigs' in this.authMaterial &&
+      !is(this.authMaterial.sessionSigs, 'Object', 'sessionSigs', this.fnName)
+    )
       return ELeft(
         new InvalidParamType(
           {
             info: {
               param: 'sessionSigs',
-              value: sessionSigs,
+              value: this.authMaterial.sessionSigs,
               functionName: this.fnName,
             },
           },
           'sessionSigs is not an object'
         )
       );
 
-    if (!sessionSigs && !authSig)
+    if (
+      !('sessionSigs' in this.authMaterial) &&
+      !('authSig' in this.authMaterial)
+    )
       return ELeft(
         new InvalidArgumentException(
           {
             info: {
               functionName: this.fnName,
-              sessionSigs,
-              authSig,
+              authMaterial: this.authMaterial,
             },
           },
           'You must pass either authSig or sessionSigs'
         )
       );
 
     // -- validate: if sessionSig and authSig exists
-    if (sessionSigs && authSig)
+    if ('sessionSigs' in this.authMaterial && 'authSig' in this.authMaterial)
       return ELeft(
         new InvalidArgumentException(
           {
             info: {
               functionName: this.fnName,
-              sessionSigs,
-              authSig,
+              authMaterial: this.authMaterial,
             },
           },
           'You cannot have both authSig and sessionSigs'
diff --git a/packages/schemas/src/lib/schemas.ts b/packages/schemas/src/lib/schemas.ts
@@ -669,39 +669,112 @@ export const EncryptSdkParamsSchema =
   });
 
 /**
- * This interface is mainly used for access control conditions & decrypt requests.
+ * These schemas are mainly used for access control conditions & decrypt requests.
  * For signing operations such as executeJs and pkpSign, only sessionSigs is used.
  */
-export const SessionSigsOrAuthSigSchema = z.object({
-  // TODO pretty sure this has to be one OR the other but at least one is required
+const SessionSigsAuthenticationSchema = z.object({
   /**
    * the session signatures to use to authorize the user with the nodes
    */
-  sessionSigs: SessionSigsSchema.optional(),
+  sessionSigs: SessionSigsSchema,
+});
+const AuthSigAuthenticationSchema = z.object({
   /**
    * This is a bare authSig generated client side by the user. It can only be used for access control conditions/encrypt/decrypt operations. It CANNOT be used for signing operation.
    */
-  authSig: AuthSigSchema.optional(),
+  authSig: AuthSigSchema,
 });
+export const SessionSigsOrAuthSigSchema = z.union([
+  SessionSigsAuthenticationSchema,
+  AuthSigAuthenticationSchema,
+]);
 
-export const DecryptRequestBaseSchema = SessionSigsOrAuthSigSchema.merge(
-  MultipleAccessControlConditionsSchema
-).extend({
+const ChainedSchema = z.object({
   /**
-   * The chain name of the chain that this contract is deployed on.  See LIT_CHAINS for currently supported chains.
+   * The chain name of the chain that will be used. See LIT_CHAINS for currently supported chains.
    */
   chain: ChainSchema,
 });
 
-export const DecryptRequestSchema = EncryptResponseSchema.merge(
-  DecryptRequestBaseSchema
-).extend({});
+export const ChainedSessionSigsOrAuthSigSchema = z.union([
+  SessionSigsAuthenticationSchema.merge(ChainedSchema.partial()),
+  AuthSigAuthenticationSchema.merge(ChainedSchema.partial()),
+]);
+
+const ChainMultipleAccessControlConditionsSchema =
+  MultipleAccessControlConditionsSchema.merge(ChainedSchema);
+const SessionSigsDecryptRequestBaseSchema =
+  ChainMultipleAccessControlConditionsSchema.merge(
+    SessionSigsAuthenticationSchema
+  );
+const AuthSigDecryptRequestBaseSchema =
+  ChainMultipleAccessControlConditionsSchema.merge(AuthSigAuthenticationSchema);
+export const DecryptRequestBaseSchema = z.union([
+  SessionSigsDecryptRequestBaseSchema,
+  AuthSigDecryptRequestBaseSchema,
+]);
+
+export const DecryptRequestSchema = z.union([
+  EncryptResponseSchema.merge(SessionSigsDecryptRequestBaseSchema),
+  EncryptResponseSchema.merge(AuthSigDecryptRequestBaseSchema),
+]);
 
 export const DecryptResponseSchema = z.object({
   // The decrypted data as a Uint8Array
   decryptedData: z.instanceof(Uint8Array),
 });
 
+const EncryptRequestBaseSchema = z.object({
+  // The data to encrypt as a Uint8Array
+  dataToEncrypt: z.instanceof(Uint8Array),
+});
+const SessionSigsEncryptRequestBaseSchema = EncryptRequestBaseSchema.merge(
+  SessionSigsDecryptRequestBaseSchema
+);
+const AuthSigEncryptRequestBaseSchema = EncryptRequestBaseSchema.merge(
+  AuthSigDecryptRequestBaseSchema
+);
+export const EncryptRequestSchema = z.union([
+  SessionSigsEncryptRequestBaseSchema,
+  AuthSigEncryptRequestBaseSchema,
+]);
+
+export const EncryptUint8ArrayRequestSchema =
+  MultipleAccessControlConditionsSchema.extend({
+    /**
+     * The uint8array that you wish to encrypt
+     */
+    dataToEncrypt: z.instanceof(Uint8Array),
+  });
+
+export const EncryptStringRequestSchema =
+  MultipleAccessControlConditionsSchema.extend({
+    /**
+     * The string that you wish to encrypt
+     */
+    dataToEncrypt: z.string(),
+  });
+
+export const EncryptFileRequestSchema =
+  MultipleAccessControlConditionsSchema.extend({
+    file: z.union([z.instanceof(File), z.instanceof(Blob)]),
+  });
+
+const EncryptToJsonPayloadBaseSchema = z.object({
+  ciphertext: z.string(),
+  dataToEncryptHash: z.string(),
+  dataType: z.enum(['string', 'file'] as const),
+});
+const SessionSigsEncryptToJsonPayloadSchema =
+  EncryptToJsonPayloadBaseSchema.merge(SessionSigsDecryptRequestBaseSchema);
+const AuthSigEncryptToJsonPayloadSchema = EncryptToJsonPayloadBaseSchema.merge(
+  AuthSigDecryptRequestBaseSchema
+);
+export const EncryptToJsonPayloadSchema = z.union([
+  SessionSigsEncryptToJsonPayloadSchema,
+  AuthSigEncryptToJsonPayloadSchema,
+]);
+
 export const SuccessNodePromisesSchema = z.object({
   success: z.literal(true),
   values: z.array(z.any()), // TODO add back generics
diff --git a/packages/types/src/lib/interfaces.ts b/packages/types/src/lib/interfaces.ts
@@ -45,6 +45,12 @@ import {
   JsonExecutionRequestSchema,
   NodeCommandResponseSchema,
   CallRequestSchema,
+  EncryptUint8ArrayRequestSchema,
+  EncryptRequestSchema,
+  EncryptStringRequestSchema,
+  EncryptFileRequestSchema,
+  EncryptToJsonPayloadSchema,
+  ChainedSessionSigsOrAuthSigSchema,
 } from '@lit-protocol/schemas';
 
 import { ILitNodeClient } from './ILitNodeClient';
@@ -55,7 +61,6 @@ import {
   Chain,
   EvmContractConditions,
   IRelayAuthStatus,
-  JsonRequest,
   ResponseStrategy,
   SolRpcConditions,
   SymmetricKey,
@@ -356,36 +361,26 @@ export interface JsonExecutionRequestTargetNode extends JsonExecutionRequest {
 
 export type JsonExecutionRequest = z.infer<typeof JsonExecutionRequestSchema>;
 
+export type ChainedSessionSigsOrAuthSig = z.infer<
+  typeof ChainedSessionSigsOrAuthSigSchema
+>;
+
 export type SessionSigsOrAuthSig = z.infer<typeof SessionSigsOrAuthSigSchema>;
 
 export type DecryptRequestBase = z.infer<typeof DecryptRequestBaseSchema>;
 export type EncryptSdkParams = z.infer<typeof EncryptSdkParamsSchema>;
 
-export interface EncryptRequest extends DecryptRequestBase {
-  // The data that you wish to encrypt as a Uint8Array
-  dataToEncrypt: Uint8Array;
-}
+export type EncryptRequest = z.infer<typeof EncryptRequestSchema>;
 
 export type EncryptResponse = z.infer<typeof EncryptResponseSchema>;
 
-export interface EncryptUint8ArrayRequest
-  extends MultipleAccessControlConditions {
-  /**
-   * The uint8array that you wish to encrypt
-   */
-  dataToEncrypt: Uint8Array;
-}
+export type EncryptUint8ArrayRequest = z.infer<
+  typeof EncryptUint8ArrayRequestSchema
+>;
 
-export interface EncryptStringRequest extends MultipleAccessControlConditions {
-  /**
-   * String that you wish to encrypt
-   */
-  dataToEncrypt: string;
-}
+export type EncryptStringRequest = z.infer<typeof EncryptStringRequestSchema>;
 
-export interface EncryptFileRequest extends DecryptRequestBase {
-  file: AcceptedFileType;
-}
+export type EncryptFileRequest = z.infer<typeof EncryptFileRequestSchema>;
 
 export type DecryptRequest = z.infer<typeof DecryptRequestSchema>;
 
@@ -603,13 +598,7 @@ export interface EncryptToJsonProps extends MultipleAccessControlConditions {
   litNodeClient: ILitNodeClient;
 }
 
-export type EncryptToJsonDataType = 'string' | 'file';
-
-export interface EncryptToJsonPayload extends DecryptRequestBase {
-  ciphertext: string;
-  dataToEncryptHash: string;
-  dataType: EncryptToJsonDataType;
-}
+export type EncryptToJsonPayload = z.infer<typeof EncryptToJsonPayloadSchema>;
 
 export interface DecryptFromJsonProps {
   // the session signatures to use to authorize the user with the nodes
