File size: 4,911 Bytes
3f69d68 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | import random
import unittest
import torch
from lightnest import rv32
from lightnest.learned_patch import (Example, LearnedPatch, MAX_WINDOW, PatchNetwork, Predictor, TRAIN_REGS, TEST_REGS, TRAIN_TEMPLATES, TEST_TEMPLATES, batch, encode, independent_expected, make_example, validate_candidate)
class ForcedNetwork(torch.nn.Module):
def __init__(self, offset, new):
super().__init__()
self.offset, self.new = offset, new
def forward(self, command, binary):
scores = torch.full((1, 257), -30.)
scores[0, self.offset] = 30.
bits = torch.full((1, 256, 8), -30.)
if self.offset < 256:
for i in range(8):
bits[0, self.offset, i] = 30. if self.new & (1 << i) else -30.
return scores, bits
class LearnedPatchTests(unittest.TestCase):
@classmethod
def setUpClass(cls):
torch.set_num_threads(2)
def test_dataset_oracle_and_semantics(self):
rng = random.Random(1)
for split in ('train', 'combined'):
for _ in range(60):
e = make_example(rng, split)
self.assertLessEqual(len(e.binary), 256)
diffs = [i for i, (a, b) in enumerate(zip(e.binary, e.target)) if a != b]
self.assertEqual(len(diffs), int(e.present))
if not e.present:
self.assertEqual(e.pointer, MAX_WINDOW)
for x in ([3, 7], [-27, 12], [0x7fffffff, 1]):
expected = independent_expected(e, x)
self.assertEqual(rv32.simulate(e.target, x), expected)
if e.present:
patch = LearnedPatch(e.pointer, e.binary[e.pointer], e.replacement)
self.assertTrue(validate_candidate(e.binary, patch, [([3, 7], rv32.simulate(e.target, [3, 7]))]))
def test_disjoint_holdouts(self):
self.assertFalse(set(TRAIN_REGS) & set(TEST_REGS))
self.assertFalse(set(TRAIN_TEMPLATES) & set(TEST_TEMPLATES))
rng = random.Random(2)
for _ in range(100):
train, test = make_example(rng), make_example(rng, 'combined')
self.assertLessEqual(train.prefix_words, 16)
self.assertGreaterEqual(test.prefix_words, 22)
self.assertTrue(set(test.registers) <= set(TEST_REGS))
self.assertIn(test.template, TEST_TEMPLATES)
def test_proposal_follows_neural_output_not_labels(self):
# A deliberately nonsensical instruction still returns forced neural bytes.
predictor = Predictor(ForcedNetwork(5, 123))
patch = predictor.propose('arbitrary unknown text', bytes(range(32)), base_offset=2**40)
self.assertEqual((patch.relative_offset, patch.old_byte, patch.new_byte), (5, 5, 123))
self.assertEqual(patch.absolute_offset, 2**40 + 5)
self.assertIn('unvalidated', predictor.last_status)
def test_none_and_reject(self):
self.assertIsNone(Predictor(ForcedNetwork(256, 0)).propose('change add to sub', b'1234'))
rng = random.Random(3)
e = make_example(rng)
bad = LearnedPatch(0, e.binary[0], 255)
self.assertFalse(validate_candidate(e.binary, bad, [([7, 4], rv32.simulate(e.target, [7, 4]))]))
self.assertFalse(validate_candidate(e.binary, bad, []))
self.assertFalse(validate_candidate(e.binary, None, [([7, 4], 3)]))
def test_resources_and_bounds(self):
model = PatchNetwork()
self.assertLess(sum(p.numel() for p in model.parameters()), 1_000_000)
e = make_example(random.Random(4))
c, b, _, _ = batch([e])
logits, bits = model(c, b)
self.assertEqual(tuple(logits.shape), (1, 257))
self.assertEqual(tuple(bits.shape), (1, 256, 8))
for command, binary in [('a'*97, b'1'), ('a', b'1'*257), ('', b'1'), ('a', b'')]:
with self.assertRaises(ValueError):
encode(command, binary)
with self.assertRaises(ValueError):
Predictor(model).propose('a', b'1', -1)
if __name__ == '__main__':
unittest.main()
def test_nonfinite_policy_fails_closed(tmp_path):
import pytest
import torch
from lightnest.learned_patch import PatchNetwork, Predictor, load_predictor
for threshold in (float('nan'), float('inf'), 0., 1.1):
with pytest.raises(ValueError):
Predictor(PatchNetwork(), threshold)
class Broken(torch.nn.Module):
def forward(self, command, binary):
return torch.full((1,257),float('nan')), torch.zeros((1,256,8))
with pytest.raises(ValueError,match='nonfinite'):
Predictor(Broken()).propose('change add to sub',b'\x13\0\0\0')
state=PatchNetwork().state_dict()
state['pointer.bias'][0]=float('nan')
path=tmp_path/'bad.pt';torch.save({'state_dict':state,'threshold':.55},path)
with pytest.raises(ValueError,match='tensor'):
load_predictor(path)
|