<commit_before>// Accounting and limiting reader

package fs

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/VividCortex/ewma"
	"github.com/tsenart/tb"
)

// Globals
var (
	Stats       = NewStats()
	tokenBucket *tb.Bucket
)

// Start the token bucket if necessary
func startTokenBucket() {
	if bwLimit > 0 {
		tokenBucket = tb.NewBucket(int64(bwLimit), 100*time.Millisecond)
		Log(nil, "Starting bandwidth limiter at %vBytes/s", &bwLimit)
	}
}

// stringSet holds a set of strings
type stringSet map[string]struct{}

// inProgress holds a synchronizes map of in progress transfers
type inProgress struct {
	mu sync.Mutex
	m  map[string]*Account
}

// newInProgress makes a new inProgress object
func newInProgress() *inProgress {
	return &inProgress{
		m: make(map[string]*Account, Config.Transfers),
	}
}

// set marks the name as in progress
func (ip *inProgress) set(name string, acc *Account) {
	ip.mu.Lock()
	defer ip.mu.Unlock()
	ip.m[name] = acc
}

// clear marks the name as no longer in progress
func (ip *inProgress) clear(name string) {
	ip.mu.Lock()
	defer ip.mu.Unlock()
	delete(ip.m, name)
}

// get gets the account for name, of nil if not found
func (ip *inProgress) get(name string) *Account {
	ip.mu.Lock()
	defer ip.mu.Unlock()
	return ip.m[name]
}

// Strings returns all the strings in the stringSet
func (ss stringSet) Strings() []string {
	strings := make([]string, 0, len(ss))
	for name := range ss {
		var out string
		if acc := Stats.inProgress.get(name); acc != nil {
			out = acc.String()
		} else {
			out = name
		}
		strings = append(strings, " * "+out)
	}
	sorted := sort.StringSlice(strings)
	sorted.Sort()
	return sorted
}

// String returns all the file names in the stringSet joined by newline
func (ss stringSet) String() string {
	return strings.Join(ss.Strings(), "\n")
}

// StatsInfo limits and accounts all transfers
type StatsInfo struct {
	lock         sync.RWMutex
	bytes        int64
	errors       int64
	checks       int64
	checking     stringSet
	transfers    int64
	transferring stringSet
	start        time.Time
	inProgress   *inProgress
}

// NewStats cretates an initialised StatsInfo
func NewStats() *StatsInfo {
	return &StatsInfo{
		checking:     make(stringSet, Config.Checkers),
		transferring: make(stringSet, Config.Transfers),
		start:        time.Now(),
		inProgress:   newInProgress(),
	}
}

// String convert the StatsInfo to a string for printing
func (s *StatsInfo) String() string {
	s.lock.RLock()
	defer s.lock.RUnlock()
	dt := time.Now().Sub(s.start)
	dtSeconds := dt.Seconds()
	speed := 0.0
	if dt > 0 {
		speed = float64(s.bytes) / 1024 / dtSeconds
	}
	dtRounded := dt - (dt % (time.Second / 10))
	buf := &bytes.Buffer{}
	fmt.Fprintf(buf, `
Transferred:   %10d Bytes (%7.2f kByte/s)
Errors:        %10d
Checks:        %10d
Transferred:   %10d
Elapsed time:  %10v
`,
		s.bytes, speed,
		s.errors,
		s.checks,
		s.transfers,
		dtRounded)
	if len(s.checking) > 0 {
		fmt.Fprintf(buf, "Checking:\n%s\n", s.checking)
	}
	if len(s.transferring) > 0 {
		fmt.Fprintf(buf, "Transferring:\n%s\n", s.transferring)
	}
	return buf.String()
}

// Log outputs the StatsInfo to the log
func (s *StatsInfo) Log() {
	log.Printf("%v\n", s)
}

// Bytes updates the stats for bytes bytes
func (s *StatsInfo) Bytes(bytes int64) {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.bytes += bytes
}

// Errors updates the stats for errors
func (s *StatsInfo) Errors(errors int64) {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.errors += errors
}

// GetErrors reads the number of errors
func (s *StatsInfo) GetErrors() int64 {
	s.lock.RLock()
	defer s.lock.RUnlock()
	return s.errors
}

// ResetCounters sets the counters (bytes, checks, errors, transfers) to 0
func (s *StatsInfo) ResetCounters() {
	s.lock.RLock()
	defer s.lock.RUnlock()
	s.bytes = 0
	s.errors = 0
	s.checks = 0
	s.transfers = 0
}

// ResetErrors sets the errors count to 0
func (s *StatsInfo) ResetErrors() {
	s.lock.RLock()
	defer s.lock.RUnlock()
	s.errors = 0
}

// Errored returns whether there have been any errors
func (s *StatsInfo) Errored() bool {
	s.lock.RLock()
	defer s.lock.RUnlock()
	return s.errors != 0
}

// Error adds a single error into the stats
func (s *StatsInfo) Error() {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.errors++
}

// Checking adds a check into the stats
func (s *StatsInfo) Checking(o Object) {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.checking[o.Remote()] = struct{}{}
}

// DoneChecking removes a check from the stats
func (s *StatsInfo) DoneChecking(o Object) {
	s.lock.Lock()
	defer s.lock.Unlock()
	delete(s.checking, o.Remote())
	s.checks++
}

// GetTransfers reads the number of transfers
func (s *StatsInfo) GetTransfers() int64 {
	s.lock.RLock()
	defer s.lock.RUnlock()
	return s.transfers
}

// Transferring adds a transfer into the stats
func (s *StatsInfo) Transferring(o Object) {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.transferring[o.Remote()] = struct{}{}
}

// DoneTransferring removes a transfer from the stats
func (s *StatsInfo) DoneTransferring(o Object) {
	s.lock.Lock()
	defer s.lock.Unlock()
	delete(s.transferring, o.Remote())
	s.transfers++
}

// Account limits and accounts for one transfer
type Account struct {
	// The mutex is to make sure Read() and Close() aren't called
	// concurrently.  Unfortunately the persistent connection loop
	// in http transport calls Read() after Do() returns on
	// CancelRequest so this race can happen when it apparently
	// shouldn't.
	mu      sync.Mutex
	in      io.ReadCloser
	size    int64
	name    string
	statmu  sync.Mutex         // Separate mutex for stat values.
	bytes   int64              // Total number of bytes read
	start   time.Time          // Start time of first read
	lpTime  time.Time          // Time of last average measurement
	lpBytes int                // Number of bytes read since last measurement
	avg     ewma.MovingAverage // Moving average of last few measurements
	closed  bool               // set if the file is closed
	exit    chan struct{}      // channel that will be closed when transfer is finished
}

// NewAccount makes a Account reader for an object
func NewAccount(in io.ReadCloser, obj Object) *Account {
	acc := &Account{
		in:     in,
		size:   obj.Size(),
		name:   obj.Remote(),
		exit:   make(chan struct{}),
		avg:    ewma.NewMovingAverage(),
		lpTime: time.Now(),
	}
	go acc.averageLoop()
	Stats.inProgress.set(acc.name, acc)
	return acc
}

func (file *Account) averageLoop() {
	tick := time.NewTicker(time.Second)
	defer tick.Stop()
	for {
		select {
		case now := <-tick.C:
			file.statmu.Lock()
			// Add average of last second.
			elapsed := now.Sub(file.lpTime).Seconds()
			avg := float64(file.lpBytes) / elapsed
			file.avg.Add(avg)
			file.lpBytes = 0
			file.lpTime = now
			// Unlock stats
			file.statmu.Unlock()
		case <-file.exit:
			return
		}
	}
}

// Read bytes from the object - see io.Reader
func (file *Account) Read(p []byte) (n int, err error) {
	file.mu.Lock()
	defer file.mu.Unlock()

	// Set start time.
	file.statmu.Lock()
	if file.start.IsZero() {
		file.start = time.Now()
	}
	file.statmu.Unlock()

	n, err = file.in.Read(p)

	// Update Stats
	file.statmu.Lock()
	file.lpBytes += n
	file.bytes += int64(n)
	file.statmu.Unlock()

	Stats.Bytes(int64(n))

	// Limit the transfer speed if required
	if tokenBucket != nil {
		tokenBucket.Wait(int64(n))
	}
	return
}

// Progress returns bytes read as well as the size.
// Size can be <= 0 if the size is unknown.
func (file *Account) Progress() (bytes, size int64) {
	if file == nil {
		return 0, 0
	}
	file.statmu.Lock()
	if bytes > size {
		size = 0
	}
	defer file.statmu.Unlock()
	return file.bytes, file.size
}

// Speed returns the speed of the current file transfer
// in bytes per second, as well a an exponentially weighted moving average
// If no read has completed yet, 0 is returned for both values.
func (file *Account) Speed() (bps, current float64) {
	if file == nil {
		return 0, 0
	}
	file.statmu.Lock()
	defer file.statmu.Unlock()
	if file.bytes == 0 {
		return 0, 0
	}
	// Calculate speed from first read.
	total := float64(time.Now().Sub(file.start)) / float64(time.Second)
	bps = float64(file.bytes) / total
	current = file.avg.Value()
	return
}

// ETA returns the ETA of the current operation,
// rounded to full seconds.
// If the ETA cannot be determined 'ok' returns false.
func (file *Account) ETA() (eta time.Duration, ok bool) {
	if file == nil || file.size <= 0 {
		return 0, false
	}
	file.statmu.Lock()
	defer file.statmu.Unlock()
	if file.bytes == 0 {
		return 0, false
	}
	left := file.size - file.bytes
	if left <= 0 {
		return 0, true
	}
	avg := file.avg.Value()
	if avg <= 0 {
		return 0, false
	}
	seconds := float64(left) / file.avg.Value()

	return time.Duration(time.Second * time.Duration(int(seconds))), true
}

// String produces stats for this file
func (file *Account) String() string {
	a, b := file.Progress()
	avg, cur := file.Speed()
	eta, etaok := file.ETA()
	etas := "-"
	if etaok {
		if eta > 0 {
			etas = fmt.Sprintf("%v", eta)
		} else {
			etas = "0s"
		}
	}
	name := []rune(file.name)
	if len(name) > 45 {
		where := len(name) - 42
		name = append([]rune{'.', '.', '.'}, name[where:]...)
	}
	if b <= 0 {
		return fmt.Sprintf("%45s: avg:%7.1f, cur: %6.1f kByte/s. ETA: %s", string(name), avg/1024, cur/1024, etas)
	}
	return fmt.Sprintf("%45s: %2d%% done. avg: %6.1f, cur: %6.1f kByte/s. ETA: %s", string(name), int(100*float64(a)/float64(b)), avg/1024, cur/1024, etas)
}

// Close the object
func (file *Account) Close() error {
	file.mu.Lock()
	defer file.mu.Unlock()
	if file.closed {
		return nil
	}
	file.closed = true
	close(file.exit)
	Stats.inProgress.clear(file.name)
	return file.in.Close()
}

// Check it satisfies the interface
var _ io.ReadCloser = &Account{}
<commit_msg>Display the transfer stats in more human readable form - fixes #428<commit_after>// Accounting and limiting reader

package fs

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/VividCortex/ewma"
	"github.com/tsenart/tb"
)

// Globals
var (
	Stats       = NewStats()
	tokenBucket *tb.Bucket
)

// Start the token bucket if necessary
func startTokenBucket() {
	if bwLimit > 0 {
		tokenBucket = tb.NewBucket(int64(bwLimit), 100*time.Millisecond)
		Log(nil, "Starting bandwidth limiter at %vBytes/s", &bwLimit)
	}
}

// stringSet holds a set of strings
type stringSet map[string]struct{}

// inProgress holds a synchronizes map of in progress transfers
type inProgress struct {
	mu sync.Mutex
	m  map[string]*Account
}

// newInProgress makes a new inProgress object
func newInProgress() *inProgress {
	return &inProgress{
		m: make(map[string]*Account, Config.Transfers),
	}
}

// set marks the name as in progress
func (ip *inProgress) set(name string, acc *Account) {
	ip.mu.Lock()
	defer ip.mu.Unlock()
	ip.m[name] = acc
}

// clear marks the name as no longer in progress
func (ip *inProgress) clear(name string) {
	ip.mu.Lock()
	defer ip.mu.Unlock()
	delete(ip.m, name)
}

// get gets the account for name, of nil if not found
func (ip *inProgress) get(name string) *Account {
	ip.mu.Lock()
	defer ip.mu.Unlock()
	return ip.m[name]
}

// Strings returns all the strings in the stringSet
func (ss stringSet) Strings() []string {
	strings := make([]string, 0, len(ss))
	for name := range ss {
		var out string
		if acc := Stats.inProgress.get(name); acc != nil {
			out = acc.String()
		} else {
			out = name
		}
		strings = append(strings, " * "+out)
	}
	sorted := sort.StringSlice(strings)
	sorted.Sort()
	return sorted
}

// String returns all the file names in the stringSet joined by newline
func (ss stringSet) String() string {
	return strings.Join(ss.Strings(), "\n")
}

// StatsInfo limits and accounts all transfers
type StatsInfo struct {
	lock         sync.RWMutex
	bytes        int64
	errors       int64
	checks       int64
	checking     stringSet
	transfers    int64
	transferring stringSet
	start        time.Time
	inProgress   *inProgress
}

// NewStats cretates an initialised StatsInfo
func NewStats() *StatsInfo {
	return &StatsInfo{
		checking:     make(stringSet, Config.Checkers),
		transferring: make(stringSet, Config.Transfers),
		start:        time.Now(),
		inProgress:   newInProgress(),
	}
}

// String convert the StatsInfo to a string for printing
func (s *StatsInfo) String() string {
	s.lock.RLock()
	defer s.lock.RUnlock()
	dt := time.Now().Sub(s.start)
	dtSeconds := dt.Seconds()
	speed := 0.0
	if dt > 0 {
		speed = float64(s.bytes) / 1024 / dtSeconds
	}
	dtRounded := dt - (dt % (time.Second / 10))
	buf := &bytes.Buffer{}
	fmt.Fprintf(buf, `
Transferred:   %10vBytes (%vByte/s)
Errors:        %10d
Checks:        %10d
Transferred:   %10d
Elapsed time:  %10v
`,
		SizeSuffix(s.bytes), SizeSuffix(speed),
		s.errors,
		s.checks,
		s.transfers,
		dtRounded)
	if len(s.checking) > 0 {
		fmt.Fprintf(buf, "Checking:\n%s\n", s.checking)
	}
	if len(s.transferring) > 0 {
		fmt.Fprintf(buf, "Transferring:\n%s\n", s.transferring)
	}
	return buf.String()
}

