feat: refactor S3 storage implementation and reduce binary size

This commit is contained in:
krau
2025-12-18 16:21:40 +08:00
parent c0b4580e34
commit 8752dd865c
5 changed files with 220 additions and 114 deletions

21
go.mod
View File

@@ -3,10 +3,6 @@ module github.com/krau/SaveAny-Bot
go 1.24.0
require (
github.com/aws/aws-sdk-go-v2 v1.40.1
github.com/aws/aws-sdk-go-v2/config v1.32.3
github.com/aws/aws-sdk-go-v2/credentials v1.19.3
github.com/aws/aws-sdk-go-v2/service/s3 v1.93.0
github.com/blang/semver v3.5.1+incompatible
github.com/celestix/gotgproto v1.0.0-beta22
github.com/cenkalti/backoff/v4 v4.3.0
@@ -31,20 +27,9 @@ require (
require (
github.com/AnimeKaizoku/cacher v1.0.3 // indirect
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.4 // indirect
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.15 // indirect
github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.15 // indirect
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.15 // indirect
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.4 // indirect
github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.15 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.4 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.6 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.15 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.15 // indirect
github.com/aws/aws-sdk-go-v2/service/signin v1.0.3 // indirect
github.com/aws/aws-sdk-go-v2/service/sso v1.30.6 // indirect
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.11 // indirect
github.com/aws/aws-sdk-go-v2/service/sts v1.41.3 // indirect
github.com/aws/aws-sdk-go-v2 v1.40.1 // indirect
github.com/aws/aws-sdk-go-v2/credentials v1.19.3 // indirect
github.com/aws/aws-sdk-go-v2/service/s3 v1.93.0 // indirect
github.com/aws/smithy-go v1.24.0 // indirect
github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect

14
go.sum
View File

@@ -8,20 +8,14 @@ github.com/aws/aws-sdk-go-v2 v1.40.1 h1:difXb4maDZkRH0x//Qkwcfpdg1XQVXEAEs2DdXld
github.com/aws/aws-sdk-go-v2 v1.40.1/go.mod h1:MayyLB8y+buD9hZqkCW3kX1AKq07Y5pXxtgB+rRFhz0=
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.4 h1:489krEF9xIGkOaaX3CE/Be2uWjiXrkCH6gUX+bZA/BU=
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.4/go.mod h1:IOAPF6oT9KCsceNTvvYMNHy0+kMF8akOjeDvPENWxp4=
github.com/aws/aws-sdk-go-v2/config v1.32.3 h1:cpz7H2uMNTDa0h/5CYL5dLUEzPSLo2g0NkbxTRJtSSU=
github.com/aws/aws-sdk-go-v2/config v1.32.3/go.mod h1:srtPKaJJe3McW6T/+GMBZyIPc+SeqJsNPJsd4mOYZ6s=
github.com/aws/aws-sdk-go-v2/credentials v1.19.3 h1:01Ym72hK43hjwDeJUfi1l2oYLXBAOR8gNSZNmXmvuas=
github.com/aws/aws-sdk-go-v2/credentials v1.19.3/go.mod h1:55nWF/Sr9Zvls0bGnWkRxUdhzKqj9uRNlPvgV1vgxKc=
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.15 h1:utxLraaifrSBkeyII9mIbVwXXWrZdlPO7FIKmyLCEcY=
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.15/go.mod h1:hW6zjYUDQwfz3icf4g2O41PHi77u10oAzJ84iSzR/lo=
github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.17.75 h1:S61/E3N01oral6B3y9hZ2E1iFDqCZPPOBoBQretCnBI=
github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.17.75/go.mod h1:bDMQbkI1vJbNjnvJYpPTSNYBkI/VIv18ngWb/K84tkk=
github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.15 h1:Y5YXgygXwDI5P4RkteB5yF7v35neH7LfJKBG+hzIons=
github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.15/go.mod h1:K+/1EpG42dFSY7CBj+Fruzm8PsCGWTXJ3jdeJ659oGQ=
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.15 h1:AvltKnW9ewxX2hFmQS0FyJH93aSvJVUEFvXfU+HWtSE=
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.15/go.mod h1:3I4oCdZdmgrREhU74qS1dK9yZ62yumob+58AbFR4cQA=
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.4 h1:WKuaxf++XKWlHWu9ECbMlha8WOEGm0OUEZqm4K/Gcfk=
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.4/go.mod h1:ZWy7j6v1vWGmPReu0iSGvRiise4YI5SkR3OHKTZ6Wuc=
github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.15 h1:NLYTEyZmVZo0Qh183sC8nC+ydJXOOeIL/qI/sS3PdLY=
github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.15/go.mod h1:Z803iB3B0bc8oJV8zH2PERLRfQUJ2n2BXISpsA4+O1M=
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.4 h1:0ryTNEdJbzUCEWkVXEXoqlXV72J5keC1GvILMOuD00E=
@@ -34,14 +28,6 @@ github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.15 h1:wsSQ4SVz5YE1c
github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.15/go.mod h1:I7sditnFGtYMIqPRU1QoHZAUrXkGp4SczmlLwrNPlD0=
github.com/aws/aws-sdk-go-v2/service/s3 v1.93.0 h1:IrbE3B8O9pm3lsg96AXIN5MXX4pECEuExh/A0Du3AuI=
github.com/aws/aws-sdk-go-v2/service/s3 v1.93.0/go.mod h1:/sJLzHtiiZvs6C1RbxS/anSAFwZD6oC6M/kotQzOiLw=
github.com/aws/aws-sdk-go-v2/service/signin v1.0.3 h1:d/6xOGIllc/XW1lzG9a4AUBMmpLA9PXcQnVPTuHHcik=
github.com/aws/aws-sdk-go-v2/service/signin v1.0.3/go.mod h1:fQ7E7Qj9GiW8y0ClD7cUJk3Bz5Iw8wZkWDHsTe8vDKs=
github.com/aws/aws-sdk-go-v2/service/sso v1.30.6 h1:8sTTiw+9yuNXcfWeqKF2x01GqCF49CpP4Z9nKrrk/ts=
github.com/aws/aws-sdk-go-v2/service/sso v1.30.6/go.mod h1:8WYg+Y40Sn3X2hioaaWAAIngndR8n1XFdRPPX+7QBaM=
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.11 h1:E+KqWoVsSrj1tJ6I/fjDIu5xoS2Zacuu1zT+H7KtiIk=
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.11/go.mod h1:qyWHz+4lvkXcr3+PoGlGHEI+3DLLiU6/GdrFfMaAhB0=
github.com/aws/aws-sdk-go-v2/service/sts v1.41.3 h1:tzMkjh0yTChUqJDgGkcDdxvZDSrJ/WB6R6ymI5ehqJI=
github.com/aws/aws-sdk-go-v2/service/sts v1.41.3/go.mod h1:T270C0R5sZNLbWUe8ueiAF42XSZxxPocTaGSgs5c/60=
github.com/aws/smithy-go v1.24.0 h1:LpilSUItNPFr1eY85RYgTIg5eIEPtvFbskaFcmmIUnk=
github.com/aws/smithy-go v1.24.0/go.mod h1:LEj2LM3rBRQJxPZTB4KuzZkaZYnZPnvgIhb4pu07mx0=
github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=

178
storage/s3/client.go Normal file
View File

@@ -0,0 +1,178 @@
package s3
import (
"context"
"crypto/hmac"
"crypto/sha256"
"encoding/hex"
"fmt"
"io"
"net/http"
"net/url"
"sort"
"strings"
"time"
storconfig "github.com/krau/SaveAny-Bot/config/storage"
)
type Client struct {
endpoint string
region string
bucket string
accessKey string
secretKey string
httpClient *http.Client
pathStyle bool
}
func NewClient(cfg storconfig.S3StorageConfig) (*Client, error) {
endpoint := cfg.Endpoint
if !strings.HasPrefix(endpoint, "http") {
if cfg.UseSSL {
endpoint = "https://" + endpoint
} else {
endpoint = "http://" + endpoint
}
}
return &Client{
endpoint: endpoint,
region: cfg.Region,
bucket: cfg.BucketName,
accessKey: cfg.AccessKeyID,
secretKey: cfg.SecretAccessKey,
pathStyle: !cfg.VirtualHost,
httpClient: http.DefaultClient,
}, nil
}
func (c *Client) HeadBucket(ctx context.Context) error {
url := c.buildURL("")
req, _ := http.NewRequestWithContext(ctx, "HEAD", url, nil)
signRequest(req, c.region, c.accessKey, c.secretKey, hashSHA256(nil))
resp, err := c.httpClient.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode >= 300 {
return fmt.Errorf("head bucket failed: %s", resp.Status)
}
return nil
}
func (c *Client) Exists(ctx context.Context, key string) bool {
req, _ := http.NewRequestWithContext(ctx, "HEAD", c.buildURL(key), nil)
signRequest(req, c.region, c.accessKey, c.secretKey, hashSHA256(nil))
resp, err := c.httpClient.Do(req)
if err != nil {
return false
}
defer resp.Body.Close()
return resp.StatusCode == http.StatusOK
}
func (c *Client) Put(ctx context.Context, key string, r io.Reader, size int64) error {
req, _ := http.NewRequestWithContext(ctx, "PUT", c.buildURL(key), r)
if size >= 0 {
req.ContentLength = size
}
signRequest(req, c.region, c.accessKey, c.secretKey, "UNSIGNED-PAYLOAD")
resp, err := c.httpClient.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode >= 300 {
return fmt.Errorf("put object failed: %s", resp.Status)
}
return nil
}
func (c *Client) buildURL(key string) string {
if c.pathStyle {
return fmt.Sprintf("%s/%s/%s", c.endpoint, c.bucket, key)
}
u, _ := url.Parse(c.endpoint)
u.Host = c.bucket + "." + u.Host
u.Path = "/" + key
return u.String()
}
func hmacSHA256(key []byte, data string) []byte {
h := hmac.New(sha256.New, key)
h.Write([]byte(data))
return h.Sum(nil)
}
func hashSHA256(data []byte) string {
sum := sha256.Sum256(data)
return hex.EncodeToString(sum[:])
}
func signRequest(req *http.Request, region, accessKey, secretKey string, payloadHash string) error {
now := time.Now().UTC()
amzDate := now.Format("20060102T150405Z")
date := now.Format("20060102")
req.Header.Set("x-amz-date", amzDate)
req.Header.Set("x-amz-content-sha256", payloadHash)
// Canonical headers
var headers []string
for k := range req.Header {
headers = append(headers, strings.ToLower(k))
}
sort.Strings(headers)
var canonicalHeaders strings.Builder
for _, k := range headers {
canonicalHeaders.WriteString(k)
canonicalHeaders.WriteString(":")
canonicalHeaders.WriteString(strings.TrimSpace(req.Header.Get(k)))
canonicalHeaders.WriteString("\n")
}
signedHeaders := strings.Join(headers, ";")
canonicalRequest := strings.Join([]string{
req.Method,
req.URL.EscapedPath(),
req.URL.RawQuery,
canonicalHeaders.String(),
signedHeaders,
payloadHash,
}, "\n")
scope := fmt.Sprintf("%s/%s/s3/aws4_request", date, region)
stringToSign := strings.Join([]string{
"AWS4-HMAC-SHA256",
amzDate,
scope,
hashSHA256([]byte(canonicalRequest)),
}, "\n")
kDate := hmacSHA256([]byte("AWS4"+secretKey), date)
kRegion := hmacSHA256(kDate, region)
kService := hmacSHA256(kRegion, "s3")
kSigning := hmacSHA256(kService, "aws4_request")
signature := hex.EncodeToString(hmacSHA256(kSigning, stringToSign))
auth := fmt.Sprintf(
"AWS4-HMAC-SHA256 Credential=%s/%s, SignedHeaders=%s, Signature=%s",
accessKey, scope, signedHeaders, signature,
)
req.Header.Set("Authorization", auth)
return nil
}

View File

@@ -4,14 +4,9 @@ import (
"context"
"fmt"
"io"
"net/url"
"path"
"strings"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/config"
"github.com/aws/aws-sdk-go-v2/credentials"
"github.com/aws/aws-sdk-go-v2/service/s3"
"github.com/charmbracelet/log"
storconfig "github.com/krau/SaveAny-Bot/config/storage"
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
@@ -21,76 +16,30 @@ import (
type S3 struct {
config storconfig.S3StorageConfig
client *s3.Client
client *Client
logger *log.Logger
}
func (m *S3) Init(ctx context.Context, cfg storconfig.StorageConfig) error {
s3Config, ok := cfg.(*storconfig.S3StorageConfig)
s3cfg, ok := cfg.(*storconfig.S3StorageConfig)
if !ok {
return fmt.Errorf("failed to cast s3 config")
}
if err := s3Config.Validate(); err != nil {
if err := s3cfg.Validate(); err != nil {
return err
}
m.config = *s3Config
m.config = *s3cfg
m.logger = log.FromContext(ctx).WithPrefix(fmt.Sprintf("s3[%s]", m.config.Name))
loadOpts := make([]config.LoadOptionsFunc, 0)
if m.config.Region != "" {
loadOpts = append(loadOpts, config.WithRegion(m.config.Region))
}
if endpoint := m.config.Endpoint; endpoint != "" {
if !strings.HasPrefix(endpoint, "http://") && !strings.HasPrefix(endpoint, "https://") {
if m.config.UseSSL {
endpoint = "https://" + endpoint
} else {
endpoint = "http://" + endpoint
}
}
if _, err := url.Parse(endpoint); err != nil {
return fmt.Errorf("invalid s3 endpoint %q: %w", m.config.Endpoint, err)
}
loadOpts = append(loadOpts, config.WithBaseEndpoint(endpoint))
}
loadOpts = append(loadOpts, config.WithCredentialsProvider(
credentials.NewStaticCredentialsProvider(
m.config.AccessKeyID,
m.config.SecretAccessKey,
"",
),
))
awsCfg, err := config.LoadDefaultConfig(
ctx,
func() []func(*config.LoadOptions) error {
// wtf aws sdk
// https://github.com/aws/aws-sdk-go-v2/issues/2193
funcs := make([]func(*config.LoadOptions) error, 0)
for _, fn := range loadOpts {
funcs = append(funcs, fn)
}
return funcs
}()...,
)
client, err := NewClient(m.config)
if err != nil {
return fmt.Errorf("failed to load AWS config: %w", err)
return fmt.Errorf("failed to create s3 client: %w", err)
}
m.client = s3.NewFromConfig(awsCfg, func(o *s3.Options) {
// Path style: https://s3.amazonaws.com/mybucket/path/to/file.jpg
// virtual hosted style: https://mybucket.s3.amazonaws.com/path/to/file.jpg
o.UsePathStyle = !m.config.VirtualHost
})
m.client = client
// Check if bucket exists
_, err = m.client.HeadBucket(ctx, &s3.HeadBucketInput{
Bucket: aws.String(m.config.BucketName),
})
if err != nil {
if err := m.client.HeadBucket(ctx); err != nil {
return fmt.Errorf("bucket %s not accessible: %w", m.config.BucketName, err)
}
return nil
}
@@ -131,18 +80,7 @@ func (m *S3) Save(ctx context.Context, r io.Reader, storagePath string) error {
}
}
// S3 PutObject needs either size or StreamingBody
input := &s3.PutObjectInput{
Bucket: aws.String(m.config.BucketName),
Key: aws.String(candidate),
Body: r,
}
if size >= 0 {
input.ContentLength = &size
}
_, err := m.client.PutObject(ctx, input)
err := m.client.Put(ctx, candidate, r, size)
if err != nil {
return fmt.Errorf("failed to upload file to S3: %w", err)
}
@@ -153,10 +91,5 @@ func (m *S3) Save(ctx context.Context, r io.Reader, storagePath string) error {
func (m *S3) Exists(ctx context.Context, storagePath string) bool {
m.logger.Debugf("Checking if file exists at %s", storagePath)
_, err := m.client.HeadObject(ctx, &s3.HeadObjectInput{
Bucket: aws.String(m.config.BucketName),
Key: aws.String(storagePath),
})
return err == nil
return m.client.Exists(ctx, storagePath)
}

View File

@@ -1,4 +1,4 @@
package s3
package s3_test
import (
"bytes"
@@ -10,6 +10,8 @@ import (
"github.com/johannesboyne/gofakes3"
"github.com/johannesboyne/gofakes3/backend/s3mem"
storconfig "github.com/krau/SaveAny-Bot/config/storage"
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
"github.com/krau/SaveAny-Bot/storage/s3"
)
func newTestContext(t *testing.T) context.Context {
@@ -19,7 +21,7 @@ func newTestContext(t *testing.T) context.Context {
return log.WithContext(ctx, logger)
}
func newFakeS3(t *testing.T) (*S3, *storconfig.S3StorageConfig) {
func newFakeS3(t *testing.T) (*s3.S3, *storconfig.S3StorageConfig) {
t.Helper()
backend := s3mem.New()
@@ -45,7 +47,7 @@ func newFakeS3(t *testing.T) (*S3, *storconfig.S3StorageConfig) {
t.Fatalf("failed to create fake bucket: %v", err)
}
s := &S3{}
s := &s3.S3{}
ctx := newTestContext(t)
if err := s.Init(ctx, cfg); err != nil {
t.Fatalf("init s3 failed: %v", err)
@@ -54,9 +56,9 @@ func newFakeS3(t *testing.T) (*S3, *storconfig.S3StorageConfig) {
return s, cfg
}
func TestS3_SaveAndExists(t *testing.T) {
func TestS3(t *testing.T) {
s, _ := newFakeS3(t)
ctx := context.Background()
ctx := t.Context()
content := []byte("hello world")
reader := bytes.NewReader(content)
@@ -69,4 +71,26 @@ func TestS3_SaveAndExists(t *testing.T) {
if !s.Exists(ctx, key) {
t.Fatalf("Exists should return true for saved key")
}
if s.Exists(ctx, "nonexistent.txt") {
t.Fatalf("Exists should return false for nonexistent key")
}
if err := s.Save(ctx, bytes.NewReader(content), key); err != nil {
t.Fatalf("Save with existing key failed: %v", err)
}
if !s.Exists(ctx, "foo/bar_1.txt") {
t.Fatalf("Exists should return true for unique renamed key")
}
var length int64 = int64(len(content))
ctx = context.WithValue(ctx, ctxkey.ContentLength, length)
if err := s.Save(ctx, bytes.NewReader(content), "size_test.txt"); err != nil {
t.Fatalf("Save with content length failed: %v", err)
}
if !s.Exists(ctx, "size_test.txt") {
t.Fatalf("Exists should return true for size_test.txt")
}
}