IPv6 support for tcp* tools (#582)

* tcpretrans: support full IPv6 addresses, fix --lossprobe

* tcpaccept: support full IPv6 addresses, fix timestamps

* tcpconnect: support full IPv6 addresses, fix timestamps

* tcpconnlat: support full IPv6 addresses, fix timestamps
diff --git a/man/man8/tcpaccept.8 b/man/man8/tcpaccept.8
index 7fd97b7..64c9d5a 100644
--- a/man/man8/tcpaccept.8
+++ b/man/man8/tcpaccept.8
@@ -55,12 +55,10 @@
 IP address family (4 or 6)
 .TP
 RADDR
-Remote IP address. IPv4 as a dotted quad, IPv6 shows "..." then the last 4
-bytes (check for newer versions of this tool for the full address).
+Remote IP address.
 .TP
 LADDR
-Local IP address. IPv4 as a dotted quad, IPv6 shows "..." then the last 4
-bytes (check for newer versions of this tool for the full address).
+Local IP address.
 .TP
 LPORT
 Local port
diff --git a/man/man8/tcpconnlat.8 b/man/man8/tcpconnlat.8
index f44af53..feb99af 100644
--- a/man/man8/tcpconnlat.8
+++ b/man/man8/tcpconnlat.8
@@ -58,12 +58,10 @@
 IP address family (4 or 6).
 .TP
 SADDR
-Source IP address. IPv4 as a dotted quad, IPv6 shows "..." then the last 4
-bytes (check for newer versions of this tool for the full address).
+Source IP address.
 .TP
 DADDR
-Destination IP address. IPv4 as a dotted quad, IPv6 shows "..." then the last 4
-bytes (check for newer versions of this tool for the full address).
+Destination IP address.
 .TP
 DPORT
 Destination port
diff --git a/man/man8/tcpretrans.8 b/man/man8/tcpretrans.8
index 9820378..8f04a15 100644
--- a/man/man8/tcpretrans.8
+++ b/man/man8/tcpretrans.8
@@ -47,8 +47,7 @@
 IP address family (4 or 6).
 .TP
 LADDR
-Local IP address. IPv4 as a dotted quad, IPv6 shows "..." then the last 4
-bytes (check for newer versions of this tool for the full address).
+Local IP address.
 .TP
 LPORT
 Local port.
@@ -57,8 +56,7 @@
 Type of event: R> == retransmit, L> == tail loss probe.
 .TP
 RADDR
-Remote IP address. IPv4 as a dotted quad, IPv6 shows "..." then the last 4
-bytes (check for newer versions of this tool for the full address).
+Remote IP address.
 .TP
 RPORT
 Remote port.
diff --git a/tools/tcpaccept.py b/tools/tcpaccept.py
index 802f94f..0c556a2 100755
--- a/tools/tcpaccept.py
+++ b/tools/tcpaccept.py
@@ -9,10 +9,6 @@
 # This uses dynamic tracing of the kernel inet_csk_accept() socket function
 # (from tcp_prot.accept), and will need to be modified to match kernel changes.
 #
-# IPv4 addresses are printed as dotted quads. For IPv6 addresses, the last four
-# bytes are printed after "..."; check for future versions with better IPv6
-# support.
-#
 # Copyright (c) 2015 Brendan Gregg.
 # Licensed under the Apache License, Version 2.0 (the "License")
 #
@@ -21,6 +17,8 @@
 
 from __future__ import print_function
 from bcc import BPF
+from socket import inet_ntop, AF_INET, AF_INET6
+from struct import pack
 import argparse
 import ctypes as ct
 
@@ -52,21 +50,20 @@
     // XXX: switch some to u32's when supported
     u64 ts_us;
     u64 pid;
-    u64 ip;
     u64 saddr;
     u64 daddr;
+    u64 ip;
     u64 lport;
     char task[TASK_COMM_LEN];
 };
 BPF_PERF_OUTPUT(ipv4_events);
 
 struct ipv6_data_t {
-    // XXX: update to transfer full ipv6 addrs
     u64 ts_us;
     u64 pid;
+    unsigned __int128 saddr;
+    unsigned __int128 daddr;
     u64 ip;
-    u64 saddr;
-    u64 daddr;
     u64 lport;
     char task[TASK_COMM_LEN];
 };