// Log outputs the StatsInfo to the log
func (s *StatsInfo) Log() {
	log.Printf("%v\n", s)
}

// Bytes updates the stats for bytes bytes
func (s *StatsInfo) Bytes(bytes int64) {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.bytes += bytes
}

// Errors updates the stats for errors
func (s *StatsInfo) Errors(errors int64) {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.errors += errors
}

// GetErrors reads the number of errors
func (s *StatsInfo) GetErrors() int64 {
	s.lock.RLock()
	defer s.lock.RUnlock()
	return s.errors
}

// ResetCounters sets the counters (bytes, checks, errors, transfers) to 0
func (s *StatsInfo) ResetCounters() {
	s.lock.RLock()
	defer s.lock.RUnlock()
	s.bytes = 0
	s.errors = 0
	s.checks = 0
	s.transfers = 0
}

// ResetErrors sets the errors count to 0
func (s *StatsInfo) ResetErrors() {
	s.lock.RLock()
	defer s.lock.RUnlock()
	s.errors = 0
}

// Errored returns whether there have been any errors
func (s *StatsInfo) Errored() bool {
	s.lock.RLock()
	defer s.lock.RUnlock()
	return s.errors != 0
}

// Error adds a single error into the stats
func (s *StatsInfo) Error() {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.errors++
}

// Checking adds a check into the stats
func (s *StatsInfo) Checking(o Object) {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.checking[o.Remote()] = struct{}{}
}

// DoneChecking removes a check from the stats
func (s *StatsInfo) DoneChecking(o Object) {
	s.lock.Lock()
	defer s.lock.Unlock()
	delete(s.checking, o.Remote())
	s.checks++
}

// GetTransfers reads the number of transfers
func (s *StatsInfo) GetTransfers() int64 {
	s.lock.RLock()
	defer s.lock.RUnlock()
	return s.transfers
}

// Transferring adds a transfer into the stats
func (s *StatsInfo) Transferring(o Object) {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.transferring[o.Remote()] = struct{}{}
}

// DoneTransferring removes a transfer from the stats
func (s *StatsInfo) DoneTransferring(o Object) {
	s.lock.Lock()
	defer s.lock.Unlock()
	delete(s.transferring, o.Remote())
	s.transfers++
}

// Account limits and accounts for one transfer
type Account struct {
	// The mutex is to make sure Read() and Close() aren't called
	// concurrently.  Unfortunately the persistent connection loop
	// in http transport calls Read() after Do() returns on
	// CancelRequest so this race can happen when it apparently
	// shouldn't.
	mu      sync.Mutex
	in      io.ReadCloser
	size    int64
	name    string
	statmu  sync.Mutex         // Separate mutex for stat values.
	bytes   int64              // Total number of bytes read
	start   time.Time          // Start time of first read
	lpTime  time.Time          // Time of last average measurement
	lpBytes int                // Number of bytes read since last measurement
	avg     ewma.MovingAverage // Moving average of last few measurements
	closed  bool               // set if the file is closed
	exit    chan struct{}      // channel that will be closed when transfer is finished
}

// NewAccount makes a Account reader for an object
func NewAccount(in io.ReadCloser, obj Object) *Account {
	acc := &Account{
		in:     in,
		size:   obj.Size(),
		name:   obj.Remote(),
		exit:   make(chan struct{}),
		avg:    ewma.NewMovingAverage(),
		lpTime: time.Now(),
	}
	go acc.averageLoop()
	Stats.inProgress.set(acc.name, acc)
	return acc
}

func (file *Account) averageLoop() {
	tick := time.NewTicker(time.Second)
	defer tick.Stop()
	for {
		select {
		case now := <-tick.C:
			file.statmu.Lock()
			// Add average of last second.
			elapsed := now.Sub(file.lpTime).Seconds()
			avg := float64(file.lpBytes) / elapsed
			file.avg.Add(avg)
			file.lpBytes = 0
			file.lpTime = now
			// Unlock stats
			file.statmu.Unlock()
		case <-file.exit:
			return
		}
	}
}

// Read bytes from the object - see io.Reader
func (file *Account) Read(p []byte) (n int, err error) {
	file.mu.Lock()
	defer file.mu.Unlock()

	// Set start time.
	file.statmu.Lock()
	if file.start.IsZero() {
		file.start = time.Now()
	}
	file.statmu.Unlock()

	n, err = file.in.Read(p)

	// Update Stats
	file.statmu.Lock()
	file.lpBytes += n
	file.bytes += int64(n)
	file.statmu.Unlock()

	Stats.Bytes(int64(n))

	// Limit the transfer speed if required
	if tokenBucket != nil {
		tokenBucket.Wait(int64(n))
	}
	return
}

// Progress returns bytes read as well as the size.
// Size can be <= 0 if the size is unknown.
func (file *Account) Progress() (bytes, size int64) {
	if file == nil {
		return 0, 0
	}
	file.statmu.Lock()
	if bytes > size {
		size = 0
	}
	defer file.statmu.Unlock()
	return file.bytes, file.size
}

// Speed returns the speed of the current file transfer
// in bytes per second, as well a an exponentially weighted moving average
// If no read has completed yet, 0 is returned for both values.
func (file *Account) Speed() (bps, current float64) {
	if file == nil {
		return 0, 0
	}
	file.statmu.Lock()
	defer file.statmu.Unlock()
	if file.bytes == 0 {
		return 0, 0
	}
	// Calculate speed from first read.
	total := float64(time.Now().Sub(file.start)) / float64(time.Second)
	bps = float64(file.bytes) / total
	current = file.avg.Value()
	return
}

// ETA returns the ETA of the current operation,
// rounded to full seconds.
// If the ETA cannot be determined 'ok' returns false.
func (file *Account) ETA() (eta time.Duration, ok bool) {
	if file == nil || file.size <= 0 {
		return 0, false
	}
	file.statmu.Lock()
	defer file.statmu.Unlock()
	if file.bytes == 0 {
		return 0, false
	}
	left := file.size - file.bytes
	if left <= 0 {
		return 0, true
	}
	avg := file.avg.Value()
	if avg <= 0 {
		return 0, false
	}
	seconds := float64(left) / file.avg.Value()

	return time.Duration(time.Second * time.Duration(int(seconds))), true
}

// String produces stats for this file
func (file *Account) String() string {
	a, b := file.Progress()
	avg, cur := file.Speed()
	eta, etaok := file.ETA()
	etas := "-"
	if etaok {
		if eta > 0 {
			etas = fmt.Sprintf("%v", eta)
		} else {
			etas = "0s"
		}
	}
	name := []rune(file.name)
	if len(name) > 45 {
		where := len(name) - 42
		name = append([]rune{'.', '.', '.'}, name[where:]...)
	}
	if b <= 0 {
		return fmt.Sprintf("%45s: avg:%7.1f, cur: %6.1f kByte/s. ETA: %s", string(name), avg/1024, cur/1024, etas)
	}
	return fmt.Sprintf("%45s: %2d%% done. avg: %6.1f, cur: %6.1f kByte/s. ETA: %s", string(name), int(100*float64(a)/float64(b)), avg/1024, cur/1024, etas)
}

// Close the object
func (file *Account) Close() error {
	file.mu.Lock()
	defer file.mu.Unlock()
	if file.closed {
		return nil
	}
	file.closed = true
	close(file.exit)
	Stats.inProgress.clear(file.name)
	return file.in.Close()
}

// Check it satisfies the interface
var _ io.ReadCloser = &Account{}
<|endoftext|>


<commit_before>package log

import (
	"encoding/json"
	"fmt"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/fatih/camelcase"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/sirupsen/logrus"
)

const (
	bucketFactor = 2.0
	bucketCount  = 20 // Which makes the max bucket 2^20 seconds or ~12 days in size
)

// This needs to be a global var, not a field on the logger, because multiple servers
// create new loggers, and the prometheus registration uses a global namespace
var reportMetricGauge prometheus.Gauge
var reportMetricsOnce sync.Once

// Logger is a helper for emitting our grpc API logs
type Logger interface {
	Log(request interface{}, response interface{}, err error, duration time.Duration)
	LogAtLevelFromDepth(request interface{}, response interface{}, err error, duration time.Duration, level logrus.Level, depth int)
}

type logger struct {
	*logrus.Entry
	histogram   map[string]*prometheus.HistogramVec
	counter     map[string]prometheus.Counter
	mutex       *sync.Mutex
	exportStats bool
}

// NewLogger creates a new logger
func NewLogger(service string) Logger {
	return newLogger(service, true)
}

// NewLogger creates a new logger for local testing (which does not report prometheus metrics)
func NewLocalLogger(service string) Logger {
	return newLogger(service, false)
}

func newLogger(service string, exportStats bool) Logger {
	l := logrus.New()
	l.Formatter = new(prettyFormatter)
	newLogger := &logger{
		l.WithFields(logrus.Fields{"service": service}),
		make(map[string]*prometheus.HistogramVec),
		make(map[string]prometheus.Counter),
		&sync.Mutex{},
		exportStats,
	}
	if exportStats {
		reportMetricsOnce.Do(func() {
			newReportMetricGauge := prometheus.NewGauge(
				prometheus.GaugeOpts{
					Namespace: "pachyderm",
					Subsystem: "pachd",
					Name:      "report_metric",
					Help:      "gauge of number of calls to ReportMetric()",
				},
			)
			if err := prometheus.Register(newReportMetricGauge); err != nil {
				entry := newLogger.WithFields(logrus.Fields{"method": "NewLogger"})
				newLogger.LogAtLevel(entry, logrus.WarnLevel, "error registering prometheus metric: %v", err)
			} else {
				reportMetricGauge = newReportMetricGauge
			}
		})
	}
	return newLogger
}

// Helper function used to log requests and responses from our GRPC method
// implementations
func (l *logger) Log(request interface{}, response interface{}, err error, duration time.Duration) {
	if err != nil {
		l.LogAtLevelFromDepth(request, response, err, duration, logrus.ErrorLevel, 4)
	} else {
		l.LogAtLevelFromDepth(request, response, err, duration, logrus.InfoLevel, 4)
	}
	// We have to grab the method's name here before we
	// enter the goro's stack
	go l.ReportMetric(getMethodName(), duration, err)
}

func getMethodName() string {
	depth := 4
	pc := make([]uintptr, depth)
	runtime.Callers(depth, pc)
	split := strings.Split(runtime.FuncForPC(pc[0]).Name(), ".")
	return split[len(split)-1]
}

func (l *logger) ReportMetric(method string, duration time.Duration, err error) {
	if !l.exportStats {
		return
	}
	// Count the number of ReportMetric() goros in case we start to leak them
	if reportMetricGauge != nil {
		reportMetricGauge.Inc()
	}
	defer func() {
		if reportMetricGauge != nil {
			reportMetricGauge.Dec()
		}
	}()
	l.mutex.Lock() // for conccurent map access (histogram,counter)
	defer l.mutex.Unlock()
	state := "started"
	if err != nil {
		state = "errored"
	} else {
		if duration.Seconds() > 0 {
			state = "finished"
		}
	}
	entry := l.WithFields(logrus.Fields{"method": method})

	var tokens []string
	for _, token := range camelcase.Split(method) {
		tokens = append(tokens, strings.ToLower(token))
	}
	rootStatName := strings.Join(tokens, "_")

	// Recording the distribution of started times is meaningless
	if state != "started" {
		runTimeName := fmt.Sprintf("%v_time", rootStatName)
		runTime, ok := l.histogram[runTimeName]
		if !ok {
			runTime = prometheus.NewHistogramVec(
				prometheus.HistogramOpts{
					Namespace: "pachyderm",
					Subsystem: "pachd",
					Name:      runTimeName,
					Help:      fmt.Sprintf("Run time of %v", method),
					Buckets:   prometheus.ExponentialBuckets(1.0, bucketFactor, bucketCount),
				},
				[]string{
					"state", // Since both finished and errored API calls can have run times
				},
			)
			if err := prometheus.Register(runTime); err != nil {
				l.LogAtLevel(entry, logrus.WarnLevel, "error registering prometheus metric: %v", err)
			} else {
				l.histogram[runTimeName] = runTime
			}
		}
		if hist, err := runTime.GetMetricWithLabelValues(state); err != nil {
			l.LogAtLevel(entry, logrus.WarnLevel, "failed to get histogram w labels: state (%v) with error %v", state, err)
		} else {
			hist.Observe(duration.Seconds())
		}
	}

	secondsCountName := fmt.Sprintf("%v_seconds_count", rootStatName)
	secondsCount, ok := l.counter[secondsCountName]
	if !ok {
		secondsCount = prometheus.NewCounter(
			prometheus.CounterOpts{
				Namespace: "pachyderm",
				Subsystem: "pachd",
				Name:      secondsCountName,
				Help:      fmt.Sprintf("cumulative number of seconds spent in %v", method),
			},
		)
		if err := prometheus.Register(secondsCount); err != nil {
			l.LogAtLevel(entry, logrus.WarnLevel, "error registering prometheus metric: %v", err)
		} else {
			l.counter[secondsCountName] = secondsCount
		}
	}
	secondsCount.Add(duration.Seconds())

}

func (l *logger) LogAtLevel(entry *logrus.Entry, level logrus.Level, args ...interface{}) {
	switch level {
	case logrus.PanicLevel:
		entry.Panic(args)
	case logrus.FatalLevel:
		entry.Fatal(args)
	case logrus.ErrorLevel:
		entry.Error(args)
	case logrus.WarnLevel:
		entry.Warn(args)
	case logrus.InfoLevel:
		entry.Info(args)
	case logrus.DebugLevel:
		entry.Debug(args)
	}
}

