Skip to content

Instantly share code, notes, and snippets.

@frispete

frispete/ctypes_mmap_ctx.py

Last active January 6, 2017 00:01
Show Gist options
  • Save frispete/97c27e24a0aae1bcaf1375e2e463d239 to your computer and use it in GitHub Desktop.
Save frispete/97c27e24a0aae1bcaf1375e2e463d239 to your computer and use it in GitHub Desktop.
Download ZIP
Demonstrate problem with context manager on mmaped ctypes structures
Raw
ctypes_mmap_ctx.py
#!/usr/bin/env python3
# -*- coding: utf8 -*
import os
import mmap
import ctypes
import logging
import weakref
from contextlib import contextmanager
log = logging.getLogger(__file__)
NOPROB = True
#NOPROB = False
WEAKREF = True
#WEAKREF = False
def align(size, alignment):
"""return size aligned to alignment"""
excess = size % alignment
if excess:
size = size - excess + alignment
return size
@contextmanager
def cstructmap(cstruct, mm, offset = 0):
# resize the mmap (and backing file), if structure exceeds mmap size
# mmap size must be aligned to mmap.PAGESIZE
cssize = ctypes.sizeof(cstruct)
if offset + cssize > mm.size():
newsize = align(offset + cssize, mmap.PAGESIZE)
mm.resize(newsize)
csinst = cstruct.from_buffer(mm, offset)
if WEAKREF:
yield weakref.proxy(csinst)
else:
yield csinst
class FileHeader(ctypes.BigEndianStructure):
Identifier = b'hfMM'[::-1]
_fields_ = [
('identifier', ctypes.c_char * 4), # 4
('offset', ctypes.c_uint64), # 8
]
FileHeaderSize = ctypes.sizeof(FileHeader)
class ItemHeader(ctypes.BigEndianStructure):
Identifier = b'hiMM'[::-1]
_fields_ = [
('identifier', ctypes.c_char * 4), # 4
('length', ctypes.c_uint32), # 4
]
ItemHeaderSize = ctypes.sizeof(ItemHeader)
class MapFile:
"""Manage memory mapped file"""
def __init__(self, filename):
self._created = False
try:
# try to create initial file
mapsize = mmap.PAGESIZE
self._fd = open(filename, 'x+b')
self._fd.write(b'\0' * mapsize)
self._created = True
except FileExistsError:
# file exists and is writable
mapsize = os.path.getsize(filename)
self._fd = open(filename, 'r+b')
# mmap this file
self._fd.seek(0)
self._mm = mmap.mmap(self._fd.fileno(), mapsize)
with cstructmap(FileHeader, self._mm) as fh:
if self._created:
fh.identifier = FileHeader.Identifier
fh.offset = fh_offset = FileHeaderSize
elif fh.identifier == FileHeader.Identifier:
fh_offset = fh.offset
else:
log.critical('%s: invalid FileHeader identifier %r',
filename, fh.identifier)
sys.exit(2)
self._offset = FileHeaderSize
if self._created:
log.debug('starting offset: 0x%x', self._offset)
return
# consistency check: move hand over hand along the items
while self._offset < self._mm.size():
with cstructmap(ItemHeader, self._mm, self._offset) as ih:
if ih.identifier == ItemHeader.Identifier:
self._offset += ih.length
else:
if self._offset != fh_offset:
log.error('%s: inconsistent header offset: %s != %s',
filename, fh_offset, self._offset)
sys.exit(3)
rest = self._mm.size() - self._offset
if rest:
overhang = ctypes.c_ubyte * rest
with cstructmap(overhang, self._mm, self._offset) as blk:
if bytes(blk) != bytes(rest):
log.error('%s: overhang not zero', filename)
sys.exit(4)
if NOPROB:
del blk
break
if NOPROB:
del ih
log.debug('starting offset: 0x%x', self._offset)
def add_data(self, data):
datasize = len(data)
log.debug('add_data: header')
with cstructmap(ItemHeader, self._mm, self._offset) as ih:
ih.identifier = ItemHeader.Identifier
ih.length = ItemHeaderSize + datasize
self._offset += ItemHeaderSize
if NOPROB:
del ih
log.debug('add_data: %s', datasize)
blktype = ctypes.c_char * datasize
with cstructmap(blktype, self._mm, self._offset) as blk:
blk.raw = data
self._offset += datasize
if NOPROB:
del blk
return ItemHeaderSize + datasize
def size(self):
return self._mm.size()
def close(self):
with cstructmap(FileHeader, self._mm) as fh:
fh.offset = self._offset
if NOPROB:
del fh
self._mm.close()
self._fd.close()
log.debug('final offset: 0x%x', self._offset)
if __name__ == '__main__':
import sys
logconfig = dict(
level = logging.DEBUG,
format = '%(levelname)5s: %(message)s',
)
logging.basicConfig(**logconfig)
mapfile = sys.argv[1:2] or 'mapfile'
datafile = sys.argv[2:3] or __file__
data = open(datafile, 'rb').read()
mf = MapFile(mapfile)
maxsize = mf.size() + 10 * mmap.PAGESIZE
while mf.size() < maxsize:
mf.add_data(data)
mf.close()
Raw
ctypes_mmap_ctx2.py
#!/usr/bin/env python3
# -*- coding: utf8 -*
import os
import mmap
import ctypes
import logging
from contextlib import contextmanager
log = logging.getLogger(__file__)
def align(size, alignment):
"""return size aligned to alignment"""
excess = size % alignment
if excess:
size = size - excess + alignment
return size
class FileHeader(ctypes.BigEndianStructure):
Identifier = b'hfMM'[::-1]
_fields_ = [
('identifier', ctypes.c_char * 4), # 4(+4 padding bytes)
('offset', ctypes.c_uint64), # 8
]
FileHeaderSize = ctypes.sizeof(FileHeader)
class ItemHeader(ctypes.BigEndianStructure):
Identifier = b'hiMM'[::-1]
_fields_ = [
('identifier', ctypes.c_char * 4), # 4
('length', ctypes.c_uint32), # 4
]
ItemHeaderSize = ctypes.sizeof(ItemHeader)
class _cstructmap:
def __init__(self, cstruct, mm, offset):
self.m = cstruct.from_buffer(mm, offset)
def close(self):
self.m = None
@contextmanager
def cstructmap(cstruct, mm, offset = 0):
# resize the mmap (and backing file), if structure exceeds mmap size
# mmap size must be aligned to mmap.PAGESIZE
cssize = ctypes.sizeof(cstruct)
if offset + cssize > mm.size():
newsize = align(offset + cssize, mmap.PAGESIZE)
mm.resize(newsize)
mapped = _cstructmap(cstruct, mm, offset)
try:
yield mapped
finally:
mapped.close()
class MapFile:
"""Manage memory mapped file"""
def __init__(self, filename):
self._created = False
try:
# try to create initial file
mapsize = mmap.PAGESIZE
self._fd = open(filename, 'x+b')
self._fd.write(b'\0' * mapsize)
self._created = True
except FileExistsError:
# file exists and is writable
mapsize = os.path.getsize(filename)
self._fd = open(filename, 'r+b')
# mmap this file
self._fd.seek(0)
self._mm = mmap.mmap(self._fd.fileno(), mapsize)
with cstructmap(FileHeader, self._mm) as fh:
if self._created:
fh.m.identifier = FileHeader.Identifier
fh.m.offset = fh_offset = FileHeaderSize
elif fh.m.identifier == FileHeader.Identifier:
fh_offset = fh.m.offset
else:
log.critical('%s: invalid FileHeader identifier %r',
filename, fh.m.identifier)
sys.exit(2)
self._offset = FileHeaderSize
if self._created:
log.debug('starting offset: 0x%x', self._offset)
return
# consistency check: rummage the items
while self._offset < self._mm.size():
with cstructmap(ItemHeader, self._mm, self._offset) as ih:
if ih.m.identifier == ItemHeader.Identifier:
self._offset += ih.m.length
else:
if self._offset != fh_offset:
log.error('%s: inconsistent header offset: %s != %s',
filename, fh_offset, self._offset)
sys.exit(3)
rest = self._mm.size() - self._offset
if rest:
if self._mm[self._offset:self._offset+rest] != bytes(rest):
log.error('%s: overhang not zero', filename)
sys.exit(4)
break
log.debug('starting offset: 0x%x', self._offset)
def add_data(self, data):
datasize = len(data)
log.debug('add_data: header')
with cstructmap(ItemHeader, self._mm, self._offset) as ih:
ih.m.identifier = ItemHeader.Identifier
ih.m.length = ItemHeaderSize + datasize
self._offset += ItemHeaderSize
log.debug('add_data: %s', datasize)
blktype = ctypes.c_char * datasize
with cstructmap(blktype, self._mm, self._offset) as blk:
blk.m.raw = data
self._offset += datasize
return ItemHeaderSize + datasize
def size(self):
return self._mm.size()
def close(self):
with cstructmap(FileHeader, self._mm) as fh:
fh.m.offset = self._offset
self._mm.close()
self._fd.close()
log.debug('final offset: 0x%x', self._offset)
if __name__ == '__main__':
import sys
logconfig = dict(
level = logging.DEBUG,
format = '%(levelname)5s: %(message)s',
)
logging.basicConfig(**logconfig)
mapfile = sys.argv[1:2] or 'mapfile'
datafile = sys.argv[2:3] or __file__
data = open(datafile, 'rb').read()
mf = MapFile(mapfile)
maxsize = mf.size() + 10 * mmap.PAGESIZE
while mf.size() < maxsize:
mf.add_data(data)
mf.close()
Raw
ctypes_mmap_ctx3.py
#!/usr/bin/env python3
# -*- coding: utf8 -*
import os
import mmap
import ctypes
import logging
from contextlib import contextmanager
log = logging.getLogger(__file__)
def align(size, alignment):
"""return size aligned to alignment"""
excess = size % alignment
if excess:
size = size - excess + alignment
return size
class FileHeader(ctypes.BigEndianStructure):
Identifier = b'hfMM'[::-1]
_fields_ = [
('identifier', ctypes.c_char * 4), # 4(+4 padding bytes)
('offset', ctypes.c_uint64), # 8
]
FileHeaderSize = ctypes.sizeof(FileHeader)
class ItemHeader(ctypes.BigEndianStructure):
Identifier = b'hiMM'[::-1]
_fields_ = [
('identifier', ctypes.c_char * 4), # 4
('length', ctypes.c_uint32), # 4
]
ItemHeaderSize = ctypes.sizeof(ItemHeader)
@contextmanager
def cstructmap(cstruct, mm, offset = 0):
# resize the mmap (and backing file), if structure exceeds mmap size
# mmap size must be aligned to mmap.PAGESIZE
cssize = ctypes.sizeof(cstruct)
if offset + cssize > mm.size():
newsize = align(offset + cssize, mmap.PAGESIZE)
mm.resize(newsize)
cmap = cstruct.from_buffer(mm, offset)
try:
yield cmap
finally:
for mv in cmap._objects.values():
if isinstance(mv, memoryview):
mv.release()
class MapFile:
"""Manage memory mapped file"""
def __init__(self, filename):
self._created = False
try:
# try to create initial file
mapsize = mmap.PAGESIZE
self._fd = open(filename, 'x+b')
self._fd.write(b'\0' * mapsize)
self._created = True
except FileExistsError:
# file exists and is writable
mapsize = os.path.getsize(filename)
self._fd = open(filename, 'r+b')
# mmap this file
self._fd.seek(0)
self._mm = mmap.mmap(self._fd.fileno(), mapsize)
with cstructmap(FileHeader, self._mm) as fh:
if self._created:
fh.identifier = FileHeader.Identifier
fh.offset = fh_offset = FileHeaderSize
elif fh.identifier == FileHeader.Identifier:
fh_offset = fh.offset
else:
log.critical('%s: invalid FileHeader identifier %r',
filename, fh.identifier)
sys.exit(2)
self._offset = FileHeaderSize
if self._created:
log.debug('starting offset: 0x%x', self._offset)
return
# consistency check: rummage the items
while self._offset < self._mm.size():
with cstructmap(ItemHeader, self._mm, self._offset) as ih:
if ih.identifier == ItemHeader.Identifier:
self._offset += ih.length
else:
if self._offset != fh_offset:
log.error('%s: inconsistent header offset: %s != %s',
filename, fh_offset, self._offset)
sys.exit(3)
rest = self._mm.size() - self._offset
if rest:
if self._mm[self._offset:self._offset+rest] != bytes(rest):
log.error('%s: overhang not zero', filename)
sys.exit(4)
break
log.debug('starting offset: 0x%x', self._offset)
def add_data(self, data):
datasize = len(data)
log.debug('add_data: header')
with cstructmap(ItemHeader, self._mm, self._offset) as ih:
ih.identifier = ItemHeader.Identifier
ih.length = ItemHeaderSize + datasize
self._offset += ItemHeaderSize
log.debug('add_data: %s', datasize)
blktype = ctypes.c_char * datasize
with cstructmap(blktype, self._mm, self._offset) as blk:
blk.raw = data
self._offset += datasize
return ItemHeaderSize + datasize
def size(self):
return self._mm.size()
def close(self):
with cstructmap(FileHeader, self._mm) as fh:
fh.offset = self._offset
self._mm.close()
self._fd.close()
log.debug('final offset: 0x%x', self._offset)
if __name__ == '__main__':
import sys
logconfig = dict(
level = logging.DEBUG,
format = '%(levelname)5s: %(message)s',
)
logging.basicConfig(**logconfig)
mapfile = sys.argv[1:2] or 'mapfile'
datafile = sys.argv[2:3] or __file__
data = open(datafile, 'rb').read()
mf = MapFile(mapfile)
maxsize = mf.size() + 10 * mmap.PAGESIZE
while mf.size() < maxsize:
mf.add_data(data)
mf.close()
Raw
mmap_test.py
import ctypes
import mmap
from contextlib import contextmanager
class T(ctypes.Structure):
_fields = [("foo", ctypes.c_uint32)]
@contextmanager
def map_struct(m, n):
m.resize(n * mmap.PAGESIZE)
yield T.from_buffer(m)
SIZE = mmap.PAGESIZE * 2
f = open("tmp.dat", "w+b")
f.write(b"\0" * SIZE)
f.seek(0)
m = mmap.mmap(f.fileno(), mmap.PAGESIZE)
with map_struct(m, 1) as a:
a.foo = 1
with map_struct(m, 2) as b:
b.foo = 2
Raw
mmap_test_weakref.py
import ctypes
import mmap
import weakref
from contextlib import contextmanager
class T(ctypes.Structure):
_fields = [("foo", ctypes.c_uint32)]
@contextmanager
def map_struct(m, n):
m.resize(n * mmap.PAGESIZE)
s = T.from_buffer(m)
yield weakref.proxy(s)
SIZE = mmap.PAGESIZE * 2
f = open("tmp.dat", "w+b")
f.write(b"\0" * SIZE)
f.seek(0)
m = mmap.mmap(f.fileno(), mmap.PAGESIZE)
with map_struct(m, 1) as a:
a.foo = 1
with map_struct(m, 2) as b:
b.foo = 2
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment