add support for Bao slices

bao.go

@@ -199,3 +199,120 @@ func BaoVerifyBuf(data, outboard []byte, group int, root [32]byte) bool {
 	ok, _ := BaoDecode(io.Discard, d, or, group, root)
 	return ok && d.Len() == 0 && o.Len() == 0 // check for trailing data
 }
+
+// BaoExtractSlice returns the slice encoding for the given offset and length.
+// When extracting from an outboard encoding, data should contain only the chunk
+// groups that will be present in the slice.
+func BaoExtractSlice(dst io.Writer, data, outboard io.Reader, group int, offset uint64, length uint64) error {
+	combinedEncoding := outboard == nil
+	if combinedEncoding {
+		outboard = data
+	}
+	groupSize := uint64(chunkSize << group)
+	buf := make([]byte, groupSize)
+	var err error
+	read := func(r io.Reader, n uint64, copy bool) {
+		if err == nil {
+			_, err = io.ReadFull(r, buf[:n])
+			if err == nil && copy {
+				_, err = dst.Write(buf[:n])
+			}
+		}
+	}
+	var pos uint64
+	var rec func(bufLen uint64)
+	rec = func(bufLen uint64) {
+		inSlice := pos < (offset+length) && offset < (pos+bufLen)
+		if err != nil {
+			return
+		} else if bufLen <= groupSize {
+			if combinedEncoding || inSlice {
+				read(data, bufLen, inSlice)
+			}
+			pos += bufLen
+			return
+		}
+		read(outboard, 64, inSlice)
+		mid := uint64(1) << (bits.Len64(bufLen-1) - 1)
+		rec(mid)
+		rec(bufLen - mid)
+	}
+	read(outboard, 8, true)
+	rec(binary.LittleEndian.Uint64(buf[:8]))
+	return err
+}
+
+// BaoDecodeSlice reads from data, which must contain a slice encoding for the
+// given offset and length, and streams verified content to dst. It returns
+// false if verification fails.
+func BaoDecodeSlice(dst io.Writer, data io.Reader, group int, offset, length uint64, root [32]byte) (bool, error) {
+	groupSize := uint64(chunkSize << group)
+	buf := make([]byte, groupSize)
+	var err error
+	read := func(n uint64) []byte {
+		if err == nil {
+			_, err = io.ReadFull(data, buf[:n])
+		}
+		return buf[:n]
+	}
+	readParent := func() (l, r [8]uint32) {
+		read(64)
+		return bytesToCV(buf[:32]), bytesToCV(buf[32:])
+	}
+	write := func(p []byte) {
+		if err == nil {
+			_, err = dst.Write(p)
+		}
+	}
+	var pos uint64
+	var rec func(cv [8]uint32, bufLen uint64, flags uint32) bool
+	rec = func(cv [8]uint32, bufLen uint64, flags uint32) bool {
+		inSlice := pos < (offset+length) && offset < (pos+bufLen)
+		if err != nil {
+			return false
+		} else if bufLen <= groupSize {
+			if !inSlice {
+				pos += bufLen
+				return true
+			}
+			n := compressGroup(read(bufLen), pos/chunkSize)
+			n.flags |= flags
+			valid := cv == chainingValue(n)
+			if valid {
+				// only write within range
+				p := buf[:bufLen]
+				if pos+bufLen > offset+length {
+					p = p[:offset+length-pos]
+				}
+				if pos < offset {
+					p = p[offset-pos:]
+				}
+				write(p)
+			}
+			pos += bufLen
+			return valid
+		}
+		if !inSlice {
+			return true
+		}
+		l, r := readParent()
+		n := parentNode(l, r, iv, flags)
+		mid := uint64(1) << (bits.Len64(bufLen-1) - 1)
+		return chainingValue(n) == cv && rec(l, mid, 0) && rec(r, bufLen-mid, 0)
+	}
+
+	dataLen := binary.LittleEndian.Uint64(read(8))
+	ok := rec(bytesToCV(root[:]), dataLen, flagRoot)
+	return ok, err
+}
+
+// BaoVerifySlice verifies the Bao slice encoding in data, returning the
+// verified bytes.
+func BaoVerifySlice(data []byte, group int, offset uint64, length uint64, root [32]byte) ([]byte, bool) {
+	d := bytes.NewBuffer(data)
+	var buf bytes.Buffer
+	if ok, _ := BaoDecodeSlice(&buf, d, group, offset, length, root); !ok || d.Len() > 0 {
+		return nil, false
+	}
+	return buf.Bytes(), true
+}

bao_test.go

@@ -190,3 +190,48 @@ func TestBaoStreaming(t *testing.T) {
 		t.Fatal("invalid data was written to buf")
 	}
 }
+
+func TestBaoSlice(t *testing.T) {
+	data := make([]byte, 1<<20)
+	blake3.New(0, nil).XOF().Read(data)
+
+	for _, test := range []struct {
+		off, len uint64
+	}{
+		{0, uint64(len(data))},
+		{0, 1024},
+		{1024, 1024},
+		{0, 10},
+		{1020, 10},
+		{1030, uint64(len(data) - 1030)},
+	} {
+		// combined encoding
+		{
+			enc, root := blake3.BaoEncodeBuf(data, 0, false)
+			var buf bytes.Buffer
+			if err := blake3.BaoExtractSlice(&buf, bytes.NewReader(enc), nil, 0, test.off, test.len); err != nil {
+				t.Error(err)
+			} else if vdata, ok := blake3.BaoVerifySlice(buf.Bytes(), 0, test.off, test.len, root); !ok {
+				t.Error("combined verify failed", test)
+			} else if !bytes.Equal(vdata, data[test.off:][:test.len]) {
+				t.Error("combined bad decode", test, vdata, data[test.off:][:test.len])
+			}
+		}
+		// outboard encoding
+		{
+			enc, root := blake3.BaoEncodeBuf(data, 0, true)
+			start, end := (test.off/1024)*1024, ((test.off+test.len+1024-1)/1024)*1024
+			if end > uint64(len(data)) {
+				end = uint64(len(data))
+			}
+			var buf bytes.Buffer
+			if err := blake3.BaoExtractSlice(&buf, bytes.NewReader(data[start:end]), bytes.NewReader(enc), 0, test.off, test.len); err != nil {
+				t.Error(err)
+			} else if vdata, ok := blake3.BaoVerifySlice(buf.Bytes(), 0, test.off, test.len, root); !ok {
+				t.Error("outboard verify failed", test)
+			} else if !bytes.Equal(vdata, data[test.off:][:test.len]) {
+				t.Error("outboard bad decode", test, vdata, data[test.off:][:test.len])
+			}
+		}
+	}
+}