<commit_before>package manssh

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/stretchr/testify/require"
)

const (
	mainConfigContent = `
Include %s/config.d/*
Host home1
    hostname 192.168.1.11
Host main1
    hostname 192.168.1.10
Host main2
    hostname 192.168.1.20
    user wen
    port 77
Host main3
    hostname 192.168.1.30
    user ROOT
	port 77
`
	testConfigContent = `
Host *
	port 22022
Host test1
    hostname 192.168.2.10
    user root
    port 22
Host test2 main2
    hostname 192.168.2.20
    port 77
Host Test3
    hostname 192.168.2.30
    user ROOT
	port 77
`
	homeConfigContent = `
Host home1
    hostname 192.168.3.10
    user ROOT
    port 77
Host home2
    hostname 192.168.3.20
    user root
    port 77
Host home3
    hostname 192.168.3.30
    user ROOT
	port 77
`
)

var (
	configRootDir  = filepath.Join(os.TempDir(), "manssh")
	mainConfigPath = filepath.Join(configRootDir, "config")
	testConfigPath = filepath.Join(configRootDir, "config.d", "test")
	homeConfigPath = filepath.Join(configRootDir, "config.d", "home")
)

func initConfig() {
	os.MkdirAll(configRootDir, os.ModePerm)
	os.MkdirAll(filepath.Join(configRootDir, "config.d"), os.ModePerm)
	if err := ioutil.WriteFile(mainConfigPath, []byte(fmt.Sprintf(mainConfigContent, configRootDir)), 0644); err != nil {
		panic(err)
	}
	if err := ioutil.WriteFile(testConfigPath, []byte(testConfigContent), 0644); err != nil {
		panic(err)
	}
	if err := ioutil.WriteFile(homeConfigPath, []byte(homeConfigContent), 0644); err != nil {
		panic(err)
	}
}

func TestList(t *testing.T) {
	initConfig()
	defer os.Remove(configRootDir)

	hosts, err := List(mainConfigPath, ListOption{})
	require.Nil(t, err)
	require.Equal(t, 10, len(hosts))
	hostMap := map[string]*HostConfig{}
	for _, host := range hosts {
		hostMap[host.Alias] = host
	}

	main2 := hostMap["main2"]
	require.NotNil(t, main2)
	require.Equal(t, 2, len(main2.OwnConfig))
	require.Equal(t, 1, len(main2.ImplicitConfig))
	require.Empty(t, main2.OwnConfig["port"])
	require.Equal(t, "22022", main2.ImplicitConfig["port"])
	require.Equal(t, "192.168.2.20", main2.OwnConfig["hostname"])

	home1 := hostMap["home1"]
	require.NotNil(t, home1)
	require.Equal(t, 3, len(home1.OwnConfig))
	require.Equal(t, 0, len(home1.ImplicitConfig))
	require.Equal(t, "77", home1.OwnConfig["port"])
	require.Equal(t, "ROOT", home1.OwnConfig["user"])
	require.Equal(t, "192.168.3.10", home1.OwnConfig["hostname"])

	hosts, err = List(mainConfigPath, ListOption{
		Keywords: []string{"Test"},
	})
	require.Equal(t, 1, len(hosts))

	hosts, err = List(mainConfigPath, ListOption{
		Keywords:   []string{"Test"},
		IgnoreCase: true,
	})
	require.Equal(t, 3, len(hosts))
}

func TestAdd(t *testing.T) {
	initConfig()
	defer os.Remove(configRootDir)

	_, err := Add(mainConfigPath, &AddOption{
		Path:    testConfigPath,
		Alias:   "test1",
		Connect: "xxx@1.2.3.4:11",
	})
	require.NotNil(t, err)

	host, err := Add(mainConfigPath, &AddOption{
		Path:    testConfigPath,
		Alias:   "test4",
		Connect: "xxx@1.2.3.4",
	})
	require.Nil(t, err)
	require.Equal(t, "22022", host.ImplicitConfig["port"])
	require.Equal(t, "1.2.3.4", host.OwnConfig["hostname"])
	require.Equal(t, "xxx", host.OwnConfig["user"])

	hosts, err := List(mainConfigPath, ListOption{})
	require.Nil(t, err)
	require.Equal(t, 11, len(hosts))
}

func TestUpdate(t *testing.T) {
	initConfig()
	defer os.Remove(configRootDir)

	_, err := Update(mainConfigPath, &UpdateOption{
		Alias:   "test4",
		Connect: "xxx@1.2.3.4:11",
	})
	require.NotNil(t, err)

	host, err := Update(mainConfigPath, &UpdateOption{
		Alias:    "test1",
		NewAlias: "test4",
		Connect:  "xxx@1.2.3.4:11",
		Config: map[string]string{
			"IdentifyFile": "~/.ssh/test4",
		},
	})
	require.Nil(t, err)
	require.Equal(t, "1.2.3.4", host.OwnConfig["hostname"])
	require.Equal(t, "xxx", host.OwnConfig["user"])
	require.Equal(t, "~/.ssh/test4", host.OwnConfig["identifyfile"])
	require.Equal(t, "22022", host.ImplicitConfig["port"])

	host, err = Update(mainConfigPath, &UpdateOption{
		Alias:   "home1",
		Connect: "1.2.3.4:11",
	})
	require.Nil(t, err)
	require.Equal(t, "1.2.3.4", host.OwnConfig["hostname"])
	require.Equal(t, "11", host.OwnConfig["port"])

	hosts, err := List(mainConfigPath, ListOption{})
	require.Nil(t, err)
	require.Equal(t, 10, len(hosts))
	hostMap := map[string]*HostConfig{}
	for _, host := range hosts {
		hostMap[host.Alias] = host
	}
	require.Nil(t, hostMap["test1"])
	require.NotNil(t, hostMap["test4"])
}

func TestDelete(t *testing.T) {
	initConfig()
	defer os.Remove(configRootDir)

	_, err := Delete(mainConfigPath, "home1", "test1", "main4")
	require.NotNil(t, err)

	hosts, err := Delete(mainConfigPath, "home1", "test1", "*")
	require.Nil(t, err)
	require.Equal(t, 3, len(hosts))

	hosts, err = List(mainConfigPath, ListOption{})
	require.Nil(t, err)
	require.Equal(t, 8, len(hosts))
	hostMap := map[string]*HostConfig{}
	for _, host := range hosts {
		hostMap[host.Alias] = host
	}
	require.Nil(t, hostMap["home1"])
	require.Nil(t, hostMap["test1"])
	require.NotNil(t, hostMap["*"])
	require.Equal(t, "22", hostMap["main1"].ImplicitConfig["port"])
}

func TestGetFilePaths(t *testing.T) {
	initConfig()
	defer os.Remove(configRootDir)
	paths, err := GetFilePaths(mainConfigPath)
	require.Nil(t, err)
	require.Equal(t, 3, len(paths))
}
<commit_msg>fix test case<commit_after>package manssh

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/stretchr/testify/require"
)

const (
	mainConfigContent = `
Include %s/config.d/*
Host home1
    hostname 192.168.1.11
Host main1
    hostname 192.168.1.10
Host main2
    hostname 192.168.1.20
    user wen
    port 77
Host main3
    hostname 192.168.1.30
    user ROOT
	port 77
`
	testConfigContent = `
Host *
	port 22022
Host test1
    hostname 192.168.2.10
    user root
    port 22
Host test2 main2
    hostname 192.168.2.20
    port 77
Host Test3
    hostname 192.168.2.30
    user ROOT
	port 77
`
	homeConfigContent = `
Host home1
    hostname 192.168.3.10
    user ROOT
    port 77
Host home2
    hostname 192.168.3.20
    user root
    port 77
Host home3
    hostname 192.168.3.30
    user ROOT
	port 77
`
)

var (
	configRootDir  = filepath.Join(os.TempDir(), "manssh")
	mainConfigPath = filepath.Join(configRootDir, "config")
	testConfigPath = filepath.Join(configRootDir, "config.d", "test")
	homeConfigPath = filepath.Join(configRootDir, "config.d", "home")
)

func initConfig() {
	os.MkdirAll(configRootDir, os.ModePerm)
	os.MkdirAll(filepath.Join(configRootDir, "config.d"), os.ModePerm)
	if err := ioutil.WriteFile(mainConfigPath, []byte(fmt.Sprintf(mainConfigContent, configRootDir)), 0644); err != nil {
		panic(err)
	}
	if err := ioutil.WriteFile(testConfigPath, []byte(testConfigContent), 0644); err != nil {
		panic(err)
	}
	if err := ioutil.WriteFile(homeConfigPath, []byte(homeConfigContent), 0644); err != nil {
		panic(err)
	}
}

func TestList(t *testing.T) {
	initConfig()
	defer os.Remove(configRootDir)

	hosts, err := List(mainConfigPath, ListOption{})
	require.Nil(t, err)
	require.Equal(t, 10, len(hosts))
	hostMap := map[string]*HostConfig{}
	for _, host := range hosts {
		hostMap[host.Alias] = host
	}

	main2 := hostMap["main2"]
	require.NotNil(t, main2)
	require.Equal(t, 2, len(main2.OwnConfig))
	require.Equal(t, 1, len(main2.ImplicitConfig))
	require.Empty(t, main2.OwnConfig["port"])
	require.Equal(t, "22022", main2.ImplicitConfig["port"])
	require.Equal(t, "192.168.2.20", main2.OwnConfig["hostname"])

	home1 := hostMap["home1"]
	require.NotNil(t, home1)
	require.Equal(t, 3, len(home1.OwnConfig))
	require.Equal(t, 0, len(home1.ImplicitConfig))
	require.Equal(t, "77", home1.OwnConfig["port"])
	require.Equal(t, "ROOT", home1.OwnConfig["user"])
	require.Equal(t, "192.168.3.10", home1.OwnConfig["hostname"])

	hosts, err = List(mainConfigPath, ListOption{
		Keywords: []string{"Test"},
	})
	require.Equal(t, 1, len(hosts))

	hosts, err = List(mainConfigPath, ListOption{
		Keywords:   []string{"Test"},
		IgnoreCase: true,
	})
	require.Equal(t, 3, len(hosts))
}

func TestAdd(t *testing.T) {
	initConfig()
	defer os.Remove(configRootDir)

	_, err := Add(mainConfigPath, &AddOption{
		Path:    testConfigPath,
		Alias:   "test1",
		Connect: "xxx@1.2.3.4:11",
		Config:  map[string]string{},
	})
	require.NotNil(t, err)

	host, err := Add(mainConfigPath, &AddOption{
		Path:    testConfigPath,
		Alias:   "test4",
		Connect: "xxx@1.2.3.4",
		Config:  map[string]string{},
	})
	require.Nil(t, err)
	require.Equal(t, "22022", host.ImplicitConfig["port"])
	require.Equal(t, "1.2.3.4", host.OwnConfig["hostname"])
	require.Equal(t, "xxx", host.OwnConfig["user"])

	hosts, err := List(mainConfigPath, ListOption{})
	require.Nil(t, err)
	require.Equal(t, 11, len(hosts))
}

func TestUpdate(t *testing.T) {
	initConfig()
	defer os.Remove(configRootDir)

	_, err := Update(mainConfigPath, &UpdateOption{
		Alias:   "test4",
		Connect: "xxx@1.2.3.4:11",
	})
	require.NotNil(t, err)

	host, err := Update(mainConfigPath, &UpdateOption{
		Alias:    "test1",
		NewAlias: "test4",
		Connect:  "xxx@1.2.3.4:11",
		Config: map[string]string{
			"IdentifyFile": "~/.ssh/test4",
		},
	})
	require.Nil(t, err)
	require.Equal(t, "1.2.3.4", host.OwnConfig["hostname"])
	require.Equal(t, "xxx", host.OwnConfig["user"])
	require.Equal(t, "~/.ssh/test4", host.OwnConfig["identifyfile"])
	require.Equal(t, "22022", host.ImplicitConfig["port"])

	host, err = Update(mainConfigPath, &UpdateOption{
		Alias:   "home1",
		Connect: "1.2.3.4:11",
		Config:  map[string]string{},
	})
	require.Nil(t, err)
	require.Equal(t, "1.2.3.4", host.OwnConfig["hostname"])
	require.Equal(t, "11", host.OwnConfig["port"])

	hosts, err := List(mainConfigPath, ListOption{})
	require.Nil(t, err)
	require.Equal(t, 10, len(hosts))
	hostMap := map[string]*HostConfig{}
	for _, host := range hosts {
		hostMap[host.Alias] = host
	}
	require.Nil(t, hostMap["test1"])
	require.NotNil(t, hostMap["test4"])
}

func TestDelete(t *testing.T) {
	initConfig()
	defer os.Remove(configRootDir)

	_, err := Delete(mainConfigPath, "home1", "test1", "main4")
	require.NotNil(t, err)

	hosts, err := Delete(mainConfigPath, "home1", "test1", "*")
	require.Nil(t, err)
	require.Equal(t, 3, len(hosts))

	hosts, err = List(mainConfigPath, ListOption{})
	require.Nil(t, err)
	require.Equal(t, 8, len(hosts))
	hostMap := map[string]*HostConfig{}
	for _, host := range hosts {
		hostMap[host.Alias] = host
	}
	require.Nil(t, hostMap["home1"])
	require.Nil(t, hostMap["test1"])
	require.NotNil(t, hostMap["*"])
	require.Equal(t, "22", hostMap["main1"].ImplicitConfig["port"])
}

