Merge pull request #774 from adotkhan/push-bin-packing

Add bin packing to MakePushPayloads

psiphon/common/push/converter/main.go

@@ -135,20 +135,18 @@ func convert(
 				})
 		}
 
-		payloads, err := push.MakePushPayloads(
-			obfuscationKey,
-			minPadding,
-			maxPadding,
-			signaturePublicKey,
-			signaturePrivateKey,
-			ttl,
-			[][]*push.PrioritizedServerEntry{
-				prioritizedServerEntries})
+		maker, err := push.NewPushPayloadMaker(
+			obfuscationKey, signaturePublicKey, signaturePrivateKey)
+		if err != nil {
+			return errors.Trace(err)
+		}
+		result, err := maker.MakePushPayloads(
+			minPadding, maxPadding, ttl, prioritizedServerEntries, 0)
 		if err != nil {
 			return errors.Trace(err)
 		}
 
-		os.Stdout.Write(payloads[0])
+		os.Stdout.Write(result.Payloads[0])
 		return nil
 	}
 

psiphon/common/push/push.go

@@ -37,6 +37,7 @@ import (
 	"crypto/rand"
 	"crypto/sha256"
 	"encoding/base64"
+	"sort"
 	"time"
 
 	"github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/errors"
@@ -48,6 +49,8 @@ import (
 const (
 	obfuscationKeySize           = 32
 	signaturePublicKeyDigestSize = 8
+	maxPaddingLimit              = 65535
+	signatureSize                = signaturePublicKeyDigestSize + ed25519.SignatureSize
 )
 
 // Payload is a push payload, consisting of a list of server entries. To
@@ -211,15 +214,51 @@ func ImportPushPayload(
 	return imported, nil
 }
 
-// MakePushPayloads generates batches of push payloads.
-func MakePushPayloads(
+// MakePushPayloadsResult is the output from MakePushPayloads.
+type MakePushPayloadsResult struct {
+	// Payloads contains generated obfuscated push payloads.
+	Payloads [][]byte
+	// PayloadEntryCounts contains the number of entries in each payload, aligned
+	// by index with Payloads.
+	PayloadEntryCounts []int
+	// SkippedIndexes contains original input indexes for entries that could not
+	// fit into a payload when max payload size is enforced.
+	SkippedIndexes []int
+}
+
+type payloadBuffers struct {
+	nonce       []byte
+	signature   []byte
+	obfuscation []byte
+	padding     []byte
+}
+
+type sortablePrioritizedServerEntry struct {
+	entry         *PrioritizedServerEntry
+	originalIndex int
+	encodedSize   int
+}
+
+// PushPayloadMaker caches expensive initialization (base64 decoding, AES-GCM
+// cipher creation, SHA256 hashing) so that multiple MakePayloads calls can
+// reuse the same state.
+//
+// PushPayloadMaker is safe for concurrent use. Each MakePayloads call
+// allocates its own mutable buffers via a fresh payloadBuffers.
+type PushPayloadMaker struct {
+	aead        cipher.AEAD
+	privateKey  ed25519.PrivateKey
+	publicKeyID []byte
+}
+
+// NewPushPayloadMaker creates a PushPayloadMaker by performing the expensive
+// one-time initialization: base64 decoding all keys, validating sizes, and
+// creating the AES-GCM cipher.
+func NewPushPayloadMaker(
 	payloadObfuscationKey string,
-	minPadding int,
-	maxPadding int,
 	payloadSignaturePublicKey string,
 	payloadSignaturePrivateKey string,
-	TTL time.Duration,
-	prioritizedServerEntries [][]*PrioritizedServerEntry) ([][]byte, error) {
+) (*PushPayloadMaker, error) {
 
 	obfuscationKey, err := base64.StdEncoding.DecodeString(
 		payloadObfuscationKey)
@@ -248,13 +287,6 @@ func MakePushPayloads(
 		return nil, errors.Trace(err)
 	}
 
-	expires := time.Now().Add(TTL).UTC()
-
-	maxPaddingLimit := 65535
-	if minPadding > maxPadding || maxPadding > maxPaddingLimit {
-		return nil, errors.TraceNew("invalid min/max padding")
-	}
-
 	blockCipher, err := aes.NewCipher(obfuscationKey)
 	if err != nil {
 		return nil, errors.Trace(err)
@@ -266,67 +298,342 @@ func MakePushPayloads(
 	}
 
 	publicKeyDigest := sha256.Sum256(publicKey)
-	publicKeyID := publicKeyDigest[:signaturePublicKeyDigestSize]
 
-	// Reuse buffers to reduce some allocations.
-	var signatureBuffer []byte
-	var obfuscationBuffer []byte
-	nonceBuffer := make([]byte, aead.NonceSize())
-	var paddingBuffer []byte
+	return &PushPayloadMaker{
+		aead:        aead,
+		privateKey:  privateKey,
+		publicKeyID: publicKeyDigest[:signaturePublicKeyDigestSize],
+	}, nil
+}
+
+// MakePushPayloads generates obfuscated push payloads from prioritized server
+// entries, reusing the cached key material and cipher from the maker.
+//
+// When maxPayloadSizeBytes <= 0, all entries are encoded into a single payload.
+//
+// When maxPayloadSizeBytes > 0, entries are packed into multiple payloads using
+// an RF(2) (random-fit with 2 candidates) strategy. Entries that cannot
+// fit by themselves under maxPayloadSizeBytes are skipped and reported in the
+// returned result metadata.
+func (m *PushPayloadMaker) MakePushPayloads(
+	minPadding int,
+	maxPadding int,
+	TTL time.Duration,
+	prioritizedServerEntries []*PrioritizedServerEntry,
+	maxPayloadSizeBytes int) (MakePushPayloadsResult, error) {
+
+	result := MakePushPayloadsResult{}
 
-	obfuscatedPayloads := [][]byte{}
+	if len(prioritizedServerEntries) == 0 {
+		return result, nil
+	}
+
+	if minPadding > maxPadding || maxPadding > maxPaddingLimit {
+		return result, errors.TraceNew("invalid min/max padding")
+	}
 
-	for _, p := range prioritizedServerEntries {
+	bufs := &payloadBuffers{
+		nonce: make([]byte, m.aead.NonceSize()),
+	}
+
+	expires := time.Now().Add(TTL).UTC()
 
-		payload := Payload{
-			Expires:                  expires,
-			PrioritizedServerEntries: p,
+	// maxPayloadSizeBytes <= 0 means no payload size cap is enforced.
+	if maxPayloadSizeBytes <= 0 {
+		paddingSize := prng.Range(minPadding, maxPadding)
+		payload, err := m.buildObfuscatedPayload(
+			bufs, prioritizedServerEntries, expires, paddingSize)
+		if err != nil {
+			return result, errors.Trace(err)
 		}
+		result.Payloads = append(result.Payloads, payload)
+		result.PayloadEntryCounts = append(
+			result.PayloadEntryCounts, len(prioritizedServerEntries))
+		return result, nil
+	}
+
+	// Pre-compute the CBOR-encoded size of the expires timestamp.
+	expiresEncoded, err := protocol.CBOREncoding.Marshal(expires)
+	if err != nil {
+		return result, errors.Trace(err)
+	}
+	expiresEncodedSize := len(expiresEncoded)
 
-		cborPayload, err := protocol.CBOREncoding.Marshal(&payload)
+	// Compute encoded sizes for each PrioritizedServerEntry.
+	serverEntries := make(
+		[]sortablePrioritizedServerEntry, 0, len(prioritizedServerEntries))
+	for i, entry := range prioritizedServerEntries {
+		encodedEntry, err := protocol.CBOREncoding.Marshal(entry)
 		if err != nil {
-			return nil, errors.Trace(err)
+			return result, errors.Trace(err)
 		}
 
-		signature := ed25519.Sign(privateKey, cborPayload)
+		serverEntries = append(serverEntries, sortablePrioritizedServerEntry{
+			entry:         entry,
+			originalIndex: i,
+			encodedSize:   len(encodedEntry),
+		})
+	}
 
-		signatureBuffer = signatureBuffer[:0]
-		signatureBuffer = append(signatureBuffer, publicKeyID...)
-		signatureBuffer = append(signatureBuffer, signature...)
+	// Sort server entries by decreasing size, this significantly
+	// increases packing quality but doesn't bias the bins themselves.
+	sort.Slice(serverEntries, func(i, j int) bool {
+		if serverEntries[i].encodedSize == serverEntries[j].encodedSize {
+			return serverEntries[i].originalIndex < serverEntries[j].originalIndex
+		}
+		return serverEntries[i].encodedSize > serverEntries[j].encodedSize
+	})
+
+	// Worst-case each PrioritizedServerEntry gets its own bin.
+	type payloadBin struct {
+		serverEntries []*PrioritizedServerEntry
+		paddingSize   int
+		// sumServerEntrySize is the total encoded size of all server
+		// entries in this bin, used to compute the obfuscated payload size.
+		sumServerEntrySize int
+	}
+	bins := make([]payloadBin, 0, len(serverEntries))
+
+	binOrder := make([]int, 0, len(serverEntries))
+
+	type candidate struct {
+		binIndex int
+		size     int
+	}
+
+	for _, sortedServerEntry := range serverEntries {
 
-		signedPayload := SignedPayload{
-			Signature: signatureBuffer,
-			Payload:   cborPayload,
+		// RF(2): randomly sample bins, collect the first 2 that fit,
+		// and pick the tightest (least remaining space).
+
+		// Grow and reset binOrder to [0..len(bins)).
+		binOrder = binOrder[:0]
+		for i := range bins {
+			binOrder = append(binOrder, i)
 		}
+		prng.Shuffle(len(binOrder), func(i, j int) {
+			binOrder[i], binOrder[j] = binOrder[j], binOrder[i]
+		})
+
+		var candidates [2]candidate
+		numCandidates := 0
+
+		for _, bi := range binOrder {
+			if numCandidates >= 2 {
+				break
+			}
 
-		// Padding is an optional part of the obfuscation layer.
-		if maxPadding > 0 {
-			paddingSize := prng.Range(minPadding, maxPadding)
-			if paddingBuffer == nil {
-				paddingBuffer = make([]byte, maxPaddingLimit)
+			// Arithmetically compute the size of the obfuscated payload size
+			// without the expensive marshalling and encryption.
+			size := m.computeObfuscatedPayloadSize(
+				expiresEncodedSize,
+				len(bins[bi].serverEntries)+1,
+				bins[bi].sumServerEntrySize+sortedServerEntry.encodedSize,
+				bins[bi].paddingSize)
+			if size <= maxPayloadSizeBytes {
+				candidates[numCandidates] = candidate{
+					binIndex: bi,
+					size:     size,
+				}
+				numCandidates++
 			}
-			if paddingSize > 0 {
-				signedPayload.Padding = paddingBuffer[0:paddingSize]
+		}
+
+		if numCandidates > 0 {
+			// Pick tightest fit (highest size).
+			best := 0
+			if numCandidates == 2 &&
+				candidates[1].size > candidates[0].size {
+				best = 1
 			}
+			bi := candidates[best].binIndex
+			bins[bi].serverEntries = append(bins[bi].serverEntries, sortedServerEntry.entry)
+			bins[bi].sumServerEntrySize += sortedServerEntry.encodedSize
+			continue
 		}
 
-		cborSignedPayload, err := protocol.CBOREncoding.
-			Marshal(&signedPayload)
+		// Server entry did not fit into existing bins,
+		// create a new bin with minPadding. Random padding is
+		// applied after packing to avoid wasting bin capacity.
+		paddingSize := minPadding
+		size := m.computeObfuscatedPayloadSize(
+			expiresEncodedSize, 1, sortedServerEntry.encodedSize, paddingSize)
+		if size > maxPayloadSizeBytes {
+			result.SkippedIndexes = append(
+				result.SkippedIndexes, sortedServerEntry.originalIndex)
+			continue
+		}
+
+		bins = append(bins, payloadBin{
+			serverEntries:      []*PrioritizedServerEntry{sortedServerEntry.entry},
+			paddingSize:        paddingSize,
+			sumServerEntrySize: sortedServerEntry.encodedSize,
+		})
+	}
+
+	// Apply random padding to each bin, respecting maxPayloadSizeBytes.
+	noPadding := minPadding == 0 && maxPadding == 0
+	if !noPadding {
+		for i := range bins {
+			randomPadding := prng.Range(minPadding, maxPadding)
+			if randomPadding <= bins[i].paddingSize {
+				continue
+			}
+			size := m.computeObfuscatedPayloadSize(
+				expiresEncodedSize, len(bins[i].serverEntries), bins[i].sumServerEntrySize, randomPadding)
+			if size <= maxPayloadSizeBytes {
+				bins[i].paddingSize = randomPadding
+			} else {
+				// Reduce padding to fit within maxPayloadSizeBytes.
+				excess := size - maxPayloadSizeBytes
+				reduced := randomPadding - excess
+				if reduced > bins[i].paddingSize {
+					bins[i].paddingSize = reduced
+				}
+			}
+		}
+	}
+
+	result.Payloads = make([][]byte, 0, len(bins))
+	result.PayloadEntryCounts = make([]int, 0, len(bins))
+
+	for _, bin := range bins {
+		payload, err := m.buildObfuscatedPayload(
+			bufs, bin.serverEntries, expires, bin.paddingSize)
 		if err != nil {
-			return nil, errors.Trace(err)
+			return result, errors.Trace(err)
 		}
+		// Apply a hard correctness check.
+		if len(payload) > maxPayloadSizeBytes {
+			return result, errors.TraceNew(
+				"internal error: payload size exceeds max")
+		}
+		result.Payloads = append(result.Payloads, payload)
+		result.PayloadEntryCounts = append(
+			result.PayloadEntryCounts, len(bin.serverEntries))
+	}
 
-		// The faster common/prng is appropriate for obfuscation.
-		prng.Read(nonceBuffer[:])
+	return result, nil
+}
 
-		obfuscationBuffer = obfuscationBuffer[:0]
-		obfuscationBuffer = append(obfuscationBuffer, nonceBuffer...)
-		obfuscationBuffer = aead.Seal(
-			obfuscationBuffer, nonceBuffer[:], cborSignedPayload, nil)
+func (m *PushPayloadMaker) buildObfuscatedPayload(
+	bufs *payloadBuffers,
+	prioritizedServerEntries []*PrioritizedServerEntry,
+	expires time.Time,
+	paddingSize int) ([]byte, error) {
 
-		obfuscatedPayloads = append(
-			obfuscatedPayloads, append([]byte(nil), obfuscationBuffer...))
+	obfuscatedPayload, err := m.makeObfuscatedPayload(
+		bufs, prioritizedServerEntries, expires, paddingSize)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+
+	return append([]byte(nil), obfuscatedPayload...), nil
+}
+
+// cborHeaderSize returns the size of a CBOR definite-length header for the
+// given count or length value.
+func cborHeaderSize(n int) int {
+	switch {
+	case n <= 23:
+		return 1
+	case n <= 255:
+		return 2
+	case n <= 65535:
+		return 3
+	default:
+		return 5
+	}
+}
+
+// computeObfuscatedPayloadSize computes the exact obfuscated payload size
+// arithmetically from pre-computed component sizes, avoiding CBOR marshaling.
+//
+// The obfuscated payload structure is:
+//
+//	nonce || AES-GCM(CBOR(SignedPayload{ Signature, CBOR(Payload), Padding })) || tag
+func (m *PushPayloadMaker) computeObfuscatedPayloadSize(
+	expiresEncodedSize int,
+	numEntries int,
+	entrySizeSum int,
+	paddingSize int) int {
+
+	// Payload = map { 1: expires, 2: array(entries) }
+	// With omitempty, the entries field is omitted when numEntries == 0.
+	payloadFields := 1 // Expires
+	payloadBody := 1 + expiresEncodedSize
+	if numEntries > 0 {
+		payloadFields++
+		payloadBody += 1 + cborHeaderSize(numEntries) + entrySizeSum
+	}
+	payloadSize := cborHeaderSize(payloadFields) + payloadBody
+
+	// SignedPayload = map { 1: bstr(signature), 2: bstr(payload), [3: bstr(padding)] }
+	sigLen := signatureSize
+	spFields := 2
+	spBody := 1 + cborHeaderSize(sigLen) + sigLen +
+		1 + cborHeaderSize(payloadSize) + payloadSize
+	if paddingSize > 0 {
+		spFields++
+		spBody += 1 + cborHeaderSize(paddingSize) + paddingSize
 	}
+	signedPayloadSize := cborHeaderSize(spFields) + spBody
+
+	return m.aead.NonceSize() + signedPayloadSize + m.aead.Overhead()
+}
+
+func (m *PushPayloadMaker) makeObfuscatedPayload(
+	bufs *payloadBuffers,
+	prioritizedServerEntries []*PrioritizedServerEntry,
+	expires time.Time,
+	paddingSize int) ([]byte, error) {
+
+	payload := Payload{
+		Expires:                  expires,
+		PrioritizedServerEntries: prioritizedServerEntries,
+	}
+
+	cborPayload, err := protocol.CBOREncoding.Marshal(&payload)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+
+	signature := ed25519.Sign(m.privateKey, cborPayload)
+
+	bufs.signature = bufs.signature[:0]
+	bufs.signature = append(bufs.signature, m.publicKeyID...)
+	bufs.signature = append(bufs.signature, signature...)
+
+	signedPayload := SignedPayload{
+		Signature: bufs.signature,
+		Payload:   cborPayload,
+	}
+
+	if paddingSize < 0 || paddingSize > maxPaddingLimit {
+		return nil, errors.TraceNew("invalid padding size")
+	}
+	if paddingSize > 0 {
+		if bufs.padding == nil {
+			bufs.padding = make([]byte, maxPaddingLimit)
+		}
+		signedPayload.Padding = bufs.padding[:paddingSize]
+	}
+
+	cborSignedPayload, err := protocol.CBOREncoding.Marshal(&signedPayload)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+
+	// The faster common/prng is appropriate for obfuscation.
+	prng.Read(bufs.nonce[:])
+
+	bufs.obfuscation = bufs.obfuscation[:0]
+	bufs.obfuscation = append(bufs.obfuscation, bufs.nonce...)
+	bufs.obfuscation = m.aead.Seal(
+		bufs.obfuscation,
+		bufs.nonce[:],
+		cborSignedPayload,
+		nil)
 
-	return obfuscatedPayloads, nil
+	return bufs.obfuscation, nil
 }

psiphon/common/push/push_test.go

@@ -38,82 +38,412 @@ func TestPush(t *testing.T) {
 	}
 }
 
-func runTestPush() error {
+func TestMakePushPayloads_RF2_RespectsMaxSize(t *testing.T) {
 
 	obfuscationKey, publicKey, privateKey, err := GenerateKeys()
 	if err != nil {
-		return errors.Trace(err)
+		t.Fatal(err)
 	}
 
-	minPadding := 0
-	maxPadding := 65535
+	entries, err := makeTestPrioritizedServerEntries(40, func(i int) int {
+		return (i % 7) * 32
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
 
-	_, incorrectPublicKey, incorrectPrivateKey, err := GenerateKeys()
+	maker, err := NewPushPayloadMaker(obfuscationKey, publicKey, privateKey)
 	if err != nil {
-		return errors.Trace(err)
+		t.Fatal(err)
 	}
 
-	var serverEntries []*PrioritizedServerEntry
+	maxSinglePayloadSize := 0
+	for _, entry := range entries {
+		result, err := maker.MakePushPayloads(
+			0, 0, 1*time.Hour,
+			[]*PrioritizedServerEntry{entry}, 0)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(result.Payloads) != 1 {
+			t.Fatalf("unexpected single-entry payload count: %d", len(result.Payloads))
+		}
+		if len(result.Payloads[0]) > maxSinglePayloadSize {
+			maxSinglePayloadSize = len(result.Payloads[0])
+		}
+	}
 
-	for i := 0; i < 128; i++ {
+	maxPayloadSizeBytes := maxSinglePayloadSize * 4
 
-		serverEntry := &protocol.ServerEntry{
-			Tag:                  prng.Base64String(32),
-			IpAddress:            fmt.Sprintf("192.0.2.%d", i),
-			SshUsername:          prng.HexString(8),
-			SshPassword:          prng.HexString(32),
-			SshHostKey:           prng.Base64String(280),
-			SshObfuscatedPort:    prng.Range(1, 65535),
-			SshObfuscatedKey:     prng.HexString(32),
-			Capabilities:         []string{"OSSH"},
-			Region:               prng.HexString(1),
-			ProviderID:           strings.ToUpper(prng.HexString(8)),
-			ConfigurationVersion: 0,
-			Signature:            prng.Base64String(80),
-		}
+	result, err := maker.MakePushPayloads(
+		0, 0, 1*time.Hour,
+		entries, maxPayloadSizeBytes)
+	if err != nil {
+		t.Fatal(err)
+	}
 
-		serverEntryFields, err := serverEntry.GetServerEntryFields()
-		if err != nil {
-			return errors.Trace(err)
+	if len(result.SkippedIndexes) != 0 {
+		t.Fatalf("unexpected skipped entries: %d", len(result.SkippedIndexes))
+	}
+
+	if len(result.Payloads) <= 1 {
+		t.Fatalf("expected multiple payloads, got %d", len(result.Payloads))
+	}
+
+	for i, payload := range result.Payloads {
+		if len(payload) > maxPayloadSizeBytes {
+			t.Fatalf("payload %d exceeded max size: %d > %d", i, len(payload), maxPayloadSizeBytes)
 		}
+	}
 
-		packed, err := protocol.EncodePackedServerEntryFields(serverEntryFields)
-		if err != nil {
-			return errors.Trace(err)
+	importedSources, err := importPayloadsAndCountSources(
+		obfuscationKey,
+		publicKey,
+		result.Payloads)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(importedSources) != len(entries) {
+		t.Fatalf("unexpected unique import count: %d", len(importedSources))
+	}
+	for source, count := range importedSources {
+		if count != 1 {
+			t.Fatalf("source %s imported %d times", source, count)
 		}
+	}
+}
 
-		serverEntries = append(serverEntries, &PrioritizedServerEntry{
-			ServerEntryFields: packed,
-			Source:            fmt.Sprintf("source-%d", i),
-			PrioritizeDial:    i < 32 || i >= 96,
-		})
+func TestMakePushPayloads_RF2_SkipsOversizeEntry(t *testing.T) {
+
+	obfuscationKey, publicKey, privateKey, err := GenerateKeys()
+	if err != nil {
+		t.Fatal(err)
 	}
 
-	// Test: successful import
+	entries, err := makeTestPrioritizedServerEntries(12, func(_ int) int {
+		return 0
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
 
-	pushServerEntries := [][]*PrioritizedServerEntry{
-		serverEntries[0:32], serverEntries[32:64],
-		serverEntries[64:96], serverEntries[96:128],
+	oversizeEntry, err := makeTestPrioritizedServerEntry(1000000, 300000)
+	if err != nil {
+		t.Fatal(err)
+	}
+	oversizeIndex := len(entries)
+	entries = append(entries, oversizeEntry)
+
+	maker, err := NewPushPayloadMaker(obfuscationKey, publicKey, privateKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	maxPayloadSizeBytes := 4096
+	result, err := maker.MakePushPayloads(
+		0, 0, 1*time.Hour, entries, maxPayloadSizeBytes)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(result.SkippedIndexes) != 1 {
+		t.Fatalf("unexpected skipped index count: %d", len(result.SkippedIndexes))
+	}
+	if result.SkippedIndexes[0] != oversizeIndex {
+		t.Fatalf("unexpected skipped index: %d", result.SkippedIndexes[0])
+	}
+
+	for i, payload := range result.Payloads {
+		if len(payload) > maxPayloadSizeBytes {
+			t.Fatalf("payload %d exceeded max size: %d > %d", i, len(payload), maxPayloadSizeBytes)
+		}
 	}
 
-	payloads, err := MakePushPayloads(
+	importedSources, err := importPayloadsAndCountSources(
 		obfuscationKey,
-		minPadding,
-		maxPadding,
 		publicKey,
-		privateKey,
-		1*time.Hour,
-		pushServerEntries)
+		result.Payloads)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(importedSources) != len(entries)-1 {
+		t.Fatalf("unexpected import count: %d", len(importedSources))
+	}
+	if _, ok := importedSources[oversizeEntry.Source]; ok {
+		t.Fatalf("oversize entry was imported")
+	}
+}
+
+func TestMakePushPayloads_RF2_StrictCapWithPadding(t *testing.T) {
+
+	obfuscationKey, publicKey, privateKey, err := GenerateKeys()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	entries, err := makeTestPrioritizedServerEntries(30, func(i int) int {
+		return (i % 5) * 16
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	maker, err := NewPushPayloadMaker(obfuscationKey, publicKey, privateKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	minPadding := 1024
+	maxPadding := 1024
+	maxSinglePayloadSize := 0
+	for _, entry := range entries {
+		result, err := maker.MakePushPayloads(
+			minPadding, maxPadding, 1*time.Hour,
+			[]*PrioritizedServerEntry{entry}, 0)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(result.Payloads[0]) > maxSinglePayloadSize {
+			maxSinglePayloadSize = len(result.Payloads[0])
+		}
+	}
+
+	maxPayloadSizeBytes := maxSinglePayloadSize * 3
+	result, err := maker.MakePushPayloads(
+		minPadding, maxPadding, 1*time.Hour,
+		entries, maxPayloadSizeBytes)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(result.SkippedIndexes) != 0 {
+		t.Fatalf("unexpected skipped entries: %d", len(result.SkippedIndexes))
+	}
+
+	for i, payload := range result.Payloads {
+		if len(payload) > maxPayloadSizeBytes {
+			t.Fatalf("payload %d exceeded max size: %d > %d", i, len(payload), maxPayloadSizeBytes)
+		}
+	}
+}
+
+func TestMakePushPayloads_MetadataIntegrity(t *testing.T) {
+
+	obfuscationKey, publicKey, privateKey, err := GenerateKeys()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	entries, err := makeTestPrioritizedServerEntries(16, func(_ int) int {
+		return 0
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	oversizeEntry, err := makeTestPrioritizedServerEntry(2000000, 300000)
+	if err != nil {
+		t.Fatal(err)
+	}
+	entries = append(entries, oversizeEntry)
+
+	maker, err := NewPushPayloadMaker(obfuscationKey, publicKey, privateKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	result, err := maker.MakePushPayloads(
+		0, 0, 1*time.Hour, entries, 4096)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(result.Payloads) != len(result.PayloadEntryCounts) {
+		t.Fatalf("payload/entry-count mismatch: %d vs %d", len(result.Payloads), len(result.PayloadEntryCounts))
+	}
+
+	totalPayloadEntries := 0
+	for _, payloadEntryCount := range result.PayloadEntryCounts {
+		totalPayloadEntries += payloadEntryCount
+	}
+
+	if totalPayloadEntries+len(result.SkippedIndexes) != len(entries) {
+		t.Fatalf("metadata does not account for all entries")
+	}
+
+	seenSkippedIndexes := make(map[int]bool)
+	for _, skippedIndex := range result.SkippedIndexes {
+		if skippedIndex < 0 || skippedIndex >= len(entries) {
+			t.Fatalf("invalid skipped index: %d", skippedIndex)
+		}
+		if seenSkippedIndexes[skippedIndex] {
+			t.Fatalf("duplicate skipped index: %d", skippedIndex)
+		}
+		seenSkippedIndexes[skippedIndex] = true
+	}
+}
+
+func TestComputeObfuscatedPayloadSize_MatchesMeasured(t *testing.T) {
+
+	obfuscationKey, publicKey, _, err := GenerateKeys()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, _, privateKey, err := GenerateKeys()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	maker, err := NewPushPayloadMaker(obfuscationKey, publicKey, privateKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	bufs := &payloadBuffers{}
+
+	expires := time.Now().Add(1 * time.Hour).UTC()
+	expiresEncoded, err := protocol.CBOREncoding.Marshal(expires)
+	if err != nil {
+		t.Fatal(err)
+	}
+	expiresEncodedSize := len(expiresEncoded)
+
+	allEntries, err := makeTestPrioritizedServerEntries(20, func(i int) int {
+		return i * 50
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	for _, paddingSize := range []int{0, 1024, 65535} {
+		for numEntries := 0; numEntries <= len(allEntries); numEntries++ {
+			entries := allEntries[:numEntries]
+
+			measured, err := maker.measureObfuscatedPayloadSize(
+				bufs, entries, expires, paddingSize)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			entrySizeSum := 0
+			for _, entry := range entries {
+				encodedEntry, err := protocol.CBOREncoding.Marshal(entry)
+				if err != nil {
+					t.Fatal(err)
+				}
+				entrySizeSum += len(encodedEntry)
+			}
+
+			computed := maker.computeObfuscatedPayloadSize(
+				expiresEncodedSize, numEntries, entrySizeSum, paddingSize)
+
+			if computed != measured {
+				t.Fatalf(
+					"mismatch: numEntries=%d paddingSize=%d computed=%d measured=%d",
+					numEntries, paddingSize, computed, measured)
+			}
+		}
+	}
+}
+
+// measureObfuscatedPayloadSize computes the obfuscated payload size by
+// performing real CBOR marshaling. This is the reference implementation used
+// to validate the arithmetic computation in computeObfuscatedPayloadSize.
+func (m *PushPayloadMaker) measureObfuscatedPayloadSize(
+	bufs *payloadBuffers,
+	prioritizedServerEntries []*PrioritizedServerEntry,
+	expires time.Time,
+	paddingSize int) (int, error) {
+
+	payload := Payload{
+		Expires:                  expires,
+		PrioritizedServerEntries: prioritizedServerEntries,
+	}
+
+	cborPayload, err := protocol.CBOREncoding.Marshal(&payload)
+	if err != nil {
+		return 0, errors.Trace(err)
+	}
+
+	signedPayload := SignedPayload{
+		Signature: make([]byte, signatureSize),
+		Payload:   cborPayload,
+	}
+
+	if paddingSize < 0 || paddingSize > maxPaddingLimit {
+		return 0, errors.TraceNew("invalid padding size")
+	}
+	if paddingSize > 0 {
+		if bufs.padding == nil {
+			bufs.padding = make([]byte, maxPaddingLimit)
+		}
+		signedPayload.Padding = bufs.padding[:paddingSize]
+	}
+
+	cborSignedPayload, err := protocol.CBOREncoding.Marshal(&signedPayload)
+	if err != nil {
+		return 0, errors.Trace(err)
+	}
+
+	return m.aead.NonceSize() + len(cborSignedPayload) + m.aead.Overhead(), nil
+}
+
+func runTestPush() error {
+
+	obfuscationKey, publicKey, privateKey, err := GenerateKeys()
 	if err != nil {
 		return errors.Trace(err)
 	}
 
-	if len(payloads) != len(pushServerEntries) {
-		return errors.TraceNew("unexpected payload count")
+	minPadding := 0
+	maxPadding := 65535
+
+	_, incorrectPublicKey, incorrectPrivateKey, err := GenerateKeys()
+	if err != nil {
+		return errors.Trace(err)
 	}
 
-	expectPrioritizeDial := true
+	serverEntries, err := makeTestPrioritizedServerEntries(128, func(_ int) int {
+		return 0
+	})
+	if err != nil {
+		return errors.Trace(err)
+	}
+
+	expectPrioritizeDial := make(map[string]bool)
+	for _, serverEntry := range serverEntries {
+		expectPrioritizeDial[serverEntry.Source] = serverEntry.PrioritizeDial
+	}
 
+	maker, err := NewPushPayloadMaker(obfuscationKey, publicKey, privateKey)
+	if err != nil {
+		return errors.Trace(err)
+	}
+
+	incorrectMaker, err := NewPushPayloadMaker(obfuscationKey, publicKey, incorrectPrivateKey)
+	if err != nil {
+		return errors.Trace(err)
+	}
+
+	// Test: successful import
+
+	result, err := maker.MakePushPayloads(
+		minPadding, maxPadding, 1*time.Hour, serverEntries, 0)
+	if err != nil {
+		return errors.Trace(err)
+	}
+
+	if len(result.Payloads) != 1 {
+		return errors.TraceNew("unexpected payload count")
+	}
+	if len(result.PayloadEntryCounts) != 1 || result.PayloadEntryCounts[0] != len(serverEntries) {
+		return errors.TraceNew("unexpected payload entry counts")
+	}
+
+	seenSources := make(map[string]int)
 	importer := func(
 		packedServerEntryFields protocol.PackedServerEntryFields,
 		source string,
@@ -126,16 +456,19 @@ func runTestPush() error {
 		if !strings.HasPrefix(serverEntryFields["ipAddress"].(string), "192.0.2") {
 			return errors.TraceNew("unexpected server entry IP address")
 		}
-		if prioritizeDial != expectPrioritizeDial {
+		expect, ok := expectPrioritizeDial[source]
+		if !ok {
+			return errors.TraceNew("unexpected source")
+		}
+		if prioritizeDial != expect {
 			return errors.TraceNew("unexpected prioritize dial")
 		}
+		seenSources[source] += 1
 		return nil
 	}
 
-	for i, payload := range payloads {
-
-		expectPrioritizeDial = i == 0 || i == 3
-
+	totalImported := 0
+	for _, payload := range result.Payloads {
 		n, err := ImportPushPayload(
 			obfuscationKey,
 			publicKey,
@@ -144,22 +477,22 @@ func runTestPush() error {
 		if err != nil {
 			return errors.Trace(err)
 		}
+		totalImported += n
+	}
 
-		if n != 32 {
-			return errors.TraceNew("unexpected import count")
+	if totalImported != len(serverEntries) {
+		return errors.TraceNew("unexpected import count")
+	}
+	for source, count := range seenSources {
+		if count != 1 {
+			return errors.Tracef("source imported unexpected number of times: %s=%d", source, count)
 		}
 	}
 
 	// Test: expired
 
-	payloads, err = MakePushPayloads(
-		obfuscationKey,
-		minPadding,
-		maxPadding,
-		publicKey,
-		privateKey,
-		1*time.Microsecond,
-		pushServerEntries)
+	result, err = maker.MakePushPayloads(
+		minPadding, maxPadding, 1*time.Microsecond, serverEntries, 0)
 	if err != nil {
 		return errors.Trace(err)
 	}
@@ -169,7 +502,7 @@ func runTestPush() error {
 	_, err = ImportPushPayload(
 		obfuscationKey,
 		publicKey,
-		payloads[0],
+		result.Payloads[0],
 		importer)
 	if err == nil {
 		return errors.TraceNew("unexpected success")
@@ -177,14 +510,8 @@ func runTestPush() error {
 
 	// Test: invalid signature
 
-	payloads, err = MakePushPayloads(
-		obfuscationKey,
-		minPadding,
-		maxPadding,
-		publicKey,
-		incorrectPrivateKey,
-		1*time.Hour,
-		pushServerEntries)
+	result, err = incorrectMaker.MakePushPayloads(
+		minPadding, maxPadding, 1*time.Hour, serverEntries, 0)
 	if err != nil {
 		return errors.Trace(err)
 	}
@@ -192,7 +519,7 @@ func runTestPush() error {
 	_, err = ImportPushPayload(
 		obfuscationKey,
 		publicKey,
-		payloads[0],
+		result.Payloads[0],
 		importer)
 	if err == nil {
 		return errors.TraceNew("unexpected success")
@@ -200,14 +527,8 @@ func runTestPush() error {
 
 	// Test: wrong signature key
 
-	payloads, err = MakePushPayloads(
-		obfuscationKey,
-		minPadding,
-		maxPadding,
-		publicKey,
-		privateKey,
-		1*time.Hour,
-		pushServerEntries)
+	result, err = maker.MakePushPayloads(
+		minPadding, maxPadding, 1*time.Hour, serverEntries, 0)
 	if err != nil {
 		return errors.Trace(err)
 	}
@@ -215,7 +536,7 @@ func runTestPush() error {
 	_, err = ImportPushPayload(
 		obfuscationKey,
 		incorrectPublicKey,
-		payloads[0],
+		result.Payloads[0],
 		importer)
 	if err == nil {
 		return errors.TraceNew("unexpected success")
@@ -223,24 +544,18 @@ func runTestPush() error {
 
 	// Test: mutate obfuscation layer
 
-	payloads, err = MakePushPayloads(
-		obfuscationKey,
-		minPadding,
-		maxPadding,
-		publicKey,
-		privateKey,
-		1*time.Hour,
-		pushServerEntries)
+	result, err = maker.MakePushPayloads(
+		minPadding, maxPadding, 1*time.Hour, serverEntries, 0)
 	if err != nil {
 		return errors.Trace(err)
 	}
 
-	payloads[0][0] = ^payloads[0][0]
+	result.Payloads[0][0] = ^result.Payloads[0][0]
 
 	_, err = ImportPushPayload(
 		obfuscationKey,
 		publicKey,
-		payloads[0],
+		result.Payloads[0],
 		importer)
 	if err == nil {
 		return errors.TraceNew("unexpected success")
@@ -248,3 +563,190 @@ func runTestPush() error {
 
 	return nil
 }
+
+func makeTestPrioritizedServerEntries(
+	count int,
+	sourceExtraBytes func(index int) int) ([]*PrioritizedServerEntry, error) {
+
+	serverEntries := make([]*PrioritizedServerEntry, 0, count)
+	for i := range count {
+		entry, err := makeTestPrioritizedServerEntry(i, sourceExtraBytes(i))
+		if err != nil {
+			return nil, errors.Trace(err)
+		}
+		serverEntries = append(serverEntries, entry)
+	}
+
+	return serverEntries, nil
+}
+
+func makeTestPrioritizedServerEntry(
+	index int,
+	sourceExtraBytes int) (*PrioritizedServerEntry, error) {
+
+	serverEntry := &protocol.ServerEntry{
+		Tag:                  prng.Base64String(32),
+		IpAddress:            fmt.Sprintf("192.0.2.%d", index%255),
+		SshUsername:          prng.HexString(8),
+		SshPassword:          prng.HexString(32),
+		SshHostKey:           prng.Base64String(280),
+		SshObfuscatedPort:    prng.Range(1, 65535),
+		SshObfuscatedKey:     prng.HexString(32),
+		Capabilities:         []string{"OSSH"},
+		Region:               prng.HexString(1),
+		ProviderID:           strings.ToUpper(prng.HexString(8)),
+		ConfigurationVersion: 0,
+		Signature:            prng.Base64String(80),
+	}
+
+	serverEntryFields, err := serverEntry.GetServerEntryFields()
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+
+	packed, err := protocol.EncodePackedServerEntryFields(serverEntryFields)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+
+	source := fmt.Sprintf("source-%d", index)
+	if sourceExtraBytes > 0 {
+		source = source + strings.Repeat("s", sourceExtraBytes)
+	}
+
+	return &PrioritizedServerEntry{
+		ServerEntryFields: packed,
+		Source:            source,
+		PrioritizeDial:    index < 32 || index >= 96,
+	}, nil
+}
+
+func importPayloadsAndCountSources(
+	obfuscationKey string,
+	signaturePublicKey string,
+	payloads [][]byte) (map[string]int, error) {
+
+	sourceCounts := make(map[string]int)
+	importer := func(
+		packedServerEntryFields protocol.PackedServerEntryFields,
+		source string,
+		_ bool) error {
+
+		_, err := protocol.DecodePackedServerEntryFields(packedServerEntryFields)
+		if err != nil {
+			return errors.Trace(err)
+		}
+		sourceCounts[source] += 1
+		return nil
+	}
+
+	for _, payload := range payloads {
+		_, err := ImportPushPayload(
+			obfuscationKey,
+			signaturePublicKey,
+			payload,
+			importer)
+		if err != nil {
+			return nil, errors.Trace(err)
+		}
+	}
+
+	return sourceCounts, nil
+}
+
+// BenchmarkMakePushPayloads_RF2_AverageBucketUtilization-16    	    8608	    138130 ns/op	         0.9204 avg_utilization	         4.000 payloads/op	   75750 B/op	     554 allocs/op
+func BenchmarkMakePushPayloads_RF2_AverageBucketUtilization(b *testing.B) {
+
+	obfuscationKey, publicKey, privateKey, err := GenerateKeys()
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	// Create 10 server entries with sizes ranging from ~700 to ~2500 bytes.
+	// Base entry is ~500-700 bytes, so add 0-2000 extra bytes to source field.
+	entries, err := makeTestPrioritizedServerEntries(10, func(i int) int {
+		// Vary size from 0 to 2000 bytes across the 10 entries.
+		return i * 200
+	})
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	maxPayloadSizeBytes := 4096
+
+	maker, err := NewPushPayloadMaker(obfuscationKey, publicKey, privateKey)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	totalUtilization := 0.0
+	totalPayloadCount := 0.0
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		result, err := maker.MakePushPayloads(
+			0, 0, 1*time.Hour, entries, maxPayloadSizeBytes)
+		if err != nil {
+			b.Fatal(err)
+		}
+
+		if len(result.SkippedIndexes) != 0 {
+			b.Fatalf("unexpected skipped entries: %d", len(result.SkippedIndexes))
+		}
+		if len(result.Payloads) == 0 {
+			b.Fatal("no payloads generated")
+		}
+
+		totalPayloadBytes := 0
+		for payloadIndex, payload := range result.Payloads {
+			if len(payload) > maxPayloadSizeBytes {
+				b.Fatalf("payload %d exceeded max size: %d > %d", payloadIndex, len(payload), maxPayloadSizeBytes)
+			}
+			totalPayloadBytes += len(payload)
+		}
+
+		averageBucketUtilization := float64(totalPayloadBytes) /
+			float64(len(result.Payloads)*maxPayloadSizeBytes)
+		totalUtilization += averageBucketUtilization
+		totalPayloadCount += float64(len(result.Payloads))
+	}
+	b.StopTimer()
+
+	b.ReportMetric(totalUtilization/float64(b.N), "avg_utilization")
+	b.ReportMetric(totalPayloadCount/float64(b.N), "payloads/op")
+}
+
+// BenchmarkMakePushPayloads-16    	    8889	    139301 ns/op	   75742 B/op	     554 allocs/op
+func BenchmarkMakePushPayloads(b *testing.B) {
+
+	obfuscationKey, publicKey, privateKey, err := GenerateKeys()
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	// Create 10 server entries with sizes ranging from ~700 to ~2500 bytes.
+	// Base entry is ~500-700 bytes, so add 0-2000 extra bytes to source field.
+	entries, err := makeTestPrioritizedServerEntries(10, func(i int) int {
+		// Vary size from 0 to 2000 bytes across the 10 entries
+		return i * 200
+	})
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	maxPayloadSizeBytes := 4096
+
+	maker, err := NewPushPayloadMaker(obfuscationKey, publicKey, privateKey)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, err := maker.MakePushPayloads(
+			0, 0, 1*time.Hour, entries, maxPayloadSizeBytes)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}