MCGMark: An Encodable and Robust Online Watermark for Tracing LLM-Generated Malicious Code

Abstract

With the advent of large language models (LLMs), numerous software service providers are developing LLMs tailored for code generation, such as CodeLlama. However, these models can be exploited by malicious developers to generate malicious code, posing severe threats to the software ecosystem. To address this issue, we first conducted an empirical study and built MCGTest , a dataset of \(406\) prompts designed to elicit malicious code from LLMs. Leveraging this dataset, we propose MCGMark , a watermarking method to trace and attribute LLM-generated malicious code. MCGMark subtly embeds user-specific information into generated code by controlling the token selection process, ensuring the watermark is imperceptible. Additionally, MCGMark dynamically adjusts the token selection range to induce the LLM to favor high-probability tokens, thus ensuring code quality. Furthermore, by leveraging code structure, MCGMark avoids embedding watermarks into regions easily modified by attackers, such as comments and variable names, enhancing robustness against tampering. Experiments on several advanced LLMs show that MCGMark successfully embeds watermarks in approximately \(85\%\) of cases, under the constraint of a \(400\) -token limit. Moreover, it maintains code quality and demonstrates strong resilience against common code modification. This approach offers a practical solution for tracing malicious code and mitigating the misuse of LLMs.