func TestGetFilePaths(t *testing.T) {
	initConfig()
	defer os.Remove(configRootDir)
	paths, err := GetFilePaths(mainConfigPath)
	require.Nil(t, err)
	require.Equal(t, 3, len(paths))
}
<|endoftext|>


<commit_before>package messagebus

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"sync/atomic"
	"time"

	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/route-registrar/config"
	"github.com/nats-io/nats"
)

//go:generate counterfeiter . MessageBus

type MessageBus interface {
	Connect(servers []config.MessageBusServer) error
	SendMessage(subject string, host string, route config.Route, privateInstanceId string) error
	Close()
}

type msgBus struct {
	natsHost *atomic.Value
	natsConn *nats.Conn
	logger   lager.Logger
}

type Message struct {
	URIs              []string          `json:"uris"`
	Host              string            `json:"host"`
	Port              int               `json:"port"`
	Tags              map[string]string `json:"tags"`
	RouteServiceUrl   string            `json:"route_service_url,omitempty"`
	PrivateInstanceId string            `json:"private_instance_id"`
}

func NewMessageBus(logger lager.Logger) MessageBus {
	return &msgBus{
		logger:   logger,
		natsHost: &atomic.Value{},
	}
}

func (m *msgBus) Connect(servers []config.MessageBusServer) error {

	var natsServers []string
	for _, server := range servers {
		natsServers = append(
			natsServers,
			fmt.Sprintf("nats://%s:%s@%s", server.User, server.Password, server.Host),
		)
	}

	opts := nats.DefaultOptions
	opts.Servers = natsServers
	opts.PingInterval = 20 * time.Second

	opts.ClosedCB = func(conn *nats.Conn) {
		m.logger.Error("nats-connection-closed", errors.New("unexpected nats conn closed"), lager.Data{"nats-host": m.natsHost.Load()})
	}

	opts.DisconnectedCB = func(conn *nats.Conn) {
		m.logger.Info("nats-connection-disconnected", lager.Data{"nats-host": m.natsHost.Load()})
	}

	opts.ReconnectedCB = func(conn *nats.Conn) {
		natsHost, err := parseNatsUrl(conn.ConnectedUrl())
		if err != nil {
			m.logger.Error("nats-url-parse-failed", err, lager.Data{"nats-host": natsHost})
		}
		m.natsHost.Store(natsHost)
		m.logger.Info("nats-connection-reconnected", lager.Data{"nats-host": m.natsHost.Load()})
	}

	natsConn, err := opts.Connect()
	if err != nil {
		m.logger.Error("nats-connection-failed", err, lager.Data{"nats-error": natsConn.LastError()})
		return err
	}

	natsHost, err := parseNatsUrl(natsConn.ConnectedUrl())
	if err != nil {
		m.logger.Error("nats-url-parse-failed", err, lager.Data{"nats-host": natsHost})
	}

	m.natsHost.Store(natsHost)
	m.logger.Info("nats-connection-successful", lager.Data{"nats-url": m.natsHost.Load()})
	m.natsConn = natsConn

	return nil
}

func (m msgBus) SendMessage(subject string, host string, route config.Route, privateInstanceId string) error {
	m.logger.Debug("creating-message", lager.Data{"subject": subject, "host": host, "route": route, "privateInstanceId": privateInstanceId})

	msg := &Message{
		URIs:              route.URIs,
		Host:              host,
		Port:              route.Port,
		Tags:              route.Tags,
		RouteServiceUrl:   route.RouteServiceUrl,
		PrivateInstanceId: privateInstanceId,
	}

	json, err := json.Marshal(msg)
	if err != nil {
		// Untested as we cannot force json.Marshal to return error.
		return err
	}

	m.logger.Debug("publishing-message", lager.Data{"msg": string(json)})

	return m.natsConn.Publish(subject, json)
}

func (m msgBus) Close() {
	m.natsConn.Close()
}

func parseNatsUrl(natsUrl string) (string, error) {
	natsURL, err := url.Parse(natsUrl)
	natsHostStr := ""
	if err != nil {
		return "", err
	} else {
		natsHostStr = natsURL.Host
	}

	return natsHostStr, nil
}
<commit_msg>Include nats hosts in failed connection messsage<commit_after>package messagebus

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"sync/atomic"
	"time"

	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/route-registrar/config"
	"github.com/nats-io/nats"
)

//go:generate counterfeiter . MessageBus

type MessageBus interface {
	Connect(servers []config.MessageBusServer) error
	SendMessage(subject string, host string, route config.Route, privateInstanceId string) error
	Close()
}

type msgBus struct {
	natsHost *atomic.Value
	natsConn *nats.Conn
	logger   lager.Logger
}

type Message struct {
	URIs              []string          `json:"uris"`
	Host              string            `json:"host"`
	Port              int               `json:"port"`
	Tags              map[string]string `json:"tags"`
	RouteServiceUrl   string            `json:"route_service_url,omitempty"`
	PrivateInstanceId string            `json:"private_instance_id"`
}

func NewMessageBus(logger lager.Logger) MessageBus {
	return &msgBus{
		logger:   logger,
		natsHost: &atomic.Value{},
	}
}

func (m *msgBus) Connect(servers []config.MessageBusServer) error {

	var natsServers []string
	var natsHosts []string
	for _, server := range servers {
		natsServers = append(
			natsServers,
			fmt.Sprintf("nats://%s:%s@%s", server.User, server.Password, server.Host),
		)
		natsHosts = append(natsHosts, server.Host)
	}

	opts := nats.DefaultOptions
	opts.Servers = natsServers
	opts.PingInterval = 20 * time.Second

	opts.ClosedCB = func(conn *nats.Conn) {
		m.logger.Error("nats-connection-closed", errors.New("unexpected nats conn closed"), lager.Data{"nats-host": m.natsHost.Load()})
	}

	opts.DisconnectedCB = func(conn *nats.Conn) {
		m.logger.Info("nats-connection-disconnected", lager.Data{"nats-host": m.natsHost.Load()})
	}

	opts.ReconnectedCB = func(conn *nats.Conn) {
		natsHost, err := parseNatsUrl(conn.ConnectedUrl())
		if err != nil {
			m.logger.Error("nats-url-parse-failed", err, lager.Data{"nats-host": natsHost})
		}
		m.natsHost.Store(natsHost)
		m.logger.Info("nats-connection-reconnected", lager.Data{"nats-host": m.natsHost.Load()})
	}

	natsConn, err := opts.Connect()
	if err != nil {
		m.logger.Error("nats-connection-failed", err, lager.Data{"nats-hosts": natsHosts})
		return err
	}

	natsHost, err := parseNatsUrl(natsConn.ConnectedUrl())
	if err != nil {
		m.logger.Error("nats-url-parse-failed", err, lager.Data{"nats-host": natsHost})
	}

	m.natsHost.Store(natsHost)
	m.logger.Info("nats-connection-successful", lager.Data{"nats-host": m.natsHost.Load()})
	m.natsConn = natsConn

	return nil
}

func (m msgBus) SendMessage(subject string, host string, route config.Route, privateInstanceId string) error {
	m.logger.Debug("creating-message", lager.Data{"subject": subject, "host": host, "route": route, "privateInstanceId": privateInstanceId})

	msg := &Message{
		URIs:              route.URIs,
		Host:              host,
		Port:              route.Port,
		Tags:              route.Tags,
		RouteServiceUrl:   route.RouteServiceUrl,
		PrivateInstanceId: privateInstanceId,
	}

	json, err := json.Marshal(msg)
	if err != nil {
		// Untested as we cannot force json.Marshal to return error.
		return err
	}

	m.logger.Debug("publishing-message", lager.Data{"msg": string(json)})

	return m.natsConn.Publish(subject, json)
}

func (m msgBus) Close() {
	m.natsConn.Close()
}

func parseNatsUrl(natsUrl string) (string, error) {
	natsURL, err := url.Parse(natsUrl)
	natsHostStr := ""
	if err != nil {
		return "", err
	} else {
		natsHostStr = natsURL.Host
	}

	return natsHostStr, nil
}
<|endoftext|>


<commit_before>package circularbuffer

import (
	"context"
	"math"
	"sync"
	"time"
)

// RateLimiter is an interface which can be used to implement
// rate limiting.
type RateLimiter interface {
	// Allow returns true if call should be allowed, false in case
	// you should rate limit.
	Allow(string) bool

	// AllowContext is like Allow but accepts an additional
	// context.Context, e.g. to support OpenTracing.
	AllowContext(context.Context, string) bool

	// Close cleans up the RateLimiter implementation.
	Close()
	Oldest(string) time.Time
	Delta(string) time.Duration
	Resize(string, int)
	// RetryAfter returns how many seconds until the next allowed request
	RetryAfter(string) int
}

// NewRateLimiter returns a new initialized RateLimitter with maxHits
// as the maximal number of hits per time.Duration d. This can be used
// to implement maximum number of requests for a backend to protect
// from a known scaling limit.
func NewRateLimiter(maxHits int, d time.Duration) RateLimiter {
	return NewCircularBuffer(maxHits, d)
}

// Allow returns true if there is a free bucket and we should not rate
// limit, if not it will return false, which means ratelimit.
func (cb *CircularBuffer) Allow(s string) bool {
	return cb.Add(time.Now())
}

// Allow returns true if there is a free bucket and we should not rate
// limit, if not it will return false, which means ratelimit.
func (cb *CircularBuffer) AllowContext(ctx context.Context, s string) bool {
	return cb.Add(time.Now())
}

// Close implements the RateLimiter interface to shutdown, nothing to
// do.
func (*CircularBuffer) Close() {}

// Oldest implements the RateLimiter interface
func (cb *CircularBuffer) Oldest(string) time.Time {
	cb.RLock()
	cur := cb.slots[cb.offset]
	cb.RUnlock()
	return cur
}

// Current implements the RateLimiter interface
func (cb *CircularBuffer) Current(string) time.Time {
	return cb.current()
}

// Delta returns the diffence between the current and the oldest value in
// the buffer, i.e. maxHits / Delta() => rate
func (cb *CircularBuffer) Delta(string) time.Duration {
	return cb.delta()
}

// Resize resizes the circular buffer to the given size. Resizing to a size
// <= 0 is not performed
func (cb *CircularBuffer) Resize(_ string, n int) {
	cb.Lock()
	cb.resize(n)
	cb.Unlock()
}

// RetryAfter returns how many seconds one should wait until the next request
// is allowed.
func (cb *CircularBuffer) RetryAfter(string) int {
	retryAfter := cb.retryAfter()
	ms := retryAfter / time.Millisecond
	secs := math.Ceil(float64(ms) / 1000)
	return int(secs)
}

// ClientRateLimiter implements the RateLimiter interface and does
// rate limiting based on the the String passed to Allow(). This can
// be used to limit per client calls to the backend. For example you
// can slow down user enumeration or dictionary attacks to /login
// APIs.
type ClientRateLimiter struct {
	sync.RWMutex
	bag        map[string]*CircularBuffer
	maxHits    int
	timeWindow time.Duration
	quitCH     chan struct{}
}

// NewRateLimiter returns a new initialized RateLimitter with maxHits is
// the maximal number of hits per time.Duration d.
func NewClientRateLimiter(maxHits int, d, cleanInterval time.Duration) *ClientRateLimiter {
	quit := make(chan struct{})
	crl := &ClientRateLimiter{
		bag:        make(map[string]*CircularBuffer),
		maxHits:    maxHits,
		timeWindow: d,
		quitCH:     quit,
	}
	go crl.startCleanerDaemon(cleanInterval)
	return crl
}

// Allow tries to add s to a circularbuffer and returns true if we have
// a free bucket, if not it will return false, which means ratelimit.
func (rl *ClientRateLimiter) Allow(s string) bool {
	var source *CircularBuffer
	var present bool

	rl.RLock()
	if source, present = rl.bag[s]; !present {
		rl.RUnlock()
		rl.Lock()
		source = NewCircularBuffer(rl.maxHits, rl.timeWindow)
		rl.bag[s] = source
		rl.Unlock()
	} else {
		rl.RUnlock()
	}
	present = source.Add(time.Now())
	return present
}

// Allow tries to add s to a circularbuffer and returns true if we have
// a free bucket, if not it will return false, which means ratelimit.
func (rl *ClientRateLimiter) AllowContext(ctx context.Context, s string) bool {
	var source *CircularBuffer
	var present bool

	rl.RLock()
	if source, present = rl.bag[s]; !present {
		rl.RUnlock()
		rl.Lock()
		source = NewCircularBuffer(rl.maxHits, rl.timeWindow)
		rl.bag[s] = source
		rl.Unlock()
	} else {
		rl.RUnlock()
	}
	present = source.Add(time.Now())
	return present
}

func (rl *ClientRateLimiter) Oldest(s string) time.Time {
	rl.RLock()
	if _, present := rl.bag[s]; !present {
		rl.RUnlock()
		return time.Time{}
	}
	delta := rl.bag[s].Oldest(s)
	rl.RUnlock()
	return delta
}

func (rl *ClientRateLimiter) Current(s string) time.Time {
	rl.RLock()
	if _, present := rl.bag[s]; !present {
		rl.RUnlock()
		return time.Time{}
	}
	delta := rl.bag[s].Current(s)
	rl.RUnlock()
	return delta
}

// Delta returns the diffence between the current and the oldest value in
// the buffer, i.e. maxHits / Delta() => rate
func (rl *ClientRateLimiter) Delta(s string) time.Duration {
	rl.RLock()
	if _, present := rl.bag[s]; !present {
		rl.RUnlock()
		return time.Duration(time.Hour * 24)
	}
	delta := rl.bag[s].delta()
	rl.RUnlock()
	return delta
}

