vdaf/prio3: validate preparation inputs

vdaf/prio3/internal/prio3/prio3.go

@@ -247,7 +247,7 @@ func (v *Prio3[M, A, T, V, E, F]) PrepInit(
 	ps PublicShare, inputShare InputShare[V, E],
 ) (*PrepState[V, E], *PrepShare[V, E], error) {
 	params := v.Params()
-	if aggID > v.shares {
+	if aggID >= v.shares {
 		return nil, nil, ErrAggID
 	}
 
@@ -264,7 +264,7 @@ func (v *Prio3[M, A, T, V, E, F]) PrepInit(
 	var jointRand V
 	jointRandLen := params.JointRandLength()
 	if jointRandLen > 0 {
-		if share.blind == nil || len(ps) == 0 {
+		if share.blind == nil || len(ps) != int(SeedSize)*int(v.shares) {
 			return nil, nil, ErrJointRand
 		}
 
@@ -280,9 +280,17 @@ func (v *Prio3[M, A, T, V, E, F]) PrepInit(
 		if err != nil {
 			return nil, nil, err
 		}
+		partStart := int(aggID) * int(SeedSize)
+		partEnd := partStart + int(SeedSize)
+		if subtle.ConstantTimeCompare(
+			ps[partStart:partEnd], prepShare.jointRandPart[:]) != 1 {
+			return nil, nil, ErrJointRand
+		}
+		correctedJointRandParts := append(PublicShare(nil), ps...)
+		copy(correctedJointRandParts[partStart:partEnd], prepShare.jointRandPart[:])
 
 		prepState.correctedJointRandSeed = &Seed{}
-		*prepState.correctedJointRandSeed, err = v.xof.jointRandSeed(ps)
+		*prepState.correctedJointRandSeed, err = v.xof.jointRandSeed(correctedJointRandParts)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -370,8 +378,13 @@ func (v *Prio3[M, A, T, V, E, F]) PrepSharesToPrep(
 func (v *Prio3[M, A, T, V, E, F]) PrepNext(
 	state *PrepState[V, E], msg *PrepMessage,
 ) (*OutShare[V, E], error) {
-	if msg != nil && state != nil && msg.joinRand != nil &&
-		state.correctedJointRandSeed != nil {
+	if state == nil {
+		return nil, ErrShare
+	}
+	if state.correctedJointRandSeed != nil {
+		if msg == nil || msg.joinRand == nil {
+			return nil, ErrJointRand
+		}
 		if subtle.ConstantTimeCompare(
 			msg.joinRand[:], state.correctedJointRandSeed[:]) != 1 {
 			return nil, ErrJointRand

vdaf/prio3/prio3_test.go

@@ -2,6 +2,7 @@ package prio3
 
 import (
 	"crypto/rand"
+	"errors"
 	"io"
 	"slices"
 	"testing"
@@ -112,6 +113,60 @@ func TestMultiHotCountVec(t *testing.T) {
 	}
 }
 
+func TestPreparationValidation(t *testing.T) {
+	h, err := histogram.New(NumShares, 4, 3, Context)
+	test.CheckNoErr(t, err, "new histogram failed")
+	params := h.Params()
+	nonce := fromReader[histogram.Nonce](t, rand.Reader)
+	verifyKey := fromReader[histogram.VerifyKey](t, rand.Reader)
+	randb := make([]byte, params.RandSize())
+	_, err = io.ReadFull(rand.Reader, randb)
+	test.CheckNoErr(t, err, "read rand bytes failed")
+	publicShare, inputShares, err := h.Shard(0, &nonce, randb)
+	test.CheckNoErr(t, err, "Shard failed")
+
+	if _, _, err = h.PrepInit(
+		&verifyKey, &nonce, NumShares, publicShare, inputShares[1],
+	); !errors.Is(err, prio3.ErrAggID) {
+		t.Fatalf("got error %v for out-of-range aggregator ID, want %v", err, prio3.ErrAggID)
+	}
+
+	invalidPublicShares := []histogram.PublicShare{
+		nil,
+		publicShare[:len(publicShare)-1],
+		append(slices.Clone(publicShare), 0),
+	}
+	for _, invalidPublicShare := range invalidPublicShares {
+		if _, _, err = h.PrepInit(
+			&verifyKey, &nonce, 0, invalidPublicShare, inputShares[0],
+		); !errors.Is(err, prio3.ErrJointRand) {
+			t.Fatalf("PrepInit with %d-byte public share: got error %v, want %v",
+				len(invalidPublicShare), err, prio3.ErrJointRand)
+		}
+	}
+
+	corruptedPublicShare := slices.Clone(publicShare)
+	corruptedPublicShare[0] ^= 1
+	if _, _, err = h.PrepInit(
+		&verifyKey, &nonce, 0, corruptedPublicShare, inputShares[0],
+	); !errors.Is(err, prio3.ErrJointRand) {
+		t.Fatalf("got error %v for corrupted joint randomness part, want %v",
+			err, prio3.ErrJointRand)
+	}
+
+	state, _, err := h.PrepInit(&verifyKey, &nonce, 0, publicShare, inputShares[0])
+	test.CheckNoErr(t, err, "PrepInit failed")
+	if _, err = h.PrepNext(state, nil); !errors.Is(err, prio3.ErrJointRand) {
+		t.Fatalf("PrepNext with nil message: got error %v, want %v", err, prio3.ErrJointRand)
+	}
+	if _, err = h.PrepNext(state, new(histogram.PrepMessage)); !errors.Is(err, prio3.ErrJointRand) {
+		t.Fatalf("PrepNext with empty message: got error %v, want %v", err, prio3.ErrJointRand)
+	}
+	if _, err = h.PrepNext(nil, nil); !errors.Is(err, prio3.ErrShare) {
+		t.Fatalf("PrepNext with nil state: got error %v, want %v", err, prio3.ErrShare)
+	}
+}
+
 func testPrio3[
 	P Prio3[
 		Measurement, Aggregate,