websocket-rpc-test.c++ (4592B)
1 // Copyright (c) 2021 Ian Denhardt and contributors 2 // Licensed under the MIT License: 3 // 4 // Permission is hereby granted, free of charge, to any person obtaining a copy 5 // of this software and associated documentation files (the "Software"), to deal 6 // in the Software without restriction, including without limitation the rights 7 // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 8 // copies of the Software, and to permit persons to whom the Software is 9 // furnished to do so, subject to the following conditions: 10 // 11 // The above copyright notice and this permission notice shall be included in 12 // all copies or substantial portions of the Software. 13 // 14 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 17 // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 19 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 20 // THE SOFTWARE. 21 22 #include "websocket-rpc.h" 23 #include <kj/test.h> 24 25 #include <capnp/test.capnp.h> 26 27 KJ_TEST("WebSocketMessageStream") { 28 kj::EventLoop loop; 29 kj::WaitScope waitScope(loop); 30 31 auto pipe = kj::newWebSocketPipe(); 32 33 auto msgStreamA = capnp::WebSocketMessageStream(*pipe.ends[0]); 34 auto msgStreamB = capnp::WebSocketMessageStream(*pipe.ends[1]); 35 36 // Make a message, fill it with some stuff 37 capnp::MallocMessageBuilder originalMsg; 38 auto object = originalMsg.initRoot<capnproto_test::capnp::test::TestAllTypes>().initStructList(10); 39 object[0].setTextField("Test"); 40 object[1].initStructField().setTextField("A string"); 41 object[2].setTextField("Another field"); 42 object[3].setInt64Field(42); 43 auto originalSegments = originalMsg.getSegmentsForOutput(); 44 45 // Send the message across the websocket, make sure it comes out unharmed. 46 auto writePromise = msgStreamA.writeMessage(nullptr, originalSegments); 47 msgStreamB.tryReadMessage(nullptr) 48 .then([&](auto maybeResult) -> kj::Promise<void> { 49 KJ_IF_MAYBE(result, maybeResult) { 50 KJ_ASSERT(result->fds.size() == 0); 51 KJ_ASSERT(result->reader->getSegment(originalSegments.size()) == nullptr); 52 for(size_t i = 0; i < originalSegments.size(); i++) { 53 auto oldSegment = originalSegments[i]; 54 auto newSegment = result->reader->getSegment(i); 55 56 KJ_ASSERT(oldSegment.size() == newSegment.size()); 57 KJ_ASSERT(memcmp( 58 &oldSegment[0], 59 &newSegment[0], 60 oldSegment.size() * sizeof(capnp::word) 61 ) == 0); 62 } 63 return kj::READY_NOW; 64 } else { 65 KJ_FAIL_ASSERT("Reading first message failed"); 66 } 67 }).wait(waitScope); 68 writePromise.wait(waitScope); 69 70 // Close the websocket, and make sure the other end gets nullptr when reading. 71 auto endPromise = msgStreamA.end(); 72 msgStreamB.tryReadMessage(nullptr).then([](auto maybe) -> kj::Promise<void> { 73 KJ_IF_MAYBE(segments, maybe) { 74 KJ_FAIL_ASSERT("Should have gotten nullptr after websocket was closed"); 75 } 76 return kj::READY_NOW; 77 }).wait(waitScope); 78 endPromise.wait(waitScope); 79 } 80 81 KJ_TEST("WebSocketMessageStreamByteCount") { 82 kj::EventLoop loop; 83 kj::WaitScope waitScope(loop); 84 85 auto pipe1 = kj::newWebSocketPipe(); 86 auto pipe2 = kj::newWebSocketPipe(); 87 88 auto msgStreamA = capnp::WebSocketMessageStream(*pipe1.ends[0]); 89 auto msgStreamB = capnp::WebSocketMessageStream(*pipe2.ends[1]); 90 91 auto pumpTask = pipe1.ends[1]->pumpTo(*pipe2.ends[0]); 92 93 capnp::MallocMessageBuilder originalMsg; 94 auto object = originalMsg.initRoot<capnproto_test::capnp::test::TestAllTypes>().initStructList(10); 95 object[0].setTextField("Test"); 96 object[1].initStructField().setTextField("A string"); 97 object[2].setTextField("Another field"); 98 object[3].setInt64Field(42); 99 auto originalSegments = originalMsg.getSegmentsForOutput(); 100 101 auto writePromise = msgStreamA.writeMessage(nullptr, originalSegments); 102 msgStreamB.tryReadMessage(nullptr).wait(waitScope); 103 writePromise.wait(waitScope); 104 105 auto endPromise = msgStreamA.end(); 106 msgStreamB.tryReadMessage(nullptr).wait(waitScope); 107 pumpTask.wait(waitScope); 108 endPromise.wait(waitScope); 109 KJ_EXPECT(pipe1.ends[0]->sentByteCount() == 2585); 110 KJ_EXPECT(pipe1.ends[1]->receivedByteCount() == 2585); 111 KJ_EXPECT(pipe2.ends[0]->sentByteCount() == 2585); 112 KJ_EXPECT(pipe2.ends[1]->receivedByteCount() == 2585); 113 }