func (l *logger) LogAtLevelFromDepth(request interface{}, response interface{}, err error, duration time.Duration, level logrus.Level, depth int) {
	pc := make([]uintptr, depth)
	runtime.Callers(depth, pc)
	split := strings.Split(runtime.FuncForPC(pc[0]).Name(), ".")
	method := split[len(split)-1]

	fields := logrus.Fields{
		"method":  method,
		"request": request,
	}
	if response != nil {
		fields["response"] = response
	}
	if err != nil {
		// "err" itself might be a code or even an empty struct
		fields["error"] = err.Error()
	}
	if duration > 0 {
		fields["duration"] = duration
	}
	l.LogAtLevel(l.WithFields(fields), level)
}

type prettyFormatter struct{}

func (f *prettyFormatter) Format(entry *logrus.Entry) ([]byte, error) {
	serialized := []byte(
		fmt.Sprintf(
			"%v %v ",
			entry.Time.Format(logrus.DefaultTimestampFormat),
			strings.ToUpper(entry.Level.String()),
		),
	)
	if entry.Data["service"] != nil {
		serialized = append(serialized, []byte(fmt.Sprintf("%v.%v ", entry.Data["service"], entry.Data["method"]))...)
	}
	if len(entry.Data) > 2 {
		delete(entry.Data, "service")
		delete(entry.Data, "method")
		if entry.Data["duration"] != nil {
			entry.Data["duration"] = entry.Data["duration"].(time.Duration).Seconds()
		}
		data, err := json.Marshal(entry.Data)
		if err != nil {
			return nil, fmt.Errorf("Failed to marshal fields to JSON, %v", err)
		}
		serialized = append(serialized, []byte(string(data))...)
		serialized = append(serialized, ' ')
	}

	serialized = append(serialized, []byte(entry.Message)...)
	serialized = append(serialized, '\n')
	return serialized, nil
}
<commit_msg>Make linter happy<commit_after>package log

import (
	"encoding/json"
	"fmt"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/fatih/camelcase"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/sirupsen/logrus"
)

const (
	bucketFactor = 2.0
	bucketCount  = 20 // Which makes the max bucket 2^20 seconds or ~12 days in size
)

// This needs to be a global var, not a field on the logger, because multiple servers
// create new loggers, and the prometheus registration uses a global namespace
var reportMetricGauge prometheus.Gauge
var reportMetricsOnce sync.Once

// Logger is a helper for emitting our grpc API logs
type Logger interface {
	Log(request interface{}, response interface{}, err error, duration time.Duration)
	LogAtLevelFromDepth(request interface{}, response interface{}, err error, duration time.Duration, level logrus.Level, depth int)
}

type logger struct {
	*logrus.Entry
	histogram   map[string]*prometheus.HistogramVec
	counter     map[string]prometheus.Counter
	mutex       *sync.Mutex
	exportStats bool
}

// NewLogger creates a new logger
func NewLogger(service string) Logger {
	return newLogger(service, true)
}

// NewLocalLogger creates a new logger for local testing (which does not report prometheus metrics)
func NewLocalLogger(service string) Logger {
	return newLogger(service, false)
}

func newLogger(service string, exportStats bool) Logger {
	l := logrus.New()
	l.Formatter = new(prettyFormatter)
	newLogger := &logger{
		l.WithFields(logrus.Fields{"service": service}),
		make(map[string]*prometheus.HistogramVec),
		make(map[string]prometheus.Counter),
		&sync.Mutex{},
		exportStats,
	}
	if exportStats {
		reportMetricsOnce.Do(func() {
			newReportMetricGauge := prometheus.NewGauge(
				prometheus.GaugeOpts{
					Namespace: "pachyderm",
					Subsystem: "pachd",
					Name:      "report_metric",
					Help:      "gauge of number of calls to ReportMetric()",
				},
			)
			if err := prometheus.Register(newReportMetricGauge); err != nil {
				entry := newLogger.WithFields(logrus.Fields{"method": "NewLogger"})
				newLogger.LogAtLevel(entry, logrus.WarnLevel, "error registering prometheus metric: %v", err)
			} else {
				reportMetricGauge = newReportMetricGauge
			}
		})
	}
	return newLogger
}

// Helper function used to log requests and responses from our GRPC method
// implementations
func (l *logger) Log(request interface{}, response interface{}, err error, duration time.Duration) {
	if err != nil {
		l.LogAtLevelFromDepth(request, response, err, duration, logrus.ErrorLevel, 4)
	} else {
		l.LogAtLevelFromDepth(request, response, err, duration, logrus.InfoLevel, 4)
	}
	// We have to grab the method's name here before we
	// enter the goro's stack
	go l.ReportMetric(getMethodName(), duration, err)
}

func getMethodName() string {
	depth := 4
	pc := make([]uintptr, depth)
	runtime.Callers(depth, pc)
	split := strings.Split(runtime.FuncForPC(pc[0]).Name(), ".")
	return split[len(split)-1]
}

func (l *logger) ReportMetric(method string, duration time.Duration, err error) {
	if !l.exportStats {
		return
	}
	// Count the number of ReportMetric() goros in case we start to leak them
	if reportMetricGauge != nil {
		reportMetricGauge.Inc()
	}
	defer func() {
		if reportMetricGauge != nil {
			reportMetricGauge.Dec()
		}
	}()
	l.mutex.Lock() // for conccurent map access (histogram,counter)
	defer l.mutex.Unlock()
	state := "started"
	if err != nil {
		state = "errored"
	} else {
		if duration.Seconds() > 0 {
			state = "finished"
		}
	}
	entry := l.WithFields(logrus.Fields{"method": method})

	var tokens []string
	for _, token := range camelcase.Split(method) {
		tokens = append(tokens, strings.ToLower(token))
	}
	rootStatName := strings.Join(tokens, "_")

	// Recording the distribution of started times is meaningless
	if state != "started" {
		runTimeName := fmt.Sprintf("%v_time", rootStatName)
		runTime, ok := l.histogram[runTimeName]
		if !ok {
			runTime = prometheus.NewHistogramVec(
				prometheus.HistogramOpts{
					Namespace: "pachyderm",
					Subsystem: "pachd",
					Name:      runTimeName,
					Help:      fmt.Sprintf("Run time of %v", method),
					Buckets:   prometheus.ExponentialBuckets(1.0, bucketFactor, bucketCount),
				},
				[]string{
					"state", // Since both finished and errored API calls can have run times
				},
			)
			if err := prometheus.Register(runTime); err != nil {
				l.LogAtLevel(entry, logrus.WarnLevel, "error registering prometheus metric: %v", err)
			} else {
				l.histogram[runTimeName] = runTime
			}
		}
		if hist, err := runTime.GetMetricWithLabelValues(state); err != nil {
			l.LogAtLevel(entry, logrus.WarnLevel, "failed to get histogram w labels: state (%v) with error %v", state, err)
		} else {
			hist.Observe(duration.Seconds())
		}
	}

	secondsCountName := fmt.Sprintf("%v_seconds_count", rootStatName)
	secondsCount, ok := l.counter[secondsCountName]
	if !ok {
		secondsCount = prometheus.NewCounter(
			prometheus.CounterOpts{
				Namespace: "pachyderm",
				Subsystem: "pachd",
				Name:      secondsCountName,
				Help:      fmt.Sprintf("cumulative number of seconds spent in %v", method),
			},
		)
		if err := prometheus.Register(secondsCount); err != nil {
			l.LogAtLevel(entry, logrus.WarnLevel, "error registering prometheus metric: %v", err)
		} else {
			l.counter[secondsCountName] = secondsCount
		}
	}
	secondsCount.Add(duration.Seconds())

}

func (l *logger) LogAtLevel(entry *logrus.Entry, level logrus.Level, args ...interface{}) {
	switch level {
	case logrus.PanicLevel:
		entry.Panic(args)
	case logrus.FatalLevel:
		entry.Fatal(args)
	case logrus.ErrorLevel:
		entry.Error(args)
	case logrus.WarnLevel:
		entry.Warn(args)
	case logrus.InfoLevel:
		entry.Info(args)
	case logrus.DebugLevel:
		entry.Debug(args)
	}
}

func (l *logger) LogAtLevelFromDepth(request interface{}, response interface{}, err error, duration time.Duration, level logrus.Level, depth int) {
	pc := make([]uintptr, depth)
	runtime.Callers(depth, pc)
	split := strings.Split(runtime.FuncForPC(pc[0]).Name(), ".")
	method := split[len(split)-1]

	fields := logrus.Fields{
		"method":  method,
		"request": request,
	}
	if response != nil {
		fields["response"] = response
	}
	if err != nil {
		// "err" itself might be a code or even an empty struct
		fields["error"] = err.Error()
	}
	if duration > 0 {
		fields["duration"] = duration
	}
	l.LogAtLevel(l.WithFields(fields), level)
}

type prettyFormatter struct{}

func (f *prettyFormatter) Format(entry *logrus.Entry) ([]byte, error) {
	serialized := []byte(
		fmt.Sprintf(
			"%v %v ",
			entry.Time.Format(logrus.DefaultTimestampFormat),
			strings.ToUpper(entry.Level.String()),
		),
	)
	if entry.Data["service"] != nil {
		serialized = append(serialized, []byte(fmt.Sprintf("%v.%v ", entry.Data["service"], entry.Data["method"]))...)
	}
	if len(entry.Data) > 2 {
		delete(entry.Data, "service")
		delete(entry.Data, "method")
		if entry.Data["duration"] != nil {
			entry.Data["duration"] = entry.Data["duration"].(time.Duration).Seconds()
		}
		data, err := json.Marshal(entry.Data)
		if err != nil {
			return nil, fmt.Errorf("Failed to marshal fields to JSON, %v", err)
		}
		serialized = append(serialized, []byte(string(data))...)
		serialized = append(serialized, ' ')
	}

	serialized = append(serialized, []byte(entry.Message)...)
	serialized = append(serialized, '\n')
	return serialized, nil
}
<|endoftext|>


<commit_before>package bridge

import (
	"sync"

	"github.com/x-cray/marathon-service-registrator/consul"
	"github.com/x-cray/marathon-service-registrator/marathon"
	"github.com/x-cray/marathon-service-registrator/types"

	log "github.com/Sirupsen/logrus"
)

type serviceGroupPair struct {
	service *types.Service
	group   *types.ServiceGroup
}

type Bridge struct {
	sync.Mutex

	scheduler              types.SchedulerAdapter
	schedulerServiceGroups map[string]*types.ServiceGroup
	registry               types.RegistryAdapter
	registryAdvertizeAddr  string
	config                 *types.Config
}

func New(c *types.Config) (*Bridge, error) {
	marathon, err := marathon.New(c.Marathon)
	if err != nil {
		return nil, err
	}

	consul, err := consul.New(c.Consul, c.DryRun)
	if err != nil {
		return nil, err
	}

	return &Bridge{
		config:    c,
		scheduler: marathon,
		registry:  consul,
	}, nil
}

func (b *Bridge) cachedServiceGroup(groupID, actionText string) *types.ServiceGroup {
	if group, ok := b.schedulerServiceGroups[groupID]; ok {
		return group
	}

	log.WithField("prefix", "bridge").Warningf(
		"Service group %s was not found in scheduler cache (has %d entries). Could not %s.",
		groupID,
		len(b.schedulerServiceGroups),
		actionText,
	)
	return nil
}

func logSkipMessage(ip string) {
	log.WithFields(log.Fields{
		"prefix": "bridge",
		"ip":     ip,
	}).Debug("Skipping event due to unrelated service host")
}

func (b *Bridge) processServiceEvent(event *types.ServiceEvent) error {
	b.Lock()
	defer b.Unlock()

	switch event.Action {
	case types.ServiceStarted:
		// New service is started, we need to refresh service cache.
		_, err := b.refreshSchedulerServices()
		if err != nil {
			return err
		}
	case types.ServiceStopped:
		// Service stopped, deregister and remove it from cache.
		// Only consider services registered on current registry's advertized address.
		if event.IP == b.registryAdvertizeAddr {
			if group := b.cachedServiceGroup(event.ServiceID, "deregister"); group != nil {
				b.registry.Deregister(group)
				delete(b.schedulerServiceGroups, event.ServiceID)
			}
		} else {
			logSkipMessage(event.IP)
		}
	case types.ServiceWentUp:
		// Service went up, register it.
		if group := b.cachedServiceGroup(event.ServiceID, "register"); group != nil {
			// Only consider services registered on current registry's advertized address.
			if group.IP == b.registryAdvertizeAddr {
				b.registry.Register(group)
			} else {
				logSkipMessage(group.IP)
			}
		}
	case types.ServiceWentDown:
		// Service went down, deregister it.
		if group := b.cachedServiceGroup(event.ServiceID, "deregister"); group != nil {
			// Only consider services registered on current registry's advertized address.
			if group.IP == b.registryAdvertizeAddr {
				b.registry.Deregister(group)
			} else {
				logSkipMessage(group.IP)
			}
		}
	}

	return nil
}

func (b *Bridge) ProcessSchedulerEvents() error {
	schedulerEvents := make(types.EventsChannel, 5)
	err := b.scheduler.ListenForEvents(schedulerEvents)
	if err != nil {
		return err
	}

	log.WithField("prefix", "bridge").Info("Registered for scheduler event stream")
	for {
		event := <-schedulerEvents
		if event.Action != types.ServiceUnchanged {
			log.WithFields(log.Fields{
				"prefix":  "bridge",
				"service": event.ServiceID,
				"action":  event.Action,
				"event":   event.OriginalEvent,
			}).Debug("Received scheduler event")
			b.processServiceEvent(event)
		}
	}

	return nil
}