// Resize resizes the given circular buffer to the given size. Resizing to a size
// <= 0 is not performed
func (rl *ClientRateLimiter) Resize(s string, n int) {
	rl.RLock()
	if _, present := rl.bag[s]; !present {
		rl.RUnlock()
		return
	}
	rl.RUnlock()
	rl.Lock()
	rl.bag[s].resize(n)
	rl.Unlock()
}

// RetryAfter returns how many seconds one should wait until the next request
// is allowed.
func (rl *ClientRateLimiter) RetryAfter(s string) int {
	rl.RLock()
	if _, present := rl.bag[s]; !present {
		rl.RUnlock()
		return 0
	}
	retryAfter := rl.bag[s].RetryAfter(s)
	rl.RUnlock()
	return retryAfter
}

// DeleteOld removes old entries from state bag
func (rl *ClientRateLimiter) DeleteOld() {
	rl.Lock()
	for k, cb := range rl.bag {
		if !cb.InUse() {
			delete(rl.bag, k)
		}
	}
	rl.Unlock()
}

// Close will stop the cleanup goroutine
func (rl *ClientRateLimiter) Close() {
	close(rl.quitCH)
}

func (rl *ClientRateLimiter) startCleanerDaemon(d time.Duration) {
	for {
		select {
		case <-rl.quitCH:
			return
		case <-time.After(d):
			rl.DeleteOld()
		}
	}
}
<commit_msg>Deprecate in godocs<commit_after>package circularbuffer

import (
	"context"
	"math"
	"sync"
	"time"
)

// RateLimiter is an interface which can be used to implement
// rate limiting.
type RateLimiter interface {
	// Allow returns true if call should be allowed, false in case
	// you should rate limit.
	//
	// Deprecated: In favour of AllowContext
	Allow(string) bool

	// AllowContext is like Allow but accepts an additional
	// context.Context, e.g. to support OpenTracing.
	AllowContext(context.Context, string) bool

	// Close cleans up the RateLimiter implementation.
	Close()
	Oldest(string) time.Time
	Delta(string) time.Duration
	Resize(string, int)
	// RetryAfter returns how many seconds until the next allowed request
	RetryAfter(string) int
}

// NewRateLimiter returns a new initialized RateLimitter with maxHits
// as the maximal number of hits per time.Duration d. This can be used
// to implement maximum number of requests for a backend to protect
// from a known scaling limit.
func NewRateLimiter(maxHits int, d time.Duration) RateLimiter {
	return NewCircularBuffer(maxHits, d)
}

// Allow returns true if there is a free bucket and we should not rate
// limit, if not it will return false, which means ratelimit.
//
// Deprecated: In favour of AllowContext
func (cb *CircularBuffer) Allow(s string) bool {
	return cb.Add(time.Now())
}

// Allow returns true if there is a free bucket and we should not rate
// limit, if not it will return false, which means ratelimit.
func (cb *CircularBuffer) AllowContext(ctx context.Context, s string) bool {
	return cb.Add(time.Now())
}

// Close implements the RateLimiter interface to shutdown, nothing to
// do.
func (*CircularBuffer) Close() {}

// Oldest implements the RateLimiter interface
func (cb *CircularBuffer) Oldest(string) time.Time {
	cb.RLock()
	cur := cb.slots[cb.offset]
	cb.RUnlock()
	return cur
}

// Current implements the RateLimiter interface
func (cb *CircularBuffer) Current(string) time.Time {
	return cb.current()
}

// Delta returns the diffence between the current and the oldest value in
// the buffer, i.e. maxHits / Delta() => rate
func (cb *CircularBuffer) Delta(string) time.Duration {
	return cb.delta()
}

// Resize resizes the circular buffer to the given size. Resizing to a size
// <= 0 is not performed
func (cb *CircularBuffer) Resize(_ string, n int) {
	cb.Lock()
	cb.resize(n)
	cb.Unlock()
}

// RetryAfter returns how many seconds one should wait until the next request
// is allowed.
func (cb *CircularBuffer) RetryAfter(string) int {
	retryAfter := cb.retryAfter()
	ms := retryAfter / time.Millisecond
	secs := math.Ceil(float64(ms) / 1000)
	return int(secs)
}

// ClientRateLimiter implements the RateLimiter interface and does
// rate limiting based on the the String passed to Allow(). This can
// be used to limit per client calls to the backend. For example you
// can slow down user enumeration or dictionary attacks to /login
// APIs.
type ClientRateLimiter struct {
	sync.RWMutex
	bag        map[string]*CircularBuffer
	maxHits    int
	timeWindow time.Duration
	quitCH     chan struct{}
}

// NewRateLimiter returns a new initialized RateLimitter with maxHits is
// the maximal number of hits per time.Duration d.
func NewClientRateLimiter(maxHits int, d, cleanInterval time.Duration) *ClientRateLimiter {
	quit := make(chan struct{})
	crl := &ClientRateLimiter{
		bag:        make(map[string]*CircularBuffer),
		maxHits:    maxHits,
		timeWindow: d,
		quitCH:     quit,
	}
	go crl.startCleanerDaemon(cleanInterval)
	return crl
}

// Allow tries to add s to a circularbuffer and returns true if we have
// a free bucket, if not it will return false, which means ratelimit.
//
// Deprecated: In favour of allow context
func (rl *ClientRateLimiter) Allow(s string) bool {
	var source *CircularBuffer
	var present bool

	rl.RLock()
	if source, present = rl.bag[s]; !present {
		rl.RUnlock()
		rl.Lock()
		source = NewCircularBuffer(rl.maxHits, rl.timeWindow)
		rl.bag[s] = source
		rl.Unlock()
	} else {
		rl.RUnlock()
	}
	present = source.Add(time.Now())
	return present
}

// AllowContext tries to add s to a circularbuffer and returns true if we have
// a free bucket, if not it will return false, which means ratelimit with an additional
// context.Context.
func (rl *ClientRateLimiter) AllowContext(ctx context.Context, s string) bool {
	var source *CircularBuffer
	var present bool

	rl.RLock()
	if source, present = rl.bag[s]; !present {
		rl.RUnlock()
		rl.Lock()
		source = NewCircularBuffer(rl.maxHits, rl.timeWindow)
		rl.bag[s] = source
		rl.Unlock()
	} else {
		rl.RUnlock()
	}
	present = source.Add(time.Now())
	return present
}

func (rl *ClientRateLimiter) Oldest(s string) time.Time {
	rl.RLock()
	if _, present := rl.bag[s]; !present {
		rl.RUnlock()
		return time.Time{}
	}
	delta := rl.bag[s].Oldest(s)
	rl.RUnlock()
	return delta
}

func (rl *ClientRateLimiter) Current(s string) time.Time {
	rl.RLock()
	if _, present := rl.bag[s]; !present {
		rl.RUnlock()
		return time.Time{}
	}
	delta := rl.bag[s].Current(s)
	rl.RUnlock()
	return delta
}

// Delta returns the diffence between the current and the oldest value in
// the buffer, i.e. maxHits / Delta() => rate
func (rl *ClientRateLimiter) Delta(s string) time.Duration {
	rl.RLock()
	if _, present := rl.bag[s]; !present {
		rl.RUnlock()
		return time.Duration(time.Hour * 24)
	}
	delta := rl.bag[s].delta()
	rl.RUnlock()
	return delta
}

// Resize resizes the given circular buffer to the given size. Resizing to a size
// <= 0 is not performed
func (rl *ClientRateLimiter) Resize(s string, n int) {
	rl.RLock()
	if _, present := rl.bag[s]; !present {
		rl.RUnlock()
		return
	}
	rl.RUnlock()
	rl.Lock()
	rl.bag[s].resize(n)
	rl.Unlock()
}

// RetryAfter returns how many seconds one should wait until the next request
// is allowed.
func (rl *ClientRateLimiter) RetryAfter(s string) int {
	rl.RLock()
	if _, present := rl.bag[s]; !present {
		rl.RUnlock()
		return 0
	}
	retryAfter := rl.bag[s].RetryAfter(s)
	rl.RUnlock()
	return retryAfter
}

// DeleteOld removes old entries from state bag
func (rl *ClientRateLimiter) DeleteOld() {
	rl.Lock()
	for k, cb := range rl.bag {
		if !cb.InUse() {
			delete(rl.bag, k)
		}
	}
	rl.Unlock()
}

// Close will stop the cleanup goroutine
func (rl *ClientRateLimiter) Close() {
	close(rl.quitCH)
}

func (rl *ClientRateLimiter) startCleanerDaemon(d time.Duration) {
	for {
		select {
		case <-rl.quitCH:
			return
		case <-time.After(d):
			rl.DeleteOld()
		}
	}
}
<|endoftext|>


<commit_before>package provider

import (
	"reflect"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/cenk/backoff"
	"github.com/containous/traefik/job"
	"github.com/containous/traefik/log"
	"github.com/containous/traefik/provider/k8s"
	"github.com/containous/traefik/safe"
	"github.com/containous/traefik/types"
	"k8s.io/client-go/1.5/pkg/api/v1"
	"k8s.io/client-go/1.5/pkg/util/intstr"
)

var _ Provider = (*Kubernetes)(nil)

// Kubernetes holds configurations of the Kubernetes provider.
type Kubernetes struct {
	BaseProvider           `mapstructure:",squash"`
	Endpoint               string         `description:"Kubernetes server endpoint"`
	DisablePassHostHeaders bool           `description:"Kubernetes disable PassHost Headers"`
	Namespaces             k8s.Namespaces `description:"Kubernetes namespaces"`
	LabelSelector          string         `description:"Kubernetes api label selector to use"`
	lastConfiguration      safe.Safe
}

func (provider *Kubernetes) newK8sClient() (k8s.Client, error) {
	if provider.Endpoint != "" {
		log.Infof("Creating in cluster Kubernetes client with endpoint %v", provider.Endpoint)
		return k8s.NewInClusterClientWithEndpoint(provider.Endpoint)
	}
	log.Info("Creating in cluster Kubernetes client")
	return k8s.NewInClusterClient()
}

// Provide allows the provider to provide configurations to traefik
// using the given configuration channel.
func (provider *Kubernetes) Provide(configurationChan chan<- types.ConfigMessage, pool *safe.Pool, constraints types.Constraints) error {
	k8sClient, err := provider.newK8sClient()
	if err != nil {
		return err
	}
	provider.Constraints = append(provider.Constraints, constraints...)

	pool.Go(func(stop chan bool) {
		operation := func() error {
			for {
				stopWatch := make(chan struct{}, 1)
				defer close(stopWatch)
				log.Debugf("Using label selector: '%s'", provider.LabelSelector)
				eventsChan, err := k8sClient.WatchAll(provider.LabelSelector, stopWatch)
				if err != nil {
					log.Errorf("Error watching kubernetes events: %v", err)
					timer := time.NewTimer(1 * time.Second)
					select {
					case <-timer.C:
						return err
					case <-stop:
						return nil
					}
				}
				for {
					select {
					case <-stop:
						return nil
					case event := <-eventsChan:
						log.Debugf("Received event from kubernetes %+v", event)
						templateObjects, err := provider.loadIngresses(k8sClient)
						if err != nil {
							return err
						}
						if reflect.DeepEqual(provider.lastConfiguration.Get(), templateObjects) {
							log.Debugf("Skipping event from kubernetes %+v", event)
						} else {
							provider.lastConfiguration.Set(templateObjects)
							configurationChan <- types.ConfigMessage{
								ProviderName:  "kubernetes",
								Configuration: provider.loadConfig(*templateObjects),
							}
						}
					}
				}
			}
		}

		notify := func(err error, time time.Duration) {
			log.Errorf("Kubernetes connection error %+v, retrying in %s", err, time)
		}
		err := backoff.RetryNotify(safe.OperationWithRecover(operation), job.NewBackOff(backoff.NewExponentialBackOff()), notify)
		if err != nil {
			log.Errorf("Cannot connect to Kubernetes server %+v", err)
		}
	})

	return nil
}

