v0.1 - initial commit

cyclic.go

@@ -0,0 +1,124 @@
+package winpwn
+
+import "bytes"
+
+const (
+	cyclicAlphabet = "abcdefghijklmnopqrstuvwxyz"
+	cyclicDefaultN = 4
+)
+
+// Cyclic returns the first length bytes of a de Bruijn sequence over the
+// lowercase alphabet with subsequence length 4 (pwntools' default), the Go
+// analogue of pwnlib.util.cyclic.cyclic -- a buffer where every 4-byte
+// window is unique, so a crash address/register read back out of it tells
+// you exactly how far into the buffer it pointed.
+func Cyclic(length int) []byte {
+	return CyclicN(length, cyclicDefaultN)
+}
+
+// CyclicN is Cyclic with an explicit subsequence length n instead of the
+// default 4 (e.g. 8 to locate an offset into a 64-bit pointer-sized
+// overwrite).
+func CyclicN(length, n int) []byte {
+	if length <= 0 {
+		return []byte{}
+	}
+	out := make([]byte, 0, length)
+	deBruijnEach(cyclicAlphabet, n, func(b byte) bool {
+		out = append(out, b)
+		return len(out) < length
+	})
+	if len(out) >= length || len(out) == 0 {
+		return out[:length]
+	}
+	// The sequence's period (len(alphabet)^n) was shorter than the
+	// requested length -- only possible for a tiny alphabet/n combination,
+	// never in practice with the default 26-letter alphabet. Wrap.
+	full := make([]byte, length)
+	for i := range full {
+		full[i] = out[i%len(out)]
+	}
+	return full
+}
+
+// CyclicFind returns the offset of subseq within the default (alphabet,
+// n=4) de Bruijn sequence, or -1 if it can't appear in it -- the analogue
+// of pwnlib.util.cyclic.cyclic_find. subseq is typically the 4 bytes read
+// back from a crashed return address/register.
+func CyclicFind(subseq []byte) int {
+	return CyclicFindN(subseq, cyclicDefaultN)
+}
+
+// CyclicFindN is CyclicFind with an explicit subsequence length n, matching
+// whatever n the buffer was generated with via CyclicN. Stops generating as
+// soon as a match is found, so this stays fast even for n=8 (where the
+// theoretical period, 26^8, is far too large to ever materialize) as long
+// as subseq actually came from a real CyclicN(_, n) buffer, which is the
+// only case this is ever used for.
+func CyclicFindN(subseq []byte, n int) int {
+	m := len(subseq)
+	if m == 0 {
+		return -1
+	}
+
+	window := make([]byte, 0, m)
+	total := 0
+	found := -1
+
+	deBruijnEach(cyclicAlphabet, n, func(b byte) bool {
+		if len(window) < m {
+			window = append(window, b)
+		} else {
+			copy(window, window[1:])
+			window[m-1] = b
+		}
+		total++
+		if len(window) == m && bytes.Equal(window, subseq) {
+			found = total - m
+			return false
+		}
+		return true
+	})
+
+	return found
+}
+
+// deBruijnEach generates the de Bruijn sequence B(k, n) over alphabet (k =
+// len(alphabet)) one byte at a time via yield, stopping as soon as yield
+// returns false -- the Go equivalent of pwntools' de_bruijn() being a
+// Python generator. This early-exit property is the entire point: B(k, n)'s
+// period is k^n, which is astronomical for n=8 even with a small alphabet
+// (26^8 ~ 2*10^11) and must never be materialized in full. Callers only
+// ever need a bounded prefix (CyclicN) or an early match (CyclicFindN), and
+// this lets both stop without walking the rest of the sequence. Uses the
+// same Fredricksen-Kessler-Maiorana algorithm pwntools' cyclic() does.
+func deBruijnEach(alphabet string, n int, yield func(byte) bool) {
+	k := len(alphabet)
+	a := make([]int, k*n)
+	stop := false
+
+	var db func(t, p int)
+	db = func(t, p int) {
+		if stop {
+			return
+		}
+		if t > n {
+			if n%p == 0 {
+				for _, idx := range a[1 : p+1] {
+					if !yield(alphabet[idx]) {
+						stop = true
+						return
+					}
+				}
+			}
+			return
+		}
+		a[t] = a[t-p]
+		db(t+1, p)
+		for j := a[t-p] + 1; j < k && !stop; j++ {
+			a[t] = j
+			db(t+1, t)
+		}
+	}
+	db(1, 1)
+}