Files
self/prover/dsc_prover.py
2024-07-18 16:16:38 +02:00

68 lines
2.4 KiB
Python

from fastapi import Request
from fastapi.responses import PlainTextResponse
from modal import Image, Mount, App, web_endpoint
import subprocess
import asyncio
import json
import hashlib
image = Image.from_dockerfile("Dockerfile")
app = App(image=image)
mount = Mount.from_local_dir("src", remote_path="/root/src")
@app.function(mounts=[mount], cpu=14)
@web_endpoint(method="POST")
async def generate_dsc_proof(request: Request):
# Read the JSON data from the request body
data = await request.json()
# Extract the signature_algorithm
signature_algorithm = data.get("signature_algorithm")
if signature_algorithm not in ["sha256_rsa", "sha1_rsa"]:
return PlainTextResponse("Invalid or missing signature_algorithm", status_code=400)
# Ensure 'inputs' key exists
if "inputs" not in data:
return PlainTextResponse("Missing 'inputs' in request data", status_code=400)
# Convert JSON data to a properly formatted string
json_data = json.dumps(data)
# Compute the hash of the input data
hash_object = hashlib.sha256(json_data.encode())
hash_hex = hash_object.hexdigest()
# Run the script and pass the JSON data to it asynchronously
process = await asyncio.create_subprocess_exec(
"/bin/bash", "/root/src/prove.sh",
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
)
stdout, stderr = await process.communicate(input=json_data.encode())
# Log detailed error message if any
if process.returncode != 0:
error_message = f"Error: {stderr.decode()}\n"
error_message += f"Standard Output: {stdout.decode()}"
return PlainTextResponse(error_message, status_code=500)
# Return the content of the unique proof and public files
try:
with open(f"/root/src/data/{hash_hex}/proof.json", "r") as proof_file:
proof_content = proof_file.read()
with open(f"/root/src/data/{hash_hex}/public.json", "r") as public_file:
public_content = public_file.read()
response_data = {
"proof": json.loads(proof_content),
"pub_signals": json.loads(public_content)
}
return PlainTextResponse(json.dumps(response_data), media_type="application/json")
except FileNotFoundError:
return PlainTextResponse("Error: proof.json or public.json not found", status_code=500)
# Run the app
if __name__ == "__main__":
app.run()