func (provider *Kubernetes) loadIngresses(k8sClient k8s.Client) (*types.Configuration, error) {
	ingresses := k8sClient.GetIngresses(provider.Namespaces)

	templateObjects := types.Configuration{
		map[string]*types.Backend{},
		map[string]*types.Frontend{},
	}
	PassHostHeader := provider.getPassHostHeader()
	for _, i := range ingresses {
		for _, r := range i.Spec.Rules {
			if r.HTTP == nil {
				log.Warnf("Error in ingress: HTTP is nil")
				continue
			}
			for _, pa := range r.HTTP.Paths {
				if _, exists := templateObjects.Backends[r.Host+pa.Path]; !exists {
					templateObjects.Backends[r.Host+pa.Path] = &types.Backend{
						Servers: make(map[string]types.Server),
						LoadBalancer: &types.LoadBalancer{
							Sticky: false,
							Method: "wrr",
						},
					}
				}
				if _, exists := templateObjects.Frontends[r.Host+pa.Path]; !exists {
					templateObjects.Frontends[r.Host+pa.Path] = &types.Frontend{
						Backend:        r.Host + pa.Path,
						PassHostHeader: PassHostHeader,
						Routes:         make(map[string]types.Route),
						Priority:       len(pa.Path),
					}
				}
				if len(r.Host) > 0 {
					rule := "Host:" + r.Host

					if strings.Contains(r.Host, "*") {
						rule = "HostRegexp:" + strings.Replace(r.Host, "*", "{subdomain:[A-Za-z0-9-_]+}", 1)
					}

					if _, exists := templateObjects.Frontends[r.Host+pa.Path].Routes[r.Host]; !exists {
						templateObjects.Frontends[r.Host+pa.Path].Routes[r.Host] = types.Route{
							Rule: rule,
						}
					}
				}
				if len(pa.Path) > 0 {
					ruleType := i.Annotations["traefik.frontend.rule.type"]

					switch strings.ToLower(ruleType) {
					case "pathprefixstrip":
						ruleType = "PathPrefixStrip"
					case "pathstrip":
						ruleType = "PathStrip"
					case "path":
						ruleType = "Path"
					case "pathprefix":
						ruleType = "PathPrefix"
					case "":
						ruleType = "PathPrefix"
					default:
						log.Warnf("Unknown RuleType %s for %s/%s, falling back to PathPrefix", ruleType, i.ObjectMeta.Namespace, i.ObjectMeta.Name)
						ruleType = "PathPrefix"
					}

					templateObjects.Frontends[r.Host+pa.Path].Routes[pa.Path] = types.Route{
						Rule: ruleType + ":" + pa.Path,
					}
				}
				service, exists, err := k8sClient.GetService(i.ObjectMeta.Namespace, pa.Backend.ServiceName)
				if err != nil || !exists {
					log.Warnf("Error retrieving service %s/%s: %v", i.ObjectMeta.Namespace, pa.Backend.ServiceName, err)
					delete(templateObjects.Frontends, r.Host+pa.Path)
					continue
				}

				if expression := service.Annotations["traefik.backend.circuitbreaker"]; expression != "" {
					templateObjects.Backends[r.Host+pa.Path].CircuitBreaker = &types.CircuitBreaker{
						Expression: expression,
					}
				}
				if service.Annotations["traefik.backend.loadbalancer.method"] == "drr" {
					templateObjects.Backends[r.Host+pa.Path].LoadBalancer.Method = "drr"
				}
				if service.Annotations["traefik.backend.loadbalancer.sticky"] == "true" {
					templateObjects.Backends[r.Host+pa.Path].LoadBalancer.Sticky = true
				}

				protocol := "http"
				for _, port := range service.Spec.Ports {
					if equalPorts(port, pa.Backend.ServicePort) {
						if port.Port == 443 {
							protocol = "https"
						}
						if service.Spec.Type == "ExternalName" {
							url := protocol + "://" + service.Spec.ExternalName
							name := url

							templateObjects.Backends[r.Host+pa.Path].Servers[name] = types.Server{
								URL:    url,
								Weight: 1,
							}
						} else {
							endpoints, exists, err := k8sClient.GetEndpoints(service.ObjectMeta.Namespace, service.ObjectMeta.Name)
							if err != nil || !exists {
								log.Errorf("Error retrieving endpoints %s/%s: %v", service.ObjectMeta.Namespace, service.ObjectMeta.Name, err)
								continue
							}
							if len(endpoints.Subsets) == 0 {
								log.Warnf("Endpoints not found for %s/%s, falling back to Service ClusterIP", service.ObjectMeta.Namespace, service.ObjectMeta.Name)
								templateObjects.Backends[r.Host+pa.Path].Servers[string(service.UID)] = types.Server{
									URL:    protocol + "://" + service.Spec.ClusterIP + ":" + strconv.Itoa(int(port.Port)),
									Weight: 1,
								}
							} else {
								for _, subset := range endpoints.Subsets {
									for _, address := range subset.Addresses {
										url := protocol + "://" + address.IP + ":" + strconv.Itoa(endpointPortNumber(port, subset.Ports))
										name := url
										if address.TargetRef != nil && address.TargetRef.Name != "" {
											name = address.TargetRef.Name
										}
										templateObjects.Backends[r.Host+pa.Path].Servers[name] = types.Server{
											URL:    url,
											Weight: 1,
										}
									}
								}
							}
						}
						break
					}
				}
			}
		}
	}
	return &templateObjects, nil
}

func endpointPortNumber(servicePort v1.ServicePort, endpointPorts []v1.EndpointPort) int {
	if len(endpointPorts) > 0 {
		//name is optional if there is only one port
		port := endpointPorts[0]
		for _, endpointPort := range endpointPorts {
			if servicePort.Name == endpointPort.Name {
				port = endpointPort
			}
		}
		return int(port.Port)
	}
	return int(servicePort.Port)
}

func equalPorts(servicePort v1.ServicePort, ingressPort intstr.IntOrString) bool {
	if int(servicePort.Port) == ingressPort.IntValue() {
		return true
	}
	if servicePort.Name != "" && servicePort.Name == ingressPort.String() {
		return true
	}
	return false
}

func (provider *Kubernetes) getPassHostHeader() bool {
	if provider.DisablePassHostHeaders {
		return false
	}
	return true
}

func (provider *Kubernetes) loadConfig(templateObjects types.Configuration) *types.Configuration {
	var FuncMap = template.FuncMap{}
	configuration, err := provider.getConfiguration("templates/kubernetes.tmpl", FuncMap, templateObjects)
	if err != nil {
		log.Error(err)
	}
	return configuration
}
<commit_msg>Added support for passHostHeader annotation on ingresses<commit_after>package provider

import (
	"reflect"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/cenk/backoff"
	"github.com/containous/traefik/job"
	"github.com/containous/traefik/log"
	"github.com/containous/traefik/provider/k8s"
	"github.com/containous/traefik/safe"
	"github.com/containous/traefik/types"
	"k8s.io/client-go/1.5/pkg/api/v1"
	"k8s.io/client-go/1.5/pkg/util/intstr"
)

var _ Provider = (*Kubernetes)(nil)

// Kubernetes holds configurations of the Kubernetes provider.
type Kubernetes struct {
	BaseProvider           `mapstructure:",squash"`
	Endpoint               string         `description:"Kubernetes server endpoint"`
	DisablePassHostHeaders bool           `description:"Kubernetes disable PassHost Headers"`
	Namespaces             k8s.Namespaces `description:"Kubernetes namespaces"`
	LabelSelector          string         `description:"Kubernetes api label selector to use"`
	lastConfiguration      safe.Safe
}

func (provider *Kubernetes) newK8sClient() (k8s.Client, error) {
	if provider.Endpoint != "" {
		log.Infof("Creating in cluster Kubernetes client with endpoint %v", provider.Endpoint)
		return k8s.NewInClusterClientWithEndpoint(provider.Endpoint)
	}
	log.Info("Creating in cluster Kubernetes client")
	return k8s.NewInClusterClient()
}

// Provide allows the provider to provide configurations to traefik
// using the given configuration channel.
func (provider *Kubernetes) Provide(configurationChan chan<- types.ConfigMessage, pool *safe.Pool, constraints types.Constraints) error {
	k8sClient, err := provider.newK8sClient()
	if err != nil {
		return err
	}
	provider.Constraints = append(provider.Constraints, constraints...)

	pool.Go(func(stop chan bool) {
		operation := func() error {
			for {
				stopWatch := make(chan struct{}, 1)
				defer close(stopWatch)
				log.Debugf("Using label selector: '%s'", provider.LabelSelector)
				eventsChan, err := k8sClient.WatchAll(provider.LabelSelector, stopWatch)
				if err != nil {
					log.Errorf("Error watching kubernetes events: %v", err)
					timer := time.NewTimer(1 * time.Second)
					select {
					case <-timer.C:
						return err
					case <-stop:
						return nil
					}
				}
				for {
					select {
					case <-stop:
						return nil
					case event := <-eventsChan:
						log.Debugf("Received event from kubernetes %+v", event)
						templateObjects, err := provider.loadIngresses(k8sClient)
						if err != nil {
							return err
						}
						if reflect.DeepEqual(provider.lastConfiguration.Get(), templateObjects) {
							log.Debugf("Skipping event from kubernetes %+v", event)
						} else {
							provider.lastConfiguration.Set(templateObjects)
							configurationChan <- types.ConfigMessage{
								ProviderName:  "kubernetes",
								Configuration: provider.loadConfig(*templateObjects),
							}
						}
					}
				}
			}
		}

		notify := func(err error, time time.Duration) {
			log.Errorf("Kubernetes connection error %+v, retrying in %s", err, time)
		}
		err := backoff.RetryNotify(safe.OperationWithRecover(operation), job.NewBackOff(backoff.NewExponentialBackOff()), notify)
		if err != nil {
			log.Errorf("Cannot connect to Kubernetes server %+v", err)
		}
	})

	return nil
}

func (provider *Kubernetes) loadIngresses(k8sClient k8s.Client) (*types.Configuration, error) {
	ingresses := k8sClient.GetIngresses(provider.Namespaces)

	templateObjects := types.Configuration{
		map[string]*types.Backend{},
		map[string]*types.Frontend{},
	}
	for _, i := range ingresses {
		for _, r := range i.Spec.Rules {
			if r.HTTP == nil {
				log.Warnf("Error in ingress: HTTP is nil")
				continue
			}
			for _, pa := range r.HTTP.Paths {
				if _, exists := templateObjects.Backends[r.Host+pa.Path]; !exists {
					templateObjects.Backends[r.Host+pa.Path] = &types.Backend{
						Servers: make(map[string]types.Server),
						LoadBalancer: &types.LoadBalancer{
							Sticky: false,
							Method: "wrr",
						},
					}
				}

				PassHostHeader := provider.getPassHostHeader()

				passHostHeaderAnnotation := i.Annotations["traefik.frontend.passHostHeader"]
				switch passHostHeaderAnnotation {
				case "true":
					PassHostHeader = true
				case "false":
					PassHostHeader = false

				}

				if _, exists := templateObjects.Frontends[r.Host+pa.Path]; !exists {
					templateObjects.Frontends[r.Host+pa.Path] = &types.Frontend{
						Backend:        r.Host + pa.Path,
						PassHostHeader: PassHostHeader,
						Routes:         make(map[string]types.Route),
						Priority:       len(pa.Path),
					}
				}
				if len(r.Host) > 0 {
					rule := "Host:" + r.Host

					if strings.Contains(r.Host, "*") {
						rule = "HostRegexp:" + strings.Replace(r.Host, "*", "{subdomain:[A-Za-z0-9-_]+}", 1)
					}

					if _, exists := templateObjects.Frontends[r.Host+pa.Path].Routes[r.Host]; !exists {
						templateObjects.Frontends[r.Host+pa.Path].Routes[r.Host] = types.Route{
							Rule: rule,
						}
					}
				}
				if len(pa.Path) > 0 {
					ruleType := i.Annotations["traefik.frontend.rule.type"]

					switch strings.ToLower(ruleType) {
					case "pathprefixstrip":
						ruleType = "PathPrefixStrip"
					case "pathstrip":
						ruleType = "PathStrip"
					case "path":
						ruleType = "Path"
					case "pathprefix":
						ruleType = "PathPrefix"
					case "":
						ruleType = "PathPrefix"
					default:
						log.Warnf("Unknown RuleType %s for %s/%s, falling back to PathPrefix", ruleType, i.ObjectMeta.Namespace, i.ObjectMeta.Name)
						ruleType = "PathPrefix"
					}

					templateObjects.Frontends[r.Host+pa.Path].Routes[pa.Path] = types.Route{
						Rule: ruleType + ":" + pa.Path,
					}
				}
				service, exists, err := k8sClient.GetService(i.ObjectMeta.Namespace, pa.Backend.ServiceName)
				if err != nil || !exists {
					log.Warnf("Error retrieving service %s/%s: %v", i.ObjectMeta.Namespace, pa.Backend.ServiceName, err)
					delete(templateObjects.Frontends, r.Host+pa.Path)
					continue
				}

				if expression := service.Annotations["traefik.backend.circuitbreaker"]; expression != "" {
					templateObjects.Backends[r.Host+pa.Path].CircuitBreaker = &types.CircuitBreaker{
						Expression: expression,
					}
				}
				if service.Annotations["traefik.backend.loadbalancer.method"] == "drr" {
					templateObjects.Backends[r.Host+pa.Path].LoadBalancer.Method = "drr"
				}
				if service.Annotations["traefik.backend.loadbalancer.sticky"] == "true" {
					templateObjects.Backends[r.Host+pa.Path].LoadBalancer.Sticky = true
				}

				protocol := "http"
				for _, port := range service.Spec.Ports {
					if equalPorts(port, pa.Backend.ServicePort) {
						if port.Port == 443 {
							protocol = "https"
						}
						if service.Spec.Type == "ExternalName" {
							url := protocol + "://" + service.Spec.ExternalName
							name := url

							templateObjects.Backends[r.Host+pa.Path].Servers[name] = types.Server{
								URL:    url,
								Weight: 1,
							}
						} else {
							endpoints, exists, err := k8sClient.GetEndpoints(service.ObjectMeta.Namespace, service.ObjectMeta.Name)
							if err != nil || !exists {
								log.Errorf("Error retrieving endpoints %s/%s: %v", service.ObjectMeta.Namespace, service.ObjectMeta.Name, err)
								continue
							}
							if len(endpoints.Subsets) == 0 {
								log.Warnf("Endpoints not found for %s/%s, falling back to Service ClusterIP", service.ObjectMeta.Namespace, service.ObjectMeta.Name)
								templateObjects.Backends[r.Host+pa.Path].Servers[string(service.UID)] = types.Server{
									URL:    protocol + "://" + service.Spec.ClusterIP + ":" + strconv.Itoa(int(port.Port)),
									Weight: 1,
								}
							} else {
								for _, subset := range endpoints.Subsets {
									for _, address := range subset.Addresses {
										url := protocol + "://" + address.IP + ":" + strconv.Itoa(endpointPortNumber(port, subset.Ports))
										name := url
										if address.TargetRef != nil && address.TargetRef.Name != "" {
											name = address.TargetRef.Name
										}
										templateObjects.Backends[r.Host+pa.Path].Servers[name] = types.Server{
											URL:    url,
											Weight: 1,
										}
									}
								}
							}
						}
						break
					}
				}
			}
		}
	}
	return &templateObjects, nil
}

