作者 | kernelxing

TCP的經(jīng)典異常問(wèn)題無(wú)非就是丟包和連接中斷，在這里我打算與各位聊一聊TCP的RST到底是什么？現(xiàn)網(wǎng)中的RST問(wèn)題有哪些模樣？我們?nèi)绾稳?yīng)對(duì)、解決？本文將從RST原理、排查手段、現(xiàn)網(wǎng)痛難點(diǎn)案例三個(gè)板塊自上而下帶給讀者一套完整的分析。

一、背景

最近一年的時(shí)間里，現(xiàn)網(wǎng)碰到RST問(wèn)題屢屢出現(xiàn)，一旦TCP連接中收到了RST包，大概率會(huì)導(dǎo)致連接中止或用戶異常。如何正確解決RST異常是較為棘手的問(wèn)題。

本文關(guān)注的不是細(xì)節(jié)，而是方法論，也確實(shí)方法更為重要。筆者始終相信，一百個(gè)人眼中的哈姆雷特最終還是一個(gè)具體的人物形象，一百個(gè)RST異常最終也會(huì)是一個(gè)簡(jiǎn)短的小問(wèn)題。

二、原理

首先，我們需要確定的RST問(wèn)題一定就是問(wèn)題嗎？如果RST發(fā)生了你會(huì)如何去解決？讀者可以嘗試問(wèn)下自己并解答這個(gè)問(wèn)題，這里”停頓、停頓、停頓“來(lái)給大家一點(diǎn)時(shí)間思考，好了，時(shí)間到，我們繼續(xù)往下看。

RST分為兩種，一種是active rst，另一種是passive rst。前者多半是指的符合預(yù)期的reset行為，此種情況多半是屬于機(jī)器自己主動(dòng)觸發(fā)，更具有先前意識(shí)，且和協(xié)議棧本身的細(xì)節(jié)關(guān)聯(lián)性不強(qiáng)；后者多半是指的機(jī)器也不清楚后面會(huì)發(fā)生什么，走一步看一步，如果不符合協(xié)議棧的if-else實(shí)現(xiàn)的RFC中條條杠杠的規(guī)則的情況下，那就只能reset重置了。

這里貼上RFC 793最經(jīng)典的最初對(duì)RST包的解釋：

1.active rst

那具體什么是active rst？如果從tcpdump抓包上來(lái)看表現(xiàn)就是（如下圖）RST的報(bào)文中含有了一串Ack標(biāo)識(shí)。

這個(gè)對(duì)應(yīng)的內(nèi)核代碼為（如果感興趣）：

tcp_send_active_reset()
    -> skb = alloc_skb(MAX_TCP_HEADER, priority);
    -> tcp_init_nondata_skb(skb, tcp_acceptable_seq(sk), TCPHDR_ACK | TCPHDR_RST);
    -> tcp_transmit_skb()

通常發(fā)生active rst的有幾種情況：

