एक्सप्लॉइट जनरेशन के लिए रिइन्फोर्समेंट लर्निंग

कुछ महीने पहले, एक CTF फाइनल के दौरान, हमारी टीम एक pwn बाइनरी पर तीन घंटे फँसी रही। स्टैक कैनरी, NX, आंशिक PIE, एक कम स्पष्ट फ़ंक्शन में ठीक 72 बाइट का ओवरफ़्लो। हमने सब कुछ आज़माया। कुछ नहीं हुआ। सुबह 03:14 बजे, हम में से एक ने लगभग मज़ाक में पाँच लाइन की एक स्क्रिप्ट चलाई जो AFL जैसे बिटफ़्लिप के साथ इनपुट को म्यूटेट करती थी और बिना सोचे फिर भेज देती थी। बीस मिनट बाद, उस स्क्रिप्ट ने शुद्ध सांख्यिकीय भाग्य से फ़्लैग जीत लिया।

उस रात हम एक सवाल लेकर घर गए जो हमारे दिमाग में घूमता रहा: यदि एक मूर्ख प्रक्रिया के बजाय हमारे पास एक ऐसी प्रक्रिया होती जो हर असफल प्रयास से सीखती, तो क्या होता?

RL कंट्रोल फ्लो हाइजैकिंग के साथ क्यों फिट होता है

एक एक्सप्लॉइट का सर्च स्पेस क्रूर है पर एकसमान नहीं। अधिकांश इनपुट एक ही व्यवहार उत्पन्न करते हैं, और केवल एक संकीर्ण क्षेत्र उस मशीन स्थिति की ओर ले जाता है जो आप चाहते हैं: इंस्ट्रक्शन रजिस्टर उस मेमोरी की ओर इशारा करता है जिसे आप नियंत्रित करते हैं। उस उप-स्थान को जो अलग करता है वह है सूचना का ग्रेडिएंट: एक इनपुट जो RIP को भ्रष्ट करने के करीब आता है, अवलोकन योग्य दुष्प्रभाव दिखाता है — आंशिक रूप से अधिलिखित रजिस्टर, इनपुट के बाइट्स से मेल खाने वाले स्टैक मान, अमान्य पतों पर छलाँग लेकिन नियंत्रित बाइट्स के पास।

यही वो चीज़ है जिसे अच्छी तरह से डिज़ाइन किए गए रिवॉर्ड फ़ंक्शन पकड़ सकते हैं।

क्षेत्र की वर्तमान स्थिति

Avgerinos और Brumley ने 2011 में सिंबोलिक एग्ज़ीक्यूशन पर पहला एंड-टू-एंड सिस्टम AEG (NDSS 2011) प्रकाशित किया। Mayhem (Cha et al., S&P 2012) ने विश्लेषण को बढ़ाया। दोनों लगभग नियतात्मक हैं पर सिंबोलिक एग्ज़ीक्यूशन की क्लासिक कीमत चुकाते हैं: पाथ एक्सप्लोज़न और दब चुके SMT सॉल्वर।

Böttinger ने Deep Reinforcement Fuzzing (2018) में DQN का उपयोग करके फ़ज़िंग को MDP के रूप में औपचारिक रूप दिया। PwnGPT (Shao et al., ACL 2025) o1-preview के साथ एक्सप्लॉइट जनरेशन दर 26.3% से 57.9% तक बढ़ाने की रिपोर्ट करता है। और PPO (Schulman et al., 2017) + RLHF (Ouyang et al., 2022) इतना परिपक्व हो गया है कि एक बाइनरी वातावरण पर एजेंट को प्रशिक्षित करना एक गंभीर बातचीत है।

हमारा सेटअप

ptrace पर Gym-शैली का वातावरण और एक eBPF साइडकार। ऑब्ज़र्वेशन एक डेंस टेन्सर है जिसमें:

• 16 जनरल रजिस्टरों (x86_64) की स्थिति, सामान्यीकृत।
• RSP के चारों ओर स्टैक विंडो (256 बाइट्स) कच्चे बाइट्स के रूप में।
• नियंत्रण फ़्लैग: NX, canary_present, RIP_corrupted, register_controlled[]।
• eBPF से ली गई अंतिम कवरेज ट्रेस का हैश (प्रत्येक बेसिक ब्लॉक पर uprobe)।

एक्शन स्पेस इनपुट बाइट्स पर अलग है। रिवॉर्ड नई कवरेज (सकारात्मक, घटती हुई), एक्सप्लॉइटेबिलिटी ह्यूरिस्टिक्स और लंबाई जुर्माना को संयोजित करता है।

import gym
from gwaihir.tools.rl_env import BinaryPwnEnv

