From 94177d594aa9ce08f05c7e7b44e5ca90c63a8cf6 Mon Sep 17 00:00:00 2001 From: Ali Younis Date: Tue, 16 May 2017 19:43:50 -0700 Subject: [PATCH] PyORAm --- PyORAM/.appveyor.yml | 18 + PyORAM/.gitignore | 22 + PyORAM/.travis.yml | 28 + PyORAM/.travis_before_install_jython.sh | 12 + PyORAM/CHANGELOG.rst | 25 + PyORAM/LICENSE.md | 21 + PyORAM/MANIFEST.in | 5 + PyORAM/Makefile | 21 + PyORAM/README.rst | 244 ++++ PyORAM/examples/aesctr_performance.py | 118 ++ PyORAM/examples/ali.py | 176 +++ PyORAM/examples/encrypted_heap_test.py | 57 + PyORAM/examples/encrypted_storage_file.py | 79 ++ PyORAM/examples/encrypted_storage_mmap.py | 79 ++ PyORAM/examples/encrypted_storage_ram.py | 85 ++ PyORAM/examples/encrypted_storage_s3.py | 92 ++ PyORAM/examples/encrypted_storage_sftp.py | 99 ++ PyORAM/examples/path_oram_file.py | 85 ++ PyORAM/examples/path_oram_mmap.py | 86 ++ PyORAM/examples/path_oram_ram.py | 91 ++ PyORAM/examples/path_oram_s3.py | 98 ++ PyORAM/examples/path_oram_sftp.py | 105 ++ PyORAM/examples/path_oram_sftp_setup.py | 115 ++ PyORAM/examples/path_oram_sftp_test.py | 99 ++ PyORAM/examples/tree_oram_test.py | 93 ++ PyORAM/requirements.txt | 6 + PyORAM/setup.cfg | 6 + PyORAM/setup.py | 72 + PyORAM/src/_cffi_src/__init__.py | 0 .../_cffi_src/virtual_heap_helper_build.py | 74 + PyORAM/src/pyoram/__about__.py | 17 + PyORAM/src/pyoram/__init__.py | 55 + PyORAM/src/pyoram/crypto/__init__.py | 1 + PyORAM/src/pyoram/crypto/aes.py | 48 + .../src/pyoram/encrypted_storage/__init__.py | 3 + .../encrypted_block_storage.py | 326 +++++ .../encrypted_heap_storage.py | 172 +++ .../top_cached_encrypted_heap_storage.py | 305 ++++ .../src/pyoram/oblivious_storage/__init__.py | 1 + .../pyoram/oblivious_storage/tree/__init__.py | 2 + .../oblivious_storage/tree/path_oram.py | 466 +++++++ .../tree/tree_oram_helper.py | 329 +++++ PyORAM/src/pyoram/storage/AliTimer.py | 64 + PyORAM/src/pyoram/storage/__init__.py | 7 + PyORAM/src/pyoram/storage/block_storage.py | 83 ++ .../src/pyoram/storage/block_storage_file.py | 394 ++++++ .../src/pyoram/storage/block_storage_mmap.py | 143 ++ .../src/pyoram/storage/block_storage_ram.py | 314 +++++ PyORAM/src/pyoram/storage/block_storage_s3.py | 396 ++++++ .../src/pyoram/storage/block_storage_sftp.py | 151 ++ PyORAM/src/pyoram/storage/boto3_s3_wrapper.py | 158 +++ PyORAM/src/pyoram/storage/heap_storage.py | 245 ++++ PyORAM/src/pyoram/tests/__init__.py | 0 .../src/pyoram/tests/baselines/exists.empty | 0 .../src/pyoram/tests/baselines/k200_h0_b1.dot | 5 + .../tests/baselines/k200_h0_b1_data.dot | 5 + .../src/pyoram/tests/baselines/k2_h3_b1.dot | 33 + .../pyoram/tests/baselines/k2_h3_b1_data.dot | 33 + .../src/pyoram/tests/baselines/k2_h3_b2.dot | 33 + .../pyoram/tests/baselines/k2_h3_b2_data.dot | 33 + .../src/pyoram/tests/baselines/k3_h3_b1.dot | 83 ++ .../pyoram/tests/baselines/k3_h3_b1_data.dot | 83 ++ .../src/pyoram/tests/baselines/k3_h3_b2.dot | 83 ++ .../pyoram/tests/baselines/k3_h3_b2_data.dot | 83 ++ PyORAM/src/pyoram/tests/test_aes.py | 108 ++ PyORAM/src/pyoram/tests/test_block_storage.py | 1023 ++++++++++++++ .../tests/test_encrypted_block_storage.py | 814 +++++++++++ .../tests/test_encrypted_heap_storage.py | 594 ++++++++ PyORAM/src/pyoram/tests/test_examples.py | 61 + PyORAM/src/pyoram/tests/test_heap_storage.py | 555 ++++++++ PyORAM/src/pyoram/tests/test_misc.py | 124 ++ PyORAM/src/pyoram/tests/test_package.py | 38 + PyORAM/src/pyoram/tests/test_path_oram.py | 851 ++++++++++++ .../test_top_cached_encrypted_heap_storage.py | 711 ++++++++++ PyORAM/src/pyoram/tests/test_virtual_heap.py | 1227 +++++++++++++++++ PyORAM/src/pyoram/util/__init__.py | 2 + PyORAM/src/pyoram/util/misc.py | 117 ++ PyORAM/src/pyoram/util/virtual_heap.py | 441 ++++++ .../light_fan_embed_benchmark/Sensor.java | 97 +- 79 files changed, 12989 insertions(+), 39 deletions(-) create mode 100644 PyORAM/.appveyor.yml create mode 100644 PyORAM/.gitignore create mode 100644 PyORAM/.travis.yml create mode 100644 PyORAM/.travis_before_install_jython.sh create mode 100644 PyORAM/CHANGELOG.rst create mode 100644 PyORAM/LICENSE.md create mode 100644 PyORAM/MANIFEST.in create mode 100644 PyORAM/Makefile create mode 100644 PyORAM/README.rst create mode 100644 PyORAM/examples/aesctr_performance.py create mode 100644 PyORAM/examples/ali.py create mode 100644 PyORAM/examples/encrypted_heap_test.py create mode 100644 PyORAM/examples/encrypted_storage_file.py create mode 100644 PyORAM/examples/encrypted_storage_mmap.py create mode 100644 PyORAM/examples/encrypted_storage_ram.py create mode 100644 PyORAM/examples/encrypted_storage_s3.py create mode 100644 PyORAM/examples/encrypted_storage_sftp.py create mode 100644 PyORAM/examples/path_oram_file.py create mode 100644 PyORAM/examples/path_oram_mmap.py create mode 100644 PyORAM/examples/path_oram_ram.py create mode 100644 PyORAM/examples/path_oram_s3.py create mode 100644 PyORAM/examples/path_oram_sftp.py create mode 100644 PyORAM/examples/path_oram_sftp_setup.py create mode 100644 PyORAM/examples/path_oram_sftp_test.py create mode 100644 PyORAM/examples/tree_oram_test.py create mode 100644 PyORAM/requirements.txt create mode 100644 PyORAM/setup.cfg create mode 100644 PyORAM/setup.py create mode 100644 PyORAM/src/_cffi_src/__init__.py create mode 100644 PyORAM/src/_cffi_src/virtual_heap_helper_build.py create mode 100644 PyORAM/src/pyoram/__about__.py create mode 100644 PyORAM/src/pyoram/__init__.py create mode 100644 PyORAM/src/pyoram/crypto/__init__.py create mode 100644 PyORAM/src/pyoram/crypto/aes.py create mode 100644 PyORAM/src/pyoram/encrypted_storage/__init__.py create mode 100644 PyORAM/src/pyoram/encrypted_storage/encrypted_block_storage.py create mode 100644 PyORAM/src/pyoram/encrypted_storage/encrypted_heap_storage.py create mode 100644 PyORAM/src/pyoram/encrypted_storage/top_cached_encrypted_heap_storage.py create mode 100644 PyORAM/src/pyoram/oblivious_storage/__init__.py create mode 100644 PyORAM/src/pyoram/oblivious_storage/tree/__init__.py create mode 100644 PyORAM/src/pyoram/oblivious_storage/tree/path_oram.py create mode 100644 PyORAM/src/pyoram/oblivious_storage/tree/tree_oram_helper.py create mode 100644 PyORAM/src/pyoram/storage/AliTimer.py create mode 100644 PyORAM/src/pyoram/storage/__init__.py create mode 100644 PyORAM/src/pyoram/storage/block_storage.py create mode 100644 PyORAM/src/pyoram/storage/block_storage_file.py create mode 100644 PyORAM/src/pyoram/storage/block_storage_mmap.py create mode 100644 PyORAM/src/pyoram/storage/block_storage_ram.py create mode 100644 PyORAM/src/pyoram/storage/block_storage_s3.py create mode 100644 PyORAM/src/pyoram/storage/block_storage_sftp.py create mode 100644 PyORAM/src/pyoram/storage/boto3_s3_wrapper.py create mode 100644 PyORAM/src/pyoram/storage/heap_storage.py create mode 100644 PyORAM/src/pyoram/tests/__init__.py create mode 100644 PyORAM/src/pyoram/tests/baselines/exists.empty create mode 100644 PyORAM/src/pyoram/tests/baselines/k200_h0_b1.dot create mode 100644 PyORAM/src/pyoram/tests/baselines/k200_h0_b1_data.dot create mode 100644 PyORAM/src/pyoram/tests/baselines/k2_h3_b1.dot create mode 100644 PyORAM/src/pyoram/tests/baselines/k2_h3_b1_data.dot create mode 100644 PyORAM/src/pyoram/tests/baselines/k2_h3_b2.dot create mode 100644 PyORAM/src/pyoram/tests/baselines/k2_h3_b2_data.dot create mode 100644 PyORAM/src/pyoram/tests/baselines/k3_h3_b1.dot create mode 100644 PyORAM/src/pyoram/tests/baselines/k3_h3_b1_data.dot create mode 100644 PyORAM/src/pyoram/tests/baselines/k3_h3_b2.dot create mode 100644 PyORAM/src/pyoram/tests/baselines/k3_h3_b2_data.dot create mode 100644 PyORAM/src/pyoram/tests/test_aes.py create mode 100644 PyORAM/src/pyoram/tests/test_block_storage.py create mode 100644 PyORAM/src/pyoram/tests/test_encrypted_block_storage.py create mode 100644 PyORAM/src/pyoram/tests/test_encrypted_heap_storage.py create mode 100644 PyORAM/src/pyoram/tests/test_examples.py create mode 100644 PyORAM/src/pyoram/tests/test_heap_storage.py create mode 100644 PyORAM/src/pyoram/tests/test_misc.py create mode 100644 PyORAM/src/pyoram/tests/test_package.py create mode 100644 PyORAM/src/pyoram/tests/test_path_oram.py create mode 100644 PyORAM/src/pyoram/tests/test_top_cached_encrypted_heap_storage.py create mode 100644 PyORAM/src/pyoram/tests/test_virtual_heap.py create mode 100644 PyORAM/src/pyoram/util/__init__.py create mode 100644 PyORAM/src/pyoram/util/misc.py create mode 100644 PyORAM/src/pyoram/util/virtual_heap.py diff --git a/PyORAM/.appveyor.yml b/PyORAM/.appveyor.yml new file mode 100644 index 0000000..51cbb6e --- /dev/null +++ b/PyORAM/.appveyor.yml @@ -0,0 +1,18 @@ +environment: + matrix: + - PYTHON: "C:\\Python27" + - PYTHON: "C:\\Python33" + - PYTHON: "C:\\Python34" + - PYTHON: "C:\\Python35" + - PYTHON: "C:\\Python27-x64" + - PYTHON: "C:\\Python35-x64" +install: + - "SET PATH=%PYTHON%;%PYTHON%\\Scripts;%PATH%" + - "%PYTHON%\\python.exe -m pip install -v -U pip setuptools" + - "%PYTHON%\\python.exe -m pip install -v -e ." + - "%PYTHON%\\python.exe -m pip install -U unittest2 nose2 cov-core codecov coverage" +build: off +test_script: + - "%PYTHON%\\python.exe -m nose2 -v --log-capture --with-coverage --coverage src --coverage examples -s src" +on_success: + - "%PYTHON%\\Scripts\\codecov.exe" diff --git a/PyORAM/.gitignore b/PyORAM/.gitignore new file mode 100644 index 0000000..043977f --- /dev/null +++ b/PyORAM/.gitignore @@ -0,0 +1,22 @@ +# Emacs +*~ +**/\#* + +# Python +*.py[cod] + +# C extensions +*.so + +# Setuptools distribution folder. +/dist/ +/build/ + +# Python egg metadata, regenerated from source files by setuptools. +*.egg-info +*.egg +*.eggs + +# nose +.coverage +coverage.xml diff --git a/PyORAM/.travis.yml b/PyORAM/.travis.yml new file mode 100644 index 0000000..db24225 --- /dev/null +++ b/PyORAM/.travis.yml @@ -0,0 +1,28 @@ +# travis CI config +language: python +matrix: + include: +# - python: 2.7 +# env: JYTHON=org.python:jython-installer:2.7.1 + - python: 2.7 + - python: 3.3 + - python: 3.4 + - python: 3.5 + - python: pypy +cache: false +before_install: + - sudo apt-get update -q + - sudo apt-get install graphviz -y + - python -m pip install -v -U pip setuptools virtualenv wheel + - if [ -n "$JYTHON" ]; then bash ./.travis_before_install_jython.sh; source $HOME/myvirtualenv/bin/activate ;fi + - if [ -n "$JYTHON" ]; then pip install jip; fi + - if [ -n "$JYTHON" ]; then export CLASSPATH=$VIRTUAL_ENV/javalib/*; fi +install: + - python -m pip install -v -e . + - python -m pip install -U unittest2 nose2 cov-core codecov coverage +script: python -m nose2 -v --log-capture --with-coverage --coverage src --coverage examples -s src +after_success: + - codecov +branches: + only: + - master diff --git a/PyORAM/.travis_before_install_jython.sh b/PyORAM/.travis_before_install_jython.sh new file mode 100644 index 0000000..8ca64a4 --- /dev/null +++ b/PyORAM/.travis_before_install_jython.sh @@ -0,0 +1,12 @@ +#!/bin/bash +set -e + +pip install jip +jip install $JYTHON +NON_GROUP_ID=${JYTHON#*:} +_JYTHON_BASENAME=${NON_GROUP_ID/:/-} +OLD_VIRTUAL_ENV=${VIRTUAL_ENV:=.} +java -jar $OLD_VIRTUAL_ENV/javalib/${_JYTHON_BASENAME}.jar -s -d $HOME/jython +$HOME/jython/bin/jython -c "import sys; print(sys.version_info)" +virtualenv --version +virtualenv -p $HOME/jython/bin/jython $HOME/myvirtualenv diff --git a/PyORAM/CHANGELOG.rst b/PyORAM/CHANGELOG.rst new file mode 100644 index 0000000..e29dfa7 --- /dev/null +++ b/PyORAM/CHANGELOG.rst @@ -0,0 +1,25 @@ +Changelog +========= + +0.3.0 - `master`_ +~~~~~~~~~~~~~~~~~ + +0.2.0 - 2016-10-18 +~~~~~~~~~~~~~~~~~~ + +* using chunking to speed up yield_blocks for SFTP +* speed up clearing entries in S3 interface by chunking delete requests +* adding helper property to access heap storage on path oram +* use a mmap to store the top-cached heap buckets +* replace the show_status_bar keywords by a global config item +* express status bar units as a memory transfer rate during setup +* tweaks to Path ORAM to make it easier to generalize to other schemes +* changing suffix of S3 index file from txt to bin +* updates to readme + +0.1.2 - 2016-05-15 +~~~~~~~~~~~~~~~~~~ + +* Initial release. + +.. _`master`: https://github.com/ghackebeil/PyORAM diff --git a/PyORAM/LICENSE.md b/PyORAM/LICENSE.md new file mode 100644 index 0000000..4d2d76c --- /dev/null +++ b/PyORAM/LICENSE.md @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) 2016 Gabriel Hackebeil + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/PyORAM/MANIFEST.in b/PyORAM/MANIFEST.in new file mode 100644 index 0000000..888e697 --- /dev/null +++ b/PyORAM/MANIFEST.in @@ -0,0 +1,5 @@ +include README.rst +include CHANGELOG.rst +include LICENSE.md + +recursive-include src/_cffi_src *.py diff --git a/PyORAM/Makefile b/PyORAM/Makefile new file mode 100644 index 0000000..cba9335 --- /dev/null +++ b/PyORAM/Makefile @@ -0,0 +1,21 @@ +all: + +.PHONY: clean +clean: + find examples -name "*.pyc" | xargs rm + find src -name "*.pyc" | xargs rm + find . -depth 1 -name "*.pyc" | xargs rm + + find examples -name "*.pyo" | xargs rm + find src -name "*.pyo" | xargs rm + find . -depth 1 -name "*.pyo" | xargs rm + + find examples -name "__pycache__" | xargs rm -r + find src -name "__pycache__" | xargs rm -r + find . -depth 1 -name "__pycache__" | xargs rm -r + + find examples -name "*~" | xargs rm + find src -name "*~" | xargs rm + find . -depth 1 -name "*~" | xargs rm + + find src -name "*.so" | xargs rm diff --git a/PyORAM/README.rst b/PyORAM/README.rst new file mode 100644 index 0000000..07c9340 --- /dev/null +++ b/PyORAM/README.rst @@ -0,0 +1,244 @@ +PyORAM +====== + +.. image:: https://travis-ci.org/ghackebeil/PyORAM.svg?branch=master + :target: https://travis-ci.org/ghackebeil/PyORAM + +.. image:: https://ci.appveyor.com/api/projects/status/898bxsvqdch1btv6/branch/master?svg=true + :target: https://ci.appveyor.com/project/ghackebeil/PyORAM?branch=master + +.. image:: https://codecov.io/github/ghackebeil/PyORAM/coverage.svg?branch=master + :target: https://codecov.io/github/ghackebeil/PyORAM?branch=master + +.. image:: https://img.shields.io/pypi/v/PyORAM.svg + :target: https://pypi.python.org/pypi/PyORAM/ + +Python-based Oblivious RAM (PyORAM) is a collection of +Oblivious RAM algorithms implemented in Python. This package +serves to enable rapid prototyping and testing of new ORAM +algorithms and ORAM-based applications tailored for the +cloud-storage setting. PyORAM is written to support as many +Python versions as possible, including Python 2.7+, Python +3.3+, and PyPy 2.6+. + +This software is copyright (c) by Gabriel A. Hackebeil (gabe.hackebeil@gmail.com). + +This software is released under the MIT software license. +This license, including disclaimer, is available in the 'LICENSE' file. + +This work was funded by the Privacy Enhancing Technologies +project under the guidance of Professor `Attila Yavuz +`_ at Oregon State +University. + +Why Python? +----------- + +This project is meant for research. It is provided mainly as +a tool for other researchers studying the applicability of +ORAM to the cloud-storage setting. In such a setting, we +observe that network latency far outweighs any overhead +introduced from switching to an interpreted language such as +Python (as opposed to C++ or Java). Thus, our hope is that +by providing a Python-based library of ORAM tools, we will +enable researchers to spend more time prototyping new and +interesting ORAM applications and less time fighting with a +compiler or chasing down segmentation faults. + +Installation +------------ + +To install the latest release of PyORAM, simply execute:: + + $ pip install PyORAM + +To install the trunk version of PyORAM, first clone the repository:: + + $ git clone https://github.com/ghackebeil/PyORAM.git + +Next, enter the directory where PyORAM has been cloned and run setup:: + + $ python setup.py install + +If you are a developer, you should instead install using:: + + $ pip install -e . + $ pip install nose2 unittest2 + +Installation Tips +----------------- + +* OS X users are recommended to work with the `homebrew + `_ version of Python2 or Python3. If you + must use the default system Python, then the best thing to + do is create a virtual environment and install PyORAM into + that. The process of creating a virtual environment that is + stored in the PyORAM directory would look something like:: + + $ sudo pip install virtualenv + $ cd + $ virtualenv local_python2.7 + + If you had already attempted to install PyORAM into the + system Python and encountered errors, it may be necessary + to delete the directories :code:`build` and :code:`dist` + from the current directory using the command:: + + $ sudo rm -rf build dist + + Once this virtual environment has been successfully + created, you can *activate* it using the command:: + + $ . local_python2.7/bin/activate + + Then, proceed with the normal installation steps to + install PyORAM into this environment. Note that you must + *activate* this environment each time you open a new + terminal if PyORAM is installed in this way. Also, note + that use of the :code:`sudo` command is no longer + necessary (and should be avoided) once a virtual + environment is activated in the current shell. + +* If you have trouble installing the cryptography package + on OS X with PyPy: `stackoverflow `_. + +* If you encounter the dreaded "unable to find + vcvarsall.bat" error when installing packages with C + extensions through pip on Windows: `blog post `_. + +Tools Available (So Far) +------------------------ + +Encrypted block storage +~~~~~~~~~~~~~~~~~~~~~~~ + +* The basic building block for any ORAM implementation. + +* Available storage interfaces include: + + - local storage using a file, a memory-mapped file, or RAM + + + Dropbox + + - cloud storage using SFTP (requires SSH access to a server) + + + Amazon EC2 + + + Microsoft Azure + + + Google Cloud Platform + + - cloud storage using Amazon Simple Storage Service (S3) + +* See Examples: + + - examples/encrypted_storage_ram.py + + - examples/encrypted_storage_mmap.py + + - examples/encrypted_storage_file.py + + - examples/encrypted_storage_sftp.py + + - examples/encrypted_storage_s3.py + +Path ORAM +~~~~~~~~~ + +* Reference: `Stefanov et al. `_ + +* Generalized to work over k-kary storage heaps. Default + settings use a binary storage heap and bucket size + parameter set to 4. Using a k-ary storage heap can reduce + the access cost; however, stash size behavior has not been + formally analyzed in this setting. + +* Tree-Top caching can be used to reduce data transmission + per access as well as reduce access latency by exploiting + parallelism across independent sub-heaps below the last + cached heap level. + +* See Examples: + + - examples/path_oram_ram.py + + - examples/path_oram_mmap.py + + - examples/path_oram_file.py + + - examples/path_oram_sftp.py + + - examples/path_oram_s3.py + +Performance Tips +---------------- + +Setup Storage Locally +~~~~~~~~~~~~~~~~~~~~~ + +Storage schemes such as BlockStorageFile ("file"), BlockStorageMMap +("mmap"), BlockStorageRAM ("ram"), and BlockStorageSFTP ("sftp") all +employ the same underlying storage format. Thus, an oblivious storage +scheme can be initialized locally and then transferred to an external +storage location and accessed via BlockStorageSFTP using SSH login +credentials. See the following pair of files for an example of this: + +* examples/path_oram_sftp_setup.py + +* examples/path_oram_sftp_test.py + +BlockStorageS3 ("s3") employs a different format whereby the +underlying blocks are stored in separate "file" objects. +This design is due to the fact that the Amazon S3 API does +not allow modifications to a specific byte range within a +file, but instead requires that the entire modified file +object be re-uploaded. Thus, any efficient block storage +scheme must use separate "file" objects for each block. + +Tree-Top Caching +~~~~~~~~~~~~~~~~ + +For schemes that employ a storage heap (such as Path ORAM), +tree-top caching provides the ability to parallelize I/O +operations across the independent sub-heaps below the last +cached heap level. The default behavior of this +implementation of Path ORAM, for instance, caches the top +three levels of the storage heap in RAM, which creates eight +independent sub-heaps across which write operations can be +asynchronous. + +If the underlying storage is being accessed through SFTP, the +tree-top cached storage heap will attempt to open an +independent SFTP session for each sub-heap using the same +SSH connection. Typically, the maximum number of allowable +sessions associated with a single SSH connection is limited +by the SSH server. For instance, the default maximum number +of sessions allowed by a server using OpenSSH is 10. Thus, +increasing the number of cached levels beyond 3 when using +a binary storage heap will attempt to generate 16 or more SFTP +sessions and result in an error such as:: + + paramiko.ssh_exception.ChannelException: (1, 'Administratively prohibited') + +There are two options for avoiding this error: + +1. If you have administrative privileges on the server, you + can increase the maximum number of allowed sessions for a + single SSH connection. For example, to set the maximum + allowed sessions to 128 on a server using OpenSSH, one + would set:: + + MaxSessions 128 + + in :code:`/etc/ssh/sshd_config`, and then run the + command :code:`sudo service ssh restart`. + +2. You can limit the number of concurrent devices that will + be created by setting the concurrency level to something + below the last cached level using the + :code:`concurrency_level` keyword. For example, the + settings :code:`cached_levels=5` and + :code:`concurrency_level=0` would cache the top 5 levels + of the storage heap locally, but all external I/O + operations would take place through a single storage + device (e.g., using 1 SFTP session). diff --git a/PyORAM/examples/aesctr_performance.py b/PyORAM/examples/aesctr_performance.py new file mode 100644 index 0000000..374862d --- /dev/null +++ b/PyORAM/examples/aesctr_performance.py @@ -0,0 +1,118 @@ +import time +import base64 + +from pyoram.crypto.aes import AES + +def runtest(label, enc_func, dec_func): + print("") + print("$"*20) + print("{0:^20}".format(label)) + print("$"*20) + for keysize in AES.key_sizes: + print("") + print("@@@@@@@@@@@@@@@@@@@@") + print(" Key Size: %s bytes" % (keysize)) + print("@@@@@@@@@@@@@@@@@@@@") + print("\nTest Bulk") + # + # generate a key + # + key = AES.KeyGen(keysize) + print("Key: %s" % (base64.b64encode(key))) + + # + # generate some plaintext + # + nblocks = 1000000 + plaintext_numbytes = AES.block_size * nblocks + print("Plaintext Size: %s MB" + % (plaintext_numbytes * 1.0e-6)) + # all zeros + plaintext = bytes(bytearray(plaintext_numbytes)) + + # + # time encryption + # + start_time = time.time() + ciphertext = enc_func(key, plaintext) + stop_time = time.time() + print("Encryption Time: %.3fs (%.3f MB/s)" + % (stop_time-start_time, + (plaintext_numbytes * 1.0e-6) / (stop_time-start_time))) + + # + # time decryption + # + start_time = time.time() + plaintext_decrypted = dec_func(key, ciphertext) + stop_time = time.time() + print("Decryption Time: %.3fs (%.3f MB/s)" + % (stop_time-start_time, + (plaintext_numbytes * 1.0e-6) / (stop_time-start_time))) + + assert plaintext_decrypted == plaintext + assert ciphertext != plaintext + # IND-CPA + assert enc_func(key, plaintext) != ciphertext + # make sure the only difference is not in the IV + assert enc_func(key, plaintext)[AES.block_size:] \ + != ciphertext[AES.block_size:] + if enc_func is AES.CTREnc: + assert len(plaintext) == \ + len(ciphertext) - AES.block_size + else: + assert enc_func is AES.GCMEnc + assert len(plaintext) == \ + len(ciphertext) - 2*AES.block_size + + del plaintext + del plaintext_decrypted + del ciphertext + + print("\nTest Chunks") + # + # generate a key + # + key = AES.KeyGen(keysize) + print("Key: %s" % (base64.b64encode(key))) + + # + # generate some plaintext + # + nblocks = 1000 + blocksize = 16000 + total_bytes = blocksize * nblocks + print("Block Size: %s KB" % (blocksize * 1.0e-3)) + print("Block Count: %s" % (nblocks)) + print("Total: %s MB" % (total_bytes * 1.0e-6)) + plaintext_blocks = [bytes(bytearray(blocksize)) + for i in range(nblocks)] + + # + # time encryption + # + start_time = time.time() + ciphertext_blocks = [enc_func(key, b) + for b in plaintext_blocks] + stop_time = time.time() + print("Encryption Time: %.3fs (%.3f MB/s)" + % (stop_time-start_time, + (total_bytes * 1.0e-6) / (stop_time-start_time))) + + # + # time decryption + # + start_time = time.time() + plaintext_decrypted_blocks = [dec_func(key, c) + for c in ciphertext_blocks] + stop_time = time.time() + print("Decryption Time: %.3fs (%.3f MB/s)" + % (stop_time-start_time, + (total_bytes * 1.0e-6) / (stop_time-start_time))) + +def main(): + runtest("AES - CTR Mode", AES.CTREnc, AES.CTRDec) + runtest("AES - GCM Mode", AES.GCMEnc, AES.GCMDec) + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/PyORAM/examples/ali.py b/PyORAM/examples/ali.py new file mode 100644 index 0000000..f758de3 --- /dev/null +++ b/PyORAM/examples/ali.py @@ -0,0 +1,176 @@ +# +# This example measures the performance of Path ORAM when +# storage is accessed through an SSH client using the Secure +# File Transfer Protocol (SFTP). +# +# In order to run this example, you must provide a host +# (server) address along with valid login credentials +# + +import os +import random +import time +import array + +import pyoram +from pyoram.util.misc import MemorySize +from pyoram.oblivious_storage.tree.path_oram import \ + PathORAM + +from pyoram.storage.AliTimer import * + +import paramiko +import tqdm + +pyoram.config.SHOW_PROGRESS_BAR = True + +# Set SSH login credentials here +# (by default, we pull these from the environment +# for testing purposes) +ssh_host = os.environ.get('PYORAM_SSH_TEST_HOST') +ssh_username = os.environ.get('PYORAM_SSH_TEST_USERNAME') +ssh_password = os.environ.get('PYORAM_SSH_TEST_PASSWORD') + +# Set the storage location and size +storage_name = "heap.bin" +# 4KB block size +block_size = 2048 +# one block per bucket in the +# storage heap of height 8 +# block_count = 2**(8+1)-1 +# block_count = 2**(12+1)-1 +# block_count = 2**(15+1)-1 +block_count = 2**(8+1)-1 + +def main(): + timer = Foo.Instance() + + print("Storage Name: %s" % (storage_name)) + print("Block Count: %s" % (block_count)) + print("Block Size: %s" % (MemorySize(block_size))) + print("Total Memory: %s" + % (MemorySize(block_size*block_count))) + print("Actual Storage Required: %s" + % (MemorySize( + PathORAM.compute_storage_size( + block_size, + block_count, + storage_type='sftp')))) + print("") + + # Start an SSH client using paramiko + print("Starting SSH Client") + with paramiko.SSHClient() as ssh: + ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + ssh.load_system_host_keys() + ssh.connect(ssh_host, + username=ssh_username, + password=ssh_password) + + print("Setting Up Path ORAM Storage") + start_time = time.time() + with PathORAM.setup(storage_name, + block_size, + block_count, + storage_type='sftp', + sshclient=ssh, + cached_levels=2, + concurrency_level = 0, + ignore_existing=True) as f: + print("Total Data Transmission: %s" % (MemorySize(f.bytes_sent + f.bytes_received))) + print("Done Initializing") + pass + + stop_time = time.time() + print("Initial Setup Processing Time: " + str(stop_time - start_time)) + print("Initial Setup Network Time: " + str(timer.getTime())) + print("") + + # print("Total Setup Time: %.2f s" + # % (time.time()-setup_start)) + # print("Current Stash Size: %s" + # % len(f.stash)) + + # print("") + + start_time = time.time() + timer.resetTimer() + # We close the device and reopen it after + # setup to reset the bytes sent and bytes + # received stats. + + print("Starting test Run...") + + + with PathORAM(storage_name, + f.stash, + f.position_map, + key=f.key, + cached_levels=2, + concurrency_level = 0, + storage_type='sftp', + sshclient=ssh) as f: + + stop_time = time.time() + print("Total Data Transmission: %s" % (MemorySize(f.bytes_sent + f.bytes_received))) + print("Test Setup Processing Time: " + str(stop_time - start_time)) + print("Test Setup Network Time: " + str(timer.getTime())) + print("") + + + + + + keys = [] + keys.extend(range(0, block_count)) + random.shuffle(keys) + + + print("Starting Ali Test 2") + timer.resetTimer() + test_count = block_count + start_time = time.time() + + for t in tqdm.tqdm(list(range(test_count)), desc="Running I/O Performance Test"): + ind = keys[t] + # ind = t + s = "a" + str(ind) + f.write_block(ind, bytearray(s.ljust(block_size, '\0'))) + print("Total Data Transmission: %s" % (MemorySize(f.bytes_sent + f.bytes_received))) + + # for t in tqdm.tqdm(list(range(test_count)), desc="Running I/O Performance Test"): + # ind = keys[t] + # # print(f.read_block(ind)) + # f.read_block(ind) + + + + stop_time = time.time() + print("Test Processing Time: " + str(stop_time - start_time)) + print("Test Network Time: " + str(timer.getTime())) + print("") + + + + + + # print("Current Stash Size: %s" + # % len(f.stash)) + # print("Access Block Avg. Data Transmitted: %s (%.3fx)" + # % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), + # (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) + # print("Fetch Block Avg. Latency: %.2f ms" + # % ((stop_time-start_time)/float(test_count)*1000)) + # print("") + # print("") + # print("") + # print("") + # print("") + # print("") + + + + + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/PyORAM/examples/encrypted_heap_test.py b/PyORAM/examples/encrypted_heap_test.py new file mode 100644 index 0000000..3ea5916 --- /dev/null +++ b/PyORAM/examples/encrypted_heap_test.py @@ -0,0 +1,57 @@ +import os +import tempfile + +from pyoram.util.virtual_heap import \ + SizedVirtualHeap +from pyoram.encrypted_storage.encrypted_heap_storage import \ + EncryptedHeapStorage + +def main(): + # + # get a unique filename in the current directory + # + fid, tmpname = tempfile.mkstemp(dir=os.getcwd()) + os.close(fid) + os.remove(tmpname) + print("Storage Name: %s" % (tmpname)) + + key_size = 32 + header_data = b'a message' + heap_base = 3 + heap_height = 2 + block_size = 8 + blocks_per_bucket=4 + initialize = lambda i: \ + bytes(bytearray([i] * block_size * blocks_per_bucket)) + vheap = SizedVirtualHeap( + heap_base, + heap_height, + blocks_per_bucket=blocks_per_bucket) + + with EncryptedHeapStorage.setup( + tmpname, + block_size, + heap_height, + key_size=key_size, + header_data=header_data, + heap_base=heap_base, + blocks_per_bucket=blocks_per_bucket, + initialize=initialize) as f: + assert tmpname == f.storage_name + assert f.header_data == header_data + print(f.read_path(vheap.random_bucket())) + key = f.key + assert os.path.exists(tmpname) + + with EncryptedHeapStorage(tmpname, key=key) as f: + assert tmpname == f.storage_name + assert f.header_data == header_data + print(f.read_path(vheap.random_bucket())) + + # + # cleanup + # + os.remove(tmpname) + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/PyORAM/examples/encrypted_storage_file.py b/PyORAM/examples/encrypted_storage_file.py new file mode 100644 index 0000000..7f6d4c1 --- /dev/null +++ b/PyORAM/examples/encrypted_storage_file.py @@ -0,0 +1,79 @@ +# +# This example measures the performance of encrypted +# storage access through a local file. +# + +import os +import random +import time + +import pyoram +from pyoram.util.misc import MemorySize +from pyoram.encrypted_storage.encrypted_block_storage import \ + EncryptedBlockStorage + +import tqdm + +pyoram.config.SHOW_PROGRESS_BAR = True + +# Set the storage location and size +storage_name = "heap.bin" +# 4KB block size +block_size = 4000 +# one block per bucket in the +# storage heap of height 8 +block_count = 2**(8+1)-1 + +def main(): + + print("Storage Name: %s" % (storage_name)) + print("Block Count: %s" % (block_count)) + print("Block Size: %s" % (MemorySize(block_size))) + print("Total Memory: %s" + % (MemorySize(block_size*block_count))) + print("Actual Storage Required: %s" + % (MemorySize( + EncryptedBlockStorage.compute_storage_size( + block_size, + block_count, + storage_type='file')))) + print("") + + print("Setting Up Encrypted Block Storage") + setup_start = time.time() + with EncryptedBlockStorage.setup(storage_name, + block_size, + block_count, + storage_type='file', + ignore_existing=True) as f: + print("Total Setup Time: %2.f s" + % (time.time()-setup_start)) + print("Total Data Transmission: %s" + % (MemorySize(f.bytes_sent + f.bytes_received))) + print("") + + # We close the device and reopen it after + # setup to reset the bytes sent and bytes + # received stats. + with EncryptedBlockStorage(storage_name, + key=f.key, + storage_type='file') as f: + + test_count = 1000 + start_time = time.time() + for t in tqdm.tqdm(list(range(test_count)), + desc="Running I/O Performance Test"): + f.read_block(random.randint(0,f.block_count-1)) + stop_time = time.time() + print("Access Block Avg. Data Transmitted: %s (%.3fx)" + % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), + (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) + print("Access Block Avg. Latency: %.2f ms" + % ((stop_time-start_time)/float(test_count)*1000)) + print("") + + # cleanup because this is a test example + os.remove(storage_name) + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/PyORAM/examples/encrypted_storage_mmap.py b/PyORAM/examples/encrypted_storage_mmap.py new file mode 100644 index 0000000..911a1b1 --- /dev/null +++ b/PyORAM/examples/encrypted_storage_mmap.py @@ -0,0 +1,79 @@ +# +# This example measures the performance of encrypted +# storage access through a local memory-mapped file. +# + +import os +import random +import time + +import pyoram +from pyoram.util.misc import MemorySize +from pyoram.encrypted_storage.encrypted_block_storage import \ + EncryptedBlockStorage + +import tqdm + +pyoram.config.SHOW_PROGRESS_BAR = True + +# Set the storage location and size +storage_name = "heap.bin" +# 4KB block size +block_size = 4000 +# one block per bucket in the +# storage heap of height 8 +block_count = 2**(8+1)-1 + +def main(): + + print("Storage Name: %s" % (storage_name)) + print("Block Count: %s" % (block_count)) + print("Block Size: %s" % (MemorySize(block_size))) + print("Total Memory: %s" + % (MemorySize(block_size*block_count))) + print("Actual Storage Required: %s" + % (MemorySize( + EncryptedBlockStorage.compute_storage_size( + block_size, + block_count, + storage_type='mmap')))) + print("") + + print("Setting Up Encrypted Block Storage") + setup_start = time.time() + with EncryptedBlockStorage.setup(storage_name, + block_size, + block_count, + storage_type='mmap', + ignore_existing=True) as f: + print("Total Setup Time: %2.f s" + % (time.time()-setup_start)) + print("Total Data Transmission: %s" + % (MemorySize(f.bytes_sent + f.bytes_received))) + print("") + + # We close the device and reopen it after + # setup to reset the bytes sent and bytes + # received stats. + with EncryptedBlockStorage(storage_name, + key=f.key, + storage_type='mmap') as f: + + test_count = 1000 + start_time = time.time() + for t in tqdm.tqdm(list(range(test_count)), + desc="Running I/O Performance Test"): + f.read_block(random.randint(0,f.block_count-1)) + stop_time = time.time() + print("Access Block Avg. Data Transmitted: %s (%.3fx)" + % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), + (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) + print("Access Block Avg. Latency: %.2f ms" + % ((stop_time-start_time)/float(test_count)*1000)) + print("") + + # cleanup because this is a test example + os.remove(storage_name) + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/PyORAM/examples/encrypted_storage_ram.py b/PyORAM/examples/encrypted_storage_ram.py new file mode 100644 index 0000000..09fef94 --- /dev/null +++ b/PyORAM/examples/encrypted_storage_ram.py @@ -0,0 +1,85 @@ +# +# This example measures the performance of encrypted +# storage access through RAM. +# + +import os +import random +import time + +import pyoram +from pyoram.util.misc import MemorySize +from pyoram.encrypted_storage.encrypted_block_storage import \ + EncryptedBlockStorage +from pyoram.storage.block_storage_ram import \ + BlockStorageRAM + +import tqdm + +pyoram.config.SHOW_PROGRESS_BAR = True + +# Set the storage location and size +storage_name = "heap.bin" +# 4KB block size +block_size = 4000 +# one block per bucket in the +# storage heap of height 8 +block_count = 2**(8+1)-1 + +def main(): + + print("Storage Name: %s" % (storage_name)) + print("Block Count: %s" % (block_count)) + print("Block Size: %s" % (MemorySize(block_size))) + print("Total Memory: %s" + % (MemorySize(block_size*block_count))) + print("Actual Storage Required: %s" + % (MemorySize( + EncryptedBlockStorage.compute_storage_size( + block_size, + block_count, + storage_type='ram')))) + print("") + + print("Setting Up Encrypted Block Storage") + setup_start = time.time() + with EncryptedBlockStorage.setup(storage_name, # RAM storage ignores this argument + block_size, + block_count, + storage_type='ram', + ignore_existing=True) as f: + print("Total Setup Time: %2.f s" + % (time.time()-setup_start)) + print("Total Data Transmission: %s" + % (MemorySize(f.bytes_sent + f.bytes_received))) + print("") + + # This must be done after closing the file to ensure the lock flag + # is set to False in the saved data. The tofile method only exists + # on BlockStorageRAM + f.raw_storage.tofile(storage_name) + + # We close the device and reopen it after + # setup to reset the bytes sent and bytes + # received stats. + with EncryptedBlockStorage(BlockStorageRAM.fromfile(storage_name), + key=f.key) as f: + + test_count = 1000 + start_time = time.time() + for t in tqdm.tqdm(list(range(test_count)), + desc="Running I/O Performance Test"): + f.read_block(random.randint(0,f.block_count-1)) + stop_time = time.time() + print("Access Block Avg. Data Transmitted: %s (%.3fx)" + % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), + (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) + print("Access Block Avg. Latency: %.2f ms" + % ((stop_time-start_time)/float(test_count)*1000)) + print("") + + # cleanup because this is a test example + os.remove(storage_name) + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/PyORAM/examples/encrypted_storage_s3.py b/PyORAM/examples/encrypted_storage_s3.py new file mode 100644 index 0000000..c61bd60 --- /dev/null +++ b/PyORAM/examples/encrypted_storage_s3.py @@ -0,0 +1,92 @@ +# +# This example measures the performance of encrypted +# storage access through Amazon Simple Storage Service +# (S3). +# +# In order to run this example, you must provide a valid +# S3 bucket name and have the following variables defined +# in your current environment: +# - AWS_ACCESS_KEY_ID +# - AWS_SECRET_ACCESS_KEY +# - AWS_DEFAULT_REGION +# These can also be set using keywords. +# + +import os +import random +import time + +import pyoram +from pyoram.util.misc import MemorySize +from pyoram.encrypted_storage.encrypted_block_storage import \ + EncryptedBlockStorage + +import tqdm + +pyoram.SHOW_PROGRESS_BAR = True + +# Set S3 bucket name here +# (by default, we pull this from the environment +# for testing purposes) +bucket_name = os.environ.get('PYORAM_AWS_TEST_BUCKET') + +# Set the storage location and size +storage_name = "heap.bin" +# 4KB block size +block_size = 4000 +# one block per bucket in the +# storage heap of height 8 +block_count = 2**(8+1)-1 + +def main(): + + print("Storage Name: %s" % (storage_name)) + print("Block Count: %s" % (block_count)) + print("Block Size: %s" % (MemorySize(block_size))) + print("Total Memory: %s" + % (MemorySize(block_size*block_count))) + print("Actual Storage Required: %s" + % (MemorySize( + EncryptedBlockStorage.compute_storage_size( + block_size, + block_count, + storage_type='s3')))) + print("") + + print("Setting Up Encrypted Block Storage") + setup_start = time.time() + with EncryptedBlockStorage.setup(storage_name, + block_size, + block_count, + storage_type='s3', + bucket_name=bucket_name, + ignore_existing=True) as f: + print("Total Setup Time: %2.f s" + % (time.time()-setup_start)) + print("Total Data Transmission: %s" + % (MemorySize(f.bytes_sent + f.bytes_received))) + print("") + + # We close the device and reopen it after + # setup to reset the bytes sent and bytes + # received stats. + with EncryptedBlockStorage(storage_name, + key=f.key, + storage_type='s3', + bucket_name=bucket_name) as f: + + test_count = 1000 + start_time = time.time() + for t in tqdm.tqdm(list(range(test_count)), + desc="Running I/O Performance Test"): + f.read_block(random.randint(0,f.block_count-1)) + stop_time = time.time() + print("Access Block Avg. Data Transmitted: %s (%.3fx)" + % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), + (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) + print("Access Block Avg. Latency: %.2f ms" + % ((stop_time-start_time)/float(test_count)*1000)) + print("") + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/PyORAM/examples/encrypted_storage_sftp.py b/PyORAM/examples/encrypted_storage_sftp.py new file mode 100644 index 0000000..cb107a5 --- /dev/null +++ b/PyORAM/examples/encrypted_storage_sftp.py @@ -0,0 +1,99 @@ +# +# This example measures the performance of encrypted storage +# access through an SSH client using the Secure File +# Transfer Protocol (SFTP). +# +# In order to run this example, you must provide a host +# (server) address along with valid login credentials +# + +import os +import random +import time + +import pyoram +from pyoram.util.misc import MemorySize +from pyoram.encrypted_storage.encrypted_block_storage import \ + EncryptedBlockStorage + +import paramiko +import tqdm + +pyoram.config.SHOW_PROGRESS_BAR = True + +# Set SSH login credentials here +# (by default, we pull these from the environment +# for testing purposes) +ssh_host = os.environ.get('PYORAM_SSH_TEST_HOST') +ssh_username = os.environ.get('PYORAM_SSH_TEST_USERNAME') +ssh_password = os.environ.get('PYORAM_SSH_TEST_PASSWORD') + +# Set the storage location and size +storage_name = "heap.bin" +# 4KB block size +block_size = 4000 +# one block per bucket in the +# storage heap of height 8 +block_count = 2**(8+1)-1 + +def main(): + + print("Storage Name: %s" % (storage_name)) + print("Block Count: %s" % (block_count)) + print("Block Size: %s" % (MemorySize(block_size))) + print("Total Memory: %s" + % (MemorySize(block_size*block_count))) + print("Actual Storage Required: %s" + % (MemorySize( + EncryptedBlockStorage.compute_storage_size( + block_size, + block_count, + storage_type='sftp')))) + print("") + + # Start an SSH client using paramiko + print("Starting SSH Client") + with paramiko.SSHClient() as ssh: + ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + ssh.load_system_host_keys() + ssh.connect(ssh_host, + username=ssh_username, + password=ssh_password) + + print("Setting Up Encrypted Block Storage") + setup_start = time.time() + with EncryptedBlockStorage.setup(storage_name, + block_size, + block_count, + storage_type='sftp', + sshclient=ssh, + ignore_existing=True) as f: + print("Total Setup Time: %2.f s" + % (time.time()-setup_start)) + print("Total Data Transmission: %s" + % (MemorySize(f.bytes_sent + f.bytes_received))) + print("") + + # We close the device and reopen it after + # setup to reset the bytes sent and bytes + # received stats. + with EncryptedBlockStorage(storage_name, + key=f.key, + storage_type='sftp', + sshclient=ssh) as f: + + test_count = 1000 + start_time = time.time() + for t in tqdm.tqdm(list(range(test_count)), + desc="Running I/O Performance Test"): + f.read_block(random.randint(0,f.block_count-1)) + stop_time = time.time() + print("Access Block Avg. Data Transmitted: %s (%.3fx)" + % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), + (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) + print("Access Block Avg. Latency: %.2f ms" + % ((stop_time-start_time)/float(test_count)*1000)) + print("") + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/PyORAM/examples/path_oram_file.py b/PyORAM/examples/path_oram_file.py new file mode 100644 index 0000000..a59fe6a --- /dev/null +++ b/PyORAM/examples/path_oram_file.py @@ -0,0 +1,85 @@ +# +# This example measures the performance of Path ORAM +# when storage is accessed through a local file. +# + +import os +import random +import time + +import pyoram +from pyoram.util.misc import MemorySize +from pyoram.oblivious_storage.tree.path_oram import \ + PathORAM + +import tqdm + +pyoram.config.SHOW_PROGRESS_BAR = True + +# Set the storage location and size +storage_name = "heap.bin" +# 4KB block size +block_size = 4000 +# one block per bucket in the +# storage heap of height 8 +block_count = 2**(8+1)-1 + +def main(): + + print("Storage Name: %s" % (storage_name)) + print("Block Count: %s" % (block_count)) + print("Block Size: %s" % (MemorySize(block_size))) + print("Total Memory: %s" + % (MemorySize(block_size*block_count))) + print("Actual Storage Required: %s" + % (MemorySize( + PathORAM.compute_storage_size( + block_size, + block_count, + storage_type='file')))) + print("") + + print("Setting Up Path ORAM Storage") + setup_start = time.time() + with PathORAM.setup(storage_name, + block_size, + block_count, + storage_type='file', + ignore_existing=True) as f: + print("Total Setup Time: %2.f s" + % (time.time()-setup_start)) + print("Current Stash Size: %s" + % len(f.stash)) + print("Total Data Transmission: %s" + % (MemorySize(f.bytes_sent + f.bytes_received))) + print("") + + # We close the device and reopen it after + # setup to reset the bytes sent and bytes + # received stats. + with PathORAM(storage_name, + f.stash, + f.position_map, + key=f.key, + storage_type='file') as f: + + test_count = 100 + start_time = time.time() + for t in tqdm.tqdm(list(range(test_count)), + desc="Running I/O Performance Test"): + f.read_block(random.randint(0,f.block_count-1)) + stop_time = time.time() + print("Current Stash Size: %s" + % len(f.stash)) + print("Access Block Avg. Data Transmitted: %s (%.3fx)" + % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), + (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) + print("Access Block Avg. Latency: %.2f ms" + % ((stop_time-start_time)/float(test_count)*1000)) + print("") + + # cleanup because this is a test example + os.remove(storage_name) + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/PyORAM/examples/path_oram_mmap.py b/PyORAM/examples/path_oram_mmap.py new file mode 100644 index 0000000..a94efa4 --- /dev/null +++ b/PyORAM/examples/path_oram_mmap.py @@ -0,0 +1,86 @@ +# +# This example measures the performance of Path ORAM +# when storage is accessed through a local memory-mapped +# file (mmap). +# + +import os +import random +import time + +import pyoram +from pyoram.util.misc import MemorySize +from pyoram.oblivious_storage.tree.path_oram import \ + PathORAM + +import tqdm + +pyoram.config.SHOW_PROGRESS_BAR = True + +# Set the storage location and size +storage_name = "heap.bin" +# 4KB block size +block_size = 4000 +# one block per bucket in the +# storage heap of height 8 +block_count = 2**(8+1)-1 + +def main(): + + print("Storage Name: %s" % (storage_name)) + print("Block Count: %s" % (block_count)) + print("Block Size: %s" % (MemorySize(block_size))) + print("Total Memory: %s" + % (MemorySize(block_size*block_count))) + print("Actual Storage Required: %s" + % (MemorySize( + PathORAM.compute_storage_size( + block_size, + block_count, + storage_type='mmap')))) + print("") + + print("Setting Up Path ORAM Storage") + setup_start = time.time() + with PathORAM.setup(storage_name, + block_size, + block_count, + storage_type='mmap', + ignore_existing=True) as f: + print("Total Setup Time: %2.f s" + % (time.time()-setup_start)) + print("Current Stash Size: %s" + % len(f.stash)) + print("Total Data Transmission: %s" + % (MemorySize(f.bytes_sent + f.bytes_received))) + print("") + + # We close the device and reopen it after + # setup to reset the bytes sent and bytes + # received stats. + with PathORAM(storage_name, + f.stash, + f.position_map, + key=f.key, + storage_type='mmap') as f: + + test_count = 100 + start_time = time.time() + for t in tqdm.tqdm(list(range(test_count)), + desc="Running I/O Performance Test"): + f.read_block(random.randint(0,f.block_count-1)) + stop_time = time.time() + print("Current Stash Size: %s" + % len(f.stash)) + print("Access Block Avg. Data Transmitted: %s (%.3fx)" + % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), + (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) + print("Access Block Avg. Latency: %.2f ms" + % ((stop_time-start_time)/float(test_count)*1000)) + print("") + + # cleanup because this is a test example + os.remove(storage_name) + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/PyORAM/examples/path_oram_ram.py b/PyORAM/examples/path_oram_ram.py new file mode 100644 index 0000000..d7b99b0 --- /dev/null +++ b/PyORAM/examples/path_oram_ram.py @@ -0,0 +1,91 @@ +# +# This example measures the performance of Path ORAM +# when storage is accessed through RAM. +# + +import os +import random +import time + +import pyoram +from pyoram.util.misc import MemorySize +from pyoram.oblivious_storage.tree.path_oram import \ + PathORAM +from pyoram.storage.block_storage_ram import \ + BlockStorageRAM + +import tqdm + +pyoram.config.SHOW_PROGRESS_BAR = True + +# Set the storage location and size +storage_name = "heap.bin" +# 4KB block size +block_size = 4000 +# one block per bucket in the +# storage heap of height 8 +block_count = 2**(8+1)-1 + +def main(): + + print("Storage Name: %s" % (storage_name)) + print("Block Count: %s" % (block_count)) + print("Block Size: %s" % (MemorySize(block_size))) + print("Total Memory: %s" + % (MemorySize(block_size*block_count))) + print("Actual Storage Required: %s" + % (MemorySize( + PathORAM.compute_storage_size( + block_size, + block_count, + storage_type='ram')))) + print("") + + print("Setting Up Path ORAM Storage") + setup_start = time.time() + with PathORAM.setup(storage_name, # RAM storage ignores this argument + block_size, + block_count, + storage_type='ram', + ignore_existing=True) as f: + print("Total Setup Time: %2.f s" + % (time.time()-setup_start)) + print("Current Stash Size: %s" + % len(f.stash)) + print("Total Data Transmission: %s" + % (MemorySize(f.bytes_sent + f.bytes_received))) + print("") + + # This must be done after closing the file to ensure the lock flag + # is set to False in the saved data. The tofile method only exists + # on BlockStorageRAM + f.raw_storage.tofile(storage_name) + + # We close the device and reopen it after + # setup to reset the bytes sent and bytes + # received stats. + with PathORAM(BlockStorageRAM.fromfile(storage_name), + f.stash, + f.position_map, + key=f.key) as f: + + test_count = 100 + start_time = time.time() + for t in tqdm.tqdm(list(range(test_count)), + desc="Running I/O Performance Test"): + f.read_block(random.randint(0,f.block_count-1)) + stop_time = time.time() + print("Current Stash Size: %s" + % len(f.stash)) + print("Access Block Avg. Data Transmitted: %s (%.3fx)" + % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), + (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) + print("Access Block Avg. Latency: %.2f ms" + % ((stop_time-start_time)/float(test_count)*1000)) + print("") + + # cleanup because this is a test example + os.remove(storage_name) + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/PyORAM/examples/path_oram_s3.py b/PyORAM/examples/path_oram_s3.py new file mode 100644 index 0000000..0f1fb36 --- /dev/null +++ b/PyORAM/examples/path_oram_s3.py @@ -0,0 +1,98 @@ +# +# This example measures the performance of Path ORAM when +# storage is accessed through Amazon Simple Storage Service +# (S3). +# +# In order to run this example, you must provide a valid +# S3 bucket name and have the following variables defined +# in your current environment: +# - AWS_ACCESS_KEY_ID +# - AWS_SECRET_ACCESS_KEY +# - AWS_DEFAULT_REGION +# These can also be set using keywords. +# + +import os +import random +import time + +import pyoram +from pyoram.util.misc import MemorySize +from pyoram.oblivious_storage.tree.path_oram import \ + PathORAM + +import tqdm + +pyoram.config.SHOW_PROGRESS_BAR = True + +# Set S3 bucket name here +# (by default, we pull this from the environment +# for testing purposes) +bucket_name = os.environ.get('PYORAM_AWS_TEST_BUCKET') + +# Set the storage location and size +storage_name = "heap.bin" +# 4KB block size +block_size = 4000 +# one block per bucket in the +# storage heap of height 8 +block_count = 2**(8+1)-1 + +def main(): + + print("Storage Name: %s" % (storage_name)) + print("Block Count: %s" % (block_count)) + print("Block Size: %s" % (MemorySize(block_size))) + print("Total Memory: %s" + % (MemorySize(block_size*block_count))) + print("Actual Storage Required: %s" + % (MemorySize( + PathORAM.compute_storage_size( + block_size, + block_count, + storage_type='s3')))) + print("") + + print("Setting Up Path ORAM Storage") + setup_start = time.time() + with PathORAM.setup(storage_name, + block_size, + block_count, + storage_type='s3', + bucket_name=bucket_name, + ignore_existing=True) as f: + print("Total Setup Time: %.2f s" + % (time.time()-setup_start)) + print("Current Stash Size: %s" + % len(f.stash)) + print("Total Data Transmission: %s" + % (MemorySize(f.bytes_sent + f.bytes_received))) + print("") + + # We close the device and reopen it after + # setup to reset the bytes sent and bytes + # received stats. + with PathORAM(storage_name, + f.stash, + f.position_map, + key=f.key, + storage_type='s3', + bucket_name=bucket_name) as f: + + test_count = 100 + start_time = time.time() + for t in tqdm.tqdm(list(range(test_count)), + desc="Running I/O Performance Test"): + f.read_block(random.randint(0,f.block_count-1)) + stop_time = time.time() + print("Current Stash Size: %s" + % len(f.stash)) + print("Access Block Avg. Data Transmitted: %s (%.3fx)" + % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), + (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) + print("Fetch Block Avg. Latency: %.2f ms" + % ((stop_time-start_time)/float(test_count)*1000)) + print("") + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/PyORAM/examples/path_oram_sftp.py b/PyORAM/examples/path_oram_sftp.py new file mode 100644 index 0000000..64c5ef2 --- /dev/null +++ b/PyORAM/examples/path_oram_sftp.py @@ -0,0 +1,105 @@ +# +# This example measures the performance of Path ORAM when +# storage is accessed through an SSH client using the Secure +# File Transfer Protocol (SFTP). +# +# In order to run this example, you must provide a host +# (server) address along with valid login credentials +# + +import os +import random +import time + +import pyoram +from pyoram.util.misc import MemorySize +from pyoram.oblivious_storage.tree.path_oram import \ + PathORAM + +import paramiko +import tqdm + +pyoram.config.SHOW_PROGRESS_BAR = True + +# Set SSH login credentials here +# (by default, we pull these from the environment +# for testing purposes) +ssh_host = os.environ.get('PYORAM_SSH_TEST_HOST') +ssh_username = os.environ.get('PYORAM_SSH_TEST_USERNAME') +ssh_password = os.environ.get('PYORAM_SSH_TEST_PASSWORD') + +# Set the storage location and size +storage_name = "heap.bin" +# 4KB block size +block_size = 4000 +# one block per bucket in the +# storage heap of height 8 +block_count = 2**(8+1)-1 + +def main(): + + print("Storage Name: %s" % (storage_name)) + print("Block Count: %s" % (block_count)) + print("Block Size: %s" % (MemorySize(block_size))) + print("Total Memory: %s" + % (MemorySize(block_size*block_count))) + print("Actual Storage Required: %s" + % (MemorySize( + PathORAM.compute_storage_size( + block_size, + block_count, + storage_type='sftp')))) + print("") + + # Start an SSH client using paramiko + print("Starting SSH Client") + with paramiko.SSHClient() as ssh: + ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + ssh.load_system_host_keys() + ssh.connect(ssh_host, + username=ssh_username, + password=ssh_password) + + print("Setting Up Path ORAM Storage") + setup_start = time.time() + with PathORAM.setup(storage_name, + block_size, + block_count, + storage_type='sftp', + sshclient=ssh, + ignore_existing=True) as f: + print("Total Setup Time: %.2f s" + % (time.time()-setup_start)) + print("Current Stash Size: %s" + % len(f.stash)) + print("Total Data Transmission: %s" + % (MemorySize(f.bytes_sent + f.bytes_received))) + print("") + + # We close the device and reopen it after + # setup to reset the bytes sent and bytes + # received stats. + with PathORAM(storage_name, + f.stash, + f.position_map, + key=f.key, + storage_type='sftp', + sshclient=ssh) as f: + + test_count = 100 + start_time = time.time() + for t in tqdm.tqdm(list(range(test_count)), + desc="Running I/O Performance Test"): + f.read_block(random.randint(0,f.block_count-1)) + stop_time = time.time() + print("Current Stash Size: %s" + % len(f.stash)) + print("Access Block Avg. Data Transmitted: %s (%.3fx)" + % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), + (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) + print("Fetch Block Avg. Latency: %.2f ms" + % ((stop_time-start_time)/float(test_count)*1000)) + print("") + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/PyORAM/examples/path_oram_sftp_setup.py b/PyORAM/examples/path_oram_sftp_setup.py new file mode 100644 index 0000000..0898224 --- /dev/null +++ b/PyORAM/examples/path_oram_sftp_setup.py @@ -0,0 +1,115 @@ +# +# This example demonstrates how to setup an instance of Path ORAM +# locally and then transfer the storage to a server using a paramiko +# SSHClient. After executing this file, path_oram_sftp_test.py can be +# executed to run simple I/O performance tests using different caching +# settings. +# +# In order to run this example, you must provide a host +# (server) address along with valid login credentials +# + +import os +import random +import time +import pickle + +import pyoram +from pyoram.util.misc import MemorySize, save_private_key +from pyoram.oblivious_storage.tree.path_oram import \ + PathORAM + +import paramiko +import tqdm + +pyoram.config.SHOW_PROGRESS_BAR = True + +# Set SSH login credentials here +# (by default, we pull these from the environment +# for testing purposes) +ssh_host = os.environ.get('PYORAM_SSH_TEST_HOST') +ssh_username = os.environ.get('PYORAM_SSH_TEST_USERNAME') +ssh_password = os.environ.get('PYORAM_SSH_TEST_PASSWORD') + +# Set the storage location and size +storage_name = "heap.bin" +# 4KB block size +block_size = 4000 +# one block per bucket in the +# storage heap of height 8 +block_count = 2**(8+1)-1 + +def main(): + + print("Storage Name: %s" % (storage_name)) + print("Block Count: %s" % (block_count)) + print("Block Size: %s" % (MemorySize(block_size))) + print("Total Memory: %s" + % (MemorySize(block_size*block_count))) + print("Actual Storage Required: %s" + % (MemorySize( + PathORAM.compute_storage_size( + block_size, + block_count, + storage_type='mmap')))) + print("") + + print("Setting Up Path ORAM Storage Locally") + setup_start = time.time() + with PathORAM.setup(storage_name, + block_size, + block_count, + storage_type='mmap', + ignore_existing=True) as f: + print("Total Setup Time: %.2f s" + % (time.time()-setup_start)) + print("Current Stash Size: %s" + % len(f.stash)) + print("Total Data Transmission: %s" + % (MemorySize(f.bytes_sent + f.bytes_received))) + print("") + + print("Saving key to file: %s.key" + % (storage_name)) + save_private_key(storage_name+".key", f.key) + print("Saving stash to file: %s.stash" + % (storage_name)) + with open(storage_name+".stash", 'wb') as fstash: + pickle.dump(f.stash, fstash) + print("Saving position map to file: %s.position" + % (storage_name)) + with open(storage_name+".position", 'wb') as fpos: + pickle.dump(f.position_map, fpos) + + # Start an SSH client using paramiko + print("Starting SSH Client") + with paramiko.SSHClient() as ssh: + ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + ssh.load_system_host_keys() + ssh.connect(ssh_host, + username=ssh_username, + password=ssh_password) + + sftp = ssh.open_sftp() + + def my_hook(t): + def inner(b, total): + t.total = total + t.update(b - inner.last_b) + inner.last_b = b + inner.last_b = 0 + return inner + with tqdm.tqdm(desc="Transferring Storage", + unit='B', + unit_scale=True, + miniters=1) as t: + sftp.put(storage_name, + storage_name, + callback=my_hook(t)) + sftp.close() + + print("Deleting Local Copy of Storage") + os.remove(storage_name) + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/PyORAM/examples/path_oram_sftp_test.py b/PyORAM/examples/path_oram_sftp_test.py new file mode 100644 index 0000000..6e0e49f --- /dev/null +++ b/PyORAM/examples/path_oram_sftp_test.py @@ -0,0 +1,99 @@ +# +# This example demonstrates how to access an existing Path ORAM +# storage space through an SSH client using the Secure File Transfer +# Protocol (SFTP). This file should not be executed until the +# path_oram_sftp_setup.py example has been executed. The user is +# encouraged to tweak the settings for 'cached_levels', +# 'concurrency_level', and 'threadpool_size' to observe their effect +# on access latency. +# +# In order to run this example, you must provide a host +# (server) address along with valid login credentials +# + +import os +import random +import time +import pickle +import multiprocessing + +import pyoram +from pyoram.util.misc import MemorySize, load_private_key +from pyoram.oblivious_storage.tree.path_oram import \ + PathORAM + +import paramiko +import tqdm + +pyoram.config.SHOW_PROGRESS_BAR = True + +# Set SSH login credentials here +# (by default, we pull these from the environment +# for testing purposes) +ssh_host = os.environ.get('PYORAM_SSH_TEST_HOST') +ssh_username = os.environ.get('PYORAM_SSH_TEST_USERNAME') +ssh_password = os.environ.get('PYORAM_SSH_TEST_PASSWORD') + +# Set the storage location and size +storage_name = "heap.bin" + +def main(): + + print("Loading key from file: %s.key" + % (storage_name)) + key = load_private_key(storage_name+".key") + print("Loading stash from file: %s.stash" + % (storage_name)) + with open(storage_name+".stash", 'rb') as fstash: + stash = pickle.load(fstash) + print("Loading position map from file: %s.position" + % (storage_name)) + with open(storage_name+".position", 'rb') as fpos: + position_map = pickle.load(fpos) + + # Start an SSH client using paramiko + print("Starting SSH Client") + with paramiko.SSHClient() as ssh: + ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + ssh.load_system_host_keys() + ssh.connect(ssh_host, + username=ssh_username, + password=ssh_password) + + with PathORAM(storage_name, + stash, + position_map, + key=key, + storage_type='sftp', + cached_levels=6, + concurrency_level=3, + threadpool_size=multiprocessing.cpu_count()*2, + sshclient=ssh) as f: + + try: + + test_count = 100 + start_time = time.time() + for t in tqdm.tqdm(list(range(test_count)), + desc="Running I/O Performance Test"): + f.read_block(random.randint(0,f.block_count-1)) + stop_time = time.time() + print("Current Stash Size: %s" + % len(f.stash)) + print("Fetch Block Avg. Latency: %.2f ms" + % ((stop_time-start_time)/float(test_count)*1000)) + print("") + + finally: + + print("Saving stash to file: %s.stash" + % (storage_name)) + with open(storage_name+".stash", 'wb') as fstash: + pickle.dump(f.stash, fstash) + print("Saving position map to file: %s.position" + % (storage_name)) + with open(storage_name+".position", 'wb') as fpos: + pickle.dump(f.position_map, fpos) + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/PyORAM/examples/tree_oram_test.py b/PyORAM/examples/tree_oram_test.py new file mode 100644 index 0000000..e9e2075 --- /dev/null +++ b/PyORAM/examples/tree_oram_test.py @@ -0,0 +1,93 @@ +import os +import struct +import random + +from pyoram.util.virtual_heap import \ + SizedVirtualHeap +from pyoram.encrypted_storage.encrypted_heap_storage import \ + EncryptedHeapStorage +from pyoram.oblivious_storage.tree.tree_oram_helper import \ + TreeORAMStorageManagerPointerAddressing + +def main(): + storage_name = "heap.bin" + print("Storage Name: %s" % (storage_name)) + + key_size = 32 + heap_base = 2 + heap_height = 2 + block_size = struct.calcsize("!?LL") + blocks_per_bucket = 2 + vheap = SizedVirtualHeap( + heap_base, + heap_height, + blocks_per_bucket=blocks_per_bucket) + + print("Block Size: %s" % (block_size)) + print("Blocks Per Bucket: %s" % (blocks_per_bucket)) + + position_map = {} + def initialize(i): + bucket = bytes() + for j in range(blocks_per_bucket): + if (i*j) % 3: + bucket += struct.pack( + "!?LL", False, 0, 0) + else: + x = vheap.Node(i) + while not vheap.is_nil_node(x): + x = x.child_node(random.randint(0, heap_base-1)) + x = x.parent_node() + bucket += struct.pack( + "!?LL", True, initialize.id_, x.bucket) + position_map[initialize.id_] = x.bucket + initialize.id_ += 1 + return bucket + initialize.id_ = 1 + + with EncryptedHeapStorage.setup( + storage_name, + block_size, + heap_height, + heap_base=heap_base, + key_size=key_size, + blocks_per_bucket=blocks_per_bucket, + initialize=initialize, + ignore_existing=True) as f: + assert storage_name == f.storage_name + stash = {} + oram = TreeORAMStorageManagerPointerAddressing(f, stash) + + b = vheap.random_bucket() + oram.load_path(b) + print("") + print(repr(vheap.Node(oram.path_stop_bucket))) + print(oram.path_block_ids) + print(oram.path_block_eviction_levels) + + oram.push_down_path() + print("") + print(repr(vheap.Node(oram.path_stop_bucket))) + print(oram.path_block_ids) + print(oram.path_block_eviction_levels) + print(oram.path_block_reordering) + + oram.evict_path() + oram.load_path(b) + print("") + print(repr(vheap.Node(oram.path_stop_bucket))) + print(oram.path_block_ids) + print(oram.path_block_eviction_levels) + + oram.push_down_path() + print("") + print(repr(vheap.Node(oram.path_stop_bucket))) + print(oram.path_block_ids) + print(oram.path_block_eviction_levels) + print(oram.path_block_reordering) + assert all(x is None for x in oram.path_block_reordering) + + os.remove(storage_name) + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/PyORAM/requirements.txt b/PyORAM/requirements.txt new file mode 100644 index 0000000..62c4b84 --- /dev/null +++ b/PyORAM/requirements.txt @@ -0,0 +1,6 @@ +cryptography +paramiko +boto3 +cffi>=1.0.0 +six +tqdm diff --git a/PyORAM/setup.cfg b/PyORAM/setup.cfg new file mode 100644 index 0000000..146c52d --- /dev/null +++ b/PyORAM/setup.cfg @@ -0,0 +1,6 @@ +[metadata] +description-file = README.rst + +[bdist_wheel] +# supports python3 +universal=1 diff --git a/PyORAM/setup.py b/PyORAM/setup.py new file mode 100644 index 0000000..fe1a966 --- /dev/null +++ b/PyORAM/setup.py @@ -0,0 +1,72 @@ +import os +import sys +import platform +from setuptools import setup, find_packages +from codecs import open + +here = os.path.abspath(os.path.dirname(__file__)) + +about = {} +with open(os.path.join("src", "pyoram", "__about__.py")) as f: + exec(f.read(), about) + +# Get the long description from the README file +def _readme(): + with open(os.path.join(here, 'README.rst'), encoding='utf-8') as f: + return f.read() + +setup_requirements = [] +requirements = ['cryptography', + 'paramiko', + 'boto3', + 'six', + 'tqdm'] + +if platform.python_implementation() == "PyPy": + if sys.pypy_version_info < (2, 6): + raise RuntimeError( + "PyORAM is not compatible with PyPy < 2.6. Please " + "upgrade PyPy to use this library.") +else: + if sys.version_info <= (2, 6): + raise RuntimeError( + "PyORAM is not compatible with Python < 2.7. Please " + "upgrade Python to use this library.") + requirements.append("cffi>=1.0.0") + setup_requirements.append("cffi>=1.0.0") + +setup( + name=about['__title__'], + version=about['__version__'], + description=about['__summary__'], + long_description=_readme(), + url=about['__uri__'], + author=about['__author__'], + author_email=about['__email__'], + license=about['__license__'], + # https://pypi.python.org/pypi?%3Aaction=list_classifiers + classifiers=[ + 'Development Status :: 4 - Beta', + 'Intended Audience :: Science/Research', + 'Topic :: Security :: Cryptography', + "Natural Language :: English", + 'License :: OSI Approved :: MIT License', + 'Programming Language :: Python :: 2.7', + 'Programming Language :: Python :: 3', + 'Programming Language :: Python :: 3.3', + 'Programming Language :: Python :: 3.4', + 'Programming Language :: Python :: 3.5', + 'Programming Language :: Python :: Implementation :: CPython', + 'Programming Language :: Python :: Implementation :: PyPy', + ], + keywords='oram, storage, privacy, cryptography, cloud storage', + package_dir={'': 'src'}, + packages=find_packages(where="src", exclude=["_cffi_src", "_cffi_src.*"]), + setup_requires=setup_requirements, + install_requires=requirements, + cffi_modules=["src/_cffi_src/virtual_heap_helper_build.py:ffi"], + # use MANIFEST.in + include_package_data=True, + test_suite='nose2.collector.collector', + tests_require=['unittest2','nose2'] +) diff --git a/PyORAM/src/_cffi_src/__init__.py b/PyORAM/src/_cffi_src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/PyORAM/src/_cffi_src/virtual_heap_helper_build.py b/PyORAM/src/_cffi_src/virtual_heap_helper_build.py new file mode 100644 index 0000000..72bfc1d --- /dev/null +++ b/PyORAM/src/_cffi_src/virtual_heap_helper_build.py @@ -0,0 +1,74 @@ +import cffi + +# +# C functions that speed up commonly +# executed heap calculations in tree-based +# orams +# + +ffi = cffi.FFI() +ffi.cdef( +""" +int calculate_bucket_level(unsigned int k, + unsigned long long b); +int calculate_last_common_level(unsigned int k, + unsigned long long b1, + unsigned long long b2); +""") + +ffi.set_source("pyoram.util._virtual_heap_helper", +""" +#include +#include + +int calculate_bucket_level(unsigned int k, + unsigned long long b) +{ + unsigned int h; + unsigned long long pow; + if (k == 2) { + // This is simply log2floor(b+1) + h = 0; + b += 1; + while (b >>= 1) {++h;} + return h; + } + b = (k - 1) * (b + 1) + 1; + h = 0; + pow = k; + while (pow < b) {++h; pow *= k;} + return h; +} + +int calculate_last_common_level(unsigned int k, + unsigned long long b1, + unsigned long long b2) +{ + int level1, level2; + level1 = calculate_bucket_level(k, b1); + level2 = calculate_bucket_level(k, b2); + if (level1 != level2) { + if (level1 > level2) { + while (level1 != level2) { + b1 = (b1 - 1)/k; + --level1; + } + } + else { + while (level2 != level1) { + b2 = (b2 - 1)/k; + --level2; + } + } + } + while (b1 != b2) { + b1 = (b1 - 1)/k; + b2 = (b2 - 1)/k; + --level1; + } + return level1; +} +""") + +if __name__ == "__main__": + ffi.compile() diff --git a/PyORAM/src/pyoram/__about__.py b/PyORAM/src/pyoram/__about__.py new file mode 100644 index 0000000..987d948 --- /dev/null +++ b/PyORAM/src/pyoram/__about__.py @@ -0,0 +1,17 @@ +__all__ = ('__title__', + '__summary__', + '__uri__', + '__version__', + '__author__', + '__email__', + '__license__', + '__copyright__') + +__title__ = 'PyORAM' +__summary__ = 'Python-based Oblivious RAM' +__uri__ = 'https://github.com/ghackebeil/PyORAM' +__version__ = '0.3.0' +__author__ = 'Gabriel A. Hackebeil' +__email__ = 'gabe.hackebeil@gmail.com' +__license__ = 'MIT' +__copyright__ = 'Copyright {0}'.format(__author__) diff --git a/PyORAM/src/pyoram/__init__.py b/PyORAM/src/pyoram/__init__.py new file mode 100644 index 0000000..2e0d133 --- /dev/null +++ b/PyORAM/src/pyoram/__init__.py @@ -0,0 +1,55 @@ +from pyoram.__about__ import __version__ + +def _configure_logging(): + import os + import logging + + log = logging.getLogger("pyoram") + formatter = logging.Formatter( + fmt=("[%(asctime)s.%(msecs)03d," + "%(name)s,%(levelname)s] %(threadName)s %(message)s"), + datefmt="%Y-%m-%d %H:%M:%S") + + level = os.environ.get("PYORAM_LOGLEVEL", "WARNING") + logfilename = os.environ.get("PYORAM_LOGFILE", None) + if len(logging.root.handlers) == 0: + # configure the logging with some sensible + # defaults. + try: + import tempfile + tempfile = tempfile.TemporaryFile(dir=".") + tempfile.close() + except OSError: + # cannot write in current directory, use the + # console logger + handler = logging.StreamHandler() + else: + if logfilename is None: + handler = logging.StreamHandler() + else: + # set up a basic logfile in current directory + handler = logging.FileHandler(logfilename) + handler.setFormatter(formatter) + handler.setLevel(level) + log.addHandler(handler) + log.setLevel(level) + log.info("PyORAM log configured using built-in " + "defaults, level=%s", level) + +_configure_logging() +del _configure_logging + +def _configure_pyoram(): + class _Configure(object): + __slots__ = ("SHOW_PROGRESS_BAR",) + def __init__(self): + self.SHOW_PROGRESS_BAR = False + return _Configure() +config = _configure_pyoram() +del _configure_pyoram + +import pyoram.util +import pyoram.crypto +import pyoram.storage +import pyoram.encrypted_storage +import pyoram.oblivious_storage diff --git a/PyORAM/src/pyoram/crypto/__init__.py b/PyORAM/src/pyoram/crypto/__init__.py new file mode 100644 index 0000000..379a24c --- /dev/null +++ b/PyORAM/src/pyoram/crypto/__init__.py @@ -0,0 +1 @@ +import pyoram.crypto.aes diff --git a/PyORAM/src/pyoram/crypto/aes.py b/PyORAM/src/pyoram/crypto/aes.py new file mode 100644 index 0000000..6a9d044 --- /dev/null +++ b/PyORAM/src/pyoram/crypto/aes.py @@ -0,0 +1,48 @@ +__all__ = ("AES",) + +import os +import cryptography.hazmat.primitives.ciphers +import cryptography.hazmat.backends + +_backend = cryptography.hazmat.backends.default_backend() +_aes = cryptography.hazmat.primitives.ciphers.algorithms.AES +_cipher = cryptography.hazmat.primitives.ciphers.Cipher +_ctrmode = cryptography.hazmat.primitives.ciphers.modes.CTR +_gcmmode = cryptography.hazmat.primitives.ciphers.modes.GCM + +class AES(object): + + key_sizes = [k//8 for k in sorted(_aes.key_sizes)] + block_size = _aes.block_size//8 + + @staticmethod + def KeyGen(size_bytes): + assert size_bytes in AES.key_sizes + return os.urandom(size_bytes) + + @staticmethod + def CTREnc(key, plaintext): + iv = os.urandom(AES.block_size) + cipher = _cipher(_aes(key), _ctrmode(iv), backend=_backend).encryptor() + return iv + cipher.update(plaintext) + cipher.finalize() + + @staticmethod + def CTRDec(key, ciphertext): + iv = ciphertext[:AES.block_size] + cipher = _cipher(_aes(key), _ctrmode(iv), backend=_backend).decryptor() + return cipher.update(ciphertext[AES.block_size:]) + \ + cipher.finalize() + + @staticmethod + def GCMEnc(key, plaintext): + iv = os.urandom(AES.block_size) + cipher = _cipher(_aes(key), _gcmmode(iv), backend=_backend).encryptor() + return iv + cipher.update(plaintext) + cipher.finalize() + cipher.tag + + @staticmethod + def GCMDec(key, ciphertext): + iv = ciphertext[:AES.block_size] + tag = ciphertext[-AES.block_size:] + cipher = _cipher(_aes(key), _gcmmode(iv, tag), backend=_backend).decryptor() + return cipher.update(ciphertext[AES.block_size:-AES.block_size]) + \ + cipher.finalize() diff --git a/PyORAM/src/pyoram/encrypted_storage/__init__.py b/PyORAM/src/pyoram/encrypted_storage/__init__.py new file mode 100644 index 0000000..8226f44 --- /dev/null +++ b/PyORAM/src/pyoram/encrypted_storage/__init__.py @@ -0,0 +1,3 @@ +import pyoram.encrypted_storage.encrypted_block_storage +import pyoram.encrypted_storage.encrypted_heap_storage +import pyoram.encrypted_storage.top_cached_encrypted_heap_storage diff --git a/PyORAM/src/pyoram/encrypted_storage/encrypted_block_storage.py b/PyORAM/src/pyoram/encrypted_storage/encrypted_block_storage.py new file mode 100644 index 0000000..f3274b1 --- /dev/null +++ b/PyORAM/src/pyoram/encrypted_storage/encrypted_block_storage.py @@ -0,0 +1,326 @@ +__all__ = ('EncryptedBlockStorage',) + +import struct +import hmac +import hashlib + +from pyoram.storage.block_storage import (BlockStorageInterface, + BlockStorageTypeFactory) +from pyoram.crypto.aes import AES + +import six + +class EncryptedBlockStorageInterface(BlockStorageInterface): + + # + # Abstract Interface + # + + @property + def key(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + @property + def raw_storage(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + +class EncryptedBlockStorage(EncryptedBlockStorageInterface): + + _index_struct_string = "!"+("x"*hashlib.sha384().digest_size)+"?" + _index_offset = struct.calcsize(_index_struct_string) + _verify_struct_string = "!LLL" + _verify_size = struct.calcsize(_verify_struct_string) + + def __init__(self, storage, **kwds): + self._key = kwds.pop('key', None) + if self._key is None: + raise ValueError( + "An encryption key is required using " + "the 'key' keyword.") + if isinstance(storage, BlockStorageInterface): + storage_owned = False + self._storage = storage + if len(kwds): + raise ValueError( + "Keywords not used when initializing " + "with a storage device: %s" + % (str(kwds))) + else: + storage_owned = True + storage_type = kwds.pop('storage_type', 'file') + self._storage = \ + BlockStorageTypeFactory(storage_type)(storage, **kwds) + + try: + header_data = AES.GCMDec(self._key, + self._storage.header_data) + (self._ismodegcm,) = struct.unpack( + self._index_struct_string, + header_data[:self._index_offset]) + self._verify_digest = header_data[:hashlib.sha384().digest_size] + + verify = hmac.HMAC( + key=self.key, + msg=struct.pack(self._verify_struct_string, + self._storage.block_size, + self._storage.block_count, + len(self._storage.header_data)), + digestmod=hashlib.sha384) + if verify.digest() != self._verify_digest: + raise ValueError( + "HMAC of plaintext index data does not match") + if self._ismodegcm: + self._encrypt_block_func = AES.GCMEnc + self._decrypt_block_func = AES.GCMDec + else: + self._encrypt_block_func = AES.CTREnc + self._decrypt_block_func = AES.CTRDec + except: + if storage_owned: + self._storage.close() + raise + + # + # Define EncryptedBlockStorageInterface Methods + # + + @property + def key(self): + return self._key + + @property + def raw_storage(self): + return self._storage + + # + # Define BlockStorageInterface Methods + # + + def clone_device(self): + return EncryptedBlockStorage(self._storage.clone_device(), + key=self.key) + + @classmethod + def compute_storage_size(cls, + block_size, + block_count, + aes_mode='ctr', + storage_type='file', + ignore_header=False, + **kwds): + assert (block_size > 0) and (block_size == int(block_size)) + assert (block_count > 0) and (block_count == int(block_count)) + assert aes_mode in ('ctr', 'gcm') + if not isinstance(storage_type, BlockStorageInterface): + storage_type = BlockStorageTypeFactory(storage_type) + + if aes_mode == 'ctr': + extra_block_data = AES.block_size + else: + assert aes_mode == 'gcm' + extra_block_data = 2 * AES.block_size + if ignore_header: + return (extra_block_data * block_count) + \ + storage_type.compute_storage_size( + block_size, + block_count, + ignore_header=True, + **kwds) + else: + return cls._index_offset + \ + 2 * AES.block_size + \ + (extra_block_data * block_count) + \ + storage_type.compute_storage_size( + block_size, + block_count, + ignore_header=False, + **kwds) + + @classmethod + def setup(cls, + storage_name, + block_size, + block_count, + aes_mode='ctr', + key_size=None, + key=None, + storage_type='file', + initialize=None, + **kwds): + + if (key is not None) and (key_size is not None): + raise ValueError( + "Only one of 'key' or 'keysize' keywords can " + "be specified at a time") + if key is None: + if key_size is None: + key_size = AES.key_sizes[-1] + if key_size not in AES.key_sizes: + raise ValueError( + "Invalid key size: %s" % (key_size)) + key = AES.KeyGen(key_size) + else: + if len(key) not in AES.key_sizes: + raise ValueError( + "Invalid key size: %s" % (len(key))) + + if (block_size <= 0) or (block_size != int(block_size)): + raise ValueError( + "Block size (bytes) must be a positive integer: %s" + % (block_size)) + + ismodegcm = None + encrypt_block_func = None + encrypted_block_size = block_size + if aes_mode == 'ctr': + ismodegcm = False + encrypt_block_func = AES.CTREnc + encrypted_block_size += AES.block_size + elif aes_mode == 'gcm': + ismodegcm = True + encrypt_block_func = AES.GCMEnc + encrypted_block_size += (2 * AES.block_size) + else: + raise ValueError( + "AES encryption mode must be one of 'ctr' or 'gcm'. " + "Invalid value: %s" % (aes_mode)) + assert ismodegcm is not None + assert encrypt_block_func is not None + + if not isinstance(storage_type, BlockStorageInterface): + storage_type = BlockStorageTypeFactory(storage_type) + + if initialize is None: + zeros = bytes(bytearray(block_size)) + initialize = lambda i: zeros + def encrypted_initialize(i): + return encrypt_block_func(key, initialize(i)) + kwds['initialize'] = encrypted_initialize + + user_header_data = kwds.get('header_data', bytes()) + if type(user_header_data) is not bytes: + raise TypeError( + "'header_data' must be of type bytes. " + "Invalid type: %s" % (type(user_header_data))) + # we generate the first time simply to + # compute the length + tmp = hmac.HMAC( + key=key, + msg=struct.pack(cls._verify_struct_string, + encrypted_block_size, + block_count, + 0), + digestmod=hashlib.sha384).digest() + header_data = bytearray(struct.pack(cls._index_struct_string, + ismodegcm)) + header_data[:hashlib.sha384().digest_size] = tmp + header_data = header_data + user_header_data + header_data = AES.GCMEnc(key, bytes(header_data)) + # now that we know the length of the header data + # being sent to the underlying storage we can + # compute the real hmac + verify_digest = hmac.HMAC( + key=key, + msg=struct.pack(cls._verify_struct_string, + encrypted_block_size, + block_count, + len(header_data)), + digestmod=hashlib.sha384).digest() + header_data = bytearray(struct.pack(cls._index_struct_string, + ismodegcm)) + header_data[:hashlib.sha384().digest_size] = verify_digest + header_data = header_data + user_header_data + kwds['header_data'] = AES.GCMEnc(key, bytes(header_data)) + + return EncryptedBlockStorage( + storage_type.setup(storage_name, + encrypted_block_size, + block_count, + **kwds), + key=key) + + @property + def header_data(self): + return AES.GCMDec(self._key, + self._storage.header_data)\ + [self._index_offset:] + + @property + def block_count(self): + return self._storage.block_count + + @property + def block_size(self): + if self._ismodegcm: + return self._storage.block_size - 2 * AES.block_size + else: + return self._storage.block_size - AES.block_size + + @property + def storage_name(self): + return self._storage.storage_name + + def update_header_data(self, new_header_data): + self._storage.update_header_data( + AES.GCMEnc( + self.key, + AES.GCMDec(self._key, + self._storage.header_data)\ + [:self._index_offset] + \ + new_header_data)) + + def close(self): + self._storage.close() + + + + + + + + + + + + + def read_block(self, i): + a = self._storage.read_block(i) + return self._decrypt_block_func(self._key,a) + + def read_blocks(self, indices, *args, **kwds): + a = self._storage.read_blocks(indices, *args, **kwds) + return [self._decrypt_block_func(self._key, b) for b in a] + + def yield_blocks(self, indices, *args, **kwds): + for b in self._storage.yield_blocks(indices, *args, **kwds): + yield self._decrypt_block_func(self._key, b) + + + def write_block(self, i, block, *args, **kwds): + a = self._encrypt_block_func(self._key, block) + self._storage.write_block(i, a,*args, **kwds) + + def write_blocks(self, indices, blocks, *args, **kwds): + enc_blocks = [] + for i, b in zip(indices, blocks): + enc_blocks.append(self._encrypt_block_func(self._key, b)) + + + self._storage.write_blocks(indices, enc_blocks, *args, **kwds) + + + + + + + + + + + + @property + def bytes_sent(self): + return self._storage.bytes_sent + + @property + def bytes_received(self): + return self._storage.bytes_received diff --git a/PyORAM/src/pyoram/encrypted_storage/encrypted_heap_storage.py b/PyORAM/src/pyoram/encrypted_storage/encrypted_heap_storage.py new file mode 100644 index 0000000..763ac93 --- /dev/null +++ b/PyORAM/src/pyoram/encrypted_storage/encrypted_heap_storage.py @@ -0,0 +1,172 @@ +__all__ = ('EncryptedHeapStorage',) + +import struct + +from pyoram.util.virtual_heap import SizedVirtualHeap +from pyoram.storage.heap_storage import \ + (HeapStorageInterface, + HeapStorage) +from pyoram.encrypted_storage.encrypted_block_storage import \ + (EncryptedBlockStorageInterface, + EncryptedBlockStorage) + +class EncryptedHeapStorageInterface(HeapStorageInterface): + + # + # Abstract Interface + # + + @property + def key(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + @property + def raw_storage(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + +class EncryptedHeapStorage(HeapStorage, + EncryptedHeapStorageInterface): + + def __init__(self, storage, **kwds): + + if isinstance(storage, EncryptedBlockStorageInterface): + if len(kwds): + raise ValueError( + "Keywords not used when initializing " + "with a storage device: %s" + % (str(kwds))) + else: + storage = EncryptedBlockStorage(storage, **kwds) + + super(EncryptedHeapStorage, self).__init__(storage) + + # + # Define EncryptedHeapStorageInterface Methods + # + + @property + def key(self): + return self._storage.key + + @property + def raw_storage(self): + return self._storage.raw_storage + + # + # Define HeapStorageInterface Methods + # (override what is defined on HeapStorage) + + def clone_device(self): + return EncryptedHeapStorage(self._storage.clone_device()) + + @classmethod + def compute_storage_size(cls, + block_size, + heap_height, + blocks_per_bucket=1, + heap_base=2, + ignore_header=False, + **kwds): + assert (block_size > 0) and (block_size == int(block_size)) + assert heap_height >= 0 + assert blocks_per_bucket >= 1 + assert heap_base >= 2 + assert 'block_count' not in kwds + vheap = SizedVirtualHeap( + heap_base, + heap_height, + blocks_per_bucket=blocks_per_bucket) + if ignore_header: + return EncryptedBlockStorage.compute_storage_size( + vheap.blocks_per_bucket * block_size, + vheap.bucket_count(), + ignore_header=True, + **kwds) + else: + return cls._header_offset + \ + EncryptedBlockStorage.compute_storage_size( + vheap.blocks_per_bucket * block_size, + vheap.bucket_count(), + ignore_header=False, + **kwds) + + @classmethod + def setup(cls, + storage_name, + block_size, + heap_height, + blocks_per_bucket=1, + heap_base=2, + **kwds): + if 'block_count' in kwds: + raise ValueError("'block_count' keyword is not accepted") + if heap_height < 0: + raise ValueError( + "heap height must be 0 or greater. Invalid value: %s" + % (heap_height)) + if blocks_per_bucket < 1: + raise ValueError( + "blocks_per_bucket must be 1 or greater. " + "Invalid value: %s" % (blocks_per_bucket)) + if heap_base < 2: + raise ValueError( + "heap base must be 2 or greater. Invalid value: %s" + % (heap_base)) + + vheap = SizedVirtualHeap( + heap_base, + heap_height, + blocks_per_bucket=blocks_per_bucket) + + user_header_data = kwds.pop('header_data', bytes()) + if type(user_header_data) is not bytes: + raise TypeError( + "'header_data' must be of type bytes. " + "Invalid type: %s" % (type(user_header_data))) + kwds['header_data'] = \ + struct.pack(cls._header_struct_string, + heap_base, + heap_height, + blocks_per_bucket) + \ + user_header_data + + return EncryptedHeapStorage( + EncryptedBlockStorage.setup( + storage_name, + vheap.blocks_per_bucket * block_size, + vheap.bucket_count(), + **kwds)) + + #@property + #def header_data(...) + + #@property + #def bucket_count(...) + + #@property + #def bucket_size(...) + + #@property + #def blocks_per_bucket(...) + + #@property + #def storage_name(...) + + #@property + #def virtual_heap(...) + + #@property + #def bucket_storage(...) + + #def update_header_data(...) + + #def close(...) + + #def read_path(...) + + #def write_path(...) + + #@property + #def bytes_sent(...) + + #@property + #def bytes_received(...) diff --git a/PyORAM/src/pyoram/encrypted_storage/top_cached_encrypted_heap_storage.py b/PyORAM/src/pyoram/encrypted_storage/top_cached_encrypted_heap_storage.py new file mode 100644 index 0000000..93d5abe --- /dev/null +++ b/PyORAM/src/pyoram/encrypted_storage/top_cached_encrypted_heap_storage.py @@ -0,0 +1,305 @@ +__all__ = ('TopCachedEncryptedHeapStorage',) + +import logging +import tempfile +import mmap + +import pyoram +from pyoram.util.virtual_heap import SizedVirtualHeap +from pyoram.encrypted_storage.encrypted_heap_storage import \ + (EncryptedHeapStorageInterface, + EncryptedHeapStorage) + +import tqdm +import six +from six.moves import xrange + +log = logging.getLogger("pyoram") + +class TopCachedEncryptedHeapStorage(EncryptedHeapStorageInterface): + """ + An encrypted block storage device for accessing memory + organized as a heap, where the top 1 or more levels can + be cached in local memory. This achieves two things: + + (1) Reduces the number of buckets that need to be read + from or written to external storage for a given + path I/O operation. + (2) Allows certain block storage devices to achieve + concurrency across path writes by partioning the + storage space into independent subheaps starting + below the cache line. + + This devices takes as input an existing encrypted heap + storage device. This class should not be cloned or used + to setup storage, but rather used as a wrapper class for + an existing heap storage device to speed up a bulk set + of I/O requests. The original heap storage device should + not be used after it is wrapped by this class. This + class will close the original device when closing + itself. + + The number of cached levels (starting from the root + bucket at level 0) can be set with the 'cached_levels' + keyword (>= 1). + + By default, this will create an independent storage + device capable of reading from and writing to the + original storage devices memory for each independent + subheap (if any) below the last cached level. The + 'concurrency_level' keyword can be used to limit the + number of concurrent devices to some level below the + cache line (>= 0, <= 'cached_levels'). + + Values for 'cached_levels' and 'concurrency_level' will + be automatically reduced when they are larger than what + is allowed by the heap size. + """ + + def __new__(cls, *args, **kwds): + if kwds.get("cached_levels", 1) == 0: + assert len(args) == 1 + storage = args[0] + storage.cached_bucket_data = bytes() + return storage + else: + return super(TopCachedEncryptedHeapStorage, cls).\ + __new__(cls) + + def __init__(self, + heap_storage, + cached_levels=1, + concurrency_level=None): + assert isinstance(heap_storage, EncryptedHeapStorage) + assert cached_levels != 0 + + + vheap = heap_storage.virtual_heap + if cached_levels < 0: + cached_levels = vheap.levels + if concurrency_level is None: + concurrency_level = cached_levels + assert concurrency_level >= 0 + cached_levels = min(vheap.levels, cached_levels) + concurrency_level = min(cached_levels, concurrency_level) + self._external_level = cached_levels + total_buckets = sum(vheap.bucket_count_at_level(l) + for l in xrange(cached_levels)) + + + print(" ILA ALI ALI ALI ALI ALI ALI ALI ALI ALI ALI ALI ALI ALI ALI ALI ALI ALI") + print(cached_levels) + print(concurrency_level) + print(" ALI ALI ALI ALI ALI ALI ALI ALI ALI ALI ALI ALI ALI ALI ALI ALI ALI ALI") + + + self._root_device = heap_storage + # clone before we download the cache so that we can + # track bytes transferred during read/write requests + # (separate from the cached download) + self._concurrent_devices = \ + {vheap.first_bucket_at_level(0): self._root_device.clone_device()} + + self._cached_bucket_count = total_buckets + self._cached_buckets_tempfile = tempfile.TemporaryFile() + self._cached_buckets_tempfile.seek(0) + with tqdm.tqdm(desc=("Downloading %s Cached Heap Buckets" + % (self._cached_bucket_count)), + total=self._cached_bucket_count*self._root_device.bucket_size, + unit="B", + unit_scale=True, + disable=not pyoram.config.SHOW_PROGRESS_BAR) as progress_bar: + for b, bucket in enumerate( + self._root_device.bucket_storage.yield_blocks( + xrange(vheap.first_bucket_at_level(cached_levels)))): + self._cached_buckets_tempfile.write(bucket) + progress_bar.update(self._root_device.bucket_size) + self._cached_buckets_tempfile.flush() + self._cached_buckets_mmap = mmap.mmap( + self._cached_buckets_tempfile.fileno(), 0) + + log.info("%s: Cloning %s sub-heap devices" + % (self.__class__.__name__, vheap.bucket_count_at_level(concurrency_level))) + # Avoid cloning devices when the cache line is at the root + # bucket or when the entire heap is cached + if (concurrency_level > 0) and \ + (concurrency_level <= vheap.last_level): + for b in xrange(vheap.first_bucket_at_level(concurrency_level), + vheap.first_bucket_at_level(concurrency_level+1)): + try: + self._concurrent_devices[b] = self._root_device.clone_device() + except: # pragma: no cover + log.error( # pragma: no cover + "%s: Exception encountered " # pragma: no cover + "while cloning device. " # pragma: no cover + "Closing storage." # pragma: no cover + % (self.__class__.__name__)) # pragma: no cover + self.close() # pragma: no cover + raise # pragma: no cover + + self._subheap_storage = {} + # Avoid populating this dictionary when the entire + # heap is cached + if self._external_level <= vheap.last_level: + for b in xrange(vheap.first_bucket_at_level(self._external_level), + vheap.first_bucket_at_level(self._external_level+1)): + node = vheap.Node(b) + while node.bucket not in self._concurrent_devices: + node = node.parent_node() + assert node.bucket >= 0 + assert node.level == concurrency_level + self._subheap_storage[b] = self._concurrent_devices[node.bucket] + + # + # Additional Methods + # + + @property + def cached_bucket_data(self): + return self._cached_buckets_mmap + + # + # Define EncryptedHeapStorageInterface Methods + # + + @property + def key(self): + return self._root_device.key + + @property + def raw_storage(self): + return self._root_device.raw_storage + + # + # Define HeapStorageInterface Methods + # + + def clone_device(self, *args, **kwds): + raise NotImplementedError( # pragma: no cover + "Class is not designed for cloning") # pragma: no cover + + @classmethod + def compute_storage_size(cls, *args, **kwds): + return EncryptedHeapStorage.compute_storage_size(*args, **kwds) + + @classmethod + def setup(cls, *args, **kwds): + raise NotImplementedError( # pragma: no cover + "Class is not designed to setup storage") # pragma: no cover + + @property + def header_data(self): + return self._root_device.header_data + + @property + def bucket_count(self): + return self._root_device.bucket_count + + @property + def bucket_size(self): + return self._root_device.bucket_size + + @property + def blocks_per_bucket(self): + return self._root_device.blocks_per_bucket + + @property + def storage_name(self): + return self._root_device.storage_name + + @property + def virtual_heap(self): + return self._root_device.virtual_heap + + @property + def bucket_storage(self): + return self._root_device.bucket_storage + + def update_header_data(self, new_header_data): + self._root_device.update_header_data(new_header_data) + + def close(self): + print("Heap Closing 1") + log.info("%s: Uploading %s cached bucket data before closing" + % (self.__class__.__name__, self._cached_bucket_count)) + with tqdm.tqdm(desc=("Uploading %s Cached Heap Buckets" + % (self._cached_bucket_count)), + total=self._cached_bucket_count*self.bucket_size, + unit="B", + unit_scale=True, + disable=not pyoram.config.SHOW_PROGRESS_BAR) as progress_bar: + self.bucket_storage.\ + write_blocks( + xrange(self._cached_bucket_count), + (self._cached_buckets_mmap[(b*self.bucket_size): + ((b+1)*self.bucket_size)] + for b in xrange(self._cached_bucket_count)), + callback=lambda i: progress_bar.update(self._root_device.bucket_size)) + for b in self._concurrent_devices: + self._concurrent_devices[b].close() + self._root_device.close() + # forces the bar to become full at close + # even if te write_blocks action was faster + # the the mininterval time + progress_bar.mininterval = 0 + + self._cached_buckets_mmap.close() + self._cached_buckets_tempfile.close() + + def read_path(self, b, level_start=0): + assert 0 <= b < self.virtual_heap.bucket_count() + bucket_list = self.virtual_heap.Node(b).bucket_path_from_root() + if len(bucket_list) <= self._external_level: + return [self._cached_buckets_mmap[(bb*self.bucket_size): + ((bb+1)*self.bucket_size)] + for bb in bucket_list[level_start:]] + elif level_start >= self._external_level: + return self._subheap_storage[bucket_list[self._external_level]].\ + bucket_storage.read_blocks(bucket_list[level_start:]) + else: + local_buckets = bucket_list[:self._external_level] + external_buckets = bucket_list[self._external_level:] + buckets = [] + for bb in local_buckets[level_start:]: + buckets.append( + self._cached_buckets_mmap[(bb*self.bucket_size): + ((bb+1)*self.bucket_size)]) + if len(external_buckets) > 0: + buckets.extend( + self._subheap_storage[external_buckets[0]].\ + bucket_storage.read_blocks(external_buckets)) + assert len(buckets) == len(bucket_list[level_start:]) + return buckets + + def write_path(self, b, buckets, level_start=0): + assert 0 <= b < self.virtual_heap.bucket_count() + bucket_list = self.virtual_heap.Node(b).bucket_path_from_root() + if len(bucket_list) <= self._external_level: + for bb, bucket in zip(bucket_list[level_start:], buckets): + self._cached_buckets_mmap[(bb*self.bucket_size): + ((bb+1)*self.bucket_size)] = bucket + elif level_start >= self._external_level: + self._subheap_storage[bucket_list[self._external_level]].\ + bucket_storage.write_blocks(bucket_list[level_start:], buckets) + else: + buckets = list(buckets) + assert len(buckets) == len(bucket_list[level_start:]) + local_buckets = bucket_list[:self._external_level] + external_buckets = bucket_list[self._external_level:] + ndx = -1 + for ndx, bb in enumerate(local_buckets[level_start:]): + self._cached_buckets_mmap[(bb*self.bucket_size): + ((bb+1)*self.bucket_size)] = buckets[ndx] + if len(external_buckets) > 0: + self._subheap_storage[external_buckets[0]].\ + bucket_storage.write_blocks(external_buckets, + buckets[(ndx+1):]) + @property + def bytes_sent(self): + return sum(device.bytes_sent for device + in self._concurrent_devices.values()) + + @property + def bytes_received(self): + return sum(device.bytes_received for device + in self._concurrent_devices.values()) diff --git a/PyORAM/src/pyoram/oblivious_storage/__init__.py b/PyORAM/src/pyoram/oblivious_storage/__init__.py new file mode 100644 index 0000000..22d6ec8 --- /dev/null +++ b/PyORAM/src/pyoram/oblivious_storage/__init__.py @@ -0,0 +1 @@ +import pyoram.oblivious_storage.tree diff --git a/PyORAM/src/pyoram/oblivious_storage/tree/__init__.py b/PyORAM/src/pyoram/oblivious_storage/tree/__init__.py new file mode 100644 index 0000000..8eb08a7 --- /dev/null +++ b/PyORAM/src/pyoram/oblivious_storage/tree/__init__.py @@ -0,0 +1,2 @@ +import pyoram.oblivious_storage.tree.tree_oram_helper +import pyoram.oblivious_storage.tree.path_oram diff --git a/PyORAM/src/pyoram/oblivious_storage/tree/path_oram.py b/PyORAM/src/pyoram/oblivious_storage/tree/path_oram.py new file mode 100644 index 0000000..17a2fe0 --- /dev/null +++ b/PyORAM/src/pyoram/oblivious_storage/tree/path_oram.py @@ -0,0 +1,466 @@ +import hashlib +import hmac +import struct +import array +import logging + +import pyoram +from pyoram.oblivious_storage.tree.tree_oram_helper import \ + (TreeORAMStorage, + TreeORAMStorageManagerExplicitAddressing) +from pyoram.encrypted_storage.encrypted_block_storage import \ + EncryptedBlockStorageInterface +from pyoram.encrypted_storage.encrypted_heap_storage import \ + (EncryptedHeapStorage, + EncryptedHeapStorageInterface) +from pyoram.encrypted_storage.top_cached_encrypted_heap_storage import \ + TopCachedEncryptedHeapStorage +from pyoram.util.virtual_heap import \ + (SizedVirtualHeap, + calculate_necessary_heap_height) + +import tqdm +import six +from six.moves import xrange + +log = logging.getLogger("pyoram") + +class PathORAM(EncryptedBlockStorageInterface): + + _header_struct_string = "!"+("x"*2*hashlib.sha384().digest_size)+"L" + _header_offset = struct.calcsize(_header_struct_string) + + def __init__(self, + storage, + stash, + position_map, + **kwds): + + + self._oram = None + self._block_count = None + + if isinstance(storage, EncryptedHeapStorageInterface): + storage_heap = storage + close_storage_heap = False + if len(kwds): + raise ValueError( + "Keywords not used when initializing " + "with a storage device: %s" + % (str(kwds))) + else: + cached_levels = kwds.pop('cached_levels', 3) + # cached_levels = kwds.pop('cached_levels', 1) + concurrency_level = kwds.pop('concurrency_level', None) + close_storage_heap = True + storage_heap = TopCachedEncryptedHeapStorage(EncryptedHeapStorage(storage, **kwds), cached_levels=cached_levels, concurrency_level=concurrency_level) + # storage_heap = EncryptedHeapStorage(storage, **kwds) + + (self._block_count,) = struct.unpack( + self._header_struct_string, + storage_heap.header_data[:self._header_offset]) + stashdigest = storage_heap.\ + header_data[:hashlib.sha384().digest_size] + positiondigest = storage_heap.\ + header_data[hashlib.sha384().digest_size:\ + (2*hashlib.sha384().digest_size)] + + try: + if stashdigest != \ + PathORAM.stash_digest( + stash, + digestmod=hmac.HMAC(key=storage_heap.key, + digestmod=hashlib.sha384)): + raise ValueError( + "Stash HMAC does not match that saved with " + "storage heap %s" % (storage_heap.storage_name)) + except: + if close_storage_heap: + storage_heap.close() + raise + + try: + if positiondigest != \ + PathORAM.position_map_digest( + position_map, + digestmod=hmac.HMAC(key=storage_heap.key, + digestmod=hashlib.sha384)): + raise ValueError( + "Position map HMAC does not match that saved with " + "storage heap %s" % (storage_heap.storage_name)) + except: + if close_storage_heap: + storage_heap.close() + raise + + self._oram = TreeORAMStorageManagerExplicitAddressing( + storage_heap, + stash, + position_map) + assert self._block_count <= \ + self._oram.storage_heap.bucket_count + + @classmethod + def _init_position_map(cls, vheap, block_count): + return array.array("L", [vheap.random_leaf_bucket() + for i in xrange(block_count)]) + + def _init_oram_block(self, id_, block): + oram_block = bytearray(self.block_size) + oram_block[self._oram.block_info_storage_size:] = block[:] + self._oram.tag_block_with_id(oram_block, id_) + return oram_block + + def _extract_virtual_block(self, block): + return block[self._oram.block_info_storage_size:] + + # + # Add some methods specific to Path ORAM + # + + @classmethod + def stash_digest(cls, stash, digestmod=None): + if digestmod is None: + digestmod = hashlib.sha1() + id_to_bytes = lambda id_: \ + struct.pack(TreeORAMStorage.block_id_storage_string, id_) + if len(stash) == 0: + digestmod.update(b'0') + else: + for id_ in sorted(stash): + if id_ < 0: + raise ValueError( + "Invalid stash id '%s'. Values must be " + "nonnegative integers." % (id_)) + digestmod.update(id_to_bytes(id_)) + digestmod.update(bytes(stash[id_])) + return digestmod.digest() + + @classmethod + def position_map_digest(cls, position_map, digestmod=None): + if digestmod is None: + digestmod = hashlib.sha1() + id_to_bytes = lambda id_: \ + struct.pack(TreeORAMStorage.block_id_storage_string, id_) + assert len(position_map) > 0 + for addr in position_map: + if addr < 0: + raise ValueError( + "Invalid position map address '%s'. Values must be " + "nonnegative integers." % (addr)) + digestmod.update(id_to_bytes(addr)) + return digestmod.digest() + + @property + def position_map(self): + return self._oram.position_map + + @property + def stash(self): + return self._oram.stash + + def access(self, id_, write_block=None): + assert 0 <= id_ <= self.block_count + bucket = self.position_map[id_] + bucket_level = self._oram.storage_heap.virtual_heap.Node(bucket).level + self.position_map[id_] = \ + self._oram.storage_heap.virtual_heap.\ + random_bucket_at_level(bucket_level) + self._oram.load_path(bucket) + block = self._oram.extract_block_from_path(id_) + + if block is None: + block = self.stash[id_] + + + if write_block is not None: + block = self._init_oram_block(id_, write_block) + + + self.stash[id_] = block + self._oram.push_down_path() + self._oram.fill_path_from_stash() + self._oram.evict_path() + if write_block is None: + return self._extract_virtual_block(block) + + @property + def heap_storage(self): + return self._oram.storage_heap + + # + # Define EncryptedBlockStorageInterface Methods + # + + @property + def key(self): + return self._oram.storage_heap.key + + @property + def raw_storage(self): + return self._oram.storage_heap.raw_storage + + # + # Define BlockStorageInterface Methods + # + + @classmethod + def compute_storage_size(cls, + block_size, + block_count, + bucket_capacity=4, + heap_base=2, + ignore_header=False, + **kwds): + assert (block_size > 0) and (block_size == int(block_size)) + assert (block_count > 0) and (block_count == int(block_count)) + assert bucket_capacity >= 1 + assert heap_base >= 2 + assert 'heap_height' not in kwds + heap_height = calculate_necessary_heap_height(heap_base, + block_count) + block_size += TreeORAMStorageManagerExplicitAddressing.\ + block_info_storage_size + if ignore_header: + return EncryptedHeapStorage.compute_storage_size( + block_size, + heap_height, + blocks_per_bucket=bucket_capacity, + heap_base=heap_base, + ignore_header=True, + **kwds) + else: + return cls._header_offset + \ + EncryptedHeapStorage.compute_storage_size( + block_size, + heap_height, + blocks_per_bucket=bucket_capacity, + heap_base=heap_base, + ignore_header=False, + **kwds) + + @classmethod + def setup(cls, + storage_name, + block_size, + block_count, + bucket_capacity=4, + heap_base=2, + cached_levels=3, + concurrency_level=None, + **kwds): + if 'heap_height' in kwds: + raise ValueError("'heap_height' keyword is not accepted") + if (bucket_capacity <= 0) or \ + (bucket_capacity != int(bucket_capacity)): + raise ValueError( + "Bucket capacity must be a positive integer: %s" + % (bucket_capacity)) + if (block_size <= 0) or (block_size != int(block_size)): + raise ValueError( + "Block size (bytes) must be a positive integer: %s" + % (block_size)) + if (block_count <= 0) or (block_count != int(block_count)): + raise ValueError( + "Block count must be a positive integer: %s" + % (block_count)) + + if heap_base < 2: + raise ValueError( + "heap base must be 2 or greater. Invalid value: %s" + % (heap_base)) + + heap_height = calculate_necessary_heap_height(heap_base, + block_count) + stash = {} + vheap = SizedVirtualHeap( + heap_base, + heap_height, + blocks_per_bucket=bucket_capacity) + position_map = cls._init_position_map(vheap, block_count) + + oram_block_size = block_size + \ + TreeORAMStorageManagerExplicitAddressing.\ + block_info_storage_size + + user_header_data = kwds.pop('header_data', bytes()) + if type(user_header_data) is not bytes: + raise TypeError( + "'header_data' must be of type bytes. " + "Invalid type: %s" % (type(user_header_data))) + + initialize = kwds.pop('initialize', None) + + header_data = struct.pack( + cls._header_struct_string, + block_count) + kwds['header_data'] = bytes(header_data) + user_header_data + empty_bucket = bytearray(oram_block_size * bucket_capacity) + empty_bucket_view = memoryview(empty_bucket) + for i in xrange(bucket_capacity): + TreeORAMStorageManagerExplicitAddressing.tag_block_as_empty( + empty_bucket_view[(i*oram_block_size):\ + ((i+1)*oram_block_size)]) + empty_bucket = bytes(empty_bucket) + + kwds['initialize'] = lambda i: empty_bucket + f = None + try: + log.info("%s: setting up encrypted heap storage" + % (cls.__name__)) + f = EncryptedHeapStorage.setup(storage_name, + oram_block_size, + heap_height, + heap_base=heap_base, + blocks_per_bucket=bucket_capacity, + **kwds) + if cached_levels != 0: + f = TopCachedEncryptedHeapStorage( + f, + cached_levels=cached_levels, + concurrency_level=concurrency_level) + elif concurrency_level is not None: + raise ValueError( # pragma: no cover + "'concurrency_level' keyword is " # pragma: no cover + "not used when no heap levels " # pragma: no cover + "are cached") # pragma: no cover + oram = TreeORAMStorageManagerExplicitAddressing( + f, stash, position_map) + if initialize is None: + zeros = bytes(bytearray(block_size)) + initialize = lambda i: zeros + initial_oram_block = bytearray(oram_block_size) + for i in tqdm.tqdm(xrange(block_count), + desc=("Initializing %s Blocks" % (cls.__name__)), + total=block_count, + disable=not pyoram.config.SHOW_PROGRESS_BAR): + + oram.tag_block_with_id(initial_oram_block, i) + initial_oram_block[oram.block_info_storage_size:] = \ + initialize(i)[:] + + bucket = oram.position_map[i] + bucket_level = vheap.Node(bucket).level + oram.position_map[i] = \ + oram.storage_heap.virtual_heap.\ + random_bucket_at_level(bucket_level) + + oram.load_path(bucket) + oram.push_down_path() + # place a copy in the stash + oram.stash[i] = bytearray(initial_oram_block) + oram.fill_path_from_stash() + oram.evict_path() + + header_data = bytearray(header_data) + stash_digest = cls.stash_digest( + oram.stash, + digestmod=hmac.HMAC(key=oram.storage_heap.key, + digestmod=hashlib.sha384)) + position_map_digest = cls.position_map_digest( + oram.position_map, + digestmod=hmac.HMAC(key=oram.storage_heap.key, + digestmod=hashlib.sha384)) + header_data[:len(stash_digest)] = stash_digest[:] + header_data[len(stash_digest):\ + (len(stash_digest)+len(position_map_digest))] = \ + position_map_digest[:] + f.update_header_data(bytes(header_data) + user_header_data) + return PathORAM(f, stash, position_map=position_map) + except: + if f is not None: + f.close() # pragma: no cover + raise + + @property + def header_data(self): + return self._oram.storage_heap.\ + header_data[self._header_offset:] + + @property + def block_count(self): + return self._block_count + + @property + def block_size(self): + return self._oram.block_size - self._oram.block_info_storage_size + + @property + def storage_name(self): + return self._oram.storage_heap.storage_name + + def update_header_data(self, new_header_data): + self._oram.storage_heap.update_header_data( + self._oram.storage_heap.header_data[:self._header_offset] + \ + new_header_data) + + def close(self): + log.info("%s: Closing" % (self.__class__.__name__)) + print("Closing") + + if self._oram is not None: + try: + stashdigest = \ + PathORAM.stash_digest( + self._oram.stash, + digestmod=hmac.HMAC(key=self._oram.storage_heap.key, + digestmod=hashlib.sha384)) + + print("Closing 1") + positiondigest = \ + PathORAM.position_map_digest( + self._oram.position_map, + digestmod=hmac.HMAC(key=self._oram.storage_heap.key, + digestmod=hashlib.sha384)) + + print("Closing 2") + new_header_data = \ + bytearray(self._oram.storage_heap.\ + header_data[:self._header_offset]) + + print("Closing 3") + new_header_data[:hashlib.sha384().digest_size] = \ + stashdigest + new_header_data[hashlib.sha384().digest_size:\ + (2*hashlib.sha384().digest_size)] = \ + positiondigest + + print("Closing 4") + self._oram.storage_heap.update_header_data( + bytes(new_header_data) + self.header_data) + print("Closing 5") + except: # pragma: no cover + log.error( # pragma: no cover + "%s: Failed to update header data with " # pragma: no cover + "current stash and position map state" # pragma: no cover + % (self.__class__.__name__)) # pragma: no cover + print("Closing ") + raise + finally: + print("Closing 6") + self._oram.storage_heap.close() + print("Closing 7") + + def read_blocks(self, indices): + blocks = [] + for i in indices: + blocks.append(self.access(i)) + return blocks + + def read_block(self, i): + return self.access(i) + + def write_blocks(self, indices, blocks): + for i, block in zip(indices, blocks): + self.access(i, write_block=block) + + def write_block(self, i, block): + self.access(i, write_block=block) + + @property + def bytes_sent(self): + return self._oram.storage_heap.bytes_sent + + @property + def bytes_received(self): + return self._oram.storage_heap.bytes_received diff --git a/PyORAM/src/pyoram/oblivious_storage/tree/tree_oram_helper.py b/PyORAM/src/pyoram/oblivious_storage/tree/tree_oram_helper.py new file mode 100644 index 0000000..560a744 --- /dev/null +++ b/PyORAM/src/pyoram/oblivious_storage/tree/tree_oram_helper.py @@ -0,0 +1,329 @@ +__all__ = ('TreeORAMStorageManagerExplicitAddressing', + 'TreeORAMStorageManagerPointerAddressing') + +import struct +import copy + +from pyoram.util.virtual_heap import \ + SizedVirtualHeap + +from six.moves import xrange + +class TreeORAMStorage(object): + + empty_block_id = -1 + + block_status_storage_string = "!?" + block_id_storage_string = "!L" + block_info_storage_string = "!?L" + + block_status_storage_size = \ + struct.calcsize(block_status_storage_string) + block_info_storage_size = \ + struct.calcsize(block_info_storage_string) + + empty_block_bytes_tag = \ + struct.pack(block_status_storage_string, False) + real_block_bytes_tag = \ + struct.pack(block_status_storage_string, True) + + def __init__(self, + storage_heap, + stash): + self.storage_heap = storage_heap + self.stash = stash + + vheap = self.storage_heap.virtual_heap + self.bucket_size = self.storage_heap.bucket_size + self.block_size = self.bucket_size // vheap.blocks_per_bucket + assert self.block_size * vheap.blocks_per_bucket == \ + self.bucket_size + + self.path_stop_bucket = None + self.path_bucket_count = 0 + self.path_byte_dataview = \ + bytearray(self.bucket_size * vheap.levels) + dataview = memoryview(self.path_byte_dataview) + self.path_bucket_dataview = \ + [dataview[(i*self.bucket_size):((i+1)*self.bucket_size)] + for i in xrange(vheap.levels)] + + self.path_block_dataview = [] + for i in xrange(vheap.levels): + bucketview = self.path_bucket_dataview[i] + for j in xrange(vheap.blocks_per_bucket): + self.path_block_dataview.append( + bucketview[(j*self.block_size):((j+1)*self.block_size)]) + + max_blocks_on_path = vheap.levels * vheap.blocks_per_bucket + assert len(self.path_block_dataview) == max_blocks_on_path + self.path_block_ids = [-1] * max_blocks_on_path + self.path_block_eviction_levels = [None] * max_blocks_on_path + self.path_block_reordering = [None] * max_blocks_on_path + self.path_blocks_inserted = [] + + def load_path(self, b): + vheap = self.storage_heap.virtual_heap + Z = vheap.blocks_per_bucket + lcl = vheap.clib.calculate_last_common_level + k = vheap.k + + read_level_start = 0 + if self.path_stop_bucket is not None: + # don't download the root and any other buckets + # that are common between the previous bucket path + # and the new one + read_level_start = lcl(k, self.path_stop_bucket, b) + assert 0 <= b < vheap.bucket_count() + self.path_stop_bucket = b + new_buckets = self.storage_heap.read_path( + self.path_stop_bucket, + level_start=read_level_start) + + self.path_bucket_count = read_level_start + len(new_buckets) + pos = 0 + for i in xrange(self.path_bucket_count): + if i >= read_level_start: + self.path_bucket_dataview[i][:] = \ + new_buckets[i-read_level_start][:] + for j in xrange(Z): + block_id, block_addr = \ + self.get_block_info(self.path_block_dataview[pos]) + self.path_block_ids[pos] = block_id + if block_id != self.empty_block_id: + self.path_block_eviction_levels[pos] = \ + lcl(k, self.path_stop_bucket, block_addr) + else: + self.path_block_eviction_levels[pos] = None + self.path_block_reordering[pos] = None + pos += 1 + + max_blocks_on_path = vheap.levels * Z + while pos != max_blocks_on_path: + self.path_block_ids[pos] = None + self.path_block_eviction_levels[pos] = None + self.path_block_reordering[pos] = None + pos += 1 + + self.path_blocks_inserted = [] + + def push_down_path(self): + vheap = self.storage_heap.virtual_heap + Z = vheap.blocks_per_bucket + + bucket_count = self.path_bucket_count + block_ids = self.path_block_ids + block_eviction_levels = self.path_block_eviction_levels + block_reordering = self.path_block_reordering + def _do_swap(write_pos, read_pos): + block_ids[write_pos], block_eviction_levels[write_pos] = \ + block_ids[read_pos], block_eviction_levels[read_pos] + block_ids[read_pos], block_eviction_levels[read_pos] = \ + self.empty_block_id, None + block_reordering[write_pos] = read_pos + block_reordering[read_pos] = -1 + + def _new_write_pos(current): + current -= 1 + if current < 0: + return None, None + while (block_eviction_levels[current] is not None): + current -= 1 + if current < 0: + return None, None + assert block_ids[current] == \ + self.empty_block_id + return current, current // Z + + def _new_read_pos(current): + current -= 1 + if current < 0: + return None + while (block_eviction_levels[current] is None): + current -= 1 + if current < 0: + return None + assert block_ids[current] != \ + self.empty_block_id + return current + + write_pos, write_level = _new_write_pos(bucket_count * Z) + while write_pos is not None: + read_pos = _new_read_pos(write_pos) + if read_pos is None: + break + while ((read_pos // Z) == write_level) or \ + (write_level > block_eviction_levels[read_pos]): + read_pos = _new_read_pos(read_pos) + if read_pos is None: + break + if read_pos is not None: + _do_swap(write_pos, read_pos) + else: + # Jump directly to the start of this + # bucket. There is not point in checking + # for other empty slots because no blocks + # can be evicted to this level. + write_pos = Z * (write_pos//Z) + write_pos, write_level = _new_write_pos(write_pos) + + def fill_path_from_stash(self): + vheap = self.storage_heap.virtual_heap + lcl = vheap.clib.calculate_last_common_level + k = vheap.k + Z = vheap.blocks_per_bucket + + bucket_count = self.path_bucket_count + stop_bucket = self.path_stop_bucket + block_ids = self.path_block_ids + block_eviction_levels = self.path_block_eviction_levels + blocks_inserted = self.path_blocks_inserted + + stash_eviction_levels = {} + largest_write_position = (bucket_count * Z) - 1 + for write_pos in xrange(largest_write_position,-1,-1): + write_level = write_pos // Z + if block_ids[write_pos] == self.empty_block_id: + del_id = None + for id_ in self.stash: + if id_ not in stash_eviction_levels: + block_id, block_addr = \ + self.get_block_info(self.stash[id_]) + assert id_ != self.empty_block_id + eviction_level = stash_eviction_levels[id_] = \ + lcl(k, stop_bucket, block_addr) + else: + eviction_level = stash_eviction_levels[id_] + if write_level <= eviction_level: + block_ids[write_pos] = id_ + block_eviction_levels[write_pos] = \ + eviction_level + blocks_inserted.append( + (write_pos, self.stash[id_])) + del_id = id_ + break + if del_id is not None: + del self.stash[del_id] + + def evict_path(self): + vheap = self.storage_heap.virtual_heap + Z = vheap.blocks_per_bucket + + bucket_count = self.path_bucket_count + stop_bucket = self.path_stop_bucket + bucket_dataview = self.path_bucket_dataview + block_dataview = self.path_block_dataview + block_reordering = self.path_block_reordering + blocks_inserted = self.path_blocks_inserted + + for i, read_pos in enumerate( + reversed(block_reordering)): + if (read_pos is not None) and \ + (read_pos != -1): + write_pos = len(block_reordering) - 1 - i + block_dataview[write_pos][:] = block_dataview[read_pos][:] + + for write_pos, read_pos in enumerate(block_reordering): + if read_pos == -1: + self.tag_block_as_empty(block_dataview[write_pos]) + + for write_pos, block in blocks_inserted: + block_dataview[write_pos][:] = block[:] + + self.storage_heap.write_path( + stop_bucket, + (bucket_dataview[i].tobytes() + for i in xrange(bucket_count))) + + def extract_block_from_path(self, id_): + block_ids = self.path_block_ids + block_dataview = self.path_block_dataview + try: + pos = block_ids.index(id_) + # make a copy + block = bytearray(block_dataview[pos]) + self._set_path_position_to_empty(pos) + return block + except ValueError: + return None + + def _set_path_position_to_empty(self, pos): + self.path_block_ids[pos] = self.empty_block_id + self.path_block_eviction_levels[pos] = None + self.path_block_reordering[pos] = -1 + + @staticmethod + def tag_block_as_empty(block): + block[:TreeORAMStorage.block_status_storage_size] = \ + TreeORAMStorage.empty_block_bytes_tag[:] + + @staticmethod + def tag_block_with_id(block, id_): + assert id_ >= 0 + struct.pack_into(TreeORAMStorage.block_info_storage_string, + block, + 0, + True, + id_) + + def get_block_info(self, block): + raise NotImplementedError # pragma: no cover + +class TreeORAMStorageManagerExplicitAddressing( + TreeORAMStorage): + """ + This class should be used to implement tree-based ORAMs + that use an explicit position map. Blocks are assumed to + begin with bytes representing the block id. + """ + + block_info_storage_string = \ + TreeORAMStorage.block_info_storage_string + block_info_storage_size = \ + struct.calcsize(block_info_storage_string) + + def __init__(self, + storage_heap, + stash, + position_map): + super(TreeORAMStorageManagerExplicitAddressing, self).\ + __init__(storage_heap, stash) + self.position_map = position_map + + def get_block_info(self, block): + real, id_ = struct.unpack_from( + self.block_info_storage_string, block) + if real: + return id_, self.position_map[id_] + else: + return self.empty_block_id, None + +class TreeORAMStorageManagerPointerAddressing( + TreeORAMStorage): + """ + This class should be used to implement tree-based ORAMs + that use a pointer-based position map stored with the + blocks. Blocks are assumed to begin with bytes + representing the block id followed by bytes representing + the blocks current heap bucket address. + """ + + block_info_storage_string = \ + TreeORAMStorage.block_info_storage_string + "L" + block_info_storage_size = \ + struct.calcsize(block_info_storage_string) + + def __init__(self, + storage_heap, + stash): + super(TreeORAMStorageManagerPointerAddressing, self).\ + __init__(storage_heap, stash) + self.position_map = None + + def get_block_info(self, block): + real, id_, addr = struct.unpack_from( + self.block_info_storage_string, block) + if not real: + return self.empty_block_id, 0 + else: + return id_, addr diff --git a/PyORAM/src/pyoram/storage/AliTimer.py b/PyORAM/src/pyoram/storage/AliTimer.py new file mode 100644 index 0000000..aab8223 --- /dev/null +++ b/PyORAM/src/pyoram/storage/AliTimer.py @@ -0,0 +1,64 @@ +import time + + + + +class Singleton: + """ + A non-thread-safe helper class to ease implementing singletons. + This should be used as a decorator -- not a metaclass -- to the + class that should be a singleton. + + The decorated class can define one `__init__` function that + takes only the `self` argument. Also, the decorated class cannot be + inherited from. Other than that, there are no restrictions that apply + to the decorated class. + + To get the singleton instance, use the `Instance` method. Trying + to use `__call__` will result in a `TypeError` being raised. + + """ + + def __init__(self, decorated): + self._decorated = decorated + + + def Instance(self): + """ + Returns the singleton instance. Upon its first call, it creates a + new instance of the decorated class and calls its `__init__` method. + On all subsequent calls, the already created instance is returned. + + """ + try: + return self._instance + except AttributeError: + self._instance = self._decorated() + return self._instance + + + + def __call__(self): + raise TypeError('Singletons must be accessed through `Instance()`.') + + def __instancecheck__(self, inst): + return isinstance(inst, self._decorated) + + +@Singleton +class Foo: + def __init__(self): + print 'Foo created' + self._totalTime = 0; + + def getTime(self): + return self._totalTime + + def resetTimer(self): + self._totalTime = 0 + + def startTimer(self): + self._startTime = time.time() + + def endTimer(self): + self._totalTime += time.time() - self._startTime diff --git a/PyORAM/src/pyoram/storage/__init__.py b/PyORAM/src/pyoram/storage/__init__.py new file mode 100644 index 0000000..5b3bab6 --- /dev/null +++ b/PyORAM/src/pyoram/storage/__init__.py @@ -0,0 +1,7 @@ +import pyoram.storage.block_storage +import pyoram.storage.block_storage_file +import pyoram.storage.block_storage_mmap +import pyoram.storage.block_storage_ram +import pyoram.storage.block_storage_sftp +import pyoram.storage.block_storage_s3 +import pyoram.storage.heap_storage diff --git a/PyORAM/src/pyoram/storage/block_storage.py b/PyORAM/src/pyoram/storage/block_storage.py new file mode 100644 index 0000000..41bd7dd --- /dev/null +++ b/PyORAM/src/pyoram/storage/block_storage.py @@ -0,0 +1,83 @@ +__all__ = ('BlockStorageTypeFactory',) + +import logging + +log = logging.getLogger("pyoram") + +def BlockStorageTypeFactory(storage_type_name): + if storage_type_name in BlockStorageTypeFactory._registered_devices: + return BlockStorageTypeFactory.\ + _registered_devices[storage_type_name] + else: + raise ValueError( + "BlockStorageTypeFactory: Unsupported storage " + "type: %s" % (storage_type_name)) +BlockStorageTypeFactory._registered_devices = {} + +def _register_device(name, type_): + if name in BlockStorageTypeFactory._registered_devices: + raise ValueError("Can not register block storage device type " + "with name '%s'. A device type is already " + "registered with that name." % (name)) + if not issubclass(type_, BlockStorageInterface): + raise TypeError("Can not register block storage device type " + "'%s'. The device must be a subclass of " + "BlockStorageInterface" % (type_)) + BlockStorageTypeFactory._registered_devices[name] = type_ +BlockStorageTypeFactory.register_device = _register_device + +class BlockStorageInterface(object): + + def __enter__(self): + return self + def __exit__(self, *args): + self.close() + + # + # Abstract Interface + # + + def clone_device(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + + @classmethod + def compute_storage_size(cls, *args, **kwds): + raise NotImplementedError # pragma: no cover + @classmethod + def setup(cls, *args, **kwds): + raise NotImplementedError # pragma: no cover + + @property + def header_data(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + @property + def block_count(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + @property + def block_size(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + @property + def storage_name(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + + def update_header_data(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + def close(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + def read_blocks(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + def yield_blocks(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + def read_block(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + def write_blocks(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + def write_block(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + + @property + def bytes_sent(self): + raise NotImplementedError # pragma: no cover + @property + def bytes_received(self): + raise NotImplementedError # pragma: no cover diff --git a/PyORAM/src/pyoram/storage/block_storage_file.py b/PyORAM/src/pyoram/storage/block_storage_file.py new file mode 100644 index 0000000..75932a2 --- /dev/null +++ b/PyORAM/src/pyoram/storage/block_storage_file.py @@ -0,0 +1,394 @@ +__all__ = ('BlockStorageFile',) + +import os +import struct +import logging +import errno +from multiprocessing.pool import ThreadPool + +import pyoram +from pyoram.storage.block_storage import \ + (BlockStorageInterface, + BlockStorageTypeFactory) + +import tqdm +import six +from six.moves import xrange + +import time +from AliTimer import * + +log = logging.getLogger("pyoram") + +class default_filesystem(object): + open = open + remove = os.remove + stat = os.stat + +class BlockStorageFile(BlockStorageInterface): + """ + A class implementing the block storage interface + using a local file. + """ + + _index_struct_string = "!LLL?" + _index_offset = struct.calcsize(_index_struct_string) + + def __init__(self, + storage_name, + threadpool_size=None, + ignore_lock=False, + _filesystem=default_filesystem): + self._timer = Foo.Instance(); + + self._bytes_sent = 0 + self._bytes_received = 0 + self._filesystem = _filesystem + self._ignore_lock = ignore_lock + self._f = None + self._pool = None + self._close_pool = True + self._async_write = None + self._storage_name = storage_name + self._f = self._filesystem.open(self.storage_name, "r+b") + self._f.seek(0) + self._block_size, self._block_count, user_header_size, locked = \ + struct.unpack( + BlockStorageFile._index_struct_string, + self._f.read(BlockStorageFile._index_offset)) + + if locked and (not self._ignore_lock): + self._f.close() + self._f = None + raise IOError( + "Can not open block storage device because it is " + "locked by another process. To ignore this check, " + "initialize this class with the keyword 'ignore_lock' " + "set to True.") + self._user_header_data = bytes() + if user_header_size > 0: + self._user_header_data = \ + self._f.read(user_header_size) + self._header_offset = BlockStorageFile._index_offset + \ + len(self._user_header_data) + + # TODO: Figure out why this is required for Python3 + # in order to prevent issues with the + # TopCachedEncryptedHeapStorage class. The + # problem has something to do with bufferedio, + # but it makes no sense why this fixes it (all + # we've done is read above these lines). As + # part of this, investigate whethor or not we + # need the call to flush after write_block(s), + # or if its simply connected to some Python3 + # bug in bufferedio. + self._f.flush() + + if not self._ignore_lock: + # turn on the locked flag + self._f.seek(0) + self._f.write( + struct.pack(BlockStorageFile._index_struct_string, + self.block_size, + self.block_count, + len(self._user_header_data), + True)) + self._f.flush() + + if threadpool_size != 0: + self._pool = ThreadPool(threadpool_size) + + def _check_async(self): + if self._async_write is not None: + self._async_write.get() + self._async_write = None + # TODO: Figure out why tests fail on Python3 without this + if six.PY3: + if self._f is None: + return + self._f.flush() + + def _schedule_async_write(self, args, callback=None): + assert self._async_write is None + if self._pool is not None: + self._async_write = \ + self._pool.apply_async(self._writev, (args, callback)) + else: + self._writev(args, callback) + + # This method is usually executed in another thread, so + # do not attempt to handle exceptions because it will + # not work. + def _writev(self, chunks, callback): + for i, block in chunks: + + # startTime = time.time(); + self._timer.startTimer(); + self._f.seek(self._header_offset + i * self.block_size) + self._f.write(block) + # self._f.flush() + self._timer.endTimer(); + + # print("Write....... " + str(time.time() - startTime)) + + if callback is not None: + callback(i) + + def _prep_for_close(self): + print("prep file close 1") + self._check_async() + print("prep file close 2") + + if self._close_pool and (self._pool is not None): + print("prep file close 3") + self._pool.close() + print("prep file close 4") + self._pool.join() + print("prep file close 5") + self._pool = None + print("prep file close 6") + + + if self._f is not None: + print("prep file close 7") + if not self._ignore_lock: + print("prep file close 8") + # turn off the locked flag + self._f.seek(0) + print("prep file close 9") + + a = struct.pack(BlockStorageFile._index_struct_string, + self.block_size, + self.block_count, + len(self._user_header_data), + False) + + print("prep file close 9.1") + + self._f.write(a) + print("prep file close 10") + self._f.flush() + print("prep file close 11") + + # + # Define BlockStorageInterface Methods + # + + def clone_device(self): + f = BlockStorageFile(self.storage_name, + threadpool_size=0, + ignore_lock=True) + f._pool = self._pool + f._close_pool = False + return f + + @classmethod + def compute_storage_size(cls, + block_size, + block_count, + header_data=None, + ignore_header=False): + assert (block_size > 0) and (block_size == int(block_size)) + assert (block_count > 0) and (block_count == int(block_count)) + if header_data is None: + header_data = bytes() + if ignore_header: + return block_size * block_count + else: + return BlockStorageFile._index_offset + \ + len(header_data) + \ + block_size * block_count + + @classmethod + def setup(cls, + storage_name, + block_size, + block_count, + initialize=None, + header_data=None, + ignore_existing=False, + threadpool_size=None, + _filesystem=default_filesystem): + + if (not ignore_existing): + _exists = True + try: + _filesystem.stat(storage_name) + except OSError as e: + if e.errno == errno.ENOENT: + _exists = False + if _exists: + raise IOError( + "Storage location already exists: %s" + % (storage_name)) + if (block_size <= 0) or (block_size != int(block_size)): + raise ValueError( + "Block size (bytes) must be a positive integer: %s" + % (block_size)) + if (block_count <= 0) or (block_count != int(block_count)): + raise ValueError( + "Block count must be a positive integer: %s" + % (block_count)) + if (header_data is not None) and \ + (type(header_data) is not bytes): + raise TypeError( + "'header_data' must be of type bytes. " + "Invalid type: %s" % (type(header_data))) + + if initialize is None: + zeros = bytes(bytearray(block_size)) + initialize = lambda i: zeros + try: + with _filesystem.open(storage_name, "wb") as f: + # create_index + if header_data is None: + f.write(struct.pack(BlockStorageFile._index_struct_string, + block_size, + block_count, + 0, + False)) + else: + f.write(struct.pack(BlockStorageFile._index_struct_string, + block_size, + block_count, + len(header_data), + False)) + f.write(header_data) + with tqdm.tqdm(total=block_count*block_size, + desc="Initializing File Block Storage Space", + unit="B", + unit_scale=True, + disable=not pyoram.config.SHOW_PROGRESS_BAR) as progress_bar: + for i in xrange(block_count): + block = initialize(i) + assert len(block) == block_size, \ + ("%s != %s" % (len(block), block_size)) + f.write(block) + progress_bar.update(n=block_size) + except: # pragma: no cover + _filesystem.remove(storage_name) # pragma: no cover + raise # pragma: no cover + + return BlockStorageFile(storage_name, + threadpool_size=threadpool_size, + _filesystem=_filesystem) + + @property + def header_data(self): + return self._user_header_data + + @property + def block_count(self): + return self._block_count + + @property + def block_size(self): + return self._block_size + + @property + def storage_name(self): + return self._storage_name + + def update_header_data(self, new_header_data): + self._check_async() + if len(new_header_data) != len(self.header_data): + raise ValueError( + "The size of header data can not change.\n" + "Original bytes: %s\n" + "New bytes: %s" % (len(self.header_data), + len(new_header_data))) + self._user_header_data = bytes(new_header_data) + + self._timer.startTimer(); + self._f.seek(BlockStorageFile._index_offset) + self._f.write(self._user_header_data) + self._timer.endTimer(); + + def close(self): + print("file close 1") + self._prep_for_close() + print("file close 2") + if self._f is not None: + try: + self._f.close() + except OSError: # pragma: no cover + pass # pragma: no cover + self._f = None + + def read_blocks(self, indices): + # print("Reading Blocks ......"); + self._check_async() + blocks = [] + for i in indices: + assert 0 <= i < self.block_count + self._bytes_received += self.block_size + + self._timer.startTimer(); + self._f.seek(self._header_offset + i * self.block_size) + a = self._f.read(self.block_size) + # self._f.flush() + self._timer.endTimer(); + + blocks.append(a) + return blocks + + def yield_blocks(self, indices): + # print("Yielding Blocks ......"); + self._check_async() + for i in indices: + assert 0 <= i < self.block_count + self._bytes_received += self.block_size + + self._timer.startTimer(); + self._f.seek(self._header_offset + i * self.block_size) + a = self._f.read(self.block_size) + # self._f.flush() + self._timer.endTimer(); + + yield a + + def read_block(self, i): + # print("Reading Block ......"); + self._check_async() + assert 0 <= i < self.block_count + self._bytes_received += self.block_size + + self._timer.startTimer(); + self._f.seek(self._header_offset + i * self.block_size) + a = self._f.read(self.block_size) + # self._f.flush() + self._timer.endTimer(); + + return a + + def write_blocks(self, indices, blocks, callback=None): + self._check_async() + chunks = [] + for i, block in zip(indices, blocks): + assert 0 <= i < self.block_count + assert len(block) == self.block_size, \ + ("%s != %s" % (len(block), self.block_size)) + self._bytes_sent += self.block_size + chunks.append((i, block)) + + self._schedule_async_write(chunks, callback=callback) + + + + def write_block(self, i, block): + self._check_async() + assert 0 <= i < self.block_count + assert len(block) == self.block_size + self._bytes_sent += self.block_size + + self._schedule_async_write(((i, block),)) + + @property + def bytes_sent(self): + return self._bytes_sent + + @property + def bytes_received(self): + return self._bytes_received + +BlockStorageTypeFactory.register_device("file", BlockStorageFile) diff --git a/PyORAM/src/pyoram/storage/block_storage_mmap.py b/PyORAM/src/pyoram/storage/block_storage_mmap.py new file mode 100644 index 0000000..abc5744 --- /dev/null +++ b/PyORAM/src/pyoram/storage/block_storage_mmap.py @@ -0,0 +1,143 @@ +__all__ = ('BlockStorageMMap',) + +import logging +import mmap + +from pyoram.storage.block_storage import \ + BlockStorageTypeFactory +from pyoram.storage.block_storage_file import \ + BlockStorageFile + +log = logging.getLogger("pyoram") + +class _BlockStorageMemoryImpl(object): + """ + This class implementents the BlockStorageInterface read/write + methods for classes with a private attribute _f that can be + accessed using __getslice__/__setslice__ notation. + """ + + def read_blocks(self, indices): + blocks = [] + for i in indices: + assert 0 <= i < self.block_count + self._bytes_received += self.block_size + pos_start = self._header_offset + i * self.block_size + pos_stop = pos_start + self.block_size + blocks.append(self._f[pos_start:pos_stop]) + return blocks + + def yield_blocks(self, indices): + for i in indices: + assert 0 <= i < self.block_count + self._bytes_received += self.block_size + pos_start = self._header_offset + i * self.block_size + pos_stop = pos_start + self.block_size + yield self._f[pos_start:pos_stop] + + def read_block(self, i): + assert 0 <= i < self.block_count + self._bytes_received += self.block_size + pos_start = self._header_offset + i * self.block_size + pos_stop = pos_start + self.block_size + return self._f[pos_start:pos_stop] + + def write_blocks(self, indices, blocks, callback=None): + for i, block in zip(indices, blocks): + assert 0 <= i < self.block_count + self._bytes_sent += self.block_size + pos_start = self._header_offset + i * self.block_size + pos_stop = pos_start + self.block_size + self._f[pos_start:pos_stop] = block + if callback is not None: + callback(i) + + def write_block(self, i, block): + assert 0 <= i < self.block_count + self._bytes_sent += self.block_size + pos_start = self._header_offset + i * self.block_size + pos_stop = pos_start + self.block_size + self._f[pos_start:pos_stop] = block + +class BlockStorageMMap(_BlockStorageMemoryImpl, + BlockStorageFile): + """ + A class implementing the block storage interface by creating a + memory map over a local file. This class uses the same storage + format as BlockStorageFile. Thus, a block storage space can be + created using this class and then, after saving the raw storage + data to disk, reopened with any other class compatible with + BlockStorageFile (and visa versa). + """ + + def __init__(self, *args, **kwds): + mm = kwds.pop('mm', None) + self._mmap_owned = True + super(BlockStorageMMap, self).__init__(*args, **kwds) + if mm is None: + self._f.flush() + mm = mmap.mmap(self._f.fileno(), 0) + else: + self._mmap_owned = False + self._f.close() + self._f = mm + + # + # Define BlockStorageInterface Methods + # (override what is defined on BlockStorageFile) + # + + #@classmethod + #def compute_storage_size(...) + + def clone_device(self): + f = BlockStorageMMap(self.storage_name, + threadpool_size=0, + mm=self._f, + ignore_lock=True) + f._pool = self._pool + f._close_pool = False + return f + + @classmethod + def setup(cls, + storage_name, + block_size, + block_count, + **kwds): + f = BlockStorageFile.setup(storage_name, + block_size, + block_count, + **kwds) + f.close() + return BlockStorageMMap(storage_name) + + #def update_header_data(...) + + def close(self): + self._prep_for_close() + if self._f is not None: + if self._mmap_owned: + try: + self._f.close() + except OSError: # pragma: no cover + pass # pragma: no cover + self._f = None + + #def read_blocks(...) + + #def yield_blocks(...) + + #def read_block(...) + + #def write_blocks(...) + + #def write_block(...) + + #@property + #def bytes_sent(...) + + #@property + #def bytes_received(...) + +BlockStorageTypeFactory.register_device("mmap", BlockStorageMMap) diff --git a/PyORAM/src/pyoram/storage/block_storage_ram.py b/PyORAM/src/pyoram/storage/block_storage_ram.py new file mode 100644 index 0000000..eb75599 --- /dev/null +++ b/PyORAM/src/pyoram/storage/block_storage_ram.py @@ -0,0 +1,314 @@ +__all__ = ('BlockStorageRAM',) + +import os +import struct +import logging +import errno +from multiprocessing.pool import ThreadPool + +import pyoram +from pyoram.storage.block_storage import \ + (BlockStorageInterface, + BlockStorageTypeFactory) +from pyoram.storage.block_storage_mmap import \ + (BlockStorageMMap, + _BlockStorageMemoryImpl) + +import tqdm +import six +from six.moves import xrange + +log = logging.getLogger("pyoram") + +class BlockStorageRAM(_BlockStorageMemoryImpl, + BlockStorageInterface): + """ + A class implementing the block storage interface where all data is + kept in RAM. This class uses the same storage format as + BlockStorageFile. Thus, a block storage space can be created using + this class and then, after saving the raw storage data to disk, + reopened with any other class compatible with BlockStorageFile + (and visa versa). + """ + + _index_struct_string = BlockStorageMMap._index_struct_string + _index_offset = struct.calcsize(_index_struct_string) + + def __init__(self, + storage_data, + threadpool_size=None, + ignore_lock=False): + + self._bytes_sent = 0 + self._bytes_received = 0 + self._ignore_lock = ignore_lock + self._f = None + self._pool = None + self._close_pool = True + if type(storage_data) is not bytearray: + raise TypeError( + "BlockStorageRAM requires input argument of type " + "'bytearray'. Invalid input type: %s" + % (type(storage_data))) + self._f = storage_data + self._block_size, self._block_count, user_header_size, locked = \ + struct.unpack( + BlockStorageRAM._index_struct_string, + self._f[:BlockStorageRAM._index_offset]) + + if locked and (not self._ignore_lock): + raise IOError( + "Can not open block storage device because it is " + "locked by another process. To ignore this check, " + "initialize this class with the keyword 'ignore_lock' " + "set to True.") + self._user_header_data = bytes() + if user_header_size > 0: + self._user_header_data = \ + bytes(self._f[BlockStorageRAM._index_offset:\ + (BlockStorageRAM._index_offset+user_header_size)]) + assert len(self._user_header_data) == user_header_size + self._header_offset = BlockStorageRAM._index_offset + \ + len(self._user_header_data) + + if not self._ignore_lock: + # turn on the locked flag + self._f[:BlockStorageRAM._index_offset] = \ + struct.pack(BlockStorageRAM._index_struct_string, + self.block_size, + self.block_count, + len(self._user_header_data), + True) + + # Although we do not use the threadpool we still + # create just in case we are the first + if threadpool_size != 0: + self._pool = ThreadPool(threadpool_size) + + # + # Add some methods specific to BlockStorageRAM + # + + @staticmethod + def fromfile(file_, + threadpool_size=None, + ignore_lock=False): + """ + Instantiate BlockStorageRAM device from a file saved in block + storage format. The file_ argument can be a file object or a + string that represents a filename. If called with a file + object, it should be opened in binary mode, and the caller is + responsible for closing the file. + + This method returns a BlockStorageRAM instance. + """ + close_file = False + if not hasattr(file_, 'read'): + file_ = open(file_, 'rb') + close_file = True + try: + header_data = file_.read(BlockStorageRAM._index_offset) + block_size, block_count, user_header_size, locked = \ + struct.unpack( + BlockStorageRAM._index_struct_string, + header_data) + if locked and (not ignore_lock): + raise IOError( + "Can not open block storage device because it is " + "locked by another process. To ignore this check, " + "call this method with the keyword 'ignore_lock' " + "set to True.") + header_offset = len(header_data) + \ + user_header_size + f = bytearray(header_offset + \ + (block_size * block_count)) + f[:header_offset] = header_data + file_.read(user_header_size) + f[header_offset:] = file_.read(block_size * block_count) + finally: + if close_file: + file_.close() + + return BlockStorageRAM(f, + threadpool_size=threadpool_size, + ignore_lock=ignore_lock) + + def tofile(self, file_): + """ + Dump all storage data to a file. The file_ argument can be a + file object or a string that represents a filename. If called + with a file object, it should be opened in binary mode, and + the caller is responsible for closing the file. + + The method should only be called after the storage device has + been closed to ensure that the locked flag has been set to + False. + """ + close_file = False + if not hasattr(file_, 'write'): + file_ = open(file_, 'wb') + close_file = True + file_.write(self._f) + if close_file: + file_.close() + + @property + def data(self): + """Access the raw bytearray""" + return self._f + + # + # Define BlockStorageInterface Methods + # + + def clone_device(self): + f = BlockStorageRAM(self._f, + threadpool_size=0, + ignore_lock=True) + f._pool = self._pool + f._close_pool = False + return f + + @classmethod + def compute_storage_size(cls, *args, **kwds): + return BlockStorageMMap.compute_storage_size(*args, **kwds) + + @classmethod + def setup(cls, + storage_name, + block_size, + block_count, + initialize=None, + header_data=None, + ignore_existing=False, + threadpool_size=None): + + # We ignore the 'storage_name' argument + # We ignore the 'ignore_existing' flag + if (block_size <= 0) or (block_size != int(block_size)): + raise ValueError( + "Block size (bytes) must be a positive integer: %s" + % (block_size)) + if (block_count <= 0) or (block_count != int(block_count)): + raise ValueError( + "Block count must be a positive integer: %s" + % (block_count)) + if (header_data is not None) and \ + (type(header_data) is not bytes): + raise TypeError( + "'header_data' must be of type bytes. " + "Invalid type: %s" % (type(header_data))) + + if initialize is None: + zeros = bytes(bytearray(block_size)) + initialize = lambda i: zeros + + # create_index + index_data = None + if header_data is None: + index_data = struct.pack(BlockStorageRAM._index_struct_string, + block_size, + block_count, + 0, + False) + header_data = bytes() + else: + index_data = struct.pack(BlockStorageRAM._index_struct_string, + block_size, + block_count, + len(header_data), + False) + header_offset = len(index_data) + len(header_data) + f = bytearray(header_offset + \ + (block_size * block_count)) + f[:header_offset] = index_data + header_data + progress_bar = tqdm.tqdm(total=block_count*block_size, + desc="Initializing File Block Storage Space", + unit="B", + unit_scale=True, + disable=not pyoram.config.SHOW_PROGRESS_BAR) + for i in xrange(block_count): + block = initialize(i) + assert len(block) == block_size, \ + ("%s != %s" % (len(block), block_size)) + pos_start = header_offset + i * block_size + pos_start = header_offset + i * block_size + pos_stop = pos_start + block_size + f[pos_start:pos_stop] = block[:] + progress_bar.update(n=block_size) + progress_bar.close() + + return BlockStorageRAM(f, threadpool_size=threadpool_size) + + @property + def header_data(self): + return self._user_header_data + + @property + def block_count(self): + return self._block_count + + @property + def block_size(self): + return self._block_size + + @property + def storage_name(self): + return None + + def update_header_data(self, new_header_data): + if len(new_header_data) != len(self.header_data): + raise ValueError( + "The size of header data can not change.\n" + "Original bytes: %s\n" + "New bytes: %s" % (len(self.header_data), + len(new_header_data))) + self._user_header_data = bytes(new_header_data) + self._f[BlockStorageRAM._index_offset:\ + (BlockStorageRAM._index_offset+len(new_header_data))] = \ + self._user_header_data + + def close(self): + if self._close_pool and (self._pool is not None): + self._pool.close() + self._pool.join() + self._pool = None + if not self._ignore_lock: + # turn off the locked flag + self._f[:BlockStorageRAM._index_offset] = \ + struct.pack(BlockStorageRAM._index_struct_string, + self.block_size, + self.block_count, + len(self._user_header_data), + False) + self._ignore_lock = True + + # + # We must cast from bytearray to bytes + # when reading from a bytearray so that this + # class works with the encryption layer. + # + + def read_blocks(self, indices): + return [bytes(block) for block + in super(BlockStorageRAM, self).read_blocks(indices)] + + def yield_blocks(self, indices): + for block in super(BlockStorageRAM, self).yield_blocks(indices): + yield bytes(block) + + def read_block(self, i): + return bytes(super(BlockStorageRAM, self).read_block(i)) + + #def write_blocks(...) + + #def write_block(...) + + @property + def bytes_sent(self): + return self._bytes_sent + + @property + def bytes_received(self): + return self._bytes_received + +BlockStorageTypeFactory.register_device("ram", BlockStorageRAM) diff --git a/PyORAM/src/pyoram/storage/block_storage_s3.py b/PyORAM/src/pyoram/storage/block_storage_s3.py new file mode 100644 index 0000000..d35f0ab --- /dev/null +++ b/PyORAM/src/pyoram/storage/block_storage_s3.py @@ -0,0 +1,396 @@ +__all__ = ('BlockStorageS3',) + +import struct +import logging +from multiprocessing.pool import ThreadPool + +import pyoram +from pyoram.storage.block_storage import \ + (BlockStorageInterface, + BlockStorageTypeFactory) +from pyoram.storage.boto3_s3_wrapper import Boto3S3Wrapper + +import tqdm +import six +from six.moves import xrange, map + +log = logging.getLogger("pyoram") + +class BlockStorageS3(BlockStorageInterface): + """ + A block storage device for Amazon Simple + Storage Service (S3). + """ + + _index_name = "PyORAMBlockStorageS3_index.bin" + _index_struct_string = "!LLL?" + _index_offset = struct.calcsize(_index_struct_string) + + def __init__(self, + storage_name, + bucket_name=None, + aws_access_key_id=None, + aws_secret_access_key=None, + region_name=None, + ignore_lock=False, + threadpool_size=None, + s3_wrapper=Boto3S3Wrapper): + + self._bytes_sent = 0 + self._bytes_received = 0 + self._storage_name = storage_name + self._bucket_name = bucket_name + self._aws_access_key_id = aws_access_key_id + self._aws_secret_access_key = aws_secret_access_key + self._region_name = region_name + self._pool = None + self._close_pool = True + self._s3 = None + self._ignore_lock = ignore_lock + self._async_write = None + self._async_write_callback = None + + if bucket_name is None: + raise ValueError("'bucket_name' keyword is required") + + if threadpool_size != 0: + self._pool = ThreadPool(threadpool_size) + + self._s3 = s3_wrapper(bucket_name, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + region_name=region_name) + self._basename = self.storage_name+"/b%d" + + index_data = self._s3.download( + self._storage_name+"/"+BlockStorageS3._index_name) + self._block_size, self._block_count, user_header_size, locked = \ + struct.unpack( + BlockStorageS3._index_struct_string, + index_data[:BlockStorageS3._index_offset]) + if locked and (not self._ignore_lock): + raise IOError( + "Can not open block storage device because it is " + "locked by another process. To ignore this check, " + "initialize this class with the keyword 'ignore_lock' " + "set to True.") + self._user_header_data = bytes() + if user_header_size > 0: + self._user_header_data = \ + index_data[BlockStorageS3._index_offset: + (BlockStorageS3._index_offset+user_header_size)] + + if not self._ignore_lock: + # turn on the locked flag + self._s3.upload((self._storage_name+"/"+BlockStorageS3._index_name, + struct.pack(BlockStorageS3._index_struct_string, + self.block_size, + self.block_count, + len(self.header_data), + True) + \ + self.header_data)) + + def _check_async(self): + if self._async_write is not None: + for i in self._async_write: + if self._async_write_callback is not None: + self._async_write_callback(i) + self._async_write = None + self._async_write_callback = None + + def _schedule_async_write(self, arglist, callback=None): + assert self._async_write is None + if self._pool is not None: + self._async_write = \ + self._pool.imap_unordered(self._s3.upload, arglist) + self._async_write_callback = callback + else: + # Note: we are using six.map which always + # behaves like imap + for i in map(self._s3.upload, arglist): + if callback is not None: + callback(i) + + def _download(self, i): + return self._s3.download(self._basename % i) + + # + # Define BlockStorageInterface Methods + # + + def clone_device(self): + f = BlockStorageS3(self.storage_name, + bucket_name=self._bucket_name, + aws_access_key_id=self._aws_access_key_id, + aws_secret_access_key=self._aws_secret_access_key, + region_name=self._region_name, + threadpool_size=0, + s3_wrapper=type(self._s3), + ignore_lock=True) + f._pool = self._pool + f._close_pool = False + return f + + @classmethod + def compute_storage_size(cls, + block_size, + block_count, + header_data=None, + ignore_header=False): + assert (block_size > 0) and (block_size == int(block_size)) + assert (block_count > 0) and (block_count == int(block_count)) + if header_data is None: + header_data = bytes() + if ignore_header: + return block_size * block_count + else: + return BlockStorageS3._index_offset + \ + len(header_data) + \ + block_size * block_count + + @classmethod + def setup(cls, + storage_name, + block_size, + block_count, + bucket_name=None, + aws_access_key_id=None, + aws_secret_access_key=None, + region_name=None, + header_data=None, + initialize=None, + threadpool_size=None, + ignore_existing=False, + s3_wrapper=Boto3S3Wrapper): + + if bucket_name is None: + raise ValueError("'bucket_name' is required") + if (block_size <= 0) or (block_size != int(block_size)): + raise ValueError( + "Block size (bytes) must be a positive integer: %s" + % (block_size)) + if (block_count <= 0) or (block_count != int(block_count)): + raise ValueError( + "Block count must be a positive integer: %s" + % (block_count)) + if (header_data is not None) and \ + (type(header_data) is not bytes): + raise TypeError( + "'header_data' must be of type bytes. " + "Invalid type: %s" % (type(header_data))) + + pool = None + if threadpool_size != 0: + pool = ThreadPool(threadpool_size) + + s3 = s3_wrapper(bucket_name, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + region_name=region_name) + exists = s3.exists(storage_name) + if (not ignore_existing) and exists: + raise IOError( + "Storage location already exists in bucket %s: %s" + % (bucket_name, storage_name)) + if exists: + log.info("Deleting objects in existing S3 entry: %s/%s" + % (bucket_name, storage_name)) + print("Clearing Existing S3 Objects With Prefix %s/%s/" + % (bucket_name, storage_name)) + s3.clear(storage_name, threadpool=pool) + + if header_data is None: + s3.upload((storage_name+"/"+BlockStorageS3._index_name, + struct.pack(BlockStorageS3._index_struct_string, + block_size, + block_count, + 0, + False))) + else: + s3.upload((storage_name+"/"+BlockStorageS3._index_name, + struct.pack(BlockStorageS3._index_struct_string, + block_size, + block_count, + len(header_data), + False) + \ + header_data)) + + if initialize is None: + zeros = bytes(bytearray(block_size)) + initialize = lambda i: zeros + basename = storage_name+"/b%d" + # NOTE: We will not be informed when a thread + # encounters an exception (e.g., when + # calling initialize(i). We must ensure + # that all iterations were processed + # by counting the results. + def init_blocks(): + for i in xrange(block_count): + yield (basename % i, initialize(i)) + def _do_upload(arg): + try: + s3.upload(arg) + except Exception as e: # pragma: no cover + log.error( # pragma: no cover + "An exception occured during S3 " # pragma: no cover + "setup when calling the block " # pragma: no cover + "initialization function: %s" # pragma: no cover + % (str(e))) # pragma: no cover + raise # pragma: no cover + total = None + progress_bar = tqdm.tqdm(total=block_count*block_size, + desc="Initializing S3 Block Storage Space", + unit="B", + unit_scale=True, + disable=not pyoram.config.SHOW_PROGRESS_BAR) + if pool is not None: + try: + for i,_ in enumerate( + pool.imap_unordered(_do_upload, init_blocks())): + total = i + progress_bar.update(n=block_size) + except Exception as e: # pragma: no cover + s3.clear(storage_name) # pragma: no cover + raise # pragma: no cover + finally: + progress_bar.close() + pool.close() + pool.join() + else: + try: + for i,_ in enumerate( + map(s3.upload, init_blocks())): + total = i + progress_bar.update(n=block_size) + except Exception as e: # pragma: no cover + s3.clear(storage_name) # pragma: no cover + raise # pragma: no cover + finally: + progress_bar.close() + + if total != block_count - 1: + s3.clear(storage_name) # pragma: no cover + if pool is not None: # pragma: no cover + pool.close() # pragma: no cover + pool.join() # pragma: no cover + raise ValueError( # pragma: no cover + "Something went wrong during S3 block" # pragma: no cover + " initialization. Check the logger " # pragma: no cover + "output for more information.") # pragma: no cover + + return BlockStorageS3(storage_name, + bucket_name=bucket_name, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + region_name=region_name, + threadpool_size=threadpool_size, + s3_wrapper=s3_wrapper) + + @property + def header_data(self): + return self._user_header_data + + @property + def block_count(self): + return self._block_count + + @property + def block_size(self): + return self._block_size + + @property + def storage_name(self): + return self._storage_name + + def update_header_data(self, new_header_data): + self._check_async() + if len(new_header_data) != len(self.header_data): + raise ValueError( + "The size of header data can not change.\n" + "Original bytes: %s\n" + "New bytes: %s" % (len(self.header_data), + len(new_header_data))) + self._user_header_data = new_header_data + + index_data = bytearray(self._s3.download( + self._storage_name+"/"+BlockStorageS3._index_name)) + lenbefore = len(index_data) + index_data[BlockStorageS3._index_offset:] = new_header_data + assert lenbefore == len(index_data) + self._s3.upload((self._storage_name+"/"+BlockStorageS3._index_name, + bytes(index_data))) + + def close(self): + self._check_async() + if self._s3 is not None: + if not self._ignore_lock: + # turn off the locked flag + self._s3.upload( + (self._storage_name+"/"+BlockStorageS3._index_name, + struct.pack(BlockStorageS3._index_struct_string, + self.block_size, + self.block_count, + len(self.header_data), + False) + \ + self.header_data)) + if self._close_pool and (self._pool is not None): + self._pool.close() + self._pool.join() + self._pool = None + + def read_blocks(self, indices): + self._check_async() + # be sure not to exhaust this if it is an iterator + # or generator + indices = list(indices) + assert all(0 <= i <= self.block_count for i in indices) + self._bytes_received += self.block_size * len(indices) + if self._pool is not None: + return self._pool.map(self._download, indices) + else: + return list(map(self._download, indices)) + + def yield_blocks(self, indices): + self._check_async() + # be sure not to exhaust this if it is an iterator + # or generator + indices = list(indices) + assert all(0 <= i <= self.block_count for i in indices) + self._bytes_received += self.block_size * len(indices) + if self._pool is not None: + return self._pool.imap(self._download, indices) + else: + return map(self._download, indices) + + def read_block(self, i): + self._check_async() + assert 0 <= i < self.block_count + self._bytes_received += self.block_size + return self._download(i) + + def write_blocks(self, indices, blocks, callback=None): + self._check_async() + # be sure not to exhaust this if it is an iterator + # or generator + indices = list(indices) + assert all(0 <= i <= self.block_count for i in indices) + self._bytes_sent += self.block_size * len(indices) + indices = (self._basename % i for i in indices) + self._schedule_async_write(zip(indices, blocks), + callback=callback) + + def write_block(self, i, block): + self._check_async() + assert 0 <= i < self.block_count + self._bytes_sent += self.block_size + self._schedule_async_write((((self._basename % i), block),)) + + @property + def bytes_sent(self): + return self._bytes_sent + + @property + def bytes_received(self): + return self._bytes_received + +BlockStorageTypeFactory.register_device("s3", BlockStorageS3) diff --git a/PyORAM/src/pyoram/storage/block_storage_sftp.py b/PyORAM/src/pyoram/storage/block_storage_sftp.py new file mode 100644 index 0000000..b6dab16 --- /dev/null +++ b/PyORAM/src/pyoram/storage/block_storage_sftp.py @@ -0,0 +1,151 @@ +__all__ = ('BlockStorageSFTP',) + +import logging +from AliTimer import * +from pyoram.util.misc import chunkiter +from pyoram.storage.block_storage import \ + BlockStorageTypeFactory +from pyoram.storage.block_storage_file import \ + BlockStorageFile + +import time + +log = logging.getLogger("pyoram") + +class BlockStorageSFTP(BlockStorageFile): + """ + A block storage device for accessing file data through + an SSH portal using Secure File Transfer Protocol (SFTP). + """ + + def __init__(self, + storage_name, + sshclient=None, + **kwds): + if sshclient is None: + raise ValueError( + "Can not open sftp block storage device " + "without an ssh client.") + super(BlockStorageSFTP, self).__init__( + storage_name, + _filesystem=sshclient.open_sftp(), + **kwds) + self._sshclient = sshclient + self._f.set_pipelined() + self._timer = Foo.Instance(); + + + + # + # Define BlockStorageInterface Methods + # + + def clone_device(self): + f = BlockStorageSFTP(self.storage_name, + sshclient=self._sshclient, + threadpool_size=0, + ignore_lock=True) + f._pool = self._pool + f._close_pool = False + return f + + #@classmethod + #def compute_storage_size(...) + + @classmethod + def setup(cls, + storage_name, + block_size, + block_count, + sshclient=None, + threadpool_size=None, + **kwds): + if sshclient is None: + raise ValueError( + "Can not setup sftp block storage device " + "without an ssh client.") + + with BlockStorageFile.setup(storage_name, + block_size, + block_count, + _filesystem=sshclient.open_sftp(), + threadpool_size=threadpool_size, + **kwds) as f: + pass + f._filesystem.close() + + return BlockStorageSFTP(storage_name, + sshclient=sshclient, + threadpool_size=threadpool_size) + + #@property + #def header_data(...) + + #@property + #def block_count(...) + + #@property + #def block_size(...) + + #@property + #def storage_name(...) + + #def update_header_data(...) + + def close(self): + print("sftp close 1") + super(BlockStorageSFTP, self).close() + print("sftp close 2") + self._filesystem.close() + print("sftp close 3") + + def read_blocks(self, indices): + self._check_async() + args = [] + for i in indices: + assert 0 <= i < self.block_count + self._bytes_received += self.block_size + args.append((self._header_offset + i * self.block_size, + self.block_size)) + + + sTime = time.time(); + self._timer.startTimer(); + a = self._f.readv(args) + self._timer.endTimer(); + + + # print("Reading Blocks SFTP......" + str(time.time() - sTime)); + return a + + def yield_blocks(self, indices, chunksize=100): + for chunk in chunkiter(indices, n=chunksize): + assert all(0 <= i <= self.block_count for i in chunk) + self._bytes_received += self.block_size * len(chunk) + args = [(self._header_offset + i * self.block_size, + self.block_size) + for i in chunk] + + + self._timer.startTimer(); + a = self._f.readv(args) + self._timer.endTimer(); + + # print("Yield SFTP......"); + + for block in a: + yield block + + #def read_block(...) + + #def write_blocks(...) + + #def write_block(...) + + #@property + #def bytes_sent(...) + + #@property + #def bytes_received(...) + +BlockStorageTypeFactory.register_device("sftp", BlockStorageSFTP) diff --git a/PyORAM/src/pyoram/storage/boto3_s3_wrapper.py b/PyORAM/src/pyoram/storage/boto3_s3_wrapper.py new file mode 100644 index 0000000..922c618 --- /dev/null +++ b/PyORAM/src/pyoram/storage/boto3_s3_wrapper.py @@ -0,0 +1,158 @@ +__all__ = ("Boto3S3Wrapper", + "MockBoto3S3Wrapper") +import os +import shutil + +import pyoram + +import tqdm +try: + import boto3 + import botocore + boto3_available = True +except: # pragma: no cover + boto3_available = False # pragma: no cover + +import six +from six.moves import xrange, map + +class Boto3S3Wrapper(object): + """ + A wrapper class for the boto3 S3 service. + """ + + def __init__(self, + bucket_name, + aws_access_key_id=None, + aws_secret_access_key=None, + region_name=None): + if not boto3_available: + raise ImportError( # pragma: no cover + "boto3 module is required to " # pragma: no cover + "use BlockStorageS3 device") # pragma: no cover + + self._s3 = boto3.session.Session( + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + region_name=region_name).resource('s3') + self._bucket = self._s3.Bucket(bucket_name) + + def exists(self, key): + try: + self._bucket.Object(key).load() + except botocore.exceptions.ClientError as e: + if e.response['Error']['Code'] == "404": + pass + else: + raise e + else: + return True + # It's not a file. Check if it's a "directory". + for obj in self._bucket.objects.filter( + Prefix=key+"/"): + return True + return False + + def download(self, key): + try: + return self._s3.meta.client.get_object( + Bucket=self._bucket.name, + Key=key)['Body'].read() + except botocore.exceptions.ClientError: + raise IOError("Can not download key: %s" + % (key)) + + def upload(self, key_block): + key, block = key_block + self._bucket.put_object(Key=key, Body=block) + + # Chunk a streamed iterator of which we do not know + # the size + def _chunks(self, objs, n=100): + assert 1 <= n <= 1000 # required by boto3 + objs = iter(objs) + try: + while (1): + chunk = [] + while len(chunk) < n: + chunk.append({'Key': six.next(objs).key}) + yield {'Objects': chunk} + except StopIteration: + pass + if len(chunk): + yield {'Objects': chunk} + + def _del(self, chunk): + self._bucket.delete_objects(Delete=chunk) + return len(chunk['Objects']) + + def clear(self, key, threadpool=None): + objs = self._bucket.objects.filter(Prefix=key+"/") + if threadpool is not None: + deliter = threadpool.imap(self._del, self._chunks(objs)) + else: + deliter = map(self._del, self._chunks(objs)) + with tqdm.tqdm(total=None, #len(objs), + desc="Clearing S3 Blocks", + unit=" objects", + disable=not pyoram.config.SHOW_PROGRESS_BAR) as progress_bar: + progress_bar.update(n=0) + for chunksize in deliter: + progress_bar.update(n=chunksize) + +class MockBoto3S3Wrapper(object): + """ + A mock class for Boto3S3Wrapper that uses the local filesystem and + treats the bucket name as a directory. + + This class is mainly used for testing, but could potentially be + used to setup storage locally that is then uploaded to S3 through + the AWS web portal. + """ + + def __init__(self, + bucket_name, + aws_access_key_id=None, + aws_secret_access_key=None, + region_name=None): + + self._bucket_name = os.path.abspath( + os.path.normpath(bucket_name)) + + # called within upload to create directory + # heirarchy on the fly + def _makedirs_if_needed(self, key): + if not os.path.exists( + os.path.dirname( + os.path.join(self._bucket_name, key))): + os.makedirs( + os.path.dirname( + os.path.join(self._bucket_name, key))) + assert not os.path.isdir( + os.path.join(self._bucket_name, key)) + + def exists(self, key): + return os.path.exists( + os.path.join(self._bucket_name, key)) + + def download(self, key): + with open(os.path.join(self._bucket_name, key), 'rb') as f: + return f.read() + + def upload(self, key_block): + key, block = key_block + self._makedirs_if_needed(key) + with open(os.path.join(self._bucket_name, key), 'wb') as f: + f.write(block) + + def clear(self, key, threadpool=None): + if os.path.exists( + os.path.join(self._bucket_name, key)): + if os.path.isdir( + os.path.join(self._bucket_name, key)): + shutil.rmtree( + os.path.join(self._bucket_name, key), + ignore_errors=True) + else: + os.remove( + os.path.join(self._bucket_name, key)) diff --git a/PyORAM/src/pyoram/storage/heap_storage.py b/PyORAM/src/pyoram/storage/heap_storage.py new file mode 100644 index 0000000..46059c2 --- /dev/null +++ b/PyORAM/src/pyoram/storage/heap_storage.py @@ -0,0 +1,245 @@ +__all__ = ('HeapStorage',) + +import struct + +from pyoram.util.virtual_heap import SizedVirtualHeap +from pyoram.storage.block_storage import (BlockStorageInterface, + BlockStorageTypeFactory) + +class HeapStorageInterface(object): + + def __enter__(self): + return self + def __exit__(self, *args): + self.close() + + # + # Abstract Interface + # + + def clone_device(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + + @classmethod + def compute_storage_size(cls, *args, **kwds): + raise NotImplementedError # pragma: no cover + @classmethod + def setup(cls, *args, **kwds): + raise NotImplementedError # pragma: no cover + + @property + def header_data(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + @property + def bucket_count(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + @property + def bucket_size(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + @property + def blocks_per_bucket(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + @property + def storage_name(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + @property + def virtual_heap(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + @property + def bucket_storage(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + + def update_header_data(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + def close(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + def read_path(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + def write_path(self, *args, **kwds): + raise NotImplementedError # pragma: no cover + + @property + def bytes_sent(self): + raise NotImplementedError # pragma: no cover + @property + def bytes_received(self): + raise NotImplementedError # pragma: no cover + +class HeapStorage(HeapStorageInterface): + + _header_struct_string = "!LLL" + _header_offset = struct.calcsize(_header_struct_string) + + def _new_storage(self, storage, **kwds): + storage_type = kwds.pop('storage_type', 'file') + + + def __init__(self, storage, **kwds): + if isinstance(storage, BlockStorageInterface): + self._storage = storage + if len(kwds): + raise ValueError( + "Keywords not used when initializing " + "with a storage device: %s" + % (str(kwds))) + else: + storage_type = kwds.pop('storage_type', 'file') + self._storage = BlockStorageTypeFactory(storage_type)\ + (storage, **kwds) + + heap_base, heap_height, blocks_per_bucket = \ + struct.unpack( + self._header_struct_string, + self._storage.header_data[:self._header_offset]) + self._vheap = SizedVirtualHeap( + heap_base, + heap_height, + blocks_per_bucket=blocks_per_bucket) + + # + # Define HeapStorageInterface Methods + # + + def clone_device(self): + return HeapStorage(self._storage.clone_device()) + + @classmethod + def compute_storage_size(cls, + block_size, + heap_height, + blocks_per_bucket=1, + heap_base=2, + ignore_header=False, + storage_type='file', + **kwds): + assert (block_size > 0) and (block_size == int(block_size)) + assert heap_height >= 0 + assert blocks_per_bucket >= 1 + assert heap_base >= 2 + assert 'block_count' not in kwds + vheap = SizedVirtualHeap( + heap_base, + heap_height, + blocks_per_bucket=blocks_per_bucket) + if ignore_header: + return BlockStorageTypeFactory(storage_type).\ + compute_storage_size( + vheap.blocks_per_bucket * block_size, + vheap.bucket_count(), + ignore_header=True, + **kwds) + else: + return cls._header_offset + \ + BlockStorageTypeFactory(storage_type).\ + compute_storage_size( + vheap.blocks_per_bucket * block_size, + vheap.bucket_count(), + ignore_header=False, + **kwds) + + @classmethod + def setup(cls, + storage_name, + block_size, + heap_height, + blocks_per_bucket=1, + heap_base=2, + storage_type='file', + **kwds): + if 'block_count' in kwds: + raise ValueError("'block_count' keyword is not accepted") + if heap_height < 0: + raise ValueError( + "heap height must be 0 or greater. Invalid value: %s" + % (heap_height)) + if blocks_per_bucket < 1: + raise ValueError( + "blocks_per_bucket must be 1 or greater. " + "Invalid value: %s" % (blocks_per_bucket)) + if heap_base < 2: + raise ValueError( + "heap base must be 2 or greater. Invalid value: %s" + % (heap_base)) + + vheap = SizedVirtualHeap( + heap_base, + heap_height, + blocks_per_bucket=blocks_per_bucket) + + user_header_data = kwds.pop('header_data', bytes()) + if type(user_header_data) is not bytes: + raise TypeError( + "'header_data' must be of type bytes. " + "Invalid type: %s" % (type(user_header_data))) + kwds['header_data'] = \ + struct.pack(cls._header_struct_string, + heap_base, + heap_height, + blocks_per_bucket) + \ + user_header_data + + return HeapStorage( + BlockStorageTypeFactory(storage_type).setup( + storage_name, + vheap.blocks_per_bucket * block_size, + vheap.bucket_count(), + **kwds)) + + @property + def header_data(self): + return self._storage.header_data[self._header_offset:] + + @property + def bucket_count(self): + return self._storage.block_count + + @property + def bucket_size(self): + return self._storage.block_size + + @property + def blocks_per_bucket(self): + return self._vheap.blocks_per_bucket + + @property + def storage_name(self): + return self._storage.storage_name + + @property + def virtual_heap(self): + return self._vheap + + @property + def bucket_storage(self): + return self._storage + + def update_header_data(self, new_header_data): + self._storage.update_header_data( + self._storage.header_data[:self._header_offset] + \ + new_header_data) + + def close(self): + print("Non enc heap close 1") + self._storage.close() + print("Non enc heap close 2") + + def read_path(self, b, level_start=0): + assert 0 <= b < self._vheap.bucket_count() + bucket_list = self._vheap.Node(b).bucket_path_from_root() + assert 0 <= level_start < len(bucket_list) + return self._storage.read_blocks(bucket_list[level_start:]) + + def write_path(self, b, buckets, level_start=0): + assert 0 <= b < self._vheap.bucket_count() + bucket_list = self._vheap.Node(b).bucket_path_from_root() + assert 0 <= level_start < len(bucket_list) + self._storage.write_blocks(bucket_list[level_start:], + buckets) + + @property + def bytes_sent(self): + return self._storage.bytes_sent + + @property + def bytes_received(self): + return self._storage.bytes_received diff --git a/PyORAM/src/pyoram/tests/__init__.py b/PyORAM/src/pyoram/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/PyORAM/src/pyoram/tests/baselines/exists.empty b/PyORAM/src/pyoram/tests/baselines/exists.empty new file mode 100644 index 0000000..e69de29 diff --git a/PyORAM/src/pyoram/tests/baselines/k200_h0_b1.dot b/PyORAM/src/pyoram/tests/baselines/k200_h0_b1.dot new file mode 100644 index 0000000..46a8c68 --- /dev/null +++ b/PyORAM/src/pyoram/tests/baselines/k200_h0_b1.dot @@ -0,0 +1,5 @@ +// Created by SizedVirtualHeap.write_as_dot(...) +digraph heaptree { +node [shape=record] + 0 [penwidth=1,label="(0, 0)}"]; +} diff --git a/PyORAM/src/pyoram/tests/baselines/k200_h0_b1_data.dot b/PyORAM/src/pyoram/tests/baselines/k200_h0_b1_data.dot new file mode 100644 index 0000000..de3a43b --- /dev/null +++ b/PyORAM/src/pyoram/tests/baselines/k200_h0_b1_data.dot @@ -0,0 +1,5 @@ +// Created by SizedVirtualHeap.write_as_dot(...) +digraph heaptree { +node [shape=record] + 0 [penwidth=1,label="{{0}}"]; +} diff --git a/PyORAM/src/pyoram/tests/baselines/k2_h3_b1.dot b/PyORAM/src/pyoram/tests/baselines/k2_h3_b1.dot new file mode 100644 index 0000000..3d92a16 --- /dev/null +++ b/PyORAM/src/pyoram/tests/baselines/k2_h3_b1.dot @@ -0,0 +1,33 @@ +// Created by SizedVirtualHeap.write_as_dot(...) +digraph heaptree { +node [shape=record] + 0 [penwidth=1,label="''}"]; + 1 [penwidth=1,label="'0'}"]; + 3 [penwidth=1,label="'00'}"]; + 7 [penwidth=1,label="'000'}"]; + 3 -> 7 ; + 8 [penwidth=1,label="'001'}"]; + 3 -> 8 ; + 1 -> 3 ; + 4 [penwidth=1,label="'01'}"]; + 9 [penwidth=1,label="'010'}"]; + 4 -> 9 ; + 10 [penwidth=1,label="'011'}"]; + 4 -> 10 ; + 1 -> 4 ; + 0 -> 1 ; + 2 [penwidth=1,label="'1'}"]; + 5 [penwidth=1,label="'10'}"]; + 11 [penwidth=1,label="'100'}"]; + 5 -> 11 ; + 12 [penwidth=1,label="'101'}"]; + 5 -> 12 ; + 2 -> 5 ; + 6 [penwidth=1,label="'11'}"]; + 13 [penwidth=1,label="'110'}"]; + 6 -> 13 ; + 14 [penwidth=1,label="'111'}"]; + 6 -> 14 ; + 2 -> 6 ; + 0 -> 2 ; +} diff --git a/PyORAM/src/pyoram/tests/baselines/k2_h3_b1_data.dot b/PyORAM/src/pyoram/tests/baselines/k2_h3_b1_data.dot new file mode 100644 index 0000000..3b5f5b5 --- /dev/null +++ b/PyORAM/src/pyoram/tests/baselines/k2_h3_b1_data.dot @@ -0,0 +1,33 @@ +// Created by SizedVirtualHeap.write_as_dot(...) +digraph heaptree { +node [shape=record] + 0 [penwidth=1,label="{{0}}"]; + 1 [penwidth=1,label="{{1}}"]; + 3 [penwidth=1,label="{{3}}"]; + 7 [penwidth=1,label="{{7}}"]; + 3 -> 7 ; + 8 [penwidth=1,label="{{8}}"]; + 3 -> 8 ; + 1 -> 3 ; + 4 [penwidth=1,label="{{4}}"]; + 9 [penwidth=1,label="{{9}}"]; + 4 -> 9 ; + 10 [penwidth=1,label="{{10}}"]; + 4 -> 10 ; + 1 -> 4 ; + 0 -> 1 ; + 2 [penwidth=1,label="{{2}}"]; + 5 [penwidth=1,label="{{5}}"]; + 11 [penwidth=1,label="{{11}}"]; + 5 -> 11 ; + 12 [penwidth=1,label="{{12}}"]; + 5 -> 12 ; + 2 -> 5 ; + 6 [penwidth=1,label="{{6}}"]; + 13 [penwidth=1,label="{{13}}"]; + 6 -> 13 ; + 14 [penwidth=1,label="{{14}}"]; + 6 -> 14 ; + 2 -> 6 ; + 0 -> 2 ; +} diff --git a/PyORAM/src/pyoram/tests/baselines/k2_h3_b2.dot b/PyORAM/src/pyoram/tests/baselines/k2_h3_b2.dot new file mode 100644 index 0000000..3d92a16 --- /dev/null +++ b/PyORAM/src/pyoram/tests/baselines/k2_h3_b2.dot @@ -0,0 +1,33 @@ +// Created by SizedVirtualHeap.write_as_dot(...) +digraph heaptree { +node [shape=record] + 0 [penwidth=1,label="''}"]; + 1 [penwidth=1,label="'0'}"]; + 3 [penwidth=1,label="'00'}"]; + 7 [penwidth=1,label="'000'}"]; + 3 -> 7 ; + 8 [penwidth=1,label="'001'}"]; + 3 -> 8 ; + 1 -> 3 ; + 4 [penwidth=1,label="'01'}"]; + 9 [penwidth=1,label="'010'}"]; + 4 -> 9 ; + 10 [penwidth=1,label="'011'}"]; + 4 -> 10 ; + 1 -> 4 ; + 0 -> 1 ; + 2 [penwidth=1,label="'1'}"]; + 5 [penwidth=1,label="'10'}"]; + 11 [penwidth=1,label="'100'}"]; + 5 -> 11 ; + 12 [penwidth=1,label="'101'}"]; + 5 -> 12 ; + 2 -> 5 ; + 6 [penwidth=1,label="'11'}"]; + 13 [penwidth=1,label="'110'}"]; + 6 -> 13 ; + 14 [penwidth=1,label="'111'}"]; + 6 -> 14 ; + 2 -> 6 ; + 0 -> 2 ; +} diff --git a/PyORAM/src/pyoram/tests/baselines/k2_h3_b2_data.dot b/PyORAM/src/pyoram/tests/baselines/k2_h3_b2_data.dot new file mode 100644 index 0000000..470cb0b --- /dev/null +++ b/PyORAM/src/pyoram/tests/baselines/k2_h3_b2_data.dot @@ -0,0 +1,33 @@ +// Created by SizedVirtualHeap.write_as_dot(...) +digraph heaptree { +node [shape=record] + 0 [penwidth=1,label="{{0}|{1}}"]; + 1 [penwidth=1,label="{{2}|{3}}"]; + 3 [penwidth=1,label="{{6}|{7}}"]; + 7 [penwidth=1,label="{{14}|{15}}"]; + 3 -> 7 ; + 8 [penwidth=1,label="{{16}|{17}}"]; + 3 -> 8 ; + 1 -> 3 ; + 4 [penwidth=1,label="{{8}|{9}}"]; + 9 [penwidth=1,label="{{18}|{19}}"]; + 4 -> 9 ; + 10 [penwidth=1,label="{{20}|{21}}"]; + 4 -> 10 ; + 1 -> 4 ; + 0 -> 1 ; + 2 [penwidth=1,label="{{4}|{5}}"]; + 5 [penwidth=1,label="{{10}|{11}}"]; + 11 [penwidth=1,label="{{22}|{23}}"]; + 5 -> 11 ; + 12 [penwidth=1,label="{{24}|{25}}"]; + 5 -> 12 ; + 2 -> 5 ; + 6 [penwidth=1,label="{{12}|{13}}"]; + 13 [penwidth=1,label="{{26}|{27}}"]; + 6 -> 13 ; + 14 [penwidth=1,label="{{28}|{29}}"]; + 6 -> 14 ; + 2 -> 6 ; + 0 -> 2 ; +} diff --git a/PyORAM/src/pyoram/tests/baselines/k3_h3_b1.dot b/PyORAM/src/pyoram/tests/baselines/k3_h3_b1.dot new file mode 100644 index 0000000..c2bb820 --- /dev/null +++ b/PyORAM/src/pyoram/tests/baselines/k3_h3_b1.dot @@ -0,0 +1,83 @@ +// Created by SizedVirtualHeap.write_as_dot(...) +digraph heaptree { +node [shape=record] + 0 [penwidth=1,label="''}"]; + 1 [penwidth=1,label="'0'}"]; + 4 [penwidth=1,label="'00'}"]; + 13 [penwidth=1,label="'000'}"]; + 4 -> 13 ; + 14 [penwidth=1,label="'001'}"]; + 4 -> 14 ; + 15 [penwidth=1,label="'002'}"]; + 4 -> 15 ; + 1 -> 4 ; + 5 [penwidth=1,label="'01'}"]; + 16 [penwidth=1,label="'010'}"]; + 5 -> 16 ; + 17 [penwidth=1,label="'011'}"]; + 5 -> 17 ; + 18 [penwidth=1,label="'012'}"]; + 5 -> 18 ; + 1 -> 5 ; + 6 [penwidth=1,label="'02'}"]; + 19 [penwidth=1,label="'020'}"]; + 6 -> 19 ; + 20 [penwidth=1,label="'021'}"]; + 6 -> 20 ; + 21 [penwidth=1,label="'022'}"]; + 6 -> 21 ; + 1 -> 6 ; + 0 -> 1 ; + 2 [penwidth=1,label="'1'}"]; + 7 [penwidth=1,label="'10'}"]; + 22 [penwidth=1,label="'100'}"]; + 7 -> 22 ; + 23 [penwidth=1,label="'101'}"]; + 7 -> 23 ; + 24 [penwidth=1,label="'102'}"]; + 7 -> 24 ; + 2 -> 7 ; + 8 [penwidth=1,label="'11'}"]; + 25 [penwidth=1,label="'110'}"]; + 8 -> 25 ; + 26 [penwidth=1,label="'111'}"]; + 8 -> 26 ; + 27 [penwidth=1,label="'112'}"]; + 8 -> 27 ; + 2 -> 8 ; + 9 [penwidth=1,label="'12'}"]; + 28 [penwidth=1,label="'120'}"]; + 9 -> 28 ; + 29 [penwidth=1,label="'121'}"]; + 9 -> 29 ; + 30 [penwidth=1,label="'122'}"]; + 9 -> 30 ; + 2 -> 9 ; + 0 -> 2 ; + 3 [penwidth=1,label="'2'}"]; + 10 [penwidth=1,label="'20'}"]; + 31 [penwidth=1,label="'200'}"]; + 10 -> 31 ; + 32 [penwidth=1,label="'201'}"]; + 10 -> 32 ; + 33 [penwidth=1,label="'202'}"]; + 10 -> 33 ; + 3 -> 10 ; + 11 [penwidth=1,label="'21'}"]; + 34 [penwidth=1,label="'210'}"]; + 11 -> 34 ; + 35 [penwidth=1,label="'211'}"]; + 11 -> 35 ; + 36 [penwidth=1,label="'212'}"]; + 11 -> 36 ; + 3 -> 11 ; + 12 [penwidth=1,label="'22'}"]; + 37 [penwidth=1,label="'220'}"]; + 12 -> 37 ; + 38 [penwidth=1,label="'221'}"]; + 12 -> 38 ; + 39 [penwidth=1,label="'222'}"]; + 12 -> 39 ; + 3 -> 12 ; + 0 -> 3 ; +} diff --git a/PyORAM/src/pyoram/tests/baselines/k3_h3_b1_data.dot b/PyORAM/src/pyoram/tests/baselines/k3_h3_b1_data.dot new file mode 100644 index 0000000..91e1802 --- /dev/null +++ b/PyORAM/src/pyoram/tests/baselines/k3_h3_b1_data.dot @@ -0,0 +1,83 @@ +// Created by SizedVirtualHeap.write_as_dot(...) +digraph heaptree { +node [shape=record] + 0 [penwidth=1,label="{{0}}"]; + 1 [penwidth=1,label="{{1}}"]; + 4 [penwidth=1,label="{{4}}"]; + 13 [penwidth=1,label="{{13}}"]; + 4 -> 13 ; + 14 [penwidth=1,label="{{14}}"]; + 4 -> 14 ; + 15 [penwidth=1,label="{{15}}"]; + 4 -> 15 ; + 1 -> 4 ; + 5 [penwidth=1,label="{{5}}"]; + 16 [penwidth=1,label="{{16}}"]; + 5 -> 16 ; + 17 [penwidth=1,label="{{17}}"]; + 5 -> 17 ; + 18 [penwidth=1,label="{{18}}"]; + 5 -> 18 ; + 1 -> 5 ; + 6 [penwidth=1,label="{{6}}"]; + 19 [penwidth=1,label="{{19}}"]; + 6 -> 19 ; + 20 [penwidth=1,label="{{20}}"]; + 6 -> 20 ; + 21 [penwidth=1,label="{{21}}"]; + 6 -> 21 ; + 1 -> 6 ; + 0 -> 1 ; + 2 [penwidth=1,label="{{2}}"]; + 7 [penwidth=1,label="{{7}}"]; + 22 [penwidth=1,label="{{22}}"]; + 7 -> 22 ; + 23 [penwidth=1,label="{{23}}"]; + 7 -> 23 ; + 24 [penwidth=1,label="{{24}}"]; + 7 -> 24 ; + 2 -> 7 ; + 8 [penwidth=1,label="{{8}}"]; + 25 [penwidth=1,label="{{25}}"]; + 8 -> 25 ; + 26 [penwidth=1,label="{{26}}"]; + 8 -> 26 ; + 27 [penwidth=1,label="{{27}}"]; + 8 -> 27 ; + 2 -> 8 ; + 9 [penwidth=1,label="{{9}}"]; + 28 [penwidth=1,label="{{28}}"]; + 9 -> 28 ; + 29 [penwidth=1,label="{{29}}"]; + 9 -> 29 ; + 30 [penwidth=1,label="{{30}}"]; + 9 -> 30 ; + 2 -> 9 ; + 0 -> 2 ; + 3 [penwidth=1,label="{{3}}"]; + 10 [penwidth=1,label="{{10}}"]; + 31 [penwidth=1,label="{{31}}"]; + 10 -> 31 ; + 32 [penwidth=1,label="{{32}}"]; + 10 -> 32 ; + 33 [penwidth=1,label="{{33}}"]; + 10 -> 33 ; + 3 -> 10 ; + 11 [penwidth=1,label="{{11}}"]; + 34 [penwidth=1,label="{{34}}"]; + 11 -> 34 ; + 35 [penwidth=1,label="{{35}}"]; + 11 -> 35 ; + 36 [penwidth=1,label="{{36}}"]; + 11 -> 36 ; + 3 -> 11 ; + 12 [penwidth=1,label="{{12}}"]; + 37 [penwidth=1,label="{{37}}"]; + 12 -> 37 ; + 38 [penwidth=1,label="{{38}}"]; + 12 -> 38 ; + 39 [penwidth=1,label="{{39}}"]; + 12 -> 39 ; + 3 -> 12 ; + 0 -> 3 ; +} diff --git a/PyORAM/src/pyoram/tests/baselines/k3_h3_b2.dot b/PyORAM/src/pyoram/tests/baselines/k3_h3_b2.dot new file mode 100644 index 0000000..c2bb820 --- /dev/null +++ b/PyORAM/src/pyoram/tests/baselines/k3_h3_b2.dot @@ -0,0 +1,83 @@ +// Created by SizedVirtualHeap.write_as_dot(...) +digraph heaptree { +node [shape=record] + 0 [penwidth=1,label="''}"]; + 1 [penwidth=1,label="'0'}"]; + 4 [penwidth=1,label="'00'}"]; + 13 [penwidth=1,label="'000'}"]; + 4 -> 13 ; + 14 [penwidth=1,label="'001'}"]; + 4 -> 14 ; + 15 [penwidth=1,label="'002'}"]; + 4 -> 15 ; + 1 -> 4 ; + 5 [penwidth=1,label="'01'}"]; + 16 [penwidth=1,label="'010'}"]; + 5 -> 16 ; + 17 [penwidth=1,label="'011'}"]; + 5 -> 17 ; + 18 [penwidth=1,label="'012'}"]; + 5 -> 18 ; + 1 -> 5 ; + 6 [penwidth=1,label="'02'}"]; + 19 [penwidth=1,label="'020'}"]; + 6 -> 19 ; + 20 [penwidth=1,label="'021'}"]; + 6 -> 20 ; + 21 [penwidth=1,label="'022'}"]; + 6 -> 21 ; + 1 -> 6 ; + 0 -> 1 ; + 2 [penwidth=1,label="'1'}"]; + 7 [penwidth=1,label="'10'}"]; + 22 [penwidth=1,label="'100'}"]; + 7 -> 22 ; + 23 [penwidth=1,label="'101'}"]; + 7 -> 23 ; + 24 [penwidth=1,label="'102'}"]; + 7 -> 24 ; + 2 -> 7 ; + 8 [penwidth=1,label="'11'}"]; + 25 [penwidth=1,label="'110'}"]; + 8 -> 25 ; + 26 [penwidth=1,label="'111'}"]; + 8 -> 26 ; + 27 [penwidth=1,label="'112'}"]; + 8 -> 27 ; + 2 -> 8 ; + 9 [penwidth=1,label="'12'}"]; + 28 [penwidth=1,label="'120'}"]; + 9 -> 28 ; + 29 [penwidth=1,label="'121'}"]; + 9 -> 29 ; + 30 [penwidth=1,label="'122'}"]; + 9 -> 30 ; + 2 -> 9 ; + 0 -> 2 ; + 3 [penwidth=1,label="'2'}"]; + 10 [penwidth=1,label="'20'}"]; + 31 [penwidth=1,label="'200'}"]; + 10 -> 31 ; + 32 [penwidth=1,label="'201'}"]; + 10 -> 32 ; + 33 [penwidth=1,label="'202'}"]; + 10 -> 33 ; + 3 -> 10 ; + 11 [penwidth=1,label="'21'}"]; + 34 [penwidth=1,label="'210'}"]; + 11 -> 34 ; + 35 [penwidth=1,label="'211'}"]; + 11 -> 35 ; + 36 [penwidth=1,label="'212'}"]; + 11 -> 36 ; + 3 -> 11 ; + 12 [penwidth=1,label="'22'}"]; + 37 [penwidth=1,label="'220'}"]; + 12 -> 37 ; + 38 [penwidth=1,label="'221'}"]; + 12 -> 38 ; + 39 [penwidth=1,label="'222'}"]; + 12 -> 39 ; + 3 -> 12 ; + 0 -> 3 ; +} diff --git a/PyORAM/src/pyoram/tests/baselines/k3_h3_b2_data.dot b/PyORAM/src/pyoram/tests/baselines/k3_h3_b2_data.dot new file mode 100644 index 0000000..7a20bc6 --- /dev/null +++ b/PyORAM/src/pyoram/tests/baselines/k3_h3_b2_data.dot @@ -0,0 +1,83 @@ +// Created by SizedVirtualHeap.write_as_dot(...) +digraph heaptree { +node [shape=record] + 0 [penwidth=1,label="{{0}|{1}}"]; + 1 [penwidth=1,label="{{2}|{3}}"]; + 4 [penwidth=1,label="{{8}|{9}}"]; + 13 [penwidth=1,label="{{26}|{27}}"]; + 4 -> 13 ; + 14 [penwidth=1,label="{{28}|{29}}"]; + 4 -> 14 ; + 15 [penwidth=1,label="{{30}|{31}}"]; + 4 -> 15 ; + 1 -> 4 ; + 5 [penwidth=1,label="{{10}|{11}}"]; + 16 [penwidth=1,label="{{32}|{33}}"]; + 5 -> 16 ; + 17 [penwidth=1,label="{{34}|{35}}"]; + 5 -> 17 ; + 18 [penwidth=1,label="{{36}|{37}}"]; + 5 -> 18 ; + 1 -> 5 ; + 6 [penwidth=1,label="{{12}|{13}}"]; + 19 [penwidth=1,label="{{38}|{39}}"]; + 6 -> 19 ; + 20 [penwidth=1,label="{{40}|{41}}"]; + 6 -> 20 ; + 21 [penwidth=1,label="{{42}|{43}}"]; + 6 -> 21 ; + 1 -> 6 ; + 0 -> 1 ; + 2 [penwidth=1,label="{{4}|{5}}"]; + 7 [penwidth=1,label="{{14}|{15}}"]; + 22 [penwidth=1,label="{{44}|{45}}"]; + 7 -> 22 ; + 23 [penwidth=1,label="{{46}|{47}}"]; + 7 -> 23 ; + 24 [penwidth=1,label="{{48}|{49}}"]; + 7 -> 24 ; + 2 -> 7 ; + 8 [penwidth=1,label="{{16}|{17}}"]; + 25 [penwidth=1,label="{{50}|{51}}"]; + 8 -> 25 ; + 26 [penwidth=1,label="{{52}|{53}}"]; + 8 -> 26 ; + 27 [penwidth=1,label="{{54}|{55}}"]; + 8 -> 27 ; + 2 -> 8 ; + 9 [penwidth=1,label="{{18}|{19}}"]; + 28 [penwidth=1,label="{{56}|{57}}"]; + 9 -> 28 ; + 29 [penwidth=1,label="{{58}|{59}}"]; + 9 -> 29 ; + 30 [penwidth=1,label="{{60}|{61}}"]; + 9 -> 30 ; + 2 -> 9 ; + 0 -> 2 ; + 3 [penwidth=1,label="{{6}|{7}}"]; + 10 [penwidth=1,label="{{20}|{21}}"]; + 31 [penwidth=1,label="{{62}|{63}}"]; + 10 -> 31 ; + 32 [penwidth=1,label="{{64}|{65}}"]; + 10 -> 32 ; + 33 [penwidth=1,label="{{66}|{67}}"]; + 10 -> 33 ; + 3 -> 10 ; + 11 [penwidth=1,label="{{22}|{23}}"]; + 34 [penwidth=1,label="{{68}|{69}}"]; + 11 -> 34 ; + 35 [penwidth=1,label="{{70}|{71}}"]; + 11 -> 35 ; + 36 [penwidth=1,label="{{72}|{73}}"]; + 11 -> 36 ; + 3 -> 11 ; + 12 [penwidth=1,label="{{24}|{25}}"]; + 37 [penwidth=1,label="{{74}|{75}}"]; + 12 -> 37 ; + 38 [penwidth=1,label="{{76}|{77}}"]; + 12 -> 38 ; + 39 [penwidth=1,label="{{78}|{79}}"]; + 12 -> 39 ; + 3 -> 12 ; + 0 -> 3 ; +} diff --git a/PyORAM/src/pyoram/tests/test_aes.py b/PyORAM/src/pyoram/tests/test_aes.py new file mode 100644 index 0000000..ba312ca --- /dev/null +++ b/PyORAM/src/pyoram/tests/test_aes.py @@ -0,0 +1,108 @@ +import unittest2 + +from pyoram.crypto.aes import AES + +class TestAES(unittest2.TestCase): + + def test_KeyGen(self): + self.assertEqual(len(AES.key_sizes), 3) + self.assertEqual(len(set(AES.key_sizes)), 3) + for keysize in AES.key_sizes: + key_list = [] + key_set = set() + for i in range(10): + k = AES.KeyGen(keysize) + self.assertEqual(len(k), keysize) + key_list.append(k) + key_set.add(k) + self.assertEqual(len(key_list), 10) + # make sure every key is unique + self.assertEqual(len(key_list), len(key_set)) + + def test_CTR(self): + self._test_Enc_Dec( + AES.CTREnc, + AES.CTRDec, + lambda i, size: bytes(bytearray([i]) * size)) + + def test_GCM(self): + self._test_Enc_Dec( + AES.GCMEnc, + AES.GCMDec, + lambda i, size: bytes(bytearray([i]) * size)) + + def _test_Enc_Dec(self, + enc_func, + dec_func, + get_plaintext): + blocksize_factor = [0.5, 1, 1.5, 2, 2.5] + plaintext_blocks = [] + for i, f in enumerate(blocksize_factor): + size = AES.block_size * f + size = int(round(size)) + if int(f) != f: + assert (size % AES.block_size) != 0 + plaintext_blocks.append(get_plaintext(i, size)) + + assert len(AES.key_sizes) > 0 + ciphertext_blocks = {} + keys = {} + for keysize in AES.key_sizes: + key = AES.KeyGen(keysize) + keys[keysize] = key + ciphertext_blocks[keysize] = [] + for block in plaintext_blocks: + ciphertext_blocks[keysize].append( + enc_func(key, block)) + + self.assertEqual(len(ciphertext_blocks), + len(AES.key_sizes)) + self.assertEqual(len(keys), + len(AES.key_sizes)) + + plaintext_decrypted_blocks = {} + for keysize in keys: + key = keys[keysize] + plaintext_decrypted_blocks[keysize] = [] + for block in ciphertext_blocks[keysize]: + plaintext_decrypted_blocks[keysize].append( + dec_func(key, block)) + + self.assertEqual(len(plaintext_decrypted_blocks), + len(AES.key_sizes)) + + for i in range(len(blocksize_factor)): + for keysize in AES.key_sizes: + self.assertEqual( + plaintext_blocks[i], + plaintext_decrypted_blocks[keysize][i]) + self.assertNotEqual( + plaintext_blocks[i], + ciphertext_blocks[keysize][i]) + if enc_func is AES.CTREnc: + self.assertEqual( + len(ciphertext_blocks[keysize][i]), + len(plaintext_blocks[i]) + AES.block_size) + else: + assert enc_func is AES.GCMEnc + self.assertEqual( + len(ciphertext_blocks[keysize][i]), + len(plaintext_blocks[i]) + 2*AES.block_size) + # check IND-CPA + key = keys[keysize] + alt_ciphertext = enc_func(key, plaintext_blocks[i]) + self.assertNotEqual( + ciphertext_blocks[keysize][i], + alt_ciphertext) + self.assertEqual( + len(ciphertext_blocks[keysize][i]), + len(alt_ciphertext)) + self.assertNotEqual( + ciphertext_blocks[keysize][i][:AES.block_size], + alt_ciphertext[:AES.block_size]) + self.assertNotEqual( + ciphertext_blocks[keysize][i][AES.block_size:], + alt_ciphertext[AES.block_size:]) + +if __name__ == "__main__": + unittest2.main() # pragma: no cover diff --git a/PyORAM/src/pyoram/tests/test_block_storage.py b/PyORAM/src/pyoram/tests/test_block_storage.py new file mode 100644 index 0000000..ea3ad0e --- /dev/null +++ b/PyORAM/src/pyoram/tests/test_block_storage.py @@ -0,0 +1,1023 @@ +import os +import shutil +import unittest2 +import tempfile +import struct + +from pyoram.storage.block_storage import \ + BlockStorageTypeFactory +from pyoram.storage.block_storage_file import \ + BlockStorageFile +from pyoram.storage.block_storage_mmap import \ + BlockStorageMMap +from pyoram.storage.block_storage_ram import \ + BlockStorageRAM +from pyoram.storage.block_storage_sftp import \ + BlockStorageSFTP +from pyoram.storage.block_storage_s3 import \ + BlockStorageS3 +from pyoram.storage.boto3_s3_wrapper import \ + (Boto3S3Wrapper, + MockBoto3S3Wrapper) + +import six +from six.moves import xrange +from six import BytesIO + +thisdir = os.path.dirname(os.path.abspath(__file__)) + +try: + import boto3 + has_boto3 = True +except: # pragma: no cover + has_boto3 = False # pragma: no cover + +class TestBlockStorageTypeFactory(unittest2.TestCase): + + def test_file(self): + self.assertIs(BlockStorageTypeFactory('file'), + BlockStorageFile) + + def test_mmap(self): + self.assertIs(BlockStorageTypeFactory('mmap'), + BlockStorageMMap) + + def test_ram(self): + self.assertIs(BlockStorageTypeFactory('ram'), + BlockStorageRAM) + + def test_sftp(self): + self.assertIs(BlockStorageTypeFactory('sftp'), + BlockStorageSFTP) + + def test_s3(self): + self.assertIs(BlockStorageTypeFactory('s3'), + BlockStorageS3) + + def test_invalid(self): + with self.assertRaises(ValueError): + BlockStorageTypeFactory(None) + + def test_register_invalid_name(self): + with self.assertRaises(ValueError): + BlockStorageTypeFactory.register_device( + 's3', BlockStorageFile) + + def test_register_invalid_type(self): + with self.assertRaises(TypeError): + BlockStorageTypeFactory.register_device( + 'new_str_type', str) + +class _TestBlockStorage(object): + + _type = None + _type_kwds = None + + @classmethod + def _read_storage(cls, storage): + with open(storage.storage_name, 'rb') as f: + return f.read() + + @classmethod + def _remove_storage(cls, name): + if os.path.exists(name): + if os.path.isdir(name): + shutil.rmtree(name, ignore_errors=True) + else: + os.remove(name) + + @classmethod + def _check_exists(cls, name): + return os.path.exists(name) + + @classmethod + def _get_empty_existing(cls): + return os.path.join(thisdir, + "baselines", + "exists.empty") + + @classmethod + def _get_dummy_noexist(cls): + fd, name = tempfile.mkstemp(dir=os.getcwd()) + os.close(fd) + return name + + def _open_teststorage(self, **kwds): + kwds.update(self._type_kwds) + return self._type(self._testfname, **kwds) + + def _reopen_storage(self, storage): + return self._type(storage.storage_name, **self._type_kwds) + + @classmethod + def setUpClass(cls): + assert cls._type is not None + assert cls._type_kwds is not None + cls._dummy_name = cls._get_dummy_noexist() + if cls._check_exists(cls._dummy_name): + cls._remove_storage(cls._dummy_name) + if os.path.exists(cls._dummy_name): + _TestBlockStorage.\ + _remove_storage(cls._dummy_name) # pragma: no cover + cls._block_size = 25 + cls._block_count = 5 + cls._testfname = cls.__name__ + "_testfile.bin" + cls._blocks = [] + f = cls._type.setup( + cls._testfname, + block_size=cls._block_size, + block_count=cls._block_count, + initialize=lambda i: bytes(bytearray([i])*cls._block_size), + ignore_existing=True, + **cls._type_kwds) + f.close() + cls._original_f = f + for i in range(cls._block_count): + data = bytearray([i])*cls._block_size + cls._blocks.append(data) + + @classmethod + def tearDownClass(cls): + cls._remove_storage(cls._testfname) + cls._remove_storage(cls._dummy_name) + + def test_setup_fails(self): + self.assertEqual(self._check_exists(self._dummy_name), False) + with self.assertRaises(IOError): + self._type.setup( + self._get_empty_existing(), + block_size=10, + block_count=10, + **self._type_kwds) + self.assertEqual(self._check_exists(self._dummy_name), False) + with self.assertRaises(IOError): + self._type.setup( + self._get_empty_existing(), + block_size=10, + block_count=10, + ignore_existing=False, + **self._type_kwds) + self.assertEqual(self._check_exists(self._dummy_name), False) + with self.assertRaises(ValueError): + self._type.setup(self._dummy_name, + block_size=0, + block_count=1, + **self._type_kwds) + self.assertEqual(self._check_exists(self._dummy_name), False) + with self.assertRaises(ValueError): + self._type.setup(self._dummy_name, + block_size=1, + block_count=0, + **self._type_kwds) + self.assertEqual(self._check_exists(self._dummy_name), False) + with self.assertRaises(TypeError): + self._type.setup(self._dummy_name, + block_size=1, + block_count=1, + header_data=2, + **self._type_kwds) + self.assertEqual(self._check_exists(self._dummy_name), False) + # TODO: The multiprocessing module is bad + # about handling exceptions raised on the + # thread's stack. + #with self.assertRaises(ValueError): + # def _init(i): + # raise ValueError + # self._type.setup(self._dummy_name, + # block_size=1, + # block_count=1, + # initialize=_init, + # **self._type_kwds) + #self.assertEqual(self._check_exists(self._dummy_name), False) + + def test_setup(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + self._remove_storage(fname) + bsize = 10 + bcount = 11 + fsetup = self._type.setup(fname, bsize, bcount, **self._type_kwds) + fsetup.close() + flen = len(self._read_storage(fsetup)) + self.assertEqual( + flen, + self._type.compute_storage_size(bsize, + bcount)) + self.assertEqual( + flen > + self._type.compute_storage_size(bsize, + bcount, + ignore_header=True), + True) + with self._reopen_storage(fsetup) as f: + self.assertEqual(f.header_data, bytes()) + self.assertEqual(fsetup.header_data, bytes()) + self.assertEqual(f.block_size, bsize) + self.assertEqual(fsetup.block_size, bsize) + self.assertEqual(f.block_count, bcount) + self.assertEqual(fsetup.block_count, bcount) + self.assertEqual(f.storage_name, fsetup.storage_name) + self.assertEqual(fsetup.storage_name, fsetup.storage_name) + if self._type is not BlockStorageRAM: + self.assertEqual(fsetup.storage_name, fname) + else: + self.assertEqual(fsetup.storage_name, None) + self._remove_storage(fname) + + def test_setup_withdata(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + self._remove_storage(fname) + bsize = 10 + bcount = 11 + header_data = bytes(bytearray([0,1,2])) + fsetup = self._type.setup(fname, + bsize, + bcount, + header_data=header_data, + **self._type_kwds) + fsetup.close() + + flen = len(self._read_storage(fsetup)) + self.assertEqual( + flen, + self._type.compute_storage_size(bsize, + bcount, + header_data=header_data)) + self.assertTrue(len(header_data) > 0) + self.assertEqual( + self._type.compute_storage_size(bsize, + bcount) < + self._type.compute_storage_size(bsize, + bcount, + header_data=header_data), + True) + self.assertEqual( + flen > + self._type.compute_storage_size(bsize, + bcount, + header_data=header_data, + ignore_header=True), + True) + with self._reopen_storage(fsetup) as f: + self.assertEqual(f.header_data, header_data) + self.assertEqual(fsetup.header_data, header_data) + self.assertEqual(f.block_size, bsize) + self.assertEqual(fsetup.block_size, bsize) + self.assertEqual(f.block_count, bcount) + self.assertEqual(fsetup.block_count, bcount) + self.assertEqual(f.storage_name, fsetup.storage_name) + self.assertEqual(fsetup.storage_name, fsetup.storage_name) + if self._type is not BlockStorageRAM: + self.assertEqual(fsetup.storage_name, fname) + else: + self.assertEqual(fsetup.storage_name, None) + + self._remove_storage(fname) + + def test_init_noexists(self): + self.assertEqual(self._check_exists(self._dummy_name), False) + with self.assertRaises(IOError): + with self._type(self._dummy_name, **self._type_kwds) as f: + pass # pragma: no cover + + def test_init_exists(self): + self.assertEqual(self._check_exists(self._testfname), True) + databefore = self._read_storage(self._original_f) + with self._open_teststorage() as f: + self.assertEqual(f.block_size, self._block_size) + self.assertEqual(f.block_count, self._block_count) + self.assertEqual(f.storage_name, self._testfname) + self.assertEqual(f.header_data, bytes()) + self.assertEqual(self._check_exists(self._testfname), True) + dataafter = self._read_storage(self._original_f) + self.assertEqual(databefore, dataafter) + + def test_read_block(self): + with self._open_teststorage() as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + for i, data in enumerate(self._blocks): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + for i, data in enumerate(self._blocks): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + for i, data in reversed(list(enumerate(self._blocks))): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + for i, data in reversed(list(enumerate(self._blocks))): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + self._block_count*self._block_size*4) + + with self._open_teststorage() as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + self.assertEqual(list(bytearray(f.read_block(0))), + list(self._blocks[0])) + self.assertEqual(list(bytearray(f.read_block(self._block_count-1))), + list(self._blocks[-1])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + self._block_size*2) + + def test_write_block(self): + data = bytearray([self._block_count])*self._block_size + self.assertEqual(len(data) > 0, True) + with self._open_teststorage() as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + for i in xrange(self._block_count): + self.assertNotEqual(list(bytearray(f.read_block(i))), + list(data)) + for i in xrange(self._block_count): + f.write_block(i, bytes(data)) + for i in xrange(self._block_count): + self.assertEqual(list(bytearray(f.read_block(i))), + list(data)) + for i, block in enumerate(self._blocks): + f.write_block(i, bytes(block)) + + self.assertEqual(f.bytes_sent, + self._block_count*self._block_size*2) + self.assertEqual(f.bytes_received, + self._block_count*self._block_size*2) + + def test_read_blocks(self): + with self._open_teststorage() as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + data = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + data = f.read_blocks([0]) + self.assertEqual(len(data), 1) + self.assertEqual(list(bytearray(data[0])), + list(self._blocks[0])) + self.assertEqual(len(self._blocks) > 1, True) + data = f.read_blocks(list(xrange(1, self._block_count)) + [0]) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data[:-1], 1): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + self.assertEqual(list(bytearray(data[-1])), + list(self._blocks[0])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + (2*self._block_count+1)*self._block_size) + + def test_yield_blocks(self): + with self._open_teststorage() as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + data = list(f.yield_blocks(list(xrange(self._block_count)))) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + data = list(f.yield_blocks([0])) + self.assertEqual(len(data), 1) + self.assertEqual(list(bytearray(data[0])), + list(self._blocks[0])) + self.assertEqual(len(self._blocks) > 1, True) + data = list(f.yield_blocks(list(xrange(1, self._block_count)) + [0])) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data[:-1], 1): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + self.assertEqual(list(bytearray(data[-1])), + list(self._blocks[0])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + (2*self._block_count+1)*self._block_size) + + def test_write_blocks(self): + data = [bytearray([self._block_count])*self._block_size + for i in xrange(self._block_count)] + with self._open_teststorage() as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + orig = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(orig), self._block_count) + for i, block in enumerate(orig): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + f.write_blocks(list(xrange(self._block_count)), + [bytes(b) for b in data]) + new = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(new), self._block_count) + for i, block in enumerate(new): + self.assertEqual(list(bytearray(block)), + list(data[i])) + f.write_blocks(list(xrange(self._block_count)), + [bytes(b) for b in self._blocks]) + orig = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(orig), self._block_count) + for i, block in enumerate(orig): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + + self.assertEqual(f.bytes_sent, + self._block_count*self._block_size*2) + self.assertEqual(f.bytes_received, + self._block_count*self._block_size*3) + + def test_update_header_data(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + self._remove_storage(fname) + bsize = 10 + bcount = 11 + header_data = bytes(bytearray([0,1,2])) + fsetup = self._type.setup(fname, + block_size=bsize, + block_count=bcount, + header_data=header_data, + **self._type_kwds) + fsetup.close() + new_header_data = bytes(bytearray([1,1,1])) + with self._reopen_storage(fsetup) as f: + self.assertEqual(f.header_data, header_data) + f.update_header_data(new_header_data) + self.assertEqual(f.header_data, new_header_data) + with self._reopen_storage(fsetup) as f: + self.assertEqual(f.header_data, new_header_data) + with self.assertRaises(ValueError): + with self._reopen_storage(fsetup) as f: + f.update_header_data(bytes(bytearray([1,1]))) + with self.assertRaises(ValueError): + with self._reopen_storage(fsetup) as f: + f.update_header_data(bytes(bytearray([1,1,1,1]))) + with self._reopen_storage(fsetup) as f: + self.assertEqual(f.header_data, new_header_data) + self._remove_storage(fname) + + def test_locked_flag(self): + with self._open_teststorage() as f: + with self.assertRaises(IOError): + with self._open_teststorage() as f1: + pass # pragma: no cover + with self.assertRaises(IOError): + with self._open_teststorage() as f1: + pass # pragma: no cover + with self._open_teststorage(ignore_lock=True) as f1: + pass + with self.assertRaises(IOError): + with self._open_teststorage() as f1: + pass # pragma: no cover + with self._open_teststorage(ignore_lock=True) as f1: + pass + with self._open_teststorage(ignore_lock=True) as f1: + pass + with self._open_teststorage(ignore_lock=True) as f: + pass + + def test_read_block_cloned(self): + with self._open_teststorage() as forig: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + with forig.clone_device() as f: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + for i, data in enumerate(self._blocks): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + for i, data in enumerate(self._blocks): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + for i, data in reversed(list(enumerate(self._blocks))): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + for i, data in reversed(list(enumerate(self._blocks))): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + self._block_count*self._block_size*4) + + with forig.clone_device() as f: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + self.assertEqual(list(bytearray(f.read_block(0))), + list(self._blocks[0])) + self.assertEqual(list(bytearray(f.read_block(self._block_count-1))), + list(self._blocks[-1])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + self._block_size*2) + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + + def test_write_block_cloned(self): + data = bytearray([self._block_count])*self._block_size + self.assertEqual(len(data) > 0, True) + with self._open_teststorage() as forig: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + with forig.clone_device() as f: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + for i in xrange(self._block_count): + self.assertNotEqual(list(bytearray(f.read_block(i))), + list(data)) + for i in xrange(self._block_count): + f.write_block(i, bytes(data)) + for i in xrange(self._block_count): + self.assertEqual(list(bytearray(f.read_block(i))), + list(data)) + for i, block in enumerate(self._blocks): + f.write_block(i, bytes(block)) + + self.assertEqual(f.bytes_sent, + self._block_count*self._block_size*2) + self.assertEqual(f.bytes_received, + self._block_count*self._block_size*2) + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + + def test_read_blocks_cloned(self): + with self._open_teststorage() as forig: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + with forig.clone_device() as f: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + data = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + data = f.read_blocks([0]) + self.assertEqual(len(data), 1) + self.assertEqual(list(bytearray(data[0])), + list(self._blocks[0])) + self.assertEqual(len(self._blocks) > 1, True) + data = f.read_blocks(list(xrange(1, self._block_count)) + [0]) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data[:-1], 1): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + self.assertEqual(list(bytearray(data[-1])), + list(self._blocks[0])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + (2*self._block_count + 1)*self._block_size) + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + + def test_yield_blocks_cloned(self): + with self._open_teststorage() as forig: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + with forig.clone_device() as f: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + data = list(f.yield_blocks(list(xrange(self._block_count)))) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + data = list(f.yield_blocks([0])) + self.assertEqual(len(data), 1) + self.assertEqual(list(bytearray(data[0])), + list(self._blocks[0])) + self.assertEqual(len(self._blocks) > 1, True) + data = list(f.yield_blocks(list(xrange(1, self._block_count)) + [0])) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data[:-1], 1): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + self.assertEqual(list(bytearray(data[-1])), + list(self._blocks[0])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + (2*self._block_count + 1)*self._block_size) + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + + def test_write_blocks_cloned(self): + data = [bytearray([self._block_count])*self._block_size + for i in xrange(self._block_count)] + with self._open_teststorage() as forig: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + with forig.clone_device() as f: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + orig = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(orig), self._block_count) + for i, block in enumerate(orig): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + f.write_blocks(list(xrange(self._block_count)), + [bytes(b) for b in data]) + new = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(new), self._block_count) + for i, block in enumerate(new): + self.assertEqual(list(bytearray(block)), + list(data[i])) + f.write_blocks(list(xrange(self._block_count)), + [bytes(b) for b in self._blocks]) + orig = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(orig), self._block_count) + for i, block in enumerate(orig): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + + self.assertEqual(f.bytes_sent, + self._block_count*self._block_size*2) + self.assertEqual(f.bytes_received, + self._block_count*self._block_size*3) + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + +class TestBlockStorageFile(_TestBlockStorage, + unittest2.TestCase): + _type = BlockStorageFile + _type_kwds = {} + +class TestBlockStorageFileNoThreadPool(_TestBlockStorage, + unittest2.TestCase): + _type = BlockStorageFile + _type_kwds = {'threadpool_size': 0} + +class TestBlockStorageFileThreadPool(_TestBlockStorage, + unittest2.TestCase): + _type = BlockStorageFile + _type_kwds = {'threadpool_size': 1} + +class TestBlockStorageMMap(_TestBlockStorage, + unittest2.TestCase): + _type = BlockStorageMMap + _type_kwds = {} + +class _TestBlockStorageRAM(_TestBlockStorage): + + @classmethod + def _read_storage(cls, storage): + return storage.data + + @classmethod + def _remove_storage(cls, name): + pass + + @classmethod + def _check_exists(cls, name): + return True + + def _open_teststorage(self, **kwds): + kwds.update(self._type_kwds) + return self._type(self._original_f.data, **kwds) + + def _reopen_storage(self, storage): + return self._type(storage.data, **self._type_kwds) + + # + # Override some of the test methods + # + + def test_setup_fails(self): + with self.assertRaises(ValueError): + self._type.setup(self._dummy_name, + block_size=0, + block_count=1, + **self._type_kwds) + with self.assertRaises(ValueError): + self._type.setup(self._dummy_name, + block_size=1, + block_count=0, + **self._type_kwds) + with self.assertRaises(TypeError): + self._type.setup(self._dummy_name, + block_size=1, + block_count=1, + header_data=2, + **self._type_kwds) + + def test_init_noexists(self): + with self.assertRaises(TypeError): + with self._type(2, **self._type_kwds) as f: + pass # pragma: no cover + with self.assertRaises(TypeError): + with self._type(None, **self._type_kwds) as f: + pass # pragma: no cover + with self.assertRaises(struct.error): + with self._type(bytearray(), **self._type_kwds) as f: + pass # pragma: no cover + + def test_init_exists(self): + databefore = self._read_storage(self._original_f) + with self._open_teststorage() as f: + self.assertEqual(f.block_size, self._block_size) + self.assertEqual(f.block_count, self._block_count) + self.assertEqual(f.storage_name, self._original_f.storage_name) + self.assertEqual(f.storage_name, None) + self.assertEqual(f.header_data, bytes()) + dataafter = self._read_storage(self._original_f) + self.assertEqual(databefore, dataafter) + + def test_tofile_fromfile_fileobj(self): + out1 = BytesIO() + self._original_f.tofile(out1) + out1.seek(0) + self.assertEqual(len(self._original_f.data) > 0, True) + self.assertEqual(self._original_f.data, out1.read()) + out1.seek(0) + in1 = self._type.fromfile(out1) + self.assertNotEqual(self._original_f.data, in1.data) + out2 = BytesIO() + in1.tofile(out2) + self.assertNotEqual(self._original_f.data, in1.data) + in1.close() + self.assertEqual(self._original_f.data, in1.data) + out2.seek(0) + with self.assertRaises(IOError): + with self._type.fromfile(out2) as in2: + pass # pragma: no cover + out2.seek(0) + with self._type.fromfile(out2, ignore_lock=True) as in2: + self.assertEqual(self._original_f.data, in1.data) + self.assertNotEqual(self._original_f.data, in2.data) + self.assertEqual(self._original_f.data, in1.data) + self.assertNotEqual(self._original_f.data, in2.data) + + def test_tofile_fromfile_filename(self): + + def _create(): + fd, out = tempfile.mkstemp() + os.close(fd) + return out + def _read(name): + with open(name, 'rb') as f: + return f.read() + + out1 = _create() + self._original_f.tofile(out1) + self.assertEqual(len(self._original_f.data) > 0, True) + self.assertEqual(self._original_f.data, _read(out1)) + in1 = self._type.fromfile(out1) + self.assertNotEqual(self._original_f.data, in1.data) + out2 = _create() + in1.tofile(out2) + self.assertNotEqual(self._original_f.data, in1.data) + in1.close() + self.assertEqual(self._original_f.data, in1.data) + with self.assertRaises(IOError): + with self._type.fromfile(out2) as in2: + pass # pragma: no cover + with self._type.fromfile(out2, ignore_lock=True) as in2: + self.assertEqual(self._original_f.data, in1.data) + self.assertNotEqual(self._original_f.data, in2.data) + self.assertEqual(self._original_f.data, in1.data) + self.assertNotEqual(self._original_f.data, in2.data) + +class TestBlockStorageRAM(_TestBlockStorageRAM, + unittest2.TestCase): + _type = BlockStorageRAM + _type_kwds = {} + +class _dummy_sftp_file(object): + def __init__(self, *args, **kwds): + self._f = open(*args, **kwds) + def __enter__(self): + return self + def __exit__(self, *args): + self._f.close() + def readv(self, chunks): + data = [] + for offset, size in chunks: + self._f.seek(offset) + data.append(self._f.read(size)) + return data + def __getattr__(self, key): + return getattr(self._f, key) + def set_pipelined(self): + pass + +class dummy_sftp(object): + remove = os.remove + stat = os.stat + @staticmethod + def open(*args, **kwds): + return _dummy_sftp_file(*args, **kwds) + @staticmethod + def close(): + pass + +class dummy_sshclient(object): + @staticmethod + def open_sftp(): + return dummy_sftp + +class TestBlockStorageSFTP(_TestBlockStorage, + unittest2.TestCase): + _type = BlockStorageSFTP + _type_kwds = {'sshclient': dummy_sshclient} + + def test_setup_fails_no_sshclient(self): + self.assertEqual(self._check_exists(self._dummy_name), False) + kwds = dict(self._type_kwds) + del kwds['sshclient'] + with self.assertRaises(ValueError): + self._type.setup(self._dummy_name, + block_size=1, + block_count=1, + **kwds) + self.assertEqual(self._check_exists(self._dummy_name), False) + + def test_init_exists_no_sshclient(self): + self.assertEqual(self._check_exists(self._testfname), True) + kwds = dict(self._type_kwds) + del kwds['sshclient'] + with self.assertRaises(ValueError): + with self._type(self._testfname, **kwds) as f: + pass # pragma: no cover + + databefore = self._read_storage(self._original_f) + with self._open_teststorage() as f: + self.assertEqual(f.block_size, self._block_size) + self.assertEqual(f.block_count, self._block_count) + self.assertEqual(f.storage_name, self._testfname) + self.assertEqual(f.header_data, bytes()) + self.assertEqual(self._check_exists(self._testfname), True) + dataafter = self._read_storage(self._original_f) + self.assertEqual(databefore, dataafter) + + +class _TestBlockStorageS3Mock(_TestBlockStorage): + _type = BlockStorageS3 + _type_kwds = {} + + @classmethod + def _read_storage(cls, storage): + import glob + data = bytearray() + name = storage.storage_name + prefix_len = len(os.path.join(name,"b")) + nblocks = max(int(bfile[prefix_len:]) for bfile in glob.glob(name+"/b*")) + 1 + with open(os.path.join(name, BlockStorageS3._index_name), 'rb') as f: + data.extend(f.read()) + for i in range(nblocks): + with open(os.path.join(name, "b"+str(i)), 'rb') as f: + data.extend(f.read()) + return data + + def test_init_exists_no_bucket(self): + self.assertEqual(self._check_exists(self._testfname), True) + databefore = self._read_storage(self._original_f) + with self._open_teststorage() as f: + self.assertEqual(f.block_size, self._block_size) + self.assertEqual(f.block_count, self._block_count) + self.assertEqual(f.storage_name, self._testfname) + self.assertEqual(f.header_data, bytes()) + self.assertEqual(self._check_exists(self._testfname), True) + dataafter = self._read_storage(self._original_f) + self.assertEqual(databefore, dataafter) + kwds = dict(self._type_kwds) + del kwds['bucket_name'] + with self.assertRaises(ValueError): + with self._type(self._testfname, **kwds) as f: + pass # pragma: no cover + dataafter = self._read_storage(self._original_f) + self.assertEqual(databefore, dataafter) + + def test_setup_fails_no_bucket(self): + self.assertEqual(self._check_exists(self._dummy_name), False) + kwds = dict(self._type_kwds) + del kwds['bucket_name'] + with self.assertRaises(ValueError): + self._type.setup(self._dummy_name, + block_size=1, + block_count=1, + **kwds) + self.assertEqual(self._check_exists(self._dummy_name), False) + + def test_setup_ignore_existing(self): + self.assertEqual(self._check_exists(self._dummy_name), False) + with self._type.setup(self._dummy_name, + block_size=1, + block_count=1, + **self._type_kwds) as f: + pass + self.assertEqual(self._check_exists(self._dummy_name), True) + with self.assertRaises(IOError): + with self._type.setup(self._dummy_name, + block_size=1, + block_count=1, + **self._type_kwds) as f: + pass # pragma: no cover + self.assertEqual(self._check_exists(self._dummy_name), True) + with self._type.setup(self._dummy_name, + block_size=1, + block_count=1, + ignore_existing=True, + **self._type_kwds) as f: + pass + self.assertEqual(self._check_exists(self._dummy_name), True) + self._remove_storage(self._dummy_name) + +class TestBlockStorageS3Mock(_TestBlockStorageS3Mock, + unittest2.TestCase): + _type_kwds = {'s3_wrapper': MockBoto3S3Wrapper, + 'bucket_name': '.'} + +class TestBlockStorageS3MockNoThreadPool(_TestBlockStorageS3Mock, + unittest2.TestCase): + _type_kwds = {'s3_wrapper': MockBoto3S3Wrapper, + 'bucket_name': '.', + 'threadpool_size': 0} + +class TestBlockStorageS3MockThreadPool(_TestBlockStorageS3Mock, + unittest2.TestCase): + _type_kwds = {'s3_wrapper': MockBoto3S3Wrapper, + 'bucket_name': '.', + 'threadpool_size': 4} + +@unittest2.skipIf((os.environ.get('PYORAM_AWS_TEST_BUCKET') is None) or \ + (not has_boto3), + "No PYORAM_AWS_TEST_BUCKET defined in environment or " + "boto3 is not available") +class TestBlockStorageS3(_TestBlockStorage, + unittest2.TestCase): + _type = BlockStorageS3 + _type_kwds = {'bucket_name': os.environ.get('PYORAM_AWS_TEST_BUCKET')} + + @classmethod + def _read_storage(cls, storage): + data = bytearray() + name = storage.storage_name + s3 = Boto3S3Wrapper(cls._type_kwds['bucket_name']) + prefix_len = len(name+"/b") + nblocks = 1 + max(int(obj.key[prefix_len:]) for obj + in s3._bucket.objects.filter(Prefix=name+"/b")) + data.extend(s3.download(name+"/"+BlockStorageS3._index_name)) + for i in range(nblocks): + data.extend(s3.download(name+"/b"+str(i))) + return data + + @classmethod + def _remove_storage(cls, name): + Boto3S3Wrapper(cls._type_kwds['bucket_name']).clear(name) + + @classmethod + def _check_exists(cls, name): + return Boto3S3Wrapper(cls._type_kwds['bucket_name']).exists(name) + + @classmethod + def _get_empty_existing(cls): + return "exists.empty" + + @classmethod + def _get_dummy_noexist(cls): + s3 = Boto3S3Wrapper(cls._type_kwds['bucket_name']) + fd, name = tempfile.mkstemp(dir=os.getcwd()) + os.close(fd) + os.remove(name) + while s3.exists(name): + fd, name = tempfile.mkstemp(dir=os.getcwd()) + os.close(fd) + os.remove(name) + return name + +if __name__ == "__main__": + unittest2.main() # pragma: no cover diff --git a/PyORAM/src/pyoram/tests/test_encrypted_block_storage.py b/PyORAM/src/pyoram/tests/test_encrypted_block_storage.py new file mode 100644 index 0000000..af9a4f0 --- /dev/null +++ b/PyORAM/src/pyoram/tests/test_encrypted_block_storage.py @@ -0,0 +1,814 @@ +import os +import unittest2 +import tempfile +import struct + +from pyoram.storage.block_storage import \ + BlockStorageTypeFactory +from pyoram.encrypted_storage.encrypted_block_storage import \ + EncryptedBlockStorage +from pyoram.crypto.aes import AES + +from six.moves import xrange + +thisdir = os.path.dirname(os.path.abspath(__file__)) + +class _TestEncryptedBlockStorage(object): + + _type_name = None + _aes_mode = None + _test_key = None + _test_key_size = None + + @classmethod + def setUpClass(cls): + assert cls._type_name is not None + assert cls._aes_mode is not None + assert not ((cls._test_key is not None) and \ + (cls._test_key_size is not None)) + fd, cls._dummy_name = tempfile.mkstemp() + os.close(fd) + try: + os.remove(cls._dummy_name) + except OSError: # pragma: no cover + pass # pragma: no cover + cls._block_size = 25 + cls._block_count = 5 + cls._testfname = cls.__name__ + "_testfile.bin" + cls._blocks = [] + f = EncryptedBlockStorage.setup( + cls._testfname, + cls._block_size, + cls._block_count, + key_size=cls._test_key_size, + key=cls._test_key, + storage_type=cls._type_name, + aes_mode=cls._aes_mode, + initialize=lambda i: bytes(bytearray([i])*cls._block_size), + ignore_existing=True) + f.close() + cls._key = f.key + for i in range(cls._block_count): + data = bytearray([i])*cls._block_size + cls._blocks.append(data) + + @classmethod + def tearDownClass(cls): + try: + os.remove(cls._testfname) + except OSError: # pragma: no cover + pass # pragma: no cover + try: + os.remove(cls._dummy_name) + except OSError: # pragma: no cover + pass # pragma: no cover + + def test_setup_fails(self): + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(IOError): + EncryptedBlockStorage.setup( + os.path.join(thisdir, + "baselines", + "exists.empty"), + block_size=10, + block_count=10, + key=self._test_key, + key_size=self._test_key_size, + aes_mode=self._aes_mode, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(IOError): + EncryptedBlockStorage.setup( + os.path.join(thisdir, + "baselines", + "exists.empty"), + block_size=10, + block_count=10, + key=self._test_key, + key_size=self._test_key_size, + storage_type=self._type_name, + aes_mode=self._aes_mode, + ignore_existing=False) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(ValueError): + EncryptedBlockStorage.setup( + self._dummy_name, + block_size=0, + block_count=1, + key=self._test_key, + key_size=self._test_key_size, + aes_mode=self._aes_mode, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(ValueError): + EncryptedBlockStorage.setup( + self._dummy_name, + block_size=1, + block_count=0, + key=self._test_key, + key_size=self._test_key_size, + aes_mode=self._aes_mode, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(TypeError): + EncryptedBlockStorage.setup( + self._dummy_name, + block_size=1, + block_count=1, + key=self._test_key, + key_size=self._test_key_size, + aes_mode=self._aes_mode, + storage_type=self._type_name, + header_data=2) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(ValueError): + EncryptedBlockStorage.setup( + self._dummy_name, + block_size=1, + block_count=1, + key=self._test_key, + key_size=self._test_key_size, + aes_mode=None, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(ValueError): + EncryptedBlockStorage.setup( + self._dummy_name, + block_size=1, + block_count=1, + key_size=-1, + aes_mode=self._aes_mode, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(TypeError): + EncryptedBlockStorage.setup( + self._dummy_name, + block_size=1, + block_count=1, + key=-1, + aes_mode=self._aes_mode, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(ValueError): + EncryptedBlockStorage.setup( + self._dummy_name, + block_size=1, + block_count=1, + key=AES.KeyGen(AES.key_sizes[-1]), + key_size=AES.key_sizes[-1], + aes_mode=self._aes_mode, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(ValueError): + EncryptedBlockStorage.setup( + self._dummy_name, + block_size=1, + block_count=1, + key=os.urandom(AES.key_sizes[-1]+100), + aes_mode=self._aes_mode, + storage_type=self._type_name) + + def test_setup(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + if os.path.exists(fname): + os.remove(fname) # pragma: no cover + bsize = 10 + bcount = 11 + fsetup = EncryptedBlockStorage.setup( + fname, + bsize, + bcount, + key=self._test_key, + key_size=self._test_key_size, + aes_mode=self._aes_mode, + storage_type=self._type_name) + fsetup.close() + self.assertEqual(type(fsetup.raw_storage), + BlockStorageTypeFactory(self._type_name)) + with open(fname, 'rb') as f: + flen = len(f.read()) + self.assertEqual( + flen, + EncryptedBlockStorage.compute_storage_size( + bsize, + bcount, + aes_mode=self._aes_mode, + storage_type=self._type_name)) + self.assertEqual( + flen > + EncryptedBlockStorage.compute_storage_size( + bsize, + bcount, + aes_mode=self._aes_mode, + storage_type=self._type_name, + ignore_header=True), + True) + with EncryptedBlockStorage(fname, + key=fsetup.key, + storage_type=self._type_name) as f: + self.assertEqual(f.header_data, bytes()) + self.assertEqual(fsetup.header_data, bytes()) + self.assertEqual(f.key, fsetup.key) + self.assertEqual(f.block_size, bsize) + self.assertEqual(fsetup.block_size, bsize) + self.assertEqual(f.block_count, bcount) + self.assertEqual(fsetup.block_count, bcount) + self.assertEqual(f.storage_name, fname) + self.assertEqual(fsetup.storage_name, fname) + # tamper with the plaintext index + with open(fname, 'r+b') as f: + f.seek(0) + f.write(struct.pack("!L",0)) + with self.assertRaises(ValueError): + with EncryptedBlockStorage(fname, + key=fsetup.key, + storage_type=self._type_name) as f: + pass # pragma: no cover + os.remove(fname) + + def test_setup_withdata(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + if os.path.exists(fname): + os.remove(fname) # pragma: no cover + bsize = 10 + bcount = 11 + header_data = bytes(bytearray([0,1,2])) + fsetup = EncryptedBlockStorage.setup( + fname, + block_size=bsize, + block_count=bcount, + key=self._test_key, + key_size=self._test_key_size, + aes_mode=self._aes_mode, + storage_type=self._type_name, + header_data=header_data) + fsetup.close() + self.assertEqual(type(fsetup.raw_storage), + BlockStorageTypeFactory(self._type_name)) + with open(fname, 'rb') as f: + flen = len(f.read()) + self.assertEqual( + flen, + EncryptedBlockStorage.compute_storage_size( + bsize, + bcount, + aes_mode=self._aes_mode, + storage_type=self._type_name, + header_data=header_data)) + self.assertTrue(len(header_data) > 0) + self.assertEqual( + EncryptedBlockStorage.compute_storage_size( + bsize, + bcount, + aes_mode=self._aes_mode, + storage_type=self._type_name) < + EncryptedBlockStorage.compute_storage_size( + bsize, + bcount, + aes_mode=self._aes_mode, + storage_type=self._type_name, + header_data=header_data), + True) + self.assertEqual( + flen > + EncryptedBlockStorage.compute_storage_size( + bsize, + bcount, + aes_mode=self._aes_mode, + storage_type=self._type_name, + header_data=header_data, + ignore_header=True), + True) + with EncryptedBlockStorage(fname, + key=fsetup.key, + storage_type=self._type_name) as f: + self.assertEqual(f.header_data, header_data) + self.assertEqual(fsetup.header_data, header_data) + self.assertEqual(f.key, fsetup.key) + self.assertEqual(f.block_size, bsize) + self.assertEqual(fsetup.block_size, bsize) + self.assertEqual(f.block_count, bcount) + self.assertEqual(fsetup.block_count, bcount) + self.assertEqual(f.storage_name, fname) + self.assertEqual(fsetup.storage_name, fname) + os.remove(fname) + + def test_init_noexists(self): + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(IOError): + with EncryptedBlockStorage( + self._dummy_name, + key=self._key, + storage_type=self._type_name) as f: + pass # pragma: no cover + + def test_init_exists(self): + self.assertEqual(os.path.exists(self._testfname), True) + with open(self._testfname, 'rb') as f: + databefore = f.read() + with self.assertRaises(ValueError): + with EncryptedBlockStorage(self._testfname, + storage_type=self._type_name) as f: + pass # pragma: no cover + with self.assertRaises(ValueError): + with BlockStorageTypeFactory(self._type_name)(self._testfname) as fb: + with EncryptedBlockStorage(fb, + key=self._key, + storage_type=self._type_name) as f: + pass # pragma: no cover + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as f: + self.assertEqual(f.key, self._key) + self.assertEqual(f.block_size, self._block_size) + self.assertEqual(f.block_count, self._block_count) + self.assertEqual(f.storage_name, self._testfname) + self.assertEqual(f.header_data, bytes()) + self.assertEqual(os.path.exists(self._testfname), True) + with open(self._testfname, 'rb') as f: + dataafter = f.read() + self.assertEqual(databefore, dataafter) + + def test_read_block(self): + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + for i, data in enumerate(self._blocks): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + for i, data in enumerate(self._blocks): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + for i, data in reversed(list(enumerate(self._blocks))): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + for i, data in reversed(list(enumerate(self._blocks))): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + self._block_count*f._storage.block_size*4) + + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + self.assertEqual(list(bytearray(f.read_block(0))), + list(self._blocks[0])) + self.assertEqual(list(bytearray(f.read_block(self._block_count-1))), + list(self._blocks[-1])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + f._storage.block_size*2) + + def test_write_block(self): + data = bytearray([self._block_count])*self._block_size + self.assertEqual(len(data) > 0, True) + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + for i in xrange(self._block_count): + self.assertNotEqual(list(bytearray(f.read_block(i))), + list(data)) + for i in xrange(self._block_count): + f.write_block(i, bytes(data)) + for i in xrange(self._block_count): + self.assertEqual(list(bytearray(f.read_block(i))), + list(data)) + for i, block in enumerate(self._blocks): + f.write_block(i, bytes(block)) + + self.assertEqual(f.bytes_sent, + self._block_count*f._storage.block_size*2) + self.assertEqual(f.bytes_received, + self._block_count*f._storage.block_size*2) + + def test_read_blocks(self): + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + data = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + data = f.read_blocks([0]) + self.assertEqual(len(data), 1) + self.assertEqual(list(bytearray(data[0])), + list(self._blocks[0])) + self.assertEqual(len(self._blocks) > 1, True) + data = f.read_blocks(list(xrange(1, self._block_count)) + [0]) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data[:-1], 1): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + self.assertEqual(list(bytearray(data[-1])), + list(self._blocks[0])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + (2*self._block_count+1)*f._storage.block_size) + + def test_yield_blocks(self): + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + data = list(f.yield_blocks(list(xrange(self._block_count)))) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + data = list(f.yield_blocks([0])) + self.assertEqual(len(data), 1) + self.assertEqual(list(bytearray(data[0])), + list(self._blocks[0])) + self.assertEqual(len(self._blocks) > 1, True) + data = list(f.yield_blocks(list(xrange(1, self._block_count)) + [0])) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data[:-1], 1): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + self.assertEqual(list(bytearray(data[-1])), + list(self._blocks[0])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + (2*self._block_count+1)*f._storage.block_size) + + def test_write_blocks(self): + data = [bytearray([self._block_count])*self._block_size + for i in xrange(self._block_count)] + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + orig = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(orig), self._block_count) + for i, block in enumerate(orig): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + f.write_blocks(list(xrange(self._block_count)), + [bytes(b) for b in data]) + new = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(new), self._block_count) + for i, block in enumerate(new): + self.assertEqual(list(bytearray(block)), + list(data[i])) + f.write_blocks(list(xrange(self._block_count)), + [bytes(b) for b in self._blocks]) + orig = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(orig), self._block_count) + for i, block in enumerate(orig): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + + self.assertEqual(f.bytes_sent, + self._block_count*f._storage.block_size*2) + self.assertEqual(f.bytes_received, + self._block_count*f._storage.block_size*3) + + def test_update_header_data(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + if os.path.exists(fname): + os.remove(fname) # pragma: no cover + bsize = 10 + bcount = 11 + header_data = bytes(bytearray([0,1,2])) + fsetup = EncryptedBlockStorage.setup( + fname, + block_size=bsize, + block_count=bcount, + key=self._test_key, + key_size=self._test_key_size, + header_data=header_data) + fsetup.close() + new_header_data = bytes(bytearray([1,1,1])) + with EncryptedBlockStorage(fname, + key=fsetup.key, + storage_type=self._type_name) as f: + self.assertEqual(f.header_data, header_data) + f.update_header_data(new_header_data) + self.assertEqual(f.header_data, new_header_data) + with EncryptedBlockStorage(fname, + key=fsetup.key, + storage_type=self._type_name) as f: + self.assertEqual(f.header_data, new_header_data) + with self.assertRaises(ValueError): + with EncryptedBlockStorage(fname, + key=fsetup.key, + storage_type=self._type_name) as f: + f.update_header_data(bytes(bytearray([1,1]))) + with self.assertRaises(ValueError): + with EncryptedBlockStorage(fname, + key=fsetup.key, + storage_type=self._type_name) as f: + f.update_header_data(bytes(bytearray([1,1,1,1]))) + with EncryptedBlockStorage(fname, + key=fsetup.key, + storage_type=self._type_name) as f: + self.assertEqual(f.header_data, new_header_data) + os.remove(fname) + + def test_locked_flag(self): + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as f: + with self.assertRaises(IOError): + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as f1: + pass # pragma: no cover + with self.assertRaises(IOError): + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as f1: + pass # pragma: no cover + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name, + ignore_lock=True) as f1: + pass + with self.assertRaises(IOError): + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as f1: + pass # pragma: no cover + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name, + ignore_lock=True) as f1: + pass + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name, + ignore_lock=True) as f1: + pass + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as f: + pass + + def test_read_block_cloned(self): + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as forig: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + with forig.clone_device() as f: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + for i, data in enumerate(self._blocks): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + for i, data in enumerate(self._blocks): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + for i, data in reversed(list(enumerate(self._blocks))): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + for i, data in reversed(list(enumerate(self._blocks))): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + self._block_count*f._storage.block_size*4) + + with forig.clone_device() as f: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + self.assertEqual(list(bytearray(f.read_block(0))), + list(self._blocks[0])) + self.assertEqual(list(bytearray(f.read_block(self._block_count-1))), + list(self._blocks[-1])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + f._storage.block_size*2) + + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + + def test_write_block_cloned(self): + data = bytearray([self._block_count])*self._block_size + self.assertEqual(len(data) > 0, True) + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as forig: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + with forig.clone_device() as f: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + for i in xrange(self._block_count): + self.assertNotEqual(list(bytearray(f.read_block(i))), + list(data)) + for i in xrange(self._block_count): + f.write_block(i, bytes(data)) + for i in xrange(self._block_count): + self.assertEqual(list(bytearray(f.read_block(i))), + list(data)) + for i, block in enumerate(self._blocks): + f.write_block(i, bytes(block)) + + self.assertEqual(f.bytes_sent, + self._block_count*f._storage.block_size*2) + self.assertEqual(f.bytes_received, + self._block_count*f._storage.block_size*2) + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + + def test_read_blocks_cloned(self): + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as forig: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + with forig.clone_device() as f: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + data = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + data = f.read_blocks([0]) + self.assertEqual(len(data), 1) + self.assertEqual(list(bytearray(data[0])), + list(self._blocks[0])) + self.assertEqual(len(self._blocks) > 1, True) + data = f.read_blocks(list(xrange(1, self._block_count)) + [0]) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data[:-1], 1): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + self.assertEqual(list(bytearray(data[-1])), + list(self._blocks[0])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + (2*self._block_count + 1)*f._storage.block_size) + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + + def test_yield_blocks_cloned(self): + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as forig: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + with forig.clone_device() as f: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + data = list(f.yield_blocks(list(xrange(self._block_count)))) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + data = list(f.yield_blocks([0])) + self.assertEqual(len(data), 1) + self.assertEqual(list(bytearray(data[0])), + list(self._blocks[0])) + self.assertEqual(len(self._blocks) > 1, True) + data = list(f.yield_blocks(list(xrange(1, self._block_count)) + [0])) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data[:-1], 1): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + self.assertEqual(list(bytearray(data[-1])), + list(self._blocks[0])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + (2*self._block_count + 1)*f._storage.block_size) + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + + def test_write_blocks_cloned(self): + data = [bytearray([self._block_count])*self._block_size + for i in xrange(self._block_count)] + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as forig: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + with forig.clone_device() as f: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + orig = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(orig), self._block_count) + for i, block in enumerate(orig): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + f.write_blocks(list(xrange(self._block_count)), + [bytes(b) for b in data]) + new = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(new), self._block_count) + for i, block in enumerate(new): + self.assertEqual(list(bytearray(block)), + list(data[i])) + f.write_blocks(list(xrange(self._block_count)), + [bytes(b) for b in self._blocks]) + orig = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(orig), self._block_count) + for i, block in enumerate(orig): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + + self.assertEqual(f.bytes_sent, + self._block_count*f._storage.block_size*2) + self.assertEqual(f.bytes_received, + self._block_count*f._storage.block_size*3) + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + +class TestEncryptedBlockStorageFileCTRKey(_TestEncryptedBlockStorage, + unittest2.TestCase): + _type_name = 'file' + _aes_mode = 'ctr' + _test_key = AES.KeyGen(16) + +class TestEncryptedBlockStorageFileCTR32(_TestEncryptedBlockStorage, + unittest2.TestCase): + _type_name = 'file' + _aes_mode = 'ctr' + _test_key_size = 16 + +class TestEncryptedBlockStorageFileGCMKey(_TestEncryptedBlockStorage, + unittest2.TestCase): + _type_name = 'file' + _aes_mode = 'gcm' + _test_key = AES.KeyGen(24) + +class TestEncryptedBlockStorageFileGCM32(_TestEncryptedBlockStorage, + unittest2.TestCase): + _type_name = 'file' + _aes_mode = 'gcm' + _test_key_size = 24 + +class TestEncryptedBlockStorageMMapFileCTRKey(_TestEncryptedBlockStorage, + unittest2.TestCase): + _type_name = 'mmap' + _aes_mode = 'ctr' + _test_key = AES.KeyGen(32) + +class TestEncryptedBlockStorageMMapFileCTR32(_TestEncryptedBlockStorage, + unittest2.TestCase): + _type_name = 'mmap' + _aes_mode = 'ctr' + _test_key_size = 32 + +class TestEncryptedBlockStorageMMapFileGCMKey(_TestEncryptedBlockStorage, + unittest2.TestCase): + _type_name = 'mmap' + _aes_mode = 'gcm' + +class TestEncryptedBlockStorageMMapFileGCM32(_TestEncryptedBlockStorage, + unittest2.TestCase): + _type_name = 'mmap' + _aes_mode = 'gcm' + +if __name__ == "__main__": + unittest2.main() # pragma: no cover diff --git a/PyORAM/src/pyoram/tests/test_encrypted_heap_storage.py b/PyORAM/src/pyoram/tests/test_encrypted_heap_storage.py new file mode 100644 index 0000000..a47acf2 --- /dev/null +++ b/PyORAM/src/pyoram/tests/test_encrypted_heap_storage.py @@ -0,0 +1,594 @@ +import os +import unittest2 +import tempfile + +from pyoram.util.virtual_heap import \ + SizedVirtualHeap +from pyoram.storage.block_storage import \ + BlockStorageTypeFactory +from pyoram.encrypted_storage.encrypted_block_storage import \ + EncryptedBlockStorage +from pyoram.encrypted_storage.encrypted_heap_storage import \ + EncryptedHeapStorage +from pyoram.crypto.aes import AES + +from six.moves import xrange + +thisdir = os.path.dirname(os.path.abspath(__file__)) + +class TestEncryptedHeapStorage(unittest2.TestCase): + + @classmethod + def setUpClass(cls): + fd, cls._dummy_name = tempfile.mkstemp() + os.close(fd) + try: + os.remove(cls._dummy_name) + except OSError: # pragma: no cover + pass # pragma: no cover + cls._block_size = 25 + cls._blocks_per_bucket = 3 + cls._heap_base = 4 + cls._heap_height = 2 + cls._bucket_count = \ + ((cls._heap_base**(cls._heap_height+1)) - 1)//(cls._heap_base-1) + cls._block_count = cls._bucket_count * \ + cls._blocks_per_bucket + cls._testfname = cls.__name__ + "_testfile.bin" + cls._buckets = [] + cls._type_name = "file" + f = EncryptedHeapStorage.setup( + cls._testfname, + block_size=cls._block_size, + heap_height=cls._heap_height, + key_size=AES.key_sizes[-1], + heap_base=cls._heap_base, + blocks_per_bucket=cls._blocks_per_bucket, + storage_type=cls._type_name, + initialize=lambda i: bytes(bytearray([i]) * \ + cls._block_size * \ + cls._blocks_per_bucket), + ignore_existing=True) + f.close() + cls._key = f.key + for i in range(cls._bucket_count): + data = bytearray([i]) * \ + cls._block_size * \ + cls._blocks_per_bucket + cls._buckets.append(data) + + @classmethod + def tearDownClass(cls): + try: + os.remove(cls._testfname) + except OSError: # pragma: no cover + pass # pragma: no cover + try: + os.remove(cls._dummy_name) + except OSError: # pragma: no cover + pass # pragma: no cover + + def test_setup_fails(self): + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(IOError): + EncryptedHeapStorage.setup( + os.path.join(thisdir, + "baselines", + "exists.empty"), + block_size=10, + heap_height=1, + key_size=AES.key_sizes[-1], + blocks_per_bucket=1, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(IOError): + EncryptedHeapStorage.setup( + os.path.join(thisdir, + "baselines", + "exists.empty"), + block_size=10, + heap_height=1, + key_size=AES.key_sizes[-1], + blocks_per_bucket=1, + storage_type=self._type_name, + ignore_existing=False) + self.assertEqual(os.path.exists(self._dummy_name), False) + # bad block_size + with self.assertRaises(ValueError): + EncryptedHeapStorage.setup( + self._dummy_name, + block_size=0, + heap_height=1, + key_size=AES.key_sizes[-1], + blocks_per_bucket=1, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + # bad heap_height + with self.assertRaises(ValueError): + EncryptedHeapStorage.setup( + self._dummy_name, + block_size=1, + heap_height=-1, + blocks_per_bucket=1, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + # bad blocks_per_bucket + with self.assertRaises(ValueError): + EncryptedHeapStorage.setup( + self._dummy_name, + block_size=1, + heap_height=1, + key_size=AES.key_sizes[-1], + blocks_per_bucket=0, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + # bad heap_base + with self.assertRaises(ValueError): + EncryptedHeapStorage.setup( + self._dummy_name, + block_size=1, + heap_height=1, + key_size=AES.key_sizes[-1], + blocks_per_bucket=1, + heap_base=1, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + # bad header_data + with self.assertRaises(TypeError): + EncryptedHeapStorage.setup( + self._dummy_name, + block_size=1, + heap_height=1, + key_size=AES.key_sizes[-1], + blocks_per_bucket=1, + storage_type=self._type_name, + header_data=2) + self.assertEqual(os.path.exists(self._dummy_name), False) + # uses block_count + with self.assertRaises(ValueError): + EncryptedHeapStorage.setup( + self._dummy_name, + block_size=1, + heap_height=1, + key_size=AES.key_sizes[-1], + blocks_per_bucket=1, + block_count=1, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + + def test_setup(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + if os.path.exists(fname): + os.remove(fname) # pragma: no cover + bsize = 10 + heap_height = 2 + blocks_per_bucket = 3 + fsetup = EncryptedHeapStorage.setup( + fname, + bsize, + heap_height, + key_size=AES.key_sizes[-1], + blocks_per_bucket=blocks_per_bucket) + fsetup.close() + self.assertEqual(type(fsetup.raw_storage), + BlockStorageTypeFactory(self._type_name)) + with open(fname, 'rb') as f: + flen = len(f.read()) + self.assertEqual( + flen, + EncryptedHeapStorage.compute_storage_size( + bsize, + heap_height, + blocks_per_bucket=blocks_per_bucket)) + self.assertEqual( + flen > + EncryptedHeapStorage.compute_storage_size( + bsize, + heap_height, + blocks_per_bucket=blocks_per_bucket, + ignore_header=True), + True) + with EncryptedHeapStorage( + fname, + key=fsetup.key, + storage_type=self._type_name) as f: + self.assertEqual(f.header_data, bytes()) + self.assertEqual(fsetup.header_data, bytes()) + self.assertEqual(f.key, fsetup.key) + self.assertEqual(f.blocks_per_bucket, + blocks_per_bucket) + self.assertEqual(fsetup.blocks_per_bucket, + blocks_per_bucket) + self.assertEqual(f.bucket_count, + 2**(heap_height+1) - 1) + self.assertEqual(fsetup.bucket_count, + 2**(heap_height+1) - 1) + self.assertEqual(f.bucket_size, + bsize * blocks_per_bucket) + self.assertEqual(fsetup.bucket_size, + bsize * blocks_per_bucket) + self.assertEqual(f.storage_name, fname) + self.assertEqual(fsetup.storage_name, fname) + os.remove(fname) + + def test_setup_withdata(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + if os.path.exists(fname): + os.remove(fname) # pragma: no cover + bsize = 10 + heap_height = 2 + blocks_per_bucket = 1 + header_data = bytes(bytearray([0,1,2])) + fsetup = EncryptedHeapStorage.setup( + fname, + bsize, + heap_height, + key_size=AES.key_sizes[-1], + blocks_per_bucket=blocks_per_bucket, + header_data=header_data) + fsetup.close() + self.assertEqual(type(fsetup.raw_storage), + BlockStorageTypeFactory(self._type_name)) + with open(fname, 'rb') as f: + flen = len(f.read()) + self.assertEqual( + flen, + EncryptedHeapStorage.compute_storage_size( + bsize, + heap_height, + header_data=header_data)) + self.assertTrue(len(header_data) > 0) + self.assertEqual( + EncryptedHeapStorage.compute_storage_size( + bsize, + heap_height, + storage_type=self._type_name) < + EncryptedHeapStorage.compute_storage_size( + bsize, + heap_height, + storage_type=self._type_name, + header_data=header_data), + True) + self.assertEqual( + flen > + EncryptedHeapStorage.compute_storage_size( + bsize, + heap_height, + storage_type=self._type_name, + header_data=header_data, + ignore_header=True), + True) + with EncryptedHeapStorage( + fname, + key=fsetup.key, + storage_type=self._type_name) as f: + self.assertEqual(f.header_data, header_data) + self.assertEqual(fsetup.header_data, header_data) + self.assertEqual(f.key, fsetup.key) + self.assertEqual(f.blocks_per_bucket, + blocks_per_bucket) + self.assertEqual(fsetup.blocks_per_bucket, + blocks_per_bucket) + self.assertEqual(f.bucket_count, + 2**(heap_height+1) - 1) + self.assertEqual(fsetup.bucket_count, + 2**(heap_height+1) - 1) + self.assertEqual(f.bucket_size, + bsize * blocks_per_bucket) + self.assertEqual(fsetup.bucket_size, + bsize * blocks_per_bucket) + self.assertEqual(f.storage_name, fname) + self.assertEqual(fsetup.storage_name, fname) + os.remove(fname) + + def test_init_noexists(self): + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(IOError): + with EncryptedHeapStorage( + self._dummy_name, + key=self._key, + storage_type=self._type_name) as f: + pass # pragma: no cover + + def test_init_exists(self): + self.assertEqual(os.path.exists(self._testfname), True) + with open(self._testfname, 'rb') as f: + databefore = f.read() + with self.assertRaises(ValueError): + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as fb: + with EncryptedHeapStorage(fb, key=self._key) as f: + pass # pragma: no cover + with EncryptedHeapStorage( + self._testfname, + key=self._key, + storage_type=self._type_name) as f: + self.assertEqual(f.key, self._key) + self.assertEqual(f.bucket_size, + self._block_size * \ + self._blocks_per_bucket) + self.assertEqual(f.bucket_count, + self._bucket_count) + self.assertEqual(f.storage_name, self._testfname) + self.assertEqual(f.header_data, bytes()) + self.assertEqual(os.path.exists(self._testfname), True) + with open(self._testfname, 'rb') as f: + dataafter = f.read() + self.assertEqual(databefore, dataafter) + + def test_read_path(self): + + with EncryptedHeapStorage( + self._testfname, + key=self._key, + storage_type=self._type_name) as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + self.assertEqual( + f.virtual_heap.first_bucket_at_level(0), 0) + self.assertNotEqual( + f.virtual_heap.last_leaf_bucket(), 0) + total_buckets = 0 + for b in range(f.virtual_heap.first_bucket_at_level(0), + f.virtual_heap.last_leaf_bucket()+1): + data = f.read_path(b) + bucket_path = f.virtual_heap.Node(b).\ + bucket_path_from_root() + total_buckets += len(bucket_path) + self.assertEqual(f.virtual_heap.Node(b).level+1, + len(bucket_path)) + for i, bucket in zip(bucket_path, data): + self.assertEqual(list(bytearray(bucket)), + list(self._buckets[i])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + total_buckets*f.bucket_storage._storage.block_size) + + def test_write_path(self): + data = [bytearray([self._bucket_count]) * \ + self._block_size * \ + self._blocks_per_bucket + for i in xrange(self._block_count)] + with EncryptedHeapStorage( + self._testfname, + key=self._key, + storage_type=self._type_name) as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + self.assertEqual( + f.virtual_heap.first_bucket_at_level(0), 0) + self.assertNotEqual( + f.virtual_heap.last_leaf_bucket(), 0) + total_buckets = 0 + for b in range(f.virtual_heap.first_bucket_at_level(0), + f.virtual_heap.last_leaf_bucket()+1): + orig = f.read_path(b) + bucket_path = f.virtual_heap.Node(b).\ + bucket_path_from_root() + total_buckets += len(bucket_path) + self.assertNotEqual(len(bucket_path), 0) + self.assertEqual(f.virtual_heap.Node(b).level+1, + len(bucket_path)) + self.assertEqual(len(orig), len(bucket_path)) + for i, bucket in zip(bucket_path, orig): + self.assertEqual(list(bytearray(bucket)), + list(self._buckets[i])) + f.write_path(b, [bytes(data[i]) + for i in bucket_path]) + + new = f.read_path(b) + self.assertEqual(len(new), len(bucket_path)) + for i, bucket in zip(bucket_path, new): + self.assertEqual(list(bytearray(bucket)), + list(data[i])) + + f.write_path(b, [bytes(self._buckets[i]) + for i in bucket_path]) + + orig = f.read_path(b) + self.assertEqual(len(orig), len(bucket_path)) + for i, bucket in zip(bucket_path, orig): + self.assertEqual(list(bytearray(bucket)), + list(self._buckets[i])) + + self.assertEqual(f.bytes_sent, + total_buckets*f.bucket_storage._storage.block_size*2) + self.assertEqual(f.bytes_received, + total_buckets*f.bucket_storage._storage.block_size*3) + + def test_update_header_data(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + if os.path.exists(fname): + os.remove(fname) # pragma: no cover + bsize = 10 + heap_height = 2 + blocks_per_bucket = 1 + header_data = bytes(bytearray([0,1,2])) + fsetup = EncryptedHeapStorage.setup( + fname, + block_size=bsize, + heap_height=heap_height, + key_size=AES.key_sizes[-1], + blocks_per_bucket=blocks_per_bucket, + header_data=header_data) + fsetup.close() + new_header_data = bytes(bytearray([1,1,1])) + with EncryptedHeapStorage( + fname, + key=fsetup.key, + storage_type=self._type_name) as f: + self.assertEqual(f.header_data, header_data) + f.update_header_data(new_header_data) + self.assertEqual(f.header_data, new_header_data) + with EncryptedHeapStorage( + fname, + key=fsetup.key, + storage_type=self._type_name) as f: + self.assertEqual(f.header_data, new_header_data) + with self.assertRaises(ValueError): + with EncryptedHeapStorage( + fname, + key=fsetup.key, + storage_type=self._type_name) as f: + f.update_header_data(bytes(bytearray([1,1]))) + with self.assertRaises(ValueError): + with EncryptedHeapStorage( + fname, + key=fsetup.key, + storage_type=self._type_name) as f: + f.update_header_data(bytes(bytearray([1,1,1,1]))) + with EncryptedHeapStorage( + fname, + key=fsetup.key, + storage_type=self._type_name) as f: + self.assertEqual(f.header_data, new_header_data) + os.remove(fname) + + def test_locked_flag(self): + with EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as f: + with self.assertRaises(IOError): + with EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as f1: + pass # pragma: no cover + with self.assertRaises(IOError): + with EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as f1: + pass # pragma: no cover + with EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._type_name, + ignore_lock=True) as f1: + pass + with self.assertRaises(IOError): + with EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as f1: + pass # pragma: no cover + with EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._type_name, + ignore_lock=True) as f1: + pass + with EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._type_name, + ignore_lock=True) as f1: + pass + with EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as f: + pass + + def test_read_path_cloned(self): + + with EncryptedHeapStorage( + self._testfname, + key=self._key, + storage_type=self._type_name) as forig: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + with forig.clone_device() as f: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + self.assertEqual( + f.virtual_heap.first_bucket_at_level(0), 0) + self.assertNotEqual( + f.virtual_heap.last_leaf_bucket(), 0) + total_buckets = 0 + for b in range(f.virtual_heap.first_bucket_at_level(0), + f.virtual_heap.last_leaf_bucket()+1): + data = f.read_path(b) + bucket_path = f.virtual_heap.Node(b).\ + bucket_path_from_root() + total_buckets += len(bucket_path) + self.assertEqual(f.virtual_heap.Node(b).level+1, + len(bucket_path)) + for i, bucket in zip(bucket_path, data): + self.assertEqual(list(bytearray(bucket)), + list(self._buckets[i])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + total_buckets*f.bucket_storage._storage.block_size) + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + + def test_write_path_cloned(self): + data = [bytearray([self._bucket_count]) * \ + self._block_size * \ + self._blocks_per_bucket + for i in xrange(self._block_count)] + with EncryptedHeapStorage( + self._testfname, + key=self._key, + storage_type=self._type_name) as forig: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + with forig.clone_device() as f: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + self.assertEqual( + f.virtual_heap.first_bucket_at_level(0), 0) + self.assertNotEqual( + f.virtual_heap.last_leaf_bucket(), 0) + total_buckets = 0 + for b in range(f.virtual_heap.first_bucket_at_level(0), + f.virtual_heap.last_leaf_bucket()+1): + orig = f.read_path(b) + bucket_path = f.virtual_heap.Node(b).\ + bucket_path_from_root() + total_buckets += len(bucket_path) + self.assertNotEqual(len(bucket_path), 0) + self.assertEqual(f.virtual_heap.Node(b).level+1, + len(bucket_path)) + self.assertEqual(len(orig), len(bucket_path)) + for i, bucket in zip(bucket_path, orig): + self.assertEqual(list(bytearray(bucket)), + list(self._buckets[i])) + f.write_path(b, [bytes(data[i]) + for i in bucket_path]) + + new = f.read_path(b) + self.assertEqual(len(new), len(bucket_path)) + for i, bucket in zip(bucket_path, new): + self.assertEqual(list(bytearray(bucket)), + list(data[i])) + + f.write_path(b, [bytes(self._buckets[i]) + for i in bucket_path]) + + orig = f.read_path(b) + self.assertEqual(len(orig), len(bucket_path)) + for i, bucket in zip(bucket_path, orig): + self.assertEqual(list(bytearray(bucket)), + list(self._buckets[i])) + + self.assertEqual(f.bytes_sent, + total_buckets*f.bucket_storage._storage.block_size*2) + self.assertEqual(f.bytes_received, + total_buckets*f.bucket_storage._storage.block_size*3) + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + +if __name__ == "__main__": + unittest2.main() # pragma: no cover diff --git a/PyORAM/src/pyoram/tests/test_examples.py b/PyORAM/src/pyoram/tests/test_examples.py new file mode 100644 index 0000000..9446c40 --- /dev/null +++ b/PyORAM/src/pyoram/tests/test_examples.py @@ -0,0 +1,61 @@ +import os +import glob +import sys +import unittest2 + +thisfile = os.path.abspath(__file__) +thisdir = os.path.dirname(thisfile) +topdir = os.path.dirname( + os.path.dirname( + os.path.dirname(thisdir))) +exdir = os.path.join(topdir, 'examples') +examples = glob.glob(os.path.join(exdir,"*.py")) + +assert os.path.exists(exdir) +assert thisfile not in examples + +tdict = {} +for fname in examples: + basename = os.path.basename(fname) + assert basename.endswith('.py') + assert len(basename) >= 3 + basename = basename[:-3] + tname = 'test_'+basename + tdict[tname] = fname, basename + +assert len(tdict) == len(examples) + +assert 'test_encrypted_storage_s3' in tdict +assert 'test_path_oram_s3' in tdict +if 'PYORAM_AWS_TEST_BUCKET' not in os.environ: + del tdict['test_encrypted_storage_s3'] + del tdict['test_path_oram_s3'] +assert 'test_encrypted_storage_sftp' in tdict +assert 'test_path_oram_sftp' in tdict +assert 'test_path_oram_sftp_setup' in tdict +assert 'test_path_oram_sftp_test' in tdict +if 'PYORAM_SSH_TEST_HOST' not in os.environ: + del tdict['test_encrypted_storage_sftp'] + del tdict['test_path_oram_sftp'] + del tdict['test_path_oram_sftp_setup'] + del tdict['test_path_oram_sftp_test'] + +def _execute_example(example_name): + filename, basename = tdict[example_name] + assert os.path.exists(filename) + try: + sys.path.insert(0, exdir) + m = __import__(basename) + m.main() + finally: + sys.path.remove(exdir) + +# this is recognized by nosetests as +# a dynamic test generator +def test_generator(): + for example_name in sorted(tdict): + yield _execute_example, example_name + +if __name__ == "__main__": + for tfunc, tname in test_generator(): # pragma: no cover + tfunc(tname) # pragma: no cover diff --git a/PyORAM/src/pyoram/tests/test_heap_storage.py b/PyORAM/src/pyoram/tests/test_heap_storage.py new file mode 100644 index 0000000..2519a00 --- /dev/null +++ b/PyORAM/src/pyoram/tests/test_heap_storage.py @@ -0,0 +1,555 @@ +import os +import unittest2 +import tempfile + +from pyoram.util.virtual_heap import \ + SizedVirtualHeap +from pyoram.storage.block_storage import \ + BlockStorageTypeFactory +from pyoram.storage.block_storage_file import \ + BlockStorageFile +from pyoram.storage.heap_storage import \ + HeapStorage + +from six.moves import xrange + +thisdir = os.path.dirname(os.path.abspath(__file__)) + +class TestHeapStorage(unittest2.TestCase): + + @classmethod + def setUpClass(cls): + fd, cls._dummy_name = tempfile.mkstemp() + os.close(fd) + try: + os.remove(cls._dummy_name) + except OSError: # pragma: no cover + pass # pragma: no cover + cls._block_size = 25 + cls._blocks_per_bucket = 3 + cls._heap_base = 4 + cls._heap_height = 2 + cls._bucket_count = \ + ((cls._heap_base**(cls._heap_height+1)) - 1)//(cls._heap_base-1) + cls._block_count = cls._bucket_count * \ + cls._blocks_per_bucket + cls._testfname = cls.__name__ + "_testfile.bin" + cls._buckets = [] + cls._type_name = "file" + f = HeapStorage.setup( + cls._testfname, + block_size=cls._block_size, + heap_height=cls._heap_height, + heap_base=cls._heap_base, + blocks_per_bucket=cls._blocks_per_bucket, + storage_type=cls._type_name, + initialize=lambda i: bytes(bytearray([i]) * \ + cls._block_size * \ + cls._blocks_per_bucket), + ignore_existing=True) + f.close() + for i in range(cls._bucket_count): + data = bytearray([i]) * \ + cls._block_size * \ + cls._blocks_per_bucket + cls._buckets.append(data) + + @classmethod + def tearDownClass(cls): + try: + os.remove(cls._testfname) + except OSError: # pragma: no cover + pass # pragma: no cover + try: + os.remove(cls._dummy_name) + except OSError: # pragma: no cover + pass # pragma: no cover + + def test_setup_fails(self): + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(IOError): + HeapStorage.setup( + os.path.join(thisdir, + "baselines", + "exists.empty"), + block_size=10, + heap_height=1, + blocks_per_bucket=1, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(IOError): + HeapStorage.setup( + os.path.join(thisdir, + "baselines", + "exists.empty"), + block_size=10, + heap_height=1, + blocks_per_bucket=1, + storage_type=self._type_name, + ignore_existing=False) + self.assertEqual(os.path.exists(self._dummy_name), False) + # bad block_size + with self.assertRaises(ValueError): + HeapStorage.setup( + self._dummy_name, + block_size=0, + heap_height=1, + blocks_per_bucket=1, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + # bad heap_height + with self.assertRaises(ValueError): + HeapStorage.setup( + self._dummy_name, + block_size=1, + heap_height=-1, + blocks_per_bucket=1, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + # bad blocks_per_bucket + with self.assertRaises(ValueError): + HeapStorage.setup( + self._dummy_name, + block_size=1, + heap_height=1, + blocks_per_bucket=0, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + # bad heap_base + with self.assertRaises(ValueError): + HeapStorage.setup( + self._dummy_name, + block_size=1, + heap_height=1, + blocks_per_bucket=1, + heap_base=1, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + # bad header_data + with self.assertRaises(TypeError): + HeapStorage.setup( + self._dummy_name, + block_size=1, + heap_height=1, + blocks_per_bucket=1, + storage_type=self._type_name, + header_data=2) + self.assertEqual(os.path.exists(self._dummy_name), False) + # uses block_count + with self.assertRaises(ValueError): + HeapStorage.setup( + self._dummy_name, + block_size=1, + heap_height=1, + blocks_per_bucket=1, + block_count=1, + storage_type=self._type_name) + self.assertEqual(os.path.exists(self._dummy_name), False) + + def test_setup(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + if os.path.exists(fname): + os.remove(fname) # pragma: no cover + bsize = 10 + heap_height = 2 + blocks_per_bucket = 3 + fsetup = HeapStorage.setup( + fname, + bsize, + heap_height, + blocks_per_bucket=blocks_per_bucket) + fsetup.close() + self.assertEqual(type(fsetup.bucket_storage), + BlockStorageTypeFactory(self._type_name)) + with open(fname, 'rb') as f: + flen = len(f.read()) + self.assertEqual( + flen, + HeapStorage.compute_storage_size( + bsize, + heap_height, + blocks_per_bucket=blocks_per_bucket)) + self.assertEqual( + flen > + HeapStorage.compute_storage_size( + bsize, + heap_height, + blocks_per_bucket=blocks_per_bucket, + ignore_header=True), + True) + with HeapStorage( + fname, + storage_type=self._type_name) as f: + self.assertEqual(f.header_data, bytes()) + self.assertEqual(fsetup.header_data, bytes()) + self.assertEqual(f.blocks_per_bucket, + blocks_per_bucket) + self.assertEqual(fsetup.blocks_per_bucket, + blocks_per_bucket) + self.assertEqual(f.bucket_count, + 2**(heap_height+1) - 1) + self.assertEqual(fsetup.bucket_count, + 2**(heap_height+1) - 1) + self.assertEqual(f.bucket_size, + bsize * blocks_per_bucket) + self.assertEqual(fsetup.bucket_size, + bsize * blocks_per_bucket) + self.assertEqual(f.storage_name, fname) + self.assertEqual(fsetup.storage_name, fname) + os.remove(fname) + + def test_setup_withdata(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + if os.path.exists(fname): + os.remove(fname) # pragma: no cover + bsize = 10 + heap_height = 2 + blocks_per_bucket = 1 + header_data = bytes(bytearray([0,1,2])) + fsetup = HeapStorage.setup( + fname, + bsize, + heap_height, + blocks_per_bucket=blocks_per_bucket, + header_data=header_data) + fsetup.close() + self.assertEqual(type(fsetup.bucket_storage), + BlockStorageTypeFactory(self._type_name)) + with open(fname, 'rb') as f: + flen = len(f.read()) + self.assertEqual( + flen, + HeapStorage.compute_storage_size( + bsize, + heap_height, + header_data=header_data)) + self.assertTrue(len(header_data) > 0) + self.assertEqual( + HeapStorage.compute_storage_size( + bsize, + heap_height, + storage_type=self._type_name) < + HeapStorage.compute_storage_size( + bsize, + heap_height, + storage_type=self._type_name, + header_data=header_data), + True) + self.assertEqual( + flen > + HeapStorage.compute_storage_size( + bsize, + heap_height, + storage_type=self._type_name, + header_data=header_data, + ignore_header=True), + True) + with HeapStorage( + fname, + storage_type=self._type_name) as f: + self.assertEqual(f.header_data, header_data) + self.assertEqual(fsetup.header_data, header_data) + self.assertEqual(f.blocks_per_bucket, + blocks_per_bucket) + self.assertEqual(fsetup.blocks_per_bucket, + blocks_per_bucket) + self.assertEqual(f.bucket_count, + 2**(heap_height+1) - 1) + self.assertEqual(fsetup.bucket_count, + 2**(heap_height+1) - 1) + self.assertEqual(f.bucket_size, + bsize * blocks_per_bucket) + self.assertEqual(fsetup.bucket_size, + bsize * blocks_per_bucket) + self.assertEqual(f.storage_name, fname) + self.assertEqual(fsetup.storage_name, fname) + os.remove(fname) + + def test_init_noexists(self): + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(IOError): + with HeapStorage( + self._dummy_name, + storage_type=self._type_name) as f: + pass # pragma: no cover + + def test_init_exists(self): + self.assertEqual(os.path.exists(self._testfname), True) + with open(self._testfname, 'rb') as f: + databefore = f.read() + with self.assertRaises(ValueError): + with BlockStorageFile(self._testfname) as fb: + with HeapStorage(fb, storage_type='file') as f: + pass # pragma: no cover + with HeapStorage( + self._testfname, + storage_type=self._type_name) as f: + self.assertEqual(f.bucket_size, + self._block_size * \ + self._blocks_per_bucket) + self.assertEqual(f.bucket_count, + self._bucket_count) + self.assertEqual(f.storage_name, self._testfname) + self.assertEqual(f.header_data, bytes()) + self.assertEqual(os.path.exists(self._testfname), True) + with open(self._testfname, 'rb') as f: + dataafter = f.read() + self.assertEqual(databefore, dataafter) + + def test_read_path(self): + + with HeapStorage( + self._testfname, + storage_type=self._type_name) as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + self.assertEqual( + f.virtual_heap.first_bucket_at_level(0), 0) + self.assertNotEqual( + f.virtual_heap.last_leaf_bucket(), 0) + total_buckets = 0 + for b in range(f.virtual_heap.first_bucket_at_level(0), + f.virtual_heap.last_leaf_bucket()+1): + data = f.read_path(b) + bucket_path = f.virtual_heap.Node(b).\ + bucket_path_from_root() + total_buckets += len(bucket_path) + self.assertEqual(f.virtual_heap.Node(b).level+1, + len(bucket_path)) + for i, bucket in zip(bucket_path, data): + self.assertEqual(list(bytearray(bucket)), + list(self._buckets[i])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + total_buckets*f.bucket_storage.block_size) + + def test_write_path(self): + data = [bytearray([self._bucket_count]) * \ + self._block_size * \ + self._blocks_per_bucket + for i in xrange(self._block_count)] + with HeapStorage( + self._testfname, + storage_type=self._type_name) as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + self.assertEqual( + f.virtual_heap.first_bucket_at_level(0), 0) + self.assertNotEqual( + f.virtual_heap.last_leaf_bucket(), 0) + total_buckets = 0 + for b in range(f.virtual_heap.first_bucket_at_level(0), + f.virtual_heap.last_leaf_bucket()+1): + orig = f.read_path(b) + bucket_path = f.virtual_heap.Node(b).\ + bucket_path_from_root() + total_buckets += len(bucket_path) + self.assertNotEqual(len(bucket_path), 0) + self.assertEqual(f.virtual_heap.Node(b).level+1, + len(bucket_path)) + self.assertEqual(len(orig), len(bucket_path)) + for i, bucket in zip(bucket_path, orig): + self.assertEqual(list(bytearray(bucket)), + list(self._buckets[i])) + f.write_path(b, [bytes(data[i]) + for i in bucket_path]) + + new = f.read_path(b) + self.assertEqual(len(new), len(bucket_path)) + for i, bucket in zip(bucket_path, new): + self.assertEqual(list(bytearray(bucket)), + list(data[i])) + + f.write_path(b, [bytes(self._buckets[i]) + for i in bucket_path]) + + orig = f.read_path(b) + self.assertEqual(len(orig), len(bucket_path)) + for i, bucket in zip(bucket_path, orig): + self.assertEqual(list(bytearray(bucket)), + list(self._buckets[i])) + + self.assertEqual(f.bytes_sent, + total_buckets*f.bucket_storage.block_size*2) + self.assertEqual(f.bytes_received, + total_buckets*f.bucket_storage.block_size*3) + + def test_update_header_data(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + if os.path.exists(fname): + os.remove(fname) # pragma: no cover + bsize = 10 + heap_height = 2 + blocks_per_bucket = 1 + header_data = bytes(bytearray([0,1,2])) + fsetup = HeapStorage.setup( + fname, + block_size=bsize, + heap_height=heap_height, + blocks_per_bucket=blocks_per_bucket, + header_data=header_data) + fsetup.close() + new_header_data = bytes(bytearray([1,1,1])) + with HeapStorage( + fname, + storage_type=self._type_name) as f: + self.assertEqual(f.header_data, header_data) + f.update_header_data(new_header_data) + self.assertEqual(f.header_data, new_header_data) + with HeapStorage( + fname, + storage_type=self._type_name) as f: + self.assertEqual(f.header_data, new_header_data) + with self.assertRaises(ValueError): + with HeapStorage( + fname, + storage_type=self._type_name) as f: + f.update_header_data(bytes(bytearray([1,1]))) + with self.assertRaises(ValueError): + with HeapStorage( + fname, + storage_type=self._type_name) as f: + f.update_header_data(bytes(bytearray([1,1,1,1]))) + with HeapStorage( + fname, + storage_type=self._type_name) as f: + self.assertEqual(f.header_data, new_header_data) + os.remove(fname) + + def test_locked_flag(self): + with HeapStorage(self._testfname, + storage_type=self._type_name) as f: + with self.assertRaises(IOError): + with HeapStorage(self._testfname, + storage_type=self._type_name) as f1: + pass # pragma: no cover + with self.assertRaises(IOError): + with HeapStorage(self._testfname, + storage_type=self._type_name) as f1: + pass # pragma: no cover + with HeapStorage(self._testfname, + storage_type=self._type_name, + ignore_lock=True) as f1: + pass + with self.assertRaises(IOError): + with HeapStorage(self._testfname, + storage_type=self._type_name) as f1: + pass # pragma: no cover + with HeapStorage(self._testfname, + storage_type=self._type_name, + ignore_lock=True) as f1: + pass + with HeapStorage(self._testfname, + storage_type=self._type_name, + ignore_lock=True) as f1: + pass + with HeapStorage(self._testfname, + storage_type=self._type_name) as f: + pass + + def test_read_path_cloned(self): + + with HeapStorage( + self._testfname, + storage_type=self._type_name) as forig: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + with forig.clone_device() as f: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + self.assertEqual( + f.virtual_heap.first_bucket_at_level(0), 0) + self.assertNotEqual( + f.virtual_heap.last_leaf_bucket(), 0) + total_buckets = 0 + for b in range(f.virtual_heap.first_bucket_at_level(0), + f.virtual_heap.last_leaf_bucket()+1): + data = f.read_path(b) + bucket_path = f.virtual_heap.Node(b).\ + bucket_path_from_root() + total_buckets += len(bucket_path) + self.assertEqual(f.virtual_heap.Node(b).level+1, + len(bucket_path)) + for i, bucket in zip(bucket_path, data): + self.assertEqual(list(bytearray(bucket)), + list(self._buckets[i])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + total_buckets*f.bucket_storage.block_size) + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + + def test_write_path_cloned(self): + data = [bytearray([self._bucket_count]) * \ + self._block_size * \ + self._blocks_per_bucket + for i in xrange(self._block_count)] + with HeapStorage( + self._testfname, + storage_type=self._type_name) as forig: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + with forig.clone_device() as f: + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + self.assertEqual( + f.virtual_heap.first_bucket_at_level(0), 0) + self.assertNotEqual( + f.virtual_heap.last_leaf_bucket(), 0) + total_buckets = 0 + for b in range(f.virtual_heap.first_bucket_at_level(0), + f.virtual_heap.last_leaf_bucket()+1): + orig = f.read_path(b) + bucket_path = f.virtual_heap.Node(b).\ + bucket_path_from_root() + total_buckets += len(bucket_path) + self.assertNotEqual(len(bucket_path), 0) + self.assertEqual(f.virtual_heap.Node(b).level+1, + len(bucket_path)) + self.assertEqual(len(orig), len(bucket_path)) + for i, bucket in zip(bucket_path, orig): + self.assertEqual(list(bytearray(bucket)), + list(self._buckets[i])) + f.write_path(b, [bytes(data[i]) + for i in bucket_path]) + + new = f.read_path(b) + self.assertEqual(len(new), len(bucket_path)) + for i, bucket in zip(bucket_path, new): + self.assertEqual(list(bytearray(bucket)), + list(data[i])) + + f.write_path(b, [bytes(self._buckets[i]) + for i in bucket_path]) + + orig = f.read_path(b) + self.assertEqual(len(orig), len(bucket_path)) + for i, bucket in zip(bucket_path, orig): + self.assertEqual(list(bytearray(bucket)), + list(self._buckets[i])) + + self.assertEqual(f.bytes_sent, + total_buckets*f.bucket_storage.block_size*2) + self.assertEqual(f.bytes_received, + total_buckets*f.bucket_storage.block_size*3) + self.assertEqual(forig.bytes_sent, 0) + self.assertEqual(forig.bytes_received, 0) + +if __name__ == "__main__": + unittest2.main() # pragma: no cover diff --git a/PyORAM/src/pyoram/tests/test_misc.py b/PyORAM/src/pyoram/tests/test_misc.py new file mode 100644 index 0000000..65b7660 --- /dev/null +++ b/PyORAM/src/pyoram/tests/test_misc.py @@ -0,0 +1,124 @@ +import os +import unittest2 +import tempfile + +import pyoram.util.misc + +class Test(unittest2.TestCase): + + def test_log2floor(self): + self.assertEqual(pyoram.util.misc.log2floor(1), 0) + self.assertEqual(pyoram.util.misc.log2floor(2), 1) + self.assertEqual(pyoram.util.misc.log2floor(3), 1) + self.assertEqual(pyoram.util.misc.log2floor(4), 2) + self.assertEqual(pyoram.util.misc.log2floor(5), 2) + self.assertEqual(pyoram.util.misc.log2floor(6), 2) + self.assertEqual(pyoram.util.misc.log2floor(7), 2) + self.assertEqual(pyoram.util.misc.log2floor(8), 3) + self.assertEqual(pyoram.util.misc.log2floor(9), 3) + + def test_log2ceil(self): + self.assertEqual(pyoram.util.misc.log2ceil(1), 0) + self.assertEqual(pyoram.util.misc.log2ceil(2), 1) + self.assertEqual(pyoram.util.misc.log2ceil(3), 2) + self.assertEqual(pyoram.util.misc.log2ceil(4), 2) + self.assertEqual(pyoram.util.misc.log2ceil(5), 3) + self.assertEqual(pyoram.util.misc.log2ceil(6), 3) + self.assertEqual(pyoram.util.misc.log2ceil(7), 3) + self.assertEqual(pyoram.util.misc.log2ceil(8), 3) + self.assertEqual(pyoram.util.misc.log2ceil(9), 4) + + def test_intdivceil(self): + + with self.assertRaises(ZeroDivisionError): + pyoram.util.misc.intdivceil(0, 0) + with self.assertRaises(ZeroDivisionError): + pyoram.util.misc.intdivceil(1, 0) + + self.assertEqual(pyoram.util.misc.intdivceil(1, 1), 1) + self.assertEqual(pyoram.util.misc.intdivceil(2, 3), 1) + self.assertEqual(2 // 3, 0) + self.assertEqual(pyoram.util.misc.intdivceil( + 123123123123123123123123123123123123123123123123, + 123123123123123123123123123123123123123123123123), 1) + self.assertEqual(pyoram.util.misc.intdivceil( + 2 * 123123123123123123123123123123123123123123123123, + 123123123123123123123123123123123123123123123123), 2) + self.assertEqual(pyoram.util.misc.intdivceil( + 2 * 123123123123123123123123123123123123123123123123 + 1, + 123123123123123123123123123123123123123123123123), 3) + self.assertEqual(pyoram.util.misc.intdivceil( + 2 * 123123123123123123123123123123123123123123123123 - 1, + 123123123123123123123123123123123123123123123123), 2) + self.assertEqual( + (2 * 123123123123123123123123123123123123123123123123 - 1) // \ + 123123123123123123123123123123123123123123123123, + 1) + + def test_MemorySize(self): + self.assertTrue("b" in str(pyoram.util.misc.MemorySize(0.1))) + self.assertTrue("B" in str(pyoram.util.misc.MemorySize(1))) + self.assertTrue("B" in str(pyoram.util.misc.MemorySize(999))) + self.assertTrue("KB" in str(pyoram.util.misc.MemorySize(1000))) + self.assertTrue("KB" in str(pyoram.util.misc.MemorySize(999999))) + self.assertTrue("MB" in str(pyoram.util.misc.MemorySize(1000000))) + self.assertTrue("MB" in str(pyoram.util.misc.MemorySize(999999999))) + self.assertTrue("GB" in str(pyoram.util.misc.MemorySize(1000000000))) + self.assertTrue("GB" in str(pyoram.util.misc.MemorySize(9999999999))) + self.assertTrue("TB" in str(pyoram.util.misc.MemorySize(1000000000000))) + self.assertTrue("b" in str(pyoram.util.misc.MemorySize(1, unit="b"))) + self.assertTrue("b" in str(pyoram.util.misc.MemorySize(2, unit="b"))) + self.assertTrue("b" in str(pyoram.util.misc.MemorySize(7.9, unit="b"))) + + self.assertTrue("B" in str(pyoram.util.misc.MemorySize(8, unit="b"))) + self.assertTrue("B" in str(pyoram.util.misc.MemorySize(1, unit="B"))) + self.assertTrue("B" in str(pyoram.util.misc.MemorySize(999, unit="B"))) + + self.assertTrue("KB" in str(pyoram.util.misc.MemorySize(1000, unit="B"))) + self.assertTrue("KB" in str(pyoram.util.misc.MemorySize(1, unit="KB"))) + self.assertTrue("KB" in str(pyoram.util.misc.MemorySize(999, unit="KB"))) + self.assertTrue("MB" in str(pyoram.util.misc.MemorySize(1000, unit="KB"))) + self.assertTrue("MB" in str(pyoram.util.misc.MemorySize(1, unit="MB"))) + self.assertTrue("MB" in str(pyoram.util.misc.MemorySize(999, unit="MB"))) + self.assertTrue("GB" in str(pyoram.util.misc.MemorySize(1000, unit="MB"))) + self.assertTrue("GB" in str(pyoram.util.misc.MemorySize(1, unit="GB"))) + self.assertTrue("GB" in str(pyoram.util.misc.MemorySize(999, unit="GB"))) + self.assertTrue("TB" in str(pyoram.util.misc.MemorySize(1000, unit="GB"))) + self.assertTrue("TB" in str(pyoram.util.misc.MemorySize(1, unit="TB"))) + + self.assertEqual(pyoram.util.misc.MemorySize(1024).KiB, 1) + self.assertEqual(pyoram.util.misc.MemorySize(1024**2).MiB, 1) + self.assertEqual(pyoram.util.misc.MemorySize(1024**3).GiB, 1) + self.assertEqual(pyoram.util.misc.MemorySize(1024**4).TiB, 1) + + def test_saveload_private_key(self): + with tempfile.NamedTemporaryFile(delete=False) as f: + filename = f.name + try: + key = os.urandom(32) + pyoram.util.misc.save_private_key(filename, key) + loaded_key = pyoram.util.misc.load_private_key(filename) + self.assertEqual(key, loaded_key) + finally: + os.remove(filename) + + def test_chunkiter(self): + self.assertEqual(list(pyoram.util.misc.chunkiter([1,2,3,4,5], 1)), + [[1],[2],[3],[4],[5]]) + self.assertEqual(list(pyoram.util.misc.chunkiter([1,2,3,4,5], 2)), + [[1,2],[3,4],[5]]) + self.assertEqual(list(pyoram.util.misc.chunkiter([1,2,3,4,5], 3)), + [[1,2,3],[4,5]]) + self.assertEqual(list(pyoram.util.misc.chunkiter([1,2,3,4,5], 4)), + [[1,2,3,4],[5]]) + self.assertEqual(list(pyoram.util.misc.chunkiter([1,2,3,4,5], 5)), + [[1,2,3,4,5]]) + self.assertEqual(list(pyoram.util.misc.chunkiter([1,2,3,4,5], 6)), + [[1,2,3,4,5]]) + self.assertEqual(list(pyoram.util.misc.chunkiter([], 1)), + []) + self.assertEqual(list(pyoram.util.misc.chunkiter([], 2)), + []) + +if __name__ == "__main__": + unittest2.main() # pragma: no cover diff --git a/PyORAM/src/pyoram/tests/test_package.py b/PyORAM/src/pyoram/tests/test_package.py new file mode 100644 index 0000000..4875d95 --- /dev/null +++ b/PyORAM/src/pyoram/tests/test_package.py @@ -0,0 +1,38 @@ +import sys +import unittest2 + +import pyoram + +is_pypy = False +try: + import __pypy__ + is_pypy = True +except ImportError: + is_pypy = False + +class Test(unittest2.TestCase): + + # See what Python versions the combined + # coverage report includes + def test_show_coverage(self): + if not is_pypy: + if sys.version_info.major == 2: + if sys.version_info.minor == 7: + print(sys.version_info) + elif sys.version_info.major == 3: + if sys.version_info.minor == 3: + print(sys.version_info) + elif sys.version_info.minor == 4: + print(sys.version_info) + elif sys.version_info.minor == 5: + print(sys.version_info) + if is_pypy: + if sys.version_info.major == 2: + if sys.version_info.minor == 7: + print(sys.version_info) + + def test_version(self): + pyoram.__version__ + +if __name__ == "__main__": + unittest2.main() # pragma: no cover diff --git a/PyORAM/src/pyoram/tests/test_path_oram.py b/PyORAM/src/pyoram/tests/test_path_oram.py new file mode 100644 index 0000000..e5875b1 --- /dev/null +++ b/PyORAM/src/pyoram/tests/test_path_oram.py @@ -0,0 +1,851 @@ +import os +import unittest2 +import tempfile + +from pyoram.oblivious_storage.tree.path_oram import \ + PathORAM +from pyoram.storage.block_storage import \ + BlockStorageTypeFactory +from pyoram.encrypted_storage.encrypted_heap_storage import \ + EncryptedHeapStorage +from pyoram.crypto.aes import AES + +from six.moves import xrange + +thisdir = os.path.dirname(os.path.abspath(__file__)) + +class _TestPathORAMBase(object): + + _type_name = None + _aes_mode = None + _test_key = None + _test_key_size = None + _bucket_capacity = None + _heap_base = None + _kwds = None + + @classmethod + def setUpClass(cls): + assert cls._type_name is not None + assert cls._aes_mode is not None + assert not ((cls._test_key is not None) and \ + (cls._test_key_size is not None)) + assert cls._bucket_capacity is not None + assert cls._heap_base is not None + assert cls._kwds is not None + fd, cls._dummy_name = tempfile.mkstemp() + os.close(fd) + try: + os.remove(cls._dummy_name) + except OSError: # pragma: no cover + pass # pragma: no cover + cls._block_size = 25 + cls._block_count = 47 + cls._testfname = cls.__name__ + "_testfile.bin" + cls._blocks = [] + f = PathORAM.setup( + cls._testfname, + cls._block_size, + cls._block_count, + bucket_capacity=cls._bucket_capacity, + heap_base=cls._heap_base, + key_size=cls._test_key_size, + key=cls._test_key, + storage_type=cls._type_name, + aes_mode=cls._aes_mode, + initialize=lambda i: bytes(bytearray([i])*cls._block_size), + ignore_existing=True, + **cls._kwds) + f.close() + cls._key = f.key + cls._stash = f.stash + cls._position_map = f.position_map + for i in range(cls._block_count): + data = bytearray([i])*cls._block_size + cls._blocks.append(data) + + @classmethod + def tearDownClass(cls): + try: + os.remove(cls._testfname) + except OSError: # pragma: no cover + pass # pragma: no cover + try: + os.remove(cls._dummy_name) + except OSError: # pragma: no cover + pass # pragma: no cover + + def test_setup_fails(self): + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(IOError): + PathORAM.setup( + os.path.join(thisdir, + "baselines", + "exists.empty"), + block_size=10, + block_count=10, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + key=self._test_key, + key_size=self._test_key_size, + aes_mode=self._aes_mode, + storage_type=self._type_name, + **self._kwds) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(IOError): + PathORAM.setup( + os.path.join(thisdir, + "baselines", + "exists.empty"), + block_size=10, + block_count=10, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + key=self._test_key, + key_size=self._test_key_size, + storage_type=self._type_name, + aes_mode=self._aes_mode, + ignore_existing=False, + **self._kwds) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(ValueError): + PathORAM.setup( + self._dummy_name, + block_size=1, + block_count=0, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + key=self._test_key, + key_size=self._test_key_size, + aes_mode=self._aes_mode, + storage_type=self._type_name, + **self._kwds) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(ValueError): + PathORAM.setup( + self._dummy_name, + block_size=0, + block_count=1, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + key=self._test_key, + key_size=self._test_key_size, + aes_mode=self._aes_mode, + storage_type=self._type_name, + **self._kwds) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(TypeError): + PathORAM.setup( + self._dummy_name, + block_size=1, + block_count=1, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + key=self._test_key, + key_size=self._test_key_size, + aes_mode=self._aes_mode, + storage_type=self._type_name, + header_data=2, + **self._kwds) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(ValueError): + PathORAM.setup( + self._dummy_name, + block_size=1, + block_count=1, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + key=self._test_key, + key_size=self._test_key_size, + aes_mode=None, + storage_type=self._type_name, + **self._kwds) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(ValueError): + PathORAM.setup( + self._dummy_name, + block_size=1, + block_count=1, + bucket_capacity=0, + heap_base=self._heap_base, + key=self._test_key, + key_size=self._test_key_size, + aes_mode=self._aes_mode, + storage_type=self._type_name, + **self._kwds) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(ValueError): + PathORAM.setup( + self._dummy_name, + block_size=1, + block_count=1, + bucket_capacity=self._bucket_capacity, + heap_base=1, + key=self._test_key, + key_size=self._test_key_size, + aes_mode=self._aes_mode, + storage_type=self._type_name, + **self._kwds) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(ValueError): + PathORAM.setup( + self._dummy_name, + block_size=1, + block_count=1, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + key_size=-1, + aes_mode=self._aes_mode, + storage_type=self._type_name, + **self._kwds) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(TypeError): + PathORAM.setup( + self._dummy_name, + block_size=1, + block_count=1, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + key=-1, + aes_mode=self._aes_mode, + storage_type=self._type_name, + **self._kwds) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(ValueError): + PathORAM.setup( + self._dummy_name, + block_size=1, + block_count=1, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + key=AES.KeyGen(AES.key_sizes[-1]), + key_size=AES.key_sizes[-1], + aes_mode=self._aes_mode, + storage_type=self._type_name, + **self._kwds) + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(ValueError): + PathORAM.setup( + self._dummy_name, + block_size=1, + block_count=1, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + key=os.urandom(AES.key_sizes[-1]+100), + aes_mode=self._aes_mode, + storage_type=self._type_name, + **self._kwds) + with self.assertRaises(ValueError): + PathORAM.setup( + self._dummy_name, + block_size=1, + block_count=1, + heap_height=1, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + key=self._key, + aes_mode=self._aes_mode, + storage_type=self._type_name, + **self._kwds) + + def test_setup(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + if os.path.exists(fname): + os.remove(fname) # pragma: no cover + bsize = 10 + bcount = 11 + fsetup = PathORAM.setup( + fname, + bsize, + bcount, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + key=self._test_key, + key_size=self._test_key_size, + aes_mode=self._aes_mode, + storage_type=self._type_name, + **self._kwds) + fsetup.close() + self.assertEqual(type(fsetup.raw_storage), + BlockStorageTypeFactory(self._type_name)) + # test that these can be called with default keyword values + fsetup.stash_digest(fsetup.stash) + fsetup.position_map_digest(fsetup.position_map) + with open(fname, 'rb') as f: + flen = len(f.read()) + self.assertEqual( + flen, + PathORAM.compute_storage_size( + bsize, + bcount, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + aes_mode=self._aes_mode, + storage_type=self._type_name)) + self.assertEqual( + flen > + PathORAM.compute_storage_size( + bsize, + bcount, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + aes_mode=self._aes_mode, + storage_type=self._type_name, + ignore_header=True), + True) + with PathORAM(fname, + fsetup.stash, + fsetup.position_map, + key=fsetup.key, + storage_type=self._type_name, + **self._kwds) as f: + self.assertEqual(f.header_data, bytes()) + self.assertEqual(fsetup.header_data, bytes()) + self.assertEqual(f.key, fsetup.key) + self.assertEqual(f.block_size, bsize) + self.assertEqual(fsetup.block_size, bsize) + self.assertEqual(f.block_count, bcount) + self.assertEqual(fsetup.block_count, bcount) + self.assertEqual(f.storage_name, fname) + self.assertEqual(fsetup.storage_name, fname) + os.remove(fname) + + def test_setup_withdata(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + if os.path.exists(fname): + os.remove(fname) # pragma: no cover + bsize = 10 + bcount = 11 + header_data = bytes(bytearray([0,1,2])) + fsetup = PathORAM.setup( + fname, + block_size=bsize, + block_count=bcount, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + key=self._test_key, + key_size=self._test_key_size, + aes_mode=self._aes_mode, + storage_type=self._type_name, + header_data=header_data, + **self._kwds) + fsetup.close() + self.assertEqual(type(fsetup.raw_storage), + BlockStorageTypeFactory(self._type_name)) + with open(fname, 'rb') as f: + flen = len(f.read()) + self.assertEqual( + flen, + PathORAM.compute_storage_size( + bsize, + bcount, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + aes_mode=self._aes_mode, + storage_type=self._type_name, + header_data=header_data)) + self.assertTrue(len(header_data) > 0) + self.assertEqual( + PathORAM.compute_storage_size( + bsize, + bcount, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + aes_mode=self._aes_mode, + storage_type=self._type_name) < + PathORAM.compute_storage_size( + bsize, + bcount, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + aes_mode=self._aes_mode, + storage_type=self._type_name, + header_data=header_data), + True) + self.assertEqual( + flen > + PathORAM.compute_storage_size( + bsize, + bcount, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + aes_mode=self._aes_mode, + storage_type=self._type_name, + header_data=header_data, + ignore_header=True), + True) + with PathORAM(fname, + fsetup.stash, + fsetup.position_map, + key=fsetup.key, + storage_type=self._type_name, + **self._kwds) as f: + self.assertEqual(f.header_data, header_data) + self.assertEqual(fsetup.header_data, header_data) + self.assertEqual(f.key, fsetup.key) + self.assertEqual(f.block_size, bsize) + self.assertEqual(fsetup.block_size, bsize) + self.assertEqual(f.block_count, bcount) + self.assertEqual(fsetup.block_count, bcount) + self.assertEqual(f.storage_name, fname) + self.assertEqual(fsetup.storage_name, fname) + os.remove(fname) + + def test_init_noexists(self): + self.assertEqual(os.path.exists(self._dummy_name), False) + with self.assertRaises(IOError): + with PathORAM( + self._dummy_name, + self._stash, + self._position_map, + key=self._key, + storage_type=self._type_name, + **self._kwds) as f: + pass # pragma: no cover + + def test_init_exists(self): + self.assertEqual(os.path.exists(self._testfname), True) + with open(self._testfname, 'rb') as f: + databefore = f.read() + # no key + with self.assertRaises(ValueError): + with PathORAM(self._testfname, + self._stash, + self._position_map, + storage_type=self._type_name, + **self._kwds) as f: + pass # pragma: no cover + # stash does not match digest + with self.assertRaises(ValueError): + with PathORAM(self._testfname, + {1: bytes()}, + self._position_map, + key=self._key, + storage_type=self._type_name, + **self._kwds) as f: + pass # pragma: no cover + # stash hash invalid key (negative) + with self.assertRaises(ValueError): + with PathORAM(self._testfname, + {-1: bytes()}, + self._position_map, + key=self._key, + storage_type=self._type_name, + **self._kwds) as f: + pass # pragma: no cover + # position map has invalid item (negative) + with self.assertRaises(ValueError): + with PathORAM(self._testfname, + self._stash, + [-1], + key=self._key, + storage_type=self._type_name, + **self._kwds) as f: + pass # pragma: no cover + # position map does not match digest + with self.assertRaises(ValueError): + with PathORAM(self._testfname, + self._stash, + [1], + key=self._key, + storage_type=self._type_name, + **self._kwds) as f: + pass # pragma: no cover + with self.assertRaises(ValueError): + with EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._type_name) as fb: + with PathORAM(fb, + self._stash, + self._position_map, + key=self._key, + storage_type=self._type_name, + **self._kwds) as f: + pass # pragma: no cover + with PathORAM(self._testfname, + self._stash, + self._position_map, + key=self._key, + storage_type=self._type_name, + **self._kwds) as f: + self.assertEqual(f.key, self._key) + self.assertEqual(f.block_size, self._block_size) + self.assertEqual(f.block_count, self._block_count) + self.assertEqual(f.storage_name, self._testfname) + self.assertEqual(f.header_data, bytes()) + self.assertEqual(os.path.exists(self._testfname), True) + with open(self._testfname, 'rb') as f: + dataafter = f.read() + self.assertEqual(databefore[-(self._block_count*self._block_size):], + dataafter[-(self._block_count*self._block_size):]) + + def test_read_block(self): + with PathORAM(self._testfname, + self._stash, + self._position_map, + key=self._key, + storage_type=self._type_name, + **self._kwds) as f: + for i, data in enumerate(self._blocks): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + for i, data in enumerate(self._blocks): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + for i, data in reversed(list(enumerate(self._blocks))): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + for i, data in reversed(list(enumerate(self._blocks))): + self.assertEqual(list(bytearray(f.read_block(i))), + list(self._blocks[i])) + with PathORAM(self._testfname, + self._stash, + self._position_map, + key=self._key, + storage_type=self._type_name, + **self._kwds) as f: + self.assertEqual(list(bytearray(f.read_block(0))), + list(self._blocks[0])) + self.assertEqual(list(bytearray(f.read_block(self._block_count-1))), + list(self._blocks[-1])) + + # test eviction behavior of the tree oram helper + with PathORAM(self._testfname, + self._stash, + self._position_map, + key=self._key, + storage_type=self._type_name, + **self._kwds) as f: + oram = f._oram + vheap = oram.storage_heap.virtual_heap + Z = vheap.blocks_per_bucket + def _has_vacancies(level): + return any(oram.path_block_ids[i] == oram.empty_block_id + for i in range(level*Z, (level+1)*Z)) + + for i in range(len(f.position_map)): + b = f.position_map[i] + f.position_map[i] = vheap.random_leaf_bucket() + oram.load_path(b) + block = oram.extract_block_from_path(i) + if block is not None: + oram.stash[i] = block + + # track where everyone should be able to move + # to, unless the bucket fills up + eviction_levels = {} + for id_, level in zip(oram.path_block_ids, + oram.path_block_eviction_levels): + eviction_levels[id_] = level + for id_ in oram.stash: + block_id, block_addr = \ + oram.get_block_info(oram.stash[id_]) + assert block_id == id_ + eviction_levels[id_] = \ + vheap.clib.calculate_last_common_level( + vheap.k, b, block_addr) + + oram.push_down_path() + oram.fill_path_from_stash() + oram.evict_path() + + # check that everyone was pushed down greedily + oram.load_path(b) + for pos, id_ in enumerate(oram.path_block_ids): + current_level = pos // Z + if (id_ != oram.empty_block_id): + eviction_level = eviction_levels[id_] + self.assertEqual(current_level <= eviction_level, True) + if current_level < eviction_level: + self.assertEqual(_has_vacancies(eviction_level), False) + for id_ in oram.stash: + self.assertEqual( + _has_vacancies(eviction_levels[id_]), False) + + def test_write_block(self): + data = bytearray([self._block_count])*self._block_size + self.assertEqual(len(data) > 0, True) + with PathORAM(self._testfname, + self._stash, + self._position_map, + key=self._key, + storage_type=self._type_name, + **self._kwds) as f: + for i in xrange(self._block_count): + self.assertNotEqual(list(bytearray(f.read_block(i))), + list(data)) + for i in xrange(self._block_count): + f.write_block(i, bytes(data)) + for i in xrange(self._block_count): + self.assertEqual(list(bytearray(f.read_block(i))), + list(data)) + for i, block in enumerate(self._blocks): + f.write_block(i, bytes(block)) + + def test_read_blocks(self): + with PathORAM(self._testfname, + self._stash, + self._position_map, + key=self._key, + storage_type=self._type_name, + **self._kwds) as f: + data = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + data = f.read_blocks([0]) + self.assertEqual(len(data), 1) + self.assertEqual(list(bytearray(data[0])), + list(self._blocks[0])) + self.assertEqual(len(self._blocks) > 1, True) + data = f.read_blocks(list(xrange(1, self._block_count)) + [0]) + self.assertEqual(len(data), self._block_count) + for i, block in enumerate(data[:-1], 1): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + self.assertEqual(list(bytearray(data[-1])), + list(self._blocks[0])) + + def test_write_blocks(self): + data = [bytearray([self._block_count])*self._block_size + for i in xrange(self._block_count)] + with PathORAM(self._testfname, + self._stash, + self._position_map, + key=self._key, + storage_type=self._type_name, + **self._kwds) as f: + orig = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(orig), self._block_count) + for i, block in enumerate(orig): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + f.write_blocks(list(xrange(self._block_count)), + [bytes(b) for b in data]) + new = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(new), self._block_count) + for i, block in enumerate(new): + self.assertEqual(list(bytearray(block)), + list(data[i])) + f.write_blocks(list(xrange(self._block_count)), + [bytes(b) for b in self._blocks]) + orig = f.read_blocks(list(xrange(self._block_count))) + self.assertEqual(len(orig), self._block_count) + for i, block in enumerate(orig): + self.assertEqual(list(bytearray(block)), + list(self._blocks[i])) + + def test_update_header_data(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + if os.path.exists(fname): + os.remove(fname) # pragma: no cover + bsize = 10 + bcount = 11 + header_data = bytes(bytearray([0,1,2])) + fsetup = PathORAM.setup( + fname, + block_size=bsize, + block_count=bcount, + bucket_capacity=self._bucket_capacity, + heap_base=self._heap_base, + key=self._test_key, + key_size=self._test_key_size, + header_data=header_data, + **self._kwds) + fsetup.close() + new_header_data = bytes(bytearray([1,1,1])) + with PathORAM(fname, + fsetup.stash, + fsetup.position_map, + key=fsetup.key, + storage_type=self._type_name, + **self._kwds) as f: + self.assertEqual(f.header_data, header_data) + f.update_header_data(new_header_data) + self.assertEqual(f.header_data, new_header_data) + with PathORAM(fname, + fsetup.stash, + fsetup.position_map, + key=fsetup.key, + storage_type=self._type_name, + **self._kwds) as f: + self.assertEqual(f.header_data, new_header_data) + with self.assertRaises(ValueError): + with PathORAM(fname, + fsetup.stash, + fsetup.position_map, + key=fsetup.key, + storage_type=self._type_name, + **self._kwds) as f: + f.update_header_data(bytes(bytearray([1,1]))) + with self.assertRaises(ValueError): + with PathORAM(fname, + fsetup.stash, + fsetup.position_map, + key=fsetup.key, + storage_type=self._type_name, + **self._kwds) as f: + f.update_header_data(bytes(bytearray([1,1,1,1]))) + with PathORAM(fname, + fsetup.stash, + fsetup.position_map, + key=fsetup.key, + storage_type=self._type_name, + **self._kwds) as f: + self.assertEqual(f.header_data, new_header_data) + os.remove(fname) + + def test_locked_flag(self): + with PathORAM(self._testfname, + self._stash, + self._position_map, + key=self._key, + storage_type=self._type_name, + **self._kwds) as f: + with self.assertRaises(IOError): + with PathORAM(self._testfname, + self._stash, + self._position_map, + key=self._key, + storage_type=self._type_name, + **self._kwds) as f1: + pass # pragma: no cover + with self.assertRaises(IOError): + with PathORAM(self._testfname, + self._stash, + self._position_map, + key=self._key, + storage_type=self._type_name, + **self._kwds) as f1: + pass # pragma: no cover + with PathORAM(self._testfname, + self._stash, + self._position_map, + key=self._key, + storage_type=self._type_name, + ignore_lock=True, + **self._kwds) as f1: + pass + with self.assertRaises(IOError): + with PathORAM(self._testfname, + self._stash, + self._position_map, + key=self._key, + storage_type=self._type_name, + **self._kwds) as f1: + pass # pragma: no cover + with PathORAM(self._testfname, + self._stash, + self._position_map, + key=self._key, + storage_type=self._type_name, + ignore_lock=True, + **self._kwds) as f1: + pass + with PathORAM(self._testfname, + self._stash, + self._position_map, + key=self._key, + storage_type=self._type_name, + ignore_lock=True, + **self._kwds) as f1: + pass + with PathORAM(self._testfname, + self._stash, + self._position_map, + key=self._key, + storage_type=self._type_name, + **self._kwds) as f: + pass + +class TestPathORAMB2Z1(_TestPathORAMBase, + unittest2.TestCase): + _type_name = 'file' + _aes_mode = 'ctr' + _bucket_capacity = 1 + _heap_base = 2 + _kwds = {'cached_levels': 0} + +class TestPathORAMB2Z2(_TestPathORAMBase, + unittest2.TestCase): + _type_name = 'mmap' + _aes_mode = 'gcm' + _bucket_capacity = 2 + _heap_base = 2 + _kwds = {'cached_levels': 0} + +class TestPathORAMB2Z3(_TestPathORAMBase, + unittest2.TestCase): + _type_name = 'mmap' + _aes_mode = 'ctr' + _bucket_capacity = 3 + _heap_base = 2 + _kwds = {'cached_levels': 1} + +class TestPathORAMB2Z4(_TestPathORAMBase, + unittest2.TestCase): + _type_name = 'file' + _aes_mode = 'gcm' + _bucket_capacity = 4 + _heap_base = 2 + _kwds = {'cached_levels': 1} + +class TestPathORAMB2Z5(_TestPathORAMBase, + unittest2.TestCase): + _type_name = 'file' + _aes_mode = 'ctr' + _bucket_capacity = 5 + _heap_base = 2 + _kwds = {'cached_levels': 2, + 'concurrency_level': 0} + +class TestPathORAMB3Z1(_TestPathORAMBase, + unittest2.TestCase): + _type_name = 'file' + _aes_mode = 'ctr' + _bucket_capacity = 1 + _heap_base = 3 + _kwds = {'cached_levels': 2, + 'concurrency_level': 1} + +class TestPathORAMB3Z2(_TestPathORAMBase, + unittest2.TestCase): + _type_name = 'mmap' + _aes_mode = 'gcm' + _bucket_capacity = 2 + _heap_base = 3 + _kwds = {} + +class TestPathORAMB3Z3(_TestPathORAMBase, + unittest2.TestCase): + _type_name = 'mmap' + _aes_mode = 'ctr' + _bucket_capacity = 3 + _heap_base = 3 + _kwds = {} + +class TestPathORAMB3Z4(_TestPathORAMBase, + unittest2.TestCase): + _type_name = 'file' + _aes_mode = 'gcm' + _bucket_capacity = 4 + _heap_base = 3 + _kwds = {} + +class TestPathORAMB3Z5(_TestPathORAMBase, + unittest2.TestCase): + _type_name = 'file' + _aes_mode = 'ctr' + _bucket_capacity = 5 + _heap_base = 3 + _kwds = {} + +if __name__ == "__main__": + unittest2.main() # pragma: no cover diff --git a/PyORAM/src/pyoram/tests/test_top_cached_encrypted_heap_storage.py b/PyORAM/src/pyoram/tests/test_top_cached_encrypted_heap_storage.py new file mode 100644 index 0000000..5c493a6 --- /dev/null +++ b/PyORAM/src/pyoram/tests/test_top_cached_encrypted_heap_storage.py @@ -0,0 +1,711 @@ +import os +import unittest2 +import tempfile +import random + +from pyoram.storage.block_storage import \ + BlockStorageTypeFactory +from pyoram.encrypted_storage.top_cached_encrypted_heap_storage import \ + TopCachedEncryptedHeapStorage +from pyoram.encrypted_storage.encrypted_block_storage import \ + EncryptedBlockStorage +from pyoram.encrypted_storage.encrypted_heap_storage import \ + EncryptedHeapStorage +from pyoram.crypto.aes import AES + +from six.moves import xrange + +thisdir = os.path.dirname(os.path.abspath(__file__)) + +class _TestTopCachedEncryptedHeapStorage(object): + + _init_kwds = None + _storage_type = None + _heap_base = None + _heap_height = None + + @classmethod + def setUpClass(cls): + assert cls._init_kwds is not None + assert cls._storage_type is not None + assert cls._heap_base is not None + assert cls._heap_height is not None + fd, cls._dummy_name = tempfile.mkstemp() + os.close(fd) + try: + os.remove(cls._dummy_name) + except OSError: # pragma: no cover + pass # pragma: no cover + cls._block_size = 50 + cls._blocks_per_bucket = 3 + cls._bucket_count = \ + ((cls._heap_base**(cls._heap_height+1)) - 1)//(cls._heap_base-1) + cls._block_count = cls._bucket_count * \ + cls._blocks_per_bucket + cls._testfname = cls.__name__ + "_testfile.bin" + cls._buckets = [] + f = EncryptedHeapStorage.setup( + cls._testfname, + cls._block_size, + cls._heap_height, + heap_base=cls._heap_base, + blocks_per_bucket=cls._blocks_per_bucket, + storage_type=cls._storage_type, + initialize=lambda i: bytes(bytearray([i]) * \ + cls._block_size * \ + cls._blocks_per_bucket), + ignore_existing=True) + f.close() + cls._key = f.key + for i in range(cls._bucket_count): + data = bytearray([i]) * \ + cls._block_size * \ + cls._blocks_per_bucket + cls._buckets.append(data) + + @classmethod + def tearDownClass(cls): + try: + os.remove(cls._testfname) + except OSError: # pragma: no cover + pass # pragma: no cover + try: + os.remove(cls._dummy_name) + except OSError: # pragma: no cover + pass # pragma: no cover + + def test_factory(self): + kwds = dict(self._init_kwds) + kwds['cached_levels'] = 0 + with EncryptedHeapStorage( + self._testfname, + key=self._key, + storage_type=self._storage_type) as f1: + with TopCachedEncryptedHeapStorage(f1, **kwds) as f2: + self.assertTrue(f1 is f2) + + def test_setup(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + if os.path.exists(fname): + os.remove(fname) # pragma: no cover + bsize = 10 + blocks_per_bucket = 3 + fsetup = EncryptedHeapStorage.setup( + fname, + bsize, + self._heap_height, + heap_base=self._heap_base, + storage_type=self._storage_type, + blocks_per_bucket=blocks_per_bucket) + fsetup.close() + self.assertEqual(type(fsetup.raw_storage), + BlockStorageTypeFactory(self._storage_type)) + with open(fname, 'rb') as f: + flen = len(f.read()) + self.assertEqual( + flen, + TopCachedEncryptedHeapStorage.compute_storage_size( + bsize, + self._heap_height, + heap_base=self._heap_base, + blocks_per_bucket=blocks_per_bucket)) + self.assertEqual( + flen > + TopCachedEncryptedHeapStorage.compute_storage_size( + bsize, + self._heap_height, + heap_base=self._heap_base, + blocks_per_bucket=blocks_per_bucket, + ignore_header=True), + True) + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage( + fname, + key=fsetup.key, + storage_type=self._storage_type), + **self._init_kwds) as f: + self.assertEqual(f.header_data, bytes()) + self.assertEqual(fsetup.header_data, bytes()) + self.assertEqual(f.key, fsetup.key) + self.assertEqual(f.blocks_per_bucket, + blocks_per_bucket) + self.assertEqual(fsetup.blocks_per_bucket, + blocks_per_bucket) + self.assertEqual(f.bucket_count, + (self._heap_base**(self._heap_height+1) - 1)//(self._heap_base-1)) + self.assertEqual(fsetup.bucket_count, + (self._heap_base**(self._heap_height+1) - 1)//(self._heap_base-1)) + self.assertEqual(f.bucket_size, + bsize * blocks_per_bucket) + self.assertEqual(fsetup.bucket_size, + bsize * blocks_per_bucket) + self.assertEqual(f.storage_name, fname) + self.assertEqual(fsetup.storage_name, fname) + os.remove(fname) + + def test_setup_withdata(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + if os.path.exists(fname): + os.remove(fname) # pragma: no cover + bsize = 10 + blocks_per_bucket = 1 + header_data = bytes(bytearray([0,1,2])) + fsetup = EncryptedHeapStorage.setup( + fname, + bsize, + self._heap_height, + heap_base=self._heap_base, + storage_type=self._storage_type, + blocks_per_bucket=blocks_per_bucket, + header_data=header_data) + fsetup.close() + self.assertEqual(type(fsetup.raw_storage), + BlockStorageTypeFactory(self._storage_type)) + with open(fname, 'rb') as f: + flen = len(f.read()) + self.assertEqual( + flen, + TopCachedEncryptedHeapStorage.compute_storage_size( + bsize, + self._heap_height, + heap_base=self._heap_base, + header_data=header_data)) + self.assertTrue(len(header_data) > 0) + self.assertEqual( + TopCachedEncryptedHeapStorage.compute_storage_size( + bsize, + self._heap_height, + heap_base=self._heap_base, + storage_type=self._storage_type) < + TopCachedEncryptedHeapStorage.compute_storage_size( + bsize, + self._heap_height, + heap_base=self._heap_base, + storage_type=self._storage_type, + header_data=header_data), + True) + self.assertEqual( + flen > + TopCachedEncryptedHeapStorage.compute_storage_size( + bsize, + self._heap_height, + heap_base=self._heap_base, + storage_type=self._storage_type, + header_data=header_data, + ignore_header=True), + True) + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage( + fname, + key=fsetup.key, + storage_type=self._storage_type), + **self._init_kwds) as f: + self.assertEqual(f.header_data, header_data) + self.assertEqual(fsetup.header_data, header_data) + self.assertEqual(f.key, fsetup.key) + self.assertEqual(f.blocks_per_bucket, + blocks_per_bucket) + self.assertEqual(fsetup.blocks_per_bucket, + blocks_per_bucket) + self.assertEqual(f.bucket_count, + (self._heap_base**(self._heap_height+1) - 1)//(self._heap_base-1)) + self.assertEqual(fsetup.bucket_count, + (self._heap_base**(self._heap_height+1) - 1)//(self._heap_base-1)) + self.assertEqual(f.bucket_size, + bsize * blocks_per_bucket) + self.assertEqual(fsetup.bucket_size, + bsize * blocks_per_bucket) + self.assertEqual(f.storage_name, fname) + self.assertEqual(fsetup.storage_name, fname) + os.remove(fname) + + def test_init_exists(self): + self.assertEqual(os.path.exists(self._testfname), True) + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._storage_type) as f: + databefore = f.read_blocks(list(range(f.block_count))) + with self.assertRaises(ValueError): + with EncryptedBlockStorage(self._testfname, + key=self._key, + storage_type=self._storage_type) as fb: + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage(fb, key=self._key), + **self._init_kwds) as f: + pass # pragma: no cover + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage( + self._testfname, + key=self._key, + storage_type=self._storage_type), + **self._init_kwds) as f: + self.assertEqual(f.key, self._key) + self.assertEqual(f.bucket_size, + self._block_size * \ + self._blocks_per_bucket) + self.assertEqual(f.bucket_count, + self._bucket_count) + self.assertEqual(f.storage_name, self._testfname) + self.assertEqual(f.header_data, bytes()) + self.assertEqual(os.path.exists(self._testfname), True) + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage( + self._testfname, + key=self._key, + storage_type=self._storage_type), + **self._init_kwds) as f: + dataafter = f.bucket_storage.read_blocks( + list(range(f.bucket_storage.block_count))) + self.assertEqual(databefore, dataafter) + + def test_read_path(self): + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage( + self._testfname, + key=self._key, + storage_type=self._storage_type), + **self._init_kwds) as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + self.assertEqual( + f.virtual_heap.first_bucket_at_level(0), 0) + self.assertNotEqual( + f.virtual_heap.last_leaf_bucket(), 0) + total_buckets = 0 + for b in range(f.virtual_heap.first_bucket_at_level(0), + f.virtual_heap.last_leaf_bucket()+1): + full_bucket_path = f.virtual_heap.Node(b).\ + bucket_path_from_root() + all_level_starts = list(range(len(full_bucket_path)+1)) + for level_start in all_level_starts: + data = f.read_path(b, level_start=level_start) + bucket_path = full_bucket_path[level_start:] + + if len(full_bucket_path) <= f._external_level: + pass + elif level_start >= f._external_level: + total_buckets += len(bucket_path) + else: + total_buckets += len(full_bucket_path[f._external_level:]) + + self.assertEqual(f.virtual_heap.Node(b).level+1-level_start, + len(bucket_path)) + for i, bucket in zip(bucket_path, data): + self.assertEqual(list(bytearray(bucket)), + list(self._buckets[i])) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, + total_buckets*f.bucket_storage._storage.block_size) + + def test_write_path(self): + data = [bytearray([self._bucket_count]) * \ + self._block_size * \ + self._blocks_per_bucket + for i in xrange(self._block_count)] + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage( + self._testfname, + key=self._key, + storage_type=self._storage_type), + **self._init_kwds) as f: + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + + self.assertEqual( + f.virtual_heap.first_bucket_at_level(0), 0) + self.assertNotEqual( + f.virtual_heap.last_leaf_bucket(), 0) + all_buckets = list(range(f.virtual_heap.first_bucket_at_level(0), + f.virtual_heap.last_leaf_bucket()+1)) + random.shuffle(all_buckets) + total_read_buckets = 0 + total_write_buckets = 0 + for b in all_buckets: + full_bucket_path = f.virtual_heap.Node(b).\ + bucket_path_from_root() + all_level_starts = list(range(len(full_bucket_path)+1)) + random.shuffle(all_level_starts) + for level_start in all_level_starts: + orig = f.read_path(b, level_start=level_start) + bucket_path = full_bucket_path[level_start:] + + if len(full_bucket_path) <= f._external_level: + pass + elif level_start >= f._external_level: + total_read_buckets += len(bucket_path) + else: + total_read_buckets += len(full_bucket_path[f._external_level:]) + + if level_start != len(full_bucket_path): + self.assertNotEqual(len(bucket_path), 0) + self.assertEqual(f.virtual_heap.Node(b).level+1-level_start, + len(bucket_path)) + self.assertEqual(len(orig), len(bucket_path)) + + for i, bucket in zip(bucket_path, orig): + self.assertEqual(list(bytearray(bucket)), + list(self._buckets[i])) + + f.write_path(b, [bytes(data[i]) + for i in bucket_path], + level_start=level_start) + if len(full_bucket_path) <= f._external_level: + pass + elif level_start >= f._external_level: + total_write_buckets += len(bucket_path) + else: + total_write_buckets += len(full_bucket_path[f._external_level:]) + + new = f.read_path(b, level_start=level_start) + if len(full_bucket_path) <= f._external_level: + pass + elif level_start >= f._external_level: + total_read_buckets += len(bucket_path) + else: + total_read_buckets += len(full_bucket_path[f._external_level:]) + + self.assertEqual(len(new), len(bucket_path)) + for i, bucket in zip(bucket_path, new): + self.assertEqual(list(bytearray(bucket)), + list(data[i])) + + f.write_path(b, [bytes(self._buckets[i]) + for i in bucket_path], + level_start=level_start) + if len(full_bucket_path) <= f._external_level: + pass + elif level_start >= f._external_level: + total_write_buckets += len(bucket_path) + else: + total_write_buckets += len(full_bucket_path[f._external_level:]) + + + orig = f.read_path(b, level_start=level_start) + if len(full_bucket_path) <= f._external_level: + pass + elif level_start >= f._external_level: + total_read_buckets += len(bucket_path) + else: + total_read_buckets += len(full_bucket_path[f._external_level:]) + + self.assertEqual(len(orig), len(bucket_path)) + for i, bucket in zip(bucket_path, orig): + self.assertEqual(list(bytearray(bucket)), + list(self._buckets[i])) + + full_orig = f.read_path(b) + if len(full_bucket_path) <= f._external_level: + pass + else: + total_read_buckets += len(full_bucket_path[f._external_level:]) + + for i, bucket in zip(full_bucket_path, full_orig): + self.assertEqual(list(bytearray(bucket)), + list(self._buckets[i])) + for c in xrange(self._heap_base): + cn = f.virtual_heap.Node(b).child_node(c) + if not f.virtual_heap.is_nil_node(cn): + cb = cn.bucket + bucket_path = f.virtual_heap.Node(cb).\ + bucket_path_from_root() + orig = f.read_path(cb) + if len(bucket_path) <= f._external_level: + pass + else: + total_read_buckets += len(bucket_path[f._external_level:]) + self.assertEqual(len(orig), len(bucket_path)) + for i, bucket in zip(bucket_path, orig): + self.assertEqual(list(bytearray(bucket)), + list(self._buckets[i])) + + self.assertEqual(f.bytes_sent, + total_write_buckets*f.bucket_storage._storage.block_size) + self.assertEqual(f.bytes_received, + total_read_buckets*f.bucket_storage._storage.block_size) + + def test_update_header_data(self): + fname = ".".join(self.id().split(".")[1:]) + fname += ".bin" + fname = os.path.join(thisdir, fname) + if os.path.exists(fname): + os.remove(fname) # pragma: no cover + bsize = 10 + blocks_per_bucket = 1 + header_data = bytes(bytearray([0,1,2])) + fsetup = EncryptedHeapStorage.setup( + fname, + bsize, + self._heap_height, + heap_base=self._heap_base, + blocks_per_bucket=blocks_per_bucket, + header_data=header_data) + fsetup.close() + new_header_data = bytes(bytearray([1,1,1])) + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage( + fname, + key=fsetup.key, + storage_type=self._storage_type), + **self._init_kwds) as f: + self.assertEqual(f.header_data, header_data) + f.update_header_data(new_header_data) + self.assertEqual(f.header_data, new_header_data) + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage( + fname, + key=fsetup.key, + storage_type=self._storage_type), + **self._init_kwds) as f: + self.assertEqual(f.header_data, new_header_data) + with self.assertRaises(ValueError): + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage( + fname, + key=fsetup.key, + storage_type=self._storage_type), + **self._init_kwds) as f: + f.update_header_data(bytes(bytearray([1,1]))) + with self.assertRaises(ValueError): + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage( + fname, + key=fsetup.key, + storage_type=self._storage_type), + **self._init_kwds) as f: + f.update_header_data(bytes(bytearray([1,1,1,1]))) + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage( + fname, + key=fsetup.key, + storage_type=self._storage_type), + **self._init_kwds) as f: + self.assertEqual(f.header_data, new_header_data) + os.remove(fname) + + def test_locked_flag(self): + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._storage_type), + **self._init_kwds) as f: + with self.assertRaises(IOError): + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._storage_type), + **self._init_kwds) as f1: + pass # pragma: no cover + with self.assertRaises(IOError): + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._storage_type), + **self._init_kwds) as f1: + pass # pragma: no cover + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._storage_type, + ignore_lock=True), + **self._init_kwds) as f1: + pass + with self.assertRaises(IOError): + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._storage_type), + **self._init_kwds) as f1: + pass # pragma: no cover + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._storage_type, + ignore_lock=True), + **self._init_kwds) as f1: + pass + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._storage_type, + ignore_lock=True), + **self._init_kwds) as f1: + pass + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._storage_type), + **self._init_kwds) as f: + pass + + def test_cache_size(self): + with TopCachedEncryptedHeapStorage( + EncryptedHeapStorage(self._testfname, + key=self._key, + storage_type=self._storage_type), + **self._init_kwds) as f: + num_cached_levels = self._init_kwds.get('cached_levels', 1) + if num_cached_levels < 0: + num_cached_levels = f.virtual_heap.levels + cache_bucket_count = 0 + for l in xrange(num_cached_levels): + if l <= f.virtual_heap.last_level: + cache_bucket_count += f.virtual_heap.bucket_count_at_level(l) + self.assertEqual(cache_bucket_count > 0, True) + self.assertEqual(len(f.cached_bucket_data), + cache_bucket_count * f.bucket_size) + + self.assertEqual(f.bytes_sent, 0) + self.assertEqual(f.bytes_received, 0) + self.assertEqual(f._root_device.bytes_sent, 0) + self.assertEqual( + f._root_device.bytes_received, + cache_bucket_count*f._root_device.bucket_storage._storage.block_size) + +class TestTopCachedEncryptedHeapStorageCacheMMapDefault( + _TestTopCachedEncryptedHeapStorage, + unittest2.TestCase): + _init_kwds = {} + _storage_type = 'mmap' + _heap_base = 2 + _heap_height = 7 + +class TestTopCachedEncryptedHeapStorageMMapCache1( + _TestTopCachedEncryptedHeapStorage, + unittest2.TestCase): + _init_kwds = {'cached_levels': 1} + _storage_type = 'mmap' + _heap_base = 2 + _heap_height = 7 + +class TestTopCachedEncryptedHeapStorageMMapCache2( + _TestTopCachedEncryptedHeapStorage, + unittest2.TestCase): + _init_kwds = {'cached_levels': 2} + _storage_type = 'mmap' + _heap_base = 2 + _heap_height = 7 + +class TestTopCachedEncryptedHeapStorageMMapCache3( + _TestTopCachedEncryptedHeapStorage, + unittest2.TestCase): + _init_kwds = {'cached_levels': 3} + _storage_type = 'mmap' + _heap_base = 2 + _heap_height = 7 + +class TestTopCachedEncryptedHeapStorageMMapCache4( + _TestTopCachedEncryptedHeapStorage, + unittest2.TestCase): + _init_kwds = {'cached_levels': 4} + _storage_type = 'mmap' + _heap_base = 2 + _heap_height = 7 + +class TestTopCachedEncryptedHeapStorageMMapCache5( + _TestTopCachedEncryptedHeapStorage, + unittest2.TestCase): + _init_kwds = {'cached_levels': 5} + _storage_type = 'mmap' + _heap_base = 2 + _heap_height = 7 + +class TestTopCachedEncryptedHeapStorageCacheFileDefault( + _TestTopCachedEncryptedHeapStorage, + unittest2.TestCase): + _init_kwds = {} + _storage_type = 'file' + _heap_base = 2 + _heap_height = 7 + +class TestTopCachedEncryptedHeapStorageFileCache1( + _TestTopCachedEncryptedHeapStorage, + unittest2.TestCase): + _init_kwds = {'cached_levels': 1} + _storage_type = 'file' + _heap_base = 2 + _heap_height = 7 + +class TestTopCachedEncryptedHeapStorageFileCache2( + _TestTopCachedEncryptedHeapStorage, + unittest2.TestCase): + _init_kwds = {'cached_levels': 2} + _storage_type = 'file' + _heap_base = 2 + _heap_height = 7 + +class TestTopCachedEncryptedHeapStorageFileCache3( + _TestTopCachedEncryptedHeapStorage, + unittest2.TestCase): + _init_kwds = {'cached_levels': 3} + _storage_type = 'file' + _heap_base = 2 + _heap_height = 7 + +class TestTopCachedEncryptedHeapStorageFileCache4( + _TestTopCachedEncryptedHeapStorage, + unittest2.TestCase): + _init_kwds = {'cached_levels': 4} + _storage_type = 'file' + _heap_base = 2 + _heap_height = 7 + +class TestTopCachedEncryptedHeapStorageFileCache5( + _TestTopCachedEncryptedHeapStorage, + unittest2.TestCase): + _init_kwds = {'cached_levels': 5} + _storage_type = 'file' + _heap_base = 2 + _heap_height = 7 + +class TestTopCachedEncryptedHeapStorageFileCacheBigConcurrency0( + _TestTopCachedEncryptedHeapStorage, + unittest2.TestCase): + _init_kwds = {'cached_levels': 20, + 'concurrency_level': 0} + _storage_type = 'file' + _heap_base = 2 + _heap_height = 7 + +class TestTopCachedEncryptedHeapStorageFileCache6Concurrency1( + _TestTopCachedEncryptedHeapStorage, + unittest2.TestCase): + _init_kwds = {'cached_levels': 6, + 'concurrency_level': 1} + _storage_type = 'file' + _heap_base = 2 + _heap_height = 7 + +class TestTopCachedEncryptedHeapStorageFileCache3ConcurrencyBig( + _TestTopCachedEncryptedHeapStorage, + unittest2.TestCase): + _init_kwds = {'cached_levels': 3, + 'concurrency_level': 20} + _storage_type = 'file' + _heap_base = 2 + _heap_height = 7 + +class TestTopCachedEncryptedHeapStorageFileCache3Concurrency1Base3( + _TestTopCachedEncryptedHeapStorage, + unittest2.TestCase): + _init_kwds = {'cached_levels': 3, + 'concurrency_level': 3} + _storage_type = 'file' + _heap_base = 3 + _heap_height = 4 + +class TestTopCachedEncryptedHeapStorageFileCacheAll( + _TestTopCachedEncryptedHeapStorage, + unittest2.TestCase): + _init_kwds = {'cached_levels': -1} + _storage_type = 'file' + _heap_base = 2 + _heap_height = 3 + +if __name__ == "__main__": + unittest2.main() # pragma: no cover diff --git a/PyORAM/src/pyoram/tests/test_virtual_heap.py b/PyORAM/src/pyoram/tests/test_virtual_heap.py new file mode 100644 index 0000000..54aacbe --- /dev/null +++ b/PyORAM/src/pyoram/tests/test_virtual_heap.py @@ -0,0 +1,1227 @@ +import os +import subprocess +import unittest2 + +import pyoram +from pyoram.util.virtual_heap import \ + (VirtualHeap, + SizedVirtualHeap, + max_k_labeled, + calculate_bucket_count_in_heap_with_height, + calculate_bucket_count_in_heap_at_level, + calculate_bucket_level, + calculate_last_common_level, + calculate_necessary_heap_height, + basek_string_to_base10_integer, + numerals, + _clib) + +from six.moves import xrange + +thisdir = os.path.dirname(os.path.abspath(__file__)) +baselinedir = os.path.join(thisdir, "baselines") + +try: + has_dot = not subprocess.call(['dot','-?'], + stdout=subprocess.PIPE) +except: + has_dot = False + +_test_bases = list(xrange(2, 15)) + [max_k_labeled+1] +_test_labeled_bases = list(xrange(2, 15)) + [max_k_labeled] + +def _do_preorder(x): + if x.level > 2: + return + yield x.bucket + for c in xrange(x.k): + for b in _do_preorder(x.child_node(c)): + yield b + +def _do_postorder(x): + if x.level > 2: + return + for c in xrange(x.k): + for b in _do_postorder(x.child_node(c)): + yield b + yield x.bucket + +def _do_inorder(x): + assert x.k == 2 + if x.level > 2: + return + for b in _do_inorder(x.child_node(0)): + yield b + yield x.bucket + for b in _do_inorder(x.child_node(1)): + yield b + +class TestVirtualHeapNode(unittest2.TestCase): + + def test_init(self): + for k in _test_bases: + vh = VirtualHeap(k) + node = vh.Node(0) + self.assertEqual(node.k, k) + self.assertEqual(node.bucket, 0) + self.assertEqual(node.level, 0) + for b in xrange(1, k+1): + node = vh.Node(b) + self.assertEqual(node.k, k) + self.assertEqual(node.bucket, b) + self.assertEqual(node.level, 1) + + def test_level(self): + Node = VirtualHeap(2).Node + self.assertEqual(Node(0).level, 0) + self.assertEqual(Node(1).level, 1) + self.assertEqual(Node(2).level, 1) + self.assertEqual(Node(3).level, 2) + self.assertEqual(Node(4).level, 2) + self.assertEqual(Node(5).level, 2) + self.assertEqual(Node(6).level, 2) + self.assertEqual(Node(7).level, 3) + + Node = VirtualHeap(3).Node + self.assertEqual(Node(0).level, 0) + self.assertEqual(Node(1).level, 1) + self.assertEqual(Node(2).level, 1) + self.assertEqual(Node(3).level, 1) + self.assertEqual(Node(4).level, 2) + self.assertEqual(Node(5).level, 2) + self.assertEqual(Node(6).level, 2) + self.assertEqual(Node(7).level, 2) + self.assertEqual(Node(8).level, 2) + self.assertEqual(Node(9).level, 2) + self.assertEqual(Node(10).level, 2) + self.assertEqual(Node(11).level, 2) + self.assertEqual(Node(12).level, 2) + self.assertEqual(Node(13).level, 3) + + def test_hash(self): + x1 = VirtualHeap(3).Node(5) + x2 = VirtualHeap(2).Node(5) + self.assertNotEqual(id(x1), id(x2)) + self.assertEqual(x1, x2) + self.assertEqual(x1, x1) + self.assertEqual(x2, x2) + + all_node_set = set() + all_node_list = list() + for k in _test_bases: + node_set = set() + node_list = list() + Node = VirtualHeap(k).Node + for height in xrange(k+2): + node = Node(height) + node_set.add(node) + all_node_set.add(node) + node_list.append(node) + all_node_list.append(node) + self.assertEqual(sorted(node_set), + sorted(node_list)) + self.assertNotEqual(sorted(all_node_set), + sorted(all_node_list)) + def test_int(self): + Node2 = VirtualHeap(2).Node + Node3 = VirtualHeap(3).Node + for b in range(100): + self.assertEqual(int(Node2(b)), b) + self.assertEqual(int(Node3(b)), b) + + def test_lt(self): + Node = VirtualHeap(3).Node + self.assertEqual(Node(5) < 4, False) + self.assertEqual(Node(5) < 5, False) + self.assertEqual(Node(5) < 6, True) + + def test_le(self): + Node = VirtualHeap(3).Node + self.assertEqual(Node(5) <= 4, False) + self.assertEqual(Node(5) <= 5, True) + self.assertEqual(Node(5) <= 6, True) + + def test_eq(self): + Node = VirtualHeap(3).Node + self.assertEqual(Node(5) == 4, False) + self.assertEqual(Node(5) == 5, True) + self.assertEqual(Node(5) == 6, False) + + def test_ne(self): + Node = VirtualHeap(3).Node + self.assertEqual(Node(5) != 4, True) + self.assertEqual(Node(5) != 5, False) + self.assertEqual(Node(5) != 6, True) + + def test_gt(self): + Node = VirtualHeap(3).Node + self.assertEqual(Node(5) > 4, True) + self.assertEqual(Node(5) > 5, False) + self.assertEqual(Node(5) > 6, False) + + def test_ge(self): + Node = VirtualHeap(3).Node + self.assertEqual(Node(5) >= 4, True) + self.assertEqual(Node(5) >= 5, True) + self.assertEqual(Node(5) >= 6, False) + + def test_last_common_level_k2(self): + Node = VirtualHeap(2).Node + n0 = Node(0) + n1 = Node(1) + n2 = Node(2) + n3 = Node(3) + n4 = Node(4) + n5 = Node(5) + n6 = Node(6) + n7 = Node(7) + self.assertEqual(n0.last_common_level(n0), 0) + self.assertEqual(n0.last_common_level(n1), 0) + self.assertEqual(n0.last_common_level(n2), 0) + self.assertEqual(n0.last_common_level(n3), 0) + self.assertEqual(n0.last_common_level(n4), 0) + self.assertEqual(n0.last_common_level(n5), 0) + self.assertEqual(n0.last_common_level(n6), 0) + self.assertEqual(n0.last_common_level(n7), 0) + + self.assertEqual(n1.last_common_level(n0), 0) + self.assertEqual(n1.last_common_level(n1), 1) + self.assertEqual(n1.last_common_level(n2), 0) + self.assertEqual(n1.last_common_level(n3), 1) + self.assertEqual(n1.last_common_level(n4), 1) + self.assertEqual(n1.last_common_level(n5), 0) + self.assertEqual(n1.last_common_level(n6), 0) + self.assertEqual(n1.last_common_level(n7), 1) + + self.assertEqual(n2.last_common_level(n0), 0) + self.assertEqual(n2.last_common_level(n1), 0) + self.assertEqual(n2.last_common_level(n2), 1) + self.assertEqual(n2.last_common_level(n3), 0) + self.assertEqual(n2.last_common_level(n4), 0) + self.assertEqual(n2.last_common_level(n5), 1) + self.assertEqual(n2.last_common_level(n6), 1) + self.assertEqual(n2.last_common_level(n7), 0) + + self.assertEqual(n3.last_common_level(n0), 0) + self.assertEqual(n3.last_common_level(n1), 1) + self.assertEqual(n3.last_common_level(n2), 0) + self.assertEqual(n3.last_common_level(n3), 2) + self.assertEqual(n3.last_common_level(n4), 1) + self.assertEqual(n3.last_common_level(n5), 0) + self.assertEqual(n3.last_common_level(n6), 0) + self.assertEqual(n3.last_common_level(n7), 2) + + self.assertEqual(n4.last_common_level(n0), 0) + self.assertEqual(n4.last_common_level(n1), 1) + self.assertEqual(n4.last_common_level(n2), 0) + self.assertEqual(n4.last_common_level(n3), 1) + self.assertEqual(n4.last_common_level(n4), 2) + self.assertEqual(n4.last_common_level(n5), 0) + self.assertEqual(n4.last_common_level(n6), 0) + self.assertEqual(n4.last_common_level(n7), 1) + + self.assertEqual(n5.last_common_level(n0), 0) + self.assertEqual(n5.last_common_level(n1), 0) + self.assertEqual(n5.last_common_level(n2), 1) + self.assertEqual(n5.last_common_level(n3), 0) + self.assertEqual(n5.last_common_level(n4), 0) + self.assertEqual(n5.last_common_level(n5), 2) + self.assertEqual(n5.last_common_level(n6), 1) + self.assertEqual(n5.last_common_level(n7), 0) + + self.assertEqual(n6.last_common_level(n0), 0) + self.assertEqual(n6.last_common_level(n1), 0) + self.assertEqual(n6.last_common_level(n2), 1) + self.assertEqual(n6.last_common_level(n3), 0) + self.assertEqual(n6.last_common_level(n4), 0) + self.assertEqual(n6.last_common_level(n5), 1) + self.assertEqual(n6.last_common_level(n6), 2) + self.assertEqual(n6.last_common_level(n7), 0) + + self.assertEqual(n7.last_common_level(n0), 0) + self.assertEqual(n7.last_common_level(n1), 1) + self.assertEqual(n7.last_common_level(n2), 0) + self.assertEqual(n7.last_common_level(n3), 2) + self.assertEqual(n7.last_common_level(n4), 1) + self.assertEqual(n7.last_common_level(n5), 0) + self.assertEqual(n7.last_common_level(n6), 0) + self.assertEqual(n7.last_common_level(n7), 3) + + def test_last_common_level_k3(self): + Node = VirtualHeap(3).Node + n0 = Node(0) + n1 = Node(1) + n2 = Node(2) + n3 = Node(3) + n4 = Node(4) + n5 = Node(5) + n6 = Node(6) + n7 = Node(7) + self.assertEqual(n0.last_common_level(n0), 0) + self.assertEqual(n0.last_common_level(n1), 0) + self.assertEqual(n0.last_common_level(n2), 0) + self.assertEqual(n0.last_common_level(n3), 0) + self.assertEqual(n0.last_common_level(n4), 0) + self.assertEqual(n0.last_common_level(n5), 0) + self.assertEqual(n0.last_common_level(n6), 0) + self.assertEqual(n0.last_common_level(n7), 0) + + self.assertEqual(n1.last_common_level(n0), 0) + self.assertEqual(n1.last_common_level(n1), 1) + self.assertEqual(n1.last_common_level(n2), 0) + self.assertEqual(n1.last_common_level(n3), 0) + self.assertEqual(n1.last_common_level(n4), 1) + self.assertEqual(n1.last_common_level(n5), 1) + self.assertEqual(n1.last_common_level(n6), 1) + self.assertEqual(n1.last_common_level(n7), 0) + + self.assertEqual(n2.last_common_level(n0), 0) + self.assertEqual(n2.last_common_level(n1), 0) + self.assertEqual(n2.last_common_level(n2), 1) + self.assertEqual(n2.last_common_level(n3), 0) + self.assertEqual(n2.last_common_level(n4), 0) + self.assertEqual(n2.last_common_level(n5), 0) + self.assertEqual(n2.last_common_level(n6), 0) + self.assertEqual(n2.last_common_level(n7), 1) + + self.assertEqual(n3.last_common_level(n0), 0) + self.assertEqual(n3.last_common_level(n1), 0) + self.assertEqual(n3.last_common_level(n2), 0) + self.assertEqual(n3.last_common_level(n3), 1) + self.assertEqual(n3.last_common_level(n4), 0) + self.assertEqual(n3.last_common_level(n5), 0) + self.assertEqual(n3.last_common_level(n6), 0) + self.assertEqual(n3.last_common_level(n7), 0) + + self.assertEqual(n4.last_common_level(n0), 0) + self.assertEqual(n4.last_common_level(n1), 1) + self.assertEqual(n4.last_common_level(n2), 0) + self.assertEqual(n4.last_common_level(n3), 0) + self.assertEqual(n4.last_common_level(n4), 2) + self.assertEqual(n4.last_common_level(n5), 1) + self.assertEqual(n4.last_common_level(n6), 1) + self.assertEqual(n4.last_common_level(n7), 0) + + self.assertEqual(n5.last_common_level(n0), 0) + self.assertEqual(n5.last_common_level(n1), 1) + self.assertEqual(n5.last_common_level(n2), 0) + self.assertEqual(n5.last_common_level(n3), 0) + self.assertEqual(n5.last_common_level(n4), 1) + self.assertEqual(n5.last_common_level(n5), 2) + self.assertEqual(n5.last_common_level(n6), 1) + self.assertEqual(n5.last_common_level(n7), 0) + + self.assertEqual(n6.last_common_level(n0), 0) + self.assertEqual(n6.last_common_level(n1), 1) + self.assertEqual(n6.last_common_level(n2), 0) + self.assertEqual(n6.last_common_level(n3), 0) + self.assertEqual(n6.last_common_level(n4), 1) + self.assertEqual(n6.last_common_level(n5), 1) + self.assertEqual(n6.last_common_level(n6), 2) + self.assertEqual(n6.last_common_level(n7), 0) + + self.assertEqual(n7.last_common_level(n0), 0) + self.assertEqual(n7.last_common_level(n1), 0) + self.assertEqual(n7.last_common_level(n2), 1) + self.assertEqual(n7.last_common_level(n3), 0) + self.assertEqual(n7.last_common_level(n4), 0) + self.assertEqual(n7.last_common_level(n5), 0) + self.assertEqual(n7.last_common_level(n6), 0) + self.assertEqual(n7.last_common_level(n7), 2) + + def test_child_node(self): + root = VirtualHeap(2).Node(0) + self.assertEqual(list(_do_preorder(root)), + [0, 1, 3, 4, 2, 5, 6]) + self.assertEqual(list(_do_postorder(root)), + [3, 4, 1, 5, 6, 2, 0]) + self.assertEqual(list(_do_inorder(root)), + [3, 1, 4, 0, 5, 2, 6]) + + root = VirtualHeap(3).Node(0) + self.assertEqual( + list(_do_preorder(root)), + [0, 1, 4, 5, 6, 2, 7, 8, 9, 3, 10, 11, 12]) + self.assertEqual( + list(_do_postorder(root)), + [4, 5, 6, 1, 7, 8, 9, 2, 10, 11, 12, 3, 0]) + + def test_parent_node(self): + Node = VirtualHeap(2).Node + self.assertEqual(Node(0).parent_node(), + None) + self.assertEqual(Node(1).parent_node(), + Node(0)) + self.assertEqual(Node(2).parent_node(), + Node(0)) + self.assertEqual(Node(3).parent_node(), + Node(1)) + self.assertEqual(Node(4).parent_node(), + Node(1)) + self.assertEqual(Node(5).parent_node(), + Node(2)) + self.assertEqual(Node(6).parent_node(), + Node(2)) + self.assertEqual(Node(7).parent_node(), + Node(3)) + + Node = VirtualHeap(3).Node + self.assertEqual(Node(0).parent_node(), + None) + self.assertEqual(Node(1).parent_node(), + Node(0)) + self.assertEqual(Node(2).parent_node(), + Node(0)) + self.assertEqual(Node(3).parent_node(), + Node(0)) + self.assertEqual(Node(4).parent_node(), + Node(1)) + self.assertEqual(Node(5).parent_node(), + Node(1)) + self.assertEqual(Node(6).parent_node(), + Node(1)) + self.assertEqual(Node(7).parent_node(), + Node(2)) + self.assertEqual(Node(8).parent_node(), + Node(2)) + self.assertEqual(Node(9).parent_node(), + Node(2)) + self.assertEqual(Node(10).parent_node(), + Node(3)) + self.assertEqual(Node(11).parent_node(), + Node(3)) + self.assertEqual(Node(12).parent_node(), + Node(3)) + self.assertEqual(Node(13).parent_node(), + Node(4)) + + def test_ancestor_node_at_level(self): + Node = VirtualHeap(2).Node + self.assertEqual(Node(0).ancestor_node_at_level(0), + Node(0)) + self.assertEqual(Node(0).ancestor_node_at_level(1), + None) + self.assertEqual(Node(1).ancestor_node_at_level(0), + Node(0)) + self.assertEqual(Node(1).ancestor_node_at_level(1), + Node(1)) + self.assertEqual(Node(1).ancestor_node_at_level(2), + None) + self.assertEqual(Node(3).ancestor_node_at_level(0), + Node(0)) + self.assertEqual(Node(3).ancestor_node_at_level(1), + Node(1)) + self.assertEqual(Node(3).ancestor_node_at_level(2), + Node(3)) + self.assertEqual(Node(3).ancestor_node_at_level(3), + None) + + Node = VirtualHeap(3).Node + self.assertEqual(Node(0).ancestor_node_at_level(0), + Node(0)) + self.assertEqual(Node(0).ancestor_node_at_level(1), + None) + self.assertEqual(Node(1).ancestor_node_at_level(0), + Node(0)) + self.assertEqual(Node(1).ancestor_node_at_level(1), + Node(1)) + self.assertEqual(Node(1).ancestor_node_at_level(2), + None) + self.assertEqual(Node(4).ancestor_node_at_level(0), + Node(0)) + self.assertEqual(Node(4).ancestor_node_at_level(1), + Node(1)) + self.assertEqual(Node(4).ancestor_node_at_level(2), + Node(4)) + self.assertEqual(Node(4).ancestor_node_at_level(3), + None) + + def test_path_to_root(self): + Node = VirtualHeap(2).Node + self.assertEqual(list(int(n) for n in Node(0).bucket_path_to_root()), + list(reversed([0]))) + self.assertEqual(list(int(n) for n in Node(7).bucket_path_to_root()), + list(reversed([0, 1, 3, 7]))) + self.assertEqual(list(int(n) for n in Node(8).bucket_path_to_root()), + list(reversed([0, 1, 3, 8]))) + self.assertEqual(list(int(n) for n in Node(9).bucket_path_to_root()), + list(reversed([0, 1, 4, 9]))) + self.assertEqual(list(int(n) for n in Node(10).bucket_path_to_root()), + list(reversed([0, 1, 4, 10]))) + self.assertEqual(list(int(n) for n in Node(11).bucket_path_to_root()), + list(reversed([0, 2, 5, 11]))) + self.assertEqual(list(int(n) for n in Node(12).bucket_path_to_root()), + list(reversed([0, 2, 5, 12]))) + self.assertEqual(list(int(n) for n in Node(13).bucket_path_to_root()), + list(reversed([0, 2, 6, 13]))) + self.assertEqual(list(int(n) for n in Node(14).bucket_path_to_root()), + list(reversed([0, 2, 6, 14]))) + + def test_path_from_root(self): + Node = VirtualHeap(2).Node + self.assertEqual(list(int(n) for n in Node(0).bucket_path_from_root()), + [0]) + self.assertEqual(list(int(n) for n in Node(7).bucket_path_from_root()), + [0, 1, 3, 7]) + self.assertEqual(list(int(n) for n in Node(8).bucket_path_from_root()), + [0, 1, 3, 8]) + self.assertEqual(list(int(n) for n in Node(9).bucket_path_from_root()), + [0, 1, 4, 9]) + self.assertEqual(list(int(n) for n in Node(10).bucket_path_from_root()), + [0, 1, 4, 10]) + self.assertEqual(list(int(n) for n in Node(11).bucket_path_from_root()), + [0, 2, 5, 11]) + self.assertEqual(list(int(n) for n in Node(12).bucket_path_from_root()), + [0, 2, 5, 12]) + self.assertEqual(list(int(n) for n in Node(13).bucket_path_from_root()), + [0, 2, 6, 13]) + self.assertEqual(list(int(n) for n in Node(14).bucket_path_from_root()), + [0, 2, 6, 14]) + + def test_bucket_path_to_root(self): + Node = VirtualHeap(2).Node + self.assertEqual(list(Node(0).bucket_path_to_root()), + list(reversed([0]))) + self.assertEqual(list(Node(7).bucket_path_to_root()), + list(reversed([0, 1, 3, 7]))) + self.assertEqual(list(Node(8).bucket_path_to_root()), + list(reversed([0, 1, 3, 8]))) + self.assertEqual(list(Node(9).bucket_path_to_root()), + list(reversed([0, 1, 4, 9]))) + self.assertEqual(list(Node(10).bucket_path_to_root()), + list(reversed([0, 1, 4, 10]))) + self.assertEqual(list(Node(11).bucket_path_to_root()), + list(reversed([0, 2, 5, 11]))) + self.assertEqual(list(Node(12).bucket_path_to_root()), + list(reversed([0, 2, 5, 12]))) + self.assertEqual(list(Node(13).bucket_path_to_root()), + list(reversed([0, 2, 6, 13]))) + self.assertEqual(list(Node(14).bucket_path_to_root()), + list(reversed([0, 2, 6, 14]))) + + def test_bucket_path_from_root(self): + Node = VirtualHeap(2).Node + self.assertEqual(Node(0).bucket_path_from_root(), + [0]) + self.assertEqual(Node(7).bucket_path_from_root(), + [0, 1, 3, 7]) + self.assertEqual(Node(8).bucket_path_from_root(), + [0, 1, 3, 8]) + self.assertEqual(Node(9).bucket_path_from_root(), + [0, 1, 4, 9]) + self.assertEqual(Node(10).bucket_path_from_root(), + [0, 1, 4, 10]) + self.assertEqual(Node(11).bucket_path_from_root(), + [0, 2, 5, 11]) + self.assertEqual(Node(12).bucket_path_from_root(), + [0, 2, 5, 12]) + self.assertEqual(Node(13).bucket_path_from_root(), + [0, 2, 6, 13]) + self.assertEqual(Node(14).bucket_path_from_root(), + [0, 2, 6, 14]) + + def test_repr(self): + Node = VirtualHeap(2).Node + self.assertEqual( + repr(Node(0)), + "VirtualHeapNode(k=2, bucket=0, level=0, label='')") + self.assertEqual( + repr(Node(7)), + "VirtualHeapNode(k=2, bucket=7, level=3, label='000')") + + Node = VirtualHeap(3).Node + self.assertEqual( + repr(Node(0)), + "VirtualHeapNode(k=3, bucket=0, level=0, label='')") + self.assertEqual( + repr(Node(7)), + "VirtualHeapNode(k=3, bucket=7, level=2, label='10')") + + Node = VirtualHeap(5).Node + self.assertEqual( + repr(Node(25)), + "VirtualHeapNode(k=5, bucket=25, level=2, label='34')") + + Node = VirtualHeap(max_k_labeled).Node + self.assertEqual( + repr(Node(0)), + ("VirtualHeapNode(k=%d, bucket=0, level=0, label='')" + % (max_k_labeled))) + self.assertEqual(max_k_labeled >= 2, True) + self.assertEqual( + repr(Node(1)), + ("VirtualHeapNode(k=%d, bucket=1, level=1, label='0')" + % (max_k_labeled))) + + Node = VirtualHeap(max_k_labeled+1).Node + self.assertEqual( + repr(Node(0)), + ("VirtualHeapNode(k=%d, bucket=0, level=0, label='')" + % (max_k_labeled+1))) + self.assertEqual( + repr(Node(1)), + ("VirtualHeapNode(k=%d, bucket=1, level=1, label='')" + % (max_k_labeled+1))) + self.assertEqual( + repr(Node(max_k_labeled+1)), + ("VirtualHeapNode(k=%d, bucket=%d, level=1, label='')" + % (max_k_labeled+1, + max_k_labeled+1))) + self.assertEqual( + repr(Node(max_k_labeled+2)), + ("VirtualHeapNode(k=%d, bucket=%d, level=2, label='')" + % (max_k_labeled+1, + max_k_labeled+2))) + + def test_str(self): + Node = VirtualHeap(2).Node + self.assertEqual( + str(Node(0)), + "(0, 0)") + self.assertEqual( + str(Node(7)), + "(3, 0)") + + Node = VirtualHeap(3).Node + self.assertEqual( + str(Node(0)), + "(0, 0)") + self.assertEqual( + str(Node(7)), + "(2, 3)") + + Node = VirtualHeap(5).Node + self.assertEqual( + str(Node(25)), + "(2, 19)") + + def test_label(self): + + Node = VirtualHeap(2).Node + self.assertEqual(Node(0).label(), "") + self.assertEqual(Node(1).label(), "0") + self.assertEqual(Node(2).label(), "1") + self.assertEqual(Node(3).label(), "00") + self.assertEqual(Node(4).label(), "01") + self.assertEqual(Node(5).label(), "10") + self.assertEqual(Node(6).label(), "11") + self.assertEqual(Node(7).label(), "000") + self.assertEqual(Node(8).label(), "001") + self.assertEqual(Node(9).label(), "010") + self.assertEqual(Node(10).label(), "011") + self.assertEqual(Node(11).label(), "100") + self.assertEqual(Node(12).label(), "101") + self.assertEqual(Node(13).label(), "110") + self.assertEqual(Node(14).label(), "111") + self.assertEqual(Node(15).label(), "0000") + self.assertEqual(Node(30).label(), "1111") + + for k in _test_labeled_bases: + Node = VirtualHeap(k).Node + for b in xrange(calculate_bucket_count_in_heap_with_height(k, 2)+1): + label = Node(b).label() + level = Node(b).level + if label == "": + self.assertEqual(b, 0) + else: + self.assertEqual( + b, + basek_string_to_base10_integer(k, label) + \ + calculate_bucket_count_in_heap_with_height(k, level-1)) + + def test_is_node_on_path(self): + Node = VirtualHeap(2).Node + self.assertEqual( + Node(0).is_node_on_path( + Node(0)), + True) + self.assertEqual( + Node(0).is_node_on_path( + Node(1)), + False) + self.assertEqual( + Node(0).is_node_on_path( + Node(2)), + False) + self.assertEqual( + Node(0).is_node_on_path( + Node(3)), + False) + + Node = VirtualHeap(5).Node + self.assertEqual( + Node(20).is_node_on_path( + Node(21)), + False) + self.assertEqual( + Node(21).is_node_on_path( + Node(4)), + True) + + Node = VirtualHeap(3).Node + self.assertEqual( + Node(7).is_node_on_path( + Node(0)), + True) + self.assertEqual( + Node(7).is_node_on_path( + Node(2)), + True) + self.assertEqual( + Node(7).is_node_on_path( + Node(7)), + True) + self.assertEqual( + Node(7).is_node_on_path( + Node(8)), + False) + +class TestVirtualHeap(unittest2.TestCase): + + def test_init(self): + vh = VirtualHeap(2, blocks_per_bucket=4) + self.assertEqual(vh.k, 2) + self.assertEqual(vh.Node.k, 2) + self.assertEqual(vh.blocks_per_bucket, 4) + vh = VirtualHeap(5, blocks_per_bucket=7) + self.assertEqual(vh.k, 5) + self.assertEqual(vh.Node.k, 5) + self.assertEqual(vh.blocks_per_bucket, 7) + + def test_node_label_to_bucket(self): + vh = VirtualHeap(2) + self.assertEqual(vh.node_label_to_bucket(""), 0) + self.assertEqual(vh.node_label_to_bucket("0"), 1) + self.assertEqual(vh.node_label_to_bucket("1"), 2) + self.assertEqual(vh.node_label_to_bucket("00"), 3) + self.assertEqual(vh.node_label_to_bucket("01"), 4) + self.assertEqual(vh.node_label_to_bucket("10"), 5) + self.assertEqual(vh.node_label_to_bucket("11"), 6) + self.assertEqual(vh.node_label_to_bucket("000"), 7) + self.assertEqual(vh.node_label_to_bucket("001"), 8) + self.assertEqual(vh.node_label_to_bucket("010"), 9) + self.assertEqual(vh.node_label_to_bucket("011"), 10) + self.assertEqual(vh.node_label_to_bucket("100"), 11) + self.assertEqual(vh.node_label_to_bucket("101"), 12) + self.assertEqual(vh.node_label_to_bucket("110"), 13) + self.assertEqual(vh.node_label_to_bucket("111"), 14) + self.assertEqual(vh.node_label_to_bucket("0000"), 15) + self.assertEqual(vh.node_label_to_bucket("1111"), + calculate_bucket_count_in_heap_with_height(2, 4)-1) + + vh = VirtualHeap(3) + self.assertEqual(vh.node_label_to_bucket(""), 0) + self.assertEqual(vh.node_label_to_bucket("0"), 1) + self.assertEqual(vh.node_label_to_bucket("1"), 2) + self.assertEqual(vh.node_label_to_bucket("2"), 3) + self.assertEqual(vh.node_label_to_bucket("00"), 4) + self.assertEqual(vh.node_label_to_bucket("01"), 5) + self.assertEqual(vh.node_label_to_bucket("02"), 6) + self.assertEqual(vh.node_label_to_bucket("10"), 7) + self.assertEqual(vh.node_label_to_bucket("11"), 8) + self.assertEqual(vh.node_label_to_bucket("12"), 9) + self.assertEqual(vh.node_label_to_bucket("20"), 10) + self.assertEqual(vh.node_label_to_bucket("21"), 11) + self.assertEqual(vh.node_label_to_bucket("22"), 12) + self.assertEqual(vh.node_label_to_bucket("000"), 13) + self.assertEqual(vh.node_label_to_bucket("222"), + calculate_bucket_count_in_heap_with_height(3, 3)-1) + + for k in xrange(2, max_k_labeled+1): + for h in xrange(5): + vh = VirtualHeap(k) + largest_symbol = numerals[k-1] + self.assertEqual(vh.k, k) + self.assertEqual(vh.node_label_to_bucket(""), 0) + self.assertEqual(vh.node_label_to_bucket(largest_symbol * h), + calculate_bucket_count_in_heap_with_height(k, h)-1) + + def test_ObjectCountAtLevel(self): + for k in _test_bases: + for height in xrange(k+2): + for blocks_per_bucket in xrange(1, 5): + vh = VirtualHeap(k, blocks_per_bucket=blocks_per_bucket) + for l in xrange(height+1): + cnt = k**l + self.assertEqual(vh.bucket_count_at_level(l), cnt) + self.assertEqual(vh.node_count_at_level(l), cnt) + self.assertEqual(vh.block_count_at_level(l), + cnt * blocks_per_bucket) + + def test_bucket_to_block(self): + for k in xrange(2, 6): + for blocks_per_bucket in xrange(1, 5): + heap = VirtualHeap(k, blocks_per_bucket=blocks_per_bucket) + for b in xrange(20): + self.assertEqual(heap.bucket_to_block(b), + blocks_per_bucket * b) + + def test_node_count_at_level(self): + self.assertEqual(VirtualHeap(2).node_count_at_level(0), 1) + self.assertEqual(VirtualHeap(2).node_count_at_level(1), 2) + self.assertEqual(VirtualHeap(2).node_count_at_level(2), 4) + self.assertEqual(VirtualHeap(2).node_count_at_level(3), 8) + self.assertEqual(VirtualHeap(2).node_count_at_level(4), 16) + + self.assertEqual(VirtualHeap(3).node_count_at_level(0), 1) + self.assertEqual(VirtualHeap(3).node_count_at_level(1), 3) + self.assertEqual(VirtualHeap(3).node_count_at_level(2), 9) + self.assertEqual(VirtualHeap(3).node_count_at_level(3), 27) + self.assertEqual(VirtualHeap(3).node_count_at_level(4), 81) + + self.assertEqual(VirtualHeap(4).node_count_at_level(0), 1) + self.assertEqual(VirtualHeap(4).node_count_at_level(1), 4) + self.assertEqual(VirtualHeap(4).node_count_at_level(2), 16) + self.assertEqual(VirtualHeap(4).node_count_at_level(3), 64) + self.assertEqual(VirtualHeap(4).node_count_at_level(4), 256) + + def test_first_node_at_level(self): + self.assertEqual(VirtualHeap(2).first_node_at_level(0), 0) + self.assertEqual(VirtualHeap(2).first_node_at_level(1), 1) + self.assertEqual(VirtualHeap(2).first_node_at_level(2), 3) + self.assertEqual(VirtualHeap(2).first_node_at_level(3), 7) + self.assertEqual(VirtualHeap(2).first_node_at_level(4), 15) + + self.assertEqual(VirtualHeap(3).first_node_at_level(0), 0) + self.assertEqual(VirtualHeap(3).first_node_at_level(1), 1) + self.assertEqual(VirtualHeap(3).first_node_at_level(2), 4) + self.assertEqual(VirtualHeap(3).first_node_at_level(3), 13) + self.assertEqual(VirtualHeap(3).first_node_at_level(4), 40) + + self.assertEqual(VirtualHeap(4).first_node_at_level(0), 0) + self.assertEqual(VirtualHeap(4).first_node_at_level(1), 1) + self.assertEqual(VirtualHeap(4).first_node_at_level(2), 5) + self.assertEqual(VirtualHeap(4).first_node_at_level(3), 21) + self.assertEqual(VirtualHeap(4).first_node_at_level(4), 85) + + def test_last_node_at_level(self): + self.assertEqual(VirtualHeap(2).last_node_at_level(0), 0) + self.assertEqual(VirtualHeap(2).last_node_at_level(1), 2) + self.assertEqual(VirtualHeap(2).last_node_at_level(2), 6) + self.assertEqual(VirtualHeap(2).last_node_at_level(3), 14) + self.assertEqual(VirtualHeap(2).last_node_at_level(4), 30) + + self.assertEqual(VirtualHeap(3).last_node_at_level(0), 0) + self.assertEqual(VirtualHeap(3).last_node_at_level(1), 3) + self.assertEqual(VirtualHeap(3).last_node_at_level(2), 12) + self.assertEqual(VirtualHeap(3).last_node_at_level(3), 39) + self.assertEqual(VirtualHeap(3).last_node_at_level(4), 120) + + self.assertEqual(VirtualHeap(4).last_node_at_level(0), 0) + self.assertEqual(VirtualHeap(4).last_node_at_level(1), 4) + self.assertEqual(VirtualHeap(4).last_node_at_level(2), 20) + self.assertEqual(VirtualHeap(4).last_node_at_level(3), 84) + self.assertEqual(VirtualHeap(4).last_node_at_level(4), 340) + + def test_random_node_up_to_level(self): + for k in xrange(2,6): + heap = VirtualHeap(k) + for l in xrange(4): + for t in xrange(2 * calculate_bucket_count_in_heap_with_height(k, l)): + node = heap.random_node_up_to_level(l) + self.assertEqual(node.level <= l, True) + + def test_random_node_at_level(self): + for k in xrange(2,6): + heap = VirtualHeap(k) + for l in xrange(4): + for t in xrange(2 * calculate_bucket_count_in_heap_at_level(k, l)): + node = heap.random_node_at_level(l) + self.assertEqual(node.level == l, True) + + def test_first_block_at_level(self): + for blocks_per_bucket in xrange(1, 5): + self.assertEqual(VirtualHeap(2, blocks_per_bucket=blocks_per_bucket).\ + first_block_at_level(0), 0 * blocks_per_bucket) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=blocks_per_bucket).\ + first_block_at_level(1), 1 * blocks_per_bucket) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=blocks_per_bucket).\ + first_block_at_level(2), 3 * blocks_per_bucket) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=blocks_per_bucket).\ + first_block_at_level(3), 7 * blocks_per_bucket) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=blocks_per_bucket).\ + first_block_at_level(4), 15 * blocks_per_bucket) + + self.assertEqual(VirtualHeap(3, blocks_per_bucket=blocks_per_bucket).\ + first_block_at_level(0), 0 * blocks_per_bucket) + self.assertEqual(VirtualHeap(3, blocks_per_bucket=blocks_per_bucket).\ + first_block_at_level(1), 1 * blocks_per_bucket) + self.assertEqual(VirtualHeap(3, blocks_per_bucket=blocks_per_bucket).\ + first_block_at_level(2), 4 * blocks_per_bucket) + self.assertEqual(VirtualHeap(3, blocks_per_bucket=blocks_per_bucket).\ + first_block_at_level(3), 13 * blocks_per_bucket) + self.assertEqual(VirtualHeap(3, blocks_per_bucket=blocks_per_bucket).\ + first_block_at_level(4), 40 * blocks_per_bucket) + + self.assertEqual(VirtualHeap(4, blocks_per_bucket=blocks_per_bucket).\ + first_block_at_level(0), 0 * blocks_per_bucket) + self.assertEqual(VirtualHeap(4, blocks_per_bucket=blocks_per_bucket).\ + first_block_at_level(1), 1 * blocks_per_bucket) + self.assertEqual(VirtualHeap(4, blocks_per_bucket=blocks_per_bucket).\ + first_block_at_level(2), 5 * blocks_per_bucket) + self.assertEqual(VirtualHeap(4, blocks_per_bucket=blocks_per_bucket).\ + first_block_at_level(3), 21 * blocks_per_bucket) + self.assertEqual(VirtualHeap(4, blocks_per_bucket=blocks_per_bucket).\ + first_block_at_level(4), 85 * blocks_per_bucket) + + def test_last_block_at_level(self): + for blocks_per_bucket in xrange(1, 5): + self.assertEqual(VirtualHeap(2, blocks_per_bucket=blocks_per_bucket).\ + last_block_at_level(0), 0 * blocks_per_bucket + (blocks_per_bucket-1)) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=blocks_per_bucket).\ + last_block_at_level(1), 2 * blocks_per_bucket + (blocks_per_bucket-1)) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=blocks_per_bucket).\ + last_block_at_level(2), 6 * blocks_per_bucket + (blocks_per_bucket-1)) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=blocks_per_bucket).\ + last_block_at_level(3), 14 * blocks_per_bucket + (blocks_per_bucket-1)) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=blocks_per_bucket).\ + last_block_at_level(4), 30 * blocks_per_bucket + (blocks_per_bucket-1)) + + self.assertEqual(VirtualHeap(3, blocks_per_bucket=blocks_per_bucket).\ + last_block_at_level(0), 0 * blocks_per_bucket + (blocks_per_bucket-1)) + self.assertEqual(VirtualHeap(3, blocks_per_bucket=blocks_per_bucket).\ + last_block_at_level(1), 3 * blocks_per_bucket + (blocks_per_bucket-1)) + self.assertEqual(VirtualHeap(3, blocks_per_bucket=blocks_per_bucket).\ + last_block_at_level(2), 12 * blocks_per_bucket + (blocks_per_bucket-1)) + self.assertEqual(VirtualHeap(3, blocks_per_bucket=blocks_per_bucket).\ + last_block_at_level(3), 39 * blocks_per_bucket + (blocks_per_bucket-1)) + self.assertEqual(VirtualHeap(3, blocks_per_bucket=blocks_per_bucket).\ + last_block_at_level(4), 120 * blocks_per_bucket + (blocks_per_bucket-1)) + + self.assertEqual(VirtualHeap(4, blocks_per_bucket=blocks_per_bucket).\ + last_block_at_level(0), 0 * blocks_per_bucket + (blocks_per_bucket-1)) + self.assertEqual(VirtualHeap(4, blocks_per_bucket=blocks_per_bucket).\ + last_block_at_level(1), 4 * blocks_per_bucket + (blocks_per_bucket-1)) + self.assertEqual(VirtualHeap(4, blocks_per_bucket=blocks_per_bucket).\ + last_block_at_level(2), 20 * blocks_per_bucket + (blocks_per_bucket-1)) + self.assertEqual(VirtualHeap(4, blocks_per_bucket=blocks_per_bucket).\ + last_block_at_level(3), 84 * blocks_per_bucket + (blocks_per_bucket-1)) + self.assertEqual(VirtualHeap(4, blocks_per_bucket=blocks_per_bucket).\ + last_block_at_level(4), 340 * blocks_per_bucket + (blocks_per_bucket-1)) + + def test_block_to_bucket(self): + self.assertEqual(VirtualHeap(2, blocks_per_bucket=1).block_to_bucket(0), 0) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=1).block_to_bucket(1), 1) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=1).block_to_bucket(2), 2) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=1).block_to_bucket(3), 3) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=1).block_to_bucket(4), 4) + + self.assertEqual(VirtualHeap(2, blocks_per_bucket=2).block_to_bucket(0), 0) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=2).block_to_bucket(1), 0) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=2).block_to_bucket(2), 1) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=2).block_to_bucket(3), 1) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=2).block_to_bucket(4), 2) + + self.assertEqual(VirtualHeap(2, blocks_per_bucket=3).block_to_bucket(0), 0) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=3).block_to_bucket(1), 0) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=3).block_to_bucket(2), 0) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=3).block_to_bucket(3), 1) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=3).block_to_bucket(4), 1) + + self.assertEqual(VirtualHeap(2, blocks_per_bucket=4).block_to_bucket(0), 0) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=4).block_to_bucket(1), 0) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=4).block_to_bucket(2), 0) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=4).block_to_bucket(3), 0) + self.assertEqual(VirtualHeap(2, blocks_per_bucket=4).block_to_bucket(4), 1) + + def test_root_node(self): + for k in range(2, 6): + for blocks_per_bucket in range(1, 5): + heap = VirtualHeap(k, blocks_per_bucket=blocks_per_bucket) + root = heap.root_node() + self.assertEqual(root, 0) + self.assertEqual(root.bucket, 0) + self.assertEqual(root.level, 0) + self.assertEqual(root.parent_node(), None) + +class TestSizedVirtualHeap(unittest2.TestCase): + + def test_init(self): + vh = SizedVirtualHeap(2, 8, blocks_per_bucket=4) + self.assertEqual(vh.k, 2) + self.assertEqual(vh.Node.k, 2) + self.assertEqual(vh.blocks_per_bucket, 4) + vh = SizedVirtualHeap(5, 9, blocks_per_bucket=7) + self.assertEqual(vh.k, 5) + self.assertEqual(vh.Node.k, 5) + self.assertEqual(vh.blocks_per_bucket, 7) + + def test_height(self): + vh = SizedVirtualHeap(2, 3, blocks_per_bucket=4) + self.assertEqual(vh.height, 3) + vh = SizedVirtualHeap(5, 6, blocks_per_bucket=7) + self.assertEqual(vh.height, 6) + + def test_levels(self): + vh = SizedVirtualHeap(2, 3, blocks_per_bucket=4) + self.assertEqual(vh.levels, 4) + vh = SizedVirtualHeap(5, 6, blocks_per_bucket=7) + self.assertEqual(vh.levels, 7) + + def test_first_level(self): + vh = SizedVirtualHeap(2, 3, blocks_per_bucket=4) + self.assertEqual(vh.first_level, 0) + vh = SizedVirtualHeap(5, 6, blocks_per_bucket=7) + self.assertEqual(vh.first_level, 0) + + def test_last_level(self): + vh = SizedVirtualHeap(2, 3, blocks_per_bucket=4) + self.assertEqual(vh.last_level, 3) + self.assertEqual(vh.last_level, vh.levels-1) + self.assertEqual(vh.last_level, vh.height) + vh = SizedVirtualHeap(5, 6, blocks_per_bucket=7) + self.assertEqual(vh.last_level, 6) + self.assertEqual(vh.last_level, vh.levels-1) + self.assertEqual(vh.last_level, vh.height) + + def test_ObjectCount(self): + for k in _test_bases: + for height in xrange(k+2): + for blocks_per_bucket in xrange(1, 5): + vh = SizedVirtualHeap(k, + height, + blocks_per_bucket=blocks_per_bucket) + cnt = (((k**(height+1))-1)//(k-1)) + self.assertEqual(vh.bucket_count(), cnt) + self.assertEqual(vh.node_count(), cnt) + self.assertEqual(vh.block_count(), cnt * blocks_per_bucket) + + def test_LeafObjectCount(self): + for k in _test_bases: + for height in xrange(k+2): + for blocks_per_bucket in xrange(1, 5): + vh = SizedVirtualHeap(k, + height, + blocks_per_bucket=blocks_per_bucket) + self.assertEqual(vh.leaf_bucket_count(), + vh.bucket_count_at_level(vh.height)) + self.assertEqual(vh.leaf_node_count(), + vh.node_count_at_level(vh.height)) + self.assertEqual(vh.leaf_block_count(), + vh.block_count_at_level(vh.height)) + + + def test_FirstLeafObject(self): + vh = SizedVirtualHeap(2, 3, blocks_per_bucket=3) + self.assertEqual(vh.first_leaf_node(), 7) + self.assertEqual(vh.first_leaf_block(), 7*3) + + def test_LastLeafObject(self): + vh = SizedVirtualHeap(2, 3, blocks_per_bucket=3) + self.assertEqual(vh.last_leaf_node(), 14) + self.assertEqual(vh.last_leaf_block(), 14*3 + 2) + + def test_random_node(self): + for k in xrange(2,6): + height = 3 + heap = SizedVirtualHeap(k, height) + for t in xrange(2 * heap.bucket_count()): + node = heap.random_node() + self.assertEqual(0 <= node.level <= height, True) + + def test_random_leaf_node(self): + for k in xrange(2,6): + height = 3 + heap = SizedVirtualHeap(k, height) + for t in xrange(2 * heap.bucket_count()): + node = heap.random_leaf_node() + self.assertEqual(node.level, height) + + def _assert_file_equals_baselines(self, fname, bname): + with open(fname)as f: + flines = f.readlines() + with open(bname) as f: + blines = f.readlines() + self.assertListEqual(flines, blines) + os.remove(fname) + + def test_write_as_dot(self): + + for k, h, b, maxl in [(2, 3, 1, None), + (2, 3, 2, None), + (3, 3, 1, None), + (3, 3, 2, None), + (3, 10, 2, 4), + (200, 0, 1, None)]: + if maxl is None: + label = "k%d_h%d_b%d" % (k, h, b) + else: + label = "k%d_h%d_b%d" % (k, maxl-1, b) + heap = SizedVirtualHeap(k, h, blocks_per_bucket=b) + + fname = label+".dot" + with open(os.path.join(thisdir, fname), "w") as f: + heap.write_as_dot(f, max_levels=maxl) + self._assert_file_equals_baselines( + os.path.join(thisdir, fname), + os.path.join(baselinedir, fname)) + + data = list(range(heap.block_count())) + fname = label+"_data.dot" + with open(os.path.join(thisdir, fname), "w") as f: + heap.write_as_dot(f, data=data, max_levels=maxl) + self._assert_file_equals_baselines( + os.path.join(thisdir, fname), + os.path.join(baselinedir, fname)) + + def test_save_image_as_pdf(self): + + for k, h, b, maxl in [(2, 3, 1, None), + (2, 3, 2, None), + (3, 3, 1, None), + (3, 3, 2, None), + (3, 10, 2, 4)]: + if maxl is None: + label = "k%d_h%d_b%d" % (k, h, b) + else: + label = "k%d_h%d_b%d" % (k, maxl-1, b) + heap = SizedVirtualHeap(k, h, blocks_per_bucket=b) + + fname = label+".pdf" + try: + os.remove(os.path.join(thisdir, fname)) + except OSError: # pragma: no cover + pass # pragma: no cover + + rc = heap.save_image_as_pdf(os.path.join(thisdir, label), + max_levels=maxl) + + if not has_dot: + self.assertEqual(rc, False) + else: + self.assertEqual(rc, True) + self.assertEqual( + os.path.exists(os.path.join(thisdir, fname)), True) + try: + os.remove(os.path.join(thisdir, fname)) + except OSError: # pragma: no cover + pass # pragma: no cover + + data = list(range(heap.block_count())) + fname = label+"_data.pdf" + try: + os.remove(os.path.join(thisdir, fname)) + except OSError: # pragma: no cover + pass # pragma: no cover + rc = heap.save_image_as_pdf(os.path.join(thisdir, fname), + data=data, + max_levels=maxl) + if not has_dot: + self.assertEqual(rc, False) + else: + self.assertEqual(rc, True) + self.assertEqual( + os.path.exists(os.path.join(thisdir, fname)), True) + try: + os.remove(os.path.join(thisdir, fname)) + except OSError: # pragma: no cover + pass # pragma: no cover + +class TestMisc(unittest2.TestCase): + + def test_calculate_bucket_level(self): + self.assertEqual(calculate_bucket_level(2, 0), 0) + self.assertEqual(calculate_bucket_level(2, 1), 1) + self.assertEqual(calculate_bucket_level(2, 2), 1) + self.assertEqual(calculate_bucket_level(2, 3), 2) + self.assertEqual(calculate_bucket_level(2, 4), 2) + self.assertEqual(calculate_bucket_level(2, 5), 2) + self.assertEqual(calculate_bucket_level(2, 6), 2) + self.assertEqual(calculate_bucket_level(2, 7), 3) + + self.assertEqual(calculate_bucket_level(3, 0), 0) + self.assertEqual(calculate_bucket_level(3, 1), 1) + self.assertEqual(calculate_bucket_level(3, 2), 1) + self.assertEqual(calculate_bucket_level(3, 3), 1) + self.assertEqual(calculate_bucket_level(3, 4), 2) + self.assertEqual(calculate_bucket_level(3, 5), 2) + self.assertEqual(calculate_bucket_level(3, 6), 2) + self.assertEqual(calculate_bucket_level(3, 7), 2) + self.assertEqual(calculate_bucket_level(3, 8), 2) + self.assertEqual(calculate_bucket_level(3, 9), 2) + self.assertEqual(calculate_bucket_level(3, 10), 2) + self.assertEqual(calculate_bucket_level(3, 11), 2) + self.assertEqual(calculate_bucket_level(3, 12), 2) + self.assertEqual(calculate_bucket_level(3, 13), 3) + + self.assertEqual(calculate_bucket_level(4, 0), 0) + self.assertEqual(calculate_bucket_level(4, 1), 1) + self.assertEqual(calculate_bucket_level(4, 2), 1) + self.assertEqual(calculate_bucket_level(4, 3), 1) + self.assertEqual(calculate_bucket_level(4, 4), 1) + + self.assertEqual(calculate_bucket_level(4, 5), 2) + self.assertEqual(calculate_bucket_level(4, 6), 2) + self.assertEqual(calculate_bucket_level(4, 7), 2) + self.assertEqual(calculate_bucket_level(4, 8), 2) + + self.assertEqual(calculate_bucket_level(4, 9), 2) + self.assertEqual(calculate_bucket_level(4, 10), 2) + self.assertEqual(calculate_bucket_level(4, 11), 2) + self.assertEqual(calculate_bucket_level(4, 12), 2) + + self.assertEqual(calculate_bucket_level(4, 13), 2) + self.assertEqual(calculate_bucket_level(4, 14), 2) + self.assertEqual(calculate_bucket_level(4, 15), 2) + self.assertEqual(calculate_bucket_level(4, 16), 2) + + self.assertEqual(calculate_bucket_level(4, 17), 2) + self.assertEqual(calculate_bucket_level(4, 18), 2) + self.assertEqual(calculate_bucket_level(4, 19), 2) + self.assertEqual(calculate_bucket_level(4, 20), 2) + + self.assertEqual(calculate_bucket_level(4, 21), 3) + + def test_clib_calculate_bucket_level(self): + for k in _test_bases: + for b in xrange(calculate_bucket_count_in_heap_with_height(k, 2)+2): + self.assertEqual(_clib.calculate_bucket_level(k, b), + calculate_bucket_level(k, b)) + for k, b in [(89, 14648774), + (89, 14648775), + (90, 14648774), + (90, 14648775)]: + self.assertEqual(_clib.calculate_bucket_level(k, b), + calculate_bucket_level(k, b)) + + def test_clib_calculate_last_common_level(self): + for k in range(2, 8): + for b1 in xrange(calculate_bucket_count_in_heap_with_height(k, 2)+2): + for b2 in xrange(calculate_bucket_count_in_heap_with_height(k, 2)+2): + self.assertEqual(_clib.calculate_last_common_level(k, b1, b2), + calculate_last_common_level(k, b1, b2)) + for k in [89,90]: + for b1 in [0, 100, 10000, 14648774, 14648775]: + for b2 in [0, 100, 10000, 14648774, 14648775]: + self.assertEqual(_clib.calculate_last_common_level(k, b1, b2), + calculate_last_common_level(k, b1, b2)) + + def test_calculate_necessary_heap_height(self): + self.assertEqual(calculate_necessary_heap_height(2, 1), 0) + self.assertEqual(calculate_necessary_heap_height(2, 2), 1) + self.assertEqual(calculate_necessary_heap_height(2, 3), 1) + self.assertEqual(calculate_necessary_heap_height(2, 4), 2) + self.assertEqual(calculate_necessary_heap_height(2, 5), 2) + self.assertEqual(calculate_necessary_heap_height(2, 6), 2) + self.assertEqual(calculate_necessary_heap_height(2, 7), 2) + self.assertEqual(calculate_necessary_heap_height(2, 8), 3) + + self.assertEqual(calculate_necessary_heap_height(3, 1), 0) + self.assertEqual(calculate_necessary_heap_height(3, 2), 1) + self.assertEqual(calculate_necessary_heap_height(3, 3), 1) + self.assertEqual(calculate_necessary_heap_height(3, 4), 1) + self.assertEqual(calculate_necessary_heap_height(3, 5), 2) + self.assertEqual(calculate_necessary_heap_height(3, 6), 2) + self.assertEqual(calculate_necessary_heap_height(3, 7), 2) + self.assertEqual(calculate_necessary_heap_height(3, 8), 2) + self.assertEqual(calculate_necessary_heap_height(3, 9), 2) + self.assertEqual(calculate_necessary_heap_height(3, 10), 2) + self.assertEqual(calculate_necessary_heap_height(3, 11), 2) + self.assertEqual(calculate_necessary_heap_height(3, 12), 2) + self.assertEqual(calculate_necessary_heap_height(3, 13), 2) + self.assertEqual(calculate_necessary_heap_height(3, 14), 3) + self.assertEqual(calculate_necessary_heap_height(3, 15), 3) + self.assertEqual(calculate_necessary_heap_height(3, 16), 3) + +if __name__ == "__main__": + unittest2.main() # pragma: no cover diff --git a/PyORAM/src/pyoram/util/__init__.py b/PyORAM/src/pyoram/util/__init__.py new file mode 100644 index 0000000..f9dcdf5 --- /dev/null +++ b/PyORAM/src/pyoram/util/__init__.py @@ -0,0 +1,2 @@ +import pyoram.util.misc +import pyoram.util.virtual_heap diff --git a/PyORAM/src/pyoram/util/misc.py b/PyORAM/src/pyoram/util/misc.py new file mode 100644 index 0000000..9ae2235 --- /dev/null +++ b/PyORAM/src/pyoram/util/misc.py @@ -0,0 +1,117 @@ +import base64 + +import six + +def log2floor(n): + """ + Returns the exact value of floor(log2(n)). + No floating point calculations are used. + Requires positive integer type. + """ + assert n > 0 + return n.bit_length() - 1 + +def log2ceil(n): + """ + Returns the exact value of ceil(log2(n)). + No floating point calculations are used. + Requires positive integer type. + """ + if n == 1: + return 0 + return log2floor(n-1) + 1 + +def intdivceil(x, y): + """ + Returns the exact value of ceil(x // y). + No floating point calculations are used. + Requires positive integer types. The result + is undefined if at least one of the inputs + is floating point. + """ + result = x // y + if (x % y): + result += 1 + return result + +def save_private_key(filename, key): + with open(filename, "wb") as f: + f.write(base64.b64encode(key)) + +def load_private_key(filename): + with open(filename, "rb") as f: + return base64.b64decode(f.read()) + +from fractions import Fraction + +class MemorySize(object): + + to_bytes = {} + to_bytes['b'] = lambda x: Fraction(x,8) + to_bytes['B'] = lambda x: Fraction(x,1) + to_bytes['KB'] = lambda x: Fraction(1000*x,1) + to_bytes['MB'] = lambda x: Fraction((1000**2)*x,1) + to_bytes['GB'] = lambda x: Fraction((1000**3)*x,1) + to_bytes['TB'] = lambda x: Fraction((1000**4)*x,1) + to_bytes['KiB'] = lambda x: Fraction(1024*x,1) + to_bytes['MiB'] = lambda x: Fraction((1024**2)*x,1) + to_bytes['GiB'] = lambda x: Fraction((1024**3)*x,1) + to_bytes['TiB'] = lambda x: Fraction((1024**4)*x,1) + + def __init__(self, size, unit='B'): + assert size >= 0 + self.numbytes = MemorySize.to_bytes[unit](Fraction.from_float(size)) + + def __str__(self): + if self.B < 1: + return "%.3f b" % (self.b) + if self.KB < 1: + return "%.3f B" % (self.B) + if self.MB < 1: + return "%.3f KB" % (self.KB) + if self.GB < 1: + return "%.3f MB" % (self.MB) + if self.TB < 1: + return "%.3f GB" % (self.GB) + return "%.3f TB" % (self.TB) + + @property + def b(self): return self.numbytes*8 + @property + def B(self): return self.numbytes + + @property + def KB(self): return self.B/1000 + @property + def MB(self): return self.KB/1000 + @property + def GB(self): return self.MB/1000 + @property + def TB(self): return self.GB/1000 + + @property + def KiB(self): return self.B/1024 + @property + def MiB(self): return self.KiB/1024 + @property + def GiB(self): return self.MiB/1024 + @property + def TiB(self): return self.GiB/1024 + +def chunkiter(objs, n=100): + """ + Chunk an iterator of unknown size. The optional + keyword 'n' sets the chunk size (default 100). + """ + + objs = iter(objs) + try: + while (1): + chunk = [] + while len(chunk) < n: + chunk.append(six.next(objs)) + yield chunk + except StopIteration: + pass + if len(chunk): + yield chunk diff --git a/PyORAM/src/pyoram/util/virtual_heap.py b/PyORAM/src/pyoram/util/virtual_heap.py new file mode 100644 index 0000000..22b83c6 --- /dev/null +++ b/PyORAM/src/pyoram/util/virtual_heap.py @@ -0,0 +1,441 @@ +__all__ = ("VirtualHeap", + "SizedVirtualHeap") + +import os +import sys +import subprocess +import random +import string +import tempfile + +from six.moves import xrange + +from pyoram.util._virtual_heap_helper import lib as _clib +from pyoram.util.misc import log2floor + +numerals = ''.join([c for c in string.printable \ + if ((c not in string.whitespace) and \ + (c != '+') and (c != '-') and \ + (c != '"') and (c != "'") and \ + (c != '\\') and (c != '/'))]) +numeral_index = dict((c,i) for i,c in enumerate(numerals)) + +# The maximum heap base for which base k labels +# can be produced. +max_k_labeled = len(numerals) + +def base10_integer_to_basek_string(k, x): + """Convert an integer into a base k string.""" + if not (2 <= k <= max_k_labeled): + raise ValueError("k must be in range [2, %d]: %s" + % (max_k_labeled, k)) + return ((x == 0) and numerals[0]) or \ + (base10_integer_to_basek_string(k, x // k).\ + lstrip(numerals[0]) + numerals[x % k]) + +def basek_string_to_base10_integer(k, x): + """Convert a base k string into an integer.""" + assert 1 < k <= max_k_labeled + return sum(numeral_index[c]*(k**i) + for i, c in enumerate(reversed(x))) + +# _clib defines a faster version of this function +def calculate_bucket_level(k, b): + """ + Calculate the level in which a 0-based bucket + lives inside of a k-ary heap. + """ + assert k >= 2 + if k == 2: + return log2floor(b+1) + v = (k - 1) * (b + 1) + 1 + h = 0 + while k**(h+1) < v: + h += 1 + return h + +# _clib defines a faster version of this function +def calculate_last_common_level(k, b1, b2): + """ + Calculate the highest level after which the + paths from the root to these buckets diverge. + """ + l1 = calculate_bucket_level(k, b1) + l2 = calculate_bucket_level(k, b2) + while l1 > l2: + b1 = (b1-1)//k + l1 -= 1 + while l2 > l1: + b2 = (b2-1)//k + l2 -= 1 + while b1 != b2: + b1 = (b1-1)//k + b2 = (b2-1)//k + l1 -= 1 + return l1 + +def calculate_necessary_heap_height(k, n): + """ + Calculate the necessary k-ary heap height + to store n buckets. + """ + assert n >= 1 + return calculate_bucket_level(k, n-1) + +def calculate_bucket_count_in_heap_with_height(k, h): + """ + Calculate the number of buckets in a + k-ary heap of height h. + """ + assert h >= 0 + return ((k**(h+1)) - 1) // (k - 1) + +def calculate_bucket_count_in_heap_at_level(k, l): + """ + Calculate the number of buckets in a + k-ary heap at level l. + """ + assert l >= 0 + return k**l + +def calculate_leaf_bucket_count_in_heap_with_height(k, h): + """ + Calculate the number of buckets in the + leaf-level of a k-ary heap of height h. + """ + return calculate_bucket_count_in_heap_at_level(k, h) + +def create_node_type(k): + + class VirtualHeapNode(object): + __slots__ = ("bucket", "level") + def __init__(self, bucket): + assert bucket >= 0 + self.bucket = bucket + self.level = _clib.calculate_bucket_level(self.k, self.bucket) + + def __hash__(self): + return self.bucket.__hash__() + def __int__(self): + return self.bucket + def __lt__(self, other): + return self.bucket < other + def __le__(self, other): + return self.bucket <= other + def __eq__(self, other): + return self.bucket == other + def __ne__(self, other): + return self.bucket != other + def __gt__(self, other): + return self.bucket > other + def __ge__(self, other): + return self.bucket >= other + def last_common_level(self, n): + return _clib.calculate_last_common_level(self.k, + self.bucket, + n.bucket) + def child_node(self, c): + assert type(c) is int + assert 0 <= c < self.k + return VirtualHeapNode(self.k * self.bucket + 1 + c) + def parent_node(self): + if self.bucket != 0: + return VirtualHeapNode((self.bucket - 1)//self.k) + return None + def ancestor_node_at_level(self, level): + if level > self.level: + return None + current = self + while current.level != level: + current = current.parent_node() + return current + def path_to_root(self): + bucket = self.bucket + yield self + while bucket != 0: + bucket = (bucket - 1)//self.k + yield type(self)(bucket) + def path_from_root(self): + return list(reversed(list(self.path_to_root()))) + def bucket_path_to_root(self): + bucket = self.bucket + yield bucket + while bucket != 0: + bucket = (bucket - 1)//self.k + yield bucket + def bucket_path_from_root(self): + return list(reversed(list(self.bucket_path_to_root()))) + + # + # Expensive Functions + # + def __repr__(self): + try: + label = self.label() + except ValueError: + # presumably, k is too large + label = "" + return ("VirtualHeapNode(k=%s, bucket=%s, level=%s, label=%r)" + % (self.k, self.bucket, self.level, label)) + def __str__(self): + """Returns a tuple (, ).""" + if self.bucket != 0: + return ("(%s, %s)" + % (self.level, + self.bucket - + calculate_bucket_count_in_heap_with_height(self.k, + self.level-1))) + assert self.level == 0 + return "(0, 0)" + + def label(self): + assert 0 <= self.bucket + if self.level == 0: + return '' + b_offset = self.bucket - \ + calculate_bucket_count_in_heap_with_height(self.k, + self.level-1) + basek = base10_integer_to_basek_string(self.k, b_offset) + return basek.zfill(self.level) + + def is_node_on_path(self, n): + if n.level <= self.level: + n_label = n.label() + if n_label == "": + return True + return self.label().startswith(n_label) + return False + + VirtualHeapNode.k = k + + return VirtualHeapNode + +class VirtualHeap(object): + + clib = _clib + random = random.SystemRandom() + + def __init__(self, k, blocks_per_bucket=1): + assert 1 < k + assert blocks_per_bucket >= 1 + self._k = k + self._blocks_per_bucket = blocks_per_bucket + self.Node = create_node_type(k) + + @property + def k(self): + return self._k + + def node_label_to_bucket(self, label): + if len(label) > 0: + return \ + (calculate_bucket_count_in_heap_with_height(self.k, + len(label)-1) + + basek_string_to_base10_integer(self.k, label)) + return 0 + + # + # Buckets (0-based integer, equivalent to block for heap + # with blocks_per_bucket=1) + # + + @property + def blocks_per_bucket(self): + return self._blocks_per_bucket + + def bucket_count_at_level(self, l): + return calculate_bucket_count_in_heap_at_level(self.k, l) + def first_bucket_at_level(self, l): + if l > 0: + return calculate_bucket_count_in_heap_with_height(self.k, l-1) + return 0 + def last_bucket_at_level(self, l): + return calculate_bucket_count_in_heap_with_height(self.k, l) - 1 + def random_bucket_up_to_level(self, l): + return self.random.randint(self.first_bucket_at_level(0), + self.last_bucket_at_level(l)) + def random_bucket_at_level(self, l): + return self.random.randint(self.first_bucket_at_level(l), + self.first_bucket_at_level(l+1)-1) + + # + # Nodes (a class that helps with heap path calculations) + # + + def root_node(self): + return self.first_node_at_level(0) + def node_count_at_level(self, l): + return self.bucket_count_at_level(l) + def first_node_at_level(self, l): + return self.Node(self.first_bucket_at_level(l)) + def last_node_at_level(self, l): + return self.Node(self.last_bucket_at_level(l)) + def random_node_up_to_level(self, l): + return self.Node(self.random_bucket_up_to_level(l)) + def random_node_at_level(self, l): + return self.Node(self.random_bucket_at_level(l)) + + # + # Block (0-based integer) + # + + def bucket_to_block(self, b): + assert b >= 0 + return b * self.blocks_per_bucket + def block_to_bucket(self, s): + assert s >= 0 + return s//self.blocks_per_bucket + def first_block_in_bucket(self, b): + return self.bucket_to_block(b) + def last_block_in_bucket(self, b): + return self.bucket_to_block(b) + self.blocks_per_bucket - 1 + def block_count_at_level(self, l): + return self.bucket_count_at_level(l) * self.blocks_per_bucket + def first_block_at_level(self, l): + return self.bucket_to_block(self.first_bucket_at_level(l)) + def last_block_at_level(self, l): + return self.bucket_to_block(self.first_bucket_at_level(l+1)) - 1 + +class SizedVirtualHeap(VirtualHeap): + + def __init__(self, k, height, blocks_per_bucket=1): + super(SizedVirtualHeap, self).\ + __init__(k, blocks_per_bucket=blocks_per_bucket) + self._height = height + + # + # Size properties + # + @property + def height(self): + return self._height + @property + def levels(self): + return self.height + 1 + @property + def first_level(self): + return 0 + @property + def last_level(self): + return self.height + + # + # Buckets (0-based integer, equivalent to block for heap + # with blocks_per_bucket=1) + # + + def bucket_count(self): + return calculate_bucket_count_in_heap_with_height(self.k, + self.height) + def leaf_bucket_count(self): + return calculate_leaf_bucket_count_in_heap_with_height(self.k, + self.height) + def first_leaf_bucket(self): + return self.first_bucket_at_level(self.height) + def last_leaf_bucket(self): + return self.last_bucket_at_level(self.height) + def random_bucket(self): + return self.random.randint(self.first_bucket_at_level(0), + self.last_leaf_bucket()) + def random_leaf_bucket(self): + return self.random_bucket_at_level(self.height) + + # + # Nodes (a class that helps with heap path calculations) + # + + def is_nil_node(self, n): + return n.bucket >= self.bucket_count() + def node_count(self): + return self.bucket_count() + def leaf_node_count(self): + return self.leaf_bucket_count() + def first_leaf_node(self): + return self.Node(self.first_leaf_bucket()) + def last_leaf_node(self): + return self.Node(self.last_leaf_bucket()) + def random_leaf_node(self): + return self.Node(self.random_leaf_bucket()) + def random_node(self): + return self.Node(self.random_bucket()) + + # + # Block (0-based integer) + # + + def block_count(self): + return self.bucket_count() * self.blocks_per_bucket + def leaf_block_count(self): + return self.leaf_bucket_count() * self.blocks_per_bucket + def first_leaf_block(self): + return self.first_block_in_bucket(self.first_leaf_bucket()) + def last_leaf_block(self): + return self.last_block_in_bucket(self.last_leaf_bucket()) + + # + # Visualization + # + + def write_as_dot(self, f, data=None, max_levels=None): + "Write the tree in the dot language format to f." + assert (max_levels is None) or (max_levels >= 0) + def visit_node(n, levels): + lbl = "{" + if data is None: + if self.k <= max_k_labeled: + lbl = repr(n.label()).\ + replace("{","\{").\ + replace("}","\}").\ + replace("|","\|").\ + replace("<","\<").\ + replace(">","\>") + else: + lbl = str(n) + else: + s = self.bucket_to_block(n.bucket) + for i in xrange(self.blocks_per_bucket): + lbl += "{%s}" % (data[s+i]) + if i + 1 != self.blocks_per_bucket: + lbl += "|" + lbl += "}" + f.write(" %s [penwidth=%s,label=\"%s\"];\n" + % (n.bucket, 1, lbl)) + levels += 1 + if (max_levels is None) or (levels <= max_levels): + for i in xrange(self.k): + cn = n.child_node(i) + if not self.is_nil_node(cn): + visit_node(cn, levels) + f.write(" %s -> %s ;\n" % (n.bucket, cn.bucket)) + + f.write("// Created by SizedVirtualHeap.write_as_dot(...)\n") + f.write("digraph heaptree {\n") + f.write("node [shape=record]\n") + + if (max_levels is None) or (max_levels > 0): + visit_node(self.root_node(), 1) + f.write("}\n") + + def save_image_as_pdf(self, filename, data=None, max_levels=None): + "Write the heap as PDF file." + assert (max_levels is None) or (max_levels >= 0) + import os + if not filename.endswith('.pdf'): + filename = filename+'.pdf' + tmpfd, tmpname = tempfile.mkstemp(suffix='dot') + with open(tmpname, 'w') as f: + self.write_as_dot(f, data=data, max_levels=max_levels) + os.close(tmpfd) + try: + subprocess.call(['dot', + tmpname, + '-Tpdf', + '-o', + ('%s'%filename)]) + except OSError: + sys.stderr.write( + "DOT -> PDF conversion failed. See DOT file: %s\n" + % (tmpname)) + return False + os.remove(tmpname) + return True diff --git a/version2/src/java/light_fan_embed_benchmark/Sensor.java b/version2/src/java/light_fan_embed_benchmark/Sensor.java index 611b664..8033677 100644 --- a/version2/src/java/light_fan_embed_benchmark/Sensor.java +++ b/version2/src/java/light_fan_embed_benchmark/Sensor.java @@ -8,57 +8,76 @@ class Sensor { Table t1 = new Table("http://dc-6.calit2.uci.edu/test.iotcloud/", "reallysecret", 361, -1); + + try { + Thread.sleep(5000); + } catch (Exception e) { + + } + + long start = System.currentTimeMillis(); t1.rebuild(); + + long stop1 = System.currentTimeMillis(); + t1.update(); + long stop2 = System.currentTimeMillis(); + System.out.println("Done......"); + System.out.println(stop1 - start); + System.out.println(stop2 - stop1); - String pingTimerKey = "sensorController"; - IoTString ipingTimerKey = new IoTString(pingTimerKey); + // t1.startTransaction(); + // t1.addKV(ipingTimerKey, ipingTimer); + // t1.addKV(ia1, senDat); + // t1.commitTransaction(); - String a1 = "sensor"; - IoTString ia1 = new IoTString(a1); - System.out.println("Starting System"); + // String pingTimerKey = "sensorController"; + // IoTString ipingTimerKey = new IoTString(pingTimerKey); + // String a1 = "sensor"; + // IoTString ia1 = new IoTString(a1); - while (true) { - try { + // System.out.println("Starting System"); - Runtime runTime = Runtime.getRuntime(); - // Process proc = runTime.exec("/opt/vc/bin/vcgencmd measure_temp | tr -d 'temp=' | tr -d \"'C\""); - Process proc = runTime.exec("/opt/vc/bin/vcgencmd measure_temp"); - BufferedReader reader = new BufferedReader(new InputStreamReader(proc.getInputStream())); - String line = null; - String dat = ""; - while ((line = reader.readLine()) != null) { - System.out.println(line); - dat = line; - } - reader.close(); + // while (true) { + // try { - String pingTimer = Long.toString(System.currentTimeMillis()); - IoTString ipingTimer = new IoTString(pingTimer); - IoTString senDat = new IoTString(dat); + // // Runtime runTime = Runtime.getRuntime(); + // // // Process proc = runTime.exec("/opt/vc/bin/vcgencmd measure_temp | tr -d 'temp=' | tr -d \"'C\""); + // // Process proc = runTime.exec("/opt/vc/bin/vcgencmd measure_temp"); + // // BufferedReader reader = new BufferedReader(new InputStreamReader(proc.getInputStream())); + // // String line = null; + // // String dat = ""; + // // while ((line = reader.readLine()) != null) { + // // System.out.println(line); + // // dat = line; + // // } + // // reader.close(); - t1.update(); - t1.startTransaction(); - t1.addKV(ipingTimerKey, ipingTimer); - t1.addKV(ia1, senDat); - t1.commitTransaction(); + // // String pingTimer = Long.toString(System.currentTimeMillis()); + // // IoTString ipingTimer = new IoTString(pingTimer); + // IoTString senDat = new IoTString(dat); + // t1.update(); + // t1.startTransaction(); + // t1.addKV(ipingTimerKey, ipingTimer); + // t1.addKV(ia1, senDat); + // t1.commitTransaction(); @@ -66,22 +85,22 @@ class Sensor { - Thread.sleep(5000); + // Thread.sleep(5000); - } catch (Error e) { - e.printStackTrace(); + // } catch (Error e) { + // e.printStackTrace(); - Runtime runTime = Runtime.getRuntime(); - runTime.exec("gpio mode 4 out"); + // Runtime runTime = Runtime.getRuntime(); + // runTime.exec("gpio mode 4 out"); - while (true) { - runTime.exec("gpio write 4 1"); - Thread.sleep(500); - runTime.exec("gpio write 4 0"); - Thread.sleep(500); - } - } - } + // while (true) { + // runTime.exec("gpio write 4 1"); + // Thread.sleep(500); + // runTime.exec("gpio write 4 0"); + // Thread.sleep(500); + // } + // } + // } } } \ No newline at end of file -- 2.34.1