@@ -106,14 +103,10 @@
     } else if (family == AF_INET6) {
         struct ipv6_data_t data6 = {.pid = pid, .ip = 6};
         data6.ts_us = bpf_ktime_get_ns() / 1000;
-        // just grab the last 4 bytes for now
-        u32 saddr = 0, daddr = 0;
-        bpf_probe_read(&saddr, sizeof(saddr),
-            &newsk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32[3]);
-        bpf_probe_read(&daddr, sizeof(daddr),
-            &newsk->__sk_common.skc_v6_daddr.in6_u.u6_addr32[3]);
-        data6.saddr = bpf_ntohl(saddr);
-        data6.daddr = bpf_ntohl(daddr);
+        bpf_probe_read(&data6.saddr, sizeof(data6.saddr),
+            &newsk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32);
+        bpf_probe_read(&data6.daddr, sizeof(data6.daddr),
+            &newsk->__sk_common.skc_v6_daddr.in6_u.u6_addr32);
         data6.lport = lport;
         bpf_get_current_comm(&data6.task, sizeof(data6.task));
         ipv6_events.perf_submit(ctx, &data6, sizeof(data6));
@@ -135,23 +128,25 @@
 
 # event data
 TASK_COMM_LEN = 16      # linux/sched.h
+
 class Data_ipv4(ct.Structure):
     _fields_ = [
         ("ts_us", ct.c_ulonglong),
         ("pid", ct.c_ulonglong),
-        ("ip", ct.c_ulonglong),
         ("saddr", ct.c_ulonglong),
         ("daddr", ct.c_ulonglong),
+        ("ip", ct.c_ulonglong),
         ("lport", ct.c_ulonglong),
         ("task", ct.c_char * TASK_COMM_LEN)
     ]
+
 class Data_ipv6(ct.Structure):
     _fields_ = [
         ("ts_us", ct.c_ulonglong),
         ("pid", ct.c_ulonglong),
+        ("saddr", (ct.c_ulonglong * 2)),
+        ("daddr", (ct.c_ulonglong * 2)),
         ("ip", ct.c_ulonglong),
-        ("saddr", ct.c_ulonglong),
-        ("daddr", ct.c_ulonglong),
         ("lport", ct.c_ulonglong),
         ("task", ct.c_char * TASK_COMM_LEN)
     ]
@@ -159,21 +154,25 @@
 # process event
 def print_ipv4_event(cpu, data, size):
     event = ct.cast(data, ct.POINTER(Data_ipv4)).contents
+    global start_ts
     if args.timestamp:
         if start_ts == 0:
             start_ts = event.ts_us
-        print("%-9.3f" % ((event.ts_us - start_ts) / 100000), end="")
+        print("%-9.3f" % ((float(event.ts_us) - start_ts) / 1000000), end="")
     print("%-6d %-12.12s %-2d %-16s %-16s %-4d" % (event.pid, event.task,
-        event.ip, inet_ntoa(event.daddr), inet_ntoa(event.saddr),
-        event.lport))
+        event.ip, inet_ntop(AF_INET, pack("I", event.daddr)),
+        inet_ntop(AF_INET, pack("I", event.saddr)), event.lport))
+
 def print_ipv6_event(cpu, data, size):
     event = ct.cast(data, ct.POINTER(Data_ipv6)).contents
+    global start_ts
     if args.timestamp:
         if start_ts == 0:
             start_ts = event.ts_us
-        print("%-9.3f" % ((event.ts_us - start_ts) / 100000), end="")
-    print("%-6d %-12.12s %-2d ...%-13x ...%-13x %-4d" % (event.pid,
-        event.task, event.ip, event.daddr, event.saddr, event.lport))
+        print("%-9.3f" % ((float(event.ts_us) - start_ts) / 1000000), end="")
+    print("%-6d %-12.12s %-2d %-16s %-16s %-4d" % (event.pid, event.task,
+        event.ip, inet_ntop(AF_INET6, event.daddr),
+        inet_ntop(AF_INET6, event.saddr), event.lport))
 
 # initialize BPF
 b = BPF(text=bpf_text)
@@ -186,15 +185,6 @@
 
 start_ts = 0
 
-def inet_ntoa(addr):
-    dq = ''
-    for i in range(0, 4):
-        dq = dq + str(addr & 0xff)
-        if (i != 3):
-            dq = dq + '.'
-        addr = addr >> 8
-    return dq
-
 # read events
 b["ipv4_events"].open_perf_buffer(print_ipv4_event)
 b["ipv6_events"].open_perf_buffer(print_ipv6_event)
diff --git a/tools/tcpaccept_example.txt b/tools/tcpaccept_example.txt
index 04cdc67..f86c439 100644
--- a/tools/tcpaccept_example.txt
+++ b/tools/tcpaccept_example.txt
@@ -9,14 +9,11 @@
 PID    COMM         IP RADDR            LADDR            LPORT
 907    sshd         4  192.168.56.1     192.168.56.102   22
 907    sshd         4  127.0.0.1        127.0.0.1        22
-5389   perl         6  ...fec0ae21      ...fec0ae21      7001
+5389   perl         6  1234:ab12:2040:5020:2299:0:5:0 1234:ab12:2040:5020:2299:0:5:0 7001
 
-This output shows three connections, two to PID 907, an "sshd" process listening
-on port 22, and one to a "perl" process listening on port 7001.
-
-The sshd connections were IPv4, and the addresses are printed as dotted quads.
-The perl connection was IPv6, and the last 4 bytes of each address is printed
-(for now; check for updated versions).
+This output shows three connections, two IPv4 connections to PID 907, an "sshd"
+process listening on port 22, and one IPv6 connection to a "perl" process
+listening on port 7001.
 
 The overhead of this tool should be negligible, since it is only tracing the
 kernel function performing accept. It is not tracing every packet and then
@@ -31,6 +28,7 @@
 # ./tcpaccept -t
 TIME(s)  PID    COMM         IP RADDR            LADDR            LPORT
 0.000    907    sshd         4  127.0.0.1        127.0.0.1        22
+0.010    5389   perl         6  1234:ab12:2040:5020:2299:0:5:0 1234:ab12:2040:5020:2299:0:5:0 7001
 0.992    907    sshd         4  127.0.0.1        127.0.0.1        22
 1.984    907    sshd         4  127.0.0.1        127.0.0.1        22
 
diff --git a/tools/tcpconnect.py b/tools/tcpconnect.py
index 02c1b3c..7615167 100755
--- a/tools/tcpconnect.py
+++ b/tools/tcpconnect.py
@@ -11,10 +11,6 @@
 # This uses dynamic tracing of kernel functions, and will need to be updated
 # to match kernel changes.
 #
-# IPv4 addresses are printed as dotted quads. For IPv6 addresses, the last four
-# bytes are printed after "..."; check for future versions with better IPv6
-# support.
-#
 # Copyright (c) 2015 Brendan Gregg.
 # Licensed under the Apache License, Version 2.0 (the "License")
 #
@@ -24,8 +20,8 @@
 from __future__ import print_function
 from bcc import BPF
 import argparse
-import re
-from struct import pack, unpack_from
+from socket import inet_ntop, AF_INET, AF_INET6
+from struct import pack
 import ctypes as ct
 
 # arguments
@@ -58,9 +54,9 @@
     // XXX: switch some to u32's when supported
     u64 ts_us;
     u64 pid;
-    u64 ip;
     u64 saddr;
     u64 daddr;
+    u64 ip;
     u64 dport;
     char task[TASK_COMM_LEN];
 };
@@ -69,9 +65,9 @@
 struct ipv6_data_t {
     u64 ts_us;
     u64 pid;
+    unsigned __int128 saddr;
+    unsigned __int128 daddr;
     u64 ip;
-    u64 saddr[2];
-    u64 daddr[2];
     u64 dport;
     char task[TASK_COMM_LEN];
 };
