From 38eff317c433cbc44e9747600f02c43d0edf2b52 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20N=C3=BA=C3=B1ez?= <david@nucypher.com>
Date: Mon, 4 Jun 2018 12:26:04 +0200
Subject: [PATCH] Fixes deserialization problems for non-default curves

* Passes the curve to BystringSplitter
---
 umbral/fragments.py | 34 ++++++++++++++++++----------------
 umbral/pre.py       | 25 ++++++++++++++-----------
 2 files changed, 32 insertions(+), 27 deletions(-)

diff --git a/umbral/fragments.py b/umbral/fragments.py
index 9034e21..262e2a9 100644
--- a/umbral/fragments.py
+++ b/umbral/fragments.py
@@ -42,14 +42,15 @@ class KFrag(object):
 
         bn_size = CurveBN.expected_bytes_length(curve)
         point_size = Point.expected_bytes_length(curve)
+        arguments = {'curve': curve}
 
         splitter = BytestringSplitter(
             bn_size,             # id
-            (CurveBN, bn_size),  # bn_key
-            (Point, point_size), # point_noninteractive
-            (Point, point_size), # point_commitment
-            (Point, point_size), # point_xcoord
-            (Signature, Signature.expected_bytes_length(curve))
+            (CurveBN, bn_size, arguments),  # bn_key
+            (Point, point_size, arguments), # point_noninteractive
+            (Point, point_size, arguments), # point_commitment
+            (Point, point_size, arguments), # point_xcoord
+            (Signature, Signature.expected_bytes_length(curve), arguments)
         )
         components = splitter(data)
 
@@ -109,14 +110,14 @@ class CorrectnessProof(object):
         curve = curve if curve is not None else default_curve()
         bn_size = CurveBN.expected_bytes_length(curve)
         point_size = Point.expected_bytes_length(curve)
-
+        arguments = {'curve': curve}
         splitter = BytestringSplitter(
-            (Point, point_size), # point_e2
-            (Point, point_size), # point_v2
-            (Point, point_size), # point_kfrag_commitment
-            (Point, point_size), # point_kfrag_pok
-            (CurveBN, bn_size),  # bn_sig
-            (Signature), # kfrag_signature
+            (Point, point_size, arguments), # point_e2
+            (Point, point_size, arguments), # point_v2
+            (Point, point_size, arguments), # point_kfrag_commitment
+            (Point, point_size, arguments), # point_kfrag_pok
+            (CurveBN, bn_size, arguments),  # bn_sig
+            (Signature, Signature.expected_bytes_length(curve), arguments), # kfrag_signature
         )
         components = splitter(data, return_remainder=True)
         metadata = components.pop(-1) or None
@@ -184,13 +185,14 @@ class CapsuleFrag(object):
 
         bn_size = CurveBN.expected_bytes_length(curve)
         point_size = Point.expected_bytes_length(curve)
+        arguments = {'curve': curve}
 
         splitter = BytestringSplitter(
-            (Point, point_size), # point_e1
-            (Point, point_size), # point_v1
+            (Point, point_size, arguments), # point_e1
+            (Point, point_size, arguments), # point_v1
             bn_size,             # kfrag_id
-            (Point, point_size), # point_noninteractive
-            (Point, point_size)  # point_xcoord
+            (Point, point_size, arguments), # point_noninteractive
+            (Point, point_size, arguments)  # point_xcoord
         )
         components = splitter(data, return_remainder=True)
 
diff --git a/umbral/pre.py b/umbral/pre.py
index 00f520b..dc316a3 100644
--- a/umbral/pre.py
+++ b/umbral/pre.py
@@ -99,20 +99,23 @@ class Capsule(object):
         point_size = Point.expected_bytes_length(curve)
 
         capsule_bytes_length = len(capsule_bytes)
-        if capsule_bytes_length == cls.expected_bytes_length(curve, activated=True):
+        expected_len_original = cls.expected_bytes_length(curve, activated=False)
+        expected_len_activated = cls.expected_bytes_length(curve, activated=True)
+        arguments = {'curve': curve}
+        if capsule_bytes_length == expected_len_original:
             splitter = BytestringSplitter(
-                (Point, point_size),  # point_e
-                (Point, point_size),  # point_v
-                (CurveBN, bn_size),  # bn_sig
-                (Point, point_size),  # point_e_prime
-                (Point, point_size),  # point_v_prime
-                (Point, point_size)  # point_noninteractive
+                (Point, point_size, arguments),  # point_e
+                (Point, point_size, arguments),  # point_v
+                (CurveBN, bn_size, arguments)  # bn_sig
             )
-        elif capsule_bytes_length == cls.expected_bytes_length(curve, activated=False):
+        elif capsule_bytes_length == expected_len_activated:
             splitter = BytestringSplitter(
-                (Point, point_size),  # point_e
-                (Point, point_size),  # point_v
-                (CurveBN, bn_size)  # bn_sig
+                (Point, point_size, arguments),  # point_e
+                (Point, point_size, arguments),  # point_v
+                (CurveBN, bn_size, arguments),  # bn_sig
+                (Point, point_size, arguments),  # point_e_prime
+                (Point, point_size, arguments),  # point_v_prime
+                (Point, point_size, arguments)  # point_noninteractive
             )
         else:
             raise ValueError("Byte string does not have a valid length for a Capsule")