// Sync performs full synchronization of scheduler tasks to service registry.
func (b *Bridge) Sync() error {
	b.Lock()
	defer b.Unlock()

	actionsPerformed := false

	// Get services from registry.
	registryServiceGroups, err := b.registry.Services()
	if err != nil {
		return err
	}

	log.WithField("prefix", "bridge").Infof("Received %d services from registry", len(registryServiceGroups))

	// Build ip:port-indexed service map.
	registryServicesMap := make(map[string]*serviceGroupPair)
	for _, group := range registryServiceGroups {
		for _, service := range group.Services {
			registryServicesMap[group.ServiceKey(service)] = &serviceGroupPair{
				service: service,
				group:   group,
			}
		}
	}

	schedulerServicesMap, err := b.refreshSchedulerServices()
	if err != nil {
		return err
	}

	// Register scheduler services absent in registry.
	for _, schedulerService := range schedulerServicesMap {
		group := schedulerService.group
		service := schedulerService.service

		// If service is not yet registered we need to register it.
		// Only consider healthy services registered on current registry's advertized address.
		if group.IP == b.registryAdvertizeAddr && service.Healthy && registryServicesMap[group.ServiceKey(service)] == nil {
			err := b.registry.Register(group)
			if err != nil {
				return err
			}
			actionsPerformed = true
		}
	}

	// Deregister dangling services (existing in registry and absent in scheduler).
	for _, registryService := range registryServicesMap {
		group := registryService.group
		service := registryService.service

		// If service is registered and we don't have it in scheduler we need to deregister it.
		if schedulerServicesMap[group.ServiceKey(service)] == nil {
			err := b.registry.Deregister(group)
			if err != nil {
				return err
			}
			actionsPerformed = true
		}
	}

	if !actionsPerformed {
		log.WithField("prefix", "bridge").Info("All services are in sync, no actions performed")
	}

	return nil
}

func (b *Bridge) refreshSchedulerServices() (map[string]*serviceGroupPair, error) {
	log.WithField("prefix", "bridge").Info("Refreshing scheduler services")

	// Get registry's advertize address.
	addr, err := b.registry.AdvertiseAddr()
	if err != nil {
		return nil, err
	}

	b.registryAdvertizeAddr = addr

	// Get services from scheduler.
	schedulerServiceGroups, err := b.scheduler.Services()
	if err != nil {
		return nil, err
	}

	log.WithField("prefix", "bridge").Infof("Registry advertize address is %s", b.registryAdvertizeAddr)

	// Build 2 maps of services:
	// ServiceID-indexed and ip:port-indexed
	ipPortServices := make(map[string]*serviceGroupPair)
	b.schedulerServiceGroups = make(map[string]*types.ServiceGroup)
	for _, group := range schedulerServiceGroups {
		b.schedulerServiceGroups[group.ID] = group
		for _, service := range group.Services {
			ipPortServices[group.ServiceKey(service)] = &serviceGroupPair{
				service: service,
				group:   group,
			}
		}
	}

	log.WithField("prefix", "bridge").Infof(
		"Received %d services from scheduler",
		len(schedulerServiceGroups),
	)

	return ipPortServices, nil
}
<commit_msg>Fix some English<commit_after>package bridge

import (
	"sync"

	"github.com/x-cray/marathon-service-registrator/consul"
	"github.com/x-cray/marathon-service-registrator/marathon"
	"github.com/x-cray/marathon-service-registrator/types"

	log "github.com/Sirupsen/logrus"
)

type serviceGroupPair struct {
	service *types.Service
	group   *types.ServiceGroup
}

type Bridge struct {
	sync.Mutex

	scheduler              types.SchedulerAdapter
	schedulerServiceGroups map[string]*types.ServiceGroup
	registry               types.RegistryAdapter
	registryAdvertizeAddr  string
	config                 *types.Config
}

func New(c *types.Config) (*Bridge, error) {
	marathon, err := marathon.New(c.Marathon)
	if err != nil {
		return nil, err
	}

	consul, err := consul.New(c.Consul, c.DryRun)
	if err != nil {
		return nil, err
	}

	return &Bridge{
		config:    c,
		scheduler: marathon,
		registry:  consul,
	}, nil
}

func (b *Bridge) cachedServiceGroup(groupID, actionText string) *types.ServiceGroup {
	if group, ok := b.schedulerServiceGroups[groupID]; ok {
		return group
	}

	log.WithField("prefix", "bridge").Warningf(
		"Service group %s was not found in scheduler cache (has %d entries). Could not %s.",
		groupID,
		len(b.schedulerServiceGroups),
		actionText,
	)
	return nil
}

func logSkipMessage(ip string) {
	log.WithFields(log.Fields{
		"prefix": "bridge",
		"ip":     ip,
	}).Debug("Skipping event due to unrelated service host")
}

func (b *Bridge) processServiceEvent(event *types.ServiceEvent) error {
	b.Lock()
	defer b.Unlock()

	switch event.Action {
	case types.ServiceStarted:
		// New service is started, we need to refresh service cache.
		_, err := b.refreshSchedulerServices()
		if err != nil {
			return err
		}
	case types.ServiceStopped:
		// Service stopped, deregister and remove it from cache.
		// Only consider services registered on current registry's advertized address.
		if event.IP == b.registryAdvertizeAddr {
			if group := b.cachedServiceGroup(event.ServiceID, "deregister"); group != nil {
				b.registry.Deregister(group)
				delete(b.schedulerServiceGroups, event.ServiceID)
			}
		} else {
			logSkipMessage(event.IP)
		}
	case types.ServiceWentUp:
		// Service went up, register it.
		if group := b.cachedServiceGroup(event.ServiceID, "register"); group != nil {
			// Only consider services registered on current registry's advertized address.
			if group.IP == b.registryAdvertizeAddr {
				b.registry.Register(group)
			} else {
				logSkipMessage(group.IP)
			}
		}
	case types.ServiceWentDown:
		// Service went down, deregister it.
		if group := b.cachedServiceGroup(event.ServiceID, "deregister"); group != nil {
			// Only consider services registered on current registry's advertized address.
			if group.IP == b.registryAdvertizeAddr {
				b.registry.Deregister(group)
			} else {
				logSkipMessage(group.IP)
			}
		}
	}

	return nil
}

func (b *Bridge) ProcessSchedulerEvents() error {
	schedulerEvents := make(types.EventsChannel, 5)
	err := b.scheduler.ListenForEvents(schedulerEvents)
	if err != nil {
		return err
	}

	log.WithField("prefix", "bridge").Info("Registered for scheduler event stream")
	for {
		event := <-schedulerEvents
		if event.Action != types.ServiceUnchanged {
			log.WithFields(log.Fields{
				"prefix":  "bridge",
				"service": event.ServiceID,
				"action":  event.Action,
				"event":   event.OriginalEvent,
			}).Debug("Received scheduler event")
			b.processServiceEvent(event)
		}
	}

	return nil
}

// Sync performs full synchronization of scheduler tasks to service registry.
func (b *Bridge) Sync() error {
	b.Lock()
	defer b.Unlock()

	actionsPerformed := false

	// Get services from registry.
	registryServiceGroups, err := b.registry.Services()
	if err != nil {
		return err
	}

	log.WithField("prefix", "bridge").Infof("Received %d services from registry", len(registryServiceGroups))

	// Build ip:port-indexed service map.
	registryServicesMap := make(map[string]*serviceGroupPair)
	for _, group := range registryServiceGroups {
		for _, service := range group.Services {
			registryServicesMap[group.ServiceKey(service)] = &serviceGroupPair{
				service: service,
				group:   group,
			}
		}
	}

	schedulerServicesMap, err := b.refreshSchedulerServices()
	if err != nil {
		return err
	}

	// Register scheduler services absent from registry.
	for _, schedulerService := range schedulerServicesMap {
		group := schedulerService.group
		service := schedulerService.service

		// If service is not yet registered we need to register it.
		// Only consider healthy services registered on current registry's advertized address.
		if group.IP == b.registryAdvertizeAddr && service.Healthy && registryServicesMap[group.ServiceKey(service)] == nil {
			err := b.registry.Register(group)
			if err != nil {
				return err
			}
			actionsPerformed = true
		}
	}

	// Deregister dangling services (existing in registry but absent from scheduler).
	for _, registryService := range registryServicesMap {
		group := registryService.group
		service := registryService.service

		// If service is registered and we don't have it in scheduler we need to deregister it.
		if schedulerServicesMap[group.ServiceKey(service)] == nil {
			err := b.registry.Deregister(group)
			if err != nil {
				return err
			}
			actionsPerformed = true
		}
	}

	if !actionsPerformed {
		log.WithField("prefix", "bridge").Info("All services are in sync, no actions performed")
	}

	return nil
}

func (b *Bridge) refreshSchedulerServices() (map[string]*serviceGroupPair, error) {
	log.WithField("prefix", "bridge").Info("Refreshing scheduler services")

	// Get registry's advertize address.
	addr, err := b.registry.AdvertiseAddr()
	if err != nil {
		return nil, err
	}

	b.registryAdvertizeAddr = addr

	// Get services from scheduler.
	schedulerServiceGroups, err := b.scheduler.Services()
	if err != nil {
		return nil, err
	}

	log.WithField("prefix", "bridge").Infof("Registry advertize address is %s", b.registryAdvertizeAddr)

	// Build 2 maps of services:
	// ServiceID-indexed and ip:port-indexed
	ipPortServices := make(map[string]*serviceGroupPair)
	b.schedulerServiceGroups = make(map[string]*types.ServiceGroup)
	for _, group := range schedulerServiceGroups {
		b.schedulerServiceGroups[group.ID] = group
		for _, service := range group.Services {
			ipPortServices[group.ServiceKey(service)] = &serviceGroupPair{
				service: service,
				group:   group,
			}
		}
	}

	log.WithField("prefix", "bridge").Infof(
		"Received %d services from scheduler",
		len(schedulerServiceGroups),
	)

	return ipPortServices, nil
}
<|endoftext|>


<commit_before>package http2

import (
	"fmt"
	hpack "github.com/ami-GS/GoHPACK"
)

type Frame interface {
	Pack()
	Parse(data []byte)
	GetWire() []byte
}

type Http2Header struct {
	Length     uint32
	Type, Flag byte
	StreamID   uint32
	HeadWire   []byte
}

func NewHttp2Header(length uint32, fType, flag byte, streamID uint32) *Http2Header {
	h := Http2Header{length, fType, flag, streamID, []byte{}}
	h.Pack()
	return &h
}

func (self *Http2Header) Pack() {
	self.HeadWire = make([]byte, 9)
	for i := 0; i < 3; i++ {
		self.HeadWire[i] = byte(self.Length >> byte((2-i)*8))
	}
	self.HeadWire[3], self.HeadWire[4] = self.Type, self.Flag
	for i := 0; i < 4; i++ {
		self.HeadWire[i+5] = byte(self.StreamID >> byte((3-i)*8))
	}
}

func (self *Http2Header) Parse(data []byte) {
	self.Length = uint32(data[0])<<16 | uint32(data[1])<<8 | uint32(data[2])
	self.Type = data[3]
	self.Flag = data[4]
	self.StreamID = uint32(data[5])<<24 | uint32(data[6])<<16 | uint32(data[7])<<8 | uint32(data[8])
}

type Data struct {
	Header *Http2Header
	Data   string
	PadLen byte
	Wire   []byte
}

func NewData(data string, streamID uint32, flag, padLen byte) *Data {
	var length uint32 = uint32(len(data))
	if flag&FLAG_PADDED == FLAG_PADDED {
		length += uint32(padLen + 1)
	}

	header := NewHttp2Header(length, TYPE_DATA, flag, streamID)
	frame := Data{header, data, padLen, []byte{}}
	frame.Pack()

	return &frame
}

func (self *Data) Pack() {
	idx := 0
	if self.Header.Flag == FLAG_PADDED {
		self.Wire = make([]byte, len(self.Data)+int(self.PadLen+1))
		self.Wire[idx] = self.PadLen
		idx++
	} else {
		self.Wire = make([]byte, uint32(len(self.Data)))
	}
	for i, d := range self.Data {
		self.Wire[idx+i] = byte(d)
	}
}

func (self *Data) Parse(data []byte) {
	if self.Header.Flag == FLAG_PADDED {
		self.PadLen = data[0]
		self.Data = string(data[1 : self.Header.Length-uint32(self.PadLen)])
	} else {
		self.Data = string(data)
	}
}

func (self *Data) GetWire() []byte {
	return append(self.Header.HeadWire, self.Wire...)
}

type Settings struct {
	Header    *Http2Header
	SettingID uint16
	Value     uint32
	Wire      []byte
}

func NewSettings(settingID uint16, value uint32, flag byte) *Settings {
	header := NewHttp2Header(uint32(6), TYPE_SETTINGS, flag, 0)
	frame := Settings{header, settingID, value, []byte{}}
	frame.Pack()

	return &frame
}

func (self *Settings) Pack() {
	self.Wire = make([]byte, 6)
	for i := 0; i < 2; i++ {
		self.Wire[i] = byte(self.SettingID >> byte((1-i)*8))
	}
	for i := 0; i < 4; i++ {
		self.Wire[2+i] = byte(self.Value >> byte((3-i)*8))
	}
}

func (self *Settings) Parse(data []byte) {
	self.SettingID = uint16(data[0])<<8 | uint16(data[1])
	self.Value = uint32(data[2])<<24 | uint32(data[3])<<16 | uint32(data[4])<<8 | uint32(data[5])
	_ = self.Header.Flag //temporally
}

func (self *Settings) GetWire() []byte {
	return append(self.Header.HeadWire, self.Wire...)
}

type Headers struct {
	Header           *Http2Header
	Headers          []hpack.Header
	block            []byte
	PadLen, Weight   byte
	E                bool
	StreamDependency uint32
	Wire             []byte
}

func NewHeaders(headers []hpack.Header, table *hpack.Table, streamID uint32, flag, padLen, weight byte, e bool, streamDependency uint32) *Headers {
	encHeaders := hpack.Encode(headers, false, false, false, table, -1)
	var length uint32 = uint32(len(encHeaders))
	if flag&FLAG_PADDED == FLAG_PADDED {
		length += uint32(padLen + 1)
	}
	if flag&FLAG_PRIORITY == FLAG_PRIORITY {
		length += 5
	}

	header := NewHttp2Header(length, TYPE_HEADERS, flag, streamID)

	frame := Headers{header, headers, encHeaders, padLen, weight, e, streamDependency, []byte{}}
	frame.Pack()

	return &frame
}

