92 lines
2.0 KiB
Go
92 lines
2.0 KiB
Go
package service
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"sync"
|
|
|
|
"github.com/aws/aws-sdk-go/aws"
|
|
"github.com/aws/aws-sdk-go/aws/credentials"
|
|
"github.com/aws/aws-sdk-go/aws/session"
|
|
"github.com/aws/aws-sdk-go/service/s3"
|
|
"tencent_ocr/pkg/errors"
|
|
)
|
|
|
|
type UploadService struct {
|
|
s3Client *s3.S3
|
|
bucket string
|
|
customDomain string
|
|
mu sync.RWMutex
|
|
}
|
|
|
|
func NewUploadService(accessKey, secretKey, bucket, endpoint, customDomain string) (*UploadService, error) {
|
|
sess, err := session.NewSession(&aws.Config{
|
|
Endpoint: aws.String(endpoint),
|
|
Region: aws.String("auto"),
|
|
Credentials: credentials.NewStaticCredentials(accessKey, secretKey, ""),
|
|
})
|
|
if err != nil {
|
|
return nil, errors.Wrap(err, "failed to create S3 session")
|
|
}
|
|
|
|
return &UploadService{
|
|
s3Client: s3.New(sess),
|
|
bucket: bucket,
|
|
customDomain: customDomain,
|
|
}, nil
|
|
}
|
|
|
|
func (s *UploadService) UploadFile(file io.Reader, fileName, contentType string) (string, error) {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
|
|
// Read file content
|
|
fileBytes, err := io.ReadAll(file)
|
|
if err != nil {
|
|
return "", errors.Wrap(err, "failed to read file")
|
|
}
|
|
|
|
// Upload file to R2
|
|
_, err = s.s3Client.PutObject(&s3.PutObjectInput{
|
|
Bucket: aws.String(s.bucket),
|
|
Key: aws.String(fileName),
|
|
Body: bytes.NewReader(fileBytes),
|
|
ContentType: aws.String(contentType),
|
|
ACL: aws.String("public-read"),
|
|
})
|
|
if err != nil {
|
|
return "", errors.Wrap(err, "failed to upload file to R2")
|
|
}
|
|
|
|
// Generate file URL
|
|
imageURL := fmt.Sprintf("https://%s/%s", s.customDomain, fileName)
|
|
return imageURL, nil
|
|
}
|
|
|
|
func (s *UploadService) IsValidFileType(contentType string) bool {
|
|
allowedTypes := []string{
|
|
"image/jpeg",
|
|
"image/png",
|
|
"image/gif",
|
|
"image/bmp",
|
|
"image/tiff",
|
|
"image/webp",
|
|
}
|
|
for _, t := range allowedTypes {
|
|
if contentType == t {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// Close implements graceful shutdown
|
|
func (s *UploadService) Close() error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
// Add any cleanup logic here if needed
|
|
return nil
|
|
}
|