From fcf08382f7f66c9cc628c50efb2994f1d333f2c0 Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@mit.edu>
Date: Sun, 10 Dec 2017 15:37:21 -0800
Subject: [PATCH] htlcswitch/failure: add Encode/Decode to ErrorEncrypter

---
 htlcswitch/failure.go | 70 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 70 insertions(+)

diff --git a/htlcswitch/failure.go b/htlcswitch/failure.go
index 0acf7b0a..f2c5ae61 100644
--- a/htlcswitch/failure.go
+++ b/htlcswitch/failure.go
@@ -3,6 +3,7 @@ package htlcswitch
 import (
 	"bytes"
 	"fmt"
+	"io"
 
 	"github.com/lightningnetwork/lightning-onion"
 	"github.com/lightningnetwork/lnd/lnwire"
@@ -45,6 +46,37 @@ type ErrorDecrypter interface {
 	DecryptError(lnwire.OpaqueReason) (*ForwardingError, error)
 }
 
+// EncrypterType establishes an enum used in serialization to indicate how to
+// decode a concrete instance of the ErrorEncrypter interface.
+type EncrypterType byte
+
+const (
+	// EncrypterTypeNone signals that no error encyrpter is present, this
+	// can happen if the htlc is originates in the switch.
+	EncrypterTypeNone EncrypterType = 0
+
+	// EncrypterTypeSphinx is used to identify a sphinx onion error
+	// encrypter instance.
+	EncrypterTypeSphinx = 1
+
+	// EncrypterTypeMock is used to identify a mock obfuscator instance.
+	EncrypterTypeMock = 2
+)
+
+// UnknownEncrypterType is an error message used to signal that an unexpected
+// EncrypterType was encountered during decoding.
+type UnknownEncrypterType EncrypterType
+
+// Error returns a formatted error indicating the invalid EncrypterType.
+func (e UnknownEncrypterType) Error() string {
+	return fmt.Sprintf("unknown error encrypter type: %d", e)
+}
+
+// ErrorEncrypterExtracter defines a function signature that extracts an
+// ErrorEncrypter from an sphinx OnionPacket.
+type ErrorEncrypterExtracter func(*sphinx.OnionPacket) (ErrorEncrypter,
+	lnwire.FailCode)
+
 // ErrorEncrypter is an interface that is used to encrypt HTLC related errors
 // at the source of the error, and also at each intermediate hop all the way
 // back to the source of the payment.
@@ -59,6 +91,16 @@ type ErrorEncrypter interface {
 	// in an additional layer of onion encryption. This process repeats
 	// until the error arrives at the source of the payment.
 	IntermediateEncrypt(lnwire.OpaqueReason) lnwire.OpaqueReason
+
+	// Type returns an enum indicating the underlying concrete instance
+	// backing this interface.
+	Type() EncrypterType
+
+	// Encode serializes the encrypter to the given io.Writer.
+	Encode(io.Writer) error
+
+	// Decode deserializes the encrypter from the given io.Reader.
+	Decode(io.Reader) error
 }
 
 // SphinxErrorEncrypter is a concrete implementation of both the ErrorEncrypter
@@ -67,6 +109,16 @@ type ErrorEncrypter interface {
 // encryption and must be treated as such accordingly.
 type SphinxErrorEncrypter struct {
 	*sphinx.OnionErrorEncrypter
+
+	ogPacket *sphinx.OnionPacket
+}
+
+// NewSphinxErrorEncrypter initializes a new sphinx error encrypter as well as
+// the embedded onion error encrypter.
+func NewSphinxErrorEncrypter() *SphinxErrorEncrypter {
+	return &SphinxErrorEncrypter{
+		OnionErrorEncrypter: &sphinx.OnionErrorEncrypter{},
+	}
 }
 
 // EncryptFirstHop transforms a concrete failure message into an encrypted
@@ -97,6 +149,24 @@ func (s *SphinxErrorEncrypter) IntermediateEncrypt(reason lnwire.OpaqueReason) l
 	return s.EncryptError(false, reason)
 }
 
+// Type returns the identifier for a sphinx error encrypter.
+func (s *SphinxErrorEncrypter) Type() EncrypterType {
+	return EncrypterTypeSphinx
+}
+
+// Encode serializes the error encrypter to the provided io.Writer.
+func (s *SphinxErrorEncrypter) Encode(w io.Writer) error {
+	return s.OnionErrorEncrypter.Encode(w)
+}
+
+// Decode reconstructs the error encrypter from the provided io.Reader.
+func (s *SphinxErrorEncrypter) Decode(r io.Reader) error {
+	if s.OnionErrorEncrypter == nil {
+		s.OnionErrorEncrypter = &sphinx.OnionErrorEncrypter{}
+	}
+	return s.OnionErrorEncrypter.Decode(r)
+}
+
 // A compile time check to ensure SphinxErrorEncrypter implements the
 // ErrorEncrypter interface.
 var _ ErrorEncrypter = (*SphinxErrorEncrypter)(nil)