func endpointPortNumber(servicePort v1.ServicePort, endpointPorts []v1.EndpointPort) int {
	if len(endpointPorts) > 0 {
		//name is optional if there is only one port
		port := endpointPorts[0]
		for _, endpointPort := range endpointPorts {
			if servicePort.Name == endpointPort.Name {
				port = endpointPort
			}
		}
		return int(port.Port)
	}
	return int(servicePort.Port)
}

func equalPorts(servicePort v1.ServicePort, ingressPort intstr.IntOrString) bool {
	if int(servicePort.Port) == ingressPort.IntValue() {
		return true
	}
	if servicePort.Name != "" && servicePort.Name == ingressPort.String() {
		return true
	}
	return false
}

func (provider *Kubernetes) getPassHostHeader() bool {
	if provider.DisablePassHostHeaders {
		return false
	}
	return true
}

func (provider *Kubernetes) loadConfig(templateObjects types.Configuration) *types.Configuration {
	var FuncMap = template.FuncMap{}
	configuration, err := provider.getConfiguration("templates/kubernetes.tmpl", FuncMap, templateObjects)
	if err != nil {
		log.Error(err)
	}
	return configuration
}
<|endoftext|>


<commit_before>/*
Package readcache implements a read-through cache where the caller supplies
the means to read-through in the form of a function.

TODO: Add a maximum size, enforce with an LRU algorithm
*/
package readcache

import (
	"sync"
	"time"
)

// Type Cacheable is something that may be stored in a cache
type Cacheable struct {
	// The item in the cache
	Value interface{}

	// The time at which this item should expire from the cache.
	ExpiresAt time.Time
}

// Type Cache defines a read-through cache.
type Cache interface {
	// Retrieve an item from the cache if available, or from a
	// backing source if it is not.
	Get(key string) interface{}

	// Forcibly expire a particular item from the cache, if it exists.
	Expire(key string)
}

// Constructs a new cache
func New(getter func(string) Cacheable) Cache {
	return &lazycache{getter, make(map[string]Cacheable), make(map[string]*sync.Once), new(sync.RWMutex), new(sync.RWMutex)}
}

// Type lazycache implements the Cache interface
type lazycache struct {
	Getter           func(string) Cacheable
	Cache            map[string]Cacheable
	ReadControls     map[string]*sync.Once
	CacheLock        *sync.RWMutex
	ReadControlsLock *sync.RWMutex
}

// Get an item from the cache, retrieving the item from the getter if necessary.
// This implemention is meant to be goroutine safe.  It assumes that updating a
// map while concurrently reading from it is unsafe, so it uses a read/write mutex
// to synchronize access to its internal maps.
func (c *lazycache) Get(key string) interface{} {
	c.CacheLock.RLock()
	cachedValue, ok := c.Cache[key]
	c.CacheLock.RUnlock()
	if ok {
		now := time.Now()
		if cachedValue.ExpiresAt.After(now) {
			return cachedValue.Value
		}
		c.CacheLock.Lock()
		// Determine if another goroutine has updated the cache before the lock
		cachedValue, ok = c.Cache[key]
		if ok && cachedValue.ExpiresAt.After(now) {
			return cachedValue
		}
		delete(c.Cache, key)
		c.CacheLock.Unlock()
	}

	c.ReadControlsLock.RLock()
	readControl, ok := c.ReadControls[key]
	c.ReadControlsLock.RUnlock()
	if !ok {
		c.ReadControlsLock.Lock()

		// Another goroutine may have created a read control, fetched an item, updated the
		// cache and cleaned up its read control by the time we reach this point.
		// Therefore, we verify that the cache still does not contain anything for the
		// given key.
		// Warning: possibility of deadlock when dealing with multiple locks.  Make sure
		//          they are always acquired in the same order.
		c.CacheLock.RLock()
		cachedValue, ok := c.Cache[key]
		c.CacheLock.RUnlock()

		if ok {
			c.ReadControlsLock.Unlock()
			return cachedValue.Value
		}

		readControl, ok = c.ReadControls[key]
		if !ok {
			readControl = new(sync.Once)
			c.ReadControls[key] = readControl
		}
		c.ReadControlsLock.Unlock()
	}

	fetchedValue := false
	readControl.Do(func() {
		defer func() {
			c.ReadControlsLock.Lock()
			delete(c.ReadControls, key)
			c.ReadControlsLock.Unlock()
		}()

		cachedValue = c.Getter(key)
		fetchedValue = true
		c.CacheLock.Lock()
		c.Cache[key] = cachedValue
		c.CacheLock.Unlock()
	})

	if !fetchedValue {
		c.CacheLock.RLock()
		cachedValue, _ = c.Cache[key]
		c.CacheLock.RUnlock()
	}

	return cachedValue.Value
}

func getFromCache(c *lazycache, key string) (value Cacheable, ok bool) {
	c.CacheLock.RLock()
	defer func() {
		c.CacheLock.RUnlock()
	}()

	value, ok = c.Cache[key]
	return
}

func (c *lazycache) Expire(key string) {
	c.CacheLock.Lock()
	delete(c.Cache, key)
	c.CacheLock.Unlock()
}
<commit_msg>Removed some unused code.  Also removed the Expire method; not sure if it is a good idea to keep.<commit_after>/*
Package readcache implements a read-through cache where the caller supplies
the means to read-through in the form of a function.

TODO: Add a maximum size, enforce with an LRU algorithm
*/
package readcache

import (
	"sync"
	"time"
)

// Type Cacheable is something that may be stored in a cache
type Cacheable struct {
	// The item in the cache
	Value interface{}

	// The time at which this item should expire from the cache.
	ExpiresAt time.Time
}

// Type Cache defines a read-through cache.
type Cache interface {
	// Retrieve an item from the cache if available, or from a
	// backing source if it is not.
	Get(key string) interface{}
}

// Constructs a new cache
func New(getter func(string) Cacheable) Cache {
	return &lazycache{getter, make(map[string]Cacheable), make(map[string]*sync.Once), new(sync.RWMutex), new(sync.RWMutex)}
}

// Type lazycache implements the Cache interface
type lazycache struct {
	Getter           func(string) Cacheable
	Cache            map[string]Cacheable
	ReadControls     map[string]*sync.Once
	CacheLock        *sync.RWMutex
	ReadControlsLock *sync.RWMutex
}

// Get an item from the cache, retrieving the item from the getter if necessary.
// This implemention is meant to be goroutine safe.  It assumes that updating a
// map while concurrently reading from it is unsafe, so it uses a read/write mutex
// to synchronize access to its internal maps.
func (c *lazycache) Get(key string) interface{} {
	c.CacheLock.RLock()
	cachedValue, ok := c.Cache[key]
	c.CacheLock.RUnlock()
	if ok {
		now := time.Now()
		if cachedValue.ExpiresAt.After(now) {
			return cachedValue.Value
		}
		c.CacheLock.Lock()
		// Determine if another goroutine has updated the cache before the lock
		cachedValue, ok = c.Cache[key]
		if ok && cachedValue.ExpiresAt.After(now) {
			return cachedValue
		}
		delete(c.Cache, key)
		c.CacheLock.Unlock()
	}

	c.ReadControlsLock.RLock()
	readControl, ok := c.ReadControls[key]
	c.ReadControlsLock.RUnlock()
	if !ok {
		c.ReadControlsLock.Lock()

		// Another goroutine may have created a read control, fetched an item, updated the
		// cache and cleaned up its read control by the time we reach this point.
		// Therefore, we verify that the cache still does not contain anything for the
		// given key.
		// Warning: possibility of deadlock when dealing with multiple locks.  Make sure
		//          they are always acquired in the same order.
		c.CacheLock.RLock()
		cachedValue, ok := c.Cache[key]
		c.CacheLock.RUnlock()

		if ok {
			c.ReadControlsLock.Unlock()
			return cachedValue.Value
		}

		readControl, ok = c.ReadControls[key]
		if !ok {
			readControl = new(sync.Once)
			c.ReadControls[key] = readControl
		}
		c.ReadControlsLock.Unlock()
	}

	fetchedValue := false
	readControl.Do(func() {
		defer func() {
			c.ReadControlsLock.Lock()
			delete(c.ReadControls, key)
			c.ReadControlsLock.Unlock()
		}()

		cachedValue = c.Getter(key)
		fetchedValue = true
		c.CacheLock.Lock()
		c.Cache[key] = cachedValue
		c.CacheLock.Unlock()
	})

	if !fetchedValue {
		c.CacheLock.RLock()
		cachedValue, _ = c.Cache[key]
		c.CacheLock.RUnlock()
	}

	return cachedValue.Value
}

<|endoftext|>


<commit_before>// Copyright (c) OpenFaaS Author(s) 2018. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

package proxy

import (
	"fmt"
	"net/http"
	"regexp"
	"testing"

	"github.com/openfaas/faas-cli/test"
	"github.com/openfaas/faas/gateway/requests"
)

func Test_GetFunctionInfo(t *testing.T) {
	s := test.MockHttpServer(t, []test.Request{
		{
			ResponseStatusCode: http.StatusOK,
			ResponseBody:       expectedGetFunctionInfoResponse,
		},
	})

	defer s.Close()
	result, err := GetFunctionInfo(s.URL, "func-test1", !tlsNoVerify)
	if err != nil {
		t.Fatalf("Error returned: %s", err)
	}
	if expectedGetFunctionInfoResponse != result {
		t.Fatalf("Want: %#v, Got: %#v", expectedGetFunctionInfoResponse, result)
	}
}

func Test_GetFunctionInfo_Not200(t *testing.T) {
	s := test.MockHttpServerStatus(t, http.StatusBadRequest)

	_, err := GetFunctionInfo(s.URL, "func-test1", tlsNoVerify)

	if err == nil {
		t.Fatalf("Error was not returned")
	}

	r := regexp.MustCompile(`(?m:server returned unexpected status code)`)
	if !r.MatchString(err.Error()) {
		t.Fatalf("Error not matched: %s", err)
	}
}

func Test_GetFunctionInfo_MissingURLPrefix(t *testing.T) {
	_, err := GetFunctionInfo("127.0.0.1:8080", "func-test", tlsNoVerify)

	if err == nil {
		t.Fatalf("Error was not returned")
	}

	expectedErrMsg := "invalid gateway URL:"
	r := regexp.MustCompile(fmt.Sprintf("(?m:%s)", expectedErrMsg))
	if !r.MatchString(err.Error()) {
		t.Fatalf("Want: %s, Got: %s", expectedErrMsg, err.Error())
	}
}

var expectedGetFunctionInfoResponse = requests.Function{
	Name:            "func-test1",
	Image:           "image-test1",
	Replicas:        1,
	InvocationCount: 1,
	EnvProcess:      "env-process test1",
}
<commit_msg>Add more unit test for GetFunctionInfo<commit_after>// Copyright (c) OpenFaaS Author(s) 2018. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

package proxy

import (
	"fmt"
	"net/http"
	"regexp"
	"testing"

	"github.com/openfaas/faas-cli/test"
	"github.com/openfaas/faas/gateway/requests"
)

func Test_GetFunctionInfo(t *testing.T) {
	s := test.MockHttpServer(t, []test.Request{
		{
			ResponseStatusCode: http.StatusOK,
			ResponseBody:       expectedGetFunctionInfoResponse,
		},
	})

	defer s.Close()
	result, err := GetFunctionInfo(s.URL, "func-test1", !tlsNoVerify)
	if err != nil {
		t.Fatalf("Error returned: %s", err)
	}
	if expectedGetFunctionInfoResponse != result {
		t.Fatalf("Want: %#v, Got: %#v", expectedGetFunctionInfoResponse, result)
	}
}

func Test_GetFunctionInfo_Not200(t *testing.T) {
	s := test.MockHttpServerStatus(t, http.StatusBadRequest)

	_, err := GetFunctionInfo(s.URL, "func-test1", tlsNoVerify)

	if err == nil {
		t.Fatalf("Error was not returned")
	}

	r := regexp.MustCompile(`(?m:server returned unexpected status code)`)
	if !r.MatchString(err.Error()) {
		t.Fatalf("Error not matched: %s", err)
	}
}

func Test_GetFunctionInfo_MissingURLPrefix(t *testing.T) {
	_, err := GetFunctionInfo("127.0.0.1:8080", "func-test", tlsNoVerify)

	if err == nil {
		t.Fatalf("Error was not returned")
	}

	expectedErrMsg := "invalid gateway URL:"
	r := regexp.MustCompile(fmt.Sprintf("(?m:%s)", expectedErrMsg))
	if !r.MatchString(err.Error()) {
		t.Fatalf("Want: %s, Got: %s", expectedErrMsg, err.Error())
	}
}

func Test_GetFunctionInfo_NotFound(t *testing.T) {
	s := test.MockHttpServerStatus(t, http.StatusNotFound)
	functionName := "funct-test"

	_, err := GetFunctionInfo(s.URL, functionName, tlsNoVerify)
	if err == nil {
		t.Fatalf("Error was not returned")
	}

	expectedErrMsg := fmt.Sprintf("No such function: %s", functionName)
	if err.Error() != expectedErrMsg {
		t.Fatalf("Want: %s, Got: %s", expectedErrMsg, err.Error())
	}

}

