| 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): |
| |
| 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) |
|
|
