Skip to content

Commit

Permalink
handler: Emit pre-finish hook for concatenated uploads
Browse files Browse the repository at this point in the history
Fixes #1169.
  • Loading branch information
Acconut committed Sep 10, 2024
1 parent dfaf9f7 commit ba7ae8b
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 19 deletions.
17 changes: 14 additions & 3 deletions pkg/handler/concat_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,14 @@ func TestConcat(t *testing.T) {
BasePath: "files",
StoreComposer: composer,
NotifyCompleteUploads: true,
PreFinishResponseCallback: func(hook HookEvent) (HTTPResponse, error) {
a.Equal("foo", hook.Upload.ID)
return HTTPResponse{
Header: HTTPHeader{
"X-Custom-Resp-Header": "hello",
},
}, nil
},
})

c := make(chan HookEvent, 1)
Expand All @@ -160,10 +168,13 @@ func TestConcat(t *testing.T) {
// A space between `final;` and the first URL should be allowed due to
// compatibility reasons, even if the specification does not define
// it. Therefore this character is included in this test case.
"Upload-Concat": "final; http://tus.io/files/a /files/b/",
"X-Custom-Header": "tada",
"Upload-Concat": "final; http://tus.io/files/a /files/b/",
"X-Custom-Req-Header": "tada",
},
Code: http.StatusCreated,
ResHeader: map[string]string{
"X-Custom-Resp-Header": "hello",
},
}).Run(handler, t)

event := <-c
Expand All @@ -178,7 +189,7 @@ func TestConcat(t *testing.T) {
req := event.HTTPRequest
a.Equal("POST", req.Method)
a.Equal("", req.URI)
a.Equal("tada", req.Header.Get("X-Custom-Header"))
a.Equal("tada", req.Header.Get("X-Custom-Req-Header"))
})

SubTest(t, "Status", func(t *testing.T, store *MockFullDataStore, composer *StoreComposer) {
Expand Down
45 changes: 29 additions & 16 deletions pkg/handler/unrouted_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -410,8 +410,10 @@ func (handler *UnroutedHandler) PostFile(w http.ResponseWriter, r *http.Request)
}
info.Offset = size

if handler.config.NotifyCompleteUploads {
handler.CompleteUploads <- newHookEvent(c, info)
resp, err = handler.emitFinishEvents(c, resp, info)
if err != nil {
handler.sendError(c, err)
return
}
}

Expand Down Expand Up @@ -936,31 +938,42 @@ func (handler *UnroutedHandler) writeChunk(c *httpContext, resp HTTPResponse, up

// 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.
// function and emit the necessary events for the hooks.
func (handler *UnroutedHandler) finishUploadIfComplete(c *httpContext, resp HTTPResponse, upload Upload, info FileInfo) (HTTPResponse, error) {
// If the upload is completed, ...
if !info.SizeIsDeferred && info.Offset == info.Size {
var err error
// ... allow the data storage to finish and cleanup the upload
if err := upload.FinishUpload(c); err != nil {
if err = upload.FinishUpload(c); err != nil {
return resp, err
}

// ... allow the hook callback to run before sending the response
if handler.config.PreFinishResponseCallback != nil {
resp2, err := handler.config.PreFinishResponseCallback(newHookEvent(c, info))
if err != nil {
return resp, err
}
resp = resp.MergeWith(resp2)
// ... and call pre-finish callback and send post-finish notification.
resp, err = handler.emitFinishEvents(c, resp, info)
if err != nil {
return resp, err
}
}

c.log.Info("UploadFinished", "size", info.Size)
handler.Metrics.incUploadsFinished()
return resp, nil
}

// ... send the info out to the channel
if handler.config.NotifyCompleteUploads {
handler.CompleteUploads <- newHookEvent(c, info)
// emitFinishEvents calls the PreFinishResponseCallback function and sends
// the necessary message on the CompleteUpload channel.
func (handler *UnroutedHandler) emitFinishEvents(c *httpContext, resp HTTPResponse, info FileInfo) (HTTPResponse, error) {
if handler.config.PreFinishResponseCallback != nil {
resp2, err := handler.config.PreFinishResponseCallback(newHookEvent(c, info))
if err != nil {
return resp, err
}
resp = resp.MergeWith(resp2)
}

c.log.Info("UploadFinished", "size", info.Size)
handler.Metrics.incUploadsFinished()

if handler.config.NotifyCompleteUploads {
handler.CompleteUploads <- newHookEvent(c, info)
}

return resp, nil
Expand Down

0 comments on commit ba7ae8b

Please sign in to comment.