From e4ba0d117537f937c50e91dca692aea04751fb96 Mon Sep 17 00:00:00 2001
From: Frederik Schwan <frederik.schwan@linux.com>
Date: Sat, 2 Jan 2021 03:14:59 +0100
Subject: [PATCH] add subfolder push

---
 internal/apns.go   | 61 +++++++++++++++++++++++++++++++++++++++++-----
 internal/socket.go | 11 +++------
 2 files changed, 58 insertions(+), 14 deletions(-)

diff --git a/internal/apns.go b/internal/apns.go
index 2aba60f..b740237 100644
--- a/internal/apns.go
+++ b/internal/apns.go
@@ -1,8 +1,11 @@
 package internal
 
 import (
+	"crypto/md5"
 	"crypto/tls"
 	"crypto/x509"
+	"encoding/hex"
+	"encoding/json"
 	"errors"
 	"github.com/freswa/dovecot-xaps-daemon/internal/config"
 	"github.com/freswa/dovecot-xaps-daemon/internal/database"
@@ -11,6 +14,7 @@ import (
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/net/http2"
 	"net/http"
+	"strings"
 	"sync"
 	"time"
 )
@@ -34,9 +38,15 @@ type Apns struct {
 	db                   *database.Database
 	mapMutex             sync.Mutex
 	delayedApns          map[database.Registration]time.Time
+	delayedMailboxes     map[database.Registration]PushedMailboxes
 	RenewTimer           *time.Timer
 }
 
+type PushedMailboxes struct {
+	mailboxes    []string
+	registration database.Registration
+}
+
 func NewApns(cfg *config.Config, db *database.Database) (apns *Apns) {
 	apns = &Apns{
 		DelayTime:            cfg.Delay,
@@ -44,6 +54,7 @@ func NewApns(cfg *config.Config, db *database.Database) (apns *Apns) {
 		db:                   db,
 		mapMutex:             sync.Mutex{},
 		delayedApns:          make(map[database.Registration]time.Time),
+		delayedMailboxes:     make(map[database.Registration]PushedMailboxes),
 	}
 	log.Debugln("APNS for non NewMessage events will be delayed for", time.Second*time.Duration(apns.DelayTime))
 	log.Debugln("Trying to get existing certs from the DB")
@@ -136,18 +147,25 @@ func (apns *Apns) checkDelayed() {
 	}
 	apns.mapMutex.Unlock()
 	for _, reg := range sendNow {
-		apns.SendNotification(reg, false)
+		apns.mapMutex.Lock()
+		mbox := apns.delayedMailboxes[reg]
+		apns.mapMutex.Unlock()
+		apns.SendNotification(reg, mbox, false)
 	}
 }
 
-func (apns *Apns) SendNotification(registration database.Registration, delayed bool) {
+func (apns *Apns) SendNotification(registration database.Registration, mailboxes PushedMailboxes, delayed bool) {
 	apns.mapMutex.Lock()
 	if delayed {
+		log.Debugln("Delaying notification for", registration.AccountId, "/", registration.DeviceToken, "for mailboxes", mailboxes.toString())
 		apns.delayedApns[registration] = time.Now()
+		apns.delayedMailboxes[registration] = mailboxes.merge(apns.delayedMailboxes[registration])
 		apns.mapMutex.Unlock()
 		return
 	} else {
 		delete(apns.delayedApns, registration)
+		mailboxes = mailboxes.merge(apns.delayedMailboxes[registration])
+		delete(apns.delayedMailboxes, registration)
 		apns.mapMutex.Unlock()
 	}
 	log.Debugln("Sending notification to", registration.AccountId, "/", registration.DeviceToken)
@@ -157,6 +175,9 @@ func (apns *Apns) SendNotification(registration database.Registration, delayed b
 	notification.Topic = apns.Topic
 	composedPayload := []byte(`{"aps":{`)
 	composedPayload = append(composedPayload, []byte(`"account-id":"`+registration.AccountId+`"`)...)
+	composedPayload = append(composedPayload, []byte(`,"m":`)...)
+	m, _ := json.Marshal(mailboxes.toHashes())
+	composedPayload = append(composedPayload, m...)
 	composedPayload = append(composedPayload, []byte(`}}`)...)
 	notification.Payload = composedPayload
 	notification.PushType = apns2.PushTypeBackground
@@ -164,10 +185,6 @@ func (apns *Apns) SendNotification(registration database.Registration, delayed b
 	// set the apns-priority
 	//notification.Priority = apns2.PriorityLow
 
-	if log.IsLevelEnabled(log.DebugLevel) {
-		dbgstr, _ := notification.MarshalJSON()
-		log.Debugf("Sending: %s", dbgstr)
-	}
 	res, err := apns.client.Push(notification)
 
 	if err != nil {
@@ -220,3 +237,35 @@ func certificateNotValidAfter(certificate *tls.Certificate) time.Time {
 	}
 	return cert.NotAfter
 }
+
+func (pm *PushedMailboxes) AddMailbox(mailbox string) {
+	present := false
+	for _, mb := range pm.mailboxes {
+		if mb == mailbox {
+			present = true
+		}
+	}
+	if !present {
+		pm.mailboxes = append(pm.mailboxes, mailbox)
+	}
+}
+
+func (pm *PushedMailboxes) toHashes() []string {
+	var mbh = make([]string, 0)
+	for _, m := range pm.mailboxes {
+		md5sum := md5.Sum([]byte(m))
+		mbh = append(mbh, hex.EncodeToString(md5sum[:]))
+	}
+	return mbh
+}
+
+func (pm *PushedMailboxes) toString() string {
+	return strings.Join(pm.mailboxes, ", ")
+}
+
+func (pm *PushedMailboxes) merge(anotherPm PushedMailboxes) PushedMailboxes {
+	for _, mb := range anotherPm.mailboxes {
+		pm.AddMailbox(mb)
+	}
+	return *pm
+}
diff --git a/internal/socket.go b/internal/socket.go
index 9dda8fb..725da12 100644
--- a/internal/socket.go
+++ b/internal/socket.go
@@ -146,13 +146,6 @@ func (httpHandler *httpHandler) handleNotify(writer http.ResponseWriter, request
 		}
 	}
 
-	// we don't know how to handle other mailboxes other than INBOX, so ignore them
-	if notify.Mailbox != "INBOX" {
-		log.Debugln("Ignoring non INBOX event for:", notify.Mailbox)
-		writer.WriteHeader(http.StatusOK)
-		return
-	}
-
 	// Find all the devices registered for this mailbox event
 	registrations, err := httpHandler.db.FindRegistrations(notify.Username, notify.Mailbox)
 	if err != nil {
@@ -179,7 +172,9 @@ func (httpHandler *httpHandler) handleNotify(writer http.ResponseWriter, request
 	// Send a notification to all registered devices. We ignore failures
 	// because there is not a lot we can do.
 	for _, registration := range registrations {
-		httpHandler.apns.SendNotification(registration, !isMessageNew)
+		pm := PushedMailboxes{}
+		pm.AddMailbox(notify.Mailbox)
+		httpHandler.apns.SendNotification(registration, pm, !isMessageNew)
 	}
 
 	writer.WriteHeader(http.StatusOK)