class CTFPwnEnv(BinaryPwnEnv):
    def step(self, action):
        mutated = self.apply_action(self.current_input, action)
        trace = self.run_with_ebpf(mutated)

        new_bbs = len(trace.basic_blocks - self.coverage_seen)
        cov_reward = np.log1p(new_bbs) * 0.3
        self.coverage_seen |= trace.basic_blocks

        expl = 0.0
        if trace.rip_controlled_bytes > 0:
            expl += 5.0 + 0.1 * trace.rip_controlled_bytes
        if trace.canary_leaked: expl += 2.5
        if trace.tainted_syscall_args: expl += 1.5

        reward = cov_reward + expl - 0.001 * len(mutated)
        done = trace.rip_fully_controlled or self.steps > 4096
        return self.observe(trace), reward, done, {"trace": trace}

PPO बनाम DQN

हमने DQN से शुरुआत की। जैसे ही एक्शन स्पेस बढ़ा, DQN नाज़ुक हो गया: हर बार जब एक्सप्लॉइटेबिलिटी रिवॉर्ड एक विशाल और अलग-थलग मान को ट्रिगर करता था तो टार्गेट नेटवर्क डायवर्ज हो जाता था। PPO अपने पॉलिसी रेशियो की क्लिपिंग के साथ उन स्पाइक्स को बिना ढहे अवशोषित करता है। GAE (λ=0.95), एन्ट्रॉपी कोएफ़ 0.01 और 4096 के बैचों के साथ PPO ने हमें एक RTX 4090 पर 24 घंटे की स्थिर ट्रेनिंग दी।

एक विवरण जो अक्सर नहीं बताया जाता: अधिकांश समय ग्रेडिएंट पर नहीं जाता, यह बाइनरी चलाने पर जाता है। प्रत्येक स्टेप के लिए fork, ptrace, eBPF संग्रह, पार्स की आवश्यकता होती है। हमने SubprocVecEnv के साथ प्रति GPU 32 envs को समानांतर किया और फिर भी बॉटलनेक I/O था।

10 CTF बाइनरीज़ पर परिणाम

बैटरी: दो babys, तीन canary leak आवश्यक के साथ, दो format string के साथ, दो न्यूनतम ROP चेन के साथ और एक heap (UAF) के साथ। एजेंट ने, प्रति बाइनरी शून्य से प्रशिक्षित होकर, 10 में से 7 हल किए। तीन जो असफल रहे: heap वाला (अपेक्षित), एक ROP जिसमें एक गैजेट की आवश्यकता थी जो मौजूद नहीं था, और एक format string जहाँ पॉलिसी पतों को हमेशा के लिए लीक करते हुए स्थानीय न्यूनतम में फँस गई।

शुद्ध angr के विरुद्ध तुलना: RL ने 5 जीते, 2 बराबरी, 3 हारे (angr ने heap और format-write पहले हल किए)। यह एक श्रेष्ठ उपकरण नहीं है, यह पूरक है।

क्रूर सीमाएँ

• प्रति बाइनरी एक पॉलिसी। हम चुनौतियों के बीच ज्ञान स्थानांतरित नहीं करते।
• घंटों का ट्रेनिंग समय उन चुनौतियों के लिए जिन्हें एक सीनियर इंसान मिनटों में हल करता है।
• नाज़ुक रिवॉर्ड शेपिंग। कस्टम एलोकेटर या गैर-मानक शमन ह्यूरिस्टिक्स को तोड़ देते हैं।
• सिंबोलिक एग्ज़ीक्यूशन अभी भी जीतता है जब पाथ नियतात्मक रूप से व्युत्पन्न होता है।

असली ट्रेड-ऑफ़: सस्ती व्यापक कवरेज (RL) बनाम महंगा गहरा रीज़निंग (सिंबोलिक)। हम उन्हें Gwaihir के अंदर समानांतर चलाते हैं।

आगे क्या आ रहा है

तीन मोर्चे: (a) exploit-db और CTF राइटअप पर प्रीट्रेनिंग, (b) RLHF शैली का रिवॉर्ड लर्निंग, (c) एजेंट को एक कोऑर्डिनेटर LLM के सब-टूल के रूप में एकीकृत करना।

अभी हम प्रशिक्षण जारी रखते हैं। और अगली बार जब कोई टीम सुबह तीन बजे तीन घंटे तक फँसी हो, हम चाहते हैं कि एजेंट ही फ़्लैग ढूँढे — भाग्य नहीं।

ग्रंथसूची

1. Avgerinos, T. et al. (2011). AEG: Automatic Exploit Generation. NDSS.
2. Cha, S.K. et al. (2012). Unleashing Mayhem. IEEE S&P.
3. Böttinger, K. et al. (2018). Deep Reinforcement Fuzzing. IEEE SPW. arXiv:1801.04589.
4. Schulman, J. et al. (2017). PPO. arXiv:1707.06347.
5. Ouyang, L. et al. (2022). InstructGPT/RLHF. NeurIPS.
6. Shao, Y. et al. (2025). PwnGPT. ACL.