apollo_starknet_os_program: move encrypt state diff to a function (#9464)

Author: AvivYossef-starkware

crates/apollo_starknet_os_program/src/cairo/starkware/starknet/core/os/encrypt.cairo

@@ -13,6 +13,46 @@ from starkware.starknet.core.os.naive_blake import (
 from starkware.cairo.common.cairo_blake2s.blake2s import blake_with_opcode
 from starkware.cairo.common.alloc import alloc
 
+// Encrypts state diff data by generating keys, outputting public keys,
+// encrypting symmetric keys, and encrypting the data.
+func encrypt_state_diff{range_check_ptr, ec_op_ptr: EcOpBuiltin*}(
+    compressed_start: felt*, compressed_dst: felt*, n_keys: felt, public_keys: felt*
+) -> (encrypted_start: felt*, encrypted_dst: felt*) {
+    alloc_locals;
+
+    // Generate random symmetric key and random starknet private keys.
+    local symmetric_key: felt;
+    local sn_private_keys: felt*;
+    %{ generate_keys_from_hash(ids.compressed_start, ids.compressed_dst, ids.n_keys) %}
+    validate_sn_private_keys(n_keys=n_keys, sn_private_keys=sn_private_keys);
+
+    local encrypted_start: felt*;
+    %{
+        if use_kzg_da:
+            ids.encrypted_start = segments.add()
+        else:
+            # Assign a temporary segment, to be relocated into the output segment.
+            ids.encrypted_start = segments.add_temp_segment()
+    %}
+
+    let encrypted_dst = encrypted_start;
+    assert encrypted_dst[0] = n_keys;
+    let encrypted_dst = &encrypted_dst[1];
+
+    with encrypted_dst {
+        output_sn_public_keys(n_keys=n_keys, sn_private_keys=sn_private_keys);
+        output_encrypted_symmetric_key(
+            n_keys=n_keys,
+            public_keys=public_keys,
+            sn_private_keys=sn_private_keys,
+            symmetric_key=symmetric_key,
+        );
+        encrypt(data_start=compressed_start, data_end=compressed_dst, symmetric_key=symmetric_key);
+    }
+
+    return (encrypted_start=encrypted_start, encrypted_dst=encrypted_dst);
+}
+
 // Validates that the private keys are within the range [1, StarkCurve.ORDER - 1] as required by
 // the Diffie-Hellman elliptic curve encryption scheme.
 func validate_sn_private_keys{range_check_ptr}(n_keys: felt, sn_private_keys: felt*) {

crates/apollo_starknet_os_program/src/cairo/starkware/starknet/core/os/output.cairo

@@ -27,6 +27,7 @@ from starkware.starknet.core.os.encrypt import (
     output_sn_public_keys,
     output_encrypted_symmetric_key,
     encrypt,
+    encrypt_state_diff,
 )
 
 // Represents the output of the OS.
@@ -270,35 +271,12 @@ func process_data_availability{range_check_ptr, ec_op_ptr: EcOpBuiltin*}(
     }
 
     // Encrypt the compressed state updates.
-    // Generate random symmetric key and random starknet private keys.
-    local symmetric_key: felt;
-    local sn_private_keys: felt*;
-    %{ generate_keys_from_hash(ids.compressed_start, ids.compressed_dst, ids.n_keys) %}
-    validate_sn_private_keys(n_keys=n_keys, sn_private_keys=sn_private_keys);
-
-    local encrypted_start: felt*;
-    %{
-        if use_kzg_da:
-            ids.encrypted_start = segments.add()
-        else:
-            # Assign a temporary segment, to be relocated into the output segment.
-            ids.encrypted_start = segments.add_temp_segment()
-    %}
-
-    let encrypted_dst = encrypted_start;
-    assert encrypted_dst[0] = n_keys;
-    let encrypted_dst = &encrypted_dst[1];
-
-    with encrypted_dst {
-        output_sn_public_keys(n_keys=n_keys, sn_private_keys=sn_private_keys);
-        output_encrypted_symmetric_key(
-            n_keys=n_keys,
-            public_keys=public_keys,
-            sn_private_keys=sn_private_keys,
-            symmetric_key=symmetric_key,
-        );
-        encrypt(data_start=compressed_start, data_end=compressed_dst, symmetric_key=symmetric_key);
-    }
+    let (encrypted_start, encrypted_dst) = encrypt_state_diff(
+        compressed_start=compressed_start,
+        compressed_dst=compressed_dst,
+        n_keys=n_keys,
+        public_keys=public_keys,
+    );
 
     return (da_start=encrypted_start, da_end=encrypted_dst);
 }

crates/apollo_starknet_os_program/src/program_hash.json

@@ -1,5 +1,5 @@
 {
-  "os": "0x3e7651b150fbcab98774184c28edbc8cfa10e0e246d612810b1ec0133c0e7d3",
-  "aggregator": "0x4b62d0cfc8c03d2b2c4e214c2f2792cdd66068e871b33010a19f2c747877fe9",
-  "aggregator_with_prefix": "0x1605e56d904eb1db8959651aa03e70828c763c842e2f433a2a25706f7dba173"
+  "os": "0x338203798f8c939bae2c9bd06afd80bb63595007601729b33d6739920416c24",
+  "aggregator": "0x110273c1890c2e09a34e3995eeeaa7e0a10fe5a8bd98c84ccf6243512f7a7e6",
+  "aggregator_with_prefix": "0x3ca5d4c4724e45c896bf36d26e28a7e585e61d680b3b06962af652cafa166d7"
 }