• 主動(dòng)方調(diào)用close()的時(shí)候，上層卻沒(méi)有取走完數(shù)據(jù)；這個(gè)屬于上層user自己犯下的錯(cuò)。
• 主動(dòng)方調(diào)用close()的時(shí)候，setsockopt設(shè)置了linger；這個(gè)標(biāo)識(shí)代表我既然設(shè)置了這個(gè)，那close就趕快結(jié)束吧。
• 主動(dòng)方調(diào)用close()的時(shí)候，發(fā)現(xiàn)全局的tcp可用的內(nèi)存不夠了（這個(gè)可以sysctl調(diào)整tcp mem第三個(gè)參數(shù)），或，發(fā)現(xiàn)已經(jīng)有太多的orphans了，這時(shí)候系統(tǒng)就是擺爛的意思：我也沒(méi)轍了”，那就只能干脆點(diǎn)長(zhǎng)痛不如短痛，結(jié)束吧。這個(gè)案例可以搜索(dmesg日志）“too many orphaned sockets”或“out of memory -- consider tuning tcp_mem”，匹配其中一個(gè)就容易中rst。

注：這里省略其他使用diag相關(guān)（如ss命令）的RST問(wèn)題。上述三類是主要的active rst問(wèn)題的情況。

2.passive rst

現(xiàn)在繼續(xù)說(shuō)說(shuō)另一種passive rst吧。如果從抓包上來(lái)看表現(xiàn)就是（如下圖）rst的報(bào)文中無(wú)ack標(biāo)識(shí)，而且RST的seq等于它否定的報(bào)文的ack號(hào)（紅色框的rst否定的黃色框的ack），當(dāng)然還有另一種極小概率出現(xiàn)的特殊情況的表現(xiàn)我這里不貼出來(lái)了，它的表現(xiàn)形式就是RST的Ack號(hào)為1。

這個(gè)對(duì)應(yīng)的內(nèi)核代碼為（如果感興趣）：

tcp_v4_send_reset()
        if (th->ack) {
                // 這里對(duì)應(yīng)的就是上圖中為何出現(xiàn)Seq==Ack
                rep.th.seq = th->ack_seq; 
        } else {
                // 極小概率，如果出現(xiàn)，那么RST包的就沒(méi)有Seq序列號(hào)
                rep.th.ack = 1;                                                   
                rep.th.ack_seq = htonl(ntohl(th->seq) + th->syn + th->fin +    
                                       skb->len - (th->doff << 2));               
        }

通常發(fā)生passive rst的有哪些情況呢？這個(gè)遠(yuǎn)比active rst更復(fù)雜，場(chǎng)景更多。具體的需要看TCP的收、發(fā)的協(xié)議，文字的描述可以參考rfc 793即可。

三、工具

我們針對(duì)線上這么多的rst如何去分析呢？首先tcpdump的抓捕是一定需要的，這個(gè)可以在整體流程上給我們縮小排查范圍，其次是，必須要手寫抓捕異常調(diào)用rst的點(diǎn)，文末我會(huì)分享一些源碼出來(lái)供參考。

那如何抓調(diào)用RST的點(diǎn)？這里只提供下思路。

1.active rst

使用bpf*相關(guān)的工具抓捕tcp_send_active_reset()函數(shù)并打印堆棧即可，通過(guò)crash現(xiàn)場(chǎng)機(jī)器并輸入“dis -l [addr]”可以得到具體的函數(shù)位置，比對(duì)源碼就可以得知了。

可以使用bpftrace進(jìn)行快速抓捕：

sudo bpftrace -e 'k:tcp_send_active_reset { @[kstack()] = count(); }'

堆棧結(jié)果如圖：

我們可以根據(jù)堆棧信息推算上下文。

2.passive rst

使用bpf*相關(guān)的工具抓捕抓捕tcp_v4_send_reset()和其他若干小的地方即可，原理同上。

sudo bpftrace -e 'k:tcp_v4_send_reset { @[kstack()] = count(); }'

效果如圖：

當(dāng)然，無(wú)論那種，我們抓到了堆棧后依然需要輸出很多的關(guān)于skb和sk的信息，這個(gè)讀者自行考慮即可。再補(bǔ)充一些抓捕小技巧，如果現(xiàn)網(wǎng)機(jī)器的rst數(shù)量較多時(shí)候，盡量使用匹配固定的ip+port方式或其它關(guān)鍵字來(lái)減少打印輸出，否則會(huì)消耗資源過(guò)多！

注：切記不能去抓捕reset tracepoint（具體函數(shù)：trace_tcp_send_reset()），這個(gè)tracepoint實(shí)現(xiàn)是有問(wèn)題的，這個(gè)問(wèn)題已經(jīng)在社區(qū)內(nèi)核中存在了7年之久！目前我正在修復(fù)中。

四、案例分析

本章節(jié)我將用現(xiàn)網(wǎng)實(shí)際碰到的三個(gè)”離譜“而且讓我非常”咬牙切齒“的case作為案例分析，當(dāng)時(shí)在查這些問(wèn)題的時(shí)候我提前告知業(yè)務(wù)”不保證有能力解決 :(“，不過(guò)最終還是用時(shí)間磨贏了bug。那么，這里讓各位讀者可以看下極為復(fù)雜的RST案例到底長(zhǎng)成什么樣？對(duì)內(nèi)核不感興趣的同學(xué)可以不用糾結(jié)具體的細(xì)節(jié)，只需要知道一個(gè)過(guò)程即可；對(duì)內(nèi)核感興趣的同學(xué)不妨可以一起構(gòu)造RST然后自己再抓取的試試。

第一個(gè)案例：小試牛刀—— close階段RST

背景：這是線上出現(xiàn)概率/次數(shù)較多的一種類型的RST，業(yè)務(wù)總是抱怨為何我的連接莫名其妙的又沒(méi)了。

我們先使用網(wǎng)絡(luò)異常檢測(cè)中最常用的工具：tcpdump。如下抓包的圖片再結(jié)合前文對(duì)RST的兩種分類（active && passive）可知，這是active rst。

好，既然知道了是active rst，我們就針對(duì)性的在線上對(duì)關(guān)鍵函數(shù)抓捕，如下：

通過(guò)crash命令找到了對(duì)應(yīng)的源碼，如下：

這時(shí)候便知是用戶設(shè)置了linger，主動(dòng)預(yù)期內(nèi)的行為觸發(fā)的rst，所以本例就解決了。不過(guò)插曲是，用戶并不認(rèn)為他設(shè)置了linger，這個(gè)怎么辦？那就再抓一次sk->sk_lingertime值就好咯，如下：

計(jì)算：socket的flag是784，第5位（從右往左）是1，這個(gè)是SO_LINGER位置位成功，但是同時(shí)linger_time為0。這個(gè)條件默認(rèn)（符合預(yù)期）觸發(fā)：上層用戶退出時(shí)候，不走四次揮手，直接RST結(jié)束。

結(jié)論：linger的默認(rèn)機(jī)制觸發(fā)了加速結(jié)束TCP連接從而RST報(bào)文發(fā)出。

第二個(gè)案例：TCP 兩個(gè)bug —— 握手與揮手的RS

背景：某重點(diǎn)業(yè)務(wù)報(bào)告他們的某重點(diǎn)用戶出現(xiàn)了莫名其妙的RST問(wèn)題，而且每一次都是出現(xiàn)在三次握手階段，復(fù)現(xiàn)概率約為——”按請(qǐng)求數(shù)來(lái)算的話差不多百萬(wàn)級(jí)別分之1的概率，概率極低“（這是來(lái)自業(yè)務(wù)的原話）。

這里需要?jiǎng)⊥敢稽c(diǎn)的是，后文提到的兩個(gè)場(chǎng)景下的rst的bug，都是由于相同的race condition導(dǎo)致的。rcu保護(hù)關(guān)注的是reader&writer的安全性（不會(huì)踩錯(cuò)地址），而不保護(hù)數(shù)據(jù)的實(shí)時(shí)性，這個(gè)很重要。所以當(dāng)rcu與hashtable結(jié)合的時(shí)候，對(duì)整個(gè)表的增刪和讀如何保證數(shù)據(jù)的絕對(duì)的同步顯得很重要！

(1) 握手階段的TCP bug

問(wèn)題的表象是，三次握手完畢后client端給server端發(fā)送了數(shù)據(jù)，結(jié)果server端卻發(fā)送了rst拒絕了。

分析：注意看上圖最左邊的第4和5這兩行的時(shí)間間隔非常短，只有11微妙，11微妙是什么概念？查一次tcp socket的hash表可能都是幾十微妙，這點(diǎn)時(shí)間完全可能會(huì)停頓在一個(gè)函數(shù)上。

當(dāng)server端看到第三行的ack的時(shí)候幾乎同時(shí)也看到了第四行的數(shù)據(jù)，詳細(xì)來(lái)說(shuō)，這時(shí)候server端在握手最后一個(gè)環(huán)節(jié)，會(huì)在socket的hash表中刪除一個(gè)老的socket（我們叫req sk），再插入一個(gè)新的socket（我們叫full sk），在刪除和插入之間的這短暫的幾微妙發(fā)生的時(shí)候，server收第行的數(shù)據(jù)的時(shí)候需要去到這個(gè)hash表中尋找（根據(jù)五元組）對(duì)應(yīng)的socket來(lái)接受這個(gè)報(bào)文，結(jié)果在這個(gè)空檔期間沒(méi)有匹配到應(yīng)該找到的socket，這時(shí)候沒(méi)辦法只能把當(dāng)時(shí)上層最初監(jiān)聽(tīng)的listener拿出來(lái)接收，這樣就出現(xiàn)了錯(cuò)誤，違背了協(xié)議棧的基本的設(shè)計(jì)：對(duì)于listener socket接收到了數(shù)據(jù)包，那么這個(gè)數(shù)據(jù)包是非預(yù)期的，應(yīng)該發(fā)送RST！

   CPU 0                           CPU 1
   -----                           -----
tcp_v4_rcv()                  syn_recv_sock()
                            inet_ehash_insert()
                            -> sk_nulls_del_node_init_rcu(osk)
__inet_lookup_established()
                            -> __sk_nulls_add_node_rcu(sk, list)

對(duì)應(yīng)上圖的cpu0就是server的第四行的讀者，cpu1就是寫者，對(duì)于cpu0而言，讀到的數(shù)據(jù)可能是三種情況：1）讀到老的sk，2）讀到新的sk，3）誰(shuí)也讀不到，前兩個(gè)都是可以接收，但是最后一個(gè)就是bug了——我們必須要找到兩者之一！如下就是一種場(chǎng)景，無(wú)法正確找到new或者old。

那如何修復(fù)這個(gè)問(wèn)題？在排查完整個(gè)握手規(guī)則后，發(fā)現(xiàn)只需要先插入新的sk到hash桶的尾部，再刪除老的sk即可，這樣就會(huì)有幾種情況：1）兩個(gè)同時(shí)都在，一定能匹配到其中一個(gè)，2）匹配到新的。如下圖，無(wú)論reader在哪里都能保證可以讀到一個(gè)。如下是正確的：

結(jié)論：第3行（client給server發(fā)生了握手最后一次ack）和第4行（client端給server發(fā)送了第一組數(shù)據(jù)）出現(xiàn)的并發(fā)問(wèn)題。

(2) 揮手階段的bug

這個(gè)問(wèn)題根因同上：rcu+hash表的使用問(wèn)題，在揮手階段發(fā)起close()的一方競(jìng)爭(zhēng)的亂序的收到了一個(gè)ack和一個(gè)fin ack觸發(fā)，導(dǎo)致socket在最后接收f(shuō)in ack時(shí)候沒(méi)有匹配到任何一個(gè)socket，又只能拿出最初監(jiān)聽(tīng)的listener來(lái)收包的時(shí)候，這時(shí)候出現(xiàn)了錯(cuò)誤。但是這個(gè)原始代碼中，是先插入新的sk再刪除了老的sk，乍一聽(tīng)沒(méi)有任何問(wèn)題，但是實(shí)際上插入新的sk出現(xiàn)了問(wèn)題，源碼中插入到頭部，這里需要插入到尾部才行！出現(xiàn)問(wèn)題的情景如下圖。

結(jié)論：這個(gè)是原生內(nèi)核長(zhǎng)達(dá)十多年的一個(gè)實(shí)現(xiàn)上的BUG，即為了性能考慮使用的RCU機(jī)制，由此必然引入的不準(zhǔn)確性導(dǎo)致并發(fā)的問(wèn)題，我定位并分析出這個(gè)問(wèn)題的并發(fā)的根因，由此提交了一份bugfix patch到社區(qū)被接收，鏈接：https://git.kernel.org/pub/scm/linux/kernel/git/netdev/net.git/commit/?id=3f4ca5fafc08881d7a57daa20449d171f2887043

第三個(gè)案例：netfilter兩個(gè)bug —— 數(shù)據(jù)傳輸RST

背景：用戶報(bào)告有兩個(gè)痛點(diǎn)問(wèn)題：偶發(fā)性出現(xiàn)1）根本無(wú)法完成三次握手連接，2）在傳輸數(shù)據(jù)的階段突然被RST異常中止。

分析：我們很容易的通過(guò)TCP的設(shè)計(jì)推測(cè)到這種情況一定不是正常的、符合預(yù)期的行為。我抓取了passive rst后發(fā)現(xiàn)原因是TCP層無(wú)法通過(guò)收到的skb包尋找到對(duì)應(yīng)的socket，要知道socket是最核心的TCP連接通信的基站，它保存了TCP應(yīng)有的信息（wscale、seq、buf等等），如果skb無(wú)法找到socket，那么就像小時(shí)候的故事小蝌蚪找媽媽但是找不到回家的路一樣。

那為什么會(huì)出現(xiàn)找不到socket？

經(jīng)過(guò)排查發(fā)現(xiàn)線上配置了DNAT規(guī)則，如下例子，凡是到達(dá)server端的1111端口或1112端口的都被轉(zhuǎn)發(fā)到80端口接收。

// iptables A port -> B port
iptables ... -p tcp --port 1111 -j REDIRECT --to-ports 80
iptables ... -p tcp --port 1112 -j REDIRECT --to-ports 80

DNAT+netfilter的流程是什么樣？

那么，有了DNAT之后，凡是進(jìn)入到server端的A port會(huì)被直接轉(zhuǎn)發(fā)到B port，最后TCP完成接收。完整的邏輯是這樣：DNAT的端口映射在ip層收包時(shí)候先進(jìn)入prerouting流程，修改skb的dst_ip:dst_port為真正的最后映射的信息，而后由ip early demux機(jī)制針對(duì)skb中的原始信息src_ip:src_port（也就是A port）修改為dst_ip:dst_port（也就是使用B port），由此4元組hash選擇一個(gè)sk，繼而成功由TCP接收才對(duì)。

(1) 兩條流沖突觸發(fā)的bug

如下，如果這時(shí)候有兩條流量想要TCP建連，二者都是由同一個(gè)client端相同的ip和port發(fā)起連接，這時(shí)候第1條連接首先發(fā)起握手那么肯定可以順利進(jìn)行，而當(dāng)?shù)?條連接發(fā)起的時(shí)候抵達(dá)到server端的1112端口最終被轉(zhuǎn)化為80端口，但是根據(jù)80端口可以發(fā)現(xiàn)我們已經(jīng)建立了連接，所以第2條流三次握手直接失敗。

