dearwolf/LocalAI
1
0
Fork 0
mirror of https://github.com/mudler/LocalAI.git synced 2025-06-29 22:20:43 +00:00
Code Issues Projects Releases Packages Wiki Activity

feat(resume download): implement range header check

Browse source 
Signed-off-by: Saarthak Verma <saarthakverma739@gmail.com>
This commit is contained in:
Saarthak Verma 2025-01-08 23:09:07 +05:30
parent fe8472c863
commit 7a46b416e5
No known key found for this signature in database
GPG key ID: 2BEFA9EF09A35950
1 changed files with 28 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

44
pkg/downloader/uri.go
View file

@ -215,6 +215,16 @@ func calculateHashForPartialFile(file *os.File) (hash.Hash, error) {
return hash, nil return hash, nil
} }
func (uri URI) checkSeverSupportsRangeHeader() (bool, error) {
url := uri.ResolveURL()
resp, err := http.Head(url)
if err != nil {
return false, err
}
defer resp.Body.Close()
return resp.Header.Get("Accept-Ranges") == "bytes", nil
}
func (uri URI) DownloadFile(filePath, sha string, fileN, total int, downloadStatus func(string, string, string, float64)) error { func (uri URI) DownloadFile(filePath, sha string, fileN, total int, downloadStatus func(string, string, string, float64)) error {
url := uri.ResolveURL() url := uri.ResolveURL()
if uri.LooksLikeOCI() { if uri.LooksLikeOCI() {
@ -282,21 +292,23 @@ func (uri URI) DownloadFile(filePath, sha string, fileN, total int, downloadStat
return fmt.Errorf("failed to create request for %q: %v", filePath, err) return fmt.Errorf("failed to create request for %q: %v", filePath, err)
} }
/* TODO
* 1. ~~Mock downloads~~
* 2. Check if server supports partial downloads
* 3. ~~Resume partial downloads~~
* 4. Ensure progressWriter accurately reflects progress if a partial file is present
* 5. MAYBE:
* a. Delete file if calculatedSHA != sha
*/
// save partial download to dedicated file // save partial download to dedicated file
tmpFilePath := filePath + ".partial" tmpFilePath := filePath + ".partial"
tmpFileInfo, err := os.Stat(tmpFilePath) tmpFileInfo, err := os.Stat(tmpFilePath)
if err == nil { if err == nil {
startPos := tmpFileInfo.Size() support, err := uri.checkSeverSupportsRangeHeader()
req.Header.Set("Range", fmt.Sprintf("bytes=%d-", startPos)) if err != nil {
return fmt.Errorf("failed to check if uri server supports range header: %v", err)
}
if support {
startPos := tmpFileInfo.Size()
req.Header.Set("Range", fmt.Sprintf("bytes=%d-", startPos))
} else {
err := removePartialFile(tmpFilePath)
if err != nil {
return err
}
}
} else if !errors.Is(err, os.ErrNotExist) { } else if !errors.Is(err, os.ErrNotExist) {
return fmt.Errorf("failed to check file %q existence: %v", filePath, err) return fmt.Errorf("failed to check file %q existence: %v", filePath, err)
} }
@ -341,6 +353,11 @@ func (uri URI) DownloadFile(filePath, sha string, fileN, total int, downloadStat
return fmt.Errorf("failed to write file %q: %v", filePath, err) return fmt.Errorf("failed to write file %q: %v", filePath, err)
} }
err = os.Rename(tmpFilePath, filePath)
if err != nil {
return fmt.Errorf("failed to rename temporary file %s -> %s: %v", tmpFilePath, filePath, err)
}
if sha != "" { if sha != "" {
// Verify SHA // Verify SHA
calculatedSHA := fmt.Sprintf("%x", progress.hash.Sum(nil)) calculatedSHA := fmt.Sprintf("%x", progress.hash.Sum(nil))
@ -352,11 +369,6 @@ func (uri URI) DownloadFile(filePath, sha string, fileN, total int, downloadStat
log.Debug().Msgf("SHA missing for %q. Skipping validation", filePath) log.Debug().Msgf("SHA missing for %q. Skipping validation", filePath)
} }
err = os.Rename(tmpFilePath, filePath)
if err != nil {
return fmt.Errorf("failed to rename temporary file %s -> %s: %v", tmpFilePath, filePath, err)
}
log.Info().Msgf("File %q downloaded and verified", filePath) log.Info().Msgf("File %q downloaded and verified", filePath)
if utils.IsArchive(filePath) { if utils.IsArchive(filePath) {
basePath := filepath.Dir(filePath) basePath := filepath.Dir(filePath)