fixup! multi: Support custom size onion packets

Author: gijswijs

cmd/main.go

@@ -70,11 +70,12 @@ func main() {
 						"data.",
 					Value: defaultHopDataPath,
 				},
-				cli.BoolFlag{
-					Name: "onion-message",
-					Usage: "Create an onion message " +
-						"packet rather than a " +
-						"payment onion.",
+				cli.IntFlag{
+					Name: "payload-size",
+					Usage: "The size for a payload for a " +
+						"single hop. Defaults to the " +
+						"max routing payload size",
+					Value: sphinx.MaxRoutingPayloadSize,
 				},
 			},
 		},
@@ -209,18 +210,10 @@ func generate(ctx *cli.Context) error {
 		return fmt.Errorf("could not peel onion spec: %v", err)
 	}
 
-	payloadSizes := []int{
-		sphinx.MaxRoutingPayloadSize,
-	}
-	if ctx.Bool("onion-message") {
-		payloadSizes = append(
-			payloadSizes,
-			sphinx.MaxOnionMessagePayloadSize,
-		)
-	}
+	payloadSize := ctx.Int("payload-size")
 	msg, err := sphinx.NewOnionPacket(
 		path, sessionKey, assocData, sphinx.DeterministicPacketFiller,
-		payloadSizes...,
+		sphinx.WithMaxPayloadSize(payloadSize),
 	)
 	if err != nil {
 		return fmt.Errorf("error creating message: %v", err)

error.go

@@ -1,6 +1,9 @@
 package sphinx
 
-import "fmt"
+import (
+	"errors"
+	"fmt"
+)
 
 var (
 	// ErrReplayedPacket is an error returned when a packet is rejected
@@ -21,7 +24,36 @@ var (
 	ErrInvalidOnionKey = fmt.Errorf("invalid onion key: pubkey isn't on " +
 		"secp256k1 curve")
 
-	// ErrLogEntryNotFound is an error returned when a packet lookup in a replay
-	// log fails because it is missing.
-	ErrLogEntryNotFound = fmt.Errorf("sphinx packet is not in log")
+	// ErrLogEntryNotFound is an error returned when a packet lookup in a
+	// replay log fails because it is missing.
+	ErrLogEntryNotFound = errors.New("sphinx packet is not in log")
+
+	// ErrPayloadSizeExceeded is returned when the payload size exceeds the
+	// configured payload size of the onion packet.
+	ErrPayloadSizeExceeded = errors.New("max payload size exceeded")
+
+	// ErrSharedSecretDerivation is returned when we fail to derive the
+	// shared secret for a hop.
+	ErrSharedSecretDerivation = errors.New("error generating shared secret")
+
+	// ErrMissingHMAC is returned when the onion packet is too small to
+	// contain a valid HMAC.
+	ErrMissingHMAC = errors.New("onion packet is too small, missing HMAC")
+
+	// ErrNegativeRoutingInfoSize is returned when a negative routing info
+	// size is specified in the Sphinx configuration.
+	ErrNegativeRoutingInfoSize = errors.New("payload size must be " +
+		"non-negative")
+
+	// ErrNegativePayloadSize is returned when a negative payload size is
+	// specified in the Sphinx configuration.
+	ErrNegativePayloadSize = errors.New("payload size must be " +
+		"non-negative")
+
+	// ErrZeroHops is returned when attempting to create a route with zero
+	// hops.
+	ErrZeroHops = errors.New("route of length zero passed in")
+
+	// ErrIOReadFull is returned when an io read full operation fails.
+	ErrIOReadFull = errors.New("io read full error")
 )

packetfiller.go

@@ -21,7 +21,8 @@ func RandPacketFiller(_ *btcec.PrivateKey, mixHeader []byte) error {
 	// after the hop payload for the final node. This mitigates a privacy
 	// leak that may reveal a lower bound on the true path length to the
 	// receiver.
-	if _, err := rand.Read(mixHeader); err != nil {
+	_, err := rand.Read(mixHeader)
+	if err != nil {
 		return err
 	}
 
@@ -55,6 +56,7 @@ func DeterministicPacketFiller(sessionKey *btcec.PrivateKey,
 	if err != nil {
 		return err
 	}
+
 	padCipher.XORKeyStream(mixHeader, mixHeader)
 
 	return nil

sphinx.go

@@ -6,7 +6,7 @@ import (
 	"crypto/sha256"
 	"fmt"
 	"io"
-	"sort"
+	"slices"
 	"sync"
 
 	"github.com/btcsuite/btcd/btcec/v2"
@@ -60,10 +60,11 @@ const (
 
 	// baseVersion represent the current supported version of onion packet.
 	baseVersion = 0
-)
 
-var (
-	ErrPayloadSizeExceeded = fmt.Errorf("max custom payload size exceeded")
+	// streamBytesMultiplier is the multiplier used to calculate the number
+	// of bytes that needs to be produced by our CSPRNG for the key stream
+	// implementing our stream cipher.
+	streamBytesMultiplier = 2
 )
 
 // OnionPacket is the onion wrapped hop-to-hop routing information necessary to
@@ -180,66 +181,59 @@ func generateSharedSecrets(paymentPath []*btcec.PublicKey,
 	return hopSharedSecrets, lastEphemeralPubKey, nil
 }
 
+type newOnionPacketCfg struct {
+	payloadSize int
+}
+
+// NewOnionPacketOpt defines the signature of a function option that can be used
+// with NewOnionPacket.
+type NewOnionPacketOpt func(cfg *newOnionPacketCfg)
+
+// WithMaxPayloadSize is a function option that can be used to set the maximum
+// payload size.
+func WithMaxPayloadSize(size int) NewOnionPacketOpt {
+	return func(cfg *newOnionPacketCfg) {
+		cfg.payloadSize = size
+	}
+}
+
 // NewOnionPacket creates a new onion packet which is capable of obliviously
 // routing a message through the mix-net path outline by 'paymentPath'. The
 // total size of the onion 'clicks' to the first value in payloadSizes that is
 // bigger than the total payload size of the path. If no size is given, it
 // defaults to the maximum routing payload size.
 func NewOnionPacket(paymentPath *PaymentPath, sessionKey *btcec.PrivateKey,
 	assocData []byte, pktFiller PacketFiller,
-	payloadSizes ...int) (*OnionPacket, error) {
+	opts ...NewOnionPacketOpt) (*OnionPacket, error) {
 
-	// If we don't actually have a partially populated route, then we'll
-	// exit early.
-	numHops := paymentPath.TrueRouteLength()
-	if numHops == 0 {
-		return nil, fmt.Errorf("route of length zero passed in")
+	cfg := &newOnionPacketCfg{}
+	for _, o := range opts {
+		o(cfg)
 	}
 
-	totalPayloadSize := paymentPath.TotalPayloadSize()
+	if cfg.payloadSize < 0 {
+		return nil, ErrNegativePayloadSize
+	}
 
 	// We default to the maximum routing payload size if the caller didn't
 	// provide any payload sizes.
-	if len(payloadSizes) == 0 {
-		payloadSizes = []int{MaxRoutingPayloadSize}
+	if cfg.payloadSize == 0 {
+		cfg.payloadSize = MaxRoutingPayloadSize
 	}
 
-	sort.Ints(payloadSizes)
-
-	// We'll now select the smallest payload size that is large enough to
-	// fit the entire onion payload. If no such size exists, then we'll
-	// return an error
-	var payloadSize int
-	found := false
-	for _, size := range payloadSizes {
-		if size >= totalPayloadSize {
-			payloadSize = size
-			found = true
-			break
-		}
+	// If we don't actually have a partially populated route, then we'll
+	// exit early.
+	numHops := paymentPath.TrueRouteLength()
+	if numHops == 0 {
+		return nil, fmt.Errorf("route of length zero passed in")
 	}
 
-	// Return an error if we couldn't find a suitable payload size.
-	if !found {
-		return nil, ErrPayloadSizeExceeded
-	}
+	totalPayloadSize := paymentPath.TotalPayloadSize()
 
-	// If the payload size is not equal to MaxRoutingPayloadSize, then we
-	// check if any of the hops have a legacy payload. If so, we return an
-	// error as legacy payloads are not supported for those payload sizes.
-	if payloadSize != MaxRoutingPayloadSize {
-		for i := 0; i < numHops; i++ {
-			hopPayload := (*paymentPath)[i].HopPayload
-			isLegacy := hopPayload.Type == PayloadLegacy
-
-			// For any onion size other than MaxRoutingPayloadSize,
-			// we only expect TLV payloads.
-			if isLegacy {
-				return nil, fmt.Errorf("hop %d has legacy "+
-					"payload, but this payload size "+
-					"requires TLV,", i)
-			}
-		}
+	// Return an error if the actual payload size exceeds the configured
+	// payload size.
+	if totalPayloadSize > cfg.payloadSize {
+		return nil, ErrPayloadSizeExceeded
 	}
 
 	// We'll force the caller to provide a packet filler, as otherwise we
@@ -253,25 +247,28 @@ func NewOnionPacket(paymentPath *PaymentPath, sessionKey *btcec.PrivateKey,
 		paymentPath.NodeKeys(), sessionKey,
 	)
 	if err != nil {
-		return nil, fmt.Errorf("error generating shared secret: %v",
-			err)
+		return nil, fmt.Errorf("%w: %w", ErrSharedSecretDerivation, err)
 	}
 
 	// Generate the padding, called "filler strings" in the paper.
-	filler := generateHeaderPadding(
-		"rho", paymentPath, hopSharedSecrets, payloadSize,
+	filler, err := generateHeaderPadding(
+		"rho", paymentPath, hopSharedSecrets, cfg.payloadSize,
 	)
+	if err != nil {
+		return nil, err
+	}
 
 	// Allocate zero'd out byte slices to store the final mix header packet
 	// and the hmac for each hop.
 	var (
-		mixHeader     = make([]byte, payloadSize)
+		mixHeader     = make([]byte, cfg.payloadSize)
 		nextHmac      [HMACSize]byte
 		hopPayloadBuf bytes.Buffer
 	)
 
 	// Fill the packet using the caller specified methodology.
-	if err := pktFiller(sessionKey, mixHeader); err != nil {
+	err = pktFiller(sessionKey, mixHeader)
+	if err != nil {
 		return nil, err
 	}
 
@@ -292,7 +289,9 @@ func NewOnionPacket(paymentPath *PaymentPath, sessionKey *btcec.PrivateKey,
 		// Next, using the key dedicated for our stream cipher, we'll
 		// generate enough bytes to obfuscate this layer of the onion
 		// packet.
-		streamBytes := generateCipherStream(rhoKey, uint(payloadSize))
+		streamBytes := generateCipherStream(
+			rhoKey, uint(cfg.payloadSize),
+		)
 		payload := paymentPath[i].HopPayload
 
 		// Before we assemble the packet, we'll shift the current
@@ -323,7 +322,7 @@ func NewOnionPacket(paymentPath *PaymentPath, sessionKey *btcec.PrivateKey,
 		// calculating the MAC, we'll also include the optional
 		// associated data which can allow higher level applications to
 		// prevent replay attacks.
-		packet := append(mixHeader, assocData...)
+		packet := slices.Concat(mixHeader, assocData)
 		nextHmac = calcMac(muKey, packet)
 
 		hopPayloadBuf.Reset()
@@ -361,7 +360,7 @@ func rightShift(slice []byte, num int) {
 // last hop.  Using this methodology, the size of the field stays constant at
 // each hop.
 func generateHeaderPadding(key string, path *PaymentPath,
-	sharedSecrets []Hash256, routingInfoLen int) []byte {
+	sharedSecrets []Hash256, routingInfoLen int) ([]byte, error) {
 
 	numHops := path.TrueRouteLength()
 
@@ -383,14 +382,18 @@ func generateHeaderPadding(key string, path *PaymentPath,
 		fillerEnd := routingInfoLen + path[i].HopPayload.NumBytes()
 
 		streamKey := generateKey(key, &sharedSecrets[i])
-		streamBytes := generateCipherStream(
-			streamKey, numStreamBytes(routingInfoLen),
-		)
+
+		streamBytesLen, err := numStreamBytes(routingInfoLen)
+		if err != nil {
+			return nil, err
+		}
+
+		streamBytes := generateCipherStream(streamKey, streamBytesLen)
 
 		xor(filler, filler, streamBytes[fillerStart:fillerEnd])
 	}
 
-	return filler
+	return filler, nil
 }
 
 // Encode serializes the raw bytes of the onion packet into the passed
@@ -407,7 +410,8 @@ func (f *OnionPacket) Encode(w io.Writer) error {
 		return err
 	}
 
-	if _, err := w.Write(f.RoutingInfo); err != nil {
+	_, err = w.Write(f.RoutingInfo)
+	if err != nil {
 		return err
 	}
 
@@ -455,7 +459,7 @@ func (f *OnionPacket) Decode(r io.Reader) error {
 
 	// The packet must have at least enough bytes for the HMAC.
 	if len(routingInfoAndMAC) < HMACSize {
-		return fmt.Errorf("onion packet is too small, missing HMAC")
+		return ErrMissingHMAC
 	}
 
 	// With the remainder of the packet read, we can now properly slice the
@@ -701,7 +705,7 @@ func unwrapPacket(onionPkt *OnionPacket, sharedSecret *Hash256,
 	// Using the derived shared secret, ensure the integrity of the routing
 	// information by checking the attached MAC without leaking timing
 	// information.
-	message := append(routeInfo, assocData...)
+	message := slices.Concat(routeInfo, assocData)
 	calculatedMac := calcMac(generateKey("mu", sharedSecret), message)
 	if !hmac.Equal(headerMac[:], calculatedMac[:]) {
 		return nil, nil, ErrInvalidOnionHMAC
@@ -710,14 +714,17 @@ func unwrapPacket(onionPkt *OnionPacket, sharedSecret *Hash256,
 	// Attach the padding zeroes in order to properly strip an encryption
 	// layer off the routing info revealing the routing information for the
 	// next hop.
+	streamBytesLen, err := numStreamBytes(routingInfoLen)
+	if err != nil {
+		return nil, nil, err
+	}
 	streamBytes := generateCipherStream(
-		generateKey("rho", sharedSecret),
-		numStreamBytes(routingInfoLen),
+		generateKey("rho", sharedSecret), streamBytesLen,
 	)
 	zeroBytes := bytes.Repeat([]byte{0}, routingInfoLen)
-	headerWithPadding := append(routeInfo, zeroBytes...)
+	headerWithPadding := slices.Concat(routeInfo, zeroBytes)
 
-	hopInfo := make([]byte, numStreamBytes(routingInfoLen))
+	hopInfo := make([]byte, streamBytesLen)
 	xor(hopInfo, headerWithPadding, streamBytes)
 
 	// Randomize the DH group element for the next hop using the
@@ -736,7 +743,8 @@ func unwrapPacket(onionPkt *OnionPacket, sharedSecret *Hash256,
 		// payload is treated as such.
 		hopPayload.Type = PayloadTLV
 	}
-	err := hopPayload.Decode(bytes.NewReader(hopInfo))
+
+	err = hopPayload.Decode(bytes.NewReader(hopInfo))
 	if err != nil {
 		return nil, nil, err
 	}
@@ -908,11 +916,15 @@ func (t *Tx) Commit() ([]ProcessedPacket, *ReplaySet, error) {
 	return t.packets, rs, err
 }
 
-// numStreamBytes is the number of bytes that needs to be produced by our CSPRG
+// numStreamBytes is the number of bytes that needs to be produced by our CSPRNG
 // for the key stream implementing our stream cipher to encrypt/decrypt the mix
 // header. The routingInfoSize bytes at the end are used to encrypt/decrypt the
 // fillers when processing the packet of generating the HMACs when creating the
 // packet.
-func numStreamBytes(routingInfoSize int) uint {
-	return uint(routingInfoSize * 2)
+func numStreamBytes(routingInfoSize int) (uint, error) {
+	if routingInfoSize < 0 {
+		return 0, ErrNegativeRoutingInfoSize
+	}
+
+	return uint(routingInfoSize) * streamBytesMultiplier, nil
 }