Repository: trailofbits/polytracker (language: python)

Issue:
DFSAN wrapper for `read` does not assign a return value label
https://github.com/trailofbits/polytracker/blob/master/polytracker/src/taint_sources/taint_sources.cpp#L143-L153
```c++
EXT_C_FUNC ssize_t __dfsw_read(int fd, void *buff, size_t size,
                               dfsan_label fd_label, dfsan_label buff_label,
                               dfsan_label size_label, dfsan_label *ret_label) {
  long read_start = lseek(fd, 0, SEEK_CUR);
  ssize_t ret_val = read(fd, buff, size);

  if (ret_val > 0)
    get_polytracker_tdag().source_taint(fd, buff, read_start, ret_val);

  return ret_val;
}
```

Does not assign to the `ret_label` causing it to return with an arbitrary value. Follow-on errors include
```
[2022-11-09 12:04:17.377] [error] Decoding invalid taint value, 0 == 0 encoded: 0
```


Add support for tracking taints from socket reads, similar to how we handle STDIN



Discussion / hints:
I have a fix pending in the upcoming PR for #6488.
Rough plan:
1. Add test case for stdin
2. Implemente stdin
3. Add test case for socket reads
4. Implement socket reads
5. Consider if taint sources shall be refactored to support multiple distinct types of taint-sources (consider comment in #6490).

Fix patch:
diff --git a/polytracker/custom_abi/dfsan_abilist.txt b/polytracker/custom_abi/dfsan_abilist.txt
index ff0e59a4..018931e8 100644
--- a/polytracker/custom_abi/dfsan_abilist.txt
+++ b/polytracker/custom_abi/dfsan_abilist.txt
@@ -170,7 +170,7 @@ fun:bind=discard
 fun:chdir=discard
 # fun:close=discard
 fun:closedir=discard
-fun:connect=discard
+# fun:connect=discard
 fun:dladdr=discard
 fun:dlclose=discard
 # fun:fclose=discard
@@ -1683,6 +1683,7 @@ fun:abort=uninstrumented
 fun:abs=uninstrumented
 fun:abs=functional
 fun:accept=uninstrumented
+fun:accept=custom
 fun:accept4=uninstrumented
 fun:access=uninstrumented
 fun:acct=uninstrumented
@@ -1890,6 +1891,7 @@ fun:conj=uninstrumented
 fun:conjf=uninstrumented
 fun:conjl=uninstrumented
 fun:connect=uninstrumented
+fun:connect=custom
 fun:copysign=uninstrumented
 fun:copysignf=uninstrumented
 fun:copysignl=uninstrumented
@@ -2241,9 +2243,13 @@ fun:getaliasent=uninstrumented
 fun:getaliasent_r=uninstrumented
 fun:getauxval=uninstrumented
 fun:getc=uninstrumented
+fun:getc=custom
 fun:getc_unlocked=uninstrumented
+fun:getc_unlocked=custom
 fun:getchar=uninstrumented
+fun:getchar=custom
 fun:getchar_unlocked=uninstrumented
+fun:getchar_unlocked=custom
 fun:getcontext=uninstrumented
 fun:getcwd=uninstrumented
 fun:getdate=uninstrumented
@@ -3126,7 +3132,9 @@ fun:realpath=uninstrumented
 fun:realpath=functional
 fun:reboot=uninstrumented
 fun:recv=uninstrumented
+fun:recv=custom
 fun:recvfrom=uninstrumented
+fun:recvfrom=custom
 fun:recvmmsg=uninstrumented
 fun:recvmsg=uninstrumented
 fun:regcomp=uninstrumented
diff --git a/polytracker/include/taintdag/polytracker.h b/polytracker/include/taintdag/polytracker.h
index a65390f2..ba7d7060 100644
--- a/polytracker/include/taintdag/polytracker.h
+++ b/polytracker/include/taintdag/polytracker.h
@@ -16,7 +16,9 @@
 #include "taintdag/fntrace.h"
 #include "taintdag/labels.h"
 #include "taintdag/sink.h"
+#include "taintdag/stream_offset.h"
 #include "taintdag/string_table.h"
+#include "taintdag/taint.h"
 #include "taintdag/taint_source.h"
 
 namespace taintdag {
@@ -78,6 +80,11 @@ class PolyTracker {
       OutputFile<Sources, Labels, StringTable, TaintSink,
                  SourceLabelIndexSection, Functions, Events>;
   ConcreteOutputFile output_file_;
+
+  // Tracking source offsets for streams (where offsets can be determined by
+  // ftell) per source index.
+  static constexpr size_t offset_capacity = size_t{max_source_index} + 1;
+  StreamOffset<offset_capacity> stream_offsets_;
 };
 
 } // namespace taintdag
diff --git a/polytracker/include/taintdag/stream_offset.h b/polytracker/include/taintdag/stream_offset.h
new file mode 100644
index 00000000..cfd17585
--- /dev/null
+++ b/polytracker/include/taintdag/stream_offset.h
@@ -0,0 +1,42 @@
+/*
+ * Copyright (c) 2022-present, Trail of Bits, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed in accordance with the terms specified in
+ * the LICENSE file found in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <array>
+#include <atomic>
+
+#include "taintdag/error.h"
+#include "taintdag/taint.h"
+
+namespace taintdag {
+
+// Helper class that tracks reads per source for a predetermined number of
+// sources, controlled by SourceCount.
+template <size_t SourceCount> class StreamOffset {
+public:
+  // Records a read of len bytes from the source at index idx. Returns the
+  // sequential offset.
+  // NOTE: Uses atomic fetch-add to update the offsets. It is assumed that
+  // concurrent read operations for the same source is the result of invalid
+  // code. Reads on the same source should be protected from running in
+  // parallell (at least for the taint sources considered here, e.g. files).
+  source_offset_t read(source_index_t idx, size_t len) {
+    if (idx >= SourceCount) {
+      error_exit("Attemted read of source index ", static_cast<uint64_t>(idx),
+                 ", only ", SourceCount, " sources available");
+    }
+    return offset[idx].fetch_add(len);
+  }
+
+private:
+  using atomic_t = std::atomic<source_offset_t>;
+  std::array<atomic_t, SourceCount> offset{0};
+};
+
+} // namespace taintdag
\ No newline at end of file
diff --git a/polytracker/include/taintdag/taint_source.h b/polytracker/include/taintdag/taint_source.h
index df44f14b..c55eec05 100644
--- a/polytracker/include/taintdag/taint_source.h
+++ b/polytracker/include/taintdag/taint_source.h
@@ -22,7 +22,7 @@ namespace taintdag {
 // it is reused to represent sinks as well.
 struct SourceEntry {
   // Representing a size that cannot be determined (e.g. for a socket).
-  static constexpr uint64_t InvalidSize = ~uint64_t(0);
+  static constexpr uint64_t InvalidSize = std::numeric_limits<uint64_t>::max();
 
   // Use this when there is no fd representing this source (e.g. argv)
   static constexpr int InvalidFD = -1;
diff --git a/polytracker/src/polytracker/main.cpp b/polytracker/src/polytracker/main.cpp
index da2269cb..9f095409 100644
--- a/polytracker/src/polytracker/main.cpp
+++ b/polytracker/src/polytracker/main.cpp
@@ -1,4 +1,3 @@
-// #include "polytracker/write_taints.h"
 #include <atomic>
 #include <fcntl.h>
 #include <fstream>
@@ -26,6 +25,7 @@ DECLARE_EARLY_CONSTRUCT(taintdag::PolyTracker, polytracker_tdag);
 DECLARE_EARLY_CONSTRUCT(std::string, polytracker_db_name);
 DECLARE_EARLY_CONSTRUCT(std::string, polytracker_stderr_sink);
 DECLARE_EARLY_CONSTRUCT(std::string, polytracker_stdout_sink);
+DECLARE_EARLY_CONSTRUCT(std::string, polytracker_stdin_source);
 
 // Controls argv being a taint source
 bool polytracker_taint_argv = false;
@@ -69,13 +69,17 @@ void polytracker_parse_env() {
     get_polytracker_stderr_sink() = err;
   }
 
+  if (auto in = getenv("POLYTRACKER_STDIN_SOURCE")) {
+    get_polytracker_stdin_source() = in;
+  }
+
   if (auto argv = getenv("POLYTRACKER_TAINT_ARGV")) {
     polytracker_taint_argv = argv[0] == '1';
   }
 }
 
 /*
-This code parses the enviornment variables and sets the globals which work as
+This code parses the environment variables and sets the globals which work as
 polytrackers settings
 
 1. Parse config if exists
@@ -86,6 +90,7 @@ void polytracker_get_settings() {
   DO_EARLY_DEFAULT_CONSTRUCT(std::string, polytracker_db_name)
   DO_EARLY_DEFAULT_CONSTRUCT(std::string, polytracker_stderr_sink);
   DO_EARLY_DEFAULT_CONSTRUCT(std::string, polytracker_stdout_sink);
+  DO_EARLY_DEFAULT_CONSTRUCT(std::string, polytracker_stdin_source);
   polytracker_parse_env();
   set_defaults();
 }
@@ -114,6 +119,11 @@ void polytracker_print_settings() {
     printf("POLYTRACKER_STDERR_SINK: %s\n",
            get_polytracker_stderr_sink().c_str());
   }
+  // stdin source flag
+  if (!get_polytracker_stdin_source().empty()) {
+    printf("POLYTRACKER_STDIN_SOURCE: %s\n",
+           get_polytracker_stdin_source().c_str());
+  }
   if (polytracker_taint_argv) {
     printf("POLYTRACKER_TAINT_ARGV: 1\n");
   }
@@ -130,12 +140,20 @@ void sink_streams() {
   }
 }
 
+// Use stdin as a taint source, if indicated by env var
+void stdin_source() {
+  if (int f = fileno(stdin); f >= 0 && get_polytracker_stdin_source() == "1") {
+    get_polytracker_tdag().open_file(f, "/dev/stdin");
+  }
+}
+
 void taint_start(void) {
   polytracker_get_settings();
   polytracker_print_settings();
   DO_EARLY_CONSTRUCT(taintdag::PolyTracker, polytracker_tdag,
                      get_polytracker_db_name());
   sink_streams();
+  stdin_source();
   // Set up the atexit call
   atexit(polytracker_end);
 }
\ No newline at end of file
diff --git a/polytracker/src/taint_sources/taint_sources.cpp b/polytracker/src/taint_sources/taint_sources.cpp
index 6e86294b..b9045d84 100644
--- a/polytracker/src/taint_sources/taint_sources.cpp
+++ b/polytracker/src/taint_sources/taint_sources.cpp
@@ -2,9 +2,11 @@
 #include "polytracker/early_construct.h"
 #include "polytracker/polytracker.h"
 
+#include "taintdag/error.h"
 #include "taintdag/polytracker.h"
 
 #include <algorithm>
+#include <arpa/inet.h>
 #include <assert.h>
 #include <fcntl.h>
 #include <iostream>
@@ -19,6 +21,7 @@
 #include <string.h>
 #include <string>
 #include <sys/mman.h>
+#include <sys/socket.h>
 #include <sys/stat.h>
 #include <thread>
 #include <time.h>
@@ -149,6 +152,7 @@ EXT_C_FUNC ssize_t __dfsw_read(int fd, void *buff, size_t size,
   if (ret_val > 0)
     get_polytracker_tdag().source_taint(fd, buff, read_start, ret_val);
 
+  *ret_label = 0;
   return ret_val;
 }
 
@@ -237,6 +241,35 @@ EXT_C_FUNC int __dfsw__IO_getc(FILE *fd, dfsan_label fd_label,
   return c;
 }
 
+EXT_C_FUNC int __dfsw_getc(FILE *fd, dfsan_label fd_label,
+                           dfsan_label *ret_label) {
+  long offset = ftell(fd);
+  int c = getc(fd);
+  *ret_label = 0;
+  if (c != EOF) {
+    auto tr =
+        get_polytracker_tdag().source_taint(fileno(fd), offset, sizeof(char));
+    if (tr) {
+      *ret_label = tr.value().first;
+    }
+  }
+  return c;
+}
+
+EXT_C_FUNC int __dfsw_getc_unlocked(FILE *fd, dfsan_label fd_label,
+                                    dfsan_label *ret_label) {
+  long offset = ftell(fd);
+  int c = getc_unlocked(fd);
+  *ret_label = 0;
+  if (c != EOF) {
+    auto tr =
+        get_polytracker_tdag().source_taint(fileno(fd), offset, sizeof(char));
+    if (tr)
+      *ret_label = tr.value().first;
+  }
+  return c;
+}
+
 EXT_C_FUNC int __dfsw_getchar(dfsan_label *ret_label) {
   long offset = ftell(stdin);
   int c = getchar();
@@ -250,6 +283,19 @@ EXT_C_FUNC int __dfsw_getchar(dfsan_label *ret_label) {
   return c;
 }
 
+EXT_C_FUNC int __dfsw_getchar_unlocked(dfsan_label *ret_label) {
+  long offset = ftell(stdin);
+  int c = getchar_unlocked();
+  *ret_label = 0;
+  if (c != EOF) {
+    auto tr = get_polytracker_tdag().source_taint(fileno(stdin), offset,
+                                                  sizeof(char));
+    if (tr)
+      *ret_label = tr.value().first;
+  }
+  return c;
+}
+
 EXT_C_FUNC char *__dfsw_fgets(char *str, int count, FILE *fd,
                               dfsan_label str_label, dfsan_label count_label,
                               dfsan_label fd_label, dfsan_label *ret_label) {
@@ -356,4 +402,107 @@ EXT_C_FUNC int __dfsw_pthread_cond_broadcast(pthread_cond_t *cond,
 
 EXT_C_FUNC void __dfsw_exit(int ret_code, dfsan_label ret_code_label) {
   exit(ret_code);
+}
+
+// Socket functions
+namespace {
+static std::optional<std::string> connect_name(int socket) {
+  sockaddr_in local_addr, remote_addr;
+  socklen_t local_len{sizeof(local_addr)}, remote_len{sizeof(remote_addr)};
+  char local_str[64], remote_str[64];
+
+  if (int ret = getsockname(socket, reinterpret_cast<sockaddr *>(&local_addr),
+                            &local_len);
+      ret != 0) {
+    taintdag::error_exit("Failed to get sockname for socket ", socket);
+  }
+
+  // The only supported address family for now is AF_INET
+  if (local_addr.sin_family != AF_INET)
+    return {};
+
+  if (int ret = getpeername(socket, reinterpret_cast<sockaddr *>(&remote_addr),
+                            &remote_len);
+      ret != 0) {
+    taintdag::error_exit("Failed to get peername for socket ", socket);
+  }
+
+  if (!inet_ntop(AF_INET, &(local_addr.sin_addr), local_str,
+                 sizeof(local_str))) {
+    taintdag::error_exit("inet_ntop failed for remote addr");
+  }
+
+  if (!inet_ntop(AF_INET, &(remote_addr.sin_addr), remote_str,
+                 sizeof(remote_str))) {
+    taintdag::error_exit("inet_ntop failed for remote addr");
+  }
+
+  std::stringstream strm;
+  strm << "socket:" << local_str << ":" << ntohs(local_addr.sin_port) << "-"
+       << remote_str << ":" << ntohs(remote_addr.sin_port);
+
+  return strm.str();
+}
+} // namespace
+
+EXT_C_FUNC int __dfsw_accept(int socket, struct sockaddr *address,
+                             socklen_t *address_len, dfsan_label socket_label,
+                             dfsan_label address_label,
+                             dfsan_label address_len_label,
+                             dfsan_label *ret_label) {
+  int client_socket = accept(socket, address, address_len);
+  if (client_socket >= 0) {
+    if (auto name = connect_name(client_socket); name) {
+      get_polytracker_tdag().open_file(client_socket, *name);
+    }
+  }
+
+  *ret_label = 0;
+  return client_socket;
+}
+
+EXT_C_FUNC int __dfsw_connect(int socket, const struct sockaddr *address,
+                              socklen_t address_len, dfsan_label socket_label,
+                              dfsan_label address_label,
+                              dfsan_label address_len_label,
+                              dfsan_label *ret_label) {
+  int status = connect(socket, address, address_len);
+  if (status == 0) {
+    if (auto name = connect_name(socket); name) {
+      get_polytracker_tdag().open_file(socket, *name);
+    }
+  }
+
+  *ret_label = 0;
+  return status;
+}
+
+EXT_C_FUNC ssize_t __dfsw_recv(int socket, void *buff, size_t length, int flags,
+                               dfsan_label socket_label, dfsan_label buff_label,
+                               dfsan_label length_label,
+                               dfsan_label flags_label,
+                               dfsan_label *ret_label) {
+  ssize_t ret_val = recv(socket, buff, length, flags);
+
+  if (ret_val > 0)
+    get_polytracker_tdag().source_taint(socket, buff, -1, ret_val);
+
+  *ret_label = 0;
+  return ret_val;
+}
+
+EXT_C_FUNC ssize_t __dfsw_recvfrom(
+    int socket, void *buffer, size_t length, int flags,
+    struct sockaddr *address, socklen_t *address_len, dfsan_label socket_label,
+    dfsan_label buffer_label, dfsan_label length_label, dfsan_label flags_label,
+    dfsan_label address_label, dfsan_label address_len_label,
+    dfsan_label *ret_label) {
+  ssize_t ret_val =
+      recvfrom(socket, buffer, length, flags, address, address_len);
+
+  if (ret_val > 0)
+    get_polytracker_tdag().source_taint(socket, buffer, -1, ret_val);
+
+  *ret_label = 0;
+  return ret_val;
 }
\ No newline at end of file
diff --git a/polytracker/src/taintdag/polytracker.cpp b/polytracker/src/taintdag/polytracker.cpp
index 77247016..9e40d325 100644
--- a/polytracker/src/taintdag/polytracker.cpp
+++ b/polytracker/src/taintdag/polytracker.cpp
@@ -70,29 +70,40 @@ taint_range_t PolyTracker::create_source_taint(source_index_t src,
 //
 // If fd is not tracked as a taint source, no labels will be assigned to the
 // corresponding mem.
+// NOTE: If offset < 0, the assumption is made that offsets aren't known and
+// per source tracking should be used. This typically happens for stdin, sockets
+// and other streams.
 std::optional<taint_range_t> PolyTracker::source_taint(int fd, void const *mem,
                                                        source_offset_t offset,
                                                        size_t length) {
   if (auto source_index = output_file_.section<Sources>().mapping_idx(fd)) {
-    return this->create_source_taint(
-        *source_index, {reinterpret_cast<uint8_t const *>(mem), length},
-        offset);
+    if (offset < 0) {
+      offset = stream_offsets_.read(*source_index, length);
+    }
+    return create_source_taint(*source_index,
+                               {reinterpret_cast<uint8_t const *>(mem), length},
+                               offset);
   }
-  // Not tracked as taint source
   return {};
 }
 
 // Introduce source taint when reading from taint source fd.
 //
-// If fd is not tracked as a taint source, no labels will be assigned to the
-// corresponding mem.
+// If fd is not tracked as a taint source, no action will be taken.
+// NOTE: If offset < 0, the assumption is made that offsets aren't known and
+// per source tracking should be used. This typically happens for stdin.
 std::optional<taint_range_t>
 PolyTracker::source_taint(int fd, source_offset_t offset, size_t length) {
   if (auto source_index = output_file_.section<Sources>().mapping_idx(fd)) {
-    return this->output_file_.section<Labels>().create_source_labels(
+    if (offset < 0) {
+      offset = stream_offsets_.read(*source_index, length);
+    }
+    auto range = output_file_.section<Labels>().create_source_labels(
         *source_index, offset, length);
+    output_file_.section<SourceLabelIndexSection>().set_range(
+        BitIndex{range.first}, BitCount{length});
+    return range;
   }
-  // Not tracked as taint source
   return {};
 }
 


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index c519aa2e..6f5f9df8 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -48,24 +48,39 @@ def input_file(tmp_path):
 
 
 @pytest.fixture
-def program_trace(input_file, monkeypatch, request):
-    # Run everything in a per-test temporary directory
-    monkeypatch.chdir(input_file.parent)
-    # Build a clean test binary to get a blight journal
+def target_source(request):
+    """Locates the target source file to instrument"""
     marker = request.node.get_closest_marker("program_trace")
     tstdir = Path(request.fspath).parent
-    target = tstdir / Path(marker.args[0])
-    binary = Path(f"{target.stem}.bin").resolve()
-    build(target, binary)
-    # Build an instrumented test binary
-    trace_file = Path(f"{target.stem}.db").resolve()
-    trace_file.unlink(missing_ok=True)
+    return tstdir / Path(marker.args[0])
+
+
+@pytest.fixture
+def instrumented_binary(tmp_path, monkeypatch, target_source):
+    """Instruments the target source and returns the instrumented binary path"""
+    monkeypatch.chdir(tmp_path)
+    binary = Path(f"{target_source.stem}.bin").resolve()
+    build(target_source, binary)
     instrument(binary.name)
-    # Run the instrumented binary to get a trace file
+    return Path(f"{binary.stem}.instrumented").resolve()
+
+
+@pytest.fixture
+def trace_file(target_source):
+    """Produces a path to the polytracker DB given a target_source"""
+    dbpath = Path(f"{target_source.stem}.db").resolve()
+    dbpath.unlink(missing_ok=True)
+    return dbpath
+
+
+@pytest.fixture
+def program_trace(input_file, trace_file, instrumented_binary, monkeypatch):
+    # Run everything in a per-test temporary directory
+    monkeypatch.chdir(input_file.parent)
     monkeypatch.setenv("POLYDB", str(trace_file))
     cmd = [
         # instrumented binary
-        Path(f"{binary.stem}.instrumented").resolve(),
+        instrumented_binary,
         # input data
         str(input_file),
     ]
diff --git a/tests/test_socket_read.cpp b/tests/test_socket_read.cpp
new file mode 100644
index 00000000..709964e8
--- /dev/null
+++ b/tests/test_socket_read.cpp
@@ -0,0 +1,117 @@
+#include <cassert>
+#include <limits>
+#include <string>
+#include <string_view>
+
+#include <arpa/inet.h>
+#include <sys/socket.h>
+#include <unistd.h>
+
+
+int read_source_data(int conn, std::string_view expected_data) {
+
+  int ret;
+  char buffer[7]{0};
+
+  // Receive source taint via read
+  ret = read(conn, buffer, 2);
+  assert(ret == 2);
+
+  // Receive source taint via recv
+  ret = recv(conn, &buffer[2], 2, 0);
+  assert(ret == 2);
+
+  // Receive source taint via recvfrom
+  ret = recvfrom(conn, &buffer[4], 2, 0, nullptr, nullptr);
+  assert(ret == 2);
+  // Ensure we received the correct data
+  assert(expected_data == buffer);
+
+  return 0;
+}
+
+int client(uint16_t port, std::string_view expected_data) {
+  int ret;
+  auto s = socket(PF_INET, SOCK_STREAM, 0);
+  assert(s >= 0);
+
+  sockaddr_in server_address {.sin_family = AF_INET, .sin_port = htons(port)};
+  ret = inet_pton(AF_INET, "127.0.0.1", &server_address.sin_addr);
+  assert(ret == 1);
+
+  ret = connect(s, reinterpret_cast<sockaddr const *>(&server_address), sizeof(server_address));
+  assert(ret == 0);
+
+  ret = read_source_data(s, expected_data);
+  assert(ret == 0);
+
+  ret = shutdown(s, SHUT_RDWR);
+  assert(ret == 0);
+  ret = close(s);
+  assert(ret == 0);
+  return 0;
+}
+
+
+int server(uint16_t port, std::string_view expected_data) {
+  int ret;
+  auto s = socket(PF_INET, SOCK_STREAM, 0);
+  assert(s >= 0);
+
+  const int enable = 1;
+  ret = setsockopt(s, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(enable));
+  assert(ret == 0);
+
+  const int disable = 0;
+  struct linger linger {.l_linger = 0, .l_onoff = 0};
+  ret = setsockopt(s, SOL_SOCKET, SO_LINGER, &linger, sizeof(linger));
+  assert(ret == 0);
+
+  sockaddr_in server_address {.sin_family = AF_INET, .sin_port = htons(port)};
+  ret = inet_pton(AF_INET, "127.0.0.1", &server_address.sin_addr);
+  assert(ret == 1);
+
+  ret = bind(s, reinterpret_cast<sockaddr const*>(&server_address), sizeof(server_address));
+  assert(ret == 0);
+
+  ret = listen(s, 5);
+  assert(ret == 0);
+
+  int client_socket = accept(s, nullptr, nullptr); // Don't care about client address
+  assert(client_socket >= 0);
+
+  ret = read_source_data(client_socket, expected_data);
+  assert(ret == 0);
+
+  // Close the client connection
+  ret = shutdown(client_socket, SHUT_RDWR);
+  assert(ret == 0);
+  ret = close(client_socket);
+  assert(ret == 0);
+
+  // Close the listening socket
+  ret = shutdown(s, SHUT_RDWR);
+  assert(ret == 0);
+  ret = close(s);
+  assert(ret == 0);
+  return 0;
+}
+
+
+int main(int argc, char* argv[]) {
+  assert(argc == 4);
+
+  std::string_view mode = argv[1];
+  auto port = std::stoul(argv[2]);
+  assert(port <= std::numeric_limits<uint16_t>::max());
+
+  std::string_view expected_data = argv[3];
+
+  if (mode == "client") {
+    return client(static_cast<uint16_t>(port), expected_data);
+  } else if (mode == "server") {
+    return server(static_cast<uint16_t>(port), expected_data);
+  } else {
+    return -1;
+  }
+}
\ No newline at end of file
diff --git a/tests/test_socket_read.py b/tests/test_socket_read.py
new file mode 100644
index 00000000..ffd8439a
--- /dev/null
+++ b/tests/test_socket_read.py
@@ -0,0 +1,93 @@
+import pytest
+import socket
+import subprocess
+import time
+
+import polytracker
+
+from pathlib import Path
+
+# The maximum duration for the test case (to prevent hangs)
+MAX_DURATION_SECS = 10
+
+
+def expected_source_name(binary_ip, binary_port, python_ip, python_port):
+    """Return taint source name given connection info"""
+    return f"socket:{binary_ip}:{binary_port}-{python_ip}:{python_port}"
+
+
+@pytest.mark.program_trace("test_socket_read.cpp")
+@pytest.mark.parametrize("mode", ["client", "server"])
+def test_socket_read(instrumented_binary: Path, trace_file: Path, mode: str):
+    """Tests reads from the socket, introducing source taint.
+
+    The mode is seen from the instrumented binary point of view. That is, client
+    is when the instrumented binary connects to this python server and server is
+    vice versa.
+    """
+    start = time.time()
+
+    # Data to write to stdin, one byte at a time
+    send_data = "abcdef".encode("utf-8")
+    port_number = 63455
+    host = "127.0.0.1"
+
+    def start_proc():
+        return subprocess.Popen(
+            [str(instrumented_binary), mode, str(port_number), send_data],
+            env={"POLYDB": str(trace_file)},
+        )
+
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+
+        if mode == "client":
+            s.bind((host, port_number))
+            s.listen()
+
+            proc = start_proc()
+
+            conn, addr = s.accept()
+            conn.sendall(send_data)
+            conn.shutdown(socket.SHUT_RDWR)
+
+            source_name = expected_source_name(addr[0], addr[1], host, port_number)
+
+        else:
+            assert mode == "server"
+            proc = start_proc()
+
+            # Retry connections until timeout
+            while True:
+                try:
+                    elapsed = int(time.time() - start)
+                    assert elapsed < MAX_DURATION_SECS
+                    s.connect((host, port_number))
+                    break
+                except ConnectionRefusedError:
+                    print("Retry connect")
+
+            client_host, client_port = s.getsockname()
+            source_name = expected_source_name(
+                host, port_number, client_host, client_port
+            )
+            s.sendall(send_data)
+            s.shutdown(socket.SHUT_RDWR)
+
+    # Wait for process to terminate, should be quick unless error
+    elapsed = int(time.time() - start)
+    if elapsed > MAX_DURATION_SECS:
+        wait_time = 0
+    else:
+        wait_time = MAX_DURATION_SECS - elapsed
+    proc.wait(wait_time)
+    assert proc.returncode == 0
+
+    program_trace = polytracker.PolyTrackerTrace.load(trace_file)
+    input_labels = list(program_trace.tdfile.input_labels())
+    inputs = list(program_trace.inputs)
+
+    assert len(input_labels) == len(send_data)
+
+    assert len(inputs) == 1
+    assert inputs[0].path == source_name
diff --git a/tests/test_stdin.cpp b/tests/test_stdin.cpp
new file mode 100644
index 00000000..c853091a
--- /dev/null
+++ b/tests/test_stdin.cpp
@@ -0,0 +1,122 @@
+#include <cassert>
+#include <cstdint>
+#include <unistd.h>
+#include <string>
+
+int stdin_read() {
+  char inbyte;
+  int sum = 0;
+
+  while (true) {
+    ssize_t ret = read(STDIN_FILENO, &inbyte, sizeof(inbyte));
+    if (ret == sizeof(inbyte)) {
+      sum += inbyte;
+    } else if (ret == 0) {
+      break;
+    }
+  }
+  return sum;
+}
+
+int stdin_fread() {
+  char inbyte;
+  int sum = 0;
+
+  while (!feof(stdin)) {
+    size_t ret = fread(&inbyte, sizeof(inbyte), 1, stdin);
+    if (ret == 1) {
+      sum += inbyte;
+    }
+  }
+  return sum;
+}
+
+int stdin_fgetc() {
+  char inbyte;
+  int sum = 0;
+
+  while (true) {
+    int ret = fgetc(stdin);
+    if (ret == EOF)
+      break;
+    sum += inbyte;
+  }
+  return sum;
+}
+
+int stdin_getc() {
+  char inbyte;
+  int sum = 0;
+
+  while (true) {
+    int ret = getc(stdin);
+    if (ret == EOF)
+      break;
+    sum += inbyte;
+  }
+  return sum;
+}
+
+int stdin_getc_unlocked() {
+  char inbyte;
+  int sum = 0;
+
+  while (true) {
+    int ret = getc_unlocked(stdin);
+    if (ret == EOF)
+      break;
+    sum += inbyte;
+  }
+  return sum;
+}
+
+int stdin_getchar() {
+  char inbyte;
+  int sum = 0;
+
+  while (true) {
+    int ret = getchar();
+    if (ret == EOF)
+      break;
+    sum += inbyte;
+  }
+  return sum;
+}
+
+int stdin_getchar_unlocked() {
+  char inbyte;
+  int sum = 0;
+
+  while (true) {
+    int ret = getchar_unlocked();
+    if (ret == EOF)
+      break;
+    sum += inbyte;
+  }
+  return sum;
+}
+
+// Reads from stdin using different methods based on argv[1]
+// the following functions can be used
+// read,
+int main(int argc, char *argv[]) {
+  assert(argc == 2);
+  std::string_view method{argv[1]};
+
+  if (method == "read") {
+    return stdin_read();
+  } else if (method == "fread") {
+    return stdin_fread();
+  } else if (method == "getc") {
+    return stdin_getc();
+  } else if (method == "getc_unlocked") {
+    return stdin_getc_unlocked();
+  } else if (method == "getchar") {
+    return stdin_getchar();
+  } else if (method == "getchar_unlocked") {
+    return stdin_getchar_unlocked();
+  } else if (method == "fgetc") {
+    return stdin_fgetc();
+  }
+  return 0;
+}
\ No newline at end of file
diff --git a/tests/test_stdin.py b/tests/test_stdin.py
new file mode 100644
index 00000000..2f4bc634
--- /dev/null
+++ b/tests/test_stdin.py
@@ -0,0 +1,44 @@
+import pytest
+import subprocess
+
+import polytracker
+from pathlib import Path
+
+
+@pytest.mark.program_trace("test_stdin.cpp")
+@pytest.mark.parametrize(
+    "method",
+    ["read", "fread", "getc", "getc_unlocked", "getchar", "getchar_unlocked", "fgetc"],
+)
+def test_stdin_read(instrumented_binary: Path, trace_file: Path, method: str):
+    # Data to write to stdin, one byte at a time
+    stdin_data = "abcdefghi\njklmnopqr"
+
+    subprocess.run(
+        [str(instrumented_binary), method],
+        input=stdin_data.encode("utf-8"),
+        env={"POLYDB": str(trace_file), "POLYTRACKER_STDIN_SOURCE": str(1)},
+    )
+    program_trace = polytracker.PolyTrackerTrace.load(trace_file)
+
+    # Ensure /dev/stdin is in the list of inputs
+    assert "/dev/stdin" in [x.path for x in program_trace.inputs]
+
+    n = 0
+    for input_label in program_trace.tdfile.input_labels():
+        src_node = program_trace.tdfile.decode_node(input_label)
+        assert isinstance(src_node, polytracker.taint_dag.TDSourceNode)
+
+        # Requires that offsets are ordered according to read
+        assert src_node.offset == n
+
+        # Ensure all source labels originate from stdin
+        assert program_trace.tdfile.fd_headers[src_node.idx][0] == Path("/dev/stdin")
+        n += 1
+
+    # Should be as many source labels as the length of stdin_data
+    assert n == len(stdin_data)
+
+
+# Ensure stdin reads in multiple ways are verified
+# examples: getc, fgetc, fread, fread_unlocked, fgetc_unlocked, gets, fgets, getdelim, __getdelim, getw
diff --git a/unittests/src/taintdag/CMakeLists.txt b/unittests/src/taintdag/CMakeLists.txt
index 5378634a..c6f1c1d5 100644
--- a/unittests/src/taintdag/CMakeLists.txt
+++ b/unittests/src/taintdag/CMakeLists.txt
@@ -5,12 +5,13 @@ add_executable(
   main.cpp
   tdag.cpp
   taintdag.cpp
+  bitmap_section.cpp
   encoding.cpp
   fnmapping.cpp
   fntrace.cpp
   union.cpp
   labeldeq.cpp
-  bitmap_section.cpp)
+  stream_offset.cpp)
 
 target_include_directories(
   ${TAINTDAG_UNITTEST} PUBLIC ${CMAKE_SOURCE_DIR}/polytracker/include
diff --git a/unittests/src/taintdag/bitmap_section.cpp b/unittests/src/taintdag/bitmap_section.cpp
index 0bebdb48..9c047c74 100644
--- a/unittests/src/taintdag/bitmap_section.cpp
+++ b/unittests/src/taintdag/bitmap_section.cpp
@@ -195,4 +195,4 @@ TEMPLATE_TEST_CASE("Atomic Bitmap", "atomicbitmap", uint8_t, uint16_t, uint32_t,
     REQUIRE_THROWS_AS(bs.set_range(capacity - 2, 3), test::ErrorExit);
   }
 }
-} // namespace taintdag
\ No newline at end of file
+} // namespace taintdag
diff --git a/unittests/src/taintdag/stream_offset.cpp b/unittests/src/taintdag/stream_offset.cpp
new file mode 100644
index 00000000..b39bbc51
--- /dev/null
+++ b/unittests/src/taintdag/stream_offset.cpp
@@ -0,0 +1,34 @@
+#include <catch2/catch.hpp>
+
+#include "taintdag/stream_offset.h"
+
+#include "utils.h"
+
+namespace taintdag {
+
+TEST_CASE("StreamOffset", "StreamOffset") {
+
+  StreamOffset<4> ofs;
+  REQUIRE(ofs.read(0, 0) == 0);
+  REQUIRE(ofs.read(0, 0) == 0);
+
+  SECTION("Reading 3 bytes twice") {
+    REQUIRE(ofs.read(0, 3) == 0);
+    REQUIRE(ofs.read(0, 3) == 3);
+  }
+
+  SECTION("Reads doesn't interfer") {
+    ofs.read(0, 99);
+    ofs.read(1, 2);
+
+    REQUIRE(ofs.read(0, 1) == 99);
+    REQUIRE(ofs.read(1, 1) == 2);
+  }
+
+  SECTION("SourceIndex out of bounds aborts") {
+    test::ErrorExitReplace errthrow;
+    REQUIRE_THROWS_AS(ofs.read(4, 1), test::ErrorExit);
+  }
+}
+
+}
\ No newline at end of file
diff --git a/unittests/src/taintdag/union.cpp b/unittests/src/taintdag/union.cpp
index 87153605..a6f3fb40 100644
--- a/unittests/src/taintdag/union.cpp
+++ b/unittests/src/taintdag/union.cpp
@@ -121,4 +121,4 @@ TEST_CASE("Union Represents Same") {
   }
 
   INFO("#labels " << nlabels << " #values " << nvalues);
-}
\ No newline at end of file
+}