1
// Copyright (c) 2017-2022 Cloudflare, Inc.
2
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
//     https://opensource.org/licenses/Apache-2.0
4
 
5
#include "channel-token.h"
6
 
7
#include <capnp/message.h>
8
#include <kj/test.h>
9
 
10
namespace workerd::server {
11
namespace {
12
 
13
kj::String strProps(const Frankenvalue& props) {
14
  capnp::MallocMessageBuilder message;
15
  auto builder = message.getRoot<rpc::Frankenvalue>();
16
  props.threadSafeClone().toCapnp(builder);
17
  return kj::str(builder.asReader());
18
}
19
 
20
struct ServiceTriplet {
21
  kj::String serviceName;
22
  kj::Maybe<kj::String> entrypoint;
23
  Frankenvalue props;
24
 
25
  ServiceTriplet(kj::StringPtr serviceName, kj::Maybe<kj::StringPtr> entrypoint, Frankenvalue props)
26
      : serviceName(kj::str(serviceName)),
27
        entrypoint(entrypoint.map([](kj::StringPtr s) { return kj::str(s); })),
28
        props(kj::mv(props)) {}
29
  ServiceTriplet(ServiceTriplet&&) = default;
30
 
31
  bool operator==(const ServiceTriplet& other) const {
32
    return serviceName == other.serviceName && entrypoint == other.entrypoint &&
33
        strProps(props) == strProps(other.props);
34
  }
35
 
36
  kj::String toString() const {
37
    return kj::str('(', serviceName, ", ", entrypoint, ", ", strProps(props), ')');
38
  }
39
};
40
 
41
class MockSubrequestChannel: public IoChannelFactory::SubrequestChannel {
42
 public:
43
  MockSubrequestChannel(ServiceTriplet triplet): triplet(kj::mv(triplet)) {}
44
  ServiceTriplet triplet;
45
 
46
  kj::Own<WorkerInterface> startRequest(IoChannelFactory::SubrequestMetadata metadata) override {
47
    KJ_UNREACHABLE;
48
  }
49
  void requireAllowsTransfer() override {
50
    KJ_UNREACHABLE;
51
  }
52
};
53
 
54
class MockActorClassChannel: public IoChannelFactory::ActorClassChannel {
55
 public:
56
  MockActorClassChannel(ServiceTriplet triplet): triplet(kj::mv(triplet)) {}
57
  ServiceTriplet triplet;
58
 
59
  void requireAllowsTransfer() override {
60
    KJ_UNREACHABLE;
61
  }
62
};
63
 
64
class MockResolver: public ChannelTokenHandler::Resolver {
65
 public:
66
  kj::Own<IoChannelFactory::SubrequestChannel> resolveEntrypoint(
67
      kj::StringPtr serviceName, kj::Maybe<kj::StringPtr> entrypoint, Frankenvalue props) override {
68
    return kj::refcounted<MockSubrequestChannel>(
69
        ServiceTriplet(serviceName, entrypoint, kj::mv(props)));
70
  }
71
 
72
  kj::Own<IoChannelFactory::ActorClassChannel> resolveActorClass(
73
      kj::StringPtr serviceName, kj::Maybe<kj::StringPtr> entrypoint, Frankenvalue props) override {
74
    return kj::refcounted<MockActorClassChannel>(
75
        ServiceTriplet(serviceName, entrypoint, kj::mv(props)));
76
  }
77
};
78
 
79
using Usage = IoChannelFactory::ChannelTokenUsage;
80
 
81
KJ_TEST("channel token basics") {
82
  MockResolver resolver;
83
  ChannelTokenHandler handler(resolver);
84
 
85
  auto props = Frankenvalue::fromJson(kj::str("{\"foo\": 123}"));
86
  auto token = handler.encodeSubrequestChannelToken(Usage::RPC, "foo", "MyEntry"_kj, props);
87
 
88
  // Decoding works.
89
  {
90
    auto channel =
91
        handler.decodeSubrequestChannelToken(Usage::RPC, token).downcast<MockSubrequestChannel>();
92
    KJ_EXPECT(channel->triplet == ServiceTriplet("foo", "MyEntry"_kj, props.clone()));
93
  }
94
 
95
  auto corruptedToken = [&](uint index) {
96
    auto copy = kj::heapArray(token.asPtr());
97
    copy[index] ^= 1;
98
    return copy;
99
  };
100
 
101
  // Corrupting any byte of the token should make it invalid.
102
  {
103
    // Corrupt the magic number.
104
    KJ_EXPECT_THROW_MESSAGE(
105
        "RPC_TOKEN_MAGIC", handler.decodeSubrequestChannelToken(Usage::RPC, corruptedToken(2)));
106
 
107
    // Corrupt the MAC.
108
    KJ_EXPECT_THROW_MESSAGE("failed authentication",
109
        handler.decodeSubrequestChannelToken(Usage::RPC, corruptedToken(token.size() - 2)));
110
 
111
    // Corrupt the IV.
112
    KJ_EXPECT_THROW_MESSAGE("failed authentication",
113
        handler.decodeSubrequestChannelToken(Usage::RPC, corruptedToken(7)));
114
 
115
    // Corrupt the key ID.
116
    KJ_EXPECT_THROW_MESSAGE("failed authentication",
117
        handler.decodeSubrequestChannelToken(Usage::RPC, corruptedToken(20)));
118
 
119
    // Corrupt the message body.
120
    KJ_EXPECT_THROW_MESSAGE("failed authentication",
121
        handler.decodeSubrequestChannelToken(Usage::RPC, corruptedToken(37)));
122
  }
123
 
124
  // Can't parse as a storage token.
125
  KJ_EXPECT_THROW_MESSAGE(
126
      "STORAGE_TOKEN_MAGIC", handler.decodeSubrequestChannelToken(Usage::STORAGE, token));
127
 
128
  // Can't use as wrong type.
129
  KJ_EXPECT_THROW_MESSAGE(
130
      "channel token type mismatch", handler.decodeActorClassChannelToken(Usage::RPC, token));
131
}
132
 
133
KJ_TEST("channel tokens for storage") {
134
  MockResolver resolver;
135
  ChannelTokenHandler handler(resolver);
136
 
137
  auto props = Frankenvalue::fromJson(kj::str("{\"foo\": 123}"));
138
  auto token = handler.encodeSubrequestChannelToken(Usage::STORAGE, "foo", "MyEntry"_kj, props);
139
 
140
  // Decoding works.
141
  {
142
    auto channel = handler.decodeSubrequestChannelToken(Usage::STORAGE, token)
143
                       .downcast<MockSubrequestChannel>();
144
    KJ_EXPECT(channel->triplet == ServiceTriplet("foo", "MyEntry"_kj, props.clone()));
145
  }
146
 
147
  auto corruptedToken = [&](uint index) {
148
    auto copy = kj::heapArray(token.asPtr());
149
    copy[index] ^= 1;
150
    return copy;
151
  };
152
 
153
  // Corrupting the magic number breaks the token.
154
  KJ_EXPECT_THROW_MESSAGE("STORAGE_TOKEN_MAGIC",
155
      handler.decodeSubrequestChannelToken(Usage::STORAGE, corruptedToken(2)));
156
 
157
  // Can't parse as an RPC token.
158
  KJ_EXPECT_THROW_MESSAGE(
159
      "RPC_TOKEN_MAGIC", handler.decodeSubrequestChannelToken(Usage::RPC, token));
160
 
161
  // Can't use as wrong type.
162
  KJ_EXPECT_THROW_MESSAGE(
163
      "channel token type mismatch", handler.decodeActorClassChannelToken(Usage::STORAGE, token));
164
}
165
 
166
KJ_TEST("actor class channel tokens") {
167
  MockResolver resolver;
168
  ChannelTokenHandler handler(resolver);
169
 
170
  auto props = Frankenvalue::fromJson(kj::str("{\"foo\": 123}"));
171
  auto token = handler.encodeActorClassChannelToken(Usage::RPC, "foo", "MyEntry"_kj, props);
172
 
173
  // Decoding works.
174
  {
175
    auto channel =
176
        handler.decodeActorClassChannelToken(Usage::RPC, token).downcast<MockActorClassChannel>();
177
    KJ_EXPECT(channel->triplet == ServiceTriplet("foo", "MyEntry"_kj, props.clone()));
178
  }
179
 
180
  // Decoding as the wrong type fails.
181
  KJ_EXPECT_THROW_MESSAGE(
182
      "channel token type mismatch", handler.decodeSubrequestChannelToken(Usage::RPC, token));
183
}
184
 
185
}  // namespace
186
}  // namespace workerd::server