var expectedGetFunctionInfoResponse = requests.Function{
	Name:            "func-test1",
	Image:           "image-test1",
	Replicas:        1,
	InvocationCount: 1,
	EnvProcess:      "env-process test1",
}
<|endoftext|>


<commit_before>package commands

import (
	"fmt"

	"github.com/camd67/moebot/moebot_bot/util/db"
	"github.com/camd67/moebot/moebot_bot/util/reddit"
)

type RandomCommand struct {
	RedditHandle *reddit.Handle
}

func (ac *RandomCommand) Execute(pack *CommPackage) {
	send, err := ac.RedditHandle.GetRandomImage("awwnime")
	if err != nil {
		pack.session.ChannelMessageSend(pack.channel.ID, "Ooops... Looks like this command isn't working right now. Sorry!")
		return
	}

	pack.session.ChannelMessageSendComplex(pack.channel.ID, send)
}

func (ac *RandomCommand) GetPermLevel() db.Permission {
	return db.PermAll
}
func (ac *RandomCommand) GetCommandKeys() []string {
	return []string{"RANDOM", "R"}
}
func (ac *RandomCommand) GetCommandHelp(commPrefix string) string {
	return fmt.Sprintf("`%[1]s r` or `%[1]s random` - Posts a cute anime character.", commPrefix)
}
<commit_msg>Update help text to only display the
andom alias<commit_after>package commands

import (
	"fmt"

	"github.com/camd67/moebot/moebot_bot/util/db"
	"github.com/camd67/moebot/moebot_bot/util/reddit"
)

type RandomCommand struct {
	RedditHandle *reddit.Handle
}

func (ac *RandomCommand) Execute(pack *CommPackage) {
	send, err := ac.RedditHandle.GetRandomImage("awwnime")
	if err != nil {
		pack.session.ChannelMessageSend(pack.channel.ID, "Ooops... Looks like this command isn't working right now. Sorry!")
		return
	}

	pack.session.ChannelMessageSendComplex(pack.channel.ID, send)
}

func (ac *RandomCommand) GetPermLevel() db.Permission {
	return db.PermAll
}
func (ac *RandomCommand) GetCommandKeys() []string {
	return []string{"RANDOM", "R"}
}
func (ac *RandomCommand) GetCommandHelp(commPrefix string) string {
	return fmt.Sprintf("`%[1]s random` - Posts a cute anime character.", commPrefix)
}
<|endoftext|>


<commit_before>// Copyright 2018 Jigsaw Operations LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package shadowsocks

import (
	"errors"
	"io"
	"net"
	"testing"
	"time"

	"github.com/Jigsaw-Code/outline-ss-server/metrics"
	onet "github.com/Jigsaw-Code/outline-ss-server/net"
	logging "github.com/op/go-logging"
)

// Simulates receiving invalid TCP connection attempts on a server with 100 ciphers.
func BenchmarkTCPFindCipherFail(b *testing.B) {
	b.StopTimer()
	b.ResetTimer()

	logging.SetLevel(logging.INFO, "")
	listener, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0})
	if err != nil {
		b.Fatalf("ListenTCP failed: %v", err)
	}

	cipherList, err := MakeTestCiphers(MakeTestSecrets(100))
	if err != nil {
		b.Fatal(err)
	}
	testPayload := MakeTestPayload(50)
	for n := 0; n < b.N; n++ {
		go func() {
			conn, err := net.Dial("tcp", listener.Addr().String())
			if err != nil {
				b.Fatalf("Failed to dial %v: %v", listener.Addr(), err)
			}
			conn.Write(testPayload)
			conn.Close()
		}()
		clientConn, err := listener.AcceptTCP()
		if err != nil {
			b.Fatalf("AcceptTCP failed: %v", err)
		}
		b.StartTimer()
		findAccessKey(clientConn, cipherList)
		b.StopTimer()
	}
}

// Fake DuplexConn
// 1-way pipe, representing the upstream flow as seen by the server.
type conn struct {
	onet.DuplexConn
	clientAddr net.Addr
	reader     io.ReadCloser
	writer     io.WriteCloser
}

func (c *conn) Read(b []byte) (int, error) {
	return c.reader.Read(b)
}

func (c *conn) Write(b []byte) (int, error) {
	// Any downstream data is ignored.
	return len(b), nil
}

func (c *conn) Close() error {
	e1 := c.reader.Close()
	e2 := c.writer.Close()
	if e1 != nil {
		return e1
	}
	return e2
}

func (c *conn) LocalAddr() net.Addr {
	return nil
}

func (c *conn) RemoteAddr() net.Addr {
	return c.clientAddr
}

func (c *conn) SetDeadline(t time.Time) error {
	return errors.New("SetDeadline is not supported")
}

func (c *conn) SetReadDeadline(t time.Time) error {
	return errors.New("SetDeadline is not supported")
}

func (c *conn) SetWriteDeadline(t time.Time) error {
	return errors.New("SetDeadline is not supported")
}

func (c *conn) CloseRead() error {
	return c.reader.Close()
}

func (c *conn) CloseWrite() error {
	return nil
}

// Simulates receiving valid TCP connection attempts from 100 different users,
// each with their own cipher and their own IP address.
func BenchmarkTCPFindCipherRepeat(b *testing.B) {
	b.StopTimer()
	b.ResetTimer()

	logging.SetLevel(logging.INFO, "")

	const numCiphers = 100 // Must be <256
	cipherList, err := MakeTestCiphers(MakeTestSecrets(numCiphers))
	if err != nil {
		b.Fatal(err)
	}
	cipherEntries := [numCiphers]*CipherEntry{}
	for cipherNumber, element := range cipherList.SnapshotForClientIP(nil) {
		cipherEntries[cipherNumber] = element.Value.(*CipherEntry)
	}
	for n := 0; n < b.N; n++ {
		cipherNumber := byte(n % numCiphers)
		reader, writer := io.Pipe()
		addr := &net.TCPAddr{IP: net.IPv4(192, 0, 2, cipherNumber), Port: 54321}
		c := conn{clientAddr: addr, reader: reader, writer: writer}
		cipher := cipherEntries[cipherNumber].Cipher
		go NewShadowsocksWriter(writer, cipher).Write(MakeTestPayload(50))
		b.StartTimer()
		_, _, _, err := findAccessKey(&c, cipherList)
		b.StopTimer()
		if err != nil {
			b.Error(err)
		}
		c.Close()
	}
}

// Stub metrics implementation for testing replay defense.
type probeTestMetrics struct {
	metrics.ShadowsocksMetrics
	probeData   []metrics.ProxyMetrics
	probeStatus []string
	closeStatus []string
}

func (m *probeTestMetrics) AddTCPProbe(clientLocation, status, drainResult string, port int, data metrics.ProxyMetrics) {
	m.probeData = append(m.probeData, data)
	m.probeStatus = append(m.probeStatus, status)
}
func (m *probeTestMetrics) AddClosedTCPConnection(clientLocation, accessKey, status string, data metrics.ProxyMetrics, timeToCipher, duration time.Duration) {
	m.closeStatus = append(m.closeStatus, status)
}

func (m *probeTestMetrics) GetLocation(net.Addr) (string, error) {
	return "", nil
}
func (m *probeTestMetrics) SetNumAccessKeys(numKeys int, numPorts int) {
}
func (m *probeTestMetrics) AddOpenTCPConnection(clientLocation string) {
}
func (m *probeTestMetrics) AddUDPPacketFromClient(clientLocation, accessKey, status string, clientProxyBytes, proxyTargetBytes int, timeToCipher time.Duration) {
}
func (m *probeTestMetrics) AddUDPPacketFromTarget(clientLocation, accessKey, status string, targetProxyBytes, proxyClientBytes int) {
}
func (m *probeTestMetrics) AddUDPNatEntry()    {}
func (m *probeTestMetrics) RemoveUDPNatEntry() {}

func TestReplayDefense(t *testing.T) {
	logging.SetLevel(logging.DEBUG, "")

	listener, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0})
	if err != nil {
		t.Fatalf("ListenTCP failed: %v", err)
	}
	cipherList, err := MakeTestCiphers(MakeTestSecrets(1))
	if err != nil {
		t.Fatal(err)
	}
	replayCache := NewReplayCache(5)
	testMetrics := &probeTestMetrics{}
	const testTimeout = 200 * time.Millisecond
	s := NewTCPService(listener, cipherList, &replayCache, testMetrics, testTimeout)
	cipherEntry := cipherList.SnapshotForClientIP(nil)[0].Value.(*CipherEntry)
	cipher := cipherEntry.Cipher
	reader, writer := io.Pipe()
	go NewShadowsocksWriter(writer, cipher).Write([]byte{0})
	preamble := make([]byte, 32+2+16)
	if _, err := io.ReadFull(reader, preamble); err != nil {
		t.Fatal(err)
	}

	run := func() net.Conn {
		conn, err := net.Dial(listener.Addr().Network(), listener.Addr().String())
		if err != nil {
			t.Fatal(err)
		}
		n, err := conn.Write(preamble)
		if n < len(preamble) {
			t.Error(err)
		}
		return conn
	}

	go s.Start()

	// First run.
	conn1 := run()
	if len(testMetrics.probeData) != 0 {
		t.Errorf("First connection should not have triggered probe detection: %v", testMetrics.probeData[0])
	}
	if len(testMetrics.closeStatus) != 0 {
		t.Errorf("First connection should not have been closed yet: %v", testMetrics.probeData[0])
	}

	// Replay.
	conn2 := run()
	// Wait for the connection to be closed by the proxy after testTimeout.
	conn2.Read(make([]byte, 1))
	if len(testMetrics.probeData) == 1 {
		data := testMetrics.probeData[0]
		if data.ClientProxy != int64(len(preamble)) {
			t.Errorf("Unexpected probe data: %v", data)
		}
		status := testMetrics.probeStatus[0]
		if status != "ERR_REPLAY" {
			t.Errorf("Unexpected TCP probe status: %s", status)
		}
	} else {
		t.Error("Replay should have triggered probe detection")
	}
	if len(testMetrics.closeStatus) == 1 {
		status := testMetrics.closeStatus[0]
		if status != "ERR_REPLAY" {
			t.Errorf("Unexpected TCP close status: %s", status)
		}
	} else {
		t.Error("Replay should have reported an error status")
	}

	conn1.Close()
	s.Stop()
}

// Test 49, 50, and 51 bytes to ensure they have the same behavior.
// 50 bytes used to be the cutoff for different behavior.
func TestTCPProbeTimeout(t *testing.T) {
	logging.SetLevel(logging.CRITICAL, "")
	probeExpectTimeout(t, 49)
	probeExpectTimeout(t, 50)
	probeExpectTimeout(t, 51)
}

func probeExpectTimeout(t *testing.T, payloadSize int) {
	const testTimeout = 200 * time.Millisecond

	listener, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0})
	if err != nil {
		t.Fatalf("ListenTCP failed: %v", err)
	}
	cipherList, err := MakeTestCiphers(MakeTestSecrets(5))
	if err != nil {
		t.Fatal(err)
	}
	testMetrics := &probeTestMetrics{}
	s := NewTCPService(listener, cipherList, nil, testMetrics, testTimeout)

	testPayload := MakeTestPayload(payloadSize)
	done := make(chan bool)
	go func() {
		defer func() { done <- true }()
		conn, err := net.Dial("tcp", listener.Addr().String())
		if err != nil {
			t.Fatalf("Failed to dial %v: %v", listener.Addr(), err)
		}
		conn.Write(testPayload)
		timerStart := time.Now()
		buf := make([]byte, 1024)
		bytesRead, err := conn.Read(buf) // will hang until connection is closed
		elapsedTime := time.Since(timerStart)
		switch {
		case err != io.EOF:
			t.Fatalf("Expected error EOF, got %v", err)
		case bytesRead > 0:
			t.Fatalf("Expected to read 0 bytes, got %v bytes", bytesRead)
		case elapsedTime < testTimeout || elapsedTime > testTimeout+10*time.Millisecond:
			t.Fatalf("Expected elapsed time close to %v, got %v", testTimeout, elapsedTime)
		default:
			// ok
		}
	}()

	go s.Start()
	<-done

	if len(testMetrics.probeData) == 1 {
		data := testMetrics.probeData[0]
		if data.ClientProxy != int64(payloadSize) {
			t.Errorf("Unexpected probe data: %v, expected %d", data, payloadSize)
		}
	} else {
		t.Error("Bad handshake should have triggered probe detection")
	}
	if len(testMetrics.probeStatus) == 1 {
		status := testMetrics.probeStatus[0]
		if status != "ERR_CIPHER" {
			t.Errorf("Unexpected TCP probe status: %s", status)
		}
	} else {
		t.Error("Bad handshake should have reported an error status")
	}
	if len(testMetrics.closeStatus) == 1 {
		status := testMetrics.closeStatus[0]
		if status != "ERR_CIPHER" {
			t.Errorf("Unexpected TCP close status: %s", status)
		}
	} else {
		t.Error("Bad handshake should have reported an error status")
	}

	s.Stop()
}
<commit_msg>Fix a race in the replay integration test<commit_after>// Copyright 2018 Jigsaw Operations LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package shadowsocks

