Proton-API-Bridge/cache.go

package proton_api_bridge

import (
	"context"
	"log"
	"sync"

	"github.com/ProtonMail/gopenpgp/v2/crypto"
	"github.com/ProtonMail/go-proton-api"
)

type cacheEntry struct {
	link *proton.Link
	kr   *crypto.KeyRing
}

type cache struct {
	data          map[string]*cacheEntry
	children      map[string]map[string]interface{}
	enableCaching bool

	sync.RWMutex
}

func newCache(enableCaching bool) *cache {
	return &cache{
		data:          make(map[string]*cacheEntry),
		children:      make(map[string]map[string]interface{}),
		enableCaching: enableCaching,
	}
}

// func (cache *cache) _debug() {
// 	if !cache.enableCaching {
// 		return
// 	}

// 	cache.RLock()
// 	defer cache.RUnlock()

// 	for k, v := range cache.data {
// 		log.Printf("data %#v %p", k, v.link)
// 	}
// 	for k, v := range cache.children {
// 		log.Printf("children %#v %#v", k, v)
// 	}
// }

func (cache *cache) _get(linkID string) *cacheEntry {
	if !cache.enableCaching {
		return nil
	}

	cache.RLock()
	defer cache.RUnlock()

	if data, ok := cache.data[linkID]; ok {
		return data
	}
	return nil
}

func (cache *cache) _insert(linkID string, link *proton.Link, kr *crypto.KeyRing) {
	if !cache.enableCaching {
		return
	}

	cache.Lock()
	defer cache.Unlock()

	cache.data[linkID] = &cacheEntry{
		link: link,
		kr:   kr,
	}

	if link != nil {
		if data, ok := cache.children[link.ParentLinkID]; ok {
			data[link.LinkID] = nil
			cache.children[link.ParentLinkID] = data
		} else {
			tmp := make(map[string]interface{})
			tmp[link.LinkID] = nil
			cache.children[link.ParentLinkID] = tmp

		}
	} else {
		// TODO: we should never have missing link though
		log.Fatalln("we should never have missing link though")
	}
}

// due to recursion, we can't perform locking here
// this function should only be called from _remove
func (cache *cache) _remove_nolock(linkID string, includingChildren bool) {
	var link *proton.Link
	if data, ok := cache.data[linkID]; ok {
		link = data.link
		delete(cache.data, linkID)
	} else {
		return
	}

	// remove linkID from parent's map
	if data, ok := cache.children[link.ParentLinkID]; ok {
		if _, ok := data[link.LinkID]; ok {
			delete(data, link.LinkID)
			cache.children[link.ParentLinkID] = data
		} else {
			log.Fatalln("we have an issue for cache inconsistency where link is not found in the parent's map")
		}
	} else {
		log.Fatalln("we have an issue for cache inconsistency where link's parent map is missing")
	}

	// we don't recursively go upward to clean up the parent's map
	// instead, we rely on periodic cache flushing
	if includingChildren {
		if data, ok := cache.children[link.LinkID]; ok {
			for k := range data {
				cache._remove_nolock(k, true)
			}
		} // else {
		// might have nothing is the link doesn't have any children
		// }
		delete(cache.children, link.LinkID)
	}
}

func (cache *cache) _remove(linkID string, includingChildren bool) {
	if !cache.enableCaching {
		return
	}

	cache.Lock()
	defer cache.Unlock()

	cache._remove_nolock(linkID, includingChildren)
}

/* The original non-caching version, which resolves the keyring recursively */
func (protonDrive *ProtonDrive) _getLinkKRByID(ctx context.Context, linkID string) (*crypto.KeyRing, error) {
	if linkID == "" {
		// most likely someone requested root link's parent link, which happen to be ""
		// return protonDrive.MainShareKR.Copy() // we need to return a deep copy since the keyring will be freed by the caller when it finishes using the keyring -> now we go through caching, so we won't clear kr

		return protonDrive.MainShareKR, nil
	}

	link, err := protonDrive.getLink(ctx, linkID)
	if err != nil {
		return nil, err
	}

	return protonDrive._getLinkKR(ctx, link)
}

