mirror of
https://github.com/tailscale/tailscale.git
synced 2024-12-05 16:05:36 +00:00
162d593514
I meant to do this in the earlier change and had a git fail. To atone, add a test too while I'm here. Updates #12486 Updates #12507 Change-Id: I4943b454a2530cb5047636f37136aa2898d2ffc7 Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
166 lines
4.3 KiB
Go
166 lines
4.3 KiB
Go
// Copyright (c) Tailscale Inc & AUTHORS
|
|
// SPDX-License-Identifier: BSD-3-Clause
|
|
//
|
|
// Original implementation (from same author) from which this was derived was:
|
|
// https://github.com/golang/groupcache/blob/5b532d6fd5efaf7fa130d4e859a2fde0fc3a9e1b/lru/lru.go
|
|
// ... which was Apache licensed:
|
|
// https://github.com/golang/groupcache/blob/master/LICENSE
|
|
|
|
// Package flowtrack contains types for tracking TCP/UDP flows by 4-tuples.
|
|
package flowtrack
|
|
|
|
import (
|
|
"container/list"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/netip"
|
|
|
|
"tailscale.com/types/ipproto"
|
|
)
|
|
|
|
// MakeTuple makes a Tuple out of netip.AddrPort values.
|
|
func MakeTuple(proto ipproto.Proto, src, dst netip.AddrPort) Tuple {
|
|
return Tuple{
|
|
proto: proto,
|
|
src: src.Addr().As16(),
|
|
srcPort: src.Port(),
|
|
dst: dst.Addr().As16(),
|
|
dstPort: dst.Port(),
|
|
}
|
|
}
|
|
|
|
// Tuple is a 5-tuple of proto, source and destination IP and port.
|
|
//
|
|
// This struct originally used netip.AddrPort, but that was about twice as slow
|
|
// when used as a map key due to the alignment and extra space for the IPv6 zone
|
|
// pointers (unneeded for all our current 2024-06-17 flowtrack needs).
|
|
//
|
|
// This struct is packed optimally and doesn't contain gaps or pointers.
|
|
type Tuple struct {
|
|
src [16]byte
|
|
dst [16]byte
|
|
srcPort uint16
|
|
dstPort uint16
|
|
proto ipproto.Proto
|
|
}
|
|
|
|
func (t Tuple) SrcAddr() netip.Addr {
|
|
return netip.AddrFrom16(t.src).Unmap()
|
|
}
|
|
|
|
func (t Tuple) DstAddr() netip.Addr {
|
|
return netip.AddrFrom16(t.dst).Unmap()
|
|
}
|
|
|
|
func (t Tuple) SrcPort() uint16 { return t.srcPort }
|
|
func (t Tuple) DstPort() uint16 { return t.dstPort }
|
|
|
|
func (t Tuple) String() string {
|
|
return fmt.Sprintf("(%v %v => %v)", t.proto,
|
|
netip.AddrPortFrom(t.SrcAddr(), t.srcPort),
|
|
netip.AddrPortFrom(t.DstAddr(), t.dstPort))
|
|
}
|
|
|
|
func (t Tuple) MarshalJSON() ([]byte, error) {
|
|
return json.Marshal(tupleOld{
|
|
Proto: t.proto,
|
|
Src: netip.AddrPortFrom(t.SrcAddr(), t.srcPort),
|
|
Dst: netip.AddrPortFrom(t.DstAddr(), t.dstPort),
|
|
})
|
|
}
|
|
|
|
func (t *Tuple) UnmarshalJSON(b []byte) error {
|
|
var ot tupleOld
|
|
if err := json.Unmarshal(b, &ot); err != nil {
|
|
return err
|
|
}
|
|
*t = MakeTuple(ot.Proto, ot.Src, ot.Dst)
|
|
return nil
|
|
}
|
|
|
|
// tupleOld is the old JSON representation of Tuple, before
|
|
// we split and rearranged the fields for efficiency. This type
|
|
// is the JSON adapter type to make sure we still generate
|
|
// the same JSON as before.
|
|
type tupleOld struct {
|
|
Proto ipproto.Proto `json:"proto"`
|
|
Src netip.AddrPort `json:"src"`
|
|
Dst netip.AddrPort `json:"dst"`
|
|
}
|
|
|
|
// Cache is an LRU cache keyed by Tuple.
|
|
//
|
|
// The zero value is valid to use.
|
|
//
|
|
// It is not safe for concurrent access.
|
|
type Cache[Value any] struct {
|
|
// MaxEntries is the maximum number of cache entries before
|
|
// an item is evicted. Zero means no limit.
|
|
MaxEntries int
|
|
|
|
ll *list.List
|
|
m map[Tuple]*list.Element // of *entry
|
|
}
|
|
|
|
// entry is the container/list element type.
|
|
type entry[Value any] struct {
|
|
key Tuple
|
|
value Value
|
|
}
|
|
|
|
// Add adds a value to the cache, set or updating its associated
|
|
// value.
|
|
//
|
|
// If MaxEntries is non-zero and the length of the cache is greater
|
|
// after any addition, the least recently used value is evicted.
|
|
func (c *Cache[Value]) Add(key Tuple, value Value) {
|
|
if c.m == nil {
|
|
c.m = make(map[Tuple]*list.Element)
|
|
c.ll = list.New()
|
|
}
|
|
if ee, ok := c.m[key]; ok {
|
|
c.ll.MoveToFront(ee)
|
|
ee.Value.(*entry[Value]).value = value
|
|
return
|
|
}
|
|
ele := c.ll.PushFront(&entry[Value]{key, value})
|
|
c.m[key] = ele
|
|
if c.MaxEntries != 0 && c.Len() > c.MaxEntries {
|
|
c.RemoveOldest()
|
|
}
|
|
}
|
|
|
|
// Get looks up a key's value from the cache, also reporting
|
|
// whether it was present.
|
|
func (c *Cache[Value]) Get(key Tuple) (value *Value, ok bool) {
|
|
if ele, hit := c.m[key]; hit {
|
|
c.ll.MoveToFront(ele)
|
|
return &ele.Value.(*entry[Value]).value, true
|
|
}
|
|
return nil, false
|
|
}
|
|
|
|
// Remove removes the provided key from the cache if it was present.
|
|
func (c *Cache[Value]) Remove(key Tuple) {
|
|
if ele, hit := c.m[key]; hit {
|
|
c.removeElement(ele)
|
|
}
|
|
}
|
|
|
|
// RemoveOldest removes the oldest item from the cache, if any.
|
|
func (c *Cache[Value]) RemoveOldest() {
|
|
if c.ll != nil {
|
|
if ele := c.ll.Back(); ele != nil {
|
|
c.removeElement(ele)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *Cache[Value]) removeElement(e *list.Element) {
|
|
c.ll.Remove(e)
|
|
delete(c.m, e.Value.(*entry[Value]).key)
|
|
}
|
|
|
|
// Len returns the number of items in the cache.
|
|
func (c *Cache[Value]) Len() int { return len(c.m) }
|