125 lines
2.2 KiB
Go
125 lines
2.2 KiB
Go
package utils
|
|
|
|
import (
|
|
"fmt"
|
|
"io"
|
|
|
|
"github.com/aws/aws-sdk-go/aws"
|
|
"github.com/aws/aws-sdk-go/service/s3"
|
|
)
|
|
|
|
const (
|
|
bufferSize = 4 * 1024 * 1024 // 4MB
|
|
)
|
|
|
|
type S3ReadSeeker struct {
|
|
client *s3.S3
|
|
bucket string
|
|
key string
|
|
offset int64
|
|
body io.ReadCloser
|
|
contentLen int64
|
|
buffer []byte
|
|
bufLen int
|
|
bufPos int
|
|
}
|
|
|
|
func NewS3ReadSeeker(client *s3.S3, bucket, key string, body io.ReadCloser, contentLen int64) *S3ReadSeeker {
|
|
return &S3ReadSeeker{
|
|
client: client,
|
|
bucket: bucket,
|
|
key: key,
|
|
body: body,
|
|
contentLen: contentLen,
|
|
buffer: make([]byte, bufferSize),
|
|
}
|
|
}
|
|
|
|
func (r *S3ReadSeeker) Read(p []byte) (n int, err error) {
|
|
if r.bufPos < r.bufLen {
|
|
n = copy(p, r.buffer[r.bufPos:r.bufLen])
|
|
r.bufPos += n
|
|
return n, nil
|
|
}
|
|
|
|
if r.body == nil {
|
|
return 0, io.EOF
|
|
}
|
|
|
|
if len(p) > bufferSize {
|
|
n, err = r.body.Read(p)
|
|
r.offset += int64(n)
|
|
return n, err
|
|
}
|
|
|
|
r.bufLen, err = r.body.Read(r.buffer)
|
|
if err != nil && err != io.EOF {
|
|
return 0, err
|
|
}
|
|
|
|
if r.bufLen > 0 {
|
|
r.bufPos = 0
|
|
n = copy(p, r.buffer[:r.bufLen])
|
|
r.bufPos = n
|
|
r.offset += int64(n)
|
|
return n, nil
|
|
}
|
|
|
|
return 0, io.EOF
|
|
}
|
|
|
|
func (r *S3ReadSeeker) Seek(offset int64, whence int) (int64, error) {
|
|
var newOffset int64
|
|
switch whence {
|
|
case io.SeekStart:
|
|
newOffset = offset
|
|
case io.SeekCurrent:
|
|
newOffset = r.offset + offset
|
|
case io.SeekEnd:
|
|
newOffset = r.contentLen + offset
|
|
default:
|
|
return 0, fmt.Errorf("invalid whence")
|
|
}
|
|
|
|
if newOffset < 0 {
|
|
return 0, fmt.Errorf("negative offset")
|
|
}
|
|
|
|
if newOffset >= r.contentLen {
|
|
return r.contentLen, nil
|
|
}
|
|
|
|
if newOffset >= r.offset-int64(r.bufLen) && newOffset < r.offset {
|
|
r.bufPos = int(newOffset - (r.offset - int64(r.bufLen)))
|
|
r.offset = newOffset
|
|
return newOffset, nil
|
|
}
|
|
|
|
if r.body != nil {
|
|
r.body.Close()
|
|
r.body = nil
|
|
}
|
|
r.bufLen = 0
|
|
r.bufPos = 0
|
|
|
|
result, err := r.client.GetObject(&s3.GetObjectInput{
|
|
Bucket: aws.String(r.bucket),
|
|
Key: aws.String(r.key),
|
|
Range: aws.String(fmt.Sprintf("bytes=%d-", newOffset)),
|
|
})
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
r.body = result.Body
|
|
r.offset = newOffset
|
|
return newOffset, nil
|
|
}
|
|
|
|
func (r *S3ReadSeeker) Close() error {
|
|
if r.body != nil {
|
|
return r.body.Close()
|
|
}
|
|
return nil
|
|
}
|