mirror of
https://gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/snowflake.git
synced 2025-10-13 20:11:19 -04:00
Implement server as a v2.1 PT Go API
This commit is contained in:
parent
e87b9175dd
commit
11f0846264
7 changed files with 552 additions and 541 deletions
211
server/lib/http.go
Normal file
211
server/lib/http.go
Normal file
|
@ -0,0 +1,211 @@
|
|||
package lib
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"git.torproject.org/pluggable-transports/snowflake.git/common/encapsulation"
|
||||
"git.torproject.org/pluggable-transports/snowflake.git/common/turbotunnel"
|
||||
"git.torproject.org/pluggable-transports/snowflake.git/common/websocketconn"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
const requestTimeout = 10 * time.Second
|
||||
|
||||
// How long to remember outgoing packets for a client, when we don't currently
|
||||
// have an active WebSocket connection corresponding to that client. Because a
|
||||
// client session may span multiple WebSocket connections, we keep packets we
|
||||
// aren't able to send immediately in memory, for a little while but not
|
||||
// indefinitely.
|
||||
const clientMapTimeout = 1 * time.Minute
|
||||
|
||||
// How big to make the map of ClientIDs to IP addresses. The map is used in
|
||||
// turbotunnelMode to store a reasonable IP address for a client session that
|
||||
// may outlive any single WebSocket connection.
|
||||
const clientIDAddrMapCapacity = 1024
|
||||
|
||||
// How long to wait for ListenAndServe or ListenAndServeTLS to return an error
|
||||
// before deciding that it's not going to return.
|
||||
const listenAndServeErrorTimeout = 100 * time.Millisecond
|
||||
|
||||
var upgrader = websocket.Upgrader{
|
||||
CheckOrigin: func(r *http.Request) bool { return true },
|
||||
}
|
||||
|
||||
// clientIDAddrMap stores short-term mappings from ClientIDs to IP addresses.
|
||||
// When we call pt.DialOr, tor wants us to provide a USERADDR string that
|
||||
// represents the remote IP address of the client (for metrics purposes, etc.).
|
||||
// This data structure bridges the gap between ServeHTTP, which knows about IP
|
||||
// addresses, and handleStream, which is what calls pt.DialOr. The common piece
|
||||
// of information linking both ends of the chain is the ClientID, which is
|
||||
// attached to the WebSocket connection and every session.
|
||||
var clientIDAddrMap = newClientIDMap(clientIDAddrMapCapacity)
|
||||
|
||||
// overrideReadConn is a net.Conn with an overridden Read method. Compare to
|
||||
// recordingConn at
|
||||
// https://dave.cheney.net/2015/05/22/struct-composition-with-go.
|
||||
type overrideReadConn struct {
|
||||
net.Conn
|
||||
io.Reader
|
||||
}
|
||||
|
||||
func (conn *overrideReadConn) Read(p []byte) (int, error) {
|
||||
return conn.Reader.Read(p)
|
||||
}
|
||||
|
||||
type HTTPHandler struct {
|
||||
// pconn is the adapter layer between stream-oriented WebSocket
|
||||
// connections and the packet-oriented KCP layer.
|
||||
pconn *turbotunnel.QueuePacketConn
|
||||
ln *SnowflakeListener
|
||||
}
|
||||
|
||||
func (handler *HTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
ws, err := upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
log.Println(err)
|
||||
return
|
||||
}
|
||||
|
||||
conn := websocketconn.New(ws)
|
||||
defer conn.Close()
|
||||
|
||||
// Pass the address of client as the remote address of incoming connection
|
||||
clientIPParam := r.URL.Query().Get("client_ip")
|
||||
addr := clientAddr(clientIPParam)
|
||||
|
||||
var token [len(turbotunnel.Token)]byte
|
||||
_, err = io.ReadFull(conn, token[:])
|
||||
if err != nil {
|
||||
// Don't bother logging EOF: that happens with an unused
|
||||
// connection, which clients make frequently as they maintain a
|
||||
// pool of proxies.
|
||||
if err != io.EOF {
|
||||
log.Printf("reading token: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
switch {
|
||||
case bytes.Equal(token[:], turbotunnel.Token[:]):
|
||||
err = turbotunnelMode(conn, addr, handler.pconn)
|
||||
default:
|
||||
// We didn't find a matching token, which means that we are
|
||||
// dealing with a client that doesn't know about such things.
|
||||
// "Unread" the token by constructing a new Reader and pass it
|
||||
// to the old one-session-per-WebSocket mode.
|
||||
conn2 := &overrideReadConn{Conn: conn, Reader: io.MultiReader(bytes.NewReader(token[:]), conn)}
|
||||
err = oneshotMode(conn2, addr, handler.ln)
|
||||
}
|
||||
if err != nil {
|
||||
log.Println(err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// oneshotMode handles clients that did not send turbotunnel.Token at the start
|
||||
// of their stream. These clients use the WebSocket as a raw pipe, and expect
|
||||
// their session to begin and end when this single WebSocket does.
|
||||
func oneshotMode(conn net.Conn, addr net.Addr, ln *SnowflakeListener) error {
|
||||
return ln.QueueConn(&SnowflakeClientConn{Conn: conn, address: addr})
|
||||
}
|
||||
|
||||
// turbotunnelMode handles clients that sent turbotunnel.Token at the start of
|
||||
// their stream. These clients expect to send and receive encapsulated packets,
|
||||
// with a long-lived session identified by ClientID.
|
||||
func turbotunnelMode(conn net.Conn, addr net.Addr, pconn *turbotunnel.QueuePacketConn) error {
|
||||
// Read the ClientID prefix. Every packet encapsulated in this WebSocket
|
||||
// connection pertains to the same ClientID.
|
||||
var clientID turbotunnel.ClientID
|
||||
_, err := io.ReadFull(conn, clientID[:])
|
||||
if err != nil {
|
||||
return fmt.Errorf("reading ClientID: %v", err)
|
||||
}
|
||||
|
||||
// Store a a short-term mapping from the ClientID to the client IP
|
||||
// address attached to this WebSocket connection. tor will want us to
|
||||
// provide a client IP address when we call pt.DialOr. But a KCP session
|
||||
// does not necessarily correspond to any single IP address--it's
|
||||
// composed of packets that are carried in possibly multiple WebSocket
|
||||
// streams. We apply the heuristic that the IP address of the most
|
||||
// recent WebSocket connection that has had to do with a session, at the
|
||||
// time the session is established, is the IP address that should be
|
||||
// credited for the entire KCP session.
|
||||
clientIDAddrMap.Set(clientID, addr.String())
|
||||
|
||||
errCh := make(chan error)
|
||||
|
||||
// The remainder of the WebSocket stream consists of encapsulated
|
||||
// packets. We read them one by one and feed them into the
|
||||
// QueuePacketConn on which kcp.ServeConn was set up, which eventually
|
||||
// leads to KCP-level sessions in the acceptSessions function.
|
||||
go func() {
|
||||
for {
|
||||
p, err := encapsulation.ReadData(conn)
|
||||
if err != nil {
|
||||
errCh <- err
|
||||
break
|
||||
}
|
||||
pconn.QueueIncoming(p, clientID)
|
||||
}
|
||||
}()
|
||||
|
||||
// At the same time, grab packets addressed to this ClientID and
|
||||
// encapsulate them into the downstream.
|
||||
go func() {
|
||||
// Buffer encapsulation.WriteData operations to keep length
|
||||
// prefixes in the same send as the data that follows.
|
||||
bw := bufio.NewWriter(conn)
|
||||
for p := range pconn.OutgoingQueue(clientID) {
|
||||
_, err := encapsulation.WriteData(bw, p)
|
||||
if err == nil {
|
||||
err = bw.Flush()
|
||||
}
|
||||
if err != nil {
|
||||
errCh <- err
|
||||
break
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// Wait until one of the above loops terminates. The closing of the
|
||||
// WebSocket connection will terminate the other one.
|
||||
<-errCh
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
type ClientMapAddr string
|
||||
|
||||
func (addr ClientMapAddr) Network() string {
|
||||
return "snowflake"
|
||||
}
|
||||
|
||||
func (addr ClientMapAddr) String() string {
|
||||
return string(addr)
|
||||
}
|
||||
|
||||
// Return a client address
|
||||
func clientAddr(clientIPParam string) net.Addr {
|
||||
if clientIPParam == "" {
|
||||
return ClientMapAddr("")
|
||||
}
|
||||
// Check if client addr is a valid IP
|
||||
clientIP := net.ParseIP(clientIPParam)
|
||||
if clientIP == nil {
|
||||
return ClientMapAddr("")
|
||||
}
|
||||
// Check if client addr is 0.0.0.0 or [::]. Some proxies erroneously
|
||||
// report an address of 0.0.0.0: https://bugs.torproject.org/33157.
|
||||
if clientIP.IsUnspecified() {
|
||||
return ClientMapAddr("")
|
||||
}
|
||||
// Add a stub port number. USERADDR requires a port number.
|
||||
return ClientMapAddr((&net.TCPAddr{IP: clientIP, Port: 1, Zone: ""}).String())
|
||||
}
|
55
server/lib/server_test.go
Normal file
55
server/lib/server_test.go
Normal file
|
@ -0,0 +1,55 @@
|
|||
package lib
|
||||
|
||||
import (
|
||||
"net"
|
||||
"strconv"
|
||||
"testing"
|
||||
|
||||
. "github.com/smartystreets/goconvey/convey"
|
||||
)
|
||||
|
||||
func TestClientAddr(t *testing.T) {
|
||||
Convey("Testing clientAddr", t, func() {
|
||||
// good tests
|
||||
for _, test := range []struct {
|
||||
input string
|
||||
expected net.IP
|
||||
}{
|
||||
{"1.2.3.4", net.ParseIP("1.2.3.4")},
|
||||
{"1:2::3:4", net.ParseIP("1:2::3:4")},
|
||||
} {
|
||||
useraddr := clientAddr(test.input).String()
|
||||
host, port, err := net.SplitHostPort(useraddr)
|
||||
if err != nil {
|
||||
t.Errorf("clientAddr(%q) → SplitHostPort error %v", test.input, err)
|
||||
continue
|
||||
}
|
||||
if !test.expected.Equal(net.ParseIP(host)) {
|
||||
t.Errorf("clientAddr(%q) → host %q, not %v", test.input, host, test.expected)
|
||||
}
|
||||
portNo, err := strconv.Atoi(port)
|
||||
if err != nil {
|
||||
t.Errorf("clientAddr(%q) → port %q", test.input, port)
|
||||
continue
|
||||
}
|
||||
if portNo == 0 {
|
||||
t.Errorf("clientAddr(%q) → port %d", test.input, portNo)
|
||||
}
|
||||
}
|
||||
|
||||
// bad tests
|
||||
for _, input := range []string{
|
||||
"",
|
||||
"abc",
|
||||
"1.2.3.4.5",
|
||||
"[12::34]",
|
||||
"0.0.0.0",
|
||||
"[::]",
|
||||
} {
|
||||
useraddr := clientAddr(input).String()
|
||||
if useraddr != "" {
|
||||
t.Errorf("clientAddr(%q) → %q, not %q", input, useraddr, "")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
242
server/lib/snowflake.go
Normal file
242
server/lib/snowflake.go
Normal file
|
@ -0,0 +1,242 @@
|
|||
package lib
|
||||
|
||||
import (
|
||||
"crypto/tls"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"git.torproject.org/pluggable-transports/snowflake.git/common/turbotunnel"
|
||||
"github.com/xtaci/kcp-go/v5"
|
||||
"github.com/xtaci/smux"
|
||||
"golang.org/x/net/http2"
|
||||
)
|
||||
|
||||
// Transport is a structure with methods that conform to the Go PT v2.1 API
|
||||
// https://github.com/Pluggable-Transports/Pluggable-Transports-spec/blob/master/releases/PTSpecV2.1/Pluggable%20Transport%20Specification%20v2.1%20-%20Go%20Transport%20API.pdf
|
||||
type Transport struct {
|
||||
getCertificate func(*tls.ClientHelloInfo) (*tls.Certificate, error)
|
||||
}
|
||||
|
||||
func NewSnowflakeServer(getCertificate func(*tls.ClientHelloInfo) (*tls.Certificate, error)) *Transport {
|
||||
|
||||
return &Transport{getCertificate: getCertificate}
|
||||
}
|
||||
|
||||
func (t *Transport) Listen(addr net.Addr) (*SnowflakeListener, error) {
|
||||
listener := &SnowflakeListener{addr: addr, queue: make(chan net.Conn, 65534)}
|
||||
|
||||
handler := HTTPHandler{
|
||||
// pconn is shared among all connections to this server. It
|
||||
// overlays packet-based client sessions on top of ephemeral
|
||||
// WebSocket connections.
|
||||
pconn: turbotunnel.NewQueuePacketConn(addr, clientMapTimeout),
|
||||
}
|
||||
server := &http.Server{
|
||||
Addr: addr.String(),
|
||||
Handler: &handler,
|
||||
ReadTimeout: requestTimeout,
|
||||
}
|
||||
// We need to override server.TLSConfig.GetCertificate--but first
|
||||
// server.TLSConfig needs to be non-nil. If we just create our own new
|
||||
// &tls.Config, it will lack the default settings that the net/http
|
||||
// package sets up for things like HTTP/2. Therefore we first call
|
||||
// http2.ConfigureServer for its side effect of initializing
|
||||
// server.TLSConfig properly. An alternative would be to make a dummy
|
||||
// net.Listener, call Serve on it, and let it return.
|
||||
// https://github.com/golang/go/issues/16588#issuecomment-237386446
|
||||
err := http2.ConfigureServer(server, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
server.TLSConfig.GetCertificate = t.getCertificate
|
||||
|
||||
// Another unfortunate effect of the inseparable net/http ListenAndServe
|
||||
// is that we can't check for Listen errors like "permission denied" and
|
||||
// "address already in use" without potentially entering the infinite
|
||||
// loop of Serve. The hack we apply here is to wait a short time,
|
||||
// listenAndServeErrorTimeout, to see if an error is returned (because
|
||||
// it's better if the error message goes to the tor log through
|
||||
// SMETHOD-ERROR than if it only goes to the snowflake log).
|
||||
errChan := make(chan error)
|
||||
go func() {
|
||||
if t.getCertificate == nil {
|
||||
// TLS is disabled
|
||||
log.Printf("listening with plain HTTP on %s", addr)
|
||||
err := server.ListenAndServe()
|
||||
if err != nil {
|
||||
log.Printf("error in ListenAndServe: %s", err)
|
||||
}
|
||||
errChan <- err
|
||||
} else {
|
||||
log.Printf("listening with HTTPS on %s", addr)
|
||||
err := server.ListenAndServeTLS("", "")
|
||||
if err != nil {
|
||||
log.Printf("error in ListenAndServeTLS: %s", err)
|
||||
}
|
||||
errChan <- err
|
||||
}
|
||||
}()
|
||||
|
||||
select {
|
||||
case err = <-errChan:
|
||||
break
|
||||
case <-time.After(listenAndServeErrorTimeout):
|
||||
break
|
||||
}
|
||||
|
||||
listener.server = server
|
||||
|
||||
// Start a KCP engine, set up to read and write its packets over the
|
||||
// WebSocket connections that arrive at the web server.
|
||||
// handler.ServeHTTP is responsible for encapsulation/decapsulation of
|
||||
// packets on behalf of KCP. KCP takes those packets and turns them into
|
||||
// sessions which appear in the acceptSessions function.
|
||||
ln, err := kcp.ServeConn(nil, 0, 0, handler.pconn)
|
||||
if err != nil {
|
||||
server.Close()
|
||||
return nil, err
|
||||
}
|
||||
go func() {
|
||||
defer ln.Close()
|
||||
err := listener.acceptSessions(ln)
|
||||
if err != nil {
|
||||
log.Printf("acceptSessions: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
listener.ln = ln
|
||||
|
||||
return listener, nil
|
||||
|
||||
}
|
||||
|
||||
type SnowflakeListener struct {
|
||||
addr net.Addr
|
||||
queue chan net.Conn
|
||||
server *http.Server
|
||||
ln *kcp.Listener
|
||||
closed chan struct{}
|
||||
closeOnce sync.Once
|
||||
}
|
||||
|
||||
// Allows the caller to accept incoming Snowflake connections
|
||||
// We accept connections from a queue to accommodate both incoming
|
||||
// smux Streams and legacy non-turbotunnel connections
|
||||
func (l *SnowflakeListener) Accept() (net.Conn, error) {
|
||||
select {
|
||||
case <-l.closed:
|
||||
//channel has been closed, no longer accepting connections
|
||||
return nil, io.ErrClosedPipe
|
||||
case conn := <-l.queue:
|
||||
return conn, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (l *SnowflakeListener) Addr() net.Addr {
|
||||
return l.addr
|
||||
}
|
||||
|
||||
func (l *SnowflakeListener) Close() error {
|
||||
// Close our HTTP server and our KCP listener
|
||||
l.closeOnce.Do(func() {
|
||||
close(l.closed)
|
||||
l.server.Close()
|
||||
l.ln.Close()
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// acceptStreams layers an smux.Session on the KCP connection and awaits streams
|
||||
// on it. Passes each stream to our SnowflakeListener accept queue.
|
||||
func (l *SnowflakeListener) acceptStreams(conn *kcp.UDPSession) error {
|
||||
// Look up the IP address associated with this KCP session, via the
|
||||
// ClientID that is returned by the session's RemoteAddr method.
|
||||
addr, ok := clientIDAddrMap.Get(conn.RemoteAddr().(turbotunnel.ClientID))
|
||||
if !ok {
|
||||
// This means that the map is tending to run over capacity, not
|
||||
// just that there was not client_ip on the incoming connection.
|
||||
// We store "" in the map in the absence of client_ip. This log
|
||||
// message means you should increase clientIDAddrMapCapacity.
|
||||
log.Printf("no address in clientID-to-IP map (capacity %d)", clientIDAddrMapCapacity)
|
||||
}
|
||||
|
||||
smuxConfig := smux.DefaultConfig()
|
||||
smuxConfig.Version = 2
|
||||
smuxConfig.KeepAliveTimeout = 10 * time.Minute
|
||||
sess, err := smux.Server(conn, smuxConfig)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for {
|
||||
stream, err := sess.AcceptStream()
|
||||
if err != nil {
|
||||
if err, ok := err.(net.Error); ok && err.Temporary() {
|
||||
continue
|
||||
}
|
||||
return err
|
||||
}
|
||||
l.QueueConn(&SnowflakeClientConn{Conn: stream, address: clientAddr(addr)})
|
||||
}
|
||||
}
|
||||
|
||||
// acceptSessions listens for incoming KCP connections and passes them to
|
||||
// acceptStreams. It is handler.ServeHTTP that provides the network interface
|
||||
// that drives this function.
|
||||
func (l *SnowflakeListener) acceptSessions(ln *kcp.Listener) error {
|
||||
for {
|
||||
conn, err := ln.AcceptKCP()
|
||||
if err != nil {
|
||||
if err, ok := err.(net.Error); ok && err.Temporary() {
|
||||
continue
|
||||
}
|
||||
return err
|
||||
}
|
||||
// Permit coalescing the payloads of consecutive sends.
|
||||
conn.SetStreamMode(true)
|
||||
// Set the maximum send and receive window sizes to a high number
|
||||
// Removes KCP bottlenecks: https://gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/snowflake/-/issues/40026
|
||||
conn.SetWindowSize(65535, 65535)
|
||||
// Disable the dynamic congestion window (limit only by the
|
||||
// maximum of local and remote static windows).
|
||||
conn.SetNoDelay(
|
||||
0, // default nodelay
|
||||
0, // default interval
|
||||
0, // default resend
|
||||
1, // nc=1 => congestion window off
|
||||
)
|
||||
go func() {
|
||||
defer conn.Close()
|
||||
err := l.acceptStreams(conn)
|
||||
if err != nil && err != io.ErrClosedPipe {
|
||||
log.Printf("acceptStreams: %v", err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
}
|
||||
|
||||
func (l *SnowflakeListener) QueueConn(conn net.Conn) error {
|
||||
select {
|
||||
case <-l.closed:
|
||||
return fmt.Errorf("accepted connection on closed listener")
|
||||
case l.queue <- conn:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// A wrapper for the underlying oneshot or turbotunnel conn
|
||||
// because we need to reference our mapping to determine the client
|
||||
// address
|
||||
type SnowflakeClientConn struct {
|
||||
net.Conn
|
||||
address net.Addr
|
||||
}
|
||||
|
||||
func (conn *SnowflakeClientConn) RemoteAddr() net.Addr {
|
||||
return conn.address
|
||||
}
|
85
server/lib/turbotunnel.go
Normal file
85
server/lib/turbotunnel.go
Normal file
|
@ -0,0 +1,85 @@
|
|||
package lib
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"git.torproject.org/pluggable-transports/snowflake.git/common/turbotunnel"
|
||||
)
|
||||
|
||||
// clientIDMap is a fixed-capacity mapping from ClientIDs to address strings.
|
||||
// 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 string
|
||||
}
|
||||
// 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 string
|
||||
}, 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 string) {
|
||||
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 string will be "".
|
||||
func (m *clientIDMap) Get(clientID turbotunnel.ClientID) (string, bool) {
|
||||
m.lock.Lock()
|
||||
defer m.lock.Unlock()
|
||||
if i, ok := m.current[clientID]; ok {
|
||||
return m.entries[i].addr, true
|
||||
} else {
|
||||
return "", false
|
||||
}
|
||||
}
|
119
server/lib/turbotunnel_test.go
Normal file
119
server/lib/turbotunnel_test.go
Normal file
|
@ -0,0 +1,119 @@
|
|||
package lib
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"testing"
|
||||
|
||||
"git.torproject.org/pluggable-transports/snowflake.git/common/turbotunnel"
|
||||
)
|
||||
|
||||
func TestClientIDMap(t *testing.T) {
|
||||
// Convert a uint64 into a ClientID.
|
||||
id := func(n uint64) turbotunnel.ClientID {
|
||||
var clientID turbotunnel.ClientID
|
||||
binary.PutUvarint(clientID[:], n)
|
||||
return clientID
|
||||
}
|
||||
|
||||
// Does m.Get(key) and checks that the output matches what is expected.
|
||||
expectGet := func(m *clientIDMap, clientID turbotunnel.ClientID, expectedAddr string, expectedOK bool) {
|
||||
t.Helper()
|
||||
addr, ok := m.Get(clientID)
|
||||
if addr != expectedAddr || ok != expectedOK {
|
||||
t.Errorf("expected (%+q, %v), got (%+q, %v)", expectedAddr, expectedOK, addr, ok)
|
||||
}
|
||||
}
|
||||
|
||||
// Checks that the len of m.current is as expected.
|
||||
expectSize := func(m *clientIDMap, expectedLen int) {
|
||||
t.Helper()
|
||||
if len(m.current) != expectedLen {
|
||||
t.Errorf("expected map len %d, got %d %+v", expectedLen, len(m.current), m.current)
|
||||
}
|
||||
}
|
||||
|
||||
// Zero-capacity map can't remember anything.
|
||||
{
|
||||
m := newClientIDMap(0)
|
||||
expectSize(m, 0)
|
||||
expectGet(m, id(0), "", false)
|
||||
expectGet(m, id(1234), "", false)
|
||||
|
||||
m.Set(id(0), "A")
|
||||
expectSize(m, 0)
|
||||
expectGet(m, id(0), "", false)
|
||||
expectGet(m, id(1234), "", false)
|
||||
|
||||
m.Set(id(1234), "A")
|
||||
expectSize(m, 0)
|
||||
expectGet(m, id(0), "", false)
|
||||
expectGet(m, id(1234), "", false)
|
||||
}
|
||||
|
||||
{
|
||||
m := newClientIDMap(1)
|
||||
expectSize(m, 0)
|
||||
expectGet(m, id(0), "", false)
|
||||
expectGet(m, id(1), "", false)
|
||||
|
||||
m.Set(id(0), "A")
|
||||
expectSize(m, 1)
|
||||
expectGet(m, id(0), "A", true)
|
||||
expectGet(m, id(1), "", false)
|
||||
|
||||
m.Set(id(1), "B") // forgets the (0, "A") entry
|
||||
expectSize(m, 1)
|
||||
expectGet(m, id(0), "", false)
|
||||
expectGet(m, id(1), "B", true)
|
||||
|
||||
m.Set(id(1), "C") // forgets the (1, "B") entry
|
||||
expectSize(m, 1)
|
||||
expectGet(m, id(0), "", false)
|
||||
expectGet(m, id(1), "C", true)
|
||||
}
|
||||
|
||||
{
|
||||
m := newClientIDMap(5)
|
||||
m.Set(id(0), "A")
|
||||
m.Set(id(1), "B")
|
||||
m.Set(id(2), "C")
|
||||
m.Set(id(0), "D") // shadows the (0, "D") entry
|
||||
m.Set(id(3), "E")
|
||||
expectSize(m, 4)
|
||||
expectGet(m, id(0), "D", true)
|
||||
expectGet(m, id(1), "B", true)
|
||||
expectGet(m, id(2), "C", true)
|
||||
expectGet(m, id(3), "E", true)
|
||||
expectGet(m, id(4), "", false)
|
||||
|
||||
m.Set(id(4), "F") // forgets the (0, "A") entry but should preserve (0, "D")
|
||||
expectSize(m, 5)
|
||||
expectGet(m, id(0), "D", true)
|
||||
expectGet(m, id(1), "B", true)
|
||||
expectGet(m, id(2), "C", true)
|
||||
expectGet(m, id(3), "E", true)
|
||||
expectGet(m, id(4), "F", true)
|
||||
|
||||
m.Set(id(5), "G") // forgets the (1, "B") entry
|
||||
m.Set(id(0), "H") // forgets the (2, "C") entry and shadows (0, "D")
|
||||
expectSize(m, 4)
|
||||
expectGet(m, id(0), "H", true)
|
||||
expectGet(m, id(1), "", false)
|
||||
expectGet(m, id(2), "", false)
|
||||
expectGet(m, id(3), "E", true)
|
||||
expectGet(m, id(4), "F", true)
|
||||
expectGet(m, id(5), "G", true)
|
||||
|
||||
m.Set(id(0), "I") // forgets the (0, "D") entry and shadows (0, "H")
|
||||
m.Set(id(0), "J") // forgets the (3, "E") entry and shadows (0, "I")
|
||||
m.Set(id(0), "K") // forgets the (4, "F") entry and shadows (0, "J")
|
||||
m.Set(id(0), "L") // forgets the (5, "G") entry and shadows (0, "K")
|
||||
expectSize(m, 1)
|
||||
expectGet(m, id(0), "L", true)
|
||||
expectGet(m, id(1), "", false)
|
||||
expectGet(m, id(2), "", false)
|
||||
expectGet(m, id(3), "", false)
|
||||
expectGet(m, id(4), "", false)
|
||||
expectGet(m, id(5), "", false)
|
||||
}
|
||||
}
|
Loading…
Add table
Add a link
Reference in a new issue