diff --git a/files/filestore.go b/files/filestore.go index fdf1978..5ed0ddc 100644 --- a/files/filestore.go +++ b/files/filestore.go @@ -12,6 +12,8 @@ type File struct { ExpiresOn time.Time `json:"expires_on"` Body io.ReadCloser + + FileSize int64 `json:"file_size"` } type FileStore interface { diff --git a/files/filestore_fs.go b/files/filestore_fs.go index ce19995..5bf6e6c 100644 --- a/files/filestore_fs.go +++ b/files/filestore_fs.go @@ -40,10 +40,12 @@ func (s *FSFileStore) Store(f *File) error { } defer dst.Close() - if _, err := io.Copy(dst, f.Body); err != nil { + n, err := io.Copy(dst, f.Body) + if err != nil { return err } s.metadata[f.ID] = metadata + s.metadata[f.ID].FileSize = n if err := s.writeMetadata(); err != nil { delete(s.metadata, f.ID) return err diff --git a/files/filestore_memory.go b/files/filestore_memory.go index 0dd7e76..0640c90 100644 --- a/files/filestore_memory.go +++ b/files/filestore_memory.go @@ -14,6 +14,7 @@ type fileData struct { MaxViews uint ExpiresOn time.Time + FileSize int64 } type MemoryFileStore struct { @@ -35,9 +36,11 @@ func (s *MemoryFileStore) Store(f *File) error { ExpiresOn: f.ExpiresOn, } - _, err := io.Copy(&data.Body, f.Body) + n, err := io.Copy(&data.Body, f.Body) _ = f.Body.Close() + data.FileSize = n + s.lock.Lock() defer s.lock.Unlock() diff --git a/files/filestore_test.go b/files/filestore_test.go index 6c7c52f..a27ea91 100644 --- a/files/filestore_test.go +++ b/files/filestore_test.go @@ -92,6 +92,7 @@ func RunPersistentFilestoreTest(newStoreFunc func() files.FileStore, t *testing. MaxViews: 5, ExpiresOn: time.Now().Add(10 * time.Minute), Body: io.NopCloser(strings.NewReader("cocks!")), + FileSize: 6, }, ExpectedData: "cocks!", }, @@ -102,6 +103,7 @@ func RunPersistentFilestoreTest(newStoreFunc func() files.FileStore, t *testing. MaxViews: 5, ExpiresOn: time.Now().Add(10 * time.Minute), Body: io.NopCloser(strings.NewReader("derps!")), + FileSize: 6, }, ExpectedData: "derps!", },