func (self *Headers) Pack() {
	idx := 0
	if self.Header.Flag&FLAG_PADDED == FLAG_PADDED {
		self.Wire = make([]byte, int(self.PadLen+1)+len(self.block))
		self.Wire[idx] = self.PadLen
		idx++
	}
	if self.Header.Flag&FLAG_PRIORITY == FLAG_PRIORITY {
		self.Wire = make([]byte, 5+len(self.block))
		for i := 0; i < 4; i++ {
			self.Wire[i] = byte(self.StreamDependency >> byte((3-i)*8))
		}
		if self.E {
			self.Wire[0] |= 0x80
		}
		self.Wire[4] = self.Weight
		idx = 5
	}
	if self.Header.Flag&FLAG_END_HEADERS == FLAG_END_HEADERS || self.Header.Flag&FLAG_END_STREAM == FLAG_END_STREAM {
		self.Wire = make([]byte, len(self.block))
	}
	/*else {
		panic("undefined flag")
	}*/
	for i, h := range self.block {
		self.Wire[idx+i] = h
	}
}

func (self *Headers) Parse(data []byte) {
	idx := 0
	if self.Header.Flag&FLAG_PADDED == FLAG_PADDED {
		self.PadLen = data[idx]
		idx++
	}
	if self.Header.Flag&FLAG_PRIORITY == FLAG_PRIORITY {
		if data[0]&0x80 > 0 {
			self.E = true
		}
		self.StreamDependency = uint32(data[0]&0xef)<<24 | uint32(data[1])<<16 | uint32(data[2])<<8 | uint32(data[3])
		self.Weight = data[4]
		idx += 5
	}
	if self.Header.Flag&FLAG_END_HEADERS == FLAG_END_HEADERS || self.Header.Flag&FLAG_END_STREAM == FLAG_END_STREAM {
		fmt.Println("change stream state")
	}
	/*else {
		panic("undefined flag")
	}*/
}

func (self *Headers) GetWire() []byte {
	return append(self.Header.HeadWire, self.Wire...)
}

type GoAway struct {
	Header       *Http2Header
	LastStreamID uint32
	ErrorCode    uint32
	Debug        string
	Wire         []byte
}

func NewGoAway(lastStreamID, errorCode uint32, debug string) *GoAway {
	var length uint32 = uint32(len(debug) + 8)
	header := NewHttp2Header(length, TYPE_GOAWAY, FLAG_NO, 0)
	frame := GoAway{header, lastStreamID, errorCode, debug, []byte{}}
	frame.Pack()
	return &frame
}

func (self *GoAway) Pack() {
	self.Wire = make([]byte, 8+len(self.Debug))
	for i := 0; i < 4; i++ {
		self.Wire[i] = byte(self.LastStreamID >> byte((3-i)*8))
		self.Wire[i+4] = byte(self.ErrorCode >> byte((3-i)*8))
	}
	for i, d := range self.Debug {
		self.Wire[i+8] = byte(d)
	}
}

func (self *GoAway) Parse(data []byte) {
	self.LastStreamID = uint32(data[0]&0xef)<<24 | uint32(data[1])<<16 | uint32(data[2])<<8 | uint32(data[3])
	self.ErrorCode = uint32(data[4])<<24 | uint32(data[5])<<16 | uint32(data[6])<<8 | uint32(data[7])
	if len(data) >= 9 {
		self.Debug = string(data[8:])
	}
}

func (self *GoAway) GetWire() []byte {
	return append(self.Header.HeadWire, self.Wire...)
}

func main() {
	table := hpack.InitTable()
	headers := []hpack.Header{hpack.Header{":method", "GET"}, hpack.Header{":scheme", "http"},
		hpack.Header{":authority", "127.0.0.1"}, hpack.Header{":path", "/"}}

	http2Header := Http2Header{Length: 12, Type: TYPE_DATA, Flag: FLAG_PADDED, StreamID: 1}
	http2Header.Pack()
	fmt.Printf("http2Header %v\n", http2Header)
	data := Data{Data: "Hello!", PadLen: 5}
	data.Pack(http2Header.Flag)
	data2 := Data{}
	data2.Parse(data.Wire, http2Header.Flag, http2Header.Length)
	fmt.Printf("data %v\n", data)
	fmt.Printf("data2 %v\n", data2)
	settings := Settings{SettingID: 0xff00, Value: 0xff00ff00}
	settings2 := Settings{}
	settings.Pack()
	settings2.Parse(settings.Wire, http2Header.Flag)
	fmt.Printf("settings %v\n", settings)
	fmt.Printf("settings2 %v\n", settings2)
	h := Headers{Headers: headers, PadLen: 5, Weight: 0, E: false}
	h2 := Headers{}
	h.Pack(http2Header.Flag, &table)
	h2.Parse(h.Wire, http2Header.Flag, &table)
	fmt.Printf("headers %v\n", h)
	fmt.Printf("headers2 %v\n", h2)
	goaway := GoAway{LastStreamID: 0xef00ff00, ErrorCode: 0xff00ff00, Debug: "DEBUG MESSAGE!!"}
	goaway2 := GoAway{}
	goaway.Pack()
	goaway2.Parse(goaway.Wire)
	fmt.Printf("goaway %v\n", goaway)
	fmt.Printf("goaway2 %v\n", goaway2)
}
<commit_msg>temporally remove<commit_after>package http2

import (
	"fmt"
	hpack "github.com/ami-GS/GoHPACK"
)

type Frame interface {
	Pack()
	Parse(data []byte)
	GetWire() []byte
}

type Http2Header struct {
	Length     uint32
	Type, Flag byte
	StreamID   uint32
	HeadWire   []byte
}

func NewHttp2Header(length uint32, fType, flag byte, streamID uint32) *Http2Header {
	h := Http2Header{length, fType, flag, streamID, []byte{}}
	h.Pack()
	return &h
}

func (self *Http2Header) Pack() {
	self.HeadWire = make([]byte, 9)
	for i := 0; i < 3; i++ {
		self.HeadWire[i] = byte(self.Length >> byte((2-i)*8))
	}
	self.HeadWire[3], self.HeadWire[4] = self.Type, self.Flag
	for i := 0; i < 4; i++ {
		self.HeadWire[i+5] = byte(self.StreamID >> byte((3-i)*8))
	}
}

func (self *Http2Header) Parse(data []byte) {
	self.Length = uint32(data[0])<<16 | uint32(data[1])<<8 | uint32(data[2])
	self.Type = data[3]
	self.Flag = data[4]
	self.StreamID = uint32(data[5])<<24 | uint32(data[6])<<16 | uint32(data[7])<<8 | uint32(data[8])
}

type Data struct {
	Header *Http2Header
	Data   string
	PadLen byte
	Wire   []byte
}

func NewData(data string, streamID uint32, flag, padLen byte) *Data {
	var length uint32 = uint32(len(data))
	if flag&FLAG_PADDED == FLAG_PADDED {
		length += uint32(padLen + 1)
	}

	header := NewHttp2Header(length, TYPE_DATA, flag, streamID)
	frame := Data{header, data, padLen, []byte{}}
	frame.Pack()

	return &frame
}

func (self *Data) Pack() {
	idx := 0
	if self.Header.Flag == FLAG_PADDED {
		self.Wire = make([]byte, len(self.Data)+int(self.PadLen+1))
		self.Wire[idx] = self.PadLen
		idx++
	} else {
		self.Wire = make([]byte, uint32(len(self.Data)))
	}
	for i, d := range self.Data {
		self.Wire[idx+i] = byte(d)
	}
}

func (self *Data) Parse(data []byte) {
	if self.Header.Flag == FLAG_PADDED {
		self.PadLen = data[0]
		self.Data = string(data[1 : self.Header.Length-uint32(self.PadLen)])
	} else {
		self.Data = string(data)
	}
}

func (self *Data) GetWire() []byte {
	return append(self.Header.HeadWire, self.Wire...)
}

type Settings struct {
	Header    *Http2Header
	SettingID uint16
	Value     uint32
	Wire      []byte
}

func NewSettings(settingID uint16, value uint32, flag byte) *Settings {
	header := NewHttp2Header(uint32(6), TYPE_SETTINGS, flag, 0)
	frame := Settings{header, settingID, value, []byte{}}
	frame.Pack()

	return &frame
}

func (self *Settings) Pack() {
	self.Wire = make([]byte, 6)
	for i := 0; i < 2; i++ {
		self.Wire[i] = byte(self.SettingID >> byte((1-i)*8))
	}
	for i := 0; i < 4; i++ {
		self.Wire[2+i] = byte(self.Value >> byte((3-i)*8))
	}
}

func (self *Settings) Parse(data []byte) {
	self.SettingID = uint16(data[0])<<8 | uint16(data[1])
	self.Value = uint32(data[2])<<24 | uint32(data[3])<<16 | uint32(data[4])<<8 | uint32(data[5])
	_ = self.Header.Flag //temporally
}

func (self *Settings) GetWire() []byte {
	return append(self.Header.HeadWire, self.Wire...)
}

type Headers struct {
	Header           *Http2Header
	Headers          []hpack.Header
	block            []byte
	PadLen, Weight   byte
	E                bool
	StreamDependency uint32
	Wire             []byte
}

func NewHeaders(headers []hpack.Header, table *hpack.Table, streamID uint32, flag, padLen, weight byte, e bool, streamDependency uint32) *Headers {
	encHeaders := hpack.Encode(headers, false, false, false, table, -1)
	var length uint32 = uint32(len(encHeaders))
	if flag&FLAG_PADDED == FLAG_PADDED {
		length += uint32(padLen + 1)
	}
	if flag&FLAG_PRIORITY == FLAG_PRIORITY {
		length += 5
	}

	header := NewHttp2Header(length, TYPE_HEADERS, flag, streamID)

	frame := Headers{header, headers, encHeaders, padLen, weight, e, streamDependency, []byte{}}
	frame.Pack()

	return &frame
}

func (self *Headers) Pack() {
	idx := 0
	if self.Header.Flag&FLAG_PADDED == FLAG_PADDED {
		self.Wire = make([]byte, int(self.PadLen+1)+len(self.block))
		self.Wire[idx] = self.PadLen
		idx++
	}
	if self.Header.Flag&FLAG_PRIORITY == FLAG_PRIORITY {
		self.Wire = make([]byte, 5+len(self.block))
		for i := 0; i < 4; i++ {
			self.Wire[i] = byte(self.StreamDependency >> byte((3-i)*8))
		}
		if self.E {
			self.Wire[0] |= 0x80
		}
		self.Wire[4] = self.Weight
		idx = 5
	}
	if self.Header.Flag&FLAG_END_HEADERS == FLAG_END_HEADERS || self.Header.Flag&FLAG_END_STREAM == FLAG_END_STREAM {
		self.Wire = make([]byte, len(self.block))
	}
	/*else {
		panic("undefined flag")
	}*/
	for i, h := range self.block {
		self.Wire[idx+i] = h
	}
}

func (self *Headers) Parse(data []byte) {
	idx := 0
	if self.Header.Flag&FLAG_PADDED == FLAG_PADDED {
		self.PadLen = data[idx]
		idx++
	}
	if self.Header.Flag&FLAG_PRIORITY == FLAG_PRIORITY {
		if data[0]&0x80 > 0 {
			self.E = true
		}
		self.StreamDependency = uint32(data[0]&0xef)<<24 | uint32(data[1])<<16 | uint32(data[2])<<8 | uint32(data[3])
		self.Weight = data[4]
		idx += 5
	}
	if self.Header.Flag&FLAG_END_HEADERS == FLAG_END_HEADERS || self.Header.Flag&FLAG_END_STREAM == FLAG_END_STREAM {
		fmt.Println("change stream state")
	}
	/*else {
		panic("undefined flag")
	}*/
}

func (self *Headers) GetWire() []byte {
	return append(self.Header.HeadWire, self.Wire...)
}

type GoAway struct {
	Header       *Http2Header
	LastStreamID uint32
	ErrorCode    uint32
	Debug        string
	Wire         []byte
}

func NewGoAway(lastStreamID, errorCode uint32, debug string) *GoAway {
	var length uint32 = uint32(len(debug) + 8)
	header := NewHttp2Header(length, TYPE_GOAWAY, FLAG_NO, 0)
	frame := GoAway{header, lastStreamID, errorCode, debug, []byte{}}
	frame.Pack()
	return &frame
}

func (self *GoAway) Pack() {
	self.Wire = make([]byte, 8+len(self.Debug))
	for i := 0; i < 4; i++ {
		self.Wire[i] = byte(self.LastStreamID >> byte((3-i)*8))
		self.Wire[i+4] = byte(self.ErrorCode >> byte((3-i)*8))
	}
	for i, d := range self.Debug {
		self.Wire[i+8] = byte(d)
	}
}

func (self *GoAway) Parse(data []byte) {
	self.LastStreamID = uint32(data[0]&0xef)<<24 | uint32(data[1])<<16 | uint32(data[2])<<8 | uint32(data[3])
	self.ErrorCode = uint32(data[4])<<24 | uint32(data[5])<<16 | uint32(data[6])<<8 | uint32(data[7])
	if len(data) >= 9 {
		self.Debug = string(data[8:])
	}
}

func (self *GoAway) GetWire() []byte {
	return append(self.Header.HeadWire, self.Wire...)
}

