mirror of
https://gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/snowflake.git
synced 2025-10-13 11:11:30 -04:00
86 lines
3 KiB
Go
86 lines
3 KiB
Go
package snowflake_server
|
|
|
|
import (
|
|
"net"
|
|
"sync"
|
|
|
|
"gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/snowflake/v2/common/turbotunnel"
|
|
)
|
|
|
|
// clientIDMap is a fixed-capacity mapping from ClientIDs to a net.Addr.
|
|
// Adding a new entry using the Set method causes the oldest existing entry to
|
|
// be forgotten.
|
|
//
|
|
// This data type is meant to be used to remember the IP address associated with
|
|
// a ClientID, during the short period of time between when a WebSocket
|
|
// connection with that ClientID began, and when a KCP session is established.
|
|
//
|
|
// The design requirements of this type are that it needs to remember a mapping
|
|
// for only a short time, and old entries should expire so as not to consume
|
|
// unbounded memory. It is not a critical error if an entry is forgotten before
|
|
// it is needed; better to forget entries than to use too much memory.
|
|
type clientIDMap struct {
|
|
lock sync.Mutex
|
|
// entries is a circular buffer of (ClientID, addr) pairs.
|
|
entries []struct {
|
|
clientID turbotunnel.ClientID
|
|
addr net.Addr
|
|
}
|
|
// oldest is the index of the oldest member of the entries buffer, the
|
|
// one that will be overwritten at the next call to Set.
|
|
oldest int
|
|
// current points to the index of the most recent entry corresponding to
|
|
// each ClientID.
|
|
current map[turbotunnel.ClientID]int
|
|
}
|
|
|
|
// newClientIDMap makes a new clientIDMap with the given capacity.
|
|
func newClientIDMap(capacity int) *clientIDMap {
|
|
return &clientIDMap{
|
|
entries: make([]struct {
|
|
clientID turbotunnel.ClientID
|
|
addr net.Addr
|
|
}, capacity),
|
|
oldest: 0,
|
|
current: make(map[turbotunnel.ClientID]int),
|
|
}
|
|
}
|
|
|
|
// Set adds a mapping from clientID to addr, replacing any previous mapping for
|
|
// clientID. It may also cause the clientIDMap to forget at most one other
|
|
// mapping, the oldest one.
|
|
func (m *clientIDMap) Set(clientID turbotunnel.ClientID, addr net.Addr) {
|
|
m.lock.Lock()
|
|
defer m.lock.Unlock()
|
|
if len(m.entries) == 0 {
|
|
// The invariant m.oldest < len(m.entries) does not hold in this
|
|
// special case.
|
|
return
|
|
}
|
|
// m.oldest is the index of the entry we're about to overwrite. If it's
|
|
// the current entry for any ClientID, we need to delete that clientID
|
|
// from the current map (that ClientID is now forgotten).
|
|
if i, ok := m.current[m.entries[m.oldest].clientID]; ok && i == m.oldest {
|
|
delete(m.current, m.entries[m.oldest].clientID)
|
|
}
|
|
// Overwrite the oldest entry.
|
|
m.entries[m.oldest].clientID = clientID
|
|
m.entries[m.oldest].addr = addr
|
|
// Add the overwritten entry to the quick-lookup map.
|
|
m.current[clientID] = m.oldest
|
|
// What was the oldest entry is now the newest.
|
|
m.oldest = (m.oldest + 1) % len(m.entries)
|
|
}
|
|
|
|
// Get returns a previously stored mapping. The second return value indicates
|
|
// whether clientID was actually present in the map. If it is false, then the
|
|
// returned address will be nil.
|
|
func (m *clientIDMap) Get(clientID turbotunnel.ClientID) (net.Addr, bool) {
|
|
m.lock.Lock()
|
|
defer m.lock.Unlock()
|
|
if i, ok := m.current[clientID]; ok {
|
|
return m.entries[i].addr, true
|
|
} else {
|
|
return nil, false
|
|
}
|
|
}
|