16. Edge Security

Chapter 16 of 18 · 20 min

Edge devices face security threats including model theft, inference manipulation, and physical tampering. Securing ML models on edge devices requires protecting model weights, validating inference inputs, and establishing secure update channels.

Model encryption at rest:

import hashlib
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
import os

def encrypt_model(model_path, output_path, key):
    """Encrypt model file for secure storage"""
    with open(model_path, 'rb') as f:
        plaintext = f.read()
    
    # Generate random nonce
    nonce = os.urandom(12)
    
    # Encrypt with AES-GCM
    aesgcm = AESGCM(key)
    ciphertext = aesgcm.encrypt(nonce, plaintext, None)
    
    # Store nonce + ciphertext
    with open(output_path, 'wb') as f:
        f.write(nonce)
        f.write(ciphertext)

def decrypt_model(model_path, key):
    """Decrypt model at runtime"""
    with open(model_path, 'rb') as f:
        nonce = f.read(12)
        ciphertext = f.read()
    
    aesgcm = AESGCM(key)
    return aesgcm.decrypt(nonce, ciphertext, None)

# Key derivation from device identity
def derive_model_key(device_id: str, master_key: bytes) -> bytes:
    """Derive per-device model encryption key"""
    from cryptography.hazmat.primitives.hashes import SHA256
    from cryptography.hazmat.primitives.kdf.hkdf import HKDF
    
    hkdf = HKDF(
        algorithm=SHA256(),
        length=32,
        salt=device_id.encode(),
        info=b"model-encryption"
    )
    return hkdf.derive(master_key)

Input validation to prevent adversarial attacks:

import numpy as np

def validate_input(input_data, expected_range=(-1.0, 1.0), max_size_mb=50):
    """Validate input before inference to prevent adversarial patterns"""
    
    # Size check
    size_bytes = input_data.nbytes
    if size_bytes > max_size_mb * 1024 * 1024:
        raise ValueError(f"Input exceeds size limit: {size_bytes}")
    
    # Range check
    min_val, max_val = expected_range
    if np.any(input_data < min_val) or np.any(input_data > max_val):
        raise ValueError("Input values outside expected range")
    
    # Distribution anomaly detection
    # Simple check: reject if >30% zeros or NaN values
    zero_ratio = np.sum(input_data == 0) / input_data.size
    if zero_ratio > 0.3:
        raise ValueError(f"Abnormal zero ratio: {zero_ratio:.2%}")
    
    return True

def secure_inference(pipeline, input_data, device_key):
    """Inference with input validation and model decryption"""
    
    # Decrypt model in memory
    decrypted_model = decrypt_model(pipeline.encrypted_model_path, device_key)
    
    # Validate input
    validate_input(input_data)
    
    # Run inference
    output = pipeline.run(decrypted_model, input_data)
    
    # Zero memory after use (defense in depth)
    decrypted_model.fill(0)
    
    return output

Secure boot verification for model authenticity:

# Verify DM-Verity integrity (Android)
adb shell dmverity mode self_verifying

# Check model file signatures
openssl dgst -sha256 -verify public.pem -signature model.sig model.tflite

# iOS model code signing verification
codesign -dvvv model.mlmodel
# Verification output includes Team ID and signing certificate

Android Keystore integration:

import android.security.keystore.KeyGenParameterSpec
import android.security.keystore.KeyProperties

class SecureModelLoader(private val context: Context) {
    
    private val keystore = KeyStore.getInstance("AndroidKeyStore").apply { load(null) }
    
    fun generateModelKey(): SecretKey {
        val keyGenerator = KeyGenerator.getInstance(
            KeyProperties.KEY_ALGORITHM_AES,
            "AndroidKeyStore"
        )
        
        val keySpec = KeyGenParameterSpec.Builder(
            "model_key",
            KeyProperties.PURPOSE_ENCRYPT or KeyProperties.PURPOSE_DECRYPT
        )
            .setBlockModes(KeyProperties.BLOCK_MODE_GCM)
            .setEncryptionPaddings(KeyProperties.ENCRYPTION_PADDING_NONE)
            .setUserAuthenticationRequired(true)
            .setUserAuthenticationValidityDurationSeconds(30)
            .build()
        
        keyGenerator.init(keySpec)
        return keyGenerator.generateKey()
    }
    
    fun decryptModel(encryptedModel: ByteArray): ByteArray {
        val key = keystore.getKey("model_key", null) as SecretKey
        val cipher = Cipher.getInstance("AES/GCM/NoPadding")
        cipher.init(Cipher.DECRYPT_MODE, key)
        
        val nonce = encryptedModel.sliceArray(0..11)
        val ciphertext = encryptedModel.sliceArray(12..encryptedModel.lastIndex)
        
        return cipher.doFinal(ciphertext)
    }
}
EXERCISE

Implement model encryption with AES-GCM, integrate Android Keystore for key management, and implement input validation that rejects adversarial patterns.