import (
	"errors"
	"io"
	"net"
	"testing"
	"time"

	"github.com/Jigsaw-Code/outline-ss-server/metrics"
	onet "github.com/Jigsaw-Code/outline-ss-server/net"
	logging "github.com/op/go-logging"
)

// Simulates receiving invalid TCP connection attempts on a server with 100 ciphers.
func BenchmarkTCPFindCipherFail(b *testing.B) {
	b.StopTimer()
	b.ResetTimer()

	logging.SetLevel(logging.INFO, "")
	listener, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0})
	if err != nil {
		b.Fatalf("ListenTCP failed: %v", err)
	}

	cipherList, err := MakeTestCiphers(MakeTestSecrets(100))
	if err != nil {
		b.Fatal(err)
	}
	testPayload := MakeTestPayload(50)
	for n := 0; n < b.N; n++ {
		go func() {
			conn, err := net.Dial("tcp", listener.Addr().String())
			if err != nil {
				b.Fatalf("Failed to dial %v: %v", listener.Addr(), err)
			}
			conn.Write(testPayload)
			conn.Close()
		}()
		clientConn, err := listener.AcceptTCP()
		if err != nil {
			b.Fatalf("AcceptTCP failed: %v", err)
		}
		b.StartTimer()
		findAccessKey(clientConn, cipherList)
		b.StopTimer()
	}
}

// Fake DuplexConn
// 1-way pipe, representing the upstream flow as seen by the server.
type conn struct {
	onet.DuplexConn
	clientAddr net.Addr
	reader     io.ReadCloser
	writer     io.WriteCloser
}

func (c *conn) Read(b []byte) (int, error) {
	return c.reader.Read(b)
}

func (c *conn) Write(b []byte) (int, error) {
	// Any downstream data is ignored.
	return len(b), nil
}

func (c *conn) Close() error {
	e1 := c.reader.Close()
	e2 := c.writer.Close()
	if e1 != nil {
		return e1
	}
	return e2
}

func (c *conn) LocalAddr() net.Addr {
	return nil
}

func (c *conn) RemoteAddr() net.Addr {
	return c.clientAddr
}

func (c *conn) SetDeadline(t time.Time) error {
	return errors.New("SetDeadline is not supported")
}

func (c *conn) SetReadDeadline(t time.Time) error {
	return errors.New("SetDeadline is not supported")
}

func (c *conn) SetWriteDeadline(t time.Time) error {
	return errors.New("SetDeadline is not supported")
}

func (c *conn) CloseRead() error {
	return c.reader.Close()
}

func (c *conn) CloseWrite() error {
	return nil
}

// Simulates receiving valid TCP connection attempts from 100 different users,
// each with their own cipher and their own IP address.
func BenchmarkTCPFindCipherRepeat(b *testing.B) {
	b.StopTimer()
	b.ResetTimer()

	logging.SetLevel(logging.INFO, "")

	const numCiphers = 100 // Must be <256
	cipherList, err := MakeTestCiphers(MakeTestSecrets(numCiphers))
	if err != nil {
		b.Fatal(err)
	}
	cipherEntries := [numCiphers]*CipherEntry{}
	for cipherNumber, element := range cipherList.SnapshotForClientIP(nil) {
		cipherEntries[cipherNumber] = element.Value.(*CipherEntry)
	}
	for n := 0; n < b.N; n++ {
		cipherNumber := byte(n % numCiphers)
		reader, writer := io.Pipe()
		addr := &net.TCPAddr{IP: net.IPv4(192, 0, 2, cipherNumber), Port: 54321}
		c := conn{clientAddr: addr, reader: reader, writer: writer}
		cipher := cipherEntries[cipherNumber].Cipher
		go NewShadowsocksWriter(writer, cipher).Write(MakeTestPayload(50))
		b.StartTimer()
		_, _, _, err := findAccessKey(&c, cipherList)
		b.StopTimer()
		if err != nil {
			b.Error(err)
		}
		c.Close()
	}
}

// Stub metrics implementation for testing replay defense.
type probeTestMetrics struct {
	metrics.ShadowsocksMetrics
	probeData   []metrics.ProxyMetrics
	probeStatus []string
	closeStatus []string
}

func (m *probeTestMetrics) AddTCPProbe(clientLocation, status, drainResult string, port int, data metrics.ProxyMetrics) {
	m.probeData = append(m.probeData, data)
	m.probeStatus = append(m.probeStatus, status)
}
func (m *probeTestMetrics) AddClosedTCPConnection(clientLocation, accessKey, status string, data metrics.ProxyMetrics, timeToCipher, duration time.Duration) {
	m.closeStatus = append(m.closeStatus, status)
}

func (m *probeTestMetrics) GetLocation(net.Addr) (string, error) {
	return "", nil
}
func (m *probeTestMetrics) SetNumAccessKeys(numKeys int, numPorts int) {
}
func (m *probeTestMetrics) AddOpenTCPConnection(clientLocation string) {
}
func (m *probeTestMetrics) AddUDPPacketFromClient(clientLocation, accessKey, status string, clientProxyBytes, proxyTargetBytes int, timeToCipher time.Duration) {
}
func (m *probeTestMetrics) AddUDPPacketFromTarget(clientLocation, accessKey, status string, targetProxyBytes, proxyClientBytes int) {
}
func (m *probeTestMetrics) AddUDPNatEntry()    {}
func (m *probeTestMetrics) RemoveUDPNatEntry() {}

func TestReplayDefense(t *testing.T) {
	logging.SetLevel(logging.DEBUG, "")

	listener, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0})
	if err != nil {
		t.Fatalf("ListenTCP failed: %v", err)
	}
	cipherList, err := MakeTestCiphers(MakeTestSecrets(1))
	if err != nil {
		t.Fatal(err)
	}
	replayCache := NewReplayCache(5)
	testMetrics := &probeTestMetrics{}
	const testTimeout = 200 * time.Millisecond
	s := NewTCPService(listener, cipherList, &replayCache, testMetrics, testTimeout)
	cipherEntry := cipherList.SnapshotForClientIP(nil)[0].Value.(*CipherEntry)
	cipher := cipherEntry.Cipher
	reader, writer := io.Pipe()
	go NewShadowsocksWriter(writer, cipher).Write([]byte{0})
	preamble := make([]byte, 32+2+16)
	if _, err := io.ReadFull(reader, preamble); err != nil {
		t.Fatal(err)
	}

	run := func() net.Conn {
		conn, err := net.Dial(listener.Addr().Network(), listener.Addr().String())
		if err != nil {
			t.Fatal(err)
		}
		n, err := conn.Write(preamble)
		if n < len(preamble) {
			t.Error(err)
		}
		return conn
	}

	go s.Start()

	// First run.
	conn1 := run()
	if len(testMetrics.probeData) != 0 {
		t.Errorf("First connection should not have triggered probe detection: %v", testMetrics.probeData[0])
	}
	if len(testMetrics.closeStatus) != 0 {
		t.Errorf("First connection should not have been closed yet: %v", testMetrics.probeData[0])
	}
	// Write a minimal invalid chunk to trigger a proxy-driven close.
	conn1.Write(make([]byte, 2+16))
	// Wait for the close.  This ensures that conn1 and conn2 can't be
	// processed out of order at the proxy.
	conn1.Read(make([]byte, 1))

	// Replay.
	conn2 := run()
	// Wait for the connection to be closed by the proxy after testTimeout.
	conn2.Read(make([]byte, 1))
	if len(testMetrics.probeData) == 1 {
		data := testMetrics.probeData[0]
		if data.ClientProxy != int64(len(preamble)) {
			t.Errorf("Unexpected probe data: %v", data)
		}
		status := testMetrics.probeStatus[0]
		if status != "ERR_REPLAY" {
			t.Errorf("Unexpected TCP probe status: %s", status)
		}
	} else {
		t.Error("Replay should have triggered probe detection")
	}
	if len(testMetrics.closeStatus) == 2 {
		status := testMetrics.closeStatus[1]
		if status != "ERR_REPLAY" {
			t.Errorf("Unexpected TCP close status: %s", status)
		}
	} else {
		t.Error("Replay should have reported an error status")
	}

	conn1.Close()
	s.Stop()
}

// Test 49, 50, and 51 bytes to ensure they have the same behavior.
// 50 bytes used to be the cutoff for different behavior.
func TestTCPProbeTimeout(t *testing.T) {
	logging.SetLevel(logging.CRITICAL, "")
	probeExpectTimeout(t, 49)
	probeExpectTimeout(t, 50)
	probeExpectTimeout(t, 51)
}

func probeExpectTimeout(t *testing.T, payloadSize int) {
	const testTimeout = 200 * time.Millisecond

	listener, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0})
	if err != nil {
		t.Fatalf("ListenTCP failed: %v", err)
	}
	cipherList, err := MakeTestCiphers(MakeTestSecrets(5))
	if err != nil {
		t.Fatal(err)
	}
	testMetrics := &probeTestMetrics{}
	s := NewTCPService(listener, cipherList, nil, testMetrics, testTimeout)

	testPayload := MakeTestPayload(payloadSize)
	done := make(chan bool)
	go func() {
		defer func() { done <- true }()
		conn, err := net.Dial("tcp", listener.Addr().String())
		if err != nil {
			t.Fatalf("Failed to dial %v: %v", listener.Addr(), err)
		}
		conn.Write(testPayload)
		timerStart := time.Now()
		buf := make([]byte, 1024)
		bytesRead, err := conn.Read(buf) // will hang until connection is closed
		elapsedTime := time.Since(timerStart)
		switch {
		case err != io.EOF:
			t.Fatalf("Expected error EOF, got %v", err)
		case bytesRead > 0:
			t.Fatalf("Expected to read 0 bytes, got %v bytes", bytesRead)
		case elapsedTime < testTimeout || elapsedTime > testTimeout+10*time.Millisecond:
			t.Fatalf("Expected elapsed time close to %v, got %v", testTimeout, elapsedTime)
		default:
			// ok
		}
	}()

	go s.Start()
	<-done

	if len(testMetrics.probeData) == 1 {
		data := testMetrics.probeData[0]
		if data.ClientProxy != int64(payloadSize) {
			t.Errorf("Unexpected probe data: %v, expected %d", data, payloadSize)
		}
	} else {
		t.Error("Bad handshake should have triggered probe detection")
	}
	if len(testMetrics.probeStatus) == 1 {
		status := testMetrics.probeStatus[0]
		if status != "ERR_CIPHER" {
			t.Errorf("Unexpected TCP probe status: %s", status)
		}
	} else {
		t.Error("Bad handshake should have reported an error status")
	}
	if len(testMetrics.closeStatus) == 1 {
		status := testMetrics.closeStatus[0]
		if status != "ERR_CIPHER" {
			t.Errorf("Unexpected TCP close status: %s", status)
		}
	} else {
		t.Error("Bad handshake should have reported an error status")
	}

	s.Stop()
}
<|endoftext|>


<commit_before>package shadowsocks

import (
	"encoding/binary"
	"fmt"
	"net"
	"strconv"
	"sync"
	"syscall"
	"time"
)

const (
	idType  = 0 // address type index
	idIP0   = 1 // ip addres start index
	idDmLen = 1 // domain address length index
	idDm0   = 2 // domain address start index

	typeIPv4 = 1 // type is ipv4 address
	typeDm   = 3 // type is domain address
	typeIPv6 = 4 // type is ipv6 address

	lenIPv4     = 1 + net.IPv4len + 2 // 1addrType + ipv4 + 2port
	lenIPv6     = 1 + net.IPv6len + 2 // 1addrType + ipv6 + 2port
	lenDmBase   = 1 + 1 + 2           // 1addrType + 1addrLen + 2port, plus addrLen
	lenHmacSha1 = 10
)

var (
	reqList            = newReqList()
	natlist            = newNatTable()
	udpTimeout         = 30 * time.Second
	reqListRefreshTime = 5 * time.Minute
)

type natTable struct {
	sync.Mutex
	conns map[string]net.PacketConn
}

func newNatTable() *natTable {
	return &natTable{conns: map[string]net.PacketConn{}}
}

func (table *natTable) DeleteAndClose(index string) {
	table.Lock()
	defer table.Unlock()
	c, ok := table.conns[index]
	if ok {
		c.Close()
		delete(table.conns, index)
	}
}

func (table *natTable) Get(index string) (c net.PacketConn, ok bool, err error) {
	table.Lock()
	defer table.Unlock()
	c, ok = table.conns[index]
	if !ok {
		c, err = net.ListenPacket("udp", "")
		if err != nil {
			return nil, false, err
		}
		table.conns[index] = c
	}
	return
}

type ReqList struct {
	sync.Mutex
	List map[string]([]byte)
}

func newReqList() *ReqList {
	ret := &ReqList{List: map[string]([]byte){}}
	go func() {
		for {
			time.Sleep(reqListRefreshTime)
			ret.Refresh()
		}
	}()
	return ret
}

func (r *ReqList) Refresh() {
	r.Lock()
	defer r.Unlock()
	for k, _ := range r.List {
		delete(r.List, k)
	}
}

func (r *ReqList) Get(dstaddr string) (req []byte, ok bool) {
	r.Lock()
	defer r.Unlock()
	req, ok = r.List[dstaddr]
	return
}

func (r *ReqList) Put(dstaddr string, req []byte) {
	r.Lock()
	defer r.Unlock()
	r.List[dstaddr] = req
	return
}

