diff --git a/wpiutil/CMakeLists.txt b/wpiutil/CMakeLists.txt index 4a66a3a05a..94e177ec4e 100644 --- a/wpiutil/CMakeLists.txt +++ b/wpiutil/CMakeLists.txt @@ -134,3 +134,7 @@ endif() file(GLOB netconsoleServer_src src/netconsoleServer/native/cpp/*.cpp) add_executable(netconsoleServer ${netconsoleServer_src}) target_link_libraries(netconsoleServer wpiutil ${LIBUTIL}) + +file(GLOB netconsoleTee_src src/netconsoleTee/native/cpp/*.cpp) +add_executable(netconsoleTee ${netconsoleTee_src}) +target_link_libraries(netconsoleTee wpiutil) diff --git a/wpiutil/build.gradle b/wpiutil/build.gradle index 41fcd7b263..e987a90fff 100644 --- a/wpiutil/build.gradle +++ b/wpiutil/build.gradle @@ -167,5 +167,19 @@ model { } } } + netconsoleTee(NativeExecutableSpec) { + targetBuildTypes 'release' + sources { + cpp { + source { + srcDirs = ['src/netconsoleTee/native/cpp'] + includes = ['**/*.cpp'] + } + } + } + binaries.all { binary -> + lib project: ':wpiutil', library: 'wpiutil', linkage: 'static' + } + } } } diff --git a/wpiutil/src/netconsoleTee/native/cpp/main.cpp b/wpiutil/src/netconsoleTee/native/cpp/main.cpp new file mode 100644 index 0000000000..d14b322bbc --- /dev/null +++ b/wpiutil/src/netconsoleTee/native/cpp/main.cpp @@ -0,0 +1,188 @@ +/*----------------------------------------------------------------------------*/ +/* Copyright (c) 2018 FIRST. All Rights Reserved. */ +/* Open Source Software - may be modified and shared by FRC teams. The code */ +/* must be accompanied by the FIRST BSD license file in the root directory of */ +/* the project. */ +/*----------------------------------------------------------------------------*/ + +#include "wpi/MathExtras.h" +#include "wpi/SmallVector.h" +#include "wpi/raw_ostream.h" +#include "wpi/raw_uv_ostream.h" +#include "wpi/timestamp.h" +#include "wpi/uv/Loop.h" +#include "wpi/uv/Tcp.h" +#include "wpi/uv/Tty.h" +#include "wpi/uv/Udp.h" +#include "wpi/uv/util.h" + +namespace uv = wpi::uv; + +static uint64_t startTime = wpi::Now(); + +static bool NewlineBuffer(std::string& rem, uv::Buffer& buf, size_t len, + wpi::SmallVectorImpl& bufs, bool tcp, + uint16_t tcpSeq) { + // scan for last newline + wpi::StringRef str(buf.base, len); + size_t idx = str.rfind('\n'); + if (idx == wpi::StringRef::npos) { + // no newline yet, just keep appending to remainder + rem += str; + return false; + } + + // build output + wpi::raw_uv_ostream out(bufs, 4096); + wpi::StringRef toCopy = str.slice(0, idx + 1); + if (tcp) { + // Header is 2 byte len, 1 byte type, 4 byte timestamp, 2 byte sequence num + uint32_t ts = wpi::FloatToBits((wpi::Now() - startTime) * 1.0e-6); + uint16_t len = rem.size() + toCopy.size() + 1 + 4 + 2; + out << wpi::ArrayRef({static_cast((len >> 8) & 0xff), + static_cast(len & 0xff), 12, + static_cast((ts >> 24) & 0xff), + static_cast((ts >> 16) & 0xff), + static_cast((ts >> 8) & 0xff), + static_cast(ts & 0xff), + static_cast((tcpSeq >> 8) & 0xff), + static_cast(tcpSeq & 0xff)}); + } + out << rem << toCopy; + + // reset remainder + rem = str.slice(idx + 1, wpi::StringRef::npos); + return true; +} + +static void CopyUdp(uv::Stream& in, std::shared_ptr out, + bool broadcast) { + sockaddr_in addr; + if (broadcast) { + out->SetBroadcast(true); + uv::NameToAddr("0.0.0.0", 6666, &addr); + } else { + uv::NameToAddr("127.0.0.1", 6666, &addr); + } + + in.data.connect( + [ rem = std::make_shared(), outPtr = out.get(), addr ]( + uv::Buffer & buf, size_t len) { + // build buffers + wpi::SmallVector bufs; + if (!NewlineBuffer(*rem, buf, len, bufs, false, 0)) return; + + // send output + outPtr->Send(addr, bufs, [](auto bufs2, uv::Error) { + for (auto buf : bufs2) buf.Deallocate(); + }); + }, + out); +} + +static void CopyTcp(uv::Stream& in, std::shared_ptr out) { + struct StreamData { + std::string rem; + uint16_t seq = 0; + }; + in.data.connect([ data = std::make_shared(), outPtr = out.get() ]( + uv::Buffer & buf, size_t len) { + // build buffers + wpi::SmallVector bufs; + if (!NewlineBuffer(data->rem, buf, len, bufs, true, data->seq++)) return; + + // send output + outPtr->Write(bufs, [](auto bufs2, uv::Error) { + for (auto buf : bufs2) buf.Deallocate(); + }); + }, + out); +} + +static void CopyStream(uv::Stream& in, std::shared_ptr out) { + in.data.connect([out](uv::Buffer& buf, size_t len) { + uv::Buffer buf2 = buf.Dup(); + buf2.len = len; + out->Write(buf2, [](auto bufs, uv::Error) { + for (auto buf : bufs) buf.Deallocate(); + }); + }); +} + +int main(int argc, char* argv[]) { + // parse arguments + int arg = 1; + bool useUdp = false; + bool broadcastUdp = false; + bool err = false; + + while (arg < argc && argv[arg][0] == '-') { + if (wpi::StringRef(argv[arg]) == "-u") { + useUdp = true; + } else if (wpi::StringRef(argv[arg]) == "-b") { + useUdp = true; + broadcastUdp = true; + } else { + wpi::errs() << "unrecognized command line option " << argv[arg] << '\n'; + err = true; + } + ++arg; + } + + if (err) { + wpi::errs() + << argv[0] << " [-ub]\n" + << " -u send udp to localhost port 6666 instead of using tcp\n" + << " -b broadcast udp to port 6666 instead of using tcp\n"; + return EXIT_FAILURE; + } + + auto loop = uv::Loop::Create(); + loop->error.connect( + [](uv::Error err) { wpi::errs() << "uv ERROR: " << err.str() << '\n'; }); + + // create ttys + auto stdinTty = uv::Tty::Create(loop, 0, true); + auto stdoutTty = uv::Tty::Create(loop, 1, false); + + // don't bother continuing if we don't have a stdin + if (!stdinTty) return EXIT_SUCCESS; + + // pass through our input to output + if (stdoutTty) CopyStream(*stdinTty, stdoutTty); + + // when our stdin closes, exit + stdinTty->end.connect([] { std::exit(EXIT_SUCCESS); }); + + if (useUdp) { + auto udp = uv::Udp::Create(loop); + // tee + CopyUdp(*stdinTty, udp, broadcastUdp); + } else { + auto tcp = uv::Tcp::Create(loop); + + // bind to listen address and port + tcp->Bind("", 1740); + + // when we get a connection, accept it + tcp->connection.connect([ srv = tcp.get(), stdinTty ] { + auto tcp = srv->Accept(); + if (!tcp) return; + + // close on error + tcp->error.connect([s = tcp.get()](wpi::uv::Error err) { s->Close(); }); + + // tee + CopyTcp(*stdinTty, tcp); + }); + + // start listening for incoming connections + tcp->Listen(); + } + + // start reading + if (stdinTty) stdinTty->StartRead(); + + // run the loop! + loop->Run(); +}