From 49569ac244c6103e190f13d626d9ef2069f8f2d2 Mon Sep 17 00:00:00 2001
From: Ethan Frey <ethanfrey@users.noreply.github.com>
Date: Sun, 19 Mar 2017 17:56:33 +0100
Subject: [PATCH] Wrap constructor to create xxxS structs and avoid recursion

---
 priv_key.go       |  7 +++++++
 pub_key.go        |  7 +++++++
 signature.go      |  7 +++++++
 signature_test.go | 34 ++++++++++++++++++++++++++++++++++
 4 files changed, 55 insertions(+)

diff --git a/priv_key.go b/priv_key.go
index f2d108ee5..7c2543abe 100644
--- a/priv_key.go
+++ b/priv_key.go
@@ -41,6 +41,13 @@ type PrivKeyS struct {
 	PrivKey
 }
 
+func WrapPrivKey(pk PrivKey) PrivKeyS {
+	for ppk, ok := pk.(PrivKeyS); ok; ppk, ok = pk.(PrivKeyS) {
+		pk = ppk.PrivKey
+	}
+	return PrivKeyS{pk}
+}
+
 func (p PrivKeyS) MarshalJSON() ([]byte, error) {
 	return privKeyMapper.ToJSON(p.PrivKey)
 }
diff --git a/pub_key.go b/pub_key.go
index 45a4db92c..ad7a54005 100644
--- a/pub_key.go
+++ b/pub_key.go
@@ -36,6 +36,13 @@ type PubKeyS struct {
 	PubKey
 }
 
+func WrapPubKey(pk PubKey) PubKeyS {
+	for ppk, ok := pk.(PubKeyS); ok; ppk, ok = pk.(PubKeyS) {
+		pk = ppk.PubKey
+	}
+	return PubKeyS{pk}
+}
+
 func (p PubKeyS) MarshalJSON() ([]byte, error) {
 	return pubKeyMapper.ToJSON(p.PubKey)
 }
diff --git a/signature.go b/signature.go
index b5fc36bf3..7a03dc6d3 100644
--- a/signature.go
+++ b/signature.go
@@ -31,6 +31,13 @@ type SignatureS struct {
 	Signature
 }
 
+func WrapSignature(sig Signature) SignatureS {
+	for ssig, ok := sig.(SignatureS); ok; ssig, ok = sig.(SignatureS) {
+		sig = ssig.Signature
+	}
+	return SignatureS{sig}
+}
+
 func (p SignatureS) MarshalJSON() ([]byte, error) {
 	return sigMapper.ToJSON(p.Signature)
 }
diff --git a/signature_test.go b/signature_test.go
index f6c8c37be..f04262b77 100644
--- a/signature_test.go
+++ b/signature_test.go
@@ -107,3 +107,37 @@ func TestSignatureEncodings(t *testing.T) {
 		assert.True(t, strings.HasPrefix(text, tc.sigName))
 	}
 }
+
+func TestWrapping(t *testing.T) {
+	assert := assert.New(t)
+
+	// construct some basic constructs
+	msg := CRandBytes(128)
+	priv := GenPrivKeyEd25519()
+	pub := priv.PubKey()
+	sig := priv.Sign(msg)
+
+	// do some wrapping
+	pubs := []PubKeyS{
+		WrapPubKey(nil),
+		WrapPubKey(pub),
+		WrapPubKey(WrapPubKey(WrapPubKey(WrapPubKey(pub)))),
+		WrapPubKey(PubKeyS{PubKeyS{PubKeyS{pub}}}),
+	}
+	for _, p := range pubs {
+		_, ok := p.PubKey.(PubKeyS)
+		assert.False(ok)
+	}
+
+	sigs := []SignatureS{
+		WrapSignature(nil),
+		WrapSignature(sig),
+		WrapSignature(WrapSignature(WrapSignature(WrapSignature(sig)))),
+		WrapSignature(SignatureS{SignatureS{SignatureS{sig}}}),
+	}
+	for _, s := range sigs {
+		_, ok := s.Signature.(SignatureS)
+		assert.False(ok)
+	}
+
+}