From 273a47db7f7a1a7d8a6b8fe6994e146fbfdf142d Mon Sep 17 00:00:00 2001 From: Marius Date: Sun, 28 Nov 2021 16:42:19 +0100 Subject: [PATCH] Allow hooks to modify response --- pkg/handler/unrouted_handler.go | 72 ++++++++++++++++++++------------- 1 file changed, 45 insertions(+), 27 deletions(-) diff --git a/pkg/handler/unrouted_handler.go b/pkg/handler/unrouted_handler.go index 54f2542..df7a3dc 100644 --- a/pkg/handler/unrouted_handler.go +++ b/pkg/handler/unrouted_handler.go @@ -46,7 +46,7 @@ func NewError(errCode string, message string, statusCode int) Error { Message: message, HTTPResponse: HTTPResponse{ StatusCode: statusCode, - Body: []byte(errCode + ": " + message + "\n"), + Body: errCode + ": " + message + "\n", Headers: HTTPHeaders{ "Content-Type": "text/plain; charset=utf-8", }, @@ -96,12 +96,9 @@ type HTTPRequest struct { type HTTPHeaders map[string]string type HTTPResponse struct { - // HTTPStatus, HTTPHeaders and HTTPBody control these details of the corresponding - // HTTP response. - // TODO: Currently only works for error responses StatusCode int Headers HTTPHeaders - Body []byte + Body string } func (resp HTTPResponse) writeTo(w http.ResponseWriter) { @@ -117,7 +114,21 @@ func (resp HTTPResponse) writeTo(w http.ResponseWriter) { w.WriteHeader(resp.StatusCode) if len(resp.Body) > 0 { - w.Write(resp.Body) + w.Write([]byte(resp.Body)) + } +} + +func (resp *HTTPResponse) MergeWith(resp2 HTTPResponse) { + if resp2.StatusCode != 0 { + resp.StatusCode = resp2.StatusCode + } + + for key, value := range resp2.Headers { + resp.Headers[key] = value + } + + if len(resp2.Body) > 0 { + resp.Body = resp2.Body } } @@ -378,11 +389,18 @@ func (handler *UnroutedHandler) PostFile(w http.ResponseWriter, r *http.Request) PartialUploads: partialUploadIDs, } + resp := HTTPResponse{ + StatusCode: http.StatusCreated, + Headers: HTTPHeaders{}, + } + if handler.config.PreUploadCreateCallback != nil { - if _, err := handler.config.PreUploadCreateCallback(newHookEvent(info, r)); err != nil { + resp2, err := handler.config.PreUploadCreateCallback(newHookEvent(info, r)) + if err != nil { handler.sendError(w, r, err) return } + resp.MergeWith(resp2) } upload, err := handler.composer.Core.NewUpload(ctx, info) @@ -402,12 +420,7 @@ func (handler *UnroutedHandler) PostFile(w http.ResponseWriter, r *http.Request) // Add the Location header directly after creating the new resource to even // include it in cases of failure when an error is returned url := handler.absFileURL(r, id) - resp := HTTPResponse{ - StatusCode: http.StatusCreated, - Headers: HTTPHeaders{ - "Location": url, - }, - } + resp.Headers["Location"] = url handler.Metrics.incUploadsCreated() handler.log("UploadCreated", "id", id, "size", i64toa(size), "url", url) @@ -440,7 +453,8 @@ func (handler *UnroutedHandler) PostFile(w http.ResponseWriter, r *http.Request) defer lock.Unlock() } - if err := handler.writeChunk(ctx, upload, info, resp, r); err != nil { + resp, err = handler.writeChunk(ctx, upload, info, resp, r) + if err != nil { handler.sendError(w, r, err) return } @@ -448,7 +462,8 @@ func (handler *UnroutedHandler) PostFile(w http.ResponseWriter, r *http.Request) // Directly finish the upload if the upload is empty (i.e. has a size of 0). // This statement is in an else-if block to avoid causing duplicate calls // to finishUploadIfComplete if an upload is empty and contains a chunk. - if err := handler.finishUploadIfComplete(ctx, upload, info, r); err != nil { + resp, err = handler.finishUploadIfComplete(ctx, upload, info, resp, r) + if err != nil { handler.sendError(w, r, err) return } @@ -620,7 +635,8 @@ func (handler *UnroutedHandler) PatchFile(w http.ResponseWriter, r *http.Request info.SizeIsDeferred = false } - if err := handler.writeChunk(ctx, upload, info, resp, r); err != nil { + resp, err = handler.writeChunk(ctx, upload, info, resp, r) + if err != nil { handler.sendError(w, r, err) return } @@ -631,7 +647,7 @@ func (handler *UnroutedHandler) PatchFile(w http.ResponseWriter, r *http.Request // writeChunk reads the body from the requests r and appends it to the upload // with the corresponding id. Afterwards, it will set the necessary response // headers but will not send the response. -func (handler *UnroutedHandler) writeChunk(ctx context.Context, upload Upload, info FileInfo, resp HTTPResponse, r *http.Request) error { +func (handler *UnroutedHandler) writeChunk(ctx context.Context, upload Upload, info FileInfo, resp HTTPResponse, r *http.Request) (HTTPResponse, error) { // Get Content-Length if possible length := r.ContentLength offset := info.Offset @@ -639,7 +655,7 @@ func (handler *UnroutedHandler) writeChunk(ctx context.Context, upload Upload, i // Test if this upload fits into the file's size if !info.SizeIsDeferred && offset+length > info.Size { - return ErrSizeExceeded + return resp, ErrSizeExceeded } maxSize := info.Size - offset @@ -719,7 +735,7 @@ func (handler *UnroutedHandler) writeChunk(ctx context.Context, upload Upload, i handler.log("ChunkWriteComplete", "id", id, "bytesWritten", i64toa(bytesWritten)) if err != nil { - return err + return resp, err } // Send new offset to client @@ -728,18 +744,18 @@ func (handler *UnroutedHandler) writeChunk(ctx context.Context, upload Upload, i handler.Metrics.incBytesReceived(uint64(bytesWritten)) info.Offset = newOffset - return handler.finishUploadIfComplete(ctx, upload, info, r) + return handler.finishUploadIfComplete(ctx, upload, info, resp, r) } // finishUploadIfComplete checks whether an upload is completed (i.e. upload offset // matches upload size) and if so, it will call the data store's FinishUpload // function and send the necessary message on the CompleteUpload channel. -func (handler *UnroutedHandler) finishUploadIfComplete(ctx context.Context, upload Upload, info FileInfo, r *http.Request) error { +func (handler *UnroutedHandler) finishUploadIfComplete(ctx context.Context, upload Upload, info FileInfo, resp HTTPResponse, r *http.Request) (HTTPResponse, error) { // If the upload is completed, ... if !info.SizeIsDeferred && info.Offset == info.Size { // ... allow custom mechanism to finish and cleanup the upload if err := upload.FinishUpload(ctx); err != nil { - return err + return resp, err } // ... send the info out to the channel @@ -750,13 +766,15 @@ func (handler *UnroutedHandler) finishUploadIfComplete(ctx context.Context, uplo handler.Metrics.incUploadsFinished() if handler.config.PreFinishResponseCallback != nil { - if _, err := handler.config.PreFinishResponseCallback(newHookEvent(info, r)); err != nil { - return err + resp2, err := handler.config.PreFinishResponseCallback(newHookEvent(info, r)) + if err != nil { + return resp, err } + resp.MergeWith(resp2) } } - return nil + return resp, nil } // GetFile handles requests to download a file using a GET request. This is not @@ -800,7 +818,7 @@ func (handler *UnroutedHandler) GetFile(w http.ResponseWriter, r *http.Request) "Content-Type": contentType, "Content-Disposition": contentDisposition, }, - Body: nil, // Body is intentionally left nil, and we copy it manually in later. + Body: "", // Body is intentionally left empty, and we copy it manually in later. } // If no data has been uploaded yet, respond with an empty "204 No Content" status. @@ -1006,7 +1024,7 @@ func (handler *UnroutedHandler) sendError(w http.ResponseWriter, r *http.Request // If we are sending the response for a HEAD request, ensure that we are not including // any response body. if r.Method == "HEAD" { - detailedErr.HTTPResponse.Body = nil + detailedErr.HTTPResponse.Body = "" } handler.sendResp(w, r, detailedErr.HTTPResponse)