Merge pull request #168 from rod-hynes/master

Throttling rules; more command line flags and config

Server/main.go

@@ -24,28 +24,76 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
+	"strings"
 
 	"github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/server"
 )
 
 func main() {
 
+	var generateServerIPaddress string
+	var generateWebServerPort, generateSSHServerPort, generateObfuscatedSSHServerPort int
+	var runConfigFilenames stringListFlag
+
+	flag.StringVar(
+		&generateServerIPaddress,
+		"ipaddress",
+		server.DEFAULT_SERVER_IP_ADDRESS,
+		"generate with this server `IP address`")
+
+	flag.IntVar(
+		&generateWebServerPort,
+		"webport",
+		server.DEFAULT_WEB_SERVER_PORT,
+		"generate with this web server `port`; 0 for no web server")
+
+	flag.IntVar(
+		&generateSSHServerPort,
+		"sshport",
+		server.DEFAULT_SSH_SERVER_PORT,
+		"generate with this SSH server `port`; 0 for no SSH server")
+
+	flag.IntVar(
+		&generateObfuscatedSSHServerPort,
+		"osshport",
+		server.DEFAULT_OBFUSCATED_SSH_SERVER_PORT,
+		"generate with this Obfuscated SSH server `port`; 0 for no Obfuscated SSH server")
+
+	flag.Var(
+		&runConfigFilenames,
+		"config",
+		"run with this config `filename`; flag may be repeated to load multiple config files")
+
+	flag.Usage = func() {
+		fmt.Fprintf(os.Stderr,
+			"Usage:\n\n"+
+				"%s <flags> generate    generates a configuration and server entry\n"+
+				"%s <flags> run         runs configured services\n\n",
+			os.Args[0], os.Args[0])
+		flag.PrintDefaults()
+	}
+
 	flag.Parse()
 
 	args := flag.Args()
 
-	// TODO: add working directory flag
 	configFilename := server.SERVER_CONFIG_FILENAME
+
 	serverEntryFilename := server.SERVER_ENTRY_FILENAME
 
 	if len(args) < 1 {
-		fmt.Errorf("usage: '%s generate' or '%s run'", os.Args[0])
+		flag.Usage()
 		os.Exit(1)
 	} else if args[0] == "generate" {
 
-		// TODO: flags to set generate params
 		configFileContents, serverEntryFileContents, err := server.GenerateConfig(
-			&server.GenerateConfigParams{})
+			&server.GenerateConfigParams{
+				ServerIPAddress:         generateServerIPaddress,
+				WebServerPort:           generateWebServerPort,
+				SSHServerPort:           generateSSHServerPort,
+				ObfuscatedSSHServerPort: generateObfuscatedSSHServerPort,
+			})
+
 		if err != nil {
 			fmt.Errorf("generate failed: %s", err)
 			os.Exit(1)
@@ -64,16 +112,37 @@ func main() {
 
 	} else if args[0] == "run" {
 
-		configFileContents, err := ioutil.ReadFile(configFilename)
-		if err != nil {
-			fmt.Errorf("error loading configuration file: %s", err)
-			os.Exit(1)
+		if len(runConfigFilenames) == 0 {
+			runConfigFilenames = []string{configFilename}
+		}
+
+		var configFileContents [][]byte
+
+		for _, configFilename := range runConfigFilenames {
+			ccontents, err := ioutil.ReadFile(configFilename)
+			if err != nil {
+				fmt.Errorf("error loading configuration file: %s", err)
+				os.Exit(1)
+			}
+
+			configFileContents = append(configFileContents, ccontents)
 		}
 
-		err = server.RunServices(configFileContents)
+		err := server.RunServices(configFileContents)
 		if err != nil {
 			fmt.Errorf("run failed: %s", err)
 			os.Exit(1)
 		}
 	}
 }
+
+type stringListFlag []string
+
+func (list *stringListFlag) String() string {
+	return strings.Join(*list, ", ")
+}
+
+func (list *stringListFlag) Set(flagValue string) error {
+	*list = append(*list, flagValue)
+	return nil
+}

psiphon/config.go

@@ -158,8 +158,8 @@ type Config struct {
 	TunnelProtocol string
 
 	// EstablishTunnelTimeoutSeconds specifies a time limit after which to halt
-	// the core tunnel controller if no tunnel has been established. By default,
-	// the controller will keep trying indefinitely.
+	// the core tunnel controller if no tunnel has been established. The default
+	// is ESTABLISH_TUNNEL_TIMEOUT_SECONDS.
 	EstablishTunnelTimeoutSeconds *int
 
 	// ListenInterface specifies which interface to listen on.  If no interface

psiphon/net.go

@@ -391,25 +391,39 @@ func IPAddressFromAddr(addr net.Addr) string {
 	return ipAddress
 }
 
-// TimeoutTCPConn wraps a net.TCPConn and sets an initial ReadDeadline. The
-// deadline is reset whenever data is received from the connection.
-type TimeoutTCPConn struct {
-	*net.TCPConn
-	deadline time.Duration
+// IdleTimeoutConn wraps a net.Conn and sets an initial ReadDeadline. The
+// deadline is extended whenever data is received from the connection.
+// Optionally, IdleTimeoutConn will also extend the deadline when data is
+// written to the connection.
+type IdleTimeoutConn struct {
+	net.Conn
+	deadline     time.Duration
+	resetOnWrite bool
 }
 
-func NewTimeoutTCPConn(tcpConn *net.TCPConn, deadline time.Duration) *TimeoutTCPConn {
-	tcpConn.SetReadDeadline(time.Now().Add(deadline))
-	return &TimeoutTCPConn{
-		TCPConn:  tcpConn,
-		deadline: deadline,
+func NewIdleTimeoutConn(
+	conn net.Conn, deadline time.Duration, resetOnWrite bool) *IdleTimeoutConn {
+
+	conn.SetReadDeadline(time.Now().Add(deadline))
+	return &IdleTimeoutConn{
+		Conn:         conn,
+		deadline:     deadline,
+		resetOnWrite: resetOnWrite,
 	}
 }
 
-func (conn *TimeoutTCPConn) Read(buffer []byte) (int, error) {
-	n, err := conn.TCPConn.Read(buffer)
+func (conn *IdleTimeoutConn) Read(buffer []byte) (int, error) {
+	n, err := conn.Conn.Read(buffer)
 	if err == nil {
-		conn.TCPConn.SetReadDeadline(time.Now().Add(conn.deadline))
+		conn.Conn.SetReadDeadline(time.Now().Add(conn.deadline))
+	}
+	return n, err
+}
+
+func (conn *IdleTimeoutConn) Write(buffer []byte) (int, error) {
+	n, err := conn.Conn.Write(buffer)
+	if err == nil && conn.resetOnWrite {
+		conn.Conn.SetReadDeadline(time.Now().Add(conn.deadline))
 	}
 	return n, err
 }

psiphon/server/config.go

@@ -26,6 +26,7 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"encoding/pem"
+	"errors"
 	"fmt"
 	"math/big"
 	"strings"
@@ -54,6 +55,7 @@ const (
 	DEFAULT_SSH_SERVER_PORT                = 2222
 	SSH_HANDSHAKE_TIMEOUT                  = 30 * time.Second
 	SSH_CONNECTION_READ_DEADLINE           = 5 * time.Minute
+	SSH_THROTTLED_PORT_FORWARD_MAX_COPY    = 32 * 1024
 	SSH_OBFUSCATED_KEY_BYTE_LENGTH         = 32
 	DEFAULT_OBFUSCATED_SSH_SERVER_PORT     = 3333
 	REDIS_POOL_MAX_IDLE                    = 50
@@ -88,7 +90,8 @@ type Config struct {
 	DiscoveryValueHMACKey string
 
 	// GeoIPDatabaseFilename is the path of the GeoIP2/GeoLite2
-	// MaxMind database file.
+	// MaxMind database file. when blank, no GeoIP lookups are
+	// performed.
 	GeoIPDatabaseFilename string
 
 	// ServerIPAddress is the public IP address of the server.
@@ -110,6 +113,10 @@ type Config struct {
 	// authenticate itself to clients.
 	WebServerPrivateKey string
 
+	// SSHServerPort is the listening port of the SSH server.
+	// When <= 0, no SSH server component is run.
+	SSHServerPort int
+
 	// SSHPrivateKey is the SSH host key. The same key is used for
 	// both the SSH and Obfuscated SSH servers.
 	SSHPrivateKey string
@@ -129,21 +136,53 @@ type Config struct {
 	// and Obfuscated SSH servers.
 	SSHPassword string
 
-	// SSHServerPort is the listening port of the SSH server.
-	// When <= 0, no SSH server component is run.
-	SSHServerPort int
+	// ObfuscatedSSHServerPort is the listening port of the Obfuscated SSH server.
+	// When <= 0, no Obfuscated SSH server component is run.
+	ObfuscatedSSHServerPort int
 
 	// ObfuscatedSSHKey is the secret key for use in the Obfuscated
 	// SSH protocol.
 	ObfuscatedSSHKey string
 
-	// ObfuscatedSSHServerPort is the listening port of the Obfuscated SSH server.
-	// When <= 0, no Obfuscated SSH server component is run.
-	ObfuscatedSSHServerPort int
-
 	// RedisServerAddress is the TCP address of a redis server. When
 	// set, redis is used to store per-session GeoIP information.
 	RedisServerAddress string
+
+	// DefaultTrafficRules specifies the traffic rules to be used when
+	// no regional-specific rules are set.
+	DefaultTrafficRules TrafficRules
+
+	// RegionalTrafficRules specifies the traffic rules for particular
+	// client regions (countries) as determined by GeoIP lookup of the
+	// client IP address. The key for each regional traffic rule entry
+	// is one or more space delimited ISO 3166-1 alpha-2 country codes.
+	RegionalTrafficRules map[string]TrafficRules
+}
+
+// TrafficRules specify the limits placed on SSH client port forward
+// traffic.
+type TrafficRules struct {
+
+	// ThrottleUpstreamSleepMilliseconds is the period to sleep
+	// between sending each chunk of client->destination traffic.
+	// The default, 0, is no sleep.
+	ThrottleUpstreamSleepMilliseconds int
+
+	// ThrottleDownstreamSleepMilliseconds is the period to sleep
+	// between sending each chunk of destination->client traffic.
+	// The default, 0, is no sleep.
+	ThrottleDownstreamSleepMilliseconds int
+
+	// IdlePortForwardTimeoutMilliseconds is the timeout period
+	// after which idle (no bytes flowing in either direction)
+	// SSH client port forwards are preemptively closed.
+	// The default, 0, is no idle timeout.
+	IdlePortForwardTimeoutMilliseconds int
+
+	// MaxClientPortForwardCount is the maximum number of port
+	// forwards each client may have open concurrently.
+	// The default, 0, is no maximum.
+	MaxClientPortForwardCount int
 }
 
 // RunWebServer indicates whether to run a web server component.
@@ -167,17 +206,61 @@ func (config *Config) UseRedis() bool {
 	return config.RedisServerAddress != ""
 }
 
-// LoadConfig loads and validates a JSON encoded server config.
-func LoadConfig(configJson []byte) (*Config, error) {
+// GetTrafficRules looks up the traffic rules for the specified country. If there
+// are no RegionalTrafficRules for the country, DefaultTrafficRules are returned.
+func (config *Config) GetTrafficRules(targetCountryCode string) TrafficRules {
+	// TODO: faster lookup?
+	for countryCodes, trafficRules := range config.RegionalTrafficRules {
+		for _, countryCode := range strings.Split(countryCodes, " ") {
+			if countryCode == targetCountryCode {
+				return trafficRules
+			}
+		}
+	}
+	return config.DefaultTrafficRules
+}
+
+// LoadConfig loads and validates a JSON encoded server config. If more than one
+// JSON config is specified, then all are loaded and values are merged together,
+// in order. Multiple configs allows for use cases like storing static, server-specific
+// values in a base config while also deploying network-wide throttling settings
+// in a secondary file that can be paved over on all server hosts.
+func LoadConfig(configJSONs [][]byte) (*Config, error) {
 
+	// Note: default values are set in GenerateConfig
 	var config Config
-	err := json.Unmarshal(configJson, &config)
-	if err != nil {
-		return nil, psiphon.ContextError(err)
+
+	for _, configJSON := range configJSONs {
+		err := json.Unmarshal(configJSON, &config)
+		if err != nil {
+			return nil, psiphon.ContextError(err)
+		}
+	}
+
+	if config.ServerIPAddress == "" {
+		return nil, errors.New("server IP address is missing from config file")
 	}
 
-	// TODO: config field validation
-	// TODO: validation case: OSSH requires extra fields
+	if config.WebServerPort > 0 && (config.WebServerSecret == "" || config.WebServerCertificate == "" ||
+		config.WebServerPrivateKey == "") {
+
+		return nil, errors.New(
+			"web server requires WebServerSecret, WebServerCertificate, WebServerPrivateKey")
+	}
+
+	if config.SSHServerPort > 0 && (config.SSHPrivateKey == "" || config.SSHServerVersion == "" ||
+		config.SSHUserName == "" || config.SSHPassword == "") {
+
+		return nil, errors.New(
+			"SSH server requires SSHPrivateKey, SSHServerVersion, SSHUserName, SSHPassword")
+	}
+
+	if config.ObfuscatedSSHServerPort > 0 && (config.SSHPrivateKey == "" || config.SSHServerVersion == "" ||
+		config.SSHUserName == "" || config.SSHPassword == "" || config.ObfuscatedSSHKey == "") {
+
+		return nil, errors.New(
+			"Obfuscated SSH server requires SSHPrivateKey, SSHServerVersion, SSHUserName, SSHPassword, ObfuscatedSSHKey")
+	}
 
 	return &config, nil
 }
@@ -293,7 +376,7 @@ func GenerateConfig(params *GenerateConfigParams) ([]byte, []byte, error) {
 		LogLevel:                DEFAULT_LOG_LEVEL,
 		SyslogAddress:           "",
 		SyslogFacility:          "",
-		SyslogTag:               "",
+		SyslogTag:               DEFAULT_SYSLOG_TAG,
 		DiscoveryValueHMACKey:   "",
 		GeoIPDatabaseFilename:   DEFAULT_GEO_IP_DATABASE_FILENAME,
 		ServerIPAddress:         serverIPaddress,

psiphon/server/log.go

@@ -77,12 +77,7 @@ var log *ContextLogger
 // goroutine.
 func InitLogging(config *Config) error {
 
-	logLevel := DEFAULT_LOG_LEVEL
-	if config.LogLevel != "" {
-		logLevel = config.LogLevel
-	}
-
-	level, err := logrus.ParseLevel(logLevel)
+	level, err := logrus.ParseLevel(config.LogLevel)
 	if err != nil {
 		return psiphon.ContextError(err)
 	}
@@ -93,16 +88,11 @@ func InitLogging(config *Config) error {
 
 	if config.SyslogAddress != "" {
 
-		tag := DEFAULT_SYSLOG_TAG
-		if config.SyslogTag != "" {
-			tag = config.SyslogTag
-		}
-
 		syslogHook, err = logrus_syslog.NewSyslogHook(
 			"udp",
 			config.SyslogAddress,
 			getSyslogPriority(config),
-			tag)
+			config.SyslogTag)
 
 		if err != nil {
 			return psiphon.ContextError(err)

psiphon/server/services.go

@@ -35,9 +35,9 @@ import (
 // redis connection pooling; and then starts the server components and runs them
 // until os.Interrupt or os.Kill signals are received. The config determines
 // which components are run.
-func RunServices(encodedConfig []byte) error {
+func RunServices(encodedConfigs [][]byte) error {
 
-	config, err := LoadConfig(encodedConfig)
+	config, err := LoadConfig(encodedConfigs)
 	if err != nil {
 		log.WithContextFields(LogFields{"error": err}).Error("load config failed")
 		return psiphon.ContextError(err)

psiphon/server/sshService.go

@@ -211,16 +211,16 @@ func (sshServer *sshServer) stopClient(client *sshClient) {
 	client.Lock()
 	log.WithContextFields(
 		LogFields{
-			"startTime":                     client.startTime,
-			"duration":                      time.Now().Sub(client.startTime),
-			"psiphonSessionID":              client.psiphonSessionID,
-			"country":                       client.geoIPData.Country,
-			"city":                          client.geoIPData.City,
-			"ISP":                           client.geoIPData.ISP,
-			"bytesUp":                       client.bytesUp,
-			"bytesDown":                     client.bytesDown,
-			"portForwardCount":              client.portForwardCount,
-			"maxConcurrentPortForwardCount": client.maxConcurrentPortForwardCount,
+			"startTime":                      client.startTime,
+			"duration":                       time.Now().Sub(client.startTime),
+			"psiphonSessionID":               client.psiphonSessionID,
+			"country":                        client.geoIPData.Country,
+			"city":                           client.geoIPData.City,
+			"ISP":                            client.geoIPData.ISP,
+			"bytesUp":                        client.bytesUp,
+			"bytesDown":                      client.bytesDown,
+			"portForwardCount":               client.portForwardCount,
+			"peakConcurrentPortForwardCount": client.peakConcurrentPortForwardCount,
 		}).Info("tunnel closed")
 	client.Unlock()
 }
@@ -244,14 +244,15 @@ func (sshServer *sshServer) handleClient(tcpConn *net.TCPConn) {
 		startTime: time.Now(),
 		geoIPData: GeoIPLookup(psiphon.IPAddressFromAddr(tcpConn.RemoteAddr())),
 	}
+	sshClient.trafficRules = sshServer.config.GetTrafficRules(sshClient.geoIPData.Country)
 
-	// Wrap the base TCP connection in a TimeoutTCPConn which will terminate
-	// the connection if it's idle for too long. This timeout is in effect for
-	// the entire duration of the SSH connection. Clients must actively use
-	// the connection or send SSH keep alive requests to keep the connection
+	// Wrap the base TCP connection with an IdleTimeoutConn which will terminate
+	// the connection if no data is received before the deadline. This timeout is
+	// in effect for the entire duration of the SSH connection. Clients must actively
+	// use the connection or send SSH keep alive requests to keep the connection
 	// active.
 
-	conn := psiphon.NewTimeoutTCPConn(tcpConn, SSH_CONNECTION_READ_DEADLINE)
+	conn := psiphon.NewIdleTimeoutConn(tcpConn, SSH_CONNECTION_READ_DEADLINE, false)
 
 	// Run the initial [obfuscated] SSH handshake in a goroutine so we can both
 	// respect shutdownBroadcast and implement a specific handshake timeout.
@@ -320,7 +321,7 @@ func (sshServer *sshServer) handleClient(tcpConn *net.TCPConn) {
 
 	clientID, ok := sshServer.registerClient(sshClient)
 	if !ok {
-		tcpConn.Close()
+		conn.Close()
 		log.WithContext().Warning("register failed")
 		return
 	}
@@ -333,16 +334,17 @@ func (sshServer *sshServer) handleClient(tcpConn *net.TCPConn) {
 
 type sshClient struct {
 	sync.Mutex
-	sshServer                     *sshServer
-	sshConn                       ssh.Conn
-	startTime                     time.Time
-	geoIPData                     GeoIPData
-	psiphonSessionID              string
-	bytesUp                       int64
-	bytesDown                     int64
-	portForwardCount              int64
-	concurrentPortForwardCount    int64
-	maxConcurrentPortForwardCount int64
+	sshServer                      *sshServer
+	sshConn                        ssh.Conn
+	startTime                      time.Time
+	geoIPData                      GeoIPData
+	trafficRules                   TrafficRules
+	psiphonSessionID               string
+	bytesUp                        int64
+	bytesDown                      int64
+	portForwardCount               int64
+	concurrentPortForwardCount     int64
+	peakConcurrentPortForwardCount int64
 }
 
 func (sshClient *sshClient) handleChannels(channels <-chan ssh.NewChannel) {
@@ -353,6 +355,18 @@ func (sshClient *sshClient) handleChannels(channels <-chan ssh.NewChannel) {
 			return
 		}
 
+		if sshClient.trafficRules.MaxClientPortForwardCount > 0 {
+			sshClient.Lock()
+			limitExceeded := sshClient.portForwardCount >= int64(sshClient.trafficRules.MaxClientPortForwardCount)
+			sshClient.Unlock()
+
+			if limitExceeded {
+				sshClient.rejectNewChannel(
+					newChannel, ssh.ResourceShortage, "maximum port forward limit exceeded")
+				return
+			}
+		}
+
 		// process each port forward concurrently
 		go sshClient.handleNewDirectTcpipChannel(newChannel)
 	}
@@ -410,8 +424,8 @@ func (sshClient *sshClient) handleNewDirectTcpipChannel(newChannel ssh.NewChanne
 	sshClient.Lock()
 	sshClient.portForwardCount += 1
 	sshClient.concurrentPortForwardCount += 1
-	if sshClient.concurrentPortForwardCount > sshClient.maxConcurrentPortForwardCount {
-		sshClient.maxConcurrentPortForwardCount = sshClient.concurrentPortForwardCount
+	if sshClient.concurrentPortForwardCount > sshClient.peakConcurrentPortForwardCount {
+		sshClient.peakConcurrentPortForwardCount = sshClient.concurrentPortForwardCount
 	}
 	sshClient.Unlock()
 
@@ -421,9 +435,19 @@ func (sshClient *sshClient) handleNewDirectTcpipChannel(newChannel ssh.NewChanne
 
 	defer fwdChannel.Close()
 
-	// relay channel to forwarded connection
+	// When idle port forward traffic rules are in place, wrap fwdConn
+	// in an IdleTimeoutConn configured to reset idle on writes as well
+	// as read. This ensures the port forward idle timeout only happens
+	// when both upstream and downstream directions are are idle.
 
-	// TODO: use a low-memory io.Copy?
+	if sshClient.trafficRules.IdlePortForwardTimeoutMilliseconds > 0 {
+		fwdConn = psiphon.NewIdleTimeoutConn(
+			fwdConn,
+			time.Duration(sshClient.trafficRules.IdlePortForwardTimeoutMilliseconds)*time.Millisecond,
+			true)
+	}
+
+	// relay channel to forwarded connection
 	// TODO: relay errors to fwdChannel.Stderr()?
 
 	var bytesUp, bytesDown int64
@@ -433,12 +457,14 @@ func (sshClient *sshClient) handleNewDirectTcpipChannel(newChannel ssh.NewChanne
 	go func() {
 		defer relayWaitGroup.Done()
 		var err error
-		bytesUp, err = io.Copy(fwdConn, fwdChannel)
+		bytesUp, err = copyWithThrottle(
+			fwdConn, fwdChannel, sshClient.trafficRules.ThrottleUpstreamSleepMilliseconds)
 		if err != nil {
 			log.WithContextFields(LogFields{"error": err}).Warning("upstream relay failed")
 		}
 	}()
-	bytesDown, err = io.Copy(fwdChannel, fwdConn)
+	bytesDown, err = copyWithThrottle(
+		fwdChannel, fwdConn, sshClient.trafficRules.ThrottleDownstreamSleepMilliseconds)
 	if err != nil {
 		log.WithContextFields(LogFields{"error": err}).Warning("downstream relay failed")
 	}
@@ -454,6 +480,28 @@ func (sshClient *sshClient) handleNewDirectTcpipChannel(newChannel ssh.NewChanne
 	log.WithContextFields(LogFields{"target": targetAddr}).Debug("exiting")
 }
 
+func copyWithThrottle(dst io.Writer, src io.Reader, throttleSleepMilliseconds int) (int64, error) {
+	// TODO: use a low-memory io.Copy?
+	if throttleSleepMilliseconds <= 0 {
+		// No throttle
+		return io.Copy(dst, src)
+	}
+	var totalBytes int64
+	for {
+		bytes, err := io.CopyN(dst, src, SSH_THROTTLED_PORT_FORWARD_MAX_COPY)
+		totalBytes += bytes
+		if err == io.EOF {
+			err = nil
+			break
+		}
+		if err != nil {
+			return totalBytes, psiphon.ContextError(err)
+		}
+		time.Sleep(time.Duration(throttleSleepMilliseconds) * time.Millisecond)
+	}
+	return totalBytes, nil
+}
+
 func (sshClient *sshClient) passwordCallback(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
 	var sshPasswordPayload struct {
 		SessionId   string `json:"SessionId"`