file_cache.go 2.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127
  1. package server
  2. import (
  3. "errors"
  4. "fmt"
  5. "heckel.io/ntfy/v2/log"
  6. "heckel.io/ntfy/v2/util"
  7. "io"
  8. "os"
  9. "path/filepath"
  10. "regexp"
  11. "sync"
  12. )
  13. var (
  14. fileIDRegex = regexp.MustCompile(fmt.Sprintf(`^[-_A-Za-z0-9]{%d}$`, messageIDLength))
  15. errInvalidFileID = errors.New("invalid file ID")
  16. errFileExists = errors.New("file exists")
  17. )
  18. type fileCache struct {
  19. dir string
  20. totalSizeCurrent int64
  21. totalSizeLimit int64
  22. mu sync.Mutex
  23. }
  24. func newFileCache(dir string, totalSizeLimit int64) (*fileCache, error) {
  25. if err := os.MkdirAll(dir, 0700); err != nil {
  26. return nil, err
  27. }
  28. size, err := dirSize(dir)
  29. if err != nil {
  30. return nil, err
  31. }
  32. return &fileCache{
  33. dir: dir,
  34. totalSizeCurrent: size,
  35. totalSizeLimit: totalSizeLimit,
  36. }, nil
  37. }
  38. func (c *fileCache) Write(id string, in io.Reader, limiters ...util.Limiter) (int64, error) {
  39. if !fileIDRegex.MatchString(id) {
  40. return 0, errInvalidFileID
  41. }
  42. log.Tag(tagFileCache).Field("message_id", id).Debug("Writing attachment")
  43. file := filepath.Join(c.dir, id)
  44. if _, err := os.Stat(file); err == nil {
  45. return 0, errFileExists
  46. }
  47. f, err := os.OpenFile(file, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
  48. if err != nil {
  49. return 0, err
  50. }
  51. defer f.Close()
  52. limiters = append(limiters, util.NewFixedLimiter(c.Remaining()))
  53. limitWriter := util.NewLimitWriter(f, limiters...)
  54. size, err := io.Copy(limitWriter, in)
  55. if err != nil {
  56. os.Remove(file)
  57. return 0, err
  58. }
  59. if err := f.Close(); err != nil {
  60. os.Remove(file)
  61. return 0, err
  62. }
  63. c.mu.Lock()
  64. c.totalSizeCurrent += size
  65. mset(metricAttachmentsTotalSize, c.totalSizeCurrent)
  66. c.mu.Unlock()
  67. return size, nil
  68. }
  69. func (c *fileCache) Remove(ids ...string) error {
  70. for _, id := range ids {
  71. if !fileIDRegex.MatchString(id) {
  72. return errInvalidFileID
  73. }
  74. log.Tag(tagFileCache).Field("message_id", id).Debug("Deleting attachment")
  75. file := filepath.Join(c.dir, id)
  76. if err := os.Remove(file); err != nil {
  77. log.Tag(tagFileCache).Field("message_id", id).Err(err).Debug("Error deleting attachment")
  78. }
  79. }
  80. size, err := dirSize(c.dir)
  81. if err != nil {
  82. return err
  83. }
  84. c.mu.Lock()
  85. c.totalSizeCurrent = size
  86. c.mu.Unlock()
  87. mset(metricAttachmentsTotalSize, size)
  88. return nil
  89. }
  90. func (c *fileCache) Size() int64 {
  91. c.mu.Lock()
  92. defer c.mu.Unlock()
  93. return c.totalSizeCurrent
  94. }
  95. func (c *fileCache) Remaining() int64 {
  96. c.mu.Lock()
  97. defer c.mu.Unlock()
  98. remaining := c.totalSizeLimit - c.totalSizeCurrent
  99. if remaining < 0 {
  100. return 0
  101. }
  102. return remaining
  103. }
  104. func dirSize(dir string) (int64, error) {
  105. entries, err := os.ReadDir(dir)
  106. if err != nil {
  107. return 0, err
  108. }
  109. var size int64
  110. for _, e := range entries {
  111. info, err := e.Info()
  112. if err != nil {
  113. return 0, err
  114. }
  115. size += info.Size()
  116. }
  117. return size, nil
  118. }