1. saddr:12345 -> daddr:80 // 正常連接
2. saddr:12345 -> daddr:1112 -> daddr:80 // NAT參與轉(zhuǎn)化
（對(duì)內(nèi)核細(xì)節(jié)不感興趣的同學(xué)可以跳過(guò)此段）

我需要補(bǔ)充的信息是：NAT轉(zhuǎn)化port分為兩次，對(duì)于上述第二條流，第一次轉(zhuǎn)化1112為80，第二次轉(zhuǎn)化12345為1112，最終此流變?yōu)閇saddr:1112 -> daddr:80]

• 第一條流：skb對(duì)應(yīng)的sk是[saddr:12345 -> daddr:80]，這個(gè)沒(méi)有NAT參與。
• 第二條流：skb在ip層這時(shí)候NAT剛完成第一次port（修改dport 1112為dport 80），然后進(jìn)入了early demux機(jī)制，此時(shí)的4元組是[saddr:12345 -> daddr:80]，所以這時(shí)候匹配上了第一條流的sk，但是系統(tǒng)并不知情有問(wèn)題了，緊接著NAT第二次改變skb的port，變?yōu)閇saddr:1112 -> daddr:80]，這個(gè)也是后續(xù)TCP層延續(xù)使用的，雖然這個(gè)4元組信息是對(duì)的，但是已經(jīng)沒(méi)有用了，因?yàn)閑arly demux階段已經(jīng)獲取、保存socket了。

