server.py

"""Server that will listen for GET and POST requests from the client."""

import time
from typing import List
from fastapi import FastAPI, File, Form, UploadFile
from fastapi.responses import JSONResponse, Response

from settings import DEPLOYMENT_PATH, SERVER_FILES
from utils.client_server_interface import MultiInputsFHEModelServer

# Load the server
FHE_SERVER = MultiInputsFHEModelServer(DEPLOYMENT_PATH)


def _get_server_file_path(name, client_id, client_type):
    """Get the correct temporary file path for the server.

    Args:
        name (str): The desired file name (either 'evaluation_key' or 'encrypted_inputs').
        client_id (int): The client ID to consider.
        client_type (str): The type of user to consider (either 'user', 'bank' or 'third_party').

    Returns:
        pathlib.Path: The file path.
    """
    return SERVER_FILES / f"{name}_{client_type}_{client_id}"


# Initialize an instance of FastAPI
app = FastAPI()

# Define the default route
@app.get("/")
def root():
    return {"message": "Welcome to Credit Card Approval Prediction server!"}


@app.post("/send_input")
def send_input(
    client_id: str = Form(),
    client_type: str = Form(),
    files: List[UploadFile] = File(),
):
    """Send the inputs to the server."""
    # Retrieve the encrypted inputs and the evaluation key paths
    encrypted_inputs_path = _get_server_file_path("encrypted_inputs", client_id, client_type)
    evaluation_key_path = _get_server_file_path("evaluation_key", client_id, client_type)
    
    # Write the files using the above paths
    with encrypted_inputs_path.open("wb") as encrypted_inputs, evaluation_key_path.open(
        "wb"
    ) as evaluation_key:
        encrypted_inputs.write(files[0].file.read())
        evaluation_key.write(files[1].file.read())


@app.post("/run_fhe")
def run_fhe(
    user_id: str = Form(),
    bank_id: str = Form(),
    third_party_id: str = Form(),
):
    """Execute the model on the encrypted inputs using FHE."""
    # Retrieve the evaluation key (from the user, as all evaluation keys should be the same)
    evaluation_key_path = _get_server_file_path("evaluation_key", user_id, "user")

    # Get the encrypted inputs
    encrypted_user_inputs_path = _get_server_file_path("encrypted_inputs", user_id, "user")
    encrypted_bank_inputs_path = _get_server_file_path("encrypted_inputs", bank_id, "bank")
    encrypted_third_party_inputs_path = _get_server_file_path("encrypted_inputs", third_party_id, "third_party")
    with (
        evaluation_key_path.open("rb") as evaluation_key_file,
        encrypted_user_inputs_path.open("rb") as encrypted_user_inputs_file,
        encrypted_bank_inputs_path.open("rb") as encrypted_bank_inputs_file,
        encrypted_third_party_inputs_path.open("rb") as encrypted_third_party_inputs_file,
    ):
        evaluation_key = evaluation_key_file.read()
        encrypted_user_input = encrypted_user_inputs_file.read()
        encrypted_bank_input = encrypted_bank_inputs_file.read()
        encrypted_third_party_input = encrypted_third_party_inputs_file.read()

    encrypted_inputs = (encrypted_user_input, encrypted_bank_input, encrypted_third_party_input)

    # Run the FHE execution
    start = time.time()
    encrypted_output = FHE_SERVER.run(*encrypted_inputs, serialized_evaluation_keys=evaluation_key)
    fhe_execution_time = round(time.time() - start, 2)

    # Retrieve the encrypted output path
    encrypted_output_path = _get_server_file_path("encrypted_output", user_id + bank_id + third_party_id, "output")

    # Write the file using the above path
    with encrypted_output_path.open("wb") as output_file:
        output_file.write(encrypted_output)

    return JSONResponse(content=fhe_execution_time)


@app.post("/get_output")
def get_output(
    user_id: str = Form(),
    bank_id: str = Form(),
    third_party_id: str = Form(),
):
    """Retrieve the encrypted output."""
    # Retrieve the encrypted output path
    encrypted_output_path = _get_server_file_path("encrypted_output", user_id + bank_id + third_party_id, "output")

    # Read the file using the above path
    with encrypted_output_path.open("rb") as encrypted_output_file:
        encrypted_output = encrypted_output_file.read()

    return Response(encrypted_output)