From 1917722b3a9dd6621ac9d0b59ad34d057434e01b Mon Sep 17 00:00:00 2001
From: namhq1989 <namhq.1989@gmail.com>
Date: Thu, 1 Dec 2022 09:14:15 +0700
Subject: [PATCH] fix download google file not found or in private mode

---
 google_drive.go | 37 +++++++++++++++++++++++++++----------
 util.go         | 13 +++++++++++++
 2 files changed, 40 insertions(+), 10 deletions(-)
 create mode 100644 util.go

diff --git a/google_drive.go b/google_drive.go
index 28f19c6..7c2c417 100644
--- a/google_drive.go
+++ b/google_drive.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"net/http"
 	"os"
+	"strings"
 
 	"github.com/tanaikech/go-getfilelist"
 	"google.golang.org/api/drive/v3"
@@ -84,19 +85,14 @@ func (g GoogleDrive) GetFilesByFolderID(folderID string, mimeType []string) ([]G
 	return result, nil
 }
 
-func (g GoogleDrive) DownloadByFileID(id, name, location string) (err error) {
-	fmt.Printf("[google drive] start download file %s \n", name)
-	defer fmt.Printf("[google drive] done download file %s \n", name)
+// DownloadByFileID download file into specific path
+func (g GoogleDrive) DownloadByFileID(id, path string) (err error) {
+	// trim "/" in the end of path
+	path = strings.TrimSuffix(path, "/")
 
+	// prepare url
 	url := fmt.Sprintf("https://drive.google.com/uc?export=download&id=%s", id)
 
-	// Create the file
-	out, err := os.Create(location)
-	if err != nil {
-		return err
-	}
-	defer out.Close()
-
 	// Get the data
 	resp, err := http.Get(url)
 	if err != nil {
@@ -109,6 +105,27 @@ func (g GoogleDrive) DownloadByFileID(id, name, location string) (err error) {
 		return fmt.Errorf("bad status: %s", resp.Status)
 	}
 
+	// get file name
+	filename := getFileNameFromHeaderContentDisposition(resp.Header.Get("Content-Disposition"))
+
+	// return err if not found
+	if filename == "" {
+		return errors.New("file is not existed or is in private mode, please check again")
+	}
+
+	// assign location with field name
+	location := fmt.Sprintf("%s/%s", path, filename)
+
+	// Create the file
+	out, err := os.Create(location)
+	if err != nil {
+		return err
+	}
+	defer out.Close()
+
+	fmt.Printf("[google drive] start download file %s with id %s \n", filename, id)
+	defer fmt.Printf("[google drive] done download file %s with id %s \n", filename, id)
+
 	// Writer the body to file
 	_, err = io.Copy(out, resp.Body)
 	if err != nil {
diff --git a/util.go b/util.go
new file mode 100644
index 0000000..55c89ac
--- /dev/null
+++ b/util.go
@@ -0,0 +1,13 @@
+package download
+
+import (
+	"mime"
+)
+
+func getFileNameFromHeaderContentDisposition(value string) string {
+	disposition, params, err := mime.ParseMediaType(value)
+	if err != nil || disposition != "attachment" {
+		return ""
+	}
+	return params["filename"]
+}