From c41ed01184dc430b16613fa6310f58deadb3131b Mon Sep 17 00:00:00 2001
From: simon987 <fortier.simon@protonmail.com>
Date: Sun, 10 Mar 2019 11:14:44 -0400
Subject: [PATCH] Allow passing the upload token as query string for reading

---
 api/slot.go       |  3 +++
 test/slot_test.go | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/api/slot.go b/api/slot.go
index 5d6f3cc..ca6adff 100644
--- a/api/slot.go
+++ b/api/slot.go
@@ -142,6 +142,9 @@ func (api *WebApi) Upload(ctx *fasthttp.RequestCtx) {
 func (api *WebApi) ReadUploadSlot(ctx *fasthttp.RequestCtx) {
 
 	tokenStr := string(ctx.Request.Header.Peek("X-Upload-Token"))
+	if tokenStr == "" {
+		tokenStr = string(ctx.Request.URI().QueryArgs().Peek("token"))
+	}
 
 	slot := UploadSlot{}
 	err := api.db.Where("token=?", tokenStr).First(&slot).Error
diff --git a/test/slot_test.go b/test/slot_test.go
index 5618df4..2123956 100644
--- a/test/slot_test.go
+++ b/test/slot_test.go
@@ -1,8 +1,13 @@
 package test
 
 import (
+	"bytes"
+	"github.com/fasthttp/websocket"
 	"github.com/simon987/ws_bucket/api"
+	"io/ioutil"
+	"net/http"
 	"testing"
+	"time"
 )
 
 func TestAllocateUploadInvalidMaxSize(t *testing.T) {
@@ -65,6 +70,33 @@ func TestDuplicateUploadSlot(t *testing.T) {
 	}
 }
 
+func TestTokenInQueryString(t *testing.T) {
+
+	if allocateUploadSlot(api.AllocateUploadSlotRequest{
+		FileName: "test.png",
+		Token:    "testquery",
+		MaxSize:  100,
+	}).Ok != true {
+		t.Error()
+	}
+
+	conn := ws("testquery")
+	conn.WriteMessage(websocket.BinaryMessage, []byte("test"))
+	conn.Close()
+
+	time.Sleep(time.Millisecond * 20)
+	r, err := http.Get("http://" + api.GetServerAddress() + "/slot?token=testquery")
+	handleErr(err)
+
+	data, err := ioutil.ReadAll(r.Body)
+	handleErr(err)
+
+	if bytes.Compare(data, []byte("test")) != 0 {
+		t.Error()
+	}
+
+}
+
 func allocateUploadSlot(request api.AllocateUploadSlotRequest) (ar *api.GenericResponse) {
 	resp := Post("/slot", request)
 	UnmarshalResponse(resp, &ar)