diff --git a/pkg/client/client.go b/pkg/client/client.go index 9bca690..885d45b 100644 --- a/pkg/client/client.go +++ b/pkg/client/client.go @@ -1,11 +1,13 @@ package client import ( + "context" "fmt" "io" "io/ioutil" "net/http" "os" + "time" ) // bufferSize is the size of the buffer for disk I/O. @@ -70,3 +72,63 @@ func GetFileWithLogin(url, user, pass, destination string) error { return nil } + +// GetFileWithAuthAndTimeout downloads a file from URL to destination using basic auth. +// It will check if the connection dropped to 0 bytes during the timeout and, in that case, cancels the download. +func GetFileWithAuthAndTimeout(url, user, pass, path string, timeout time.Duration) error { + ctx, cancel := context.WithCancel(context.Background()) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) // Will be added in Go 1.13 + if err != nil { + return fmt.Errorf("error creating request: %s", err) + } + req.SetBasicAuth(user, pass) + + f, err := os.Create(path) + if err != nil { + return fmt.Errorf("error creating file in path \"%s\": %s", path, err) + } + defer f.Close() + + go checkFileSizeChanging(path, timeout, cancel) + + resp, err := HttpClient.Do(req) + if err != nil { + return fmt.Errorf("download request failed: %s", err) + } + defer resp.Body.Close() + + if _, err = io.CopyBuffer(f, resp.Body, make([]byte, bufferSize)); err != nil { + // If failed, force close and try to remove it + _ = f.Close() + _ = os.Remove(path) + return fmt.Errorf("error while saving file in path \"%s\": %s", path, err) + } + + if err = f.Close(); err != nil { + return fmt.Errorf("error found closing file in path \"%s\": %s", path, err) + } + + return nil +} + +// checkFileSizeChanging will check if the size of the file from the path provided have changed every timeout duration, +// if it haven't, it calls cancel() and exits. +func checkFileSizeChanging(path string, timeout time.Duration, cancel func()) { + var size int64 = -1 + + for { + stat, err := os.Stat(path) + if err != nil { + cancel() + return + } + + if stat.Size() == size { + cancel() + return + } + + size = stat.Size() + time.Sleep(timeout) + } +}