server.go 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735
  1. package server
  2. import (
  3. "bytes"
  4. "context"
  5. "embed" // required for go:embed
  6. "encoding/json"
  7. firebase "firebase.google.com/go"
  8. "firebase.google.com/go/messaging"
  9. "fmt"
  10. "google.golang.org/api/option"
  11. "heckel.io/ntfy/util"
  12. "html/template"
  13. "io"
  14. "log"
  15. "net"
  16. "net/http"
  17. "regexp"
  18. "strconv"
  19. "strings"
  20. "sync"
  21. "time"
  22. )
  23. // TODO add "max messages in a topic" limit
  24. // TODO implement "since=<ID>"
  25. // Server is the main server, providing the UI and API for ntfy
  26. type Server struct {
  27. config *Config
  28. topics map[string]*topic
  29. visitors map[string]*visitor
  30. firebase subscriber
  31. messages int64
  32. cache cache
  33. mu sync.Mutex
  34. }
  35. // errHTTP is a generic HTTP error for any non-200 HTTP error
  36. type errHTTP struct {
  37. Code int
  38. Status string
  39. }
  40. func (e errHTTP) Error() string {
  41. return fmt.Sprintf("http: %s", e.Status)
  42. }
  43. type indexPage struct {
  44. Topic string
  45. CacheDuration time.Duration
  46. }
  47. type sinceTime time.Time
  48. func (t sinceTime) IsAll() bool {
  49. return t == sinceAllMessages
  50. }
  51. func (t sinceTime) IsNone() bool {
  52. return t == sinceNoMessages
  53. }
  54. func (t sinceTime) Time() time.Time {
  55. return time.Time(t)
  56. }
  57. var (
  58. sinceAllMessages = sinceTime(time.Unix(0, 0))
  59. sinceNoMessages = sinceTime(time.Unix(1, 0))
  60. )
  61. var (
  62. topicRegex = regexp.MustCompile(`^/[-_A-Za-z0-9]{1,64}$`) // Regex must match JS & Android app!
  63. jsonRegex = regexp.MustCompile(`^/[-_A-Za-z0-9]{1,64}(,[-_A-Za-z0-9]{1,64})*/json$`)
  64. sseRegex = regexp.MustCompile(`^/[-_A-Za-z0-9]{1,64}(,[-_A-Za-z0-9]{1,64})*/sse$`)
  65. rawRegex = regexp.MustCompile(`^/[-_A-Za-z0-9]{1,64}(,[-_A-Za-z0-9]{1,64})*/raw$`)
  66. sendRegex = regexp.MustCompile(`^/[-_A-Za-z0-9]{1,64}(,[-_A-Za-z0-9]{1,64})*/(publish|send|trigger)$`)
  67. staticRegex = regexp.MustCompile(`^/static/.+`)
  68. docsRegex = regexp.MustCompile(`^/docs(|/.*)$`)
  69. disallowedTopics = []string{"docs", "static"}
  70. templateFnMap = template.FuncMap{
  71. "durationToHuman": util.DurationToHuman,
  72. }
  73. //go:embed "index.gohtml"
  74. indexSource string
  75. indexTemplate = template.Must(template.New("index").Funcs(templateFnMap).Parse(indexSource))
  76. //go:embed "example.html"
  77. exampleSource string
  78. //go:embed static
  79. webStaticFs embed.FS
  80. webStaticFsCached = &util.CachingEmbedFS{ModTime: time.Now(), FS: webStaticFs}
  81. //go:embed docs
  82. docsStaticFs embed.FS
  83. docsStaticCached = &util.CachingEmbedFS{ModTime: time.Now(), FS: docsStaticFs}
  84. errHTTPBadRequest = &errHTTP{http.StatusBadRequest, http.StatusText(http.StatusBadRequest)}
  85. errHTTPNotFound = &errHTTP{http.StatusNotFound, http.StatusText(http.StatusNotFound)}
  86. errHTTPTooManyRequests = &errHTTP{http.StatusTooManyRequests, http.StatusText(http.StatusTooManyRequests)}
  87. )
  88. const (
  89. firebaseControlTopic = "~control" // See Android if changed
  90. )
  91. // New instantiates a new Server. It creates the cache and adds a Firebase
  92. // subscriber (if configured).
  93. func New(conf *Config) (*Server, error) {
  94. var firebaseSubscriber subscriber
  95. if conf.FirebaseKeyFile != "" {
  96. var err error
  97. firebaseSubscriber, err = createFirebaseSubscriber(conf)
  98. if err != nil {
  99. return nil, err
  100. }
  101. }
  102. cache, err := createCache(conf)
  103. if err != nil {
  104. return nil, err
  105. }
  106. topics, err := cache.Topics()
  107. if err != nil {
  108. return nil, err
  109. }
  110. return &Server{
  111. config: conf,
  112. cache: cache,
  113. firebase: firebaseSubscriber,
  114. topics: topics,
  115. visitors: make(map[string]*visitor),
  116. }, nil
  117. }
  118. func createCache(conf *Config) (cache, error) {
  119. if conf.CacheDuration == 0 {
  120. return newNopCache(), nil
  121. } else if conf.CacheFile != "" {
  122. return newSqliteCache(conf.CacheFile)
  123. }
  124. return newMemCache(), nil
  125. }
  126. func createFirebaseSubscriber(conf *Config) (subscriber, error) {
  127. fb, err := firebase.NewApp(context.Background(), nil, option.WithCredentialsFile(conf.FirebaseKeyFile))
  128. if err != nil {
  129. return nil, err
  130. }
  131. msg, err := fb.Messaging(context.Background())
  132. if err != nil {
  133. return nil, err
  134. }
  135. return func(m *message) error {
  136. var data map[string]string // Matches https://ntfy.sh/docs/subscribe/api/#json-message-format
  137. switch m.Event {
  138. case keepaliveEvent, openEvent:
  139. data = map[string]string{
  140. "id": m.ID,
  141. "time": fmt.Sprintf("%d", m.Time),
  142. "event": m.Event,
  143. "topic": m.Topic,
  144. }
  145. case messageEvent:
  146. data = map[string]string{
  147. "id": m.ID,
  148. "time": fmt.Sprintf("%d", m.Time),
  149. "event": m.Event,
  150. "topic": m.Topic,
  151. "priority": fmt.Sprintf("%d", m.Priority),
  152. "tags": strings.Join(m.Tags, ","),
  153. "title": m.Title,
  154. "message": m.Message,
  155. }
  156. }
  157. _, err := msg.Send(context.Background(), &messaging.Message{
  158. Topic: m.Topic,
  159. Data: data,
  160. })
  161. return err
  162. }, nil
  163. }
  164. // Run executes the main server. It listens on HTTP (+ HTTPS, if configured), and starts
  165. // a manager go routine to print stats and prune messages.
  166. func (s *Server) Run() error {
  167. go s.runManager()
  168. go s.runAtSender()
  169. go s.runFirebaseKeepliver()
  170. listenStr := fmt.Sprintf("%s/http", s.config.ListenHTTP)
  171. if s.config.ListenHTTPS != "" {
  172. listenStr += fmt.Sprintf(" %s/https", s.config.ListenHTTPS)
  173. }
  174. log.Printf("Listening on %s", listenStr)
  175. http.HandleFunc("/", s.handle)
  176. errChan := make(chan error)
  177. go func() {
  178. errChan <- http.ListenAndServe(s.config.ListenHTTP, nil)
  179. }()
  180. if s.config.ListenHTTPS != "" {
  181. go func() {
  182. errChan <- http.ListenAndServeTLS(s.config.ListenHTTPS, s.config.CertFile, s.config.KeyFile, nil)
  183. }()
  184. }
  185. return <-errChan
  186. }
  187. func (s *Server) handle(w http.ResponseWriter, r *http.Request) {
  188. if err := s.handleInternal(w, r); err != nil {
  189. if e, ok := err.(*errHTTP); ok {
  190. s.fail(w, r, e.Code, e)
  191. } else {
  192. s.fail(w, r, http.StatusInternalServerError, err)
  193. }
  194. }
  195. }
  196. func (s *Server) handleInternal(w http.ResponseWriter, r *http.Request) error {
  197. if r.Method == http.MethodGet && r.URL.Path == "/" {
  198. return s.handleHome(w, r)
  199. } else if r.Method == http.MethodGet && r.URL.Path == "/example.html" {
  200. return s.handleExample(w, r)
  201. } else if r.Method == http.MethodHead && r.URL.Path == "/" {
  202. return s.handleEmpty(w, r)
  203. } else if r.Method == http.MethodGet && staticRegex.MatchString(r.URL.Path) {
  204. return s.handleStatic(w, r)
  205. } else if r.Method == http.MethodGet && docsRegex.MatchString(r.URL.Path) {
  206. return s.handleDocs(w, r)
  207. } else if r.Method == http.MethodOptions {
  208. return s.handleOptions(w, r)
  209. } else if r.Method == http.MethodGet && topicRegex.MatchString(r.URL.Path) {
  210. return s.handleHome(w, r)
  211. } else if (r.Method == http.MethodPut || r.Method == http.MethodPost) && topicRegex.MatchString(r.URL.Path) {
  212. return s.withRateLimit(w, r, s.handlePublish)
  213. } else if r.Method == http.MethodGet && sendRegex.MatchString(r.URL.Path) {
  214. return s.withRateLimit(w, r, s.handlePublish)
  215. } else if r.Method == http.MethodGet && jsonRegex.MatchString(r.URL.Path) {
  216. return s.withRateLimit(w, r, s.handleSubscribeJSON)
  217. } else if r.Method == http.MethodGet && sseRegex.MatchString(r.URL.Path) {
  218. return s.withRateLimit(w, r, s.handleSubscribeSSE)
  219. } else if r.Method == http.MethodGet && rawRegex.MatchString(r.URL.Path) {
  220. return s.withRateLimit(w, r, s.handleSubscribeRaw)
  221. }
  222. return errHTTPNotFound
  223. }
  224. func (s *Server) handleHome(w http.ResponseWriter, r *http.Request) error {
  225. return indexTemplate.Execute(w, &indexPage{
  226. Topic: r.URL.Path[1:],
  227. CacheDuration: s.config.CacheDuration,
  228. })
  229. }
  230. func (s *Server) handleEmpty(_ http.ResponseWriter, _ *http.Request) error {
  231. return nil
  232. }
  233. func (s *Server) handleExample(w http.ResponseWriter, _ *http.Request) error {
  234. _, err := io.WriteString(w, exampleSource)
  235. return err
  236. }
  237. func (s *Server) handleStatic(w http.ResponseWriter, r *http.Request) error {
  238. http.FileServer(http.FS(webStaticFsCached)).ServeHTTP(w, r)
  239. return nil
  240. }
  241. func (s *Server) handleDocs(w http.ResponseWriter, r *http.Request) error {
  242. http.FileServer(http.FS(docsStaticCached)).ServeHTTP(w, r)
  243. return nil
  244. }
  245. func (s *Server) handlePublish(w http.ResponseWriter, r *http.Request, _ *visitor) error {
  246. t, err := s.topicFromPath(r.URL.Path)
  247. if err != nil {
  248. return err
  249. }
  250. reader := io.LimitReader(r.Body, int64(s.config.MessageLimit))
  251. b, err := io.ReadAll(reader)
  252. if err != nil {
  253. return err
  254. }
  255. m := newDefaultMessage(t.ID, strings.TrimSpace(string(b)))
  256. cache, firebase, err := s.parseParams(r, m)
  257. if err != nil {
  258. return err
  259. }
  260. if m.Message == "" {
  261. m.Message = "triggered"
  262. }
  263. delayed := m.Time > time.Now().Unix()
  264. if !delayed {
  265. if err := t.Publish(m); err != nil {
  266. return err
  267. }
  268. }
  269. if s.firebase != nil && firebase && !delayed {
  270. go func() {
  271. if err := s.firebase(m); err != nil {
  272. log.Printf("Unable to publish to Firebase: %v", err.Error())
  273. }
  274. }()
  275. }
  276. if cache {
  277. if err := s.cache.AddMessage(m); err != nil {
  278. return err
  279. }
  280. }
  281. w.Header().Set("Content-Type", "application/json")
  282. w.Header().Set("Access-Control-Allow-Origin", "*") // CORS, allow cross-origin requests
  283. if err := json.NewEncoder(w).Encode(m); err != nil {
  284. return err
  285. }
  286. s.inc(&s.messages)
  287. return nil
  288. }
  289. func (s *Server) parseParams(r *http.Request, m *message) (cache bool, firebase bool, err error) {
  290. cache = readParam(r, "x-cache", "cache") != "no"
  291. firebase = readParam(r, "x-firebase", "firebase") != "no"
  292. m.Title = readParam(r, "x-title", "title", "t")
  293. messageStr := readParam(r, "x-message", "message", "m")
  294. if messageStr != "" {
  295. m.Message = messageStr
  296. }
  297. m.Priority, err = util.ParsePriority(readParam(r, "x-priority", "priority", "prio", "p"))
  298. if err != nil {
  299. return false, false, errHTTPBadRequest
  300. }
  301. tagsStr := readParam(r, "x-tags", "tags", "tag", "ta")
  302. if tagsStr != "" {
  303. m.Tags = make([]string, 0)
  304. for _, s := range util.SplitNoEmpty(tagsStr, ",") {
  305. m.Tags = append(m.Tags, strings.TrimSpace(s))
  306. }
  307. }
  308. delayStr := readParam(r, "x-delay", "delay", "x-at", "at", "x-in", "in")
  309. if delayStr != "" {
  310. if !cache {
  311. return false, false, errHTTPBadRequest
  312. }
  313. delay, err := util.ParseFutureTime(delayStr, time.Now())
  314. if err != nil {
  315. return false, false, errHTTPBadRequest
  316. } else if delay.Unix() < time.Now().Add(s.config.MinDelay).Unix() {
  317. return false, false, errHTTPBadRequest
  318. } else if delay.Unix() > time.Now().Add(s.config.MaxDelay).Unix() {
  319. return false, false, errHTTPBadRequest
  320. }
  321. m.Time = delay.Unix()
  322. }
  323. return cache, firebase, nil
  324. }
  325. func readParam(r *http.Request, names ...string) string {
  326. for _, name := range names {
  327. value := r.Header.Get(name)
  328. if value != "" {
  329. return strings.TrimSpace(value)
  330. }
  331. }
  332. for _, name := range names {
  333. value := r.URL.Query().Get(strings.ToLower(name))
  334. if value != "" {
  335. return strings.TrimSpace(value)
  336. }
  337. }
  338. return ""
  339. }
  340. func (s *Server) handleSubscribeJSON(w http.ResponseWriter, r *http.Request, v *visitor) error {
  341. encoder := func(msg *message) (string, error) {
  342. var buf bytes.Buffer
  343. if err := json.NewEncoder(&buf).Encode(&msg); err != nil {
  344. return "", err
  345. }
  346. return buf.String(), nil
  347. }
  348. return s.handleSubscribe(w, r, v, "json", "application/x-ndjson", encoder)
  349. }
  350. func (s *Server) handleSubscribeSSE(w http.ResponseWriter, r *http.Request, v *visitor) error {
  351. encoder := func(msg *message) (string, error) {
  352. var buf bytes.Buffer
  353. if err := json.NewEncoder(&buf).Encode(&msg); err != nil {
  354. return "", err
  355. }
  356. if msg.Event != messageEvent {
  357. return fmt.Sprintf("event: %s\ndata: %s\n", msg.Event, buf.String()), nil // Browser's .onmessage() does not fire on this!
  358. }
  359. return fmt.Sprintf("data: %s\n", buf.String()), nil
  360. }
  361. return s.handleSubscribe(w, r, v, "sse", "text/event-stream", encoder)
  362. }
  363. func (s *Server) handleSubscribeRaw(w http.ResponseWriter, r *http.Request, v *visitor) error {
  364. encoder := func(msg *message) (string, error) {
  365. if msg.Event == messageEvent { // only handle default events
  366. return strings.ReplaceAll(msg.Message, "\n", " ") + "\n", nil
  367. }
  368. return "\n", nil // "keepalive" and "open" events just send an empty line
  369. }
  370. return s.handleSubscribe(w, r, v, "raw", "text/plain", encoder)
  371. }
  372. func (s *Server) handleSubscribe(w http.ResponseWriter, r *http.Request, v *visitor, format string, contentType string, encoder messageEncoder) error {
  373. if err := v.AddSubscription(); err != nil {
  374. return errHTTPTooManyRequests
  375. }
  376. defer v.RemoveSubscription()
  377. topicsStr := strings.TrimSuffix(r.URL.Path[1:], "/"+format) // Hack
  378. topicIDs := util.SplitNoEmpty(topicsStr, ",")
  379. topics, err := s.topicsFromIDs(topicIDs...)
  380. if err != nil {
  381. return err
  382. }
  383. poll := readParam(r, "x-poll", "poll", "po") == "1"
  384. scheduled := readParam(r, "x-scheduled", "scheduled", "sched") == "1"
  385. since, err := parseSince(r, poll)
  386. if err != nil {
  387. return err
  388. }
  389. messageFilter, titleFilter, priorityFilter, tagsFilter, err := parseQueryFilters(r)
  390. if err != nil {
  391. return err
  392. }
  393. var wlock sync.Mutex
  394. sub := func(msg *message) error {
  395. if !passesQueryFilter(msg, messageFilter, titleFilter, priorityFilter, tagsFilter) {
  396. return nil
  397. }
  398. m, err := encoder(msg)
  399. if err != nil {
  400. return err
  401. }
  402. wlock.Lock()
  403. defer wlock.Unlock()
  404. if _, err := w.Write([]byte(m)); err != nil {
  405. return err
  406. }
  407. if fl, ok := w.(http.Flusher); ok {
  408. fl.Flush()
  409. }
  410. return nil
  411. }
  412. w.Header().Set("Access-Control-Allow-Origin", "*") // CORS, allow cross-origin requests
  413. w.Header().Set("Content-Type", contentType+"; charset=utf-8") // Android/Volley client needs charset!
  414. if poll {
  415. return s.sendOldMessages(topics, since, scheduled, sub)
  416. }
  417. subscriberIDs := make([]int, 0)
  418. for _, t := range topics {
  419. subscriberIDs = append(subscriberIDs, t.Subscribe(sub))
  420. }
  421. defer func() {
  422. for i, subscriberID := range subscriberIDs {
  423. topics[i].Unsubscribe(subscriberID) // Order!
  424. }
  425. }()
  426. if err := sub(newOpenMessage(topicsStr)); err != nil { // Send out open message
  427. return err
  428. }
  429. if err := s.sendOldMessages(topics, since, scheduled, sub); err != nil {
  430. return err
  431. }
  432. for {
  433. select {
  434. case <-r.Context().Done():
  435. return nil
  436. case <-time.After(s.config.KeepaliveInterval):
  437. v.Keepalive()
  438. if err := sub(newKeepaliveMessage(topicsStr)); err != nil { // Send keepalive message
  439. return err
  440. }
  441. }
  442. }
  443. }
  444. func parseQueryFilters(r *http.Request) (messageFilter string, titleFilter string, priorityFilter []int, tagsFilter []string, err error) {
  445. messageFilter = readParam(r, "x-message", "message", "m")
  446. titleFilter = readParam(r, "x-title", "title", "t")
  447. tagsFilter = util.SplitNoEmpty(readParam(r, "x-tags", "tags", "tag", "ta"), ",")
  448. priorityFilter = make([]int, 0)
  449. for _, p := range util.SplitNoEmpty(readParam(r, "x-priority", "priority", "prio", "p"), ",") {
  450. priority, err := util.ParsePriority(p)
  451. if err != nil {
  452. return "", "", nil, nil, err
  453. }
  454. priorityFilter = append(priorityFilter, priority)
  455. }
  456. return
  457. }
  458. func passesQueryFilter(msg *message, messageFilter string, titleFilter string, priorityFilter []int, tagsFilter []string) bool {
  459. if msg.Event != messageEvent {
  460. return true // filters only apply to messages
  461. }
  462. if messageFilter != "" && msg.Message != messageFilter {
  463. return false
  464. }
  465. if titleFilter != "" && msg.Title != titleFilter {
  466. return false
  467. }
  468. messagePriority := msg.Priority
  469. if messagePriority == 0 {
  470. messagePriority = 3 // For query filters, default priority (3) is the same as "not set" (0)
  471. }
  472. if len(priorityFilter) > 0 && !util.InIntList(priorityFilter, messagePriority) {
  473. return false
  474. }
  475. if len(tagsFilter) > 0 && !util.InStringListAll(msg.Tags, tagsFilter) {
  476. return false
  477. }
  478. return true
  479. }
  480. func (s *Server) sendOldMessages(topics []*topic, since sinceTime, scheduled bool, sub subscriber) error {
  481. if since.IsNone() {
  482. return nil
  483. }
  484. for _, t := range topics {
  485. messages, err := s.cache.Messages(t.ID, since, scheduled)
  486. if err != nil {
  487. return err
  488. }
  489. for _, m := range messages {
  490. if err := sub(m); err != nil {
  491. return err
  492. }
  493. }
  494. }
  495. return nil
  496. }
  497. // parseSince returns a timestamp identifying the time span from which cached messages should be received.
  498. //
  499. // Values in the "since=..." parameter can be either a unix timestamp or a duration (e.g. 12h), or
  500. // "all" for all messages.
  501. func parseSince(r *http.Request, poll bool) (sinceTime, error) {
  502. since := readParam(r, "x-since", "since", "si")
  503. if since == "" {
  504. if poll {
  505. return sinceAllMessages, nil
  506. }
  507. return sinceNoMessages, nil
  508. }
  509. if since == "all" {
  510. return sinceAllMessages, nil
  511. } else if s, err := strconv.ParseInt(since, 10, 64); err == nil {
  512. return sinceTime(time.Unix(s, 0)), nil
  513. } else if d, err := time.ParseDuration(since); err == nil {
  514. return sinceTime(time.Now().Add(-1 * d)), nil
  515. }
  516. return sinceNoMessages, errHTTPBadRequest
  517. }
  518. func (s *Server) handleOptions(w http.ResponseWriter, _ *http.Request) error {
  519. w.Header().Set("Access-Control-Allow-Origin", "*") // CORS, allow cross-origin requests
  520. w.Header().Set("Access-Control-Allow-Methods", "GET, PUT, POST")
  521. return nil
  522. }
  523. func (s *Server) topicFromPath(path string) (*topic, error) {
  524. parts := strings.Split(path, "/")
  525. if len(parts) < 2 {
  526. return nil, errHTTPBadRequest
  527. }
  528. topics, err := s.topicsFromIDs(parts[1])
  529. if err != nil {
  530. return nil, err
  531. }
  532. return topics[0], nil
  533. }
  534. func (s *Server) topicsFromIDs(ids ...string) ([]*topic, error) {
  535. s.mu.Lock()
  536. defer s.mu.Unlock()
  537. topics := make([]*topic, 0)
  538. for _, id := range ids {
  539. if util.InStringList(disallowedTopics, id) {
  540. return nil, errHTTPBadRequest
  541. }
  542. if _, ok := s.topics[id]; !ok {
  543. if len(s.topics) >= s.config.GlobalTopicLimit {
  544. return nil, errHTTPTooManyRequests
  545. }
  546. s.topics[id] = newTopic(id)
  547. }
  548. topics = append(topics, s.topics[id])
  549. }
  550. return topics, nil
  551. }
  552. func (s *Server) updateStatsAndPrune() {
  553. s.mu.Lock()
  554. defer s.mu.Unlock()
  555. // Expire visitors from rate visitors map
  556. for ip, v := range s.visitors {
  557. if v.Stale() {
  558. delete(s.visitors, ip)
  559. }
  560. }
  561. // Prune message cache
  562. olderThan := time.Now().Add(-1 * s.config.CacheDuration)
  563. if err := s.cache.Prune(olderThan); err != nil {
  564. log.Printf("error pruning cache: %s", err.Error())
  565. }
  566. // Prune old topics, remove subscriptions without subscribers
  567. var subscribers, messages int
  568. for _, t := range s.topics {
  569. subs := t.Subscribers()
  570. msgs, err := s.cache.MessageCount(t.ID)
  571. if err != nil {
  572. log.Printf("cannot get stats for topic %s: %s", t.ID, err.Error())
  573. continue
  574. }
  575. if msgs == 0 && subs == 0 {
  576. delete(s.topics, t.ID)
  577. continue
  578. }
  579. subscribers += subs
  580. messages += msgs
  581. }
  582. // Print stats
  583. log.Printf("Stats: %d message(s) published, %d topic(s) active, %d subscriber(s), %d message(s) buffered, %d visitor(s)",
  584. s.messages, len(s.topics), subscribers, messages, len(s.visitors))
  585. }
  586. func (s *Server) runManager() {
  587. func() {
  588. ticker := time.NewTicker(s.config.ManagerInterval)
  589. for {
  590. <-ticker.C
  591. s.updateStatsAndPrune()
  592. }
  593. }()
  594. }
  595. func (s *Server) runAtSender() {
  596. ticker := time.NewTicker(s.config.AtSenderInterval)
  597. for {
  598. <-ticker.C
  599. if err := s.sendDelayedMessages(); err != nil {
  600. log.Printf("error sending scheduled messages: %s", err.Error())
  601. }
  602. }
  603. }
  604. func (s *Server) runFirebaseKeepliver() {
  605. if s.firebase == nil {
  606. return
  607. }
  608. ticker := time.NewTicker(s.config.FirebaseKeepaliveInterval)
  609. for {
  610. <-ticker.C
  611. if err := s.firebase(newKeepaliveMessage(firebaseControlTopic)); err != nil {
  612. log.Printf("error sending Firebase keepalive message: %s", err.Error())
  613. }
  614. }
  615. }
  616. func (s *Server) sendDelayedMessages() error {
  617. s.mu.Lock()
  618. defer s.mu.Unlock()
  619. messages, err := s.cache.MessagesDue()
  620. if err != nil {
  621. return err
  622. }
  623. for _, m := range messages {
  624. t, ok := s.topics[m.Topic] // If no subscribers, just mark message as published
  625. if ok {
  626. if err := t.Publish(m); err != nil {
  627. log.Printf("unable to publish message %s to topic %s: %v", m.ID, m.Topic, err.Error())
  628. }
  629. if s.firebase != nil {
  630. if err := s.firebase(m); err != nil {
  631. log.Printf("unable to publish to Firebase: %v", err.Error())
  632. }
  633. }
  634. }
  635. if err := s.cache.MarkPublished(m); err != nil {
  636. return err
  637. }
  638. }
  639. return nil
  640. }
  641. func (s *Server) withRateLimit(w http.ResponseWriter, r *http.Request, handler func(w http.ResponseWriter, r *http.Request, v *visitor) error) error {
  642. v := s.visitor(r)
  643. if err := v.RequestAllowed(); err != nil {
  644. return err
  645. }
  646. return handler(w, r, v)
  647. }
  648. // visitor creates or retrieves a rate.Limiter for the given visitor.
  649. // This function was taken from https://www.alexedwards.net/blog/how-to-rate-limit-http-requests (MIT).
  650. func (s *Server) visitor(r *http.Request) *visitor {
  651. s.mu.Lock()
  652. defer s.mu.Unlock()
  653. remoteAddr := r.RemoteAddr
  654. ip, _, err := net.SplitHostPort(remoteAddr)
  655. if err != nil {
  656. ip = remoteAddr // This should not happen in real life; only in tests.
  657. }
  658. if s.config.BehindProxy && r.Header.Get("X-Forwarded-For") != "" {
  659. ip = r.Header.Get("X-Forwarded-For")
  660. }
  661. v, exists := s.visitors[ip]
  662. if !exists {
  663. s.visitors[ip] = newVisitor(s.config)
  664. return s.visitors[ip]
  665. }
  666. v.Keepalive()
  667. return v
  668. }
  669. func (s *Server) inc(counter *int64) {
  670. s.mu.Lock()
  671. defer s.mu.Unlock()
  672. *counter++
  673. }
  674. func (s *Server) fail(w http.ResponseWriter, r *http.Request, code int, err error) {
  675. log.Printf("[%s] %s - %d - %s", r.RemoteAddr, r.Method, code, err.Error())
  676. w.WriteHeader(code)
  677. _, _ = io.WriteString(w, fmt.Sprintf("%s\n", http.StatusText(code)))
  678. }