From 851c72b215419970ca0a68bc0a1e01f0956c7b90 Mon Sep 17 00:00:00 2001
From: Chris Kennedy <chris.kennedy@ltnglobal.com>
Date: Fri, 5 Apr 2024 01:43:54 -0400
Subject: [PATCH] fix timestamp extraction

copy to each packet till it changes, so each packet knows the
timestamp it is from.
---
 scripts/probe.sh   |  4 ++--
 src/stream_data.rs | 39 ++++++++++++++++++++++++++++++++++-----
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/scripts/probe.sh b/scripts/probe.sh
index 97f4eec8..5b26c42f 100644
--- a/scripts/probe.sh
+++ b/scripts/probe.sh
@@ -1,11 +1,11 @@
 #!/bin/bash
 #
-RUST_LOG=info \
+sudo RUST_LOG=info \
     target/release/probe \
     --pcap-stats \
     --source-ip 224.0.0.200 \
     --source-port 10000 \
-    --source-device en0  \
+    --source-device eth0  \
     --target-ip 0.0.0.0 \
     --send-null-packets \
     --target-port 5556 \
diff --git a/src/stream_data.rs b/src/stream_data.rs
index b1b63bd6..c6863839 100644
--- a/src/stream_data.rs
+++ b/src/stream_data.rs
@@ -633,6 +633,11 @@ pub fn process_packet(
                 Arc::make_mut(&mut stream_data)
                     .set_continuity_counter(stream_data_packet.continuity_counter);
             }
+
+            if stream_data_packet.timestamp != 0 {
+                Arc::make_mut(&mut stream_data).timestamp = stream_data_packet.timestamp;
+            }
+
             // calculate uptime using the arrival time as SystemTime and start_time as u64
             let uptime = stream_data.capture_time - stream_data.start_time;
 
@@ -659,6 +664,9 @@ pub fn process_packet(
             stream_data_packet.program_number = stream_data.program_number;
             stream_data_packet.error_count = stream_data.error_count;
             stream_data_packet.current_error_count = stream_data.current_error_count;
+            if stream_data_packet.timestamp == 0 {
+                stream_data_packet.timestamp = stream_data.timestamp;
+            }
 
             // write the stream_data back to the pid_map with modified values
             pid_map.insert(pid, stream_data);
@@ -1056,11 +1064,7 @@ pub fn process_mpegts_packet(
             let stream_type = determine_stream_type(pid);
             let stream_type_number = determine_stream_type_number(pid);
             let stream_program_number = determine_stream_program_number(pid);
-            let timestamp = ((chunk[4] as u64) << 25)
-                | ((chunk[5] as u64) << 17)
-                | ((chunk[6] as u64) << 9)
-                | ((chunk[7] as u64) << 1)
-                | ((chunk[8] as u64) >> 7);
+            let timestamp = extract_timestamp(chunk);
             let continuity_counter = chunk[3] & 0x0F;
 
             let mut stream_data = StreamData::new(
@@ -1090,3 +1094,28 @@ pub fn process_mpegts_packet(
 
     streams
 }
+
+fn extract_timestamp(chunk: &[u8]) -> u64 {
+    let adaptation_field_control = (chunk[3] & 0x30) >> 4;
+
+    if adaptation_field_control == 0b10 || adaptation_field_control == 0b11 {
+        // Adaptation field is present
+        let adaptation_field_length = chunk[4] as usize;
+
+        if adaptation_field_length > 0 && (chunk[5] & 0x10) != 0 {
+            // PCR is present
+            let pcr_base = ((chunk[6] as u64) << 25)
+                | ((chunk[7] as u64) << 17)
+                | ((chunk[8] as u64) << 9)
+                | ((chunk[9] as u64) << 1)
+                | ((chunk[10] as u64) >> 7);
+            let pcr_ext = (((chunk[10] as u64) & 0x01) << 8) | (chunk[11] as u64);
+            let pcr = pcr_base * 300 + pcr_ext;
+            pcr
+        } else {
+            0 // Default value when PCR is not present
+        }
+    } else {
+        0 // Default value when adaptation field is not present
+    }
+}