/*
func main() {
	table := hpack.InitTable()
	headers := []hpack.Header{hpack.Header{":method", "GET"}, hpack.Header{":scheme", "http"},
		hpack.Header{":authority", "127.0.0.1"}, hpack.Header{":path", "/"}}

	http2Header := Http2Header{Length: 12, Type: TYPE_DATA, Flag: FLAG_PADDED, StreamID: 1}
	http2Header.Pack()
	fmt.Printf("http2Header %v\n", http2Header)
	data := Data{Data: "Hello!", PadLen: 5}
	data.Pack(http2Header.Flag)
	data2 := Data{}
	data2.Parse(data.Wire, http2Header.Flag, http2Header.Length)
	fmt.Printf("data %v\n", data)
	fmt.Printf("data2 %v\n", data2)
	settings := Settings{SettingID: 0xff00, Value: 0xff00ff00}
	settings2 := Settings{}
	settings.Pack()
	settings2.Parse(settings.Wire, http2Header.Flag)
	fmt.Printf("settings %v\n", settings)
	fmt.Printf("settings2 %v\n", settings2)
	h := Headers{Headers: headers, PadLen: 5, Weight: 0, E: false}
	h2 := Headers{}
	h.Pack(http2Header.Flag, &table)
	h2.Parse(h.Wire, http2Header.Flag, &table)
	fmt.Printf("headers %v\n", h)
	fmt.Printf("headers2 %v\n", h2)
	goaway := GoAway{LastStreamID: 0xef00ff00, ErrorCode: 0xff00ff00, Debug: "DEBUG MESSAGE!!"}
	goaway2 := GoAway{}
	goaway.Pack()
	goaway2.Parse(goaway.Wire)
	fmt.Printf("goaway %v\n", goaway)
	fmt.Printf("goaway2 %v\n", goaway2)
}
*/
<|endoftext|>


<commit_before>package transport

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"fmt"
	"io"
	"log"
	"net"
	"time"

	"github.com/Randomsock5/tcptunnel/constants"
)

type AESConn struct {
	conn net.Conn
	r    *cipher.StreamReader
	w    *cipher.StreamWriter
}

func (ac *AESConn) Read(dst []byte) (int, error) {
	n, err := ac.r.Read(dst)
	return n, err
}

func (ac *AESConn) Write(src []byte) (int, error) {
	n, err := ac.w.Write(src)
	return n, err
}

func (ac *AESConn) Close() error {
	return ac.conn.Close()
}

func (ac *AESConn) LocalAddr() net.Addr {
	return ac.conn.LocalAddr()
}

func (ac *AESConn) RemoteAddr() net.Addr {
	return ac.conn.RemoteAddr()
}

func (ac *AESConn) SetDeadline(t time.Time) error {
	return ac.conn.SetDeadline(t)
}

func (ac *AESConn) SetReadDeadline(t time.Time) error {
	return ac.conn.SetReadDeadline(t)
}

func (ac *AESConn) SetWriteDeadline(t time.Time) error {
	return ac.conn.SetWriteDeadline(t)
}

func NewAESConn(key string, iv [aes.BlockSize]byte, conn net.Conn) (*AESConn, error) {
	keyBin := []byte(key)
	t := time.Now().UTC()

	hashKey := sha256.Sum256(keyBin)
	salt := sha256.Sum256(
		[]byte(fmt.Sprintf("%d-%02d-%02d", t.Year(), t.Month(), t.Day())))

	hashSalt := append(hashKey[:], salt[:]...)

	aesKey := sha256.Sum256(hashSalt)

	rBlock, err := aes.NewCipher(aesKey[:])
	if err != nil {
		return nil, err
	}

	wBlock, err := aes.NewCipher(aesKey[:])
	if err != nil {
		return nil, err
	}

	rStream := cipher.NewOFB(rBlock, iv[:])
	wStream := cipher.NewOFB(wBlock, iv[:])

	aesConn := &AESConn{
		conn: conn,
		r:    &cipher.StreamReader{S: rStream, R: conn},
		w:    &cipher.StreamWriter{S: wStream, W: conn},
	}
	return aesConn, err
}

type Listener struct {
	listener net.Listener
	key      string
}

func (l *Listener) Accept() (net.Conn, error) {
	conn, err := l.listener.Accept()
	if err != nil {
		return nil, err
	}
	err = conn.SetReadDeadline(time.Now().Add(constants.ConnTimeout))
	if err != nil {
		return nil, err
	}

	ivVector := make([]byte, constants.IVLength)
	if _, err := io.ReadFull(conn, ivVector); err != nil {
		log.Fatal(err)
	}

	password := sha256.Sum256([]byte(l.key))
	longIv := append(password[:], ivVector...)

	hash := sha256.Sum256(longIv)

	var iv [aes.BlockSize]byte
	copy(iv[:], hash[:aes.BlockSize])

	aesConn, err := NewAESConn(l.key, iv, conn)
	err = conn.SetReadDeadline(time.Time{})

	return aesConn, err
}

func (l *Listener) Close() error {
	return l.listener.Close()
}

func (l *Listener) Addr() net.Addr {
	return l.listener.Addr()
}

func Listen(listenAddress string, key string) (*Listener, error) {
	listen, err := net.Listen("tcp", listenAddress)
	if err != nil {
		return nil, err
	}

	l := &Listener{
		listener: listen,
		key:      key,
	}
	return l, nil
}

func Dial(address string, key string, timeout time.Duration) (net.Conn, error) {
	conn, err := net.DialTimeout("tcp", address, timeout)
	if err != nil {
		return nil, err
	}

	ivVector := make([]byte, constants.IVLength)
	_, err = rand.Read(ivVector[:])

	if err != nil {
		return nil, err
	}

	password := sha256.Sum256([]byte(key))
	longIv := append(password[:], ivVector...)
	hash := sha256.Sum256(longIv)

	buf := bytes.NewBuffer(ivVector)
	if _, err := io.CopyN(conn, buf, constants.IVLength); err != nil {
		log.Fatal(err)
	}

	var iv [aes.BlockSize]byte
	copy(iv[:], hash[:aes.BlockSize])

	aesConn, err := NewAESConn(key, iv, conn)
	return aesConn, err
}
<commit_msg>set tcp keep alive<commit_after>package transport

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"fmt"
	"io"
	"log"
	"net"
	"time"

	"github.com/Randomsock5/tcptunnel/constants"
)

type AESConn struct {
	conn net.Conn
	r    *cipher.StreamReader
	w    *cipher.StreamWriter
}

func (ac *AESConn) Read(dst []byte) (int, error) {
	n, err := ac.r.Read(dst)
	return n, err
}

func (ac *AESConn) Write(src []byte) (int, error) {
	n, err := ac.w.Write(src)
	return n, err
}

func (ac *AESConn) Close() error {
	return ac.conn.Close()
}

func (ac *AESConn) LocalAddr() net.Addr {
	return ac.conn.LocalAddr()
}

func (ac *AESConn) RemoteAddr() net.Addr {
	return ac.conn.RemoteAddr()
}

func (ac *AESConn) SetDeadline(t time.Time) error {
	return ac.conn.SetDeadline(t)
}

func (ac *AESConn) SetReadDeadline(t time.Time) error {
	return ac.conn.SetReadDeadline(t)
}

func (ac *AESConn) SetWriteDeadline(t time.Time) error {
	return ac.conn.SetWriteDeadline(t)
}

func NewAESConn(key string, iv [aes.BlockSize]byte, conn net.Conn) (*AESConn, error) {
	keyBin := []byte(key)
	t := time.Now().UTC()

	hashKey := sha256.Sum256(keyBin)
	salt := sha256.Sum256(
		[]byte(fmt.Sprintf("%d-%02d-%02d", t.Year(), t.Month(), t.Day())))

	hashSalt := append(hashKey[:], salt[:]...)

	aesKey := sha256.Sum256(hashSalt)

	rBlock, err := aes.NewCipher(aesKey[:])
	if err != nil {
		return nil, err
	}

	wBlock, err := aes.NewCipher(aesKey[:])
	if err != nil {
		return nil, err
	}

	rStream := cipher.NewOFB(rBlock, iv[:])
	wStream := cipher.NewOFB(wBlock, iv[:])

	aesConn := &AESConn{
		conn: conn,
		r:    &cipher.StreamReader{S: rStream, R: conn},
		w:    &cipher.StreamWriter{S: wStream, W: conn},
	}
	return aesConn, err
}

type Listener struct {
	listener net.Listener
	key      string
}

func (l *Listener) Accept() (net.Conn, error) {
	conn, err := l.listener.Accept()
	if err != nil {
		return nil, err
	}
	err = conn.SetReadDeadline(time.Now().Add(constants.ConnTimeout))
	if err != nil {
		return nil, err
	}

	ivVector := make([]byte, constants.IVLength)
	if _, err := io.ReadFull(conn, ivVector); err != nil {
		log.Fatal(err)
	}

	password := sha256.Sum256([]byte(l.key))
	longIv := append(password[:], ivVector...)

	hash := sha256.Sum256(longIv)

	var iv [aes.BlockSize]byte
	copy(iv[:], hash[:aes.BlockSize])

	aesConn, err := NewAESConn(l.key, iv, conn)
	err = conn.SetReadDeadline(time.Time{})

	return aesConn, err
}

func (l *Listener) Close() error {
	return l.listener.Close()
}

func (l *Listener) Addr() net.Addr {
	return l.listener.Addr()
}

func Listen(listenAddress string, key string) (*Listener, error) {
	listen, err := net.Listen("tcp", listenAddress)
	if err != nil {
		return nil, err
	}

	l := &Listener{
		listener: listen,
		key:      key,
	}
	return l, nil
}

func Dial(address string, key string, timeout time.Duration) (net.Conn, error) {
	conn, err := net.DialTimeout("tcp", address, timeout)
	if err != nil {
		return nil, err
	}

	err = conn.(*net.TCPConn).SetKeepAlive(true)
	if err != nil {
		return nil, err
	}

	err = conn.(*net.TCPConn).SetKeepAlivePeriod(2 * timeout)
	if err != nil {
		return nil, err
	}

	ivVector := make([]byte, constants.IVLength)
	_, err = rand.Read(ivVector[:])

	if err != nil {
		return nil, err
	}

	password := sha256.Sum256([]byte(key))
	longIv := append(password[:], ivVector...)
	hash := sha256.Sum256(longIv)

	buf := bytes.NewBuffer(ivVector)
	if _, err := io.CopyN(conn, buf, constants.IVLength); err != nil {
		log.Fatal(err)
	}

	var iv [aes.BlockSize]byte
	copy(iv[:], hash[:aes.BlockSize])

	aesConn, err := NewAESConn(key, iv, conn)
	return aesConn, err
}
<|endoftext|>


<commit_before>package transport

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"fmt"
	"io"
	"log"
	"net"
	"time"

	"github.com/Randomsock5/tcptunnel/constants"
)

type AESConn struct {
	conn net.Conn
	r    *cipher.StreamReader
	w    *cipher.StreamWriter
}

func (ac *AESConn) Read(dst []byte) (int, error) {
	n, err := ac.r.Read(dst)
	return n, err
}

func (ac *AESConn) Write(src []byte) (int, error) {
	n, err := ac.w.Write(src)
	return n, err
}

func (ac *AESConn) Close() error {
	return ac.conn.Close()
}

func (ac *AESConn) LocalAddr() net.Addr {
	return ac.conn.LocalAddr()
}

func (ac *AESConn) RemoteAddr() net.Addr {
	return ac.conn.RemoteAddr()
}

func (ac *AESConn) SetDeadline(t time.Time) error {
	return ac.conn.SetDeadline(t)
}

func (ac *AESConn) SetReadDeadline(t time.Time) error {
	return ac.conn.SetReadDeadline(t)
}

func (ac *AESConn) SetWriteDeadline(t time.Time) error {
	return ac.conn.SetWriteDeadline(t)
}

func NewAESConn(key string, iv [aes.BlockSize]byte, conn net.Conn) (*AESConn, error) {
	keyBin := []byte(key)
	t := time.Now().UTC()

	hashKey := sha256.Sum256(keyBin)
	salt := sha256.Sum256(
		[]byte(fmt.Sprintf("%d-%02d-%02d", t.Year(), t.Month(), t.Day())))

	hashSalt := append(hashKey[:], salt[:]...)

	aesKey := sha256.Sum256(hashSalt)

	rBlock, err := aes.NewCipher(aesKey[:])
	if err != nil {
		return nil, err
	}

	wBlock, err := aes.NewCipher(aesKey[:])
	if err != nil {
		return nil, err
	}

	rStream := cipher.NewOFB(rBlock, iv[:])
	wStream := cipher.NewOFB(wBlock, iv[:])

	aesConn := &AESConn{
		conn: conn,
		r:    &cipher.StreamReader{S: rStream, R: conn},
		w:    &cipher.StreamWriter{S: wStream, W: conn},
	}
	return aesConn, err
}

type Listener struct {
	listener net.Listener
	key      string
}

func (l *Listener) Accept() (net.Conn, error) {
	conn, err := l.listener.Accept()
	if err != nil {
		return nil, err
	}
	err = conn.SetReadDeadline(time.Now().Add(constants.ConnTimeout))
	if err != nil {
		return nil, err
	}

	ivVector := make([]byte, constants.IVLength)
	if _, err := io.ReadFull(conn, ivVector); err != nil {
		log.Fatal(err)
	}

	password := sha256.Sum256([]byte(l.key))
	longIv := append(password[:], ivVector...)

	hash := sha256.Sum256(longIv)

	var iv [aes.BlockSize]byte
	copy(iv[:], hash[:aes.BlockSize])

	aesConn, err := NewAESConn(l.key, iv, conn)
	return aesConn, err
}

func (l *Listener) Close() error {
	return l.listener.Close()
}

func (l *Listener) Addr() net.Addr {
	return l.listener.Addr()
}

func Listen(listenAddress string, key string) (*Listener, error) {
	listen, err := net.Listen("tcp", listenAddress)
	if err != nil {
		return nil, err
	}

	l := &Listener{
		listener: listen,
		key:      key,
	}
	return l, nil
}

func Dial(address string, key string, timeout time.Duration) (net.Conn, error) {
	conn, err := net.DialTimeout("tcp", address, timeout)
	if err != nil {
		return nil, err
	}

	ivVector := make([]byte, constants.IVLength)
	_, err = rand.Read(ivVector[:])

	if err != nil {
		return nil, err
	}

	password := sha256.Sum256([]byte(key))
	longIv := append(password[:], ivVector...)
	hash := sha256.Sum256(longIv)

	buf := bytes.NewBuffer(ivVector)
	if _, err := io.CopyN(conn, buf, constants.IVLength); err != nil {
		log.Fatal(err)
	}

	var iv [aes.BlockSize]byte
	copy(iv[:], hash[:aes.BlockSize])

	aesConn, err := NewAESConn(key, iv, conn)
	return aesConn, err
}
<commit_msg>cancel readtimeout<commit_after>package transport

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"fmt"
	"io"
	"log"
	"net"
	"time"

	"github.com/Randomsock5/tcptunnel/constants"
)