注：內(nèi)核修復(fù)后，對(duì)于第二條流就是放棄early demux階段選擇的4元組，而是安心等待NAT完成兩輪port的轉(zhuǎn)化之后，使用[saddr:1112 -> daddr:80]來(lái)匹配socket，這時(shí)候發(fā)現(xiàn)沒(méi)有對(duì)應(yīng)的socket，就找到了listener socket，從而完成三次連接。

結(jié)論：這個(gè)是early demux+DNAT的bug，它未能解決沖突問(wèn)題，導(dǎo)致了異常RST的發(fā)生。

(2) 特殊skb觸發(fā)的bug

注：在這個(gè)場(chǎng)景里面多了一個(gè)中間的gateway。

在本例中，我發(fā)現(xiàn)依然是熟知的一幕，skb無(wú)法lookup尋找到對(duì)應(yīng)的socket，此時(shí)我們要相信一定不會(huì)lookup算法出錯(cuò)，因?yàn)榇怂惴▋H僅是做簡(jiǎn)單的4元組的hash計(jì)算與匹配。所以追溯異常的skb和socket的四元組信息是頭等事情，經(jīng)過(guò)對(duì)比果然發(fā)現(xiàn)skb的端口信息未能成功被iptables轉(zhuǎn)化為B port，所以使用了含有A port的四元組信息去找socket，而socket當(dāng)初的建立是使用了B port，所以skb與sk的相遇就這么擦身而過(guò)了。

