110 lines
2.6 KiB
Go
110 lines
2.6 KiB
Go
package marian
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
)
|
|
|
|
// RemoteFile is a model file to download, checked against its size and
|
|
// SHA-256 digest.
|
|
type RemoteFile struct {
|
|
Name string
|
|
URL string
|
|
Size int64
|
|
SHA256 string
|
|
}
|
|
|
|
// Installed reports whether dir holds every file of a model.
|
|
func Installed(dir string) bool {
|
|
for _, f := range Files {
|
|
if st, err := os.Stat(filepath.Join(dir, f)); err != nil || st.Size() == 0 {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
// Download fetches files into dir. Each file is written next to its
|
|
// final name and renamed once its digest is checked, so an interrupted
|
|
// download never leaves a half-written model behind. progress receives
|
|
// the bytes downloaded so far and the total.
|
|
func Download(ctx context.Context, client *http.Client, dir string, files []RemoteFile, progress func(done, total int64)) error {
|
|
if err := os.MkdirAll(dir, 0o755); err != nil {
|
|
return err
|
|
}
|
|
var total, done int64
|
|
for _, f := range files {
|
|
total += f.Size
|
|
}
|
|
for _, f := range files {
|
|
n, err := download(ctx, client, dir, f, func(n int64) {
|
|
if progress != nil {
|
|
progress(done+n, total)
|
|
}
|
|
})
|
|
if err != nil {
|
|
return fmt.Errorf("downloading %s: %w", f.Name, err)
|
|
}
|
|
done += n
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func download(ctx context.Context, client *http.Client, dir string, f RemoteFile, progress func(int64)) (int64, error) {
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, f.URL, nil)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode != http.StatusOK {
|
|
return 0, fmt.Errorf("server answered %s", resp.Status)
|
|
}
|
|
tmp, err := os.CreateTemp(dir, "."+f.Name+".*")
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
defer os.Remove(tmp.Name()) // no-op once renamed
|
|
hash := sha256.New()
|
|
w := &countingWriter{w: io.MultiWriter(tmp, hash), progress: progress}
|
|
_, err = io.Copy(w, io.LimitReader(resp.Body, f.Size+1))
|
|
if cerr := tmp.Close(); err == nil {
|
|
err = cerr
|
|
}
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
if w.n != f.Size {
|
|
return 0, fmt.Errorf("got %d bytes, expected %d", w.n, f.Size)
|
|
}
|
|
if sum := hex.EncodeToString(hash.Sum(nil)); sum != f.SHA256 {
|
|
return 0, fmt.Errorf("checksum mismatch")
|
|
}
|
|
if err := os.Chmod(tmp.Name(), 0o644); err != nil {
|
|
return 0, err
|
|
}
|
|
return w.n, os.Rename(tmp.Name(), filepath.Join(dir, f.Name))
|
|
}
|
|
|
|
type countingWriter struct {
|
|
w io.Writer
|
|
n int64
|
|
progress func(int64)
|
|
}
|
|
|
|
func (c *countingWriter) Write(p []byte) (int, error) {
|
|
n, err := c.w.Write(p)
|
|
c.n += int64(n)
|
|
c.progress(c.n)
|
|
return n, err
|
|
}
|