refactor: bao verifier needs to manage a data buffer to allow any number of bytes read while verifying in chunks
This commit is contained in:
parent
9f5b676f47
commit
c8a6570e8f
59
bao/bao.go
59
bao/bao.go
|
@ -30,28 +30,47 @@ type Verifier struct {
|
||||||
r io.ReadCloser
|
r io.ReadCloser
|
||||||
proof Result
|
proof Result
|
||||||
read uint64
|
read uint64
|
||||||
|
buffer *bytes.Buffer
|
||||||
}
|
}
|
||||||
|
|
||||||
func (v Verifier) Read(p []byte) (n int, err error) {
|
func (v *Verifier) Read(p []byte) (int, error) {
|
||||||
var buf [VERIFY_CHUNK_SIZE]byte
|
// Initial attempt to read from the buffer
|
||||||
|
n, err := v.buffer.Read(p)
|
||||||
n, err = v.r.Read(buf[:])
|
if n == len(p) {
|
||||||
if err != nil {
|
// If the buffer already had enough data to fulfill the request, return immediately
|
||||||
return 0, err
|
|
||||||
}
|
|
||||||
|
|
||||||
if !bao.Verify(buf[:n], v.read, v.proof.Proof, v.proof.Hash) {
|
|
||||||
return 0, ErrVerifyFailed
|
|
||||||
}
|
|
||||||
|
|
||||||
v.read += uint64(n)
|
|
||||||
|
|
||||||
copy(p, buf[:n])
|
|
||||||
|
|
||||||
return n, nil
|
return n, nil
|
||||||
|
} else if err != nil && err != io.EOF {
|
||||||
|
// For errors other than EOF, return the error immediately
|
||||||
|
return n, err
|
||||||
}
|
}
|
||||||
|
|
||||||
func (v Verifier) Close() error {
|
buf := make([]byte, VERIFY_CHUNK_SIZE)
|
||||||
|
// Continue reading from the source and verifying until we have enough data or hit an error
|
||||||
|
for v.buffer.Len() < len(p)-n {
|
||||||
|
bytesRead, err := v.r.Read(buf)
|
||||||
|
if err != nil && err != io.EOF {
|
||||||
|
return n, err // Return any read error immediately
|
||||||
|
}
|
||||||
|
|
||||||
|
if !bao.Verify(buf[:bytesRead], v.read, v.proof.Proof, v.proof.Hash) {
|
||||||
|
return n, ErrVerifyFailed
|
||||||
|
}
|
||||||
|
|
||||||
|
v.read += uint64(bytesRead)
|
||||||
|
v.buffer.Write(buf[:bytesRead]) // Append new data to the buffer
|
||||||
|
|
||||||
|
if err == io.EOF {
|
||||||
|
// If EOF, break the loop as no more data can be read
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Attempt to read the remainder of the data from the buffer
|
||||||
|
additionalBytes, _ := v.buffer.Read(p[n:])
|
||||||
|
return n + additionalBytes, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (v *Verifier) Close() error {
|
||||||
return v.r.Close()
|
return v.r.Close()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -145,5 +164,9 @@ func Hash(r io.Reader) (*Result, error) {
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewVerifier(r io.ReadCloser, proof Result) *Verifier {
|
func NewVerifier(r io.ReadCloser, proof Result) *Verifier {
|
||||||
return &Verifier{r: r, proof: proof}
|
return &Verifier{
|
||||||
|
r: r,
|
||||||
|
proof: proof,
|
||||||
|
buffer: new(bytes.Buffer),
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue