From 7b0cf5b63f109cfbface19fde487714d7ce6aa8c Mon Sep 17 00:00:00 2001
From: Kittycat <74904820+Lioncat2002@users.noreply.github.com>
Date: Wed, 31 May 2023 18:49:48 +0530
Subject: [PATCH] Added a Raw File Upload method (#3)

* feat: added a rawfileuplaod method

* refactor: moved into one method

* Set `MallocNanoZone=0` env var to fix actions

GitHub actions due to a SIGABRT received on OSX jobs, caused by: golang/go#49138

---------

Co-authored-by: Wayback Archiver <66856220+waybackarchiver@users.noreply.github.com>
---
 .github/workflows/testing.yml |  3 ++
 catbox.go                     | 74 +++++++++++++++++++++++++++++------
 catbox_test.go                | 22 +++++++++++
 3 files changed, 86 insertions(+), 13 deletions(-)

diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml
index 2c47a5e..79357b3 100644
--- a/.github/workflows/testing.yml
+++ b/.github/workflows/testing.yml
@@ -18,6 +18,9 @@ on:
       - "go.sum"
       - ".github/workflows/testing.yml"
 
+env:
+  MallocNanoZone: 0
+
 jobs:
   test:
     name: Testing
diff --git a/catbox.go b/catbox.go
index af9b4e4..177b167 100644
--- a/catbox.go
+++ b/catbox.go
@@ -42,23 +42,71 @@ func New(client *http.Client) *Catbox {
 }
 
 // Upload file or URI to the Catbox. It returns an URL string and error.
-func (cat *Catbox) Upload(path string) (string, error) {
-	parse := func(s string, _ error) (string, error) {
-		uri, err := url.Parse(s)
+func (cat *Catbox) Upload(v ...interface{}) (string, error) {
+	if len(v) == 0 {
+		return "", fmt.Errorf(`must specify file path or byte slice`)
+	}
+
+	switch t := v[0].(type) {
+	case string:
+		path := t
+		parse := func(s string, _ error) (string, error) {
+			uri, err := url.Parse(s)
+			if err != nil {
+				return "", err
+			}
+			return uri.String(), nil
+		}
+		switch {
+		case helper.IsURL(path):
+			return parse(cat.urlUpload(path))
+		case helper.Exists(path):
+			return parse(cat.fileUpload(path))
+		default:
+			return "", errors.New(`path invalid`)
+		}
+	case []byte:
+		if len(v) != 2 {
+			return "", fmt.Errorf(`must specify file name`)
+		}
+		return cat.rawUpload(t, v[1].(string))
+	}
+	return "", fmt.Errorf(`unhandled`)
+}
+
+func (cat *Catbox) rawUpload(b []byte, name string) (string, error) {
+	r, w := io.Pipe()
+	m := multipart.NewWriter(w)
+
+	go func() {
+		defer w.Close()
+		defer m.Close()
+
+		m.WriteField("reqtype", "fileupload")
+		m.WriteField("userhash", cat.Userhash)
+		part, err := m.CreateFormFile("fileToUpload", filepath.Base(name))
 		if err != nil {
-			return "", err
+			return
+		}
+		if _, err = io.Copy(part, bytes.NewBuffer(b)); err != nil {
+			return
 		}
-		return uri.String(), nil
+	}()
+	req, _ := http.NewRequest(http.MethodPost, ENDPOINT, r)
+	req.Header.Add("Content-Type", m.FormDataContentType())
+
+	resp, err := cat.Client.Do(req)
+	if err != nil {
+		return "", err
 	}
+	defer resp.Body.Close()
 
-	switch {
-	case helper.IsURL(path):
-		return parse(cat.urlUpload(path))
-	case helper.Exists(path):
-		return parse(cat.fileUpload(path))
-	default:
-		return "", errors.New(`path invalid`)
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return "", err
 	}
+
+	return string(body), nil
 }
 
 func (cat *Catbox) fileUpload(path string) (string, error) {
@@ -69,7 +117,7 @@ func (cat *Catbox) fileUpload(path string) (string, error) {
 	defer file.Close()
 
 	if size := helper.FileSize(path); size > 209715200 {
-		return "", fmt.Errorf("File too large, size: %d MB", size/1024/1024)
+		return "", fmt.Errorf("file too large, size: %d MB", size/1024/1024)
 	}
 
 	r, w := io.Pipe()
diff --git a/catbox_test.go b/catbox_test.go
index 31972f8..bf055a7 100644
--- a/catbox_test.go
+++ b/catbox_test.go
@@ -27,6 +27,28 @@ func TestFileUpload(t *testing.T) {
 	}
 }
 
+func TestRawUpload(t *testing.T) {
+	content := make([]byte, 5000)
+	tmpfile, err := ioutil.TempFile("", "go-catbox-")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.Remove(tmpfile.Name())
+	b, err := ioutil.ReadFile(tmpfile.Name())
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if _, err := tmpfile.Write(content); err != nil {
+		t.Fatal(err)
+	}
+
+	if _, err := New(nil).Upload(b, "test"); err != nil {
+		t.Fatal(err)
+	}
+}
+
 func TestURLUpload(t *testing.T) {
 	url := "https://www.gstatic.com/webp/gallery/1.webp"
 	if _, err := New(nil).Upload(url); err != nil {