func ParseHeader(addr net.Addr) ([]byte, int) {
	// if the request address type is domain, it cannot be reverselookuped
	ip, port, err := net.SplitHostPort(addr.String())
	if err != nil {
		return nil, 0
	}
	buf := make([]byte, 20)
	IP := net.ParseIP(ip)
	b1 := IP.To4()
	iplen := 0
	if b1 == nil { //ipv6
		b1 = IP.To16()
		buf[0] = typeIPv6
		iplen = net.IPv6len
	} else { //ipv4
		buf[0] = typeIPv4
		iplen = net.IPv4len
	}
	copy(buf[1:], b1)
	port_i, _ := strconv.Atoi(port)
	binary.BigEndian.PutUint16(buf[1+iplen:], uint16(port_i))
	return buf[:1+iplen+2], 1 + iplen + 2
}

func Pipeloop(ss *SecurePacketConn, addr net.Addr, in net.PacketConn, compatiblemode bool) {
	buf := leakyBuf.Get()
	defer leakyBuf.Put(buf)
	defer in.Close()
	for {
		in.SetDeadline(time.Now().Add(udpTimeout))
		n, raddr, err := in.ReadFrom(buf)
		if err != nil {
			if ne, ok := err.(*net.OpError); ok {
				if ne.Err == syscall.EMFILE || ne.Err == syscall.ENFILE {
					// log too many open file error
					// EMFILE is process reaches open file limits, ENFILE is system limit
					Debug.Println("[udp]read error:", err)
				}
			}
			Debug.Printf("[udp]closed pipe %s<-%s\n", addr, in.LocalAddr())
			return
		}
		// need improvement here
		if req, ok := reqList.Get(raddr.String()); ok {
			if compatiblemode {
				ss.ForceOTAWriteTo(append(req, buf[:n]...), addr)
			} else {
				ss.WriteTo(append(req, buf[:n]...), addr)
			}
		} else {
			header, hlen := ParseHeader(raddr)
			if compatiblemode {
				ss.ForceOTAWriteTo(append(header[:hlen], buf[:n]...), addr)
			} else {
				ss.WriteTo(append(header[:hlen], buf[:n]...), addr)
			}

		}
	}
}

func handleUDPConnection(handle *SecurePacketConn, n int, src net.Addr, receive []byte) {
	var dstIP net.IP
	var reqLen int
	var ota bool
	addrType := receive[idType]
	defer leakyBuf.Put(receive)

	if addrType&OneTimeAuthMask > 0 {
		ota = true
	}
	if handle.IsOta() && !ota {
		Debug.Println("[udp]incoming connection dropped, due to ota enforcement")
		return
	}
	compatiblemode := !handle.IsOta() && ota

	switch addrType & AddrMask {
	case typeIPv4:
		reqLen = lenIPv4
		dstIP = net.IP(receive[idIP0 : idIP0+net.IPv4len])
	case typeIPv6:
		reqLen = lenIPv6
		dstIP = net.IP(receive[idIP0 : idIP0+net.IPv6len])
	case typeDm:
		reqLen = int(receive[idDmLen]) + lenDmBase
		dIP, err := net.ResolveIPAddr("ip", string(receive[idDm0:idDm0+int(receive[idDmLen])])) // carefully with const type
		if err != nil {
			Debug.Printf("[udp]failed to resolve domain name: %s\n", string(receive[idDm0:idDm0+receive[idDmLen]]))
			return
		}
		dstIP = dIP.IP
	default:
		Debug.Printf("[udp]addrType %d not supported", addrType)
		return
	}
	dst := &net.UDPAddr{
		IP:   dstIP,
		Port: int(binary.BigEndian.Uint16(receive[reqLen-2 : reqLen])),
	}
	if _, ok := reqList.Get(dst.String()); !ok {
		req := make([]byte, reqLen)
		copy(req, receive)
		reqList.Put(dst.String(), req)
	}

	remote, exist, err := natlist.Get(src.String())
	if err != nil {
		return
	}
	if !exist {
		Debug.Printf("[udp]new client %s->%s via %s ota=%v\n", src, dst, remote.LocalAddr(), ota)
		go func() {
			Pipeloop(handle, src, remote, compatiblemode)
			natlist.DeleteAndClose(src.String())
		}()
	} else {
		Debug.Printf("[udp]using cached client %s->%s via %s ota=%v\n", src, dst, remote.LocalAddr(), ota)
	}
	if remote == nil {
		fmt.Println("WTF")
	}
	remote.SetDeadline(time.Now().Add(udpTimeout))
	_, err = remote.WriteTo(receive[reqLen:n], dst)
	if err != nil {
		if ne, ok := err.(*net.OpError); ok && (ne.Err == syscall.EMFILE || ne.Err == syscall.ENFILE) {
			// log too many open file error
			// EMFILE is process reaches open file limits, ENFILE is system limit
			Debug.Println("[udp]write error:", err)
		} else {
			Debug.Println("[udp]error connecting to:", dst, err)
		}
		natlist.DeleteAndClose(src.String())
	}
	// Pipeloop
	return
}

func ReadAndHandleUDPReq(c *SecurePacketConn) {
	buf := leakyBuf.Get()
	n, src, err := c.ReadFrom(buf[0:])
	if err != nil {
		return
	}
	go handleUDPConnection(c, n, src, buf)
}
<commit_msg>Avoid syscall panic on windows. When receiving false data.<commit_after>package shadowsocks

import (
	"encoding/binary"
	"fmt"
	"net"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
)

const (
	idType  = 0 // address type index
	idIP0   = 1 // ip addres start index
	idDmLen = 1 // domain address length index
	idDm0   = 2 // domain address start index

	typeIPv4 = 1 // type is ipv4 address
	typeDm   = 3 // type is domain address
	typeIPv6 = 4 // type is ipv6 address

	lenIPv4     = 1 + net.IPv4len + 2 // 1addrType + ipv4 + 2port
	lenIPv6     = 1 + net.IPv6len + 2 // 1addrType + ipv6 + 2port
	lenDmBase   = 1 + 1 + 2           // 1addrType + 1addrLen + 2port, plus addrLen
	lenHmacSha1 = 10
)

var (
	reqList            = newReqList()
	natlist            = newNatTable()
	udpTimeout         = 30 * time.Second
	reqListRefreshTime = 5 * time.Minute
)

type natTable struct {
	sync.Mutex
	conns map[string]net.PacketConn
}

func newNatTable() *natTable {
	return &natTable{conns: map[string]net.PacketConn{}}
}

func (table *natTable) DeleteAndClose(index string) {
	table.Lock()
	defer table.Unlock()
	c, ok := table.conns[index]
	if ok {
		c.Close()
		delete(table.conns, index)
	}
}

func (table *natTable) Get(index string) (c net.PacketConn, ok bool, err error) {
	table.Lock()
	defer table.Unlock()
	c, ok = table.conns[index]
	if !ok {
		c, err = net.ListenPacket("udp", "")
		if err != nil {
			return nil, false, err
		}
		table.conns[index] = c
	}
	return
}

type ReqList struct {
	sync.Mutex
	List map[string]([]byte)
}

func newReqList() *ReqList {
	ret := &ReqList{List: map[string]([]byte){}}
	go func() {
		for {
			time.Sleep(reqListRefreshTime)
			ret.Refresh()
		}
	}()
	return ret
}

func (r *ReqList) Refresh() {
	r.Lock()
	defer r.Unlock()
	for k, _ := range r.List {
		delete(r.List, k)
	}
}

func (r *ReqList) Get(dstaddr string) (req []byte, ok bool) {
	r.Lock()
	defer r.Unlock()
	req, ok = r.List[dstaddr]
	return
}

func (r *ReqList) Put(dstaddr string, req []byte) {
	r.Lock()
	defer r.Unlock()
	r.List[dstaddr] = req
	return
}

func ParseHeader(addr net.Addr) ([]byte, int) {
	// if the request address type is domain, it cannot be reverselookuped
	ip, port, err := net.SplitHostPort(addr.String())
	if err != nil {
		return nil, 0
	}
	buf := make([]byte, 20)
	IP := net.ParseIP(ip)
	b1 := IP.To4()
	iplen := 0
	if b1 == nil { //ipv6
		b1 = IP.To16()
		buf[0] = typeIPv6
		iplen = net.IPv6len
	} else { //ipv4
		buf[0] = typeIPv4
		iplen = net.IPv4len
	}
	copy(buf[1:], b1)
	port_i, _ := strconv.Atoi(port)
	binary.BigEndian.PutUint16(buf[1+iplen:], uint16(port_i))
	return buf[:1+iplen+2], 1 + iplen + 2
}

func Pipeloop(ss *SecurePacketConn, addr net.Addr, in net.PacketConn, compatiblemode bool) {
	buf := leakyBuf.Get()
	defer leakyBuf.Put(buf)
	defer in.Close()
	for {
		in.SetDeadline(time.Now().Add(udpTimeout))
		n, raddr, err := in.ReadFrom(buf)
		if err != nil {
			if ne, ok := err.(*net.OpError); ok {
				if ne.Err == syscall.EMFILE || ne.Err == syscall.ENFILE {
					// log too many open file error
					// EMFILE is process reaches open file limits, ENFILE is system limit
					Debug.Println("[udp]read error:", err)
				}
			}
			Debug.Printf("[udp]closed pipe %s<-%s\n", addr, in.LocalAddr())
			return
		}
		// need improvement here
		if req, ok := reqList.Get(raddr.String()); ok {
			if compatiblemode {
				ss.ForceOTAWriteTo(append(req, buf[:n]...), addr)
			} else {
				ss.WriteTo(append(req, buf[:n]...), addr)
			}
		} else {
			header, hlen := ParseHeader(raddr)
			if compatiblemode {
				ss.ForceOTAWriteTo(append(header[:hlen], buf[:n]...), addr)
			} else {
				ss.WriteTo(append(header[:hlen], buf[:n]...), addr)
			}

		}
	}
}

func handleUDPConnection(handle *SecurePacketConn, n int, src net.Addr, receive []byte) {
	var dstIP net.IP
	var reqLen int
	var ota bool
	addrType := receive[idType]
	defer leakyBuf.Put(receive)

	if addrType&OneTimeAuthMask > 0 {
		ota = true
	}
	if handle.IsOta() && !ota {
		Debug.Println("[udp]incoming connection dropped, due to ota enforcement")
		return
	}
	compatiblemode := !handle.IsOta() && ota

	switch addrType & AddrMask {
	case typeIPv4:
		reqLen = lenIPv4
		if len(receive) < reqLen {
			Debug.Println("[udp]invalid received message.")
		}
		dstIP = net.IP(receive[idIP0 : idIP0+net.IPv4len])
	case typeIPv6:
		reqLen = lenIPv6
		if len(receive) < reqLen {
			Debug.Println("[udp]invalid received message.")
		}
		dstIP = net.IP(receive[idIP0 : idIP0+net.IPv6len])
	case typeDm:
		reqLen = int(receive[idDmLen]) + lenDmBase
		if len(receive) < reqLen {
			Debug.Println("[udp]invalid received message.")
		}
		name := string(receive[idDm0 : idDm0+int(receive[idDmLen])])
		// avoid panic: syscall: string with NUL passed to StringToUTF16 on windows.
		if strings.ContainsRune(name, 0x00) {
			fmt.Println("[udp]invalid domain name.")
			return
		}
		dIP, err := net.ResolveIPAddr("ip", name) // carefully with const type
		if err != nil {
			Debug.Printf("[udp]failed to resolve domain name: %s\n", string(receive[idDm0:idDm0+receive[idDmLen]]))
			return
		}
		dstIP = dIP.IP
	default:
		Debug.Printf("[udp]addrType %d not supported", addrType)
		return
	}
	dst := &net.UDPAddr{
		IP:   dstIP,
		Port: int(binary.BigEndian.Uint16(receive[reqLen-2 : reqLen])),
	}
	if _, ok := reqList.Get(dst.String()); !ok {
		req := make([]byte, reqLen)
		copy(req, receive)
		reqList.Put(dst.String(), req)
	}

	remote, exist, err := natlist.Get(src.String())
	if err != nil {
		return
	}
	if !exist {
		Debug.Printf("[udp]new client %s->%s via %s ota=%v\n", src, dst, remote.LocalAddr(), ota)
		go func() {
			Pipeloop(handle, src, remote, compatiblemode)
			natlist.DeleteAndClose(src.String())
		}()
	} else {
		Debug.Printf("[udp]using cached client %s->%s via %s ota=%v\n", src, dst, remote.LocalAddr(), ota)
	}
	if remote == nil {
		fmt.Println("WTF")
	}
	remote.SetDeadline(time.Now().Add(udpTimeout))
	_, err = remote.WriteTo(receive[reqLen:n], dst)
	if err != nil {
		if ne, ok := err.(*net.OpError); ok && (ne.Err == syscall.EMFILE || ne.Err == syscall.ENFILE) {
			// log too many open file error
			// EMFILE is process reaches open file limits, ENFILE is system limit
			Debug.Println("[udp]write error:", err)
		} else {
			Debug.Println("[udp]error connecting to:", dst, err)
		}
		natlist.DeleteAndClose(src.String())
	}
	// Pipeloop
	return
}

func ReadAndHandleUDPReq(c *SecurePacketConn) {
	buf := leakyBuf.Get()
	n, src, err := c.ReadFrom(buf[0:])
	if err != nil {
		return
	}
	go handleUDPConnection(c, n, src, buf)
}
<|endoftext|>


<commit_before><commit_msg>fixing lang errors ##<commit_after><|endoftext|>


<commit_before><commit_msg>Replace self defined const StatusTooManyRequests with http.StatusTooManyRequests.<commit_after><|endoftext|>