fix(drive): validate download formats

This commit is contained in:
salmonumbrella 2026-01-30 19:15:37 -08:00 committed by Peter Steinberger
parent eb0bc8e695
commit 0c4978d2f7
3 changed files with 53 additions and 3 deletions

View File

@ -1052,6 +1052,14 @@ func guessMimeType(path string) string {
func downloadDriveFile(ctx context.Context, svc *drive.Service, meta *drive.File, destPath string, format string) (string, int64, error) {
isGoogleDoc := strings.HasPrefix(meta.MimeType, "application/vnd.google-apps.")
normalizedFormat := strings.ToLower(strings.TrimSpace(format))
if normalizedFormat == "auto" {
normalizedFormat = ""
}
if !isGoogleDoc && normalizedFormat != "" {
return "", 0, fmt.Errorf("--format %q not supported for non-Google Workspace files (mimeType=%q); file can only be downloaded as-is", format, meta.MimeType)
}
var (
resp *http.Response
@ -1061,11 +1069,11 @@ func downloadDriveFile(ctx context.Context, svc *drive.Service, meta *drive.File
if isGoogleDoc {
var exportMimeType string
if strings.TrimSpace(format) == "" {
if normalizedFormat == "" {
exportMimeType = driveExportMimeType(meta.MimeType)
} else {
var mimeErr error
exportMimeType, mimeErr = driveExportMimeTypeForFormat(meta.MimeType, format)
exportMimeType, mimeErr = driveExportMimeTypeForFormat(meta.MimeType, normalizedFormat)
if mimeErr != nil {
return "", 0, mimeErr
}
@ -1129,7 +1137,7 @@ func driveExportMimeType(googleMimeType string) string {
func driveExportMimeTypeForFormat(googleMimeType string, format string) (string, error) {
format = strings.ToLower(strings.TrimSpace(format))
if format == "" {
if format == "" || format == "auto" {
return driveExportMimeType(googleMimeType), nil
}

View File

@ -57,6 +57,36 @@ func TestDownloadDriveFile_NonGoogleDoc(t *testing.T) {
}
}
func TestDownloadDriveFile_NonGoogleDocFormatRejected(t *testing.T) {
origDownload := driveDownload
t.Cleanup(func() { driveDownload = origDownload })
called := false
driveDownload = func(context.Context, *drive.Service, string) (*http.Response, error) {
called = true
return &http.Response{
Status: "200 OK",
StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader("ok")),
}, nil
}
dest := filepath.Join(t.TempDir(), "file.html")
_, _, err := downloadDriveFile(context.Background(), &drive.Service{}, &drive.File{Id: "id1", MimeType: "application/pdf"}, dest, "html")
if err == nil {
t.Fatalf("expected error")
}
if !strings.Contains(err.Error(), "non-Google Workspace") {
t.Fatalf("unexpected error: %v", err)
}
if called {
t.Fatalf("download should not be called on format error")
}
if _, statErr := os.Stat(dest); !os.IsNotExist(statErr) {
t.Fatalf("expected no file written, stat=%v", statErr)
}
}
func TestDownloadDriveFile_GoogleDocExport(t *testing.T) {
body := "exported"
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

View File

@ -28,6 +28,12 @@ func TestDriveExportMimeTypeForFormat(t *testing.T) {
format: "",
wantMime: "application/pdf",
},
{
name: "doc_auto",
googleMime: "application/vnd.google-apps.document",
format: "auto",
wantMime: "application/pdf",
},
{
name: "doc_pdf",
googleMime: "application/vnd.google-apps.document",
@ -59,6 +65,12 @@ func TestDriveExportMimeTypeForFormat(t *testing.T) {
format: "",
wantMime: "text/csv",
},
{
name: "sheet_auto",
googleMime: "application/vnd.google-apps.spreadsheet",
format: "auto",
wantMime: "text/csv",
},
{
name: "sheet_pdf",
googleMime: "application/vnd.google-apps.spreadsheet",