etan39
/
cloudstack
mirror of https://github.com/apache/cloudstack.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
cloudstack/systemvm/debian/opt/cloud/bin/image_server.py

679 lines
24 KiB
Python
Raw Blame History

#!/usr/bin/env python3
"""
POC "imageio-like" HTTP server backed by NBD over TCP.
How to run
----------
- Install dependency:
dnf install python3-libnbd
or
apt install python3-libnbd
- Run server:
python image_server.py --listen 0.0.0.0 --port 54323
Example curl commands
--------------------
- OPTIONS:
curl -i -X OPTIONS http://127.0.0.1:54323/images/demo
- GET full image:
curl -v http://127.0.0.1:54323/images/demo -o demo.img
- GET a byte range:
curl -v -H "Range: bytes=0-1048575" http://127.0.0.1:54323/images/demo -o first_1MiB.bin
- PUT full image (Content-Length must equal export size exactly):
curl -v -T demo.img http://127.0.0.1:54323/images/demo
- GET extents (POC-level; may return a single allocated extent):
curl -s http://127.0.0.1:54323/images/demo/extents | jq .
- POST flush:
curl -s -X POST http://127.0.0.1:54323/images/demo/flush | jq .
"""
from __future__ import annotations
import argparse
import json
import logging
import os
import socket
import threading
import time
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any, Dict, Optional, Tuple
import nbd
CHUNK_SIZE = 256 * 1024 # 256 KiB
# Concurrency limits across ALL images.
MAX_PARALLEL_READS = 8
MAX_PARALLEL_WRITES = 2
_READ_SEM = threading.Semaphore(MAX_PARALLEL_READS)
_WRITE_SEM = threading.Semaphore(MAX_PARALLEL_WRITES)
# In-memory per-image lock: single lock gates both read and write.
_IMAGE_LOCKS: Dict[str, threading.Lock] = {}
_IMAGE_LOCKS_GUARD = threading.Lock()
# Dynamic image_id(transferId) -> NBD export mapping:
# CloudStack writes a JSON file at /tmp/imagetransfer/<transferId> with:
# {"host": "...", "port": 10809, "export": "vda"}
#
# This server reads that file on-demand.
_CFG_DIR = "/tmp/imagetransfer"
_CFG_CACHE: Dict[str, Tuple[float, Dict[str, Any]]] = {}
_CFG_CACHE_GUARD = threading.Lock()
def _json_bytes(obj: Any) -> bytes:
return json.dumps(obj, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
def _get_image_lock(image_id: str) -> threading.Lock:
with _IMAGE_LOCKS_GUARD:
lock = _IMAGE_LOCKS.get(image_id)
if lock is None:
lock = threading.Lock()
_IMAGE_LOCKS[image_id] = lock
return lock
def _now_s() -> float:
return time.monotonic()
def _safe_transfer_id(image_id: str) -> Optional[str]:
"""
Only allow a single filename component to avoid path traversal.
We intentionally keep validation simple: reject anything containing '/' or '\\'.
"""
if not image_id:
return None
if image_id != os.path.basename(image_id):
return None
if "/" in image_id or "\\" in image_id:
return None
if image_id in (".", ".."):
return None
return image_id
def _load_image_cfg(image_id: str) -> Optional[Dict[str, Any]]:
safe_id = _safe_transfer_id(image_id)
if safe_id is None:
return None
cfg_path = os.path.join(_CFG_DIR, safe_id)
try:
st = os.stat(cfg_path)
except FileNotFoundError:
return None
except OSError as e:
logging.warning("cfg stat failed image_id=%s err=%r", image_id, e)
return None
with _CFG_CACHE_GUARD:
cached = _CFG_CACHE.get(safe_id)
if cached is not None:
cached_mtime, cached_cfg = cached
# Use cached config if the file hasn't changed.
if float(st.st_mtime) == float(cached_mtime):
return cached_cfg
try:
with open(cfg_path, "rb") as f:
raw = f.read(4096)
except OSError as e:
logging.warning("cfg read failed image_id=%s err=%r", image_id, e)
return None
try:
obj = json.loads(raw.decode("utf-8"))
except Exception as e:
logging.warning("cfg parse failed image_id=%s err=%r", image_id, e)
return None
if not isinstance(obj, dict):
logging.warning("cfg invalid type image_id=%s type=%s", image_id, type(obj).__name__)
return None
host = obj.get("host")
port = obj.get("port")
export = obj.get("export")
if not isinstance(host, str) or not host:
logging.warning("cfg missing/invalid host image_id=%s", image_id)
return None
try:
port_i = int(port)
except Exception:
logging.warning("cfg missing/invalid port image_id=%s", image_id)
return None
if port_i <= 0 or port_i > 65535:
logging.warning("cfg out-of-range port image_id=%s port=%r", image_id, port)
return None
if export is not None and (not isinstance(export, str) or not export):
logging.warning("cfg missing/invalid export image_id=%s", image_id)
return None
cfg: Dict[str, Any] = {"host": host, "port": port_i, "export": export}
with _CFG_CACHE_GUARD:
_CFG_CACHE[safe_id] = (float(st.st_mtime), cfg)
return cfg
class _NbdConn:
"""
Small helper to connect to NBD using an already-open TCP socket.
Opens a fresh handle per request, per POC requirements.
"""
def __init__(self, host: str, port: int, export: Optional[str]):
self._sock = socket.create_connection((host, port))
self._nbd = nbd.NBD()
# Select export name if supported/needed.
if export and hasattr(self._nbd, "set_export_name"):
self._nbd.set_export_name(export)
self._connect_existing_socket(self._sock)
def _connect_existing_socket(self, sock: socket.socket) -> None:
# Requirement: attach libnbd to an existing socket / FD (no qemu-nbd).
# libnbd python API varies slightly by version, so try common options.
last_err: Optional[BaseException] = None
if hasattr(self._nbd, "connect_socket"):
try:
self._nbd.connect_socket(sock)
return
except Exception as e: # pragma: no cover (depends on binding)
last_err = e
try:
self._nbd.connect_socket(sock.fileno())
return
except Exception as e2: # pragma: no cover
last_err = e2
if hasattr(self._nbd, "connect_fd"):
try:
self._nbd.connect_fd(sock.fileno())
return
except Exception as e: # pragma: no cover
last_err = e
raise RuntimeError(
"Unable to connect libnbd using existing socket/fd; "
f"binding missing connect_socket/connect_fd or call failed: {last_err!r}"
)
def size(self) -> int:
return int(self._nbd.get_size())
def pread(self, length: int, offset: int) -> bytes:
# Expected signature: pread(length, offset)
try:
return self._nbd.pread(length, offset)
except TypeError: # pragma: no cover (binding differences)
return self._nbd.pread(offset, length)
def pwrite(self, buf: bytes, offset: int) -> None:
# Expected signature: pwrite(buf, offset)
try:
self._nbd.pwrite(buf, offset)
except TypeError: # pragma: no cover (binding differences)
self._nbd.pwrite(offset, buf)
def flush(self) -> None:
if hasattr(self._nbd, "flush"):
self._nbd.flush()
return
if hasattr(self._nbd, "fsync"):
self._nbd.fsync()
return
raise RuntimeError("libnbd binding has no flush/fsync method")
def close(self) -> None:
# Best-effort; bindings may differ.
try:
if hasattr(self._nbd, "shutdown"):
self._nbd.shutdown()
except Exception:
pass
try:
if hasattr(self._nbd, "close"):
self._nbd.close()
except Exception:
pass
try:
self._sock.close()
except Exception:
pass
def __enter__(self) -> "_NbdConn":
return self
def __exit__(self, exc_type, exc, tb) -> None:
self.close()
class Handler(BaseHTTPRequestHandler):
server_version = "imageio-poc/0.1"
# Keep BaseHTTPRequestHandler from printing noisy default logs
def log_message(self, fmt: str, *args: Any) -> None:
logging.info("%s - - %s", self.address_string(), fmt % args)
def _send_imageio_headers(self) -> None:
# Include these headers for compatibility with the imageio contract.
self.send_header("Access-Control-Allow-Origin", "*")
self.send_header("Access-Control-Allow-Methods", "GET, PUT, OPTIONS")
self.send_header(
"Access-Control-Allow-Headers",
"Range, Content-Range, Content-Type, Content-Length",
)
self.send_header(
"Access-Control-Expose-Headers",
"Accept-Ranges, Content-Range, Content-Length",
)
self.send_header("Accept-Ranges", "bytes")
def _send_json(self, status: int, obj: Any) -> None:
body = _json_bytes(obj)
self.send_response(status)
self._send_imageio_headers()
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
try:
self.wfile.write(body)
except BrokenPipeError:
pass
def _send_error_json(self, status: int, message: str) -> None:
self._send_json(status, {"error": message})
def _send_range_not_satisfiable(self, size: int) -> None:
# RFC 7233: reply with Content-Range: bytes */<size>
self.send_response(HTTPStatus.REQUESTED_RANGE_NOT_SATISFIABLE)
self._send_imageio_headers()
self.send_header("Content-Type", "application/json")
self.send_header("Content-Range", f"bytes */{size}")
body = _json_bytes({"error": "range not satisfiable"})
self.send_header("Content-Length", str(len(body)))
self.end_headers()
try:
self.wfile.write(body)
except BrokenPipeError:
pass
def _parse_single_range(self, range_header: str, size: int) -> Tuple[int, int]:
"""
Parse a single HTTP byte range (RFC 7233) and return (start, end_inclusive).
Supported:
- Range: bytes=START-END
- Range: bytes=START-
- Range: bytes=-SUFFIX
Raises ValueError for invalid headers. Caller handles 416 vs 400.
"""
if size < 0:
raise ValueError("invalid size")
if not range_header:
raise ValueError("empty Range")
if "," in range_header:
raise ValueError("multiple ranges not supported")
prefix = "bytes="
if not range_header.startswith(prefix):
raise ValueError("only bytes ranges supported")
spec = range_header[len(prefix) :].strip()
if "-" not in spec:
raise ValueError("invalid bytes range")
left, right = spec.split("-", 1)
left = left.strip()
right = right.strip()
if left == "":
# Suffix range: last N bytes.
if right == "":
raise ValueError("invalid suffix range")
try:
suffix_len = int(right, 10)
except ValueError as e:
raise ValueError("invalid suffix length") from e
if suffix_len <= 0:
raise ValueError("invalid suffix length")
if size == 0:
# Nothing to serve
raise ValueError("unsatisfiable")
if suffix_len >= size:
return 0, size - 1
return size - suffix_len, size - 1
# START is present
try:
start = int(left, 10)
except ValueError as e:
raise ValueError("invalid range start") from e
if start < 0:
raise ValueError("invalid range start")
if start >= size:
raise ValueError("unsatisfiable")
if right == "":
# START-
return start, size - 1
try:
end = int(right, 10)
except ValueError as e:
raise ValueError("invalid range end") from e
if end < start:
raise ValueError("unsatisfiable")
if end >= size:
end = size - 1
return start, end
def _parse_route(self) -> Tuple[Optional[str], Optional[str]]:
# Returns (image_id, tail) where tail is:
# None => /images/{id}
# "extents" => /images/{id}/extents
# "flush" => /images/{id}/flush
path = self.path.split("?", 1)[0]
parts = [p for p in path.split("/") if p]
if len(parts) < 2 or parts[0] != "images":
return None, None
image_id = parts[1]
tail = parts[2] if len(parts) >= 3 else None
if len(parts) > 3:
return None, None
return image_id, tail
def _image_cfg(self, image_id: str) -> Optional[Dict[str, Any]]:
return _load_image_cfg(image_id)
def do_OPTIONS(self) -> None:
image_id, tail = self._parse_route()
if image_id is None or tail is not None:
self._send_error_json(HTTPStatus.NOT_FOUND, "not found")
return
if self._image_cfg(image_id) is None:
self._send_error_json(HTTPStatus.NOT_FOUND, "unknown image_id")
return
self.send_response(HTTPStatus.OK)
self._send_imageio_headers()
self.send_header("Content-Length", "0")
self.end_headers()
def do_GET(self) -> None:
image_id, tail = self._parse_route()
if image_id is None:
self._send_error_json(HTTPStatus.NOT_FOUND, "not found")
return
cfg = self._image_cfg(image_id)
if cfg is None:
self._send_error_json(HTTPStatus.NOT_FOUND, "unknown image_id")
return
if tail == "extents":
self._handle_get_extents(image_id, cfg)
return
if tail is not None:
self._send_error_json(HTTPStatus.NOT_FOUND, "not found")
return
range_header = self.headers.get("Range")
self._handle_get_image(image_id, cfg, range_header)
def do_PUT(self) -> None:
image_id, tail = self._parse_route()
if image_id is None or tail is not None:
self._send_error_json(HTTPStatus.NOT_FOUND, "not found")
return
cfg = self._image_cfg(image_id)
if cfg is None:
self._send_error_json(HTTPStatus.NOT_FOUND, "unknown image_id")
return
if self.headers.get("Range") is not None or self.headers.get("Content-Range") is not None:
self._send_error_json(
HTTPStatus.BAD_REQUEST, "Range/Content-Range not supported; full writes only"
)
return
content_length_hdr = self.headers.get("Content-Length")
if content_length_hdr is None:
self._send_error_json(HTTPStatus.BAD_REQUEST, "Content-Length required")
return
try:
content_length = int(content_length_hdr)
except ValueError:
self._send_error_json(HTTPStatus.BAD_REQUEST, "Invalid Content-Length")
return
if content_length < 0:
self._send_error_json(HTTPStatus.BAD_REQUEST, "Invalid Content-Length")
return
self._handle_put_image(image_id, cfg, content_length)
def do_POST(self) -> None:
image_id, tail = self._parse_route()
if image_id is None:
self._send_error_json(HTTPStatus.NOT_FOUND, "not found")
return
cfg = self._image_cfg(image_id)
if cfg is None:
self._send_error_json(HTTPStatus.NOT_FOUND, "unknown image_id")
return
if tail == "flush":
self._handle_post_flush(image_id, cfg)
return
self._send_error_json(HTTPStatus.NOT_FOUND, "not found")
def _handle_get_image(
self, image_id: str, cfg: Dict[str, Any], range_header: Optional[str]
) -> None:
lock = _get_image_lock(image_id)
if not lock.acquire(blocking=False):
self._send_error_json(HTTPStatus.CONFLICT, "image busy")
return
if not _READ_SEM.acquire(blocking=False):
lock.release()
self._send_error_json(HTTPStatus.SERVICE_UNAVAILABLE, "too many parallel reads")
return
start = _now_s()
bytes_sent = 0
try:
logging.info("GET start image_id=%s range=%s", image_id, range_header or "-")
with _NbdConn(cfg["host"], int(cfg["port"]), cfg.get("export")) as conn:
size = conn.size()
start_off = 0
end_off_incl = size - 1 if size > 0 else -1
status = HTTPStatus.OK
content_length = size
if range_header is not None:
try:
start_off, end_off_incl = self._parse_single_range(range_header, size)
except ValueError as e:
if str(e) == "unsatisfiable":
self._send_range_not_satisfiable(size)
return
if "unsatisfiable" in str(e):
self._send_range_not_satisfiable(size)
return
self._send_error_json(HTTPStatus.BAD_REQUEST, "invalid Range header")
return
status = HTTPStatus.PARTIAL_CONTENT
content_length = (end_off_incl - start_off) + 1
self.send_response(status)
self._send_imageio_headers()
self.send_header("Content-Type", "application/octet-stream")
self.send_header("Content-Length", str(content_length))
if status == HTTPStatus.PARTIAL_CONTENT:
self.send_header("Content-Range", f"bytes {start_off}-{end_off_incl}/{size}")
self.end_headers()
offset = start_off
end_excl = end_off_incl + 1
while offset < end_excl:
to_read = min(CHUNK_SIZE, end_excl - offset)
data = conn.pread(to_read, offset)
if not data:
raise RuntimeError("backend returned empty read")
try:
self.wfile.write(data)
except BrokenPipeError:
logging.info("GET client disconnected image_id=%s at=%d", image_id, offset)
break
offset += len(data)
bytes_sent += len(data)
except Exception as e:
# If headers already sent, we can't return JSON reliably; just log.
logging.warning("GET error image_id=%s err=%r", image_id, e)
try:
if not self.wfile.closed:
self.close_connection = True
except Exception:
pass
finally:
_READ_SEM.release()
lock.release()
dur = _now_s() - start
logging.info(
"GET end image_id=%s bytes=%d duration_s=%.3f", image_id, bytes_sent, dur
)
def _handle_put_image(self, image_id: str, cfg: Dict[str, Any], content_length: int) -> None:
lock = _get_image_lock(image_id)
if not lock.acquire(blocking=False):
self._send_error_json(HTTPStatus.CONFLICT, "image busy")
return
if not _WRITE_SEM.acquire(blocking=False):
lock.release()
self._send_error_json(HTTPStatus.SERVICE_UNAVAILABLE, "too many parallel writes")
return
start = _now_s()
bytes_written = 0
try:
logging.info("PUT start image_id=%s content_length=%d", image_id, content_length)
with _NbdConn(cfg["host"], int(cfg["port"]), cfg.get("export")) as conn:
size = conn.size()
if content_length != size:
self._send_error_json(
HTTPStatus.BAD_REQUEST,
f"Content-Length must equal image size ({size})",
)
return
offset = 0
remaining = content_length
while remaining > 0:
chunk = self.rfile.read(min(CHUNK_SIZE, remaining))
if not chunk:
self._send_error_json(
HTTPStatus.BAD_REQUEST,
f"request body ended early at {offset} bytes",
)
return
conn.pwrite(chunk, offset)
offset += len(chunk)
remaining -= len(chunk)
bytes_written += len(chunk)
# POC-level: do not auto-flush on PUT; expose explicit /flush endpoint.
self._send_json(HTTPStatus.OK, {"ok": True, "bytes_written": bytes_written})
except Exception as e:
logging.warning("PUT error image_id=%s err=%r", image_id, e)
self._send_error_json(HTTPStatus.INTERNAL_SERVER_ERROR, "backend error")
finally:
_WRITE_SEM.release()
lock.release()
dur = _now_s() - start
logging.info(
"PUT end image_id=%s bytes=%d duration_s=%.3f", image_id, bytes_written, dur
)
def _handle_get_extents(self, image_id: str, cfg: Dict[str, Any]) -> None:
# Keep deterministic and simple (POC): report entire image allocated.
# No per-image lock required by spec, but we still take it to avoid racing
# with a write and to keep behavior consistent.
lock = _get_image_lock(image_id)
if not lock.acquire(blocking=False):
self._send_error_json(HTTPStatus.CONFLICT, "image busy")
return
start = _now_s()
try:
logging.info("EXTENTS start image_id=%s", image_id)
with _NbdConn(cfg["host"], int(cfg["port"]), cfg.get("export")) as conn:
size = conn.size()
self._send_json(
HTTPStatus.OK,
[{"start": 0, "length": size, "allocated": True}],
)
except Exception as e:
logging.warning("EXTENTS error image_id=%s err=%r", image_id, e)
self._send_error_json(HTTPStatus.INTERNAL_SERVER_ERROR, "backend error")
finally:
lock.release()
dur = _now_s() - start
logging.info("EXTENTS end image_id=%s duration_s=%.3f", image_id, dur)
def _handle_post_flush(self, image_id: str, cfg: Dict[str, Any]) -> None:
lock = _get_image_lock(image_id)
if not lock.acquire(blocking=False):
self._send_error_json(HTTPStatus.CONFLICT, "image busy")
return
start = _now_s()
try:
logging.info("FLUSH start image_id=%s", image_id)
with _NbdConn(cfg["host"], int(cfg["port"]), cfg.get("export")) as conn:
conn.flush()
self._send_json(HTTPStatus.OK, {"ok": True})
except Exception as e:
logging.warning("FLUSH error image_id=%s err=%r", image_id, e)
self._send_error_json(HTTPStatus.INTERNAL_SERVER_ERROR, "backend error")
finally:
lock.release()
dur = _now_s() - start
logging.info("FLUSH end image_id=%s duration_s=%.3f", image_id, dur)
def main() -> None:
parser = argparse.ArgumentParser(description="POC imageio-like HTTP server backed by NBD")
parser.add_argument("--listen", default="127.0.0.1", help="Address to bind")
parser.add_argument("--port", type=int, default=54323, help="Port to listen on")
args = parser.parse_args()
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s %(levelname)s %(message)s",
)
addr = (args.listen, args.port)
httpd = ThreadingHTTPServer(addr, Handler)
logging.info("listening on http://%s:%d", args.listen, args.port)
logging.info("image configs are read from %s/<transferId>", _CFG_DIR)
httpd.serve_forever()
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink