docs: add document for AI models security enhancement

Signed-off-by: Dong, Xiaocheng <[email protected]>

docs/Security_Enhancement_For_AI_Models.md

@@ -0,0 +1,249 @@
+# Security Enhancement For AI Models in Cloud Native Environment
+
+AI models often rely on large and sensitive datasets for training, and sometimes can be considered valuable intellectual property. Protecting the AI models is essential to prevent data breaches and unauthorized access to sensitive information. Typically, a common way to protect AI models is encryption, and some encryption algorithms (such as AES GCM) can ensure the confidentiality and integrity of the data.
+
+However, although a model has been encrypted, it would be loaded after decryption, and the runtime environment is untrusted most times, especially for cloud native environments. This document describes a design and the implementation to enhance the security of the flow in a Trusted Execution Environment (TEE).
+
+## 1. Design
+
+![](secure-model-design.png)
+
+
+### 1.1 Model Preprocessing
+
+To protect the model and its encryption key, the following preprocessing steps are taken:
+
+- Generate a key, a AES GCM key can be used for confidentiality and integrity.
+- Encrypt the model by the key.
+- Upload the encrypted model to the model server.
+- Register key to Key Broker Service (KBS), and KBS communicates with key management service (KMS) to store the key.
+
+
+### 1.2 Confidential Cloud Native Primitives (CCNP)
+
+[CCNP](https://github.com/intel/confidential-cloud-native-primitives) is the solution targeted on simplifying the use of Trusted Execution Environment (TEE) in cloud-native environment. It comprises two main components: the services and the SDK.
+
+- Service is designed to hide the complexity of different TEE platforms and provides common interfaces and scalability for cloud-native environment.
+
+- SDK is to simplify the use of the service interface for development, it covers communication to the service and parses the results from the services.
+
+The service supports attestation, measurement fetching and event log collecting of various platforms including Intel Trusted Domain Extensions (TDX), Trusted Platform Modules (TPM) and AMD SEV-SNP.
+
+To get the key to decrypt the model, we need provide the quote of TEE for attestation, CCNP is a good choice to get the quote and it hides the complexity and is easy to use:
+
+```Python
+from ccnp import Quote
+quote=Quote.get_quote()
+```
+
+Detail documentation can be found [here](https://intel.github.io/confidential-cloud-native-primitives/).
+
+### 1.3 Inference Flow
+
+The inference service flow of getting the model should be updated since the model has been encrypted:
+
+- Deploy CCNP as a [DaemonSet](https://intel.github.io/confidential-cloud-native-primitives/_rst/service.quote.html#deploy-as-daemonset-in-kubernetes), it will detect the TEE environment.
+- Get the quote and request the key from KBS.
+- Get the encrypted model from model server.
+- Decrypt the model and load the model for inference.
+
+## 2. Implementation
+
+Current Cloud Native AI Pipeline (CNAP) already implemented the new design to protect the models, this section will describe the details of each components and introduce how to enable a new model server or KBS interface.
+
+To protect the model, the first thing is to generate a key for model encryption, encrypt the model and register the key to KBS. KBS will submit further request to store the key to KMS.
+
+To fetch and decrypt the model, the runtime environment should be verified by attestation. A key broker client should get a quote and send the quote to KBS to verify that it is running on an authentic TEE platform. Upon successful attestation, KBS will then respond with a wrapped key for model decryption.
+
+### 2.1 Mode Encryption
+
+AES-GCM is a widely used block cipher mode of operation for securing data, particularly in network communications and data storage. It combines two important cryptographic techniques: AES and GCM. AES is a symmetric encryption algorithm that uses a shared secret key for both encryption and decryption. GCM is based on a combination of CTR and GMAC, this combination ensures that not only the data is encrypted, but it also includes an authentication tag that allows the recipient to verify the integrity of the data.
+
+The key features and benefits of AES GCM include confidentiality, data integrity and authentication. The Initialization Vector (IV) is a random value used as an input alongside the encryption key to initialize the encryption process, and the authentication tag is a small piece of data generated during encryption and appended to the ciphertext.
+
+Typically, the AES 256 GCM encrypted data format should be:
+
+12 bytes header | [12] bytes IV | encrypted data | [16] bytes tag
+---|---|---|---
+
+and the 12 bytes header:
+
+uint32 IV length | uint32 tag length | uint32 data length
+---|---|---
+
+```
+
+Some sample codes for encyrption:
+
+```Python
+import secrets
+import struct
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+
+def encrypt_file(input, output, key):
+    with open(args.input, 'rb') as infile:
+        data = infile.read()
+        aesgcm = AESGCM(key)
+        IV = secrets.token_bytes(12)
+        data_enc = aesgcm.encrypt(IV, data, None)
+        with open(output, 'wb') as outfile:
+            header = struct.pack('<3I', len(IV), 16, len(data_enc))
+            outfile.write(header)
+            outfile.write(IV)
+            outfile.write(data_enc)
+```
+
+### 2.2 Key Registration
+
+After encrypting the model, the key should be registered with KBS. The KBS handles the attestation process for a client to retrieve the key.
+A typical flow of key registration:
+
+- Connect to the KBS with authentication information to get a token.
+- Create the key transfer policy with the token, the policy should includes the attestation type and attestation polices.
+- Register the key with the key transfer policy ID and the token.
+
+### 2.3 Model Server
+
+An [abstract class](https://github.com/intel/cloud-native-ai-pipeline/blob/main/cnap/core/modelprovider.py) `ModelProvider` was introduced in CNAP project, it includes the common interface of getting the model information from a model server (service). And the `SimpleModelProvider` is an implementaion of `ModelProvider`, which shows an example to get a model from the model server.
+
+A simple model server includes a interface to get the model information, such as id, framework, model URL, KBS URL etc. here is an example in CNAP:
+
+- Request
+
+```Python
+# HTTP Headers
+headers = {
+            "Accept": "application/json",
+            "Content-Type": "application/json"
+}
+# Request URL
+req_url = http://{model_info_url}/{model_id}
+```
+
+- Response
+
+```JSON
+{
+    "id": "c8b019e0-f4d8-4831-8936-f7f64ad99509",
+    "framework": "tensorflow",
+    "target": "object-detection",
+    "url": "http://{model_server_url}/tensorflow/ssdmobilenet_v10.pb.enc",
+    "name": "ssdmobilenet",
+    "version": "1.0",
+    "dtype": "int8",
+    "encrypted": true,
+    "kbs": "simple_kbs",
+    "kbs_url": "https://{key_brocker_server_url}",
+    "key_id": "88b3a4b3-ba4b-4e2b-b00e-bc32ca3d56be"
+}
+```
+
+A model provider should be implemented to get models information and download the model from model URL, if the model is encrypted, then invoke the key broker client to get the key from the KBS and decrypt the model.
+
+### 2.4 Key Broker Service
+
+An [abstract class](https://github.com/intel/cloud-native-ai-pipeline/blob/main/cnap/core/keybroker.py) `KeyBrokerClientBase` was introduced in CNAP project, it has an abstract method `get_key` to get a key from Key Broker Service (KBS). For a new KBS, a key broker client should be implemented to connect to KBS and get the key. `SimpleKeyBrokerClient` is an example to get a key from a simple KBS.
+
+An example flow of a simple KBS:
+
+- Accept a quote and a public key from client.
+- Verify the quote and do attestation, return if verify failed.
+- Get the user key from key management server (KMS), generate a symmetric wrapping
+    key (SWK) to encrypt the user key (wrapped_key).
+- Encrypt the SWK by the public key from client (wrapped_swk).
+
+For a key broker client, here is an example flow to get a key from KBS:
+
+- Generate 2048 bit RSA key pair (a public key and a private key).
+- Encode the public key to base64 for transferring (user_data).
+- Get quote in the TEE with the hash of the public key for measurement (quote).
+- Request wrapped_key and wrapped_swk from KBS with quote and user_data.
+- Decrypt the user key by the SWK.
+
+An example interface to communicate to the KBS:
+
+- Request
+
+```Python
+# HTTP Headers
+headers = {
+    "Accept": "application/json",
+    "Content-Type": "application/json",
+    "Attestation-Type": "TDX"
+}
+# HTTP Body
+req_body = {
+    "quote": quote,
+    "user_data": user_data
+}
+```
+
+- Response
+
+```JSON
+{
+    "wrapped_key":"",
+    "wrapped_swk":""
+}
+```
+
+### 2.5 Attestation
+
+CCNP provides an easy way to get the quote for attestation. Before getting the quote, a RSA key pair need to be generated for wrapping purpose and the public key will be measured as the user data input to the quote fetching process.
+
+The public key is used as parameter to get the key, the KBS is to encyrpt the key wrapped key by the public key.
+
+The sample code to get the quote with user data:
+
+```Python
+import base64
+
+from ccnp import Quote
+from cryptography.hazmat.primitives.asymmetric import rsa
+
+private_key = rsa.generate_private_key(public_exponent=65537, key_size=3072)
+pubkey = private_key.public_key()
+pubkey_der = pubkey.public_bytes(encoding=serialization.Encoding.DER,
+                                 format=serialization.PublicFormat.SubjectPublicKeyInfo)
+user_data = base64.b64encode(pubkey_der).decode('utf-8')
+quote = Quote.get_quote(user_data=user_data)
+```
+
+### 2.6 Model Decryption
+
+Just like the encryption, the AES 256 GCM encrypted data format should be:
+
+12 bytes header | [12] bytes IV | encrypted data | [16] bytes tag
+---|---|---|---
+
+and the 12 bytes header:
+
+uint32 IV length | uint32 tag length | uint32 data length
+---|---|---
+
+To decrypt the data, here are some sample codes:
+
+```
+
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+
+def decrypt_data(encrypted_data, key) -> bytes:
+
+    if encrypted_data is None or key is None:
+        raise ValueError("Encrypted data or key is empty")
+
+    header_len = 12
+    iv_len, tag_len, data_len = struct.unpack('<3I', encrypted_data[:header_len])
+    iv = encrypted_data[header_len : (iv_len + header_len)]
+    data = encrypted_data[(iv_len + header_len) : -tag_len]
+    tag = encrypted_data[-tag_len:]
+
+    decryptor = Cipher(algorithms.AES(key), modes.GCM(iv, tag)).decryptor()
+    decrypted_data = decryptor.update(data) + decryptor.finalize()
+    return decrypted_data
+```
+
+# References
+
+Model Provider: https://github.com/intel/cloud-native-ai-pipeline/blob/main/cnap/core/keybroker.py
+Key broker client: https://github.com/intel/cloud-native-ai-pipeline/blob/main/cnap/core/keybroker.py