Add TCP socket implementation for output stream. (#53)

6 files changed, 267 insertions, 5 deletions

diff --git a/src/track/CMakeLists.txt b/src/track/CMakeLists.txt
index 6ad255d..6dd29c5 100644
--- a/src/track/CMakeLists.txt
+++ b/src/track/CMakeLists.txt
@@ -20,6 +20,7 @@ add_library(heaptrack_preload MODULE
     libheaptrack.cpp
     outstream/outstream.cpp
     outstream/outstream_file.cpp
+    outstream/outstream_socket.cpp
 )
 
 target_compile_options(heaptrack_preload PRIVATE "-ftls-model=initial-exec")
@@ -48,6 +49,7 @@ add_library(heaptrack_inject MODULE
     libheaptrack.cpp
     outstream/outstream.cpp
     outstream/outstream_file.cpp
+    outstream/outstream_socket.cpp
 )
 
 target_link_libraries(heaptrack_inject LINK_PRIVATE
diff --git a/src/track/libheaptrack.cpp b/src/track/libheaptrack.cpp
index d50cc1d..f93acfb 100644
--- a/src/track/libheaptrack.cpp
+++ b/src/track/libheaptrack.cpp
@@ -1,4 +1,4 @@
-/*
+/*
  * Copyright 2014-2017 Milian Wolff <mail@milianw.de>
  *
  * This library is free software; you can redistribute it and/or
@@ -54,6 +54,7 @@
 #include "util/config.h"
 #include "util/libunwind_config.h"
 #include "outstream/outstream_file.h"
+#include "outstream/outstream_socket.h"
 
 /**
  * uncomment this to get extended debug code for known pointers
@@ -139,6 +140,32 @@ outStream* createFile(const char* fileName)
     } else if (outputFileName == "stderr") {
         debugLog<VerboseOutput>("%s", "will write to stderr");
         return OpenStream<outStreamFILE, FILE*>(stderr);
+    } else if (outputFileName == "socket") {
+        uint16_t Port = outStreamSOCKET::DefaultSocketPort;
+        char *env = nullptr;
+        env = getenv("DUMP_HEAPTRACK_SOCKET");
+        if (env) {
+            try {
+                int tmpPort = std::stoi(std::string(env));
+                if (tmpPort < outStreamSOCKET::MinAllowedSocketPort
+                    || tmpPort > outStreamSOCKET::MaxAllowedSocketPort) {
+                    fprintf(stderr, "WARNING: DUMP_HEAPTRACK_SOCKET socket port is out of allowed range.\n");
+                    throw std::out_of_range("DUMP_HEAPTRACK_SOCKET socket port is out of allowed range");
+                }
+                Port = static_cast<uint16_t>(tmpPort);
+            } catch (...) {
+                // do nothing, use default port
+                fprintf(stderr,
+                        "WARNING: DUMP_HEAPTRACK_SOCKET should be number in %i-%i range\n",
+                        outStreamSOCKET::MinAllowedSocketPort,
+                        outStreamSOCKET::MaxAllowedSocketPort);
+                fprintf(stderr, "WARNING: switched to default port %i\n",
+                        static_cast<int>(outStreamSOCKET::DefaultSocketPort));
+            }
+            unsetenv("DUMP_HEAPTRACK_SOCKET");
+        }
+        debugLog<VerboseOutput>("%s", "will write to socket");
+        return OpenStream<outStreamSOCKET, uint16_t>(Port);
     }
 
     if (outputFileName.empty()) {
diff --git a/src/track/outstream/outstream.cpp b/src/track/outstream/outstream.cpp
index ebc4175..3cdbd89 100644
--- a/src/track/outstream/outstream.cpp
+++ b/src/track/outstream/outstream.cpp
@@ -38,7 +38,7 @@ int fprintf(outStream *stream, const char* format, ...) noexcept
     }
 
     int ret = stream->Puts(Buf.get());
-    if (ret > 0) {
+    if (ret >= 0) {
         // make proper return code, since it different from fputs()
         ret = tmpStrSize;
     }
diff --git a/src/track/outstream/outstream_file.cpp b/src/track/outstream/outstream_file.cpp
index 7b32567..54ab758 100644
--- a/src/track/outstream/outstream_file.cpp
+++ b/src/track/outstream/outstream_file.cpp
@@ -42,7 +42,7 @@ int outStreamFILE::Putc(int Char) noexcept
 {

     if (!Stream_) {

         errno = EIO;

-        return -1;

+        return EOF;

     }

     return fputc(Char, Stream_);

 }

@@ -51,10 +51,10 @@ int outStreamFILE::Puts(const char *String) noexcept
 {

     if (!Stream_) {

         errno = EIO;

-        return -1;

+        return EOF;

     } else if (!String) {

         errno = EINVAL;

-        return -1;

+        return EOF;

     }

     return fputs(String, Stream_);

 }

diff --git a/src/track/outstream/outstream_socket.cpp b/src/track/outstream/outstream_socket.cpp
new file mode 100644
index 0000000..2311cbe
--- /dev/null
+++ b/src/track/outstream/outstream_socket.cpp
@@ -0,0 +1,165 @@
+#include <arpa/inet.h>

+#include <cassert>

+#include <netinet/in.h>

+#include <netinet/tcp.h>

+#include <stdexcept>

+#include <string.h>

+#include <unistd.h>

+

+#include "outstream_socket.h"

+

+outStreamSOCKET::outStreamSOCKET(uint16_t Port) :

+    Socket_(-1),

+    BufferUsedSize_(0),

+    Buffer_(new char[BufferCapacity_])

+{

+    int tmpSocketID = -1;

+    auto HandleError = [&tmpSocketID] (const char *ErrorText, int Error) {

+        if (tmpSocketID != -1) {

+            close(tmpSocketID);

+        }

+        fprintf(stderr, "WARNING! %s: %s\n", ErrorText, strerror(Error));

+        throw std::runtime_error(ErrorText);

+    };

+

+    tmpSocketID = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);

+    if (tmpSocketID == -1) {

+        HandleError("socket()", errno);

+    }

+

+    int on = 1;

+    if (setsockopt(tmpSocketID, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on)) == -1

+        || setsockopt(tmpSocketID, IPPROTO_TCP, TCP_NODELAY, &on, sizeof(on)) == -1) {

+        HandleError("setsockopt()", errno);

+    }

+

+    struct sockaddr_in tmpServerAddr;

+    memset((char*)&tmpServerAddr, 0, sizeof(tmpServerAddr));

+    tmpServerAddr.sin_family = AF_INET;

+    tmpServerAddr.sin_port = htons(Port);

+    tmpServerAddr.sin_addr.s_addr = INADDR_ANY;

+    if (bind(tmpSocketID, (struct sockaddr *) &tmpServerAddr, sizeof(tmpServerAddr)) == -1) {

+        HandleError("bind()", errno);

+    }

+

+    if (listen(tmpSocketID, 1) == -1) {

+        HandleError("listen()", errno);

+    }

+

+    struct sockaddr_storage tmpServerStorage;

+    socklen_t addr_size = sizeof tmpServerStorage;

+    Socket_ = accept(tmpSocketID, (struct sockaddr*)&tmpServerStorage, &addr_size);

+    if (Socket_ == -1) {

+        HandleError("accept()", errno);

+    }

+

+    close(tmpSocketID);

+}

+

+outStreamSOCKET::~outStreamSOCKET()

+{

+    if (Socket_ == -1) {

+        return;

+    }

+

+    FlushBuffer();

+    close(Socket_);

+}

+

+bool outStreamSOCKET::SocketErrorDetected() noexcept

+{

+    int SocketErrno;

+    unsigned int ErrnoSize = sizeof(SocketErrno);

+    if (getsockopt(Socket_, SOL_SOCKET, SO_ERROR, &SocketErrno, &ErrnoSize) == -1) {

+        return true;

+    }

+

+    if (SocketErrno) {

+        close(Socket_);

+        Socket_ = -1;

+        fprintf(stderr, "WARNING! Unable to use socket: %s", strerror(SocketErrno));

+        errno = SocketErrno;

+        return true;

+    }

+

+    return false;

+}

+

+bool outStreamSOCKET::BufferedWriteToSocket(const void *Data, size_t Count) noexcept

+{

+    if (Count > BufferCapacity_) {

+        if (!FlushBuffer()) {

+            return false;

+        }

+        return SendToSocket(Data, Count);

+    }

+

+    if (Count > AvailableSpace()) {

+        if (!FlushBuffer()) {

+            return false;

+        }

+    }

+

+    CopyToBuffer(Data, Count);

+    return true;

+}

+

+bool outStreamSOCKET::SendToSocket(const void *Data, size_t Count) noexcept

+{

+    if (SocketErrorDetected()) {

+        return false;

+    }

+    if (Count == 0) {

+        return true;

+    }

+    return send(Socket_, Data, Count, MSG_NOSIGNAL) != -1;

+}

+

+void outStreamSOCKET::CopyToBuffer(const void *Data, size_t Count) noexcept

+{

+    memcpy(BufferPos(), Data, Count);

+    BufferUsedSize_ += Count;

+}

+

+bool outStreamSOCKET::FlushBuffer() noexcept

+{

+    if (Socket_ == -1) {

+        errno = EIO;

+        return false;

+    }

+    bool ret = SendToSocket(Buffer_.get(), BufferUsedSize_);

+    BufferUsedSize_ = 0;

+    return ret;

+}

+

+int outStreamSOCKET::Putc(int Char) noexcept

+{

+    if (Socket_ == -1) {

+        errno = EIO;

+        return EOF;

+    }

+

+    // same behavior as for fputc()

+    unsigned char tmpChar = static_cast<unsigned char>(Char);

+    if (BufferedWriteToSocket(&tmpChar, sizeof(unsigned char)))

+        return Char;

+

+    return EOF;

+}

+

+int outStreamSOCKET::Puts(const char *String) noexcept

+{

+    if (Socket_ == -1) {

+        errno = EIO;

+        return EOF;

+    } else if (!String) {

+        errno = EINVAL;

+        return EOF;

+    }

+

+    // same behavior as for fputs()

+    if (BufferedWriteToSocket(String, strlen(String)))

+        return 1; // return a nonnegative number on success

+

+    return EOF;

+}

diff --git a/src/track/outstream/outstream_socket.h b/src/track/outstream/outstream_socket.h
new file mode 100644
index 0000000..c02de54
--- /dev/null
+++ b/src/track/outstream/outstream_socket.h
@@ -0,0 +1,68 @@
+#ifndef OUTSTREAMSOCKET_H
+#define OUTSTREAMSOCKET_H
+
+#include <memory>
+#include <cassert>
+#include "outstream.h"
+
+class outStreamSOCKET final : public outStream
+{
+public:
+    outStreamSOCKET() = delete;
+    explicit outStreamSOCKET(uint16_t Port);
+    ~outStreamSOCKET();
+
+    outStreamSOCKET(const outStreamSOCKET &) = delete;
+    outStreamSOCKET &operator = (const outStreamSOCKET &) = delete;
+
+    outStreamSOCKET(outStreamSOCKET &&other) :
+        Socket_(other.Socket_),
+        BufferUsedSize_(other.BufferUsedSize_),
+        Buffer_(std::move(other.Buffer_))
+    {
+        other.Socket_ = -1;
+        other.BufferUsedSize_ = 0;
+    }
+    outStreamSOCKET &operator = (outStreamSOCKET &&other) {
+        Socket_ = other.Socket_;
+        other.Socket_ = -1;
+        BufferUsedSize_ = other.BufferUsedSize_;
+        other.BufferUsedSize_ = 0;
+        Buffer_ = std::move(other.Buffer_);
+        return *this;
+    }
+
+    int Putc(int Char) noexcept override;
+    int Puts(const char *String) noexcept override;
+
+    static constexpr int MinAllowedSocketPort = 1;
+    static constexpr int MaxAllowedSocketPort = 65535;
+    static constexpr uint16_t DefaultSocketPort = 5050;
+
+private:
+    bool SocketErrorDetected() noexcept;
+    bool BufferedWriteToSocket(const void *Data, size_t Count) noexcept;
+    bool SendToSocket(const void *Data, size_t Count) noexcept;
+    void CopyToBuffer(const void *Data, size_t Count) noexcept;
+    bool FlushBuffer() noexcept;
+
+    size_t AvailableSpace() const noexcept
+    {
+        assert(BufferCapacity_ >= BufferUsedSize_);
+        return BufferCapacity_ - BufferUsedSize_;
+    }
+
+    char *BufferPos() const noexcept
+    {
+        assert(BufferCapacity_ >= BufferUsedSize_);
+        return Buffer_.get() + BufferUsedSize_;
+    }
+
+    int Socket_;
+    // mainly, heaptrack send small blocks (about 1-50 bytes each)
+    static constexpr size_t BufferCapacity_ = 4096;
+    size_t BufferUsedSize_;
+    std::unique_ptr<char[]> Buffer_;
+};
+
+#endif // OUTSTREAMSOCKET_H