1
#pragma once
2
 
3
#include <capnp/membrane.h>
4
 
5
namespace workerd {
6
 
7
// A membrane applied which detects when no capabilities are held any longer, at which point it
8
// fulfills a fulfiller.
9
//
10
// TODO(cleanup): This is generally useful, should it be part of capnp?
11
class CompletionMembrane final: public capnp::MembranePolicy, public kj::Refcounted {
12
 public:
13
  explicit CompletionMembrane(kj::Own<kj::PromiseFulfiller<void>> doneFulfiller)
14
      : doneFulfiller(kj::mv(doneFulfiller)) {}
15
  ~CompletionMembrane() noexcept(false) {
16
    doneFulfiller->fulfill();
17
  }
18
 
19
  kj::Maybe<capnp::Capability::Client> inboundCall(
20
      uint64_t interfaceId, uint16_t methodId, capnp::Capability::Client target) override {
21
    return kj::none;
22
  }
23
 
24
  kj::Maybe<capnp::Capability::Client> outboundCall(
25
      uint64_t interfaceId, uint16_t methodId, capnp::Capability::Client target) override {
26
    return kj::none;
27
  }
28
 
29
  kj::Own<MembranePolicy> addRef() override {
30
    return kj::addRef(*this);
31
  }
32
 
33
  void reject(kj::Exception&& e) {
34
    doneFulfiller->reject(kj::mv(e));
35
  }
36
 
37
 private:
38
  kj::Own<kj::PromiseFulfiller<void>> doneFulfiller;
39
};
40
 
41
// A membrane which revokes when some Promise is fulfilled.
42
//
43
// TODO(cleanup): This is generally useful, should it be part of capnp?
44
class RevokerMembrane final: public capnp::MembranePolicy, public kj::Refcounted {
45
 public:
46
  explicit RevokerMembrane(kj::Promise<void> promise): promise(promise.fork()) {}
47
 
48
  kj::Maybe<capnp::Capability::Client> inboundCall(
49
      uint64_t interfaceId, uint16_t methodId, capnp::Capability::Client target) override {
50
    return kj::none;
51
  }
52
 
53
  kj::Maybe<capnp::Capability::Client> outboundCall(
54
      uint64_t interfaceId, uint16_t methodId, capnp::Capability::Client target) override {
55
    return kj::none;
56
  }
57
 
58
  kj::Own<MembranePolicy> addRef() override {
59
    return kj::addRef(*this);
60
  }
61
 
62
  kj::Maybe<kj::Promise<void>> onRevoked() override {
63
    return promise.addBranch();
64
  }
65
 
66
 private:
67
  kj::ForkedPromise<void> promise;
68
};
69
 
70
}  // namespace workerd