GDrive provider support

README.md

@@ -2,7 +2,7 @@
 
 Easy and fast file sharing from the command-line. This code contains the server with everything you need to create your own instance.
 
-Transfer.sh support currently the s3 (Amazon S3) provider and local file system (local).
+Transfer.sh support currently the s3 (Amazon S3) and gdrive (Google Drive) providers and local file system (local).
 
 ## Usage
 
@@ -46,12 +46,13 @@ tls-cert-file | path to tls certificate | |
 tls-private-key | path to tls private key | |
 temp-path | path to temp folder | system temp |
 web-path | path to static web files (for development) | |
-provider | which storage provider to use | (s3 or local) |
+provider | which storage provider to use | (s3, grdrive or local) |
 aws-access-key | aws access key | | AWS_ACCESS_KEY
 aws-secret-key | aws access key | | AWS_SECRET_KEY
 bucket | aws bucket | | BUCKET
-basedir | path storage for local provider| | 
-lets-encrypt-hosts | hosts to use for lets encrypt certificates (comma seperated) | | 
+basedir | path storage for local/gdrive provider| |
+gdrive-client-json-filepath | path to client json config for gdrive provider| |
+lets-encrypt-hosts | hosts to use for lets encrypt certificates (comma seperated) | |
 log | path to log file| | 
 
 If you want to use TLS using lets encrypt certificates, set lets-encrypt-hosts to your domain, set tls-listener to :443 and enable force-https.

cmd/cmd.go

@@ -7,7 +7,7 @@ import (
 
 	"strings"
 
-	"github.com/dutchcoders/transfer.sh/server"
+	"transfer.sh/server"
 	"github.com/fatih/color"
 	"github.com/minio/cli"
 )
