llmgw/internal/guard/detect/injection.go

package detect

import (
	"encoding/base64"
	"encoding/hex"
	"net/url"
	"regexp"
	"strings"
	"unicode"

	"golang.org/x/text/unicode/norm"
)

// Normalize is text as a model reads it, with what hides it from a pattern
// taken off: compatibility forms folded (NFKC: fullwidth letters,
// ligatures), invisible and direction-changing characters removed, common
// look-alike letters from other scripts folded to Latin, and case lowered.
func Normalize(s string) string {
	s = norm.NFKC.String(s)
	var b strings.Builder
	b.Grow(len(s))
	for _, r := range s {
		if invisible(r) {
			continue
		}
		if l, ok := homoglyphs[r]; ok {
			r = l
		}
		b.WriteRune(unicode.ToLower(r))
	}
	return b.String()
}

// invisible reports characters that change what text says without being
// seen: zero-width ones, direction overrides, and Unicode tags, which carry
// a whole hidden ASCII message.
func invisible(r rune) bool {
	switch {
	case r >= 0x200B && r <= 0x200F, r >= 0x202A && r <= 0x202E, r >= 0x2060 && r <= 0x2069,
		r == 0xFEFF, r == 0x00AD, r >= 0xE0000 && r <= 0xE007F:
		return true
	}
	return false
}

// homoglyphs are Cyrillic and Greek letters drawn like Latin ones.
var homoglyphs = map[rune]rune{
	'а': 'a', 'е': 'e', 'о': 'o', 'р': 'p', 'с': 'c', 'у': 'y', 'х': 'x', 'і': 'i', 'ј': 'j', 'ѕ': 's',
	'А': 'a', 'В': 'b', 'Е': 'e', 'К': 'k', 'М': 'm', 'Н': 'h', 'О': 'o', 'Р': 'p', 'С': 'c', 'Т': 't', 'Х': 'x',
	'α': 'a', 'ο': 'o', 'ρ': 'p', 'ν': 'v', 'ι': 'i', 'κ': 'k', 'τ': 't', 'υ': 'u',
	'Α': 'a', 'Β': 'b', 'Ε': 'e', 'Η': 'h', 'Ι': 'i', 'Κ': 'k', 'Μ': 'm', 'Ν': 'n', 'Ο': 'o', 'Ρ': 'p', 'Τ': 't', 'Χ': 'x',
}

// Variants are the text and what it decodes to — base64, hex and URL
// encoding, to depth two — normalized: each a way the same words can reach
// a model.
func Variants(s string) []string {
	out := []string{Normalize(s)}
	frontier := []string{s}
	for depth := 0; depth < 2; depth++ {
		var next []string
		for _, t := range frontier {
			for _, d := range decodings(t) {
				out = append(out, Normalize(d))
				next = append(next, d)
			}
		}
		frontier = next
	}
	return out
}

var (
	b64Re = regexp.MustCompile(`[A-Za-z0-9+/_\-]{24,}={0,2}`)
	hexRe = regexp.MustCompile(`\b(?:[0-9a-fA-F]{2}){12,}\b`)
)

func decodings(s string) []string {
	var out []string
	for _, m := range b64Re.FindAllString(s, 8) {
		for _, enc := range []*base64.Encoding{base64.StdEncoding, base64.RawStdEncoding, base64.URLEncoding, base64.RawURLEncoding} {
			if b, err := enc.DecodeString(m); err == nil && printable(b) {
				out = append(out, string(b))
				break
			}
		}
	}
	for _, m := range hexRe.FindAllString(s, 8) {
		if b, err := hex.DecodeString(m); err == nil && printable(b) {
			out = append(out, string(b))
		}
	}
	if strings.Contains(s, "%") {
		if u, err := url.QueryUnescape(s); err == nil && u != s {
			out = append(out, u)
		}
	}
	return out
}

// printable reports whether decoded bytes read as text, so random tokens that
// happen to be valid base64 are not taken for hidden instructions.
func printable(b []byte) bool {
	if len(b) < 8 {
		return false
	}
	n := 0
	for _, c := range string(b) {
		if unicode.IsPrint(c) || c == '\n' || c == '\t' {
			n++
		}
	}
	return n*10 >= len([]rune(string(b)))*9
}

var phraseRes = []*regexp.Regexp{
	regexp.MustCompile(`\b(ignore|disregard|forget|override)\b.{0,30}\b(previous|prior|above|earlier|all|your|system)\b.{0,30}\b(instructions?|prompts?|rules|guidelines|directions)\b`),
	regexp.MustCompile(`\byou are now\b.{0,40}\b(unrestricted|jailbroken|dan|developer mode|free|without (any )?(rules|limits|restrictions))\b`),
	regexp.MustCompile(`\b(reveal|print|show|repeat|output)\b.{0,30}\b(your|the) (system prompt|initial instructions|hidden instructions)\b`),
	regexp.MustCompile(`\b(do anything now|developer mode enabled|jailbreak mode)\b`),
	// Turkish.
	regexp.MustCompile(`(önceki|yukarıdaki|tüm|bütün)\s.{0,30}(talimat|kural|yönerge)\S*\s.{0,20}(yok say|görmezden gel|unut)`),
	regexp.MustCompile(`sistem (istemini|promptunu|talimatlarını)\s.{0,20}(göster|yaz|söyle|tekrarla)`),
}

// templateTokens are chat templates' control tokens: in a user's text they
// try to open a new system or assistant turn the template did not.
var templateTokens = []string{
	"<|im_start|>", "<|im_end|>", "<|system|>", "<|assistant|>", "<|user|>", "<|eot_id|>",
	"<|start_header_id|>", "<|end_header_id|>", "[inst]", "[/inst]", "<<sys>>", "<</sys>>",
	"<start_of_turn>", "<end_of_turn>", "<|endoftext|>",
}

// Injection reports the kinds of prompt-injection signs in text: in what it
// says once normalized and decoded, and — for invisible text and encoded
// blobs — in how it is written.
func Injection(text string) []string {
	var kinds []string
	add := func(k string) {
		for _, have := range kinds {
			if have == k {
				return
			}
		}
		kinds = append(kinds, k)
	}
	for _, r := range text {
		if invisible(r) && r != 0x00AD {
			add(InjectionInvisible)
			break
		}
	}
	for _, m := range b64Re.FindAllString(text, -1) {
		if len(m) >= 200 {
			add(InjectionEncoded)
			break
		}
	}
	for _, v := range Variants(text) {
		for _, re := range phraseRes {
			if re.MatchString(v) {
				add(InjectionPhrase)
			}
		}
		for _, t := range templateTokens {
			if strings.Contains(v, t) {
				add(InjectionTemplate)
			}
		}
	}
	return kinds
}