From 45e83d878ac8889552618338f8cb49dd6306edef Mon Sep 17 00:00:00 2001 From: Marius Kleidl Date: Wed, 23 Aug 2023 09:04:40 +0200 Subject: [PATCH] handler: Allow custom response when upload is stopped --- pkg/handler/datastore.go | 16 ++++++++-------- pkg/handler/patch_test.go | 14 +++++++++++--- pkg/handler/unrouted_handler.go | 29 ++++++++++++++++++----------- pkg/hooks/hooks.go | 3 +-- 4 files changed, 38 insertions(+), 24 deletions(-) diff --git a/pkg/handler/datastore.go b/pkg/handler/datastore.go index 7218f9e5e..327b87ccf 100644 --- a/pkg/handler/datastore.go +++ b/pkg/handler/datastore.go @@ -33,20 +33,20 @@ type FileInfo struct { // store is used. This map may also be nil. Storage map[string]string - // stopUpload is the cancel function for the upload's context.Context. When - // invoked it will interrupt the writes to DataStore#WriteChunk. - stopUpload context.CancelFunc + // stopUpload is a channel for communicating that an upload should by stopped + // and interrupt the writes to DataStore#WriteChunk. + stopUpload chan HTTPResponse } -// StopUpload interrupts an running upload from the server-side. This means that +// StopUpload interrupts a running upload from the server-side. This means that // the current request body is closed, so that the data store does not get any // more data. Furthermore, a response is sent to notify the client of the // interrupting and the upload is terminated (if supported by the data store), -// so the upload cannot be resumed anymore. -// TODO: Allow passing in a HTTP Response -func (f FileInfo) StopUpload() { +// so the upload cannot be resumed anymore. The response to the client can be +// optionally modified by providing values in the HTTPResponse struct. +func (f FileInfo) StopUpload(response HTTPResponse) { if f.stopUpload != nil { - f.stopUpload() + f.stopUpload <- response } } diff --git a/pkg/handler/patch_test.go b/pkg/handler/patch_test.go index 42c7b95e0..2942990cc 100644 --- a/pkg/handler/patch_test.go +++ b/pkg/handler/patch_test.go @@ -600,7 +600,14 @@ func TestPatch(t *testing.T) { event := <-c info := event.Upload - info.StopUpload() + // Include a custom response + info.StopUpload(HTTPResponse{ + StatusCode: http.StatusPaymentRequired, + Body: "upload is stopped because you didn't pay", + Headers: HTTPHeaders{ + "X-Foo": "bar", + }, + }) // Wait a short time to ensure that the goroutine in the PATCH // handler has received and processed the stop event. @@ -624,11 +631,12 @@ func TestPatch(t *testing.T) { "Upload-Offset": "0", }, ReqBody: reader, - Code: http.StatusBadRequest, + Code: http.StatusPaymentRequired, ResHeader: map[string]string{ "Upload-Offset": "", + "X-Foo": "bar", }, - ResBody: "ERR_UPLOAD_STOPPED: upload has been stopped by server\n", + ResBody: "upload is stopped because you didn't pay", }).Run(handler, t) _, more := <-c diff --git a/pkg/handler/unrouted_handler.go b/pkg/handler/unrouted_handler.go index e48fb1ee6..473122110 100644 --- a/pkg/handler/unrouted_handler.go +++ b/pkg/handler/unrouted_handler.go @@ -843,25 +843,31 @@ func (handler *UnroutedHandler) writeChunk(c *httpContext, resp HTTPResponse, up // Limit the data read from the request's body to the allowed maximum c.body = newBodyReader(r.Body, maxSize) - // We use a context object to allow the hook system to cancel an upload - uploadCtx, stopUpload := context.WithCancel(context.Background()) - info.stopUpload = stopUpload + // We use a channel to allow the hook system to cancel an upload. The channel + // is closed, so that the goroutine can exit when the upload completes normally. + info.stopUpload = make(chan HTTPResponse) + defer close(info.stopUpload) // terminateUpload specifies whether the upload should be deleted after // the write has finished terminateUpload := false + var terminateUploadResponse HTTPResponse + // serverShutDown specifies whether the upload was stopped because the server closes down. serverShutDown := false - // Cancel the context when the function exits to ensure that the goroutine - // is properly cleaned up - defer stopUpload() - go func() { select { - case <-uploadCtx.Done(): - // uploadCtx is done if the upload is stopped by a post-receive hook + case resp, ok := <-info.stopUpload: + // If the channel is closed, the request completed (successfully or not) and so + // we can stop waiting on the channels. + if !ok { + return + } + + // Otherwise, the upload is stopped by a post-receive hook and resp contains the response. terminateUpload = true + terminateUploadResponse = resp case <-handler.serverCtx: // serverCtx is closed if the server is being shut down serverShutDown = true @@ -895,9 +901,10 @@ func (handler *UnroutedHandler) writeChunk(c *httpContext, resp HTTPResponse, up } // If the upload was stopped by the server, send an error response indicating this. - // TODO: Include a custom reason for the end user why the upload was stopped. if terminateUpload { - err = ErrUploadStoppedByServer + stoppedErr := ErrUploadStoppedByServer + stoppedErr.HTTPResponse = stoppedErr.HTTPResponse.MergeWith(terminateUploadResponse) + err = stoppedErr } // If the server is closing down, send an error response indicating this. diff --git a/pkg/hooks/hooks.go b/pkg/hooks/hooks.go index e639bd287..0009a3118 100644 --- a/pkg/hooks/hooks.go +++ b/pkg/hooks/hooks.go @@ -139,8 +139,7 @@ func postReceiveCallback(event handler.HookEvent, hookHandler HookHandler) { if hookRes.StopUpload { slog.Info("HookStopUpload", "id", event.Upload.ID) - // TODO: Control response for PATCH request - event.Upload.StopUpload() + event.Upload.StopUpload(hookRes.HTTPResponse) } }