type AESConn struct {
	conn net.Conn
	r    *cipher.StreamReader
	w    *cipher.StreamWriter
}

func (ac *AESConn) Read(dst []byte) (int, error) {
	n, err := ac.r.Read(dst)
	return n, err
}

func (ac *AESConn) Write(src []byte) (int, error) {
	n, err := ac.w.Write(src)
	return n, err
}

func (ac *AESConn) Close() error {
	return ac.conn.Close()
}

func (ac *AESConn) LocalAddr() net.Addr {
	return ac.conn.LocalAddr()
}

func (ac *AESConn) RemoteAddr() net.Addr {
	return ac.conn.RemoteAddr()
}

func (ac *AESConn) SetDeadline(t time.Time) error {
	return ac.conn.SetDeadline(t)
}

func (ac *AESConn) SetReadDeadline(t time.Time) error {
	return ac.conn.SetReadDeadline(t)
}

func (ac *AESConn) SetWriteDeadline(t time.Time) error {
	return ac.conn.SetWriteDeadline(t)
}

func NewAESConn(key string, iv [aes.BlockSize]byte, conn net.Conn) (*AESConn, error) {
	keyBin := []byte(key)
	t := time.Now().UTC()

	hashKey := sha256.Sum256(keyBin)
	salt := sha256.Sum256(
		[]byte(fmt.Sprintf("%d-%02d-%02d", t.Year(), t.Month(), t.Day())))

	hashSalt := append(hashKey[:], salt[:]...)

	aesKey := sha256.Sum256(hashSalt)

	rBlock, err := aes.NewCipher(aesKey[:])
	if err != nil {
		return nil, err
	}

	wBlock, err := aes.NewCipher(aesKey[:])
	if err != nil {
		return nil, err
	}

	rStream := cipher.NewOFB(rBlock, iv[:])
	wStream := cipher.NewOFB(wBlock, iv[:])

	aesConn := &AESConn{
		conn: conn,
		r:    &cipher.StreamReader{S: rStream, R: conn},
		w:    &cipher.StreamWriter{S: wStream, W: conn},
	}
	return aesConn, err
}

type Listener struct {
	listener net.Listener
	key      string
}

func (l *Listener) Accept() (net.Conn, error) {
	conn, err := l.listener.Accept()
	if err != nil {
		return nil, err
	}
	err = conn.SetReadDeadline(time.Now().Add(constants.ConnTimeout))
	if err != nil {
		return nil, err
	}

	ivVector := make([]byte, constants.IVLength)
	if _, err := io.ReadFull(conn, ivVector); err != nil {
		log.Fatal(err)
	}

	password := sha256.Sum256([]byte(l.key))
	longIv := append(password[:], ivVector...)

	hash := sha256.Sum256(longIv)

	var iv [aes.BlockSize]byte
	copy(iv[:], hash[:aes.BlockSize])

	aesConn, err := NewAESConn(l.key, iv, conn)
	err = conn.SetReadDeadline(time.Time{})

	return aesConn, err
}

func (l *Listener) Close() error {
	return l.listener.Close()
}

func (l *Listener) Addr() net.Addr {
	return l.listener.Addr()
}

func Listen(listenAddress string, key string) (*Listener, error) {
	listen, err := net.Listen("tcp", listenAddress)
	if err != nil {
		return nil, err
	}

	l := &Listener{
		listener: listen,
		key:      key,
	}
	return l, nil
}

func Dial(address string, key string, timeout time.Duration) (net.Conn, error) {
	conn, err := net.DialTimeout("tcp", address, timeout)
	if err != nil {
		return nil, err
	}

	ivVector := make([]byte, constants.IVLength)
	_, err = rand.Read(ivVector[:])

	if err != nil {
		return nil, err
	}

	password := sha256.Sum256([]byte(key))
	longIv := append(password[:], ivVector...)
	hash := sha256.Sum256(longIv)

	buf := bytes.NewBuffer(ivVector)
	if _, err := io.CopyN(conn, buf, constants.IVLength); err != nil {
		log.Fatal(err)
	}

	var iv [aes.BlockSize]byte
	copy(iv[:], hash[:aes.BlockSize])

	aesConn, err := NewAESConn(key, iv, conn)
	return aesConn, err
}
<|endoftext|>


<commit_before>package routing

import (
	"net/http"
	"sync"
	"time"

	"github.com/matrix-org/dendrite/clientapi/jsonerror"
	"github.com/matrix-org/dendrite/internal/config"
	"github.com/matrix-org/util"
)

type rateLimits struct {
	limits           map[string]chan struct{}
	limitsMutex      sync.RWMutex
	enabled          bool
	requestThreshold int64
	cooloffDuration  time.Duration
}

func newRateLimits(cfg *config.RateLimiting) *rateLimits {
	l := &rateLimits{
		limits:           make(map[string]chan struct{}),
		enabled:          cfg.Enabled,
		requestThreshold: cfg.Threshold,
		cooloffDuration:  time.Duration(cfg.CooloffMS) * time.Millisecond,
	}
	if l.enabled {
		go l.clean()
	}
	return l
}

func (l *rateLimits) clean() {
	for {
		// On a 30 second interval, we'll take an exclusive write
		// lock of the entire map and see if any of the channels are
		// empty. If they are then we will close and delete them,
		// freeing up memory.
		time.Sleep(time.Second * 30)
		l.limitsMutex.Lock()
		for k, c := range l.limits {
			if len(c) == 0 {
				close(c)
				delete(l.limits, k)
			}
		}
		l.limitsMutex.Unlock()
	}
}

