scripts/space_server.py

"""Single-process public gateway. Backend is loopback-only; uploads are session-owned."""
import asyncio, contextlib, hashlib, hmac, json, os, re, secrets, shutil, subprocess, sys, tempfile, time
from contextlib import asynccontextmanager
from pathlib import Path
import httpx
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, RedirectResponse, Response
from fastapi.staticfiles import StaticFiles
ROOT=Path(__file__).resolve().parents[1]
PUBLIC=Path(os.environ.get('ACE_PUBLIC_DIR',ROOT/'public'))
KEY=secrets.token_bytes(32);owners={};busy=set();gate=asyncio.Semaphore(1);waiting=0
inflight=0
TTL=7200;MAX_UPLOAD=50*1024*1024
backend_url=os.environ.get('ACE_BACKEND_URL','http://127.0.0.1:9877')
client=None;store=None

def sign(n):return hmac.new(KEY,n.encode(),hashlib.sha256).hexdigest()
def session(req):
    token=req.cookies.get('ace_session','');bits=token.split('.')
    if len(bits)==2 and re.fullmatch('[a-f0-9]{64}',bits[0]) and hmac.compare_digest(sign(bits[0]),bits[1]):return bits[0],None
    n=secrets.token_hex(32);return n,n+'.'+sign(n)

def remove_job(j):
    if store is None:return
    target=(store/j).resolve()
    if target.parent!=store.resolve() or not re.fullmatch('[a-f0-9]{32}',j):raise ValueError('Invalid cleanup target')
    if target.exists():shutil.rmtree(target)
    owners.pop(j,None)

async def cleanup():
    while True:
        await asyncio.sleep(60)
        for j,entry in list(owners.items()):
            if j not in busy and time.time()-entry['created']>TTL:remove_job(j)

@asynccontextmanager
async def lifespan(app):
    global client,store
    with tempfile.TemporaryDirectory(prefix='ace-public-') as temp:
        store=Path(temp)
        env={**os.environ,'ACE_PORT':'9877','ACE_STORE':temp,'PYTHONUNBUFFERED':'1'}
        proc=subprocess.Popen([sys.executable,str(ROOT/'scripts/casting_app.py')],env=env)
        client=httpx.AsyncClient(base_url=backend_url,timeout=420)
        task=None
        try:
            for _ in range(240):
                if proc.poll() is not None:raise RuntimeError('Inference backend failed to start')
                try:
                    r=await client.get('/api/model')
                    if r.status_code==200:break
                except httpx.HTTPError:pass
                await asyncio.sleep(1)
            else:raise RuntimeError('Inference backend startup timed out')
            task=asyncio.create_task(cleanup());yield
        finally:
            if task:
                task.cancel()
                with contextlib.suppress(asyncio.CancelledError):await task
            await client.aclose();proc.terminate()
            try:proc.wait(timeout=10)
            except subprocess.TimeoutExpired:proc.kill();proc.wait()

app=FastAPI(lifespan=lifespan,docs_url=None,redoc_url=None,openapi_url=None)
@app.middleware('http')
async def identity(req,call_next):
    global inflight
    limited=req.method=='POST' and req.url.path.startswith('/api/')
    if limited and inflight>=4:return JSONResponse({'error':'Service busy; please retry shortly.'},status_code=429)
    if limited:inflight+=1
    sid,cookie=session(req);req.state.sid=sid
    try:response=await call_next(req)
    finally:
        if limited:inflight-=1
    if cookie:response.set_cookie('ace_session',cookie,httponly=True,secure=req.url.scheme=='https',samesite='lax',max_age=TTL)
    if req.url.path.startswith('/api'):response.headers['Cache-Control']='no-store'
    response.headers['X-Content-Type-Options']='nosniff'
    return response

@app.get('/')
async def home():return RedirectResponse('/artifacts/whitepaper/index.html')
@app.get('/app')
async def app_root():return RedirectResponse('/app/')
@app.get('/healthz')
async def health():return {'status':'ready'}

def err(message,code=400):return JSONResponse({'error':message},status_code=code)
def owned(j,sid):return j in owners and owners[j]['owner']==sid and time.time()-owners[j]['created']<=TTL

def forward(r):
    # httpx decompresses; do not forward Content-Encoding/Length from upstream.
    return Response(r.content,status_code=r.status_code,media_type=r.headers.get('content-type','application/json'))

@app.api_route('/api/{route:path}',methods=['GET','POST'])
async def api(route,req:Request):
    global waiting
    sid=req.state.sid
    if req.method=='GET':
        if route=='model':return forward(await client.get('/api/model'))
        if route=='recent':return [{'id':j,'name':e['name']} for j,e in owners.items() if owned(j,sid)]
        if not re.fullmatch(r'(job|geometry)/[a-f0-9]{32}',route):return err('Not found',404)
        if not owned(route.split('/')[-1],sid):return err('Part not found or session expired',404)
        return forward(await client.get('/api/'+route))
    # Browser requests must originate from this host; no cross-origin upload API.
    origin=req.headers.get('origin')
    if origin and origin.rstrip('/')!=str(req.base_url).rstrip('/'):return err('Origin not allowed',403)
    if route not in ['upload','resume','predict']:return err('Not found',404)
    raw=bytearray();limit=MAX_UPLOAD if route=='upload' else 65536
    async for chunk in req.stream():
        raw.extend(chunk)
        if len(raw)>limit:return err('Request exceeds size limit',413)
    if route=='upload':
        if len(owners)>=24 or sum(owned(j,sid) for j in owners)>=3:return err('Upload limit reached. Temporary parts expire after two hours.',429)
        if b'ISO-10303-21' not in raw[:4096].upper():return err('Expected a STEP/STP file')
    else:
        try:opt=json.loads(raw);j=opt['id']
        except (ValueError,KeyError,TypeError):return err('Invalid request')
        if not isinstance(j,str) or not owned(j,sid):return err('Part not found or session expired',404)
        if route=='resume':return forward(await client.post('/api/resume',content=bytes(raw)))
    if waiting>=4:return err('Prediction service is busy. Please retry shortly.',429)
    waiting+=1;j=None
    try:
        async with gate:
            headers={'Accept':req.headers.get('accept','application/json')}
            if route=='upload':
                if len(owners)>=24 or sum(owned(k,sid) for k in owners)>=3:return err('Upload limit reached',429)
                name=req.headers.get('x-filename','Uploaded STEP')[:200];headers['X-Filename']=name
                r=await client.post('/api/upload',content=bytes(raw),headers=headers)
                if r.status_code!=200:return forward(r)
                result=r.json();j=result['id'];owners[j]={'owner':sid,'created':time.time(),'name':name};busy.add(j)
                # Keep CAD and inference serialized; browser shows its preparing state.
                for _ in range(390):
                    state=await client.get('/api/job/'+j)
                    if state.json().get('status')!='processing':return forward(state)
                    await asyncio.sleep(1)
                return err('Geometry preparation timed out',504)
            j=opt['id']
            if not owned(j,sid):return err('Part expired',404)
            busy.add(j)
            return forward(await client.post('/api/predict',content=bytes(raw),headers=headers))
    except httpx.HTTPError:return err('Prediction service unavailable; please retry',503)
    finally:
        waiting-=1
        if j:busy.discard(j)

# Only explicitly packaged assets are served. No repository or job-directory mounts.
app.mount('/app',StaticFiles(directory=PUBLIC/'artifacts/casting_app',html=True),name='app')
app.mount('/artifacts',StaticFiles(directory=PUBLIC/'artifacts'),name='showcase')