Fix: change version handling in query membership

balloon/balloon.go

@@ -255,7 +255,9 @@ func (b Balloon) QueryDigestMembership(keyDigest hashing.Digest, version uint64)
 	proof.KeyDigest = keyDigest
 	proof.QueryVersion = version
 	proof.CurrentVersion = b.version - 1
-
+	if version > b.version {
+		version = b.version
+	}
 	leaf, err := b.store.Get(storage.IndexPrefix, proof.KeyDigest)
 	if err != nil {
 		return nil, fmt.Errorf("No leaf with digest %v", proof.KeyDigest)
@@ -264,14 +266,14 @@ func (b Balloon) QueryDigestMembership(keyDigest hashing.Digest, version uint64)
 	proof.Exists = true
 	proof.ActualVersion = util.BytesAsUint64(leaf.Value)
 
-	if proof.ActualVersion <= proof.QueryVersion {
+	if proof.ActualVersion <= version {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			historyProof, historyErr = b.historyTree.ProveMembership(proof.ActualVersion, proof.QueryVersion)
+			historyProof, historyErr = b.historyTree.ProveMembership(proof.ActualVersion, version)
 		}()
 	} else {
-		return nil, fmt.Errorf("Query version %d is not on history tree which version is %d", proof.QueryVersion, proof.ActualVersion)
+		return nil, fmt.Errorf("Query version %d is not on history tree which version is %d", version, proof.ActualVersion)
 	}
 
 	hyperProof, hyperErr = b.hyperTree.QueryMembership(leaf.Key, leaf.Value)