func (l *rateLimits) rateLimit(req *http.Request) *util.JSONResponse {
	// If rate limiting is disabled then do nothing.
	if !l.enabled {
		return nil
	}

	// Lock the map long enough to check for rate limiting. We hold it
	// for longer here than we really need to but it makes sure that we
	// also don't conflict with the cleaner goroutine which might clean
	// up a channel after we have retrieved it otherwise.
	l.limitsMutex.RLock()
	defer l.limitsMutex.RUnlock()

	// First of all, work out if X-Forwarded-For was sent to us. If not
	// then we'll just use the IP address of the caller.
	caller := req.RemoteAddr
	if forwardedFor := req.Header.Get("X-Forwarded-For"); forwardedFor != "" {
		caller = forwardedFor
	}

	// Look up the caller's channel, if they have one. If they don't then
	// let's create one.
	rateLimit, ok := l.limits[caller]
	if !ok {
		l.limits[caller] = make(chan struct{}, l.requestThreshold)
		rateLimit = l.limits[caller]
	}

	// Check if the user has got free resource slots for this request.
	// If they don't then we'll return an error.
	select {
	case rateLimit <- struct{}{}:
	default:
		// We hit the rate limit. Tell the client to back off.
		return &util.JSONResponse{
			Code: http.StatusTooManyRequests,
			JSON: jsonerror.LimitExceeded("You are sending too many requests too quickly!", l.cooloffDuration.Milliseconds()),
		}
	}

	// After the time interval, drain a resource from the rate limiting
	// channel. This will free up space in the channel for new requests.
	go func() {
		<-time.After(l.cooloffDuration)
		<-rateLimit
	}()
	return nil
}
<commit_msg>Take write lock for rate limit map (#1532)<commit_after>package routing

import (
	"net/http"
	"sync"
	"time"

	"github.com/matrix-org/dendrite/clientapi/jsonerror"
	"github.com/matrix-org/dendrite/internal/config"
	"github.com/matrix-org/util"
)

type rateLimits struct {
	limits           map[string]chan struct{}
	limitsMutex      sync.RWMutex
	cleanMutex       sync.RWMutex
	enabled          bool
	requestThreshold int64
	cooloffDuration  time.Duration
}

func newRateLimits(cfg *config.RateLimiting) *rateLimits {
	l := &rateLimits{
		limits:           make(map[string]chan struct{}),
		enabled:          cfg.Enabled,
		requestThreshold: cfg.Threshold,
		cooloffDuration:  time.Duration(cfg.CooloffMS) * time.Millisecond,
	}
	if l.enabled {
		go l.clean()
	}
	return l
}

func (l *rateLimits) clean() {
	for {
		// On a 30 second interval, we'll take an exclusive write
		// lock of the entire map and see if any of the channels are
		// empty. If they are then we will close and delete them,
		// freeing up memory.
		time.Sleep(time.Second * 30)
		l.cleanMutex.Lock()
		l.limitsMutex.Lock()
		for k, c := range l.limits {
			if len(c) == 0 {
				close(c)
				delete(l.limits, k)
			}
		}
		l.limitsMutex.Unlock()
		l.cleanMutex.Unlock()
	}
}

func (l *rateLimits) rateLimit(req *http.Request) *util.JSONResponse {
	// If rate limiting is disabled then do nothing.
	if !l.enabled {
		return nil
	}

	// Take a read lock out on the cleaner mutex. The cleaner expects to
	// be able to take a write lock, which isn't possible while there are
	// readers, so this has the effect of blocking the cleaner goroutine
	// from doing its work until there are no requests in flight.
	l.cleanMutex.RLock()
	defer l.cleanMutex.RUnlock()

	// First of all, work out if X-Forwarded-For was sent to us. If not
	// then we'll just use the IP address of the caller.
	caller := req.RemoteAddr
	if forwardedFor := req.Header.Get("X-Forwarded-For"); forwardedFor != "" {
		caller = forwardedFor
	}

	// Look up the caller's channel, if they have one.
	l.limitsMutex.RLock()
	rateLimit, ok := l.limits[caller]
	l.limitsMutex.RUnlock()

	// If the caller doesn't have a channel, create one and write it
	// back to the map.
	if !ok {
		rateLimit = make(chan struct{}, l.requestThreshold)

		l.limitsMutex.Lock()
		l.limits[caller] = rateLimit
		l.limitsMutex.Unlock()
	}

	// Check if the user has got free resource slots for this request.
	// If they don't then we'll return an error.
	select {
	case rateLimit <- struct{}{}:
	default:
		// We hit the rate limit. Tell the client to back off.
		return &util.JSONResponse{
			Code: http.StatusTooManyRequests,
			JSON: jsonerror.LimitExceeded("You are sending too many requests too quickly!", l.cooloffDuration.Milliseconds()),
		}
	}

	// After the time interval, drain a resource from the rate limiting
	// channel. This will free up space in the channel for new requests.
	go func() {
		<-time.After(l.cooloffDuration)
		<-rateLimit
	}()
	return nil
}
<|endoftext|>


<commit_before>package media_library

import (
	"bytes"
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"html/template"
	"image"
	"io"
	"mime/multipart"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"

	"github.com/disintegration/imaging"
	"github.com/jinzhu/gorm"
)

var ErrNotImplemented = errors.New("not implemented")

type CropOption struct {
	X, Y, Width, Height int
}

type fileHeader interface {
	Open() (multipart.File, error)
}

type fileWrapper struct {
	*os.File
}

func (fileWrapper *fileWrapper) Open() (multipart.File, error) {
	return fileWrapper.File, nil
}

type Base struct {
	FileName    string
	Url         string
	CropOptions map[string]*CropOption `json:",omitempty"`
	Crop        bool                   `json:"-"`
	Valid       bool                   `json:"-"`
	FileHeader  fileHeader             `json:"-"`
	Reader      io.Reader              `json:"-"`
}

func (b *Base) Scan(data interface{}) (err error) {
	switch values := data.(type) {
	case *os.File:
		b.FileHeader = &fileWrapper{values}
		b.FileName = path.Base(values.Name())
		b.Valid = true
	case []*multipart.FileHeader:
		if len(values) > 0 {
			file := values[0]
			b.FileHeader, b.FileName, b.Valid = file, file.Filename, true
		}
	case []byte:
		if err = json.Unmarshal(values, b); err == nil {
			b.Valid = true
		}
		var doCrop struct{ Crop bool }
		if err = json.Unmarshal(values, &doCrop); err == nil && doCrop.Crop {
			b.Crop = true
		}
	case string:
		b.Scan([]byte(values))
	case []string:
		for _, str := range values {
			b.Scan(str)
		}
	default:
		err = errors.New("unsupported driver -> Scan pair for MediaLibrary")
	}
	return
}

func (b Base) Value() (driver.Value, error) {
	if b.Valid {
		result, err := json.Marshal(b)
		return string(result), err
	}
	return nil, nil
}

func (b Base) URL(styles ...string) string {
	if len(styles) > 0 {
		ext := path.Ext(b.Url)
		return fmt.Sprintf("%v.%v%v", strings.TrimSuffix(b.Url, ext), styles[0], ext)
	}
	return b.Url
}

func (b Base) String() string {
	return b.URL()
}

func (b Base) GetFileName() string {
	return b.FileName
}

func (b Base) GetFileHeader() fileHeader {
	return b.FileHeader
}

func (b Base) GetURLTemplate(option *Option) (path string) {
	if path = option.Get("URL"); path == "" {
		path = "/system/{{class}}/{{primary_key}}/{{column}}/{{filename_with_hash}}"
	}
	return
}

func getFuncMap(scope *gorm.Scope, field *gorm.Field, filename string) template.FuncMap {
	hash := func() string { return strings.Replace(time.Now().Format("20060102150506.000000000"), ".", "", -1) }
	return template.FuncMap{
		"class":       scope.TableName,
		"primary_key": func() string { return fmt.Sprintf("%v", scope.PrimaryKeyValue()) },
		"column":      func() string { return field.Name },
		"filename":    func() string { return filename },
		"basename":    func() string { return strings.TrimSuffix(path.Base(filename), path.Ext(filename)) },
		"hash":        hash,
		"filename_with_hash": func() string {
			return fmt.Sprintf("%v.%v%v", strings.TrimSuffix(filename, path.Ext(filename)), hash(), path.Ext(filename))
		},
		"extension": func() string { return strings.TrimPrefix(path.Ext(filename), ".") },
	}
}

func (b Base) GetURL(option *Option, scope *gorm.Scope, field *gorm.Field, templater URLTemplater) string {
	if path := templater.GetURLTemplate(option); path != "" {
		tmpl := template.New("").Funcs(getFuncMap(scope, field, b.GetFileName()))
		if tmpl, err := tmpl.Parse(path); err == nil {
			var result = bytes.NewBufferString("")
			if err := tmpl.Execute(result, scope.Value); err == nil {
				return result.String()
			}
		}
	}
	return ""
}

func (b *Base) NeedCrop() bool {
	return b.Crop
}

func (b *Base) GetCropOption(name string) *image.Rectangle {
	if cropOption := b.CropOptions[name]; cropOption != nil {
		return &image.Rectangle{
			Min: image.Point{X: cropOption.X, Y: cropOption.Y},
			Max: image.Point{X: cropOption.X + cropOption.Width, Y: cropOption.Y + cropOption.Height},
		}
	} else {
		return nil
	}
}

func (b Base) Retrieve(url string) (*os.File, error) {
	return nil, ErrNotImplemented
}

func (b Base) GetSizes() map[string]Size {
	return map[string]Size{}
}

func (b Base) IsImage() bool {
	_, err := getImageFormat(b.URL())
	return err == nil
}

func getImageFormat(url string) (*imaging.Format, error) {
	formats := map[string]imaging.Format{
		".jpg":  imaging.JPEG,
		".jpeg": imaging.JPEG,
		".png":  imaging.PNG,
		".tif":  imaging.TIFF,
		".tiff": imaging.TIFF,
		".bmp":  imaging.BMP,
		".gif":  imaging.GIF,
	}

	ext := strings.ToLower(filepath.Ext(url))
	if f, ok := formats[ext]; ok {
		return &f, nil
	} else {
		return nil, imaging.ErrUnsupportedFormat
	}
}
<commit_msg>Bigger size for media library option<commit_after>package media_library

import (
	"bytes"
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"html/template"
	"image"
	"io"
	"mime/multipart"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"

	"github.com/disintegration/imaging"
	"github.com/jinzhu/gorm"
)

var ErrNotImplemented = errors.New("not implemented")

type CropOption struct {
	X, Y, Width, Height int
}

type fileHeader interface {
	Open() (multipart.File, error)
}

type fileWrapper struct {
	*os.File
}

func (fileWrapper *fileWrapper) Open() (multipart.File, error) {
	return fileWrapper.File, nil
}

type Base struct {
	FileName    string `sql:"size:10240"`
	Url         string
	CropOptions map[string]*CropOption `json:",omitempty"`
	Crop        bool                   `json:"-"`
	Valid       bool                   `json:"-"`
	FileHeader  fileHeader             `json:"-"`
	Reader      io.Reader              `json:"-"`
}

func (b *Base) Scan(data interface{}) (err error) {
	switch values := data.(type) {
	case *os.File:
		b.FileHeader = &fileWrapper{values}
		b.FileName = path.Base(values.Name())
		b.Valid = true
	case []*multipart.FileHeader:
		if len(values) > 0 {
			file := values[0]
			b.FileHeader, b.FileName, b.Valid = file, file.Filename, true
		}
	case []byte:
		if err = json.Unmarshal(values, b); err == nil {
			b.Valid = true
		}
		var doCrop struct{ Crop bool }
		if err = json.Unmarshal(values, &doCrop); err == nil && doCrop.Crop {
			b.Crop = true
		}
	case string:
		b.Scan([]byte(values))
	case []string:
		for _, str := range values {
			b.Scan(str)
		}
	default:
		err = errors.New("unsupported driver -> Scan pair for MediaLibrary")
	}
	return
}

func (b Base) Value() (driver.Value, error) {
	if b.Valid {
		result, err := json.Marshal(b)
		return string(result), err
	}
	return nil, nil
}

func (b Base) URL(styles ...string) string {
	if len(styles) > 0 {
		ext := path.Ext(b.Url)
		return fmt.Sprintf("%v.%v%v", strings.TrimSuffix(b.Url, ext), styles[0], ext)
	}
	return b.Url
}

func (b Base) String() string {
	return b.URL()
}

func (b Base) GetFileName() string {
	return b.FileName
}

func (b Base) GetFileHeader() fileHeader {
	return b.FileHeader
}

func (b Base) GetURLTemplate(option *Option) (path string) {
	if path = option.Get("URL"); path == "" {
		path = "/system/{{class}}/{{primary_key}}/{{column}}/{{filename_with_hash}}"
	}
	return
}

func getFuncMap(scope *gorm.Scope, field *gorm.Field, filename string) template.FuncMap {
	hash := func() string { return strings.Replace(time.Now().Format("20060102150506.000000000"), ".", "", -1) }
	return template.FuncMap{
		"class":       scope.TableName,
		"primary_key": func() string { return fmt.Sprintf("%v", scope.PrimaryKeyValue()) },
		"column":      func() string { return field.Name },
		"filename":    func() string { return filename },
		"basename":    func() string { return strings.TrimSuffix(path.Base(filename), path.Ext(filename)) },
		"hash":        hash,
		"filename_with_hash": func() string {
			return fmt.Sprintf("%v.%v%v", strings.TrimSuffix(filename, path.Ext(filename)), hash(), path.Ext(filename))
		},
		"extension": func() string { return strings.TrimPrefix(path.Ext(filename), ".") },
	}
}

func (b Base) GetURL(option *Option, scope *gorm.Scope, field *gorm.Field, templater URLTemplater) string {
	if path := templater.GetURLTemplate(option); path != "" {
		tmpl := template.New("").Funcs(getFuncMap(scope, field, b.GetFileName()))
		if tmpl, err := tmpl.Parse(path); err == nil {
			var result = bytes.NewBufferString("")
			if err := tmpl.Execute(result, scope.Value); err == nil {
				return result.String()
			}
		}
	}
	return ""
}

func (b *Base) NeedCrop() bool {
	return b.Crop
}

func (b *Base) GetCropOption(name string) *image.Rectangle {
	if cropOption := b.CropOptions[name]; cropOption != nil {
		return &image.Rectangle{
			Min: image.Point{X: cropOption.X, Y: cropOption.Y},
			Max: image.Point{X: cropOption.X + cropOption.Width, Y: cropOption.Y + cropOption.Height},
		}
	} else {
		return nil
	}
}

func (b Base) Retrieve(url string) (*os.File, error) {
	return nil, ErrNotImplemented
}

func (b Base) GetSizes() map[string]Size {
	return map[string]Size{}
}

func (b Base) IsImage() bool {
	_, err := getImageFormat(b.URL())
	return err == nil
}

func getImageFormat(url string) (*imaging.Format, error) {
	formats := map[string]imaging.Format{
		".jpg":  imaging.JPEG,
		".jpeg": imaging.JPEG,
		".png":  imaging.PNG,
		".tif":  imaging.TIFF,
		".tiff": imaging.TIFF,
		".bmp":  imaging.BMP,
		".gif":  imaging.GIF,
	}

	ext := strings.ToLower(filepath.Ext(url))
	if f, ok := formats[ext]; ok {
		return &f, nil
	} else {
		return nil, imaging.ErrUnsupportedFormat
	}
}
<|endoftext|>


<commit_before>package esmlt

import (
	"bytes"
	"fmt"
	"net/url"
	"strconv"
	"strings"

	"github.com/belogik/goes"
)

// Application Version
const AppVersion = "0.1.0"

// SearchConnection is an interface that is satisfied by goes.Connection and
// can be satisfied by mock index connections for testing
type SearchConnection interface {
	Search(query map[string]interface{}, indexList []string, typeList []string, extraArgs url.Values) (goes.Response, error)
}

// ParseIndicesShift parses strings like `2,4,5` into an int slice and adds a `shift`
func ParseIndicesShift(s string, shift int) ([]int, error) {
	parts := strings.Split(s, ",")
	var indices []int
	for _, p := range parts {
		value := strings.TrimSpace(p)
		if value == "" {
			continue
		}
		i, err := strconv.ParseInt(value, 10, 0)
		if err != nil {
			return nil, err
		}
		indices = append(indices, int(i)+shift)
	}
	return indices, nil
}

// ParseIndices parses strings like `2,4,5` into an int slice
func ParseIndices(s string) ([]int, error) {
	return ParseIndicesShift(s, -1)
}

// ConcatenateValuesNull extracts values according to indices slice and concatenates them.
// Values that equal a given `nullValue` are ignored.
func ConcatenateValuesNull(values []string, indices []int, nullValue string) (string, error) {
	var buffer bytes.Buffer
	for _, i := range indices {
		if i > len(values)-1 {
			return "", fmt.Errorf("index %d exceeds array", i)
		}
		if values[i] == nullValue {
			buffer.WriteString("")
			continue
		}
		buffer.WriteString(values[i])
		buffer.WriteString(" ")
	}
	return strings.TrimSpace(buffer.String()), nil
}

// ConcatenateValues extracts values according to indices slice and concatenates them,
// uses a default `nullValue` of `<NULL>`
func ConcatenateValues(values []string, indices []int) (string, error) {
	return ConcatenateValuesNull(values, indices, "<NULL>")
}

// Value returns the value in a nested map according to a key in dot notation
func Value(key string, doc map[string]interface{}) interface{} {
	keys := strings.Split(key, ".")
	for _, k := range keys {
		value := doc[k]
		if value == nil {
			return nil
		}
		switch value.(type) {
		case map[string]interface{}:
			return Value(strings.Join(keys[1:], "."), value.(map[string]interface{}))
		case []interface{}:
			if len(value.([]interface{})) == 0 {
				return nil
			}
			first := value.([]interface{})[0]
			switch first.(type) {
			case map[string]interface{}:
				return Value(strings.Join(keys[1:], "."), first.(map[string]interface{}))
			case string:
				return first
				continue
			}
		default:
			return value
		}
	}
	return nil
}
<commit_msg>updated doc<commit_after>package esmlt

import (
	"bytes"
	"fmt"
	"net/url"
	"strconv"
	"strings"

	"github.com/belogik/goes"
)

// Application Version
const AppVersion = "0.1.0"

// SearchConnection is an interface that is satisfied by goes.Connection and
// can be satisfied by mock index connections for testing
type SearchConnection interface {
	Search(query map[string]interface{}, indexList []string, typeList []string, extraArgs url.Values) (goes.Response, error)
}

// ParseIndicesShift parses strings like `2,4,5` into an int slice and adds a `shift`
func ParseIndicesShift(s string, shift int) ([]int, error) {
	parts := strings.Split(s, ",")
	var indices []int
	for _, p := range parts {
		value := strings.TrimSpace(p)
		if value == "" {
			continue
		}
		i, err := strconv.ParseInt(value, 10, 0)
		if err != nil {
			return nil, err
		}
		indices = append(indices, int(i)+shift)
	}
	return indices, nil
}

// ParseIndices parses strings like `2,4,5` into an int slice
func ParseIndices(s string) ([]int, error) {
	return ParseIndicesShift(s, -1)
}

// ConcatenateValuesNull extracts values according to indices slice and concatenates them.
// Values that equal a given `nullValue` are ignored.
func ConcatenateValuesNull(values []string, indices []int, nullValue string) (string, error) {
	var buffer bytes.Buffer
	for _, i := range indices {
		if i > len(values)-1 {
			return "", fmt.Errorf("index %d exceeds array", i)
		}
		if values[i] == nullValue {
			buffer.WriteString("")
			continue
		}
		buffer.WriteString(values[i])
		buffer.WriteString(" ")
	}
	return strings.TrimSpace(buffer.String()), nil
}

// ConcatenateValues extracts values according to indices slice and concatenates them,
// uses a default `nullValue` of `<NULL>`
func ConcatenateValues(values []string, indices []int) (string, error) {
	return ConcatenateValuesNull(values, indices, "<NULL>")
}

// Value returns the value in a (nested) map according to a key in dot notation.
// If the value is a slice, only the first element is considered.
func Value(key string, doc map[string]interface{}) interface{} {
	keys := strings.Split(key, ".")
	for _, k := range keys {
		value := doc[k]
		if value == nil {
			return nil
		}
		switch value.(type) {
		case map[string]interface{}:
			return Value(strings.Join(keys[1:], "."), value.(map[string]interface{}))
		case []interface{}:
			if len(value.([]interface{})) == 0 {
				return nil
			}
			first := value.([]interface{})[0]
			switch first.(type) {
			case map[string]interface{}:
				return Value(strings.Join(keys[1:], "."), first.(map[string]interface{}))
			case string:
				return first
				continue
			}
		default:
			return value
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package espsdk

import (
	"encoding/json"
	"time"

	"github.com/Sirupsen/logrus"
	prefixed "github.com/x-cray/logrus-prefixed-formatter"
)

var Log = logrus.New()

func init() {
	Log.Formatter = &prefixed.TextFormatter{TimestampFormat: time.RFC3339}
}

// A Token is a string representation of an OAuth2 token. It grants a user
// access to the ESP API for a limited time.
type Token string

// A Serializable object can be serialized to a byte stream such as JSON.
type serializable interface {
	Marshal() ([]byte, error)
}

func indentedJSON(obj interface{}) ([]byte, error) {
	return json.MarshalIndent(obj, "", "\t")
}
<commit_msg>docs: obey the linter<commit_after>package espsdk

import (
	"encoding/json"
	"time"

	"github.com/Sirupsen/logrus"
	prefixed "github.com/x-cray/logrus-prefixed-formatter"
)

// Log is an instance of logrus.(*Logger) that allows easy access to
// nicely-formatted log output that includes properties of completed HTTP
// requests such as the response time and status code.
var Log = logrus.New()

func init() {
	Log.Formatter = &prefixed.TextFormatter{TimestampFormat: time.RFC3339}
}

// A Token is a string representation of an OAuth2 token. It grants a user
// access to the ESP API for a limited time.
type Token string

// A Serializable object can be serialized to a byte stream such as JSON.
type serializable interface {
	Marshal() ([]byte, error)
}

func indentedJSON(obj interface{}) ([]byte, error) {
	return json.MarshalIndent(obj, "", "\t")
}
<|endoftext|>