（對(duì)內(nèi)核細(xì)節(jié)不感興趣的同學(xué)可以跳過(guò)后面大段）

那么為什么會(huì)DNAT無(wú)法轉(zhuǎn)化？

我們先看下，異常未被轉(zhuǎn)化的skb和應(yīng)當(dāng)能接收的socket的4元組信息：

// 2.2.2.2是去敏后的server端ip地址，另外兩個(gè)是client的ip
sk info: 1.1.1.1:1111 <-> 2.2.2.2:80 // 我們可以知道真實(shí)的socket的建立是使用了80端口
skb info: 1.1.1.2:2222 <-> 2.2.2.2:1112 // 異常的skb未成功將1112端口轉(zhuǎn)化為80端口

client->gw->server的流程中，由于gw側(cè)發(fā)送了一些unknown skb再加上client端發(fā)送了一些out-of-window的包，導(dǎo)致進(jìn)入到server的netfilter階段會(huì)被識(shí)別出來(lái)INVALID異常，這個(gè)異常被識(shí)別后直接清除netfilter保持的該有的流信息，繼而異常的skb抵達(dá)DNAT階段后無(wú)法轉(zhuǎn)化端口（因?yàn)榕袛噢D(zhuǎn)化的流信息沒(méi)有了），最終skb無(wú)法成功轉(zhuǎn)化port端口號(hào)。

