pomerium/authorize/internal/store/index.go

package store

import (
	"sync"

	"github.com/kentik/patricia"
	"github.com/kentik/patricia/string_tree"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/structpb"
)

const (
	indexField = "$index"
	cidrField  = "cidr"
)

type index struct {
	mu     sync.RWMutex
	byType map[string]*recordIndex
}

func newIndex() *index {
	idx := new(index)
	idx.clear()
	return idx
}

func (idx *index) clear() {
	idx.mu.Lock()
	defer idx.mu.Unlock()
	idx.byType = map[string]*recordIndex{}
}

func (idx *index) delete(typeURL, id string) {
	idx.mu.Lock()
	defer idx.mu.Unlock()

	ridx, ok := idx.byType[typeURL]
	if !ok {
		return
	}
	ridx.delete(id)

	if len(ridx.byID) == 0 {
		delete(idx.byType, typeURL)
	}
}

func (idx *index) find(typeURL, id string) proto.Message {
	idx.mu.RLock()
	defer idx.mu.RUnlock()

	ridx, ok := idx.byType[typeURL]
	if !ok {
		return nil
	}
	return ridx.find(id)
}

func (idx *index) get(typeURL, id string) proto.Message {
	idx.mu.RLock()
	defer idx.mu.RUnlock()

	ridx, ok := idx.byType[typeURL]
	if !ok {
		return nil
	}
	return ridx.get(id)
}

func (idx *index) set(typeURL, id string, msg proto.Message) {
	idx.mu.Lock()
	defer idx.mu.Unlock()

	ridx, ok := idx.byType[typeURL]
	if !ok {
		ridx = newRecordIndex()
		idx.byType[typeURL] = ridx
	}
	ridx.set(id, msg)
}

// a recordIndex indexes records for of a specific type
type recordIndex struct {
	byID     map[string]proto.Message
	byCIDRV4 *string_tree.TreeV4
	byCIDRV6 *string_tree.TreeV6
}

// newRecordIndex creates a new record index.
func newRecordIndex() *recordIndex {
	return &recordIndex{
		byID:     map[string]proto.Message{},
		byCIDRV4: string_tree.NewTreeV4(),
		byCIDRV6: string_tree.NewTreeV6(),
	}
}

func (idx *recordIndex) delete(id string) {
	r, ok := idx.byID[id]
	if !ok {
		return
	}

	delete(idx.byID, id)

	addr4, addr6 := getIndexCIDR(r)
	if addr4 != nil {
		idx.byCIDRV4.Delete(*addr4, func(payload, val string) bool {
			return payload == val
		}, id)
	}
	if addr6 != nil {
		idx.byCIDRV6.Delete(*addr6, func(payload, val string) bool {
			return payload == val
		}, id)
	}
}

func (idx *recordIndex) find(idOrString string) proto.Message {
	r, ok := idx.byID[idOrString]
	if ok {
		return r
	}

	addrv4, addrv6, _ := patricia.ParseIPFromString(idOrString)
	if addrv4 != nil {
		found, id := idx.byCIDRV4.FindDeepestTag(*addrv4)
		if found {
			return idx.byID[id]
		}
	}
	if addrv6 != nil {
		found, id := idx.byCIDRV6.FindDeepestTag(*addrv6)
		if found {
			return idx.byID[id]
		}
	}

	return nil
}

func (idx *recordIndex) get(id string) proto.Message {
	return idx.byID[id]
}

func (idx *recordIndex) set(id string, msg proto.Message) {
	_, ok := idx.byID[id]
	if ok {
		idx.delete(id)
	}

	idx.byID[id] = msg
	addr4, addr6 := getIndexCIDR(msg)
	if addr4 != nil {
		idx.byCIDRV4.Set(*addr4, id)
	}
	if addr6 != nil {
		idx.byCIDRV6.Set(*addr6, id)
	}
}

func getIndexCIDR(msg proto.Message) (*patricia.IPv4Address, *patricia.IPv6Address) {
	var s *structpb.Struct
	if sv, ok := msg.(*structpb.Value); ok {
		s = sv.GetStructValue()
	} else {
		s, _ = msg.(*structpb.Struct)
	}
	if s == nil {
		return nil, nil
	}

	f, ok := s.Fields[indexField]
	if !ok {
		return nil, nil
	}

	obj := f.GetStructValue()
	if obj == nil {
		return nil, nil
	}

	cf, ok := obj.Fields[cidrField]
	if !ok {
		return nil, nil
	}

	c := cf.GetStringValue()
	if c == "" {
		return nil, nil
	}

	addr4, addr6, _ := patricia.ParseIPFromString(c)
	return addr4, addr6
}