Spaces:
Running on Zero
Running on Zero
| import json | |
| import os | |
| import sys | |
| import tempfile | |
| import time | |
| import unittest | |
| from pathlib import Path | |
| from types import SimpleNamespace | |
| from unittest import mock | |
| ROOT = Path(__file__).resolve().parent.parent | |
| sys.path.insert(0, str(ROOT)) | |
| os.environ["LABS_FALLBACK_SIGNING_KEY"] = "unit-signing-key" | |
| os.environ.pop("SPACE_ID", None) | |
| import access_limits as al # noqa: E402 | |
| import quota_fallback as qf # noqa: E402 | |
| PARENT = qf.EMBED_PARENTS[0] | |
| SPACE_HOST = "broyang-3dai.hf.space" | |
| def request(referer="", forwarded="", host=SPACE_HOST, client="10.0.0.9"): | |
| headers = {"host": host} | |
| if referer: | |
| headers["referer"] = referer | |
| if forwarded: | |
| headers["x-forwarded-for"] = forwarded | |
| return SimpleNamespace(headers=headers, client=SimpleNamespace(host=client)) | |
| def embed_request(**kwargs): | |
| return request(referer=f"https://{SPACE_HOST}/?__theme=light&__labs_parent={PARENT}", **kwargs) | |
| def public_request(forwarded="203.0.113.7"): | |
| return request(referer="https://3dai.skyebrowse.com/", forwarded=forwarded) | |
| def opened(ticket): | |
| return json.loads(qf._b64url_decode(ticket.split(".")[0])) | |
| def grant_for(nonce, mode="visitor", **overrides): | |
| payload = {"v": 1, "purpose": "grant", "space": qf.space_identifier(), "nonce": nonce, | |
| "mode": mode, "exp": int(time.time()) + 180} | |
| payload.update(overrides) | |
| return qf.sign_document(payload) | |
| class GrantBookTests(unittest.TestCase): | |
| def setUp(self): | |
| self.book = al.GrantBook() | |
| def issue(self, session="session-a"): | |
| receipt = self.book.request(session) | |
| return receipt, opened(receipt["ticket"])["nonce"] | |
| def test_grant_requests_are_signed_grant_operations_that_cannot_unlock_the_fallback(self): | |
| receipt, nonce = self.issue() | |
| claim = opened(receipt["ticket"]) | |
| self.assertEqual((claim["operation"], claim["space"]), (qf.GRANT, qf.space_identifier())) | |
| self.assertRegex(nonce, r"^[0-9a-f]{32}$") | |
| self.assertIsNone(qf.verify_receipt(receipt["ticket"])) | |
| def test_verified_grants_run_natively_once(self): | |
| receipt, nonce = self.issue() | |
| self.book.settle("session-a", {"ticket": receipt["ticket"], "grant": grant_for(nonce)}, receipt) | |
| access = self.book.take("session-a") | |
| self.assertEqual((access.mode, access.nonce, access.notice), ("native", nonce, None)) | |
| self.assertEqual(self.book.take("session-a").mode, "denied") | |
| def test_forged_mismatched_or_missing_grants_deny_the_generation(self): | |
| receipt, nonce = self.issue() | |
| forged = [ | |
| None, | |
| "not-a-grant", | |
| grant_for("b" * 32), | |
| grant_for(nonce, purpose="execute"), | |
| grant_for(nonce, mode="unlimited"), | |
| grant_for(nonce, mode="shared"), | |
| grant_for(nonce, space="broyang/other"), | |
| grant_for(nonce, exp=int(time.time()) - 1), | |
| ] | |
| for grant in forged: | |
| receipt, nonce = self.issue() | |
| self.book.settle("session-a", {"ticket": receipt["ticket"], "grant": grant}, receipt) | |
| self.assertEqual(self.book.take("session-a").mode, "denied") | |
| receipt, nonce = self.issue() | |
| self.book.settle("session-b", {"ticket": receipt["ticket"], "grant": grant_for(nonce)}, receipt) | |
| self.assertEqual(self.book.take("session-a").mode, "denied") | |
| self.assertEqual(self.book.take("session-b").mode, "denied") | |
| def test_limit_refusal_shows_the_upgrade_message(self): | |
| receipt, _nonce = self.issue() | |
| message = "You have used your 3 free 3D AI generations. Upgrade to a Commercial plan to keep generating." | |
| self.book.settle("session-a", {"ticket": receipt["ticket"], "reason": al.FREE_LIMIT_REASON, | |
| "error": message}, receipt) | |
| access = self.book.take("session-a") | |
| self.assertEqual(access.mode, "denied") | |
| self.assertIn("3 free 3D AI generations", access.notice) | |
| def test_stale_grant_requests_expire(self): | |
| receipt, nonce = self.issue() | |
| with mock.patch.object(al.time, "time", return_value=time.time() + al.GRANT_TTL_SECONDS + 1): | |
| self.book.settle("session-a", {"ticket": receipt["ticket"], "grant": grant_for(nonce)}, receipt) | |
| self.assertEqual(self.book.take("session-a").mode, "denied") | |
| def test_grant_requests_are_bounded_per_session_and_in_total(self): | |
| for _ in range(al.MAX_GRANTS_PER_SESSION): | |
| self.assertIsNotNone(self.book.request("session-a")) | |
| self.assertIsNone(self.book.request("session-a")) | |
| self.assertIsNotNone(self.book.request("session-b")) | |
| book = al.GrantBook(max_entries=1) | |
| self.assertIsNotNone(book.request("one")) | |
| self.assertIsNone(book.request("two")) | |
| def settle(self, receipt, nonce, session="session-a"): | |
| self.book.settle(session, {"ticket": receipt["ticket"], "grant": grant_for(nonce)}, receipt) | |
| def test_a_settled_grant_survives_a_long_gpu_queue(self): | |
| receipt, nonce = self.issue() | |
| self.settle(receipt, nonce) | |
| with mock.patch.object(al.time, "time", return_value=time.time() + al.GRANT_TTL_SECONDS + 60): | |
| access = self.book.take("session-a") | |
| self.assertEqual((access.mode, access.nonce), ("native", nonce)) | |
| def test_a_lapsed_grant_is_denied_with_its_nonce_for_a_refund(self): | |
| receipt, nonce = self.issue() | |
| self.settle(receipt, nonce) | |
| with mock.patch.object(al.time, "time", return_value=time.time() + al.SETTLED_GRANT_TTL_SECONDS + 1): | |
| access = self.book.take("session-a") | |
| self.assertEqual((access.mode, access.nonce), ("denied", nonce)) | |
| self.assertIsNone(self.book.take("session-a").nonce) | |
| def test_overlapping_runs_in_one_session_each_keep_their_grant(self): | |
| first, first_nonce = self.issue() | |
| second, second_nonce = self.issue() | |
| self.settle(first, first_nonce) | |
| self.settle(second, second_nonce) | |
| self.assertEqual(self.book.take("session-a").nonce, first_nonce) | |
| self.assertEqual(self.book.take("session-a").nonce, second_nonce) | |
| self.assertEqual(self.book.take("session-a").mode, "denied") | |
| def test_overlapping_runs_keep_their_own_refusals(self): | |
| first, first_nonce = self.issue() | |
| second, _nonce = self.issue() | |
| self.settle(first, first_nonce) | |
| self.book.settle("session-a", {"ticket": second["ticket"], "reason": al.FREE_LIMIT_REASON, "error": ""}, second) | |
| self.assertEqual(self.book.take("session-a").nonce, first_nonce) | |
| self.assertIn(al.EMBED_LIMIT_MESSAGE, self.book.take("session-a").notice) | |
| def test_an_ack_for_an_unknown_nonce_marks_no_other_grant(self): | |
| receipt, nonce = self.issue() | |
| stranger = {"ticket": qf.issue_ticket("e" * 32, qf.GRANT)} | |
| self.book.settle("session-a", {"ticket": stranger["ticket"], "grant": grant_for("e" * 32)}, stranger) | |
| self.settle(receipt, nonce) | |
| self.assertEqual(self.book.take("session-a").nonce, nonce) | |
| class PublicAllowanceTests(unittest.TestCase): | |
| def setUp(self): | |
| self.directory = tempfile.TemporaryDirectory() | |
| self.addCleanup(self.directory.cleanup) | |
| self.path = Path(self.directory.name) / "usage.json" | |
| def allowance(self, limit=3, window=3600): | |
| return al.PublicAllowance(limit, window, self.path) | |
| def test_three_uses_per_visitor_then_refusal_and_refunds(self): | |
| allowance = self.allowance() | |
| key = allowance.key_for(public_request()) | |
| self.assertEqual([allowance.try_consume(key) for _ in range(4)], [True, True, True, False]) | |
| allowance.refund(key) | |
| self.assertTrue(allowance.try_consume(key)) | |
| other = allowance.key_for(public_request("198.51.100.4")) | |
| self.assertNotEqual(key, other) | |
| self.assertTrue(allowance.try_consume(other)) | |
| def test_counts_survive_a_reload_and_never_store_the_raw_address(self): | |
| allowance = self.allowance() | |
| key = allowance.key_for(public_request()) | |
| for _ in range(3): | |
| allowance.try_consume(key) | |
| self.assertNotIn("203.0.113.7", self.path.read_text()) | |
| self.assertFalse(self.allowance().try_consume(key)) | |
| def test_uses_age_out_of_the_window(self): | |
| allowance = self.allowance(window=60) | |
| key = allowance.key_for(public_request()) | |
| for _ in range(3): | |
| allowance.try_consume(key) | |
| with mock.patch.object(al.time, "time", return_value=time.time() + 61): | |
| self.assertTrue(allowance.try_consume(key)) | |
| def test_first_forwarded_address_identifies_the_visitor(self): | |
| allowance = self.allowance() | |
| self.assertEqual(allowance.key_for(public_request("203.0.113.7, 10.1.1.1")), | |
| allowance.key_for(public_request("203.0.113.7"))) | |
| def test_a_missing_forwarded_address_never_shares_one_bucket(self): | |
| allowance = self.allowance() | |
| with self.assertLogs(al.logger, "WARNING"): | |
| key = allowance.key_for(public_request("")) | |
| self.assertIsNone(key) | |
| self.assertTrue(all(allowance.try_consume(key) for _ in range(5))) | |
| class AccessPolicyTests(unittest.TestCase): | |
| def setUp(self): | |
| self.directory = tempfile.TemporaryDirectory() | |
| self.addCleanup(self.directory.cleanup) | |
| self.policy = al.AccessPolicy(al.GrantBook(), al.PublicAllowance(3, 3600, Path(self.directory.name) / "u.json")) | |
| def test_public_browser_visits_get_three_generations_then_a_signup_link(self): | |
| modes = [self.policy.authorize(public_request(), "s").mode for _ in range(4)] | |
| self.assertEqual(modes, ["native", "native", "native", "denied"]) | |
| notice = self.policy.authorize(public_request(), "s").notice | |
| self.assertIn('href="https://app.skyebrowse.com/register"', notice) | |
| self.assertIn("Sign up free", notice) | |
| def test_server_side_callers_without_a_browser_referer_are_not_counted(self): | |
| for _ in range(5): | |
| self.assertEqual(self.policy.authorize(request(forwarded="203.0.113.7"), "s").mode, "native") | |
| self.assertEqual(self.policy.public.used(self.policy.public.key_for(public_request())), 0) | |
| def test_embedded_generations_need_a_grant_only_when_grants_are_enabled(self): | |
| with mock.patch.dict(os.environ, {"LABS_EMBED_GRANTS": ""}): | |
| self.assertEqual(self.policy.authorize(embed_request(), "s").mode, "native") | |
| with mock.patch.dict(os.environ, {"LABS_EMBED_GRANTS": "1"}): | |
| self.assertEqual(self.policy.authorize(embed_request(), "s").mode, "denied") | |
| receipt = self.policy.grants.request("s") | |
| nonce = opened(receipt["ticket"])["nonce"] | |
| self.policy.grants.settle("s", {"ticket": receipt["ticket"], "grant": grant_for(nonce)}, receipt) | |
| self.assertEqual(self.policy.authorize(embed_request(), "s").nonce, nonce) | |
| self.assertEqual(self.policy.public.used(self.policy.public.key_for(embed_request())), 0) | |
| def test_grants_can_be_enforced_for_listed_parents_only(self): | |
| with mock.patch.dict(os.environ, {"LABS_EMBED_GRANTS": f"http://localhost:8080, {PARENT}"}): | |
| self.assertTrue(al.grants_required(embed_request())) | |
| self.assertFalse(al.grants_required(public_request())) | |
| with mock.patch.dict(os.environ, {"LABS_EMBED_GRANTS": "http://localhost:8080"}): | |
| self.assertFalse(al.grants_required(embed_request())) | |
| def test_quota_denial_refunds_the_public_use_and_points_to_signup(self): | |
| access = self.policy.authorize(public_request(), "s") | |
| notice = self.policy.quota_denied(access) | |
| self.assertIn("app.skyebrowse.com/register", notice) | |
| self.assertEqual(self.policy.public.used(access.visitor_key), 0) | |
| self.assertIsNone(self.policy.quota_denied(al.Access())) | |
| def test_failed_generations_issue_one_signed_refund_and_return_the_public_use(self): | |
| public = self.policy.authorize(public_request(), "p") | |
| self.policy.generation_failed("p", public) | |
| self.assertEqual(self.policy.public.used(public.visitor_key), 0) | |
| self.assertIsNone(self.policy.refund_receipt("p")) | |
| self.policy.generation_failed("e", al.Access(nonce="c" * 32)) | |
| refund = self.policy.refund_receipt("e") | |
| claim = opened(refund["ticket"]) | |
| self.assertEqual((claim["operation"], claim["nonce"]), (qf.REFUND, "c" * 32)) | |
| self.assertIsNone(qf.verify_receipt(refund["ticket"])) | |
| self.assertIsNone(self.policy.refund_receipt("e")) | |
| def grant(self, session="e"): | |
| receipt = self.policy.grants.request(session) | |
| nonce = opened(receipt["ticket"])["nonce"] | |
| self.policy.grants.settle(session, {"ticket": receipt["ticket"], "grant": grant_for(nonce)}, receipt) | |
| return nonce | |
| def test_overlapping_failed_generations_are_each_refunded(self): | |
| with mock.patch.dict(os.environ, {"LABS_EMBED_GRANTS": "1"}): | |
| nonces = [self.grant(), self.grant()] | |
| for _ in nonces: | |
| access = self.policy.authorize(embed_request(), "e") | |
| self.assertEqual(access.mode, "native") | |
| self.policy.generation_failed("e", access) | |
| self.assertEqual([opened(self.policy.refund_receipt("e")["ticket"])["nonce"] for _ in nonces], nonces) | |
| self.assertIsNone(self.policy.refund_receipt("e")) | |
| def test_a_lapsed_grant_is_refunded_when_its_run_is_denied(self): | |
| with mock.patch.dict(os.environ, {"LABS_EMBED_GRANTS": "1"}): | |
| nonce = self.grant() | |
| with mock.patch.object(al.time, "time", return_value=time.time() + al.SETTLED_GRANT_TTL_SECONDS + 1): | |
| self.assertEqual(self.policy.authorize(embed_request(), "e").mode, "denied") | |
| self.assertEqual(opened(self.policy.refund_receipt("e")["ticket"])["nonce"], nonce) | |
| def test_notices_are_delivered_once(self): | |
| self.policy.remember("s", "<div>notice</div>") | |
| self.assertEqual(self.policy.pop_notice("s"), "<div>notice</div>") | |
| self.assertIsNone(self.policy.pop_notice("s")) | |
| class EmbedParentTests(unittest.TestCase): | |
| def test_only_localhost_extras_extend_the_parent_allow_list(self): | |
| extras = "http://localhost:8109, https://evil.example, http://localhost:8109, javascript:alert(1)" | |
| with mock.patch.dict(os.environ, {"LABS_EXTRA_EMBED_PARENTS": extras}): | |
| self.assertEqual(qf.embed_parents(), qf.EMBED_PARENTS + ("http://localhost:8109",)) | |
| for name in ("grant_bridge.js", "refund_bridge.js"): | |
| bridge = al.bridge_script(ROOT / name) | |
| self.assertIn('"http://localhost:8109"', bridge) | |
| self.assertNotIn("evil.example", bridge) | |
| self.assertNotIn("__LABS_EMBED_PARENTS__", al.bridge_script(ROOT / "fallback_bridge.js")) | |
| if __name__ == "__main__": | |
| unittest.main() | |