/* The original non-caching version, which resolves the keyring recursively */
func (protonDrive *ProtonDrive) _getLinkKR(ctx context.Context, link *proton.Link) (*crypto.KeyRing, error) {
	if link.ParentLinkID == "" { // link is rootLink
		signatureVerificationKR, err := protonDrive.getSignatureVerificationKeyring([]string{link.SignatureEmail})
		if err != nil {
			return nil, err
		}
		nodeKR, err := link.GetKeyRing(protonDrive.MainShareKR, signatureVerificationKR)
		if err != nil {
			return nil, err
		}

		return nodeKR, nil
	}

	parentLink, err := protonDrive.getLink(ctx, link.ParentLinkID)
	if err != nil {
		return nil, err
	}

	// parentNodeKR is used to decrypt the current node's KR, as each node has its keyring, which can be decrypted by its parent
	parentNodeKR, err := protonDrive._getLinkKR(ctx, parentLink)
	if err != nil {
		return nil, err
	}

	signatureVerificationKR, err := protonDrive.getSignatureVerificationKeyring([]string{link.SignatureEmail})
	if err != nil {
		return nil, err
	}
	nodeKR, err := link.GetKeyRing(parentNodeKR, signatureVerificationKR)
	if err != nil {
		return nil, err
	}

	return nodeKR, nil
}

func (protonDrive *ProtonDrive) getLink(ctx context.Context, linkID string) (*proton.Link, error) {
	if linkID == "" {
		// this is only possible when doing rootLink's parent, which should be handled beforehand
		return nil, ErrWrongUsageOfGetLink
	}

	// attempt to get from cache first
	if data := protonDrive.cache._get(linkID); data != nil && data.link != nil {
		return data.link, nil
	}

	// no cached data, fetch
	link, err := protonDrive.c.GetLink(ctx, protonDrive.MainShare.ShareID, linkID)
	if err != nil {
		return nil, err
	}

	// populate cache
	protonDrive.cache._insert(linkID, &link, nil)

	return &link, nil
}

func (protonDrive *ProtonDrive) getLinkKR(ctx context.Context, link *proton.Link) (*crypto.KeyRing, error) {
	if !protonDrive.cache.enableCaching {
		return protonDrive._getLinkKR(ctx, link)
	}

	if link == nil {
		return nil, ErrWrongUsageOfGetLinkKR
	}

	// attempt to get from cache first
	if data := protonDrive.cache._get(link.LinkID); data != nil && data.link != nil {
		if data.kr != nil {
			return data.kr, nil
		}

		// decrypt keyring and cache it
		parentNodeKR, err := protonDrive.getLinkKRByID(ctx, data.link.ParentLinkID)
		if err != nil {
			return nil, err
		}

		signatureVerificationKR, err := protonDrive.getSignatureVerificationKeyring([]string{data.link.SignatureEmail})
		if err != nil {
			return nil, err
		}
		kr, err := data.link.GetKeyRing(parentNodeKR, signatureVerificationKR)
		if err != nil {
			return nil, err
		}
		data.kr = kr
		return data.kr, nil
	}

	// no cached data, fetch
	protonDrive.cache._insert(link.LinkID, link, nil)

	return protonDrive.getLinkKR(ctx, link)
}

func (protonDrive *ProtonDrive) getLinkKRByID(ctx context.Context, linkID string) (*crypto.KeyRing, error) {
	if !protonDrive.cache.enableCaching {
		return protonDrive._getLinkKRByID(ctx, linkID)
	}

	if linkID == "" {
		return protonDrive.MainShareKR, nil
	}

	// attempt to get from cache first
	if data := protonDrive.cache._get(linkID); data != nil && data.link != nil {
		return protonDrive.getLinkKR(ctx, data.link)
	}

	// log.Println("Not from cache")
	// no cached data, fetch
	link, err := protonDrive.getLink(ctx, linkID)
	if err != nil {
		return nil, err
	}

	return protonDrive.getLinkKR(ctx, link)
}

func (protonDrive *ProtonDrive) removeLinkIDFromCache(linkID string, includingChildren bool) {
	if !protonDrive.cache.enableCaching {
		return
	}

	// log.Println("===================================")
	// log.Println(linkID, includingChildren)
	// protonDrive.cache._debug()
	protonDrive.cache._remove(linkID, includingChildren)
	// protonDrive.cache._debug()
	// log.Println("===================================")
}

func (protonDrive *ProtonDrive) ClearCache() {
	if !protonDrive.cache.enableCaching {
		return
	}

	protonDrive.cache.Lock()
	defer protonDrive.cache.Unlock()

	// TODO: we come back to fix this later
	// for _, entry := range protonDrive.cache.data {
	// 	entry.kr.ClearPrivateParams()
	// }

	protonDrive.cache.data = make(map[string]*cacheEntry)
	protonDrive.cache.children = make(map[string]map[string]interface{})
}