handler: Allow custom response when upload is stopped

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
 	}
 }
 

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

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.

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)
 	}
 }