2 files changed, 67 insertions, 61 deletions
diff --git a/py11/mutex.py b/py11/mutex.py
new file mode 100644
index 0000000..5ad30ba
--- /dev/null
+++ b/py11/mutex.py
@@ -0,0 +1,63 @@
+# Optional mutex implementation.
+
+from struct import pack, unpack
+
+# This controls how big our mutex handles are.
+encoded_format = "=L"
+
+# These are all determined by encoded_format, don't touch.
+
+encoded_length = len(pack(encoded_format, 0))
+encoded_type   = CK_BYTE * encoded_length
+handle_max     = unpack(encoded_format, chr(0xff) * encoded_length)[0]
+
+class Mutex(object):
+
+  def __init__(self, handle):
+    from threading import Lock
+    self.encoded = encoded_type(*handle)
+    self.lock    = Lock()
+
+def p11_callback(func):
+  from threading import ThreadError
+  def wrapper(self, arg):
+    try:
+      func(self, arg)
+    except ThreadError:
+      return CKR_MUTEX_NOT_LOCKED
+    except:
+      return CKR_FUNCTION_FAILED
+    else:
+      return CKR_OK
+  return wrapper
+
+class MutexDB(object):
+
+  def __init__(self):
+    self.mutexes = {}
+    self.next_handle = 0
+
+  def find_free_handle(self):
+    if len(self.mutexes) > handle_max:
+      raise RuntimeError
+    while self.next_handle in self.mutexes:
+      self.next_handle = (self.next_handle + 1) & handle_max
+    return pack(encoded_format, self.next_handle)
+
+  @p11_callback
+  def create(self, handle_ptr):
+    handle = self.find_free_handle()
+    self.mutexes[handle] = self.mutex_type(handle)
+    handle_ptr[0] = self.mutexes[handle].encoded
+
+  @p11_callback
+  def destroy(self, handle):
+    del self.mutexes[handle[:encoded_length]]
+
+  @p11_callback
+  def lock(self, handle):
+    self.mutexes[handle[:encoded_length]].lock.acquire()
+
+  @p11_callback
+  def unlock(self, handle):
+    self.mutexes[handle[:encoded_length]].lock.release()
diff --git a/scripts/py11-test.py b/scripts/py11-test.py
index a5e9c88..1719f8b 100644
--- a/scripts/py11-test.py
+++ b/scripts/py11-test.py
@@ -71,72 +71,15 @@ if False:
   print "Not using MUTEXes at all"
   p11.C_Initialize()
 
-elif True:
+elif False:
   print "Using OS MUTEXes"
   p11.C_Initialize(CKF_OS_LOCKING_OK)
 
 else:
   print "Using our own pseudo-MUTEXes"
-
-  # This class probably belongs in the library.
-
-  class Mutexes(object):
-
-    format = "=L"
-    length = len(pack(format, 0))
-
-    def __init__(self, verbose = True, flags = 0):
-      self.init_args = (flags, self.create, self.destroy, self.lock, self.unlock)
-      self.next_mutex = 0
-      self.mutexes = {}
-      self.verbose = verbose
-
-    def find(self):
-      if len(self.mutexes) > 0xFFFFFFFF:
-        raise RuntimeError
-      while self.next_mutex in self.mutexes:
-        self.next_mutex = (self.next_mutex + 1) & 0xFFFFFFFF
-
-    def encode(self, n):
-      return (CK_BYTE * self.length)(*pack(self.format, n))
-
-    def decode(self, m):
-      return unpack(self.format, m[:self.length])[0]
-
-    def create(self, mm):
-      self.find()
-      if self.verbose:
-        print "Creating mutex", self.next_mutex
-      m = self.encode(self.next_mutex)
-      self.mutexes[self.next_mutex] = [False, m]
-      mm[0] = m
-      return CKR_OK
-
-    def destroy(self, m):
-      mutex = self.decode(m)
-      if self.verbose:
-        print "Destroying mutex", mutex
-      del self.mutexes[mutex]
-      return CKR_OK
-
-    def lock(self, m):
-      mutex = self.decode(m)
-      if self.verbose:
-        print "Locking mutex", mutex
-      self.mutexes[mutex][0] = True
-      return CKR_OK
-
-    def unlock(self, m):
-      mutex = self.decode(m)
-      if self.verbose:
-        print "Unlocking mutex", mutex
-      self.mutexes[mutex][0] = False
-      return CKR_OK
-
-  mutexes = Mutexes()
-  p11.C_Initialize(*mutexes.init_args)
-
-# Dun with MUTEX insanit and C_Initialize() call
+  from py11.mutex import MutexDB
+  mdb = MutexDB()
+  p11.C_Initialize(0, mdb.create, mdb.destroy, mdb.lock, mdb.unlock)
 
 slots = p11.C_GetSlotList()