Merge pull request #469 from rod-hynes/master

Add meek rate limiter

psiphon/server/meek.go

@@ -31,6 +31,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"runtime"
 	"strconv"
 	"strings"
 	"sync"
@@ -93,16 +94,20 @@ const (
 // HTTP payload traffic for a given session into net.Conn conforming Read()s and Write()s via
 // the meekConn struct.
 type MeekServer struct {
-	support       *SupportServices
-	listener      net.Listener
-	tlsConfig     *utls.Config
-	clientHandler func(clientTunnelProtocol string, clientConn net.Conn)
-	openConns     *common.Conns
-	stopBroadcast <-chan struct{}
-	sessionsLock  sync.RWMutex
-	sessions      map[string]*meekSession
-	checksumTable *crc64.Table
-	bufferPool    *CachedResponseBufferPool
+	support           *SupportServices
+	listener          net.Listener
+	tlsConfig         *utls.Config
+	clientHandler     func(clientTunnelProtocol string, clientConn net.Conn)
+	openConns         *common.Conns
+	stopBroadcast     <-chan struct{}
+	sessionsLock      sync.RWMutex
+	sessions          map[string]*meekSession
+	checksumTable     *crc64.Table
+	bufferPool        *CachedResponseBufferPool
+	rateLimitLock     sync.Mutex
+	rateLimitHistory  map[string][]monotime.Time
+	rateLimitCount    int
+	rateLimitSignalGC chan struct{}
 }
 
 // NewMeekServer initializes a new meek server.
@@ -128,14 +133,16 @@ func NewMeekServer(
 	bufferPool := NewCachedResponseBufferPool(bufferLength, bufferCount)
 
 	meekServer := &MeekServer{
-		support:       support,
-		listener:      listener,
-		clientHandler: clientHandler,
-		openConns:     new(common.Conns),
-		stopBroadcast: stopBroadcast,
-		sessions:      make(map[string]*meekSession),
-		checksumTable: checksumTable,
-		bufferPool:    bufferPool,
+		support:           support,
+		listener:          listener,
+		clientHandler:     clientHandler,
+		openConns:         new(common.Conns),
+		stopBroadcast:     stopBroadcast,
+		sessions:          make(map[string]*meekSession),
+		checksumTable:     checksumTable,
+		bufferPool:        bufferPool,
+		rateLimitHistory:  make(map[string][]monotime.Time),
+		rateLimitSignalGC: make(chan struct{}, 1),
 	}
 
 	if useTLS {
@@ -158,12 +165,11 @@ func NewMeekServer(
 // signal specified in NewMeekServer.
 func (server *MeekServer) Run() error {
 
-	// Expire sessions
+	waitGroup := new(sync.WaitGroup)
 
-	reaperWaitGroup := new(sync.WaitGroup)
-	reaperWaitGroup.Add(1)
+	waitGroup.Add(1)
 	go func() {
-		defer reaperWaitGroup.Done()
+		defer waitGroup.Done()
 		ticker := time.NewTicker(MEEK_MAX_SESSION_STALENESS / 2)
 		defer ticker.Stop()
 		for {
@@ -176,9 +182,14 @@ func (server *MeekServer) Run() error {
 		}
 	}()
 
-	// Serve HTTP or HTTPS
+	waitGroup.Add(1)
+	go func() {
+		defer waitGroup.Done()
+		server.rateLimitWorker()
+	}()
 
-	// Notes:
+	// Serve HTTP or HTTPS
+	//
 	// - WriteTimeout may include time awaiting request, as per:
 	//   https://blog.cloudflare.com/the-complete-guide-to-golang-net-http-timeouts
 	// - Legacy meek-server wrapped each client HTTP connection with an explicit idle
@@ -222,7 +233,7 @@ func (server *MeekServer) Run() error {
 	server.listener.Close()
 	server.openConns.CloseAll()
 
-	reaperWaitGroup.Wait()
+	waitGroup.Wait()
 
 	return err
 }
@@ -504,37 +515,19 @@ func checkRangeHeader(request *http.Request) (int, bool) {
 func (server *MeekServer) getSessionOrEndpoint(
 	request *http.Request, meekCookie *http.Cookie) (string, *meekSession, string, string, error) {
 
-	// Check for an existing session
+	// Check for an existing session.
 
 	server.sessionsLock.RLock()
 	existingSessionID := meekCookie.Value
 	session, ok := server.sessions[existingSessionID]
 	server.sessionsLock.RUnlock()
 	if ok {
+		// TODO: can multiple http client connections using same session cookie
+		// cause race conditions on session struct?
 		session.touch()
 		return existingSessionID, session, "", "", nil
 	}
 
-	// TODO: can multiple http client connections using same session cookie
-	// cause race conditions on session struct?
-
-	// The session is new (or expired). Treat the cookie value as a new meek
-	// cookie, extract the payload, and create a new session.
-
-	payloadJSON, err := getMeekCookiePayload(server.support, meekCookie.Value)
-	if err != nil {
-		return "", nil, "", "", common.ContextError(err)
-	}
-
-	// Note: this meek server ignores legacy values PsiphonClientSessionId
-	// and PsiphonServerAddress.
-	var clientSessionData protocol.MeekCookieData
-
-	err = json.Unmarshal(payloadJSON, &clientSessionData)
-	if err != nil {
-		return "", nil, "", "", common.ContextError(err)
-	}
-
 	// Determine the client remote address, which is used for geolocation
 	// and stats. When an intermediate proxy or CDN is in use, we may be
 	// able to determine the original client address by inspecting HTTP
@@ -560,6 +553,27 @@ func (server *MeekServer) getSessionOrEndpoint(
 		}
 	}
 
+	if server.rateLimit(clientIP) {
+		return "", nil, "", "", common.ContextError(errors.New("rate limit exceeded"))
+	}
+
+	// The session is new (or expired). Treat the cookie value as a new meek
+	// cookie, extract the payload, and create a new session.
+
+	payloadJSON, err := getMeekCookiePayload(server.support, meekCookie.Value)
+	if err != nil {
+		return "", nil, "", "", common.ContextError(err)
+	}
+
+	// Note: this meek server ignores legacy values PsiphonClientSessionId
+	// and PsiphonServerAddress.
+	var clientSessionData protocol.MeekCookieData
+
+	err = json.Unmarshal(payloadJSON, &clientSessionData)
+	if err != nil {
+		return "", nil, "", "", common.ContextError(err)
+	}
+
 	// Handle endpoints before enforcing the GetEstablishTunnels check.
 	// Currently, endpoints are tactics requests, and we allow these to be
 	// handled by servers which would otherwise reject new tunnels.
@@ -637,6 +651,118 @@ func (server *MeekServer) getSessionOrEndpoint(
 	return sessionID, session, "", "", nil
 }
 
+func (server *MeekServer) rateLimit(clientIP string) bool {
+
+	historySize, thresholdSeconds, regions, GCTriggerCount, _ :=
+		server.support.TrafficRulesSet.GetMeekRateLimiterConfig()
+
+	if historySize == 0 {
+		return false
+	}
+
+	if len(regions) > 0 {
+		// TODO: avoid redundant GeoIP lookups?
+		if !common.Contains(regions, server.support.GeoIPService.Lookup(clientIP).Country) {
+			return false
+		}
+	}
+
+	limit := true
+	triggerGC := false
+
+	now := monotime.Now()
+	threshold := now.Add(-time.Duration(thresholdSeconds) * time.Second)
+
+	server.rateLimitLock.Lock()
+
+	history, ok := server.rateLimitHistory[clientIP]
+	if !ok || len(history) != historySize {
+		history = make([]monotime.Time, historySize)
+		server.rateLimitHistory[clientIP] = history
+	}
+
+	for i := 0; i < len(history); i++ {
+		if history[i] == 0 || history[i].Before(threshold) {
+			limit = false
+		}
+		if i == len(history)-1 {
+			history[i] = now
+		} else {
+			history[i] = history[i+1]
+		}
+	}
+
+	if limit {
+
+		server.rateLimitCount += 1
+
+		if server.rateLimitCount >= GCTriggerCount {
+			triggerGC = true
+			server.rateLimitCount = 0
+		}
+	}
+
+	server.rateLimitLock.Unlock()
+
+	if triggerGC {
+		select {
+		case server.rateLimitSignalGC <- *new(struct{}):
+		default:
+		}
+	}
+
+	return limit
+}
+
+func (server *MeekServer) rateLimitWorker() {
+
+	_, _, _, _, reapFrequencySeconds :=
+		server.support.TrafficRulesSet.GetMeekRateLimiterConfig()
+
+	timer := time.NewTimer(time.Duration(reapFrequencySeconds) * time.Second)
+	defer timer.Stop()
+
+	for {
+		select {
+		case <-timer.C:
+
+			_, thresholdSeconds, _, _, reapFrequencySeconds :=
+				server.support.TrafficRulesSet.GetMeekRateLimiterConfig()
+
+			server.rateLimitLock.Lock()
+
+			threshold := monotime.Now().Add(-time.Duration(thresholdSeconds) * time.Second)
+
+			for key, history := range server.rateLimitHistory {
+				reap := true
+				for i := 0; i < len(history); i++ {
+					if history[i] != 0 && !history[i].Before(threshold) {
+						reap = false
+					}
+				}
+				if reap {
+					delete(server.rateLimitHistory, key)
+				}
+			}
+
+			// Enable rate limit history map to be garbage collected when possible.
+			if len(server.rateLimitHistory) == 0 {
+				server.rateLimitHistory = make(map[string][]monotime.Time)
+			}
+
+			server.rateLimitLock.Unlock()
+
+			timer.Reset(time.Duration(reapFrequencySeconds) * time.Second)
+
+		case <-server.rateLimitSignalGC:
+			runtime.GC()
+
+		case <-server.stopBroadcast:
+			return
+		}
+	}
+}
+
 func (server *MeekServer) deleteSession(sessionID string) {
 
 	// Don't obtain the server.sessionsLock write lock until modifying

psiphon/server/meek_test.go

@@ -240,6 +240,7 @@ func TestMeekResiliency(t *testing.T) {
 			MeekObfuscatedKey:              meekObfuscatedKey,
 			MeekCookieEncryptionPrivateKey: meekCookieEncryptionPrivateKey,
 		},
+		TrafficRulesSet: &TrafficRulesSet{},
 	}
 
 	listener, err := net.Listen("tcp", "127.0.0.1:0")
@@ -371,3 +372,152 @@ func (interruptor *fileDescriptorInterruptor) BindToDevice(fileDescriptor int) (
 	})
 	return "", nil
 }
+
+func TestMeekRateLimiter(t *testing.T) {
+
+	allowedConnections := 5
+	testDurationSeconds := 10
+
+	// Run meek server
+
+	rawMeekCookieEncryptionPublicKey, rawMeekCookieEncryptionPrivateKey, err := box.GenerateKey(crypto_rand.Reader)
+	if err != nil {
+		t.Fatalf("box.GenerateKey failed: %s", err)
+	}
+	meekCookieEncryptionPublicKey := base64.StdEncoding.EncodeToString(rawMeekCookieEncryptionPublicKey[:])
+	meekCookieEncryptionPrivateKey := base64.StdEncoding.EncodeToString(rawMeekCookieEncryptionPrivateKey[:])
+	meekObfuscatedKey, err := common.MakeSecureRandomStringHex(SSH_OBFUSCATED_KEY_BYTE_LENGTH)
+	if err != nil {
+		t.Fatalf("common.MakeSecureRandomStringHex failed: %s", err)
+	}
+
+	mockSupport := &SupportServices{
+		Config: &Config{
+			MeekObfuscatedKey:              meekObfuscatedKey,
+			MeekCookieEncryptionPrivateKey: meekCookieEncryptionPrivateKey,
+		},
+		TrafficRulesSet: &TrafficRulesSet{
+			MeekRateLimiterHistorySize:                   allowedConnections,
+			MeekRateLimiterThresholdSeconds:              testDurationSeconds,
+			MeekRateLimiterGarbageCollectionTriggerCount: 1,
+			MeekRateLimiterReapHistoryFrequencySeconds:   1,
+		},
+	}
+
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("net.Listen failed: %s", err)
+	}
+	defer listener.Close()
+
+	serverAddress := listener.Addr().String()
+
+	stopBroadcast := make(chan struct{})
+
+	server, err := NewMeekServer(
+		mockSupport,
+		listener,
+		false,
+		false,
+		func(_ string, conn net.Conn) {
+			go func() {
+				for {
+					buffer := make([]byte, 1)
+					n, err := conn.Read(buffer)
+					if err == nil && n == 1 {
+						_, err = conn.Write(buffer)
+					}
+					if err != nil {
+						conn.Close()
+						break
+					}
+				}
+			}()
+		},
+		stopBroadcast)
+	if err != nil {
+		t.Fatalf("NewMeekServer failed: %s", err)
+	}
+
+	serverWaitGroup := new(sync.WaitGroup)
+
+	serverWaitGroup.Add(1)
+	go func() {
+		defer serverWaitGroup.Done()
+		err := server.Run()
+		if err != nil {
+			t.Fatalf("MeekServer.Run failed: %s", err)
+		}
+	}()
+
+	// Run meek clients:
+	// For 10 seconds, connect once per second vs. rate limit of 5-per-10 seconds,
+	// so about half of the connections should be rejected by the rate limiter.
+
+	stopTime := time.Now().Add(time.Duration(testDurationSeconds) * time.Second)
+
+	totalConnections := 0
+	totalFailures := 0
+
+	for {
+
+		dialConfig := &psiphon.DialConfig{}
+
+		clientParameters, err := parameters.NewClientParameters(nil)
+		if err != nil {
+			t.Fatalf("NewClientParameters failed: %s", err)
+		}
+
+		meekConfig := &psiphon.MeekConfig{
+			ClientParameters:              clientParameters,
+			DialAddress:                   serverAddress,
+			HostHeader:                    "example.com",
+			MeekCookieEncryptionPublicKey: meekCookieEncryptionPublicKey,
+			MeekObfuscatedKey:             meekObfuscatedKey,
+		}
+
+		ctx, cancelFunc := context.WithTimeout(
+			context.Background(), 500*time.Millisecond)
+		defer cancelFunc()
+
+		clientConn, err := psiphon.DialMeek(ctx, meekConfig, dialConfig)
+
+		if err == nil {
+			_, err = clientConn.Write([]byte{0})
+		}
+		if err == nil {
+			buffer := make([]byte, 1)
+			_, err = clientConn.Read(buffer)
+		}
+
+		if clientConn != nil {
+			clientConn.Close()
+		}
+
+		if err != nil {
+			totalFailures += 1
+		} else {
+			totalConnections += 1
+		}
+
+		if !time.Now().Before(stopTime) {
+			break
+		}
+
+		time.Sleep(1 * time.Second)
+	}
+
+	if totalConnections != allowedConnections || totalFailures == 0 {
+		t.Fatalf(
+			"Unexpected results: %d connections, %d failures",
+			totalConnections, totalFailures)
+	}
+
+	// Graceful shutdown
+
+	listener.Close()
+	close(stopBroadcast)
+
+	// This wait will hang if shutdown is broken, and the test will ultimately panic
+	serverWaitGroup.Wait()
+}

psiphon/server/server_test.go

@@ -1037,7 +1037,12 @@ func paveTrafficRulesFile(
                 "WriteBytesPerSecond": 16384
             },
             "AllowTCPPorts" : [0],
-            "AllowUDPPorts" : [0]
+            "AllowUDPPorts" : [0],
+            "MeekRateLimiterHistorySize" : 10,
+            "MeekRateLimiterThresholdSeconds" : 1,
+            "MeekRateLimiterGarbageCollectionTriggerCount" : 1,
+            "MeekRateLimiterReapHistoryFrequencySeconds" : 1,
+            "MeekRateLimiterRegions" : []
         },
         "FilteredRules" : [
             {

psiphon/server/trafficRules.go

@@ -21,6 +21,7 @@ package server
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"net"
 
@@ -28,12 +29,14 @@ import (
 )
 
 const (
-	DEFAULT_IDLE_TCP_PORT_FORWARD_TIMEOUT_MILLISECONDS = 30000
-	DEFAULT_IDLE_UDP_PORT_FORWARD_TIMEOUT_MILLISECONDS = 30000
-	DEFAULT_DIAL_TCP_PORT_FORWARD_TIMEOUT_MILLISECONDS = 10000
-	DEFAULT_MAX_TCP_DIALING_PORT_FORWARD_COUNT         = 64
-	DEFAULT_MAX_TCP_PORT_FORWARD_COUNT                 = 512
-	DEFAULT_MAX_UDP_PORT_FORWARD_COUNT                 = 32
+	DEFAULT_IDLE_TCP_PORT_FORWARD_TIMEOUT_MILLISECONDS        = 30000
+	DEFAULT_IDLE_UDP_PORT_FORWARD_TIMEOUT_MILLISECONDS        = 30000
+	DEFAULT_DIAL_TCP_PORT_FORWARD_TIMEOUT_MILLISECONDS        = 10000
+	DEFAULT_MAX_TCP_DIALING_PORT_FORWARD_COUNT                = 64
+	DEFAULT_MAX_TCP_PORT_FORWARD_COUNT                        = 512
+	DEFAULT_MAX_UDP_PORT_FORWARD_COUNT                        = 32
+	DEFAULT_MEEK_RATE_LIMITER_GARBAGE_COLLECTOR_TRIGGER_COUNT = 5000
+	DEFAULT_MEEK_RATE_LIMITER_REAP_HISTORY_FREQUENCY_SECONDS  = 600
 )
 
 // TrafficRulesSet represents the various traffic rules to
@@ -59,6 +62,44 @@ type TrafficRulesSet struct {
 		Filter TrafficRulesFilter
 		Rules  TrafficRules
 	}
+
+	// MeekRateLimiterHistorySize enables the late-stage meek rate limiter and
+	// sets its history size. The late-stage meek rate limiter acts on client
+	// IPs relayed in MeekProxyForwardedForHeaders, and so it must wait for
+	// the HTTP headers to be read. This rate limiter immediately terminates
+	// any client endpoint request or any request to create a new session, but
+	// not any meek request for an existing session, if the
+	// MeekRateLimiterHistorySize requests occur in
+	// MeekRateLimiterThresholdSeconds. The scope of rate limiting may be
+	// limited using LimitMeekRateLimiterRegions.
+	//
+	// Hot reloading a new history size will result in existing history being
+	// truncated.
+	MeekRateLimiterHistorySize int
+
+	// MeekRateLimiterThresholdSeconds is part of the meek rate limiter
+	// specification and must be set when MeekRateLimiterHistorySize is set.
+	MeekRateLimiterThresholdSeconds int
+
+	// MeekRateLimiterRegions, if set, limits application of the meek
+	// late-stage rate limiter to clients in the specified list of GeoIP
+	// countries. When omitted or empty, meek rate limiting, if configured,
+	// is applied to all clients.
+	MeekRateLimiterRegions []string
+
+	// MeekRateLimiterGarbageCollectionTriggerCount specifies the number of
+	// rate limit events after which garbage collection is manually triggered
+	// in order to reclaim memory used by rate limited and other rejected
+	// requests.
+	// A default of 5000 is used when
+	// MeekRateLimiterGarbageCollectionTriggerCount is 0.
+	MeekRateLimiterGarbageCollectionTriggerCount int
+
+	// MeekRateLimiterReapHistoryFrequencySeconds specifies a schedule for
+	// reaping old records from the rate limit history.
+	// A default of 600 is used when
+	// MeekRateLimiterReapHistoryFrequencySeconds is 0.
+	MeekRateLimiterReapHistoryFrequencySeconds int
 }
 
 // TrafficRulesFilter defines a filter to match against client attributes.
@@ -232,7 +273,37 @@ func NewTrafficRulesSet(filename string) (*TrafficRulesSet, error) {
 // Validate checks for correct input formats in a TrafficRulesSet.
 func (set *TrafficRulesSet) Validate() error {
 
+	if set.MeekRateLimiterHistorySize < 0 ||
+		set.MeekRateLimiterThresholdSeconds < 0 ||
+		set.MeekRateLimiterGarbageCollectionTriggerCount < 0 ||
+		set.MeekRateLimiterReapHistoryFrequencySeconds < 0 {
+		return common.ContextError(
+			errors.New("MeekRateLimiter values must be >= 0"))
+	}
+
+	if set.MeekRateLimiterHistorySize > 0 {
+		if set.MeekRateLimiterThresholdSeconds <= 0 {
+			return common.ContextError(
+				errors.New("MeekRateLimiterThresholdSeconds must be > 0"))
+		}
+	}
+
 	validateTrafficRules := func(rules *TrafficRules) error {
+
+		if (rules.RateLimits.ReadUnthrottledBytes != nil && *rules.RateLimits.ReadUnthrottledBytes < 0) ||
+			(rules.RateLimits.ReadBytesPerSecond != nil && *rules.RateLimits.ReadBytesPerSecond < 0) ||
+			(rules.RateLimits.WriteUnthrottledBytes != nil && *rules.RateLimits.WriteUnthrottledBytes < 0) ||
+			(rules.RateLimits.WriteBytesPerSecond != nil && *rules.RateLimits.WriteBytesPerSecond < 0) ||
+			(rules.DialTCPPortForwardTimeoutMilliseconds != nil && *rules.DialTCPPortForwardTimeoutMilliseconds < 0) ||
+			(rules.IdleTCPPortForwardTimeoutMilliseconds != nil && *rules.IdleTCPPortForwardTimeoutMilliseconds < 0) ||
+			(rules.IdleUDPPortForwardTimeoutMilliseconds != nil && *rules.IdleUDPPortForwardTimeoutMilliseconds < 0) ||
+			(rules.MaxTCPDialingPortForwardCount != nil && *rules.MaxTCPDialingPortForwardCount < 0) ||
+			(rules.MaxTCPPortForwardCount != nil && *rules.MaxTCPPortForwardCount < 0) ||
+			(rules.MaxUDPPortForwardCount != nil && *rules.MaxUDPPortForwardCount < 0) {
+			return common.ContextError(
+				errors.New("TrafficRules values must be >= 0"))
+		}
+
 		for _, subnet := range rules.AllowSubnets {
 			_, _, err := net.ParseCIDR(subnet)
 			if err != nil {
@@ -240,6 +311,7 @@ func (set *TrafficRulesSet) Validate() error {
 					fmt.Errorf("invalid subnet: %s %s", subnet, err))
 			}
 		}
+
 		return nil
 	}
 
@@ -369,6 +441,10 @@ func (set *TrafficRulesSet) GetTrafficRules(
 		trafficRules.AllowUDPPorts = make([]int, 0)
 	}
 
+	if trafficRules.AllowSubnets == nil {
+		trafficRules.AllowSubnets = make([]string, 0)
+	}
+
 	// TODO: faster lookup?
 	for _, filteredRules := range set.FilteredRules {
 
@@ -496,6 +572,10 @@ func (set *TrafficRulesSet) GetTrafficRules(
 			trafficRules.AllowUDPPorts = filteredRules.Rules.AllowUDPPorts
 		}
 
+		if filteredRules.Rules.AllowSubnets != nil {
+			trafficRules.AllowSubnets = filteredRules.Rules.AllowSubnets
+		}
+
 		break
 	}
 
@@ -508,3 +588,28 @@ func (set *TrafficRulesSet) GetTrafficRules(
 
 	return trafficRules
 }
+
+// GetMeekRateLimiterConfig gets a snapshot of the meek rate limiter
+// configuration values.
+func (set *TrafficRulesSet) GetMeekRateLimiterConfig() (int, int, []string, int, int) {
+
+	set.ReloadableFile.RLock()
+	defer set.ReloadableFile.RUnlock()
+
+	GCTriggerCount := set.MeekRateLimiterGarbageCollectionTriggerCount
+	if GCTriggerCount <= 0 {
+		GCTriggerCount = DEFAULT_MEEK_RATE_LIMITER_GARBAGE_COLLECTOR_TRIGGER_COUNT
+	}
+
+	reapFrequencySeconds := set.MeekRateLimiterReapHistoryFrequencySeconds
+	if reapFrequencySeconds <= 0 {
+		reapFrequencySeconds = DEFAULT_MEEK_RATE_LIMITER_REAP_HISTORY_FREQUENCY_SECONDS
+
+	}
+
+	return set.MeekRateLimiterHistorySize,
+		set.MeekRateLimiterThresholdSeconds,
+		set.MeekRateLimiterRegions,
+		GCTriggerCount,
+		reapFrequencySeconds
+}