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)

sphinx.go

@@ -6,7 +6,6 @@ import (
 	"crypto/sha256"
 	"fmt"
 	"io"
-	"sort"
 	"sync"
 
 	"github.com/btcsuite/btcd/btcec/v2"
@@ -180,14 +179,37 @@ func generateSharedSecrets(paymentPath []*btcec.PublicKey,
 	return hopSharedSecrets, lastEphemeralPubKey, nil
 }
 
+type newOnionPacketCfg struct {
+	payloadSize int
+}
+
+type NewOnionPacketOpt func(cfg *newOnionPacketCfg)
+
+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) {
+
+	cfg := &newOnionPacketCfg{}
+	for _, o := range opts {
+		o(cfg)
+	}
+
+	// We default to the maximum routing payload size if the caller didn't
+	// provide any payload sizes.
+	if cfg.payloadSize == 0 {
+		cfg.payloadSize = MaxRoutingPayloadSize
+	}
 
 	// If we don't actually have a partially populated route, then we'll
 	// exit early.
@@ -198,29 +220,9 @@ func NewOnionPacket(paymentPath *PaymentPath, sessionKey *btcec.PrivateKey,
 
 	totalPayloadSize := paymentPath.TotalPayloadSize()
 
-	// We default to the maximum routing payload size if the caller didn't
-	// provide any payload sizes.
-	if len(payloadSizes) == 0 {
-		payloadSizes = []int{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
-		}
-	}
-
-	// Return an error if we couldn't find a suitable payload size.
-	if !found {
+	// Return an error if the actual payload size exceeds the configured
+	// payload size.
+	if totalPayloadSize > cfg.payloadSize {
 		return nil, ErrPayloadSizeExceeded
 	}
 
@@ -241,13 +243,13 @@ func NewOnionPacket(paymentPath *PaymentPath, sessionKey *btcec.PrivateKey,
 
 	// Generate the padding, called "filler strings" in the paper.
 	filler := generateHeaderPadding(
-		"rho", paymentPath, hopSharedSecrets, payloadSize,
+		"rho", paymentPath, hopSharedSecrets, cfg.payloadSize,
 	)
 
 	// 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
 	)
@@ -274,7 +276,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

sphinx_test.go

@@ -275,11 +275,20 @@ func newOnionMessageRoute(numHops int) (*OnionPacket, *PaymentPath, []*Router,
 		}
 	}
 
+	// According to BOLT04 we SHOULD set onion_message_packet len to 1366 or
+	// 32834. This means a payload size of 1300 (MaxRoutingPayloadSize) or
+	// 32768 (MaxOnionMessagePayloadSize) bytes. By checking the total
+	// payload size of the route we can determine which one to use.
+	payloadSize := MaxRoutingPayloadSize
+	if route.TotalPayloadSize() > MaxRoutingPayloadSize {
+		payloadSize = MaxOnionMessagePayloadSize
+	}
+
 	// Generate the onion packet.
 	sessionKey, _ := btcec.NewPrivateKey()
 	onionPacket, err := NewOnionPacket(
 		&route, sessionKey, nil, DeterministicPacketFiller,
-		MaxRoutingPayloadSize, MaxOnionMessagePayloadSize,
+		WithMaxPayloadSize(payloadSize),
 	)
 	if err != nil {
 		return nil, nil, nil, fmt.Errorf("unable to create onion "+
@@ -898,7 +907,7 @@ func TestCustomPayloadSize(t *testing.T) {
 	// Now, create the onion packet.
 	onionPacket, err := NewOnionPacket(
 		path, sessionKey, nil, DeterministicPacketFiller,
-		customPayloadSize,
+		WithMaxPayloadSize(customPayloadSize),
 	)
 	require.NoError(t, err)