@@ -125,15 +121,10 @@
     } else /* 6 */ {
         struct ipv6_data_t data6 = {.pid = pid, .ip = ipver};
         data6.ts_us = bpf_ktime_get_ns() / 1000;
-        // just grab the last 4 bytes for now
-        bpf_probe_read(&data6.saddr[0], sizeof(data6.saddr[0]),
-            &skp->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32[0]);
-        bpf_probe_read(&data6.saddr[1], sizeof(data6.saddr[1]),
-            &skp->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32[2]);
-        bpf_probe_read(&data6.daddr[0], sizeof(data6.daddr[0]),
-            &skp->__sk_common.skc_v6_daddr.in6_u.u6_addr32[0]);
-        bpf_probe_read(&data6.daddr[1], sizeof(data6.daddr[1]),
-            &skp->__sk_common.skc_v6_daddr.in6_u.u6_addr32[2]);
+        bpf_probe_read(&data6.saddr, sizeof(data6.saddr),
+            &skp->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32);
+        bpf_probe_read(&data6.daddr, sizeof(data6.daddr),
+            &skp->__sk_common.skc_v6_daddr.in6_u.u6_addr32);
         data6.dport = ntohs(dport);
         bpf_get_current_comm(&data6.task, sizeof(data6.task));
         ipv6_events.perf_submit(ctx, &data6, sizeof(data6));
@@ -166,23 +157,25 @@
 
 # event data
 TASK_COMM_LEN = 16      # linux/sched.h
+
 class Data_ipv4(ct.Structure):
     _fields_ = [
         ("ts_us", ct.c_ulonglong),
         ("pid", ct.c_ulonglong),
-        ("ip", ct.c_ulonglong),
         ("saddr", ct.c_ulonglong),
         ("daddr", ct.c_ulonglong),
+        ("ip", ct.c_ulonglong),
         ("dport", ct.c_ulonglong),
         ("task", ct.c_char * TASK_COMM_LEN)
     ]
+
 class Data_ipv6(ct.Structure):
     _fields_ = [
         ("ts_us", ct.c_ulonglong),
         ("pid", ct.c_ulonglong),
+        ("saddr", (ct.c_ulonglong * 2)),
+        ("daddr", (ct.c_ulonglong * 2)),
         ("ip", ct.c_ulonglong),
-        ("saddr", ct.c_ulonglong * 2),
-        ("daddr", ct.c_ulonglong * 2),
         ("dport", ct.c_ulonglong),
         ("task", ct.c_char * TASK_COMM_LEN)
     ]
@@ -190,25 +183,25 @@
 # process event
 def print_ipv4_event(cpu, data, size):
     event = ct.cast(data, ct.POINTER(Data_ipv4)).contents
+    global start_ts
     if args.timestamp:
         if start_ts == 0:
             start_ts = event.ts_us
-        print("%-9.3f" % ((event.ts_us - start_ts) / 100000), end="")
+        print("%-9.3f" % ((float(event.ts_us) - start_ts) / 1000000), end="")
     print("%-6d %-12.12s %-2d %-16s %-16s %-4d" % (event.pid, event.task,
-        event.ip, inet_ntoa(event.saddr), inet_ntoa(event.daddr),
-        event.dport))
+        event.ip, inet_ntop(AF_INET, pack("I", event.saddr)),
+        inet_ntop(AF_INET, pack("I", event.daddr)), event.dport))
 
 def print_ipv6_event(cpu, data, size):
     event = ct.cast(data, ct.POINTER(Data_ipv6)).contents
+    global start_ts
     if args.timestamp:
         if start_ts == 0:
             start_ts = event.ts_us
-        print("%-9.3f" % ((event.ts_us - start_ts) / 100000), end="")
+        print("%-9.3f" % ((float(event.ts_us) - start_ts) / 1000000), end="")
     print("%-6d %-12.12s %-2d %-16s %-16s %-4d" % (event.pid,
-        event.task, event.ip,
-        inet6_ntoa(event.saddr[1] << 64 | event.saddr[0]),
-        inet6_ntoa(event.daddr[1] << 64 | event.daddr[0]),
-        event.dport))
+        event.task, event.ip, inet_ntop(AF_INET6, event.saddr),
+        inet_ntop(AF_INET6, event.daddr), event.dport))
 
 # initialize BPF
 b = BPF(text=bpf_text)
@@ -225,25 +218,6 @@
 
 start_ts = 0
 
-def inet_ntoa(addr):
-    # u32 to dotted quad string
-    dq = ''
-    for i in range(0, 4):
-        dq = dq + str(addr & 0xff)
-        if (i != 3):
-            dq = dq + '.'
-        addr = addr >> 8
-    return dq
-
-def inet6_ntoa(addr):
-    # see RFC4291 summary in RFC5952 section 2
-    s = ":".join(["%x" % x for x in unpack_from(">HHHHHHHH",
-        pack("QQ", addr & 0xffffffff, addr >> 64))])
-
-    # compress left-most zero run only (change to most for RFC5952):
-    s = re.sub(r'(^|:)0:(0:)+', r'::', s, 1)
-    return s
-
 # read events
 b["ipv4_events"].open_perf_buffer(print_ipv4_event)
 b["ipv6_events"].open_perf_buffer(print_ipv6_event)
