github.com/vektra/tachyon@v0.0.0-20150921164542-0da4f3861aef/net/s3/s3.go (about)

     1  package s3
     2  
     3  import (
     4  	"bytes"
     5  	"compress/gzip"
     6  	"crypto/md5"
     7  	"encoding/hex"
     8  	"fmt"
     9  	"github.com/crowdmob/goamz/aws"
    10  	"github.com/crowdmob/goamz/s3"
    11  	"github.com/vektra/tachyon"
    12  	"io"
    13  	"os"
    14  	"time"
    15  )
    16  
    17  type S3 struct {
    18  	Bucket      string `tachyon:"bucket,required"`
    19  	PutFile     string `tachyon:"put_file"`
    20  	GetFile     string `tachyon:"get_file"`
    21  	At          string `tachyon:"at"`
    22  	Public      bool   `tachyon:"public"`
    23  	ContentType string `tachyon:"content_type"`
    24  	Writable    bool   `tachyon:"writable"`
    25  	GZip        bool   `tachyon:"gzip"`
    26  }
    27  
    28  func (s *S3) Run(env *tachyon.CommandEnv) (*tachyon.Result, error) {
    29  	auth, err := aws.GetAuth("", "", "", time.Time{})
    30  	if err != nil {
    31  		return nil, err
    32  	}
    33  
    34  	c := s3.New(auth, aws.USWest2)
    35  	b := c.Bucket(s.Bucket)
    36  
    37  	res := tachyon.NewResult(true)
    38  
    39  	res.Add("bucket", s.Bucket)
    40  	res.Add("remote", s.At)
    41  
    42  	if s.PutFile != "" {
    43  		path := env.Paths.File(s.PutFile)
    44  
    45  		f, err := os.Open(path)
    46  		if err != nil {
    47  			return nil, err
    48  		}
    49  
    50  		if f == nil {
    51  			return nil, fmt.Errorf("Unknown local file %s", s.PutFile)
    52  		}
    53  
    54  		defer f.Close()
    55  
    56  		var perm s3.ACL
    57  
    58  		if s.Public {
    59  			if s.Writable {
    60  				perm = s3.PublicReadWrite
    61  			} else {
    62  				perm = s3.PublicRead
    63  			}
    64  		} else {
    65  			perm = s3.Private
    66  		}
    67  
    68  		ct := s.ContentType
    69  		if ct == "" {
    70  			ct = "application/octet-stream"
    71  		}
    72  
    73  		fi, err := f.Stat()
    74  		if err != nil {
    75  			return nil, err
    76  		}
    77  
    78  		var (
    79  			input io.Reader
    80  			opts  s3.Options
    81  			size  int64
    82  		)
    83  
    84  		h := md5.New()
    85  
    86  		if s.GZip {
    87  			var buf bytes.Buffer
    88  
    89  			z := gzip.NewWriter(io.MultiWriter(h, &buf))
    90  
    91  			_, err = io.Copy(z, f)
    92  			if err != nil {
    93  				return nil, err
    94  			}
    95  
    96  			z.Close()
    97  
    98  			opts.ContentEncoding = "gzip"
    99  
   100  			input = &buf
   101  			size = int64(buf.Len())
   102  		} else {
   103  			input = io.TeeReader(f, h)
   104  			size = fi.Size()
   105  		}
   106  
   107  		err = b.PutReader(s.At, input, size, ct, perm, opts)
   108  
   109  		rep, err := b.Head(s.At, nil)
   110  		if err != nil {
   111  			return nil, err
   112  		}
   113  
   114  		localMD5 := hex.EncodeToString(h.Sum(nil))
   115  
   116  		res.Add("wrote", size)
   117  		res.Add("local", s.PutFile)
   118  		res.Add("md5", localMD5)
   119  
   120  		etag := rep.Header.Get("ETag")
   121  		if etag != "" {
   122  			etag = etag[1 : len(etag)-1]
   123  
   124  			if localMD5 != etag {
   125  				return nil, fmt.Errorf("corruption uploading file detected")
   126  			}
   127  		}
   128  
   129  	} else if s.GetFile != "" {
   130  		f, err := os.OpenFile(s.GetFile, os.O_CREATE|os.O_WRONLY, 0644)
   131  		if err != nil {
   132  			return nil, err
   133  		}
   134  
   135  		defer f.Close()
   136  
   137  		i, err := b.GetReader(s.At)
   138  		if err != nil {
   139  			return nil, err
   140  		}
   141  
   142  		defer i.Close()
   143  
   144  		n, err := io.Copy(f, i)
   145  		if err != nil {
   146  			return nil, err
   147  		}
   148  
   149  		res.Add("read", n)
   150  		res.Add("local", s.GetFile)
   151  	} else {
   152  		return nil, fmt.Errorf("Specify put_file or get_file")
   153  	}
   154  
   155  	return res, nil
   156  }
   157  
   158  func init() {
   159  	tachyon.RegisterCommand("s3", &S3{})
   160  }