From 742a7c7e78b599eaa214740f863c94aa56857209 Mon Sep 17 00:00:00 2001
From: William Woodruff <william@trailofbits.com>
Date: Sat, 21 Mar 2020 10:29:46 -0400
Subject: [PATCH] Add support for DW_LNE_set_discriminator  (#282)

---
 elftools/dwarf/constants.py    |  3 +++
 elftools/dwarf/lineprogram.py  |  9 ++++++++-
 test/test_dwarf_lineprogram.py | 23 ++++++++++++++++++++++-
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/elftools/dwarf/constants.py b/elftools/dwarf/constants.py
index 09264f7..79db975 100644
--- a/elftools/dwarf/constants.py
+++ b/elftools/dwarf/constants.py
@@ -156,6 +156,9 @@ DW_LNS_set_isa = 0x0c
 DW_LNE_end_sequence = 0x01
 DW_LNE_set_address = 0x02
 DW_LNE_define_file = 0x03
+DW_LNE_set_discriminator = 0x04
+DW_LNE_lo_user = 0x80
+DW_LNE_hi_user = 0xff
 
 
 # Call frame instructions
diff --git a/elftools/dwarf/lineprogram.py b/elftools/dwarf/lineprogram.py
index 8996b5c..a79b37d 100644
--- a/elftools/dwarf/lineprogram.py
+++ b/elftools/dwarf/lineprogram.py
@@ -58,12 +58,14 @@ class LineState(object):
         self.prologue_end = False
         self.epilogue_begin = False
         self.isa = 0
+        self.discriminator = 0
 
     def __repr__(self):
         a = ['<LineState %x:' % id(self)]
         a.append('  address = 0x%x' % self.address)
         for attr in ('file', 'line', 'column', 'is_stmt', 'basic_block',
-                     'end_sequence', 'prologue_end', 'epilogue_begin', 'isa'):
+                     'end_sequence', 'prologue_end', 'epilogue_begin', 'isa',
+                     'discriminator'):
             a.append('  %s = %s' % (attr, getattr(self, attr)))
         return '\n'.join(a) + '>\n'
 
@@ -130,6 +132,7 @@ class LineProgram(object):
             # After adding, clear some state registers.
             entries.append(LineProgramEntry(
                 cmd, is_extended, args, copy.copy(state)))
+            state.discriminator = 0
             state.basic_block = False
             state.prologue_end = False
             state.epilogue_begin = False
@@ -188,6 +191,10 @@ class LineProgram(object):
                         self.structs.Dwarf_lineprog_file_entry, self.stream)
                     self['file_entry'].append(operand)
                     add_entry_old_state(ex_opcode, [operand], is_extended=True)
+                elif ex_opcode == DW_LNE_set_discriminator:
+                    operand = struct_parse(self.structs.Dwarf_uleb128(''),
+                                           self.stream)
+                    state.discriminator = operand
                 else:
                     # Unknown, but need to roll forward the stream because the
                     # length is specified. Seek forward inst_len - 1 because
diff --git a/test/test_dwarf_lineprogram.py b/test/test_dwarf_lineprogram.py
index 5a5c65a..2a0a19e 100644
--- a/test/test_dwarf_lineprogram.py
+++ b/test/test_dwarf_lineprogram.py
@@ -18,7 +18,7 @@ class TestLineProgram(unittest.TestCase):
         """
         ds = DWARFStructs(little_endian=True, dwarf_format=32, address_size=4)
         header = ds.Dwarf_lineprog_header.parse(
-            b'\x04\x10\x00\x00' +    # initial lenght
+            b'\x04\x10\x00\x00' +    # initial length
             b'\x03\x00' +            # version
             b'\x20\x00\x00\x00' +    # header length
             b'\x01\x01\x01\x0F' +    # flags
@@ -100,6 +100,27 @@ class TestLineProgram(unittest.TestCase):
         self.assertLineState(linetable[7].state, address=0x24b, line=7, end_sequence=False)
         self.assertLineState(linetable[9].state, address=0x24d, line=7, end_sequence=True)
 
+    def test_lne_set_discriminator(self):
+        """
+        Tests the handling of DWARFv4's new DW_LNE_set_discriminator opcode.
+        """
+        s = BytesIO()
+        s.write(
+            b'\x00\x02\x04\x05' +  # DW_LNE_set_discriminator (discriminator=0x05)
+            b'\x01' +              # DW_LNS_copy
+            b'\x00\x01\x01'        # DW_LNE_end_sequence
+        )
+
+        lp = self._make_program_in_stream(s)
+        linetable = lp.get_entries()
+
+        # We expect two entries, since DW_LNE_set_discriminator does not add
+        # an entry of its own.
+        self.assertEqual(len(linetable), 2)
+        self.assertEqual(linetable[0].command, DW_LNS_copy)
+        self.assertLineState(linetable[0].state, discriminator=0x05)
+        self.assertLineState(linetable[1].state, discriminator=0x00, end_sequence=True)
+
 
 if __name__ == '__main__':
     unittest.main()
-- 
2.30.2