150 lines
2.8 KiB
Go
150 lines
2.8 KiB
Go
package bao
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
_ "embed"
|
|
"errors"
|
|
"io"
|
|
"math"
|
|
"os"
|
|
"os/exec"
|
|
|
|
"github.com/docker/go-units"
|
|
"github.com/hashicorp/go-plugin"
|
|
)
|
|
|
|
//go:generate buf generate
|
|
//go:generate bash -c "cd rust && cargo build -r"
|
|
//go:embed rust/target/release/rust
|
|
var pluginBin []byte
|
|
|
|
var bao Bao
|
|
var client *plugin.Client
|
|
|
|
var _ io.ReadCloser = (*Verifier)(nil)
|
|
|
|
var ErrVerifyFailed = errors.New("verification failed")
|
|
|
|
type Verifier struct {
|
|
r io.ReadCloser
|
|
proof Result
|
|
read uint64
|
|
}
|
|
|
|
func (v Verifier) Read(p []byte) (n int, err error) {
|
|
var buf [VERIFY_CHUNK_SIZE]byte
|
|
|
|
n, err = io.ReadAtLeast(v.r, buf[:], VERIFY_CHUNK_SIZE)
|
|
if err != nil {
|
|
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
|
|
}
|
|
|
|
func (v Verifier) Close() error {
|
|
return v.r.Close()
|
|
}
|
|
|
|
func init() {
|
|
temp, err := os.CreateTemp(os.TempDir(), "bao")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
err = temp.Chmod(1755)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
_, err = io.Copy(temp, bytes.NewReader(pluginBin))
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
err = temp.Close()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
clientInst := plugin.NewClient(&plugin.ClientConfig{
|
|
HandshakeConfig: plugin.HandshakeConfig{
|
|
ProtocolVersion: 1,
|
|
},
|
|
Plugins: plugin.PluginSet{
|
|
"bao": &BaoPlugin{},
|
|
},
|
|
Cmd: exec.Command(temp.Name()),
|
|
AllowedProtocols: []plugin.Protocol{plugin.ProtocolGRPC},
|
|
})
|
|
|
|
rpcClient, err := clientInst.Client()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
pluginInst, err := rpcClient.Dispense("bao")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
bao = pluginInst.(Bao)
|
|
}
|
|
|
|
func Shutdown() {
|
|
client.Kill()
|
|
}
|
|
|
|
func Hash(r io.Reader) (*Result, error) {
|
|
hasherId := bao.NewHasher()
|
|
initialSize := 4 * units.KiB
|
|
maxSize := 3.5 * units.MiB
|
|
bufSize := initialSize
|
|
|
|
reader := bufio.NewReaderSize(r, bufSize)
|
|
var totalReadSize int
|
|
|
|
buf := make([]byte, bufSize)
|
|
for {
|
|
|
|
n, err := reader.Read(buf)
|
|
if err != nil {
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
return nil, err
|
|
}
|
|
totalReadSize += n
|
|
|
|
if !bao.Hash(hasherId, buf[:n]) {
|
|
return nil, errors.New("hashing failed")
|
|
}
|
|
|
|
// Adaptively adjust buffer size based on read patterns
|
|
if n == bufSize && float64(bufSize) < maxSize {
|
|
// If buffer was fully used, consider increasing buffer size
|
|
bufSize = int(math.Min(float64(bufSize*2), float64(maxSize))) // Double the buffer size, up to a maximum
|
|
buf = make([]byte, bufSize) // Apply new buffer size
|
|
reader = bufio.NewReaderSize(r, bufSize) // Apply new buffer size
|
|
}
|
|
}
|
|
|
|
result := bao.Finish(hasherId)
|
|
result.Length = uint(totalReadSize)
|
|
|
|
return &result, nil
|
|
}
|
|
|
|
func NewVerifier(r io.ReadCloser, proof Result) *Verifier {
|
|
return &Verifier{r: r, proof: proof}
|
|
}
|