Skip to content

Commit

Permalink
feat(handler, s3store): implement ContentServerDataStore for direct c…
Browse files Browse the repository at this point in the history
…ontent serving, closes tus#1064

- Add ServerDataStore interface
- Implement ContentServerDataStore in S3Store with streaming support
- Add Range header support for partial content requests
- Update StoreComposer to support ContentServer capability
- Add tests for new ContentServerDataStore functionality
- Update Go version to 1.22.1
  • Loading branch information
pcfreak30 committed Nov 26, 2024
1 parent 9779a84 commit 259cd68
Show file tree
Hide file tree
Showing 6 changed files with 211 additions and 1 deletion.
3 changes: 2 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,8 @@ module github.com/tus/tusd/v2
// Specify the Go version needed for the Heroku deployment
// See https://github.com/heroku/heroku-buildpack-go#go-module-specifics
// +heroku goVersion go1.22
go 1.21.0
go 1.22.1

toolchain go1.22.7

require (
Expand Down
6 changes: 6 additions & 0 deletions pkg/handler/composer.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@ type StoreComposer struct {
Concater ConcaterDataStore
UsesLengthDeferrer bool
LengthDeferrer LengthDeferrerDataStore
ContentServer ContentServerDataStore
UsesContentServer bool
}

// NewStoreComposer creates a new and empty store composer.
Expand Down Expand Up @@ -85,3 +87,7 @@ func (store *StoreComposer) UseLengthDeferrer(ext LengthDeferrerDataStore) {
store.UsesLengthDeferrer = ext != nil
store.LengthDeferrer = ext
}
func (store *StoreComposer) UseContentServer(ext ContentServerDataStore) {
store.UsesContentServer = ext != nil
store.ContentServer = ext
}
11 changes: 11 additions & 0 deletions pkg/handler/datastore.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package handler
import (
"context"
"io"
"net/http"
)

type MetaData map[string]string
Expand Down Expand Up @@ -121,6 +122,16 @@ type DataStore interface {
GetUpload(ctx context.Context, id string) (upload Upload, err error)
}

// ServableUpload defines the method for serving content directly
type ServableUpload interface {
ServeContent(ctx context.Context, w http.ResponseWriter, r *http.Request) error
}

// ContentServerDataStore is the interface for data stores that can serve content directly
type ContentServerDataStore interface {
AsServableUpload(upload Upload) ServableUpload
}

type TerminatableUpload interface {
// Terminate an upload so any further requests to the upload resource will
// return the ErrNotFound error.
Expand Down
11 changes: 11 additions & 0 deletions pkg/handler/unrouted_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -1013,6 +1013,17 @@ func (handler *UnroutedHandler) GetFile(w http.ResponseWriter, r *http.Request)
return
}

// If the data store implements ContentServerDataStore, use the ServableUpload interface
if handler.composer.UsesContentServer {
servableUpload := handler.composer.ContentServer.AsServableUpload(upload)
err = servableUpload.ServeContent(c, w, r)
if err != nil {
handler.sendError(c, err)
}
return
}

// Fall back to the existing GetReader implementation if ContentServerDataStore is not implemented
contentType, contentDisposition := filterContentType(info)
resp := HTTPResponse{
StatusCode: http.StatusOK,
Expand Down
50 changes: 50 additions & 0 deletions pkg/s3store/s3store.go
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,7 @@ import (
"net/http"
"os"
"regexp"
"strconv"
"strings"
"sync"
"time"
Expand Down Expand Up @@ -376,6 +377,55 @@ func (store S3Store) AsConcatableUpload(upload handler.Upload) handler.Concatabl
return upload.(*s3Upload)
}

func (store S3Store) AsServableUpload(upload handler.Upload) handler.ServableUpload {
return upload.(*s3Upload)
}

func (su *s3Upload) ServeContent(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
// Get file info
info, err := su.GetInfo(ctx)
if err != nil {
return err
}

// Prepare GetObject input
input := &s3.GetObjectInput{
Bucket: aws.String(su.store.Bucket),
Key: su.store.keyWithPrefix(su.objectId),
}

// Forward the Range header if present
if rangeHeader := r.Header.Get("Range"); rangeHeader != "" {
input.Range = aws.String(rangeHeader)
}

// Let S3 handle the request
result, err := su.store.Service.GetObject(ctx, input)
if err != nil {
return err
}
defer result.Body.Close()

// Set headers
w.Header().Set("Content-Length", strconv.FormatInt(info.Size, 10))
w.Header().Set("Content-Type", info.MetaData["filetype"])
w.Header().Set("ETag", *result.ETag)

// Add Content-Disposition if present in S3 response
if result.ContentDisposition != nil {
w.Header().Set("Content-Disposition", *result.ContentDisposition)
}

// Add Content-Encoding if present in S3 response
if result.ContentEncoding != nil {
w.Header().Set("Content-Encoding", *result.ContentEncoding)
}

// Stream the content
_, err = io.Copy(w, result.Body)
return err
}

func (upload *s3Upload) writeInfo(ctx context.Context, info handler.FileInfo) error {
store := upload.store

Expand Down
131 changes: 131 additions & 0 deletions pkg/s3store/s3store_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,11 @@ package s3store
import (
"bytes"
"context"
"errors"
"fmt"
"io"
"net/http"
"net/http/httptest"
"os"
"strings"
"testing"
Expand Down Expand Up @@ -1468,3 +1471,131 @@ func TestWriteChunkCleansUpTempFiles(t *testing.T) {
assert.Nil(err)
assert.Equal(len(files), 0)
}

func TestS3StoreAsServerDataStore(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
assert := assert.New(t)

s3obj := NewMockS3API(mockCtrl)
store := New("bucket", s3obj)

upload := &s3Upload{
store: &store,
info: &handler.FileInfo{},
objectId: "uploadId",
multipartId: "multipartId",
}

servableUpload := store.AsServableUpload(upload)
assert.NotNil(servableUpload)
assert.IsType(&s3Upload{}, servableUpload)
}

func TestS3ServableUploadServeContent(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
assert := assert.New(t)

s3obj := NewMockS3API(mockCtrl)
store := New("bucket", s3obj)

upload := &s3Upload{
store: &store,
info: &handler.FileInfo{Size: 100, Offset: 100, MetaData: map[string]string{"filetype": "text/plain"}},
objectId: "uploadId",
multipartId: "multipartId",
}

s3obj.EXPECT().GetObject(gomock.Any(), &s3.GetObjectInput{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
}).Return(&s3.GetObjectOutput{
Body: io.NopCloser(strings.NewReader("test content")),
ContentLength: aws.Int64(100),
ETag: aws.String("etag123"),
}, nil)

servableUpload := store.AsServableUpload(upload)

w := httptest.NewRecorder()
r := httptest.NewRequest("GET", "/", nil)

err := servableUpload.ServeContent(context.Background(), w, r)
assert.Nil(err)

assert.Equal(http.StatusOK, w.Code)
assert.Equal("100", w.Header().Get("Content-Length"))
assert.Equal("text/plain", w.Header().Get("Content-Type"))
assert.Equal("etag123", w.Header().Get("ETag"))
assert.Equal("test content", w.Body.String())
}

func TestS3ServableUploadServeContentWithRange(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
assert := assert.New(t)

s3obj := NewMockS3API(mockCtrl)
store := New("bucket", s3obj)

upload := &s3Upload{
store: &store,
info: &handler.FileInfo{Size: 100, Offset: 100, MetaData: map[string]string{"filetype": "text/plain"}},
objectId: "uploadId",
multipartId: "multipartId",
}

s3obj.EXPECT().GetObject(gomock.Any(), &s3.GetObjectInput{
Bucket: aws.String("bucket"),
Key: aws.String("uploadId"),
Range: aws.String("bytes=10-19"),
}).Return(&s3.GetObjectOutput{
Body: io.NopCloser(strings.NewReader("0123456789")),
ContentLength: aws.Int64(10),
ETag: aws.String("etag123"),
}, nil)

servableUpload := store.AsServableUpload(upload)

w := httptest.NewRecorder()
r := httptest.NewRequest("GET", "/", nil)
r.Header.Set("Range", "bytes=10-19")

err := servableUpload.ServeContent(context.Background(), w, r)
assert.Nil(err)

assert.Equal(http.StatusPartialContent, w.Code)
assert.Equal("10", w.Header().Get("Content-Length"))
assert.Equal("text/plain", w.Header().Get("Content-Type"))
assert.Equal("etag123", w.Header().Get("ETag"))
assert.Equal("bytes 10-19/100", w.Header().Get("Content-Range"))
assert.Equal("0123456789", w.Body.String())
}

func TestS3ServableUploadServeContentError(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
assert := assert.New(t)

s3obj := NewMockS3API(mockCtrl)
store := New("bucket", s3obj)

upload := &s3Upload{
store: &store,
info: &handler.FileInfo{Size: 100, Offset: 100, MetaData: map[string]string{"filetype": "text/plain"}},
objectId: "uploadId",
multipartId: "multipartId",
}

expectedError := errors.New("S3 error")
s3obj.EXPECT().GetObject(gomock.Any(), gomock.Any()).Return(nil, expectedError)

servableUpload := store.AsServableUpload(upload)

w := httptest.NewRecorder()
r := httptest.NewRequest("GET", "/", nil)

err := servableUpload.ServeContent(context.Background(), w, r)
assert.Equal(expectedError, err)
}

0 comments on commit 259cd68

Please sign in to comment.