154 lines
3.5 KiB
Go
154 lines
3.5 KiB
Go
package artifacts
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/base64"
|
|
"encoding/hex"
|
|
"encoding/xml"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"sync"
|
|
)
|
|
|
|
type PullResult struct {
|
|
SHA256Sum string
|
|
Size int64
|
|
}
|
|
type PushResult struct {
|
|
SHA256Sum string
|
|
Size int64
|
|
}
|
|
|
|
func PushBlob(ctx context.Context, url string, data io.Reader, chunkSize, maxParallel int) (int, string, error) {
|
|
var empty string
|
|
if chunkSize <= 0 {
|
|
chunkSize = 4 * 1024 * 1024 // 4MB default
|
|
}
|
|
if maxParallel <= 0 {
|
|
maxParallel = 4
|
|
}
|
|
|
|
type chunk struct {
|
|
index int
|
|
data []byte
|
|
}
|
|
wrkinp := make(chan chunk)
|
|
wrkerr := make(chan error, maxParallel)
|
|
blockIDs := make([]string, 0)
|
|
var blockMu sync.Mutex
|
|
var wg sync.WaitGroup
|
|
|
|
// Worker pool for parallel uploading
|
|
for i := 0; i < maxParallel; i++ {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
for c := range wrkinp {
|
|
id := base64.StdEncoding.EncodeToString([]byte(fmt.Sprintf("%06d", c.index)))
|
|
address := fmt.Sprintf("%s&comp=block&blockid=%s", url, id)
|
|
req, err := http.NewRequestWithContext(ctx, "PUT", address, bytes.NewReader(c.data))
|
|
if err != nil {
|
|
wrkerr <- err
|
|
return
|
|
}
|
|
req.ContentLength = int64(len(c.data))
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
wrkerr <- err
|
|
return
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode != http.StatusCreated {
|
|
body, _ := io.ReadAll(resp.Body)
|
|
wrkerr <- fmt.Errorf("chunk %d upload failed: %s body: %s", c.index, resp.Status, string(body))
|
|
return
|
|
}
|
|
|
|
blockMu.Lock()
|
|
defer blockMu.Unlock()
|
|
blockIDs = append(blockIDs, id)
|
|
}
|
|
}()
|
|
}
|
|
|
|
// Chunk the data and send to workers
|
|
index := 0
|
|
hash := sha256.New()
|
|
for {
|
|
buf := make([]byte, chunkSize)
|
|
n, err := io.ReadFull(data, buf)
|
|
if err == io.EOF || err == io.ErrUnexpectedEOF {
|
|
if n > 0 {
|
|
hash.Write(buf[:n])
|
|
wrkinp <- chunk{index: index, data: buf[:n]}
|
|
index++
|
|
}
|
|
break
|
|
}
|
|
if err != nil {
|
|
close(wrkinp)
|
|
wg.Wait()
|
|
return empty, err
|
|
}
|
|
hash.Write(buf)
|
|
wrkinp <- chunk{index: index, data: buf}
|
|
index++
|
|
}
|
|
close(wrkinp)
|
|
wg.Wait()
|
|
|
|
select {
|
|
case err := <-wrkerr:
|
|
return empty, err
|
|
default:
|
|
}
|
|
|
|
type BlockList struct {
|
|
Latest []string `xml:"Latest"`
|
|
}
|
|
|
|
type Latest struct {
|
|
Value string `xml:",chardata"`
|
|
}
|
|
|
|
// Sort block IDs to ensure correct order
|
|
blockMu.Lock()
|
|
defer blockMu.Unlock()
|
|
if len(blockIDs) == 0 {
|
|
return empty, fmt.Errorf("no blocks uploaded")
|
|
}
|
|
blockList := &BlockList{}
|
|
for i := 0; i < len(blockIDs); i++ {
|
|
blockList.Latest = append(blockList.Blocks, xmlBlock{XMLName: xml.Name{Local: "Latest"}, ID: blockIDs[i]})
|
|
}
|
|
blockListData, err := xml.Marshal(blockList)
|
|
if err != nil {
|
|
return empty, err
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, "PUT", fmt.Sprintf("%s&comp=blocklist", url), bytes.NewReader(blockListData))
|
|
if err != nil {
|
|
return empty, err
|
|
}
|
|
req.Header.Set("Content-Type", "application/xml")
|
|
req.ContentLength = int64(len(blockListData))
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
return empty, err
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode != http.StatusCreated {
|
|
body, _ := io.ReadAll(resp.Body)
|
|
return empty, fmt.Errorf("blocklist commit failed: %s body: %s", resp.Status, string(body))
|
|
}
|
|
return hex.EncodeToString(hash.Sum(nil)), nil
|
|
}
|
|
|
|
func PushBlob_(ctx context.Context, r io.Reader, blobURL string, chunkSize int64, parallelism int) (*PushResult, error) {
|
|
}
|
|
|
|
func PullBlob_(ctx context.Context, w io.Writer, blobURL string, chunkSize int64, parallelism int) (*PullResult, error) {
|
|
}
|