這個(gè)是netfilter+DNAT的設(shè)計(jì)上的bug，我認(rèn)為：無(wú)論是否有netfilter，都不應(yīng)當(dāng)是TCP的行為被改變，所以如果netfilter識(shí)別到了問(wèn)題所在，1）要么忽視，直接傳給TCP，交給TCP處理，2）要么丟棄，這樣也能避免RST的發(fā)生。但是，就這么一個(gè)小小的細(xì)節(jié)上，我和社區(qū)的幾個(gè)維護(hù)者拉鋸戰(zhàn)的battle了三百回合（鏈接：https://lore.kernel.org/all/20240311070550.7438-1-kerneljasonxing@gmail.com/），可惜雖然有一個(gè)維護(hù)者ACK了我的補(bǔ)丁，但是另外的維護(hù)者考慮netfilter不適合用于丟包功能，所以讓用戶去使用iptables --log功能、檢測(cè)出invalid異常包、繼而用iptables配置主動(dòng)丟棄。就憑這點(diǎn)，我認(rèn)為嚴(yán)重違背了user friendly的初衷，這些應(yīng)該是default默認(rèn)功能才對(duì)。此時(shí)的我雖然表面打不過(guò)，但是在內(nèi)心世界里很顯然我battle贏了...

結(jié)論：netfilter識(shí)別異常的skb未能成功保留DNAT信息，導(dǎo)致最后port端口不能成功被轉(zhuǎn)化，從而觸發(fā)了TCP的RST行為。

五、小結(jié)

RST問(wèn)題并不可怕，只要思路理清楚，先判斷類型，再抓取對(duì)應(yīng)代碼，繼而翻出RFC協(xié)議，最后分析源碼就能搞定，僅僅四步就可以了 ：）

希望這篇文章對(duì)大家有用！

六、附錄

這里列一下bcc的工具源碼，感興趣的同學(xué)可以自行查閱。如下是針對(duì)4.14內(nèi)核寫的，如果是更高版本需要調(diào)整一些python與c對(duì)照的格式問(wèn)題。

