133 lines
3.5 KiB
Go
133 lines
3.5 KiB
Go
// Package objectstore talks to S3-compatible object storage (RustFS) via AWS SDK v2.
|
|
package objectstore
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/aws/aws-sdk-go-v2/aws"
|
|
"github.com/aws/aws-sdk-go-v2/credentials"
|
|
"github.com/aws/aws-sdk-go-v2/service/s3"
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
const (
|
|
defaultRegion = "us-east-1"
|
|
defaultTimeout = 15 * time.Second
|
|
)
|
|
|
|
// Store is a RustFS / S3 client. Nil or !Ready means object storage is not configured.
|
|
type Store struct {
|
|
client *s3.Client
|
|
presign *s3.PresignClient
|
|
bucket string
|
|
log *zap.Logger
|
|
}
|
|
|
|
// Options configure Open. Empty Endpoint skips the connection.
|
|
type Options struct {
|
|
Endpoint string
|
|
PublicEndpoint string
|
|
AccessKey string
|
|
SecretKey string
|
|
Bucket string
|
|
Region string
|
|
UseSSL bool
|
|
ChatExpireDays int
|
|
Log *zap.Logger
|
|
}
|
|
|
|
func Open(o *Options) (*Store, error) {
|
|
if o == nil || strings.TrimSpace(o.Endpoint) == "" {
|
|
return &Store{}, nil
|
|
}
|
|
client, presignBase, err := newClients(o)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
bucket := strings.TrimSpace(o.Bucket)
|
|
if bucket == "" {
|
|
bucket = "spark"
|
|
}
|
|
s := &Store{client: client, presign: s3.NewPresignClient(presignBase), bucket: bucket, log: o.Log}
|
|
if s.log == nil {
|
|
s.log = zap.NewNop()
|
|
}
|
|
region := strings.TrimSpace(o.Region)
|
|
if region == "" {
|
|
region = defaultRegion
|
|
}
|
|
ctx, cancel := context.WithTimeout(context.Background(), defaultTimeout)
|
|
defer cancel()
|
|
if err := s.ensureBucket(ctx, region, o.ChatExpireDays); err != nil {
|
|
return nil, err
|
|
}
|
|
return s, nil
|
|
}
|
|
|
|
func (s *Store) Ready() bool {
|
|
return s != nil && s.client != nil && s.presign != nil && s.bucket != ""
|
|
}
|
|
|
|
func (s *Store) Ping(ctx context.Context) error {
|
|
if !s.Ready() {
|
|
return nil
|
|
}
|
|
_, err := s.client.HeadBucket(ctx, &s3.HeadBucketInput{Bucket: aws.String(s.bucket)})
|
|
return err
|
|
}
|
|
|
|
func newClients(o *Options) (*s3.Client, *s3.Client, error) {
|
|
region := strings.TrimSpace(o.Region)
|
|
if region == "" {
|
|
region = defaultRegion
|
|
}
|
|
host, secure, err := parseEndpoint(o.Endpoint, o.UseSSL)
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
client := newS3Client(region, endpointURL(host, secure), o.AccessKey, o.SecretKey)
|
|
if pub := strings.TrimSpace(o.PublicEndpoint); pub != "" {
|
|
phost, psecure, err := parseEndpoint(pub, true)
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
return client, newS3Client(region, endpointURL(phost, psecure), o.AccessKey, o.SecretKey), nil
|
|
}
|
|
return client, client, nil
|
|
}
|
|
|
|
func newS3Client(region, baseURL, access, secret string) *s3.Client {
|
|
return s3.New(s3.Options{
|
|
Region: region,
|
|
Credentials: credentials.NewStaticCredentialsProvider(access, secret, ""),
|
|
BaseEndpoint: aws.String(baseURL),
|
|
UsePathStyle: true,
|
|
RequestChecksumCalculation: aws.RequestChecksumCalculationWhenRequired,
|
|
ResponseChecksumValidation: aws.ResponseChecksumValidationWhenRequired,
|
|
})
|
|
}
|
|
|
|
func endpointURL(host string, secure bool) string {
|
|
if secure {
|
|
return "https://" + host
|
|
}
|
|
return "http://" + host
|
|
}
|
|
|
|
func parseEndpoint(raw string, useSSL bool) (host string, secure bool, err error) {
|
|
raw = strings.TrimSpace(strings.TrimRight(raw, "/"))
|
|
switch {
|
|
case raw == "":
|
|
return "", false, fmt.Errorf("objectstore: empty endpoint")
|
|
case strings.HasPrefix(raw, "https://"):
|
|
return strings.TrimPrefix(raw, "https://"), true, nil
|
|
case strings.HasPrefix(raw, "http://"):
|
|
return strings.TrimPrefix(raw, "http://"), false, nil
|
|
default:
|
|
return raw, useSSL, nil
|
|
}
|
|
}
|