youp0m/backend_s3.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"image"
	"image/jpeg"
	"io"
	"log"
	"net/http"
	"os"
	"path"

	"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"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
)

var (
	s3_endpoint   string
	s3_region     = "us"
	s3_bucket     string
	s3_access_key string
	s3_secret_key string
	s3_path_style bool
)

func init() {
	existing_backends = append(existing_backends, "s3")

	if endpoint, ok := os.LookupEnv("S3_ENDPOINT"); ok {
		backend = "s3"
		s3_endpoint = endpoint
	}
	if region, ok := os.LookupEnv("S3_REGION"); ok {
		backend = "s3"
		s3_region = region
	}
	s3_bucket, _ = os.LookupEnv("S3_BUCKET")
	s3_access_key, _ = os.LookupEnv("S3_ACCESS_KEY")
	s3_secret_key, _ = os.LookupEnv("S3_SECRET_KEY")
	if path_style, ok := os.LookupEnv("S3_PATH_STYLE"); ok {
		s3_path_style = path_style == "1" || path_style == "ON" || path_style == "on" || path_style == "TRUE" || path_style == "true" || path_style == "yes" || path_style == "YES"
	}

	flag.StringVar(&s3_endpoint, "s3-endpoint", s3_endpoint, "When using S3 backend, endpoint to use")
	flag.StringVar(&s3_region, "s3-region", s3_region, "When using S3 backend, region to use")
	flag.StringVar(&s3_bucket, "s3-bucket", s3_bucket, "When using S3 backend, bucket to use")
	flag.StringVar(&s3_access_key, "s3-access-key", s3_access_key, "When using S3 backend, Access Key")
	flag.StringVar(&s3_secret_key, "s3-secret-key", s3_secret_key, "When using S3 backend, Secret Key")
	flag.BoolVar(&s3_path_style, "s3-path-style", s3_path_style, "When using S3 backend, force path style (when using minio)")
}

type S3FileBackend struct {
	Endpoint  string
	Region    string
	Bucket    string
	AccessKey string
	SecretKey string
	PathStyle bool
	BaseDir   string
}

func (l *S3FileBackend) getPath(box, name string) string {
	if l.BaseDir == "" {
		return path.Join(box, name+".jpg")
	} else {
		return path.Join(l.BaseDir, box, name+".jpg")
	}
}

func (l *S3FileBackend) newSession() (*session.Session, error) {
	return session.NewSession(&aws.Config{
		Credentials:      credentials.NewStaticCredentials(l.AccessKey, l.SecretKey, ""),
		Endpoint:         &l.Endpoint,
		Region:           &l.Region,
		S3ForcePathStyle: &l.PathStyle,
	})
}

func (l *S3FileBackend) DeletePicture(box, name string) error {
	s, err := l.newSession()
	if err != nil {
		return err
	}

	log.Println(l.getPath(box, name))

	input := &s3.DeleteObjectsInput{
		Bucket: aws.String(l.Bucket),
		Delete: &s3.Delete{
			Objects: []*s3.ObjectIdentifier{
				{
					Key: aws.String(l.getPath(box, name)),
				},
			},
		},
	}

	_, err = s3.New(s).DeleteObjects(input)
	return err
}

func (l *S3FileBackend) ServeFile() http.Handler {
	return l
}

func (l *S3FileBackend) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	s, err := l.newSession()
	if err != nil {
		return
	}

	result, err := s3.New(s).GetObject(&s3.GetObjectInput{
		Bucket: aws.String(l.Bucket),
		Key:    aws.String(r.URL.Path),
	})
	if err != nil {
		return
	}

	io.Copy(w, result.Body)
}

func (l *S3FileBackend) createPictureInfo(local_path, name string, file *s3.Object) *Picture {
	return &Picture{
		local_path,                           // Path
		name,                                 // Basename
		name[:len(name)-len(path.Ext(name))], // Sanitized filename
		*file.LastModified,                   // UploadTime
	}
}

func (l *S3FileBackend) GetPicture(box, name string, w io.Writer) error {
	s, err := l.newSession()
	if err != nil {
		return err
	}

	s3_path := l.getPath(box, name)

	input := &s3.GetObjectInput{
		Bucket: aws.String(l.Bucket),
		Key:    aws.String(s3_path),
	}

	result, err := s3.New(s).GetObject(input)
	if err != nil {
		return err
	}

	_, err = io.Copy(w, result.Body)

	return err
}

func (l *S3FileBackend) GetPictureInfo(box, name string) (*Picture, error) {
	s, err := l.newSession()
	if err != nil {
		return nil, err
	}

	input := &s3.ListObjectsInput{
		Bucket: aws.String(l.Bucket),
		Prefix: aws.String(l.getPath(box, name)),
	}

	result, err := s3.New(s).ListObjects(input)
	if err != nil {
		return nil, err
	}

	var pictures []*Picture
	for _, file := range result.Contents {
		pictures = append(pictures, l.createPictureInfo(*file.Key, path.Base(*file.Key), file))
		break
	}

	if len(pictures) == 0 {
		return nil, fmt.Errorf("Object not found")
	}

	return pictures[0], nil
}

func (l *S3FileBackend) ListPictures(box string) ([]*Picture, error) {
	s, err := l.newSession()
	if err != nil {
		return nil, err
	}

	input := &s3.ListObjectsInput{
		Bucket: aws.String(l.Bucket),
	}

	if l.BaseDir == "" {
		input.Prefix = aws.String(box)
	} else {
		input.Prefix = aws.String(path.Join(l.BaseDir, box))
	}

	result, err := s3.New(s).ListObjects(input)
	if err != nil {
		return nil, err
	}

	pictures := make([]*Picture, 0)
	for _, file := range result.Contents {
		pictures = append(pictures, l.createPictureInfo(path.Join(box, *file.Key), path.Base(*file.Key), file))
	}

	return pictures, nil
}

func (l *S3FileBackend) MovePicture(box_from, box_to, name string) error {
	s, err := l.newSession()
	if err != nil {
		return err
	}

	log.Println(l.getPath(box_from, name), l.getPath(box_to, name))

	_, err = s3.New(s).CopyObject(&s3.CopyObjectInput{
		Bucket:     aws.String(l.Bucket),
		CopySource: aws.String(path.Join("", l.Bucket, l.getPath(box_from, name))),
		Key:        aws.String(l.getPath(box_to, name)),
	})
	if err != nil {
		return err
	}

	log.Println(l.getPath(box_from, name))

	input := &s3.DeleteObjectsInput{
		Bucket: aws.String(l.Bucket),
		Delete: &s3.Delete{
			Objects: []*s3.ObjectIdentifier{
				{
					Key: aws.String(l.getPath(box_from, name)),
				},
			},
		},
	}

	_, err = s3.New(s).DeleteObjects(input)
	return err
}

func (l *S3FileBackend) PutPicture(box, name string, img *image.Image) error {
	s, err := l.newSession()
	if err != nil {
		return err
	}

	bbuf := new(bytes.Buffer)
	err = jpeg.Encode(bbuf, *img, nil)
	if err != nil {
		return err
	}

	_, err = s3manager.NewUploader(s).Upload(&s3manager.UploadInput{
		Bucket: aws.String(l.Bucket),
		ACL:    aws.String("public-read"),
		Key:    aws.String(l.getPath(box, name)),
		Body:   bbuf,
	})

	return err
}