#!/usr/bin/env python

from __future__ import print_function
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
from time import sleep
from bcc import tcp

# arguments
examples = """examples:
    ./tcpdrop           # trace kernel TCP drops
"""
parser = argparse.ArgumentParser(
    description="Trace TCP drops by the kernel",
    formatter_class=argparse.RawDescriptionHelpFormatter,
    epilog=examples)
parser.add_argument("--ebpf", action="store_true",
    help=argparse.SUPPRESS)
args = parser.parse_args()
debug = 0

# define BPF program
bpf_text = """
#include <uapi/linux/ptrace.h>
#include <uapi/linux/tcp.h>
#include <uapi/linux/ip.h>
#include <net/sock.h>
#include <bcc/proto.h>

BPF_STACK_TRACE(stack_traces, 1024);

struct ipv4_data_t {
    u32 pid;
    u64 is_sknull;
    u32 saddr;
    u32 daddr;
    u16 sport;
    u16 dport;
    u8 state;
    u8 tcpflags;
    u32 stack_id;
};
BPF_PERF_OUTPUT(ipv4_events);

struct active_data_t {
    u32 pid;
    u32 saddr;
    u32 daddr;
    u16 sport;
    u16 dport;
    u32 stack_id;
};
BPF_PERF_OUTPUT(active_events);

static struct tcphdr *skb_to_tcphdr(const struct sk_buff *skb)
{
    // unstable API. verify logic in tcp_hdr() -> skb_transport_header().
    return (struct tcphdr *)(skb->head + skb->transport_header);
}

static inline struct iphdr *skb_to_iphdr(const struct sk_buff *skb)
{
    // unstable API. verify logic in ip_hdr() -> skb_network_header().
    return (struct iphdr *)(skb->head + skb->network_header);
}

// from include/net/tcp.h:
#ifndef tcp_flag_byte
#define tcp_flag_byte(th) (((u_int8_t *)th)[13])
#endif

int trace_tcp_v4_send_reset(struct pt_regs *ctx, struct sock *sk, struct sk_buff *skb)
{
    u8 is_sk_null = sk ? 0 : 1;
    u8 state = sk ? (u8)sk->__sk_common.skc_state : 1;
    u32 pid = bpf_get_current_pid_tgid();
    struct iphdr *ip = skb_to_iphdr(skb);
    u32 daddr = ip->daddr;
    u32 saddr = ip->saddr;

    // pull in details from the packet headers and the sock struct
    u16 family = sk->__sk_common.skc_family;
    u16 sport = 0, dport = 0;
    struct tcphdr *tcp = skb_to_tcphdr(skb);
    u8 tcpflags = ((u_int8_t *)tcp)[13];
    sport = tcp->source;
    dport = tcp->dest;
    sport = ntohs(sport);
    dport = ntohs(dport);

    if (family == AF_INET &&
        (saddr == 16777343 && daddr == 16777343) &&
        (sport == 8004 || dport == 8004)) {
        struct ipv4_data_t data4 = {};
        data4.pid = pid;
        data4.saddr = saddr;
        data4.daddr = daddr;
        data4.dport = dport;
        data4.sport = sport;
        data4.state = state;
        data4.tcpflags = tcpflags;
        data4.stack_id = stack_traces.get_stackid(ctx, 0);
        ipv4_events.perf_submit(ctx, &data4, sizeof(data4));

    }

    return 0;
}

int trace_tcp_send_active_reset(struct pt_regs *ctx, struct sock *sk, unsigned int priority)
{
    u32 pid = bpf_get_current_pid_tgid() >> 32;
    u32 saddr = 0, daddr = 0;
    u16 family = AF_INET;
    u16 sport = 0, dport = 0;

    // sport is not right
    sport = sk->__sk_common.skc_num;
    dport = sk->__sk_common.skc_dport;
    dport = ntohs(dport);

    saddr = sk->__sk_common.skc_rcv_saddr;
    daddr = sk->__sk_common.skc_daddr;

    if (family == AF_INET && (saddr == 16777343 && daddr == 16777343)) {
        struct active_data_t data4 = {};
        data4.pid = pid;
        data4.saddr = saddr;
        data4.daddr = daddr;
        data4.dport = dport;
        data4.sport = sport;
        data4.stack_id = stack_traces.get_stackid(ctx, 0);
        active_events.perf_submit(ctx, &data4, sizeof(data4));
    }

    return 0;
}
"""

