Official cwtch.im peer and server implementations. https://cwtch.im
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

153 lines
4.1KB

  1. package storage
  2. import (
  3. "bufio"
  4. "cwtch.im/cwtch/protocol"
  5. "cwtch.im/cwtch/server/metrics"
  6. "encoding/json"
  7. "fmt"
  8. "git.openprivacy.ca/openprivacy/log"
  9. "os"
  10. "path"
  11. "sync"
  12. )
  13. const (
  14. fileStorePartitions = 10
  15. fileStoreFilename = "cwtch.messages"
  16. directory = "messages"
  17. )
  18. // MessageStoreInterface defines an interface to interact with a store of cwtch messages.
  19. type MessageStoreInterface interface {
  20. AddMessage(protocol.GroupMessage)
  21. FetchMessages() []*protocol.GroupMessage
  22. }
  23. // MessageStore is a file-backed implementation of MessageStoreInterface
  24. type MessageStore struct {
  25. activeLogFile *os.File
  26. filePos int
  27. storeDirectory string
  28. lock sync.Mutex
  29. messages []*protocol.GroupMessage
  30. messageCounter metrics.Counter
  31. maxBufferLines int
  32. bufferPos int
  33. bufferRotated bool
  34. }
  35. // Close closes the message store and underlying resources.
  36. func (ms *MessageStore) Close() {
  37. ms.lock.Lock()
  38. ms.messages = nil
  39. ms.activeLogFile.Close()
  40. ms.lock.Unlock()
  41. }
  42. func (ms *MessageStore) updateBuffer(gm *protocol.GroupMessage) {
  43. ms.messages[ms.bufferPos] = gm
  44. ms.bufferPos++
  45. if ms.bufferPos == ms.maxBufferLines {
  46. ms.bufferPos = 0
  47. ms.bufferRotated = true
  48. }
  49. }
  50. func (ms *MessageStore) initAndLoadFiles() error {
  51. ms.activeLogFile = nil
  52. for i := fileStorePartitions - 1; i >= 0; i-- {
  53. ms.filePos = 0
  54. filename := path.Join(ms.storeDirectory, fmt.Sprintf("%s.%d", fileStoreFilename, i))
  55. f, err := os.OpenFile(filename, os.O_CREATE|os.O_APPEND|os.O_RDWR, 0600)
  56. if err != nil {
  57. log.Errorf("MessageStore could not open: %v: %v", filename, err)
  58. continue
  59. }
  60. ms.activeLogFile = f
  61. scanner := bufio.NewScanner(f)
  62. for scanner.Scan() {
  63. gms := scanner.Text()
  64. ms.filePos++
  65. gm := &protocol.GroupMessage{}
  66. err := json.Unmarshal([]byte(gms), gm)
  67. if err == nil {
  68. ms.updateBuffer(gm)
  69. }
  70. }
  71. }
  72. if ms.activeLogFile == nil {
  73. return fmt.Errorf("Could not create log file to write to in %s", ms.storeDirectory)
  74. }
  75. return nil
  76. }
  77. func (ms *MessageStore) updateFile(gm *protocol.GroupMessage) {
  78. s, err := json.Marshal(gm)
  79. if err != nil {
  80. log.Errorf("Failed to unmarshal group message %v\n", err)
  81. }
  82. fmt.Fprintf(ms.activeLogFile, "%s\n", s)
  83. ms.filePos++
  84. if ms.filePos >= ms.maxBufferLines/fileStorePartitions {
  85. ms.rotateFileStore()
  86. }
  87. }
  88. func (ms *MessageStore) rotateFileStore() {
  89. ms.activeLogFile.Close()
  90. os.Remove(path.Join(ms.storeDirectory, fmt.Sprintf("%s.%d", fileStoreFilename, fileStorePartitions-1)))
  91. for i := fileStorePartitions - 2; i >= 0; i-- {
  92. os.Rename(path.Join(ms.storeDirectory, fmt.Sprintf("%s.%d", fileStoreFilename, i)), path.Join(ms.storeDirectory, fmt.Sprintf("%s.%d", fileStoreFilename, i+1)))
  93. }
  94. f, err := os.OpenFile(path.Join(ms.storeDirectory, fmt.Sprintf("%s.%d", fileStoreFilename, 0)), os.O_CREATE|os.O_APPEND|os.O_RDWR, 0600)
  95. if err != nil {
  96. log.Errorf("Could not open new message store file in: %s", ms.storeDirectory)
  97. }
  98. ms.filePos = 0
  99. ms.activeLogFile = f
  100. }
  101. // Init sets up a MessageStore of size maxBufferLines (# of messages) backed by filename
  102. func (ms *MessageStore) Init(appDirectory string, maxBufferLines int, messageCounter metrics.Counter) error {
  103. ms.storeDirectory = path.Join(appDirectory, directory)
  104. os.Mkdir(ms.storeDirectory, 0700)
  105. ms.bufferPos = 0
  106. ms.maxBufferLines = maxBufferLines
  107. ms.messages = make([]*protocol.GroupMessage, maxBufferLines)
  108. ms.bufferRotated = false
  109. ms.messageCounter = messageCounter
  110. err := ms.initAndLoadFiles()
  111. return err
  112. }
  113. // FetchMessages returns all messages from the backing file.
  114. func (ms *MessageStore) FetchMessages() (messages []*protocol.GroupMessage) {
  115. ms.lock.Lock()
  116. if !ms.bufferRotated {
  117. messages = make([]*protocol.GroupMessage, ms.bufferPos)
  118. copy(messages, ms.messages[0:ms.bufferPos])
  119. } else {
  120. messages = make([]*protocol.GroupMessage, ms.maxBufferLines)
  121. copy(messages, ms.messages[ms.bufferPos:ms.maxBufferLines])
  122. copy(messages[ms.bufferPos:], ms.messages[0:ms.bufferPos])
  123. }
  124. ms.lock.Unlock()
  125. return
  126. }
  127. // AddMessage adds a GroupMessage to the store
  128. func (ms *MessageStore) AddMessage(gm protocol.GroupMessage) {
  129. ms.messageCounter.Add(1)
  130. ms.lock.Lock()
  131. ms.updateBuffer(&gm)
  132. ms.updateFile(&gm)
  133. ms.lock.Unlock()
  134. }