@@ -74,7 +74,7 @@ var globalFlags = []cli.Flag{
 	},
 	cli.StringFlag{
 		Name:  "provider",
-		Usage: "s3|local",
+		Usage: "s3|gdrive|local",
 		Value: "",
 	},
 	cli.StringFlag{
@@ -101,6 +101,13 @@ var globalFlags = []cli.Flag{
 		Value:  "",
 		EnvVar: "BUCKET",
 	},
+	cli.StringFlag{
+		Name:   "gdrive-client-json-filepath",
+		Usage:  "",
+		Value:  "",
+		EnvVar: "",
+	},
+
 	cli.IntFlag{
 		Name:   "rate-limit",
 		Usage:  "requests per minute",
@@ -233,6 +240,16 @@ func New() *Cmd {
 			} else {
 				options = append(options, server.UseStorage(storage))
 			}
+		case "gdrive":
+			if clientJsonFilepath := c.String("gdrive-client-json-filepath"); clientJsonFilepath == "" {
+				panic("client-json-filepath not set.")
+			} else if basedir := c.String("basedir"); basedir == "" {
+				panic("basedir not set.")
+			} else if storage, err := server.NewGDriveStorage(clientJsonFilepath, basedir); err != nil {
+				panic(err)
+			} else {
+				options = append(options, server.UseStorage(storage))
+			}
 		case "local":
 			if v := c.String("basedir"); v == "" {
 				panic("basedir not set.")

main.go

@@ -1,6 +1,6 @@
 package main
 
-import "github.com/dutchcoders/transfer.sh/cmd"
+import "transfer.sh/cmd"
 
 func main() {
 	app := cmd.New()

server/storage.go

@@ -12,6 +12,16 @@ import (
 	"sync"
 
 	"github.com/goamz/goamz/s3"
+	"encoding/json"
+
+	"golang.org/x/oauth2"
+	"golang.org/x/net/context"
+	"golang.org/x/oauth2/google"
+	"google.golang.org/api/drive/v3"
+	"google.golang.org/api/googleapi"
+	"net/http"
+	"io/ioutil"
+	"time"
 )
 
 type Storage interface {
@@ -284,3 +294,386 @@ func (s *S3Storage) Put(token string, filename string, reader io.Reader, content
 
 	return
 }
+
+type GDrive struct {
+	service *drive.Service
+	basedir string
+}
+
+func NewGDriveStorage(clientJsonFilepath string, basedir string) (*GDrive, error) {
+	b, err := ioutil.ReadFile(clientJsonFilepath)
+	if err != nil {
+		return nil, err
+	}
+
+	// If modifying these scopes, delete your previously saved client_secret.json.
+	config, err := google.ConfigFromJSON(b, drive.DriveScope, drive.DriveMetadataScope)
+	if err != nil {
+		return nil, err
+	}
+
+	srv, err := drive.New(getGDriveClient(config))
+	if err != nil {
+		return nil, err
+	}
+
+	return &GDrive{service: srv, basedir: basedir}, nil
+}
+
+const GDriveTimeoutTimerInterval = time.Second * 10
+const GDriveDirectoryMimeType = "application/vnd.google-apps.folder"
+
+type gDriveTimeoutReaderWrapper func(io.Reader) io.Reader
+
+func (s *GDrive) getTimeoutReader(r io.Reader, cancel context.CancelFunc, timeout time.Duration) io.Reader {
+	return &GDriveTimeoutReader{
+		reader:         r,
+		cancel:         cancel,
+		mutex:          &sync.Mutex{},
+		maxIdleTimeout: timeout,
+	}
+}
+
+type GDriveTimeoutReader struct {
+	reader         io.Reader
+	cancel         context.CancelFunc
+	lastActivity   time.Time
+	timer          *time.Timer
+	mutex          *sync.Mutex
+	maxIdleTimeout time.Duration
+	done           bool
+}
+
+func (r *GDriveTimeoutReader) Read(p []byte) (int, error) {
+	if r.timer == nil {
+		r.startTimer()
+	}
+
+	r.mutex.Lock()
+
+	// Read
+	n, err := r.reader.Read(p)
+
+	r.lastActivity = time.Now()
+	r.done = (err != nil)
+
+	r.mutex.Unlock()
+
+	if r.done {
+		r.stopTimer()
+	}
+
+	return n, err
+}
+
+func (r *GDriveTimeoutReader) Close() error {
+	return r.reader.(io.ReadCloser).Close()
+}
+
+func (r *GDriveTimeoutReader) startTimer() {
+	r.mutex.Lock()
+	defer r.mutex.Unlock()
+
+	if !r.done {
+		r.timer = time.AfterFunc(GDriveTimeoutTimerInterval, r.timeout)
+	}
+}
+
+func (r *GDriveTimeoutReader) stopTimer() {
+	r.mutex.Lock()
+	defer r.mutex.Unlock()
+
+	if r.timer != nil {
+		r.timer.Stop()
+	}
+}
+
+func (r *GDriveTimeoutReader) timeout() {
+	r.mutex.Lock()
+
+	if r.done {
+		r.mutex.Unlock()
+		return
+	}
+
+	if time.Since(r.lastActivity) > r.maxIdleTimeout {
+		r.cancel()
+		r.mutex.Unlock()
+		return
+	}
+
+	r.mutex.Unlock()
+	r.startTimer()
+}
+
+func (s *GDrive) getTimeoutReaderWrapperContext(timeout time.Duration) (gDriveTimeoutReaderWrapper, context.Context) {
+	ctx, cancel := context.WithCancel(context.TODO())
+	wrapper := func(r io.Reader) io.Reader {
+		// Return untouched reader if timeout is 0
+		if timeout == 0 {
+			return r
+		}
+
+		return s.getTimeoutReader(r, cancel, timeout)
+	}
+	return wrapper, ctx
+}
+
+func (s *GDrive) hasChecksum(f *drive.File) bool {
+	return f.Md5Checksum != ""
+}
+
+func (s *GDrive) list(nextPageToken string, q string) (*drive.FileList, error){
+	return s.service.Files.List().Fields("nextPageToken, files(id, name, mimeType)").Q(q).PageToken(nextPageToken).Do()
+}
+
+func (s *GDrive) findId(filename string, token string) (string, error) {
+	fileId, rootId, tokenId, nextPageToken := "", "", "", ""
+
+	q := fmt.Sprintf("name='%s' and trashed=false", s.basedir)
+	l, err := s.list(nextPageToken, q)
+	for 0 < len(l.Files) {
+		if err != nil {
+			return "", err
+		}
+
+		for _, fi := range l.Files {
+			rootId = fi.Id
+			break
+		}
+
+		if l.NextPageToken == "" {
+			break
+		}
+
+		l, err = s.list(l.NextPageToken, q)
+	}
+
+	if token == "" {
+		if rootId == "" {
+			return "", fmt.Errorf("Cannot find file %s/%s", token, filename)
+		}
+
+		return rootId, nil
+	}
+
+	q = fmt.Sprintf("'%s' in parents and name='%s' and mimeType='%s' and trashed=false", rootId, token, GDriveDirectoryMimeType)
+	l, err = s.list(nextPageToken, q)
+	for 0 < len(l.Files) {
+		if err != nil {
+			return "", err
+		}
+
+		for _, fi := range l.Files {
+			tokenId = fi.Id
+			break
+		}
+
+		if l.NextPageToken == "" {
+			break
+		}
+
+		l, err = s.list(l.NextPageToken, q)
+	}
+
+	if filename == "" {
+		return tokenId, nil
+	} else if tokenId == "" {
+		return "", fmt.Errorf("Cannot find file %s/%s", token, filename)
+	}
+
+	q = fmt.Sprintf("'%s' in parents and name='%s' and mimeType!='%s' and trashed=false", tokenId, filename, GDriveDirectoryMimeType)
+	l, err = s.list(nextPageToken, q)
+
+	for 0 < len(l.Files) {
+		if err != nil {
+			return "", err
+		}
+
+		for _, fi := range l.Files {
+
+			fileId = fi.Id
+			break
+		}
+
+		if l.NextPageToken == "" {
+			break
+		}
+
+		l, err = s.list(l.NextPageToken, q)
+	}
+
+
+	if fileId == "" {
+		return "", fmt.Errorf("Cannot find file %s/%s", token, filename)
+	}
+
+	return fileId, nil
+}
+
+func (s *GDrive) Type() string {
+	return "gdrive"
+}
+
+func (s *GDrive) Head(token string, filename string) (contentType string, contentLength uint64, err error) {
+	var fileId string
+	fileId, err = s.findId(filename, token)
+	if err != nil {
+		return
+	}
+
+	var fi *drive.File
+	if fi, err = s.service.Files.Get(fileId).Fields("mimeType", "size").Do(); err != nil {
+		return
+	}
+
+	contentLength = uint64(fi.Size)
+
+	contentType = mime.TypeByExtension(fi.MimeType)
+
+	return
+}
+
+func (s *GDrive) Get(token string, filename string) (reader io.ReadCloser, contentType string, contentLength uint64, err error) {
+	var fileId string
+	fileId, err = s.findId(filename, token)
+	if err != nil {
+		return
+	}
+
+	var fi *drive.File
+	fi, err = s.service.Files.Get(fileId).Fields("mimeType", "size", "md5Checksum").Do()
+	if !s.hasChecksum(fi) {
+		return
+	}
+
+
+	contentLength = uint64(fi.Size)
+
+	contentType = mime.TypeByExtension(fi.MimeType)
+
+	// Get timeout reader wrapper and context
+	timeoutReaderWrapper, ctx := s.getTimeoutReaderWrapperContext(time.Duration(10))
+
+	var res *http.Response
+	res, err = s.service.Files.Get(fileId).Context(ctx).Download()
+	if err != nil {
+		return
+	}
+
+	reader = timeoutReaderWrapper(res.Body).(io.ReadCloser)
+
+	return
+}
+
+func (s *GDrive) IsNotExist(err error) bool {
+	if err == nil {
+		return false
+	}
+
+	if err != nil {
+		if e, ok := err.(*googleapi.Error); ok {
+			return e.Code == http.StatusNotFound
+		}
+	}
+
+	return false
+}
+
+func (s *GDrive) Put(token string, filename string, reader io.Reader, contentType string, contentLength uint64) error {
+	rootId, err := s.findId("", "")
+	if err != nil {
+		return err
+	}
+
+	dirId, err := s.findId("", token)
+	if err != nil {
+		return err
+	}
+
+
+	if dirId == "" {
+		dir := &drive.File{
+			Name:        token,
+			Parents: 	 []string{rootId},
+			MimeType:    GDriveDirectoryMimeType,
+		}
+
+		di, err := s.service.Files.Create(dir).Fields("id").Do()
+		if err != nil {
+			return err
+		}
+
+		dirId = di.Id
+	}
+
+	// Wrap reader in timeout reader
+	timeoutReaderWrapper, ctx := s.getTimeoutReaderWrapperContext(time.Duration(10))
+
+	// Instantiate empty drive file
+	dst := &drive.File{
+		Name: filename,
+		Parents: []string{dirId},
+		MimeType: contentType,
+	}
+
+	_, err = s.service.Files.Create(dst).Context(ctx).Media(timeoutReaderWrapper(reader)).Do()
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+
+// Retrieve a token, saves the token, then returns the generated client.
+func getGDriveClient(config *oauth2.Config) *http.Client {
+	tokenFile := "token.json"
+	tok, err := gDriveTokenFromFile(tokenFile)
+	if err != nil {
+		tok = getGDriveTokenFromWeb(config)
+		saveGDriveToken(tokenFile, tok)
+	}
+	return config.Client(context.Background(), tok)
+}
+
+// Request a token from the web, then returns the retrieved token.
+func getGDriveTokenFromWeb(config *oauth2.Config) *oauth2.Token {
+	authURL := config.AuthCodeURL("state-token", oauth2.AccessTypeOffline)
+	fmt.Printf("Go to the following link in your browser then type the "+
+		"authorization code: \n%v\n", authURL)
+
+	var authCode string
+	if _, err := fmt.Scan(&authCode); err != nil {
+		log.Fatalf("Unable to read authorization code %v", err)
+	}
+
+	tok, err := config.Exchange(oauth2.NoContext, authCode)
+	if err != nil {
+		log.Fatalf("Unable to retrieve token from web %v", err)
+	}
+	return tok
+}
+
+// Retrieves a token from a local file.
+func gDriveTokenFromFile(file string) (*oauth2.Token, error) {
+	f, err := os.Open(file)
+	defer f.Close()
+	if err != nil {
+		return nil, err
+	}
+	tok := &oauth2.Token{}
+	err = json.NewDecoder(f).Decode(tok)
+	return tok, err
+}
+
+// Saves a token to a file path.
+func saveGDriveToken(path string, token *oauth2.Token) {
+	fmt.Printf("Saving credential file to: %s\n", path)
+	f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0600)
+	defer f.Close()
+	if err != nil {
+		log.Fatalf("Unable to cache oauth token: %v", err)
+	}
+	json.NewEncoder(f).Encode(token)
+}