if debug or args.ebpf:
    print(bpf_text)
    if args.ebpf:
        exit()

# event data
class Data_ipv4(ct.Structure):
    _fields_ = [
        ("pid", ct.c_uint),
        ("is_sknull", ct.c_ulonglong),
        ("saddr", ct.c_uint),
        ("daddr", ct.c_uint),
        ("sport", ct.c_ushort),
        ("dport", ct.c_ushort),
        ("state", ct.c_ubyte),
        ("tcpflags", ct.c_ubyte),
        ("stack_id", ct.c_ulong)
    ]

class Data_active(ct.Structure):
    _fields_ = [
        ("pid", ct.c_uint),
        ("saddr", ct.c_uint),
        ("daddr", ct.c_uint),
        ("sport", ct.c_ushort),
        ("dport", ct.c_ushort),
        ("stack_id", ct.c_ulong)
    ]

# process event
def print_ipv4_event(cpu, data, size):
    event = ct.cast(data, ct.POINTER(Data_ipv4)).contents
    if event.is_sknull is 1:
        print("%-8s %-7d %-20s > %-20s %s (%s)" % (
            strftime("%H:%M:%S"), event.pid,
            "%s:%d" % (inet_ntop(AF_INET, pack('I', event.saddr)), event.sport),
            "%s:%s" % (inet_ntop(AF_INET, pack('I', event.daddr)), event.dport),
            "sk-is-null", tcp.flags2str(event.tcpflags)))
    else:
        print("%-8s %-7d %-20s > %-20s %s (%s)" % (
            strftime("%H:%M:%S"), event.pid,
            "%s:%d" % (inet_ntop(AF_INET, pack('I', event.saddr)), event.sport),
            "%s:%s" % (inet_ntop(AF_INET, pack('I', event.daddr)), event.dport),
            tcp.tcpstate[event.state], tcp.flags2str(event.tcpflags)))
    for addr in stack_traces.walk(event.stack_id):
        sym = b.ksym(addr, show_offset=True)
        print("\t%s" % sym)
    print("")

def print_active_event(cpu, data, size):
    event = ct.cast(data, ct.POINTER(Data_active)).contents
    print("%-8s %-7d %-20s > %-20s" % (
        strftime("%H:%M:%S"), event.pid,
        "%s:%d" % (inet_ntop(AF_INET, pack('I', event.saddr)), event.sport),
        "%s:%d" % (inet_ntop(AF_INET, pack('I', event.daddr)), event.dport)))

    for addr in stack_traces.walk(event.stack_id):
        sym = b.ksym(addr, show_offset=True)
        print("\t%s" % sym)
    print("")

# initialize BPF
b = BPF(text=bpf_text)
if b.get_kprobe_functions(b"tcp_v4_send_reset"):
    b.attach_kprobe(event="tcp_v4_send_reset", fn_name="trace_tcp_v4_send_reset")
else:
    print("ERROR: tcp_drop() kernel function not found or traceable. "
        "Older kernel versions not supported.")
    exit()

if b.get_kprobe_functions(b"tcp_send_active_reset"):
    b.attach_kprobe(event="tcp_send_active_reset", fn_name="trace_tcp_send_active_reset")
else:
    print("ERROR: tcp_v4_send_reset() kernel function")
    exit()

stack_traces = b.get_table("stack_traces")

# header
print("%-8s %-6s %-2s %-20s > %-20s %s (%s)" % ("TIME", "PID", "IP",
    "SADDR:SPORT", "DADDR:DPORT", "STATE", "FLAGS"))

# read events
b["ipv4_events"].open_perf_buffer(print_ipv4_event)
#b["active_events"].open_perf_buffer(print_active_event)
while 1:
    try:
        b.perf_buffer_poll()
    except KeyboardInterrupt:
        exit()