diff --git a/tools/tcpconnlat.py b/tools/tcpconnlat.py
index 26add2a..8a8377f 100755
--- a/tools/tcpconnlat.py
+++ b/tools/tcpconnlat.py
@@ -9,10 +9,6 @@
 # This uses dynamic tracing of kernel functions, and will need to be updated
 # to match kernel changes.
 #
-# IPv4 addresses are printed as dotted quads. For IPv6 addresses, the last four
-# bytes are printed after "..."; check for future versions with better IPv6
-# support.
-#
 # Copyright 2016 Netflix, Inc.
 # Licensed under the Apache License, Version 2.0 (the "License")
 #
@@ -20,6 +16,8 @@
 
 from __future__ import print_function
 from bcc import BPF
+from socket import inet_ntop, AF_INET, AF_INET6
+from struct import pack
 import argparse
 import ctypes as ct
 
@@ -44,6 +42,7 @@
 bpf_text = """
 #include <uapi/linux/ptrace.h>
 #include <net/sock.h>
+#include <net/tcp_states.h>
 #include <bcc/proto.h>
 
 struct info_t {
@@ -58,9 +57,9 @@
     // XXX: switch some to u32's when supported
     u64 ts_us;
     u64 pid;
-    u64 ip;
     u64 saddr;
     u64 daddr;
+    u64 ip;
     u64 dport;
     u64 delta_us;
     char task[TASK_COMM_LEN];
@@ -68,12 +67,11 @@
 BPF_PERF_OUTPUT(ipv4_events);
 
 struct ipv6_data_t {
-    // XXX: update to transfer full ipv6 addrs
     u64 ts_us;
     u64 pid;
+    unsigned __int128 saddr;
+    unsigned __int128 daddr;
     u64 ip;
-    u64 saddr;
-    u64 daddr;
     u64 dport;
     u64 delta_us;
     char task[TASK_COMM_LEN];
@@ -132,14 +130,10 @@
     } else /* AF_INET6 */ {
         struct ipv6_data_t data6 = {.pid = infop->pid, .ip = 6};
         data6.ts_us = now / 1000;
-        // just grab the last 4 bytes for now
-        u32 saddr = 0, daddr = 0;
-        bpf_probe_read(&saddr, sizeof(saddr),
-            &skp->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32[3]);
-        bpf_probe_read(&daddr, sizeof(daddr),
-            &skp->__sk_common.skc_v6_daddr.in6_u.u6_addr32[3]);
-        data6.saddr = bpf_ntohl(saddr);
-        data6.daddr = bpf_ntohl(daddr);
+        bpf_probe_read(&data6.saddr, sizeof(data6.saddr),
+            &skp->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32);
+        bpf_probe_read(&data6.daddr, sizeof(data6.daddr),
+            &skp->__sk_common.skc_v6_daddr.in6_u.u6_addr32);
         data6.dport = ntohs(dport);
         data6.delta_us = (now - ts) / 1000;
         __builtin_memcpy(&data6.task, infop->task, sizeof(data6.task));
@@ -170,60 +164,56 @@
 
 # event data
 TASK_COMM_LEN = 16      # linux/sched.h
+
 class Data_ipv4(ct.Structure):
     _fields_ = [
         ("ts_us", ct.c_ulonglong),
         ("pid", ct.c_ulonglong),
-        ("ip", ct.c_ulonglong),
         ("saddr", ct.c_ulonglong),
         ("daddr", ct.c_ulonglong),
+        ("ip", ct.c_ulonglong),
         ("dport", ct.c_ulonglong),
         ("delta_us", ct.c_ulonglong),
         ("task", ct.c_char * TASK_COMM_LEN)
     ]
+
 class Data_ipv6(ct.Structure):
     _fields_ = [
         ("ts_us", ct.c_ulonglong),
         ("pid", ct.c_ulonglong),
+        ("saddr", (ct.c_ulonglong * 2)),
+        ("daddr", (ct.c_ulonglong * 2)),
         ("ip", ct.c_ulonglong),
-        ("saddr", ct.c_ulonglong),
-        ("daddr", ct.c_ulonglong),
         ("dport", ct.c_ulonglong),
         ("delta_us", ct.c_ulonglong),
         ("task", ct.c_char * TASK_COMM_LEN)
     ]
 
-# functions
-def inet_ntoa(addr):
-    dq = ''
-    for i in range(0, 4):
-        dq = dq + str(addr & 0xff)
-        if (i != 3):
-            dq = dq + '.'
-        addr = addr >> 8
-    return dq
-
 # process event
 start_ts = 0
+
 def print_ipv4_event(cpu, data, size):
     event = ct.cast(data, ct.POINTER(Data_ipv4)).contents
     global start_ts
     if args.timestamp:
         if start_ts == 0:
             start_ts = event.ts_us
-        print("%-9.3f" % ((event.ts_us - start_ts) / 100000), end="")
+        print("%-9.3f" % ((float(event.ts_us) - start_ts) / 1000000), end="")
     print("%-6d %-12.12s %-2d %-16s %-16s %-5d %.2f" % (event.pid, event.task,
-        event.ip, inet_ntoa(event.saddr), inet_ntoa(event.daddr),
-        event.dport, float(event.delta_us) / 1000))
+        event.ip, inet_ntop(AF_INET, pack("I", event.saddr)),
+        inet_ntop(AF_INET, pack("I", event.daddr)), event.dport,
+        float(event.delta_us) / 1000))
+
 def print_ipv6_event(cpu, data, size):
     event = ct.cast(data, ct.POINTER(Data_ipv6)).contents
     global start_ts
     if args.timestamp:
         if start_ts == 0:
             start_ts = event.ts_us
-        print("%-9.3f" % ((event.ts_us - start_ts) / 100000), end="")
-    print("%-6d %-12.12s %-2d ...%-13x ...%-13x %-5d %.2f" % (event.pid,
-        event.task, event.ip, event.saddr, event.daddr, event.dport,
+        print("%-9.3f" % ((float(event.ts_us) - start_ts) / 1000000), end="")
+    print("%-6d %-12.12s %-2d %-16s %-16s %-5d %.2f" % (event.pid, event.task,
+        event.ip, inet_ntop(AF_INET6, event.saddr),
+        inet_ntop(AF_INET6, event.daddr), event.dport,
         float(event.delta_us) / 1000))
 
 # header
diff --git a/tools/tcpconnlat_example.txt b/tools/tcpconnlat_example.txt
index c5b289e..f179bfc 100644
--- a/tools/tcpconnlat_example.txt
+++ b/tools/tcpconnlat_example.txt
@@ -15,6 +15,8 @@
 1690   wget         4  10.153.223.157   66.220.156.68    443   0.95
 1690   wget         4  10.153.223.157   66.220.156.68    443   0.99
 2852   curl         4  10.153.223.157   23.101.17.61     80    250.86
+20337  python2.7    6  1234:ab12:2040:5020:2299:0:5:0 1234:ab12:20:9f1d:2299:dde9:0:f5 7001  62.20
+21588  nc           6  ::1              ::1              80    0.05
 [...]
 
 The first line shows a connection from the "wget" process to the IPv4
diff --git a/tools/tcpretrans.py b/tools/tcpretrans.py
index 1b67ca5..01de704 100755
--- a/tools/tcpretrans.py
+++ b/tools/tcpretrans.py
@@ -9,10 +9,6 @@
 # This uses dynamic tracing of kernel functions, and will need to be updated
 # to match kernel changes.
 #
-# IPv4 addresses are printed as dotted quads. For IPv6 addresses, the last four
-# bytes are printed after "..."; check for future versions with better IPv6
-# support.
-#
 # Copyright 2016 Netflix, Inc.
 # Licensed under the Apache License, Version 2.0 (the "License")
 #
@@ -22,6 +18,8 @@
 from bcc import BPF
 import argparse
 from time import strftime
+from socket import inet_ntop, AF_INET, AF_INET6
+from struct import pack
 import ctypes as ct
 
 # arguments
@@ -61,11 +59,10 @@
 BPF_PERF_OUTPUT(ipv4_events);
 
 struct ipv6_data_t {
-    // XXX: update to transfer full ipv6 addrs
     u64 pid;
     u64 ip;
-    u64 saddr;
-    u64 daddr;
+    unsigned __int128 saddr;
+    unsigned __int128 daddr;
     u64 lport;
     u64 dport;
     u64 state;
@@ -102,14 +99,10 @@
 
     } else if (family == AF_INET6) {
         struct ipv6_data_t data6 = {.pid = pid, .ip = 6, .type = type};
-        // just grab the last 4 bytes for now
-        u32 saddr = 0, daddr = 0;
-        bpf_probe_read(&saddr, sizeof(saddr),
-            &skp->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32[3]);
-        bpf_probe_read(&daddr, sizeof(daddr),
-            &skp->__sk_common.skc_v6_daddr.in6_u.u6_addr32[3]);
-        data6.saddr = bpf_ntohl(saddr);
-        data6.daddr = bpf_ntohl(daddr);
+        bpf_probe_read(&data6.saddr, sizeof(data6.saddr),
+            &skp->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32);
+        bpf_probe_read(&data6.daddr, sizeof(data6.daddr),
+            &skp->__sk_common.skc_v6_daddr.in6_u.u6_addr32);
         data6.lport = lport;
         data6.dport = dport;
         data6.state = state;
@@ -134,7 +127,6 @@
 """
 
 # event data
-TASK_COMM_LEN = 16      # linux/sched.h
 class Data_ipv4(ct.Structure):
     _fields_ = [
         ("pid", ct.c_ulonglong),
@@ -146,12 +138,13 @@
         ("state", ct.c_ulonglong),
         ("type", ct.c_ulonglong)
     ]
+
 class Data_ipv6(ct.Structure):
     _fields_ = [
         ("pid", ct.c_ulonglong),
         ("ip", ct.c_ulonglong),
-        ("saddr", ct.c_ulonglong),
-        ("daddr", ct.c_ulonglong),
+        ("saddr", (ct.c_ulonglong * 2)),
+        ("daddr", (ct.c_ulonglong * 2)),
         ("lport", ct.c_ulonglong),
         ("dport", ct.c_ulonglong),
         ("state", ct.c_ulonglong),
@@ -163,15 +156,6 @@
 type[1] = 'R'
 type[2] = 'L'
 
-def inet_ntoa(addr):
-    dq = ''
-    for i in range(0, 4):
-        dq = dq + str(addr & 0xff)
-        if (i != 3):
-            dq = dq + '.'
-        addr = addr >> 8
-    return dq
-
 # from include/net/tcp_states.h:
 tcpstate = {}
 tcpstate[1] = 'ESTABLISHED'
@@ -192,23 +176,25 @@
     event = ct.cast(data, ct.POINTER(Data_ipv4)).contents
     print("%-8s %-6d %-2d %-20s %1s> %-20s %s" % (
         strftime("%H:%M:%S"), event.pid, event.ip,
-        "%s:%s" % (inet_ntoa(event.saddr), event.lport),
+        "%s:%d" % (inet_ntop(AF_INET, pack('I', event.saddr)), event.lport),
         type[event.type],
-        "%s:%s" % (inet_ntoa(event.daddr), event.dport),
+        "%s:%s" % (inet_ntop(AF_INET, pack('I', event.daddr)), event.dport),
         tcpstate[event.state]))
+
 def print_ipv6_event(cpu, data, size):
     event = ct.cast(data, ct.POINTER(Data_ipv6)).contents
     print("%-8s %-6d %-2d %-20s %1s> %-20s %s" % (
         strftime("%H:%M:%S"), event.pid, event.ip,
-        "...%x:%d" % (event.saddr, event.lport),
+        "%s:%d" % (inet_ntop(AF_INET6, event.saddr), event.lport),
         type[event.type],
-        "...%x:%d" % (event.daddr, event.dport),
+        "%s:%d" % (inet_ntop(AF_INET6, event.daddr), event.dport),
         tcpstate[event.state]))
 
 # initialize BPF
 b = BPF(text=bpf_text)
 b.attach_kprobe(event="tcp_retransmit_skb", fn_name="trace_retransmit")
-b.attach_kprobe(event="tcp_send_loss_probe", fn_name="trace_tlp")
+if args.lossprobe:
+    b.attach_kprobe(event="tcp_send_loss_probe", fn_name="trace_tlp")
 
 # header
 print("%-8s %-6s %-2s %-20s %1s> %-20s %-4s" % ("TIME", "PID", "IP",