<commit_before>// Copyright 2020 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package logstream

import (
	"bytes"
	"context"
	"errors"
	"expvar"
	"io"
	"os"
	"sync"
	"syscall"
	"time"

	"github.com/golang/glog"
	"github.com/google/mtail/internal/logline"
	"github.com/google/mtail/internal/waker"
)

// fileTruncates counts the truncations of a file stream.
var fileTruncates = expvar.NewMap("file_truncates_total")

// fileStream streams log lines from a regular file on the file system.  These
// log files are appended to by another process, and are either rotated or
// truncated by that (or yet another) process.  Rotation implies that a new
// inode with the same name has been created, the old file descriptor will be
// valid until EOF at which point it's considered completed.  A truncation means
// the same file descriptor is used but the file offset will be reset to 0.
// The latter is potentially lossy as far as mtail is concerned, if the last
// logs are not read before truncation occurs.  When an EOF is read, the
// goroutine tests for both truncation and inode change and resets or spins off
// a new goroutine and closes itself down.  The shared context is used for
// cancellation.
type fileStream struct {
	ctx   context.Context
	lines chan<- *logline.LogLine

	pathname string // Given name for the underlying file on the filesystem

	mu           sync.RWMutex // protects following fields.
	lastReadTime time.Time    // Last time a log line was read from this file
	completed    bool         // The filestream is completed and can no longer be used.

	stopOnce sync.Once     // Ensure stopChan only closed once.
	stopChan chan struct{} // Close to start graceful shutdown.
}

// newFileStream creates a new log stream from a regular file.
func newFileStream(ctx context.Context, wg *sync.WaitGroup, waker waker.Waker, pathname string, fi os.FileInfo, lines chan<- *logline.LogLine, streamFromStart bool) (LogStream, error) {
	fs := &fileStream{ctx: ctx, pathname: pathname, lastReadTime: time.Now(), lines: lines, stopChan: make(chan struct{})}
	if err := fs.stream(ctx, wg, waker, fi, streamFromStart); err != nil {
		return nil, err
	}
	return fs, nil
}

func (fs *fileStream) LastReadTime() time.Time {
	fs.mu.RLock()
	defer fs.mu.RUnlock()
	return fs.lastReadTime
}

func (fs *fileStream) stream(ctx context.Context, wg *sync.WaitGroup, waker waker.Waker, fi os.FileInfo, streamFromStart bool) error {
	fd, err := os.OpenFile(fs.pathname, os.O_RDONLY, 0o600)
	if err != nil {
		logErrors.Add(fs.pathname, 1)
		return err
	}
	logOpens.Add(fs.pathname, 1)
	glog.V(2).Infof("%v: opened new file", fd)
	if !streamFromStart {
		if _, err := fd.Seek(0, io.SeekEnd); err != nil {
			logErrors.Add(fs.pathname, 1)
			if err := fd.Close(); err != nil {
				logErrors.Add(fs.pathname, 1)
				glog.Info(err)
			}
			return err
		}
		glog.V(2).Infof("%v: seeked to end", fd)
	}
	b := make([]byte, defaultReadBufferSize)
	var lastBytes []byte
	partial := bytes.NewBufferString("")
	started := make(chan struct{})
	var total int
	wg.Add(1)
	go func() {
		defer wg.Done()
		defer func() {
			glog.V(2).Infof("%v: read total %d bytes from %s", fd, total, fs.pathname)
			glog.V(2).Infof("%v: closing file descriptor", fd)
			if err := fd.Close(); err != nil {
				logErrors.Add(fs.pathname, 1)
				glog.Info(err)
			}
			logCloses.Add(fs.pathname, 1)
		}()
		close(started)
		for {
			// Blocking read but regular files will return EOF straight away.
			count, err := fd.Read(b)
			glog.V(2).Infof("%v: read %d bytes, err is %v", fd, count, err)

			if count > 0 {
				total += count
				glog.V(2).Infof("%v: decode and send", fd)
				needSend := append(lastBytes, b[:count]...)
				sendCount := decodeAndSend(ctx, fs.lines, fs.pathname, len(needSend), needSend, partial)
				if sendCount < len(needSend) {
					lastBytes = append([]byte{}, needSend[sendCount:]...)
				} else {
					lastBytes = []byte{}
				}
				fs.mu.Lock()
				fs.lastReadTime = time.Now()
				fs.mu.Unlock()
			}

			if err != nil && err != io.EOF {
				logErrors.Add(fs.pathname, 1)
				// TODO: This could be generalised to check for any retryable
				// errors, and end on unretriables; e.g. ESTALE looks
				// retryable.
				if errors.Is(err, syscall.ESTALE) {
					glog.Infof("%v: reopening stream due to %s", fd, err)
					if nerr := fs.stream(ctx, wg, waker, fi, true); nerr != nil {
						glog.Info(nerr)
					}
					// Close this stream.
					return
				}
				glog.Info(err)
			}

			// If we have read no bytes and are at EOF, check for truncation and rotation.
			if err == io.EOF && count == 0 {
				glog.V(2).Infof("%v: eof an no bytes", fd)
				// Both rotation and truncation need to stat, so check for
				// rotation first.  It is assumed that rotation is the more
				// common change pattern anyway.
				newfi, serr := os.Stat(fs.pathname)
				if serr != nil {
					glog.Info(serr)
					// If this is a NotExist error, then we should wrap up this
					// goroutine. The Tailer will create a new logstream if the
					// file is in the middle of a rotation and gets recreated
					// in the next moment.  We can't rely on the Tailer to tell
					// us we're deleted because the tailer can only tell us to
					// Stop, which ends up causing us to race here against
					// detection of IsCompleted.
					if os.IsNotExist(serr) {
						glog.V(2).Infof("%v: source no longer exists, exiting", fd)
						if partial.Len() > 0 {
							sendLine(ctx, fs.pathname, partial, fs.lines)
						}
						fs.mu.Lock()
						fs.completed = true
						fs.mu.Unlock()
						return
					}
					logErrors.Add(fs.pathname, 1)
					goto Sleep
				}
				if !os.SameFile(fi, newfi) {
					glog.V(2).Infof("%v: adding a new file routine", fd)
					if err := fs.stream(ctx, wg, waker, newfi, true); err != nil {
						glog.Info(err)
					}
					// We're at EOF so there's nothing left to read here.
					return
				}
				currentOffset, serr := fd.Seek(0, io.SeekCurrent)
				if serr != nil {
					logErrors.Add(fs.pathname, 1)
					glog.Info(serr)
					continue
				}
				glog.V(2).Infof("%v: current seek is %d", fd, currentOffset)
				glog.V(2).Infof("%v: new size is %d", fd, newfi.Size())
				// We know that newfi is from the current file.  Truncation can
				// only be detected if the new file is currently shorter than
				// the current seek offset.  In test this can be a race, but in
				// production it's unlikely that a new file writes more bytes
				// than the previous after rotation in the time it takes for
				// mtail to notice.
				if newfi.Size() < currentOffset {
					glog.V(2).Infof("%v: truncate? currentoffset is %d and size is %d", fd, currentOffset, newfi.Size())
					// About to lose all remaining data because of the truncate so flush the accumulator.
					if partial.Len() > 0 {
						sendLine(ctx, fs.pathname, partial, fs.lines)
					}
					p, serr := fd.Seek(0, io.SeekStart)
					if serr != nil {
						logErrors.Add(fs.pathname, 1)
						glog.Info(serr)
					}
					glog.V(2).Infof("%v: Seeked to %d", fd, p)
					fileTruncates.Add(fs.pathname, 1)
					continue
				}
			}

			// No error implies there is more to read in this file so go
			// straight back to read unless it looks like context is Done.
			if err == nil && ctx.Err() == nil {
				continue
			}

		Sleep:
			// If we get here it's because we've stalled.  First test to see if it's
			// time to exit.
			if err == io.EOF || ctx.Err() != nil {
				select {
				case <-fs.stopChan:
					glog.V(2).Infof("%v: stream has been stopped, exiting", fd)
					if partial.Len() > 0 {
						sendLine(ctx, fs.pathname, partial, fs.lines)
					}
					fs.mu.Lock()
					fs.completed = true
					fs.mu.Unlock()
					return
				case <-ctx.Done():
					glog.V(2).Infof("%v: stream has been cancelled, exiting", fd)
					if partial.Len() > 0 {
						sendLine(ctx, fs.pathname, partial, fs.lines)
					}
					fs.mu.Lock()
					fs.completed = true
					fs.mu.Unlock()
					return
				default:
					// keep going
				}
			}

			// Don't exit, instead yield and wait for a termination signal or
			// wakeup.
			glog.V(2).Infof("%v: waiting", fd)
			select {
			case <-fs.stopChan:
				// We may have started waiting here when the stop signal
				// arrives, but since that wait the file may have been
				// written to.  The file is not technically yet at EOF so
				// we need to go back and try one more read.  We'll exit
				// the stream in the select stanza above.
				glog.V(2).Infof("%v: Stopping after next read", fd)
			case <-ctx.Done():
				// Same for cancellation; this makes tests stable, but
				// could argue exiting immediately is less surprising.
				// Assumption is that this doesn't make a difference in
				// production.
				glog.V(2).Infof("%v: Cancelled after next read", fd)
			case <-waker.Wake():
				// sleep until next Wake()
				glog.V(2).Infof("%v: Wake received", fd)
			}
		}
	}()

	<-started
	return nil
}

func (fs *fileStream) IsComplete() bool {
	fs.mu.RLock()
	defer fs.mu.RUnlock()
	return fs.completed
}

// Stop implements the LogStream interface.
func (fs *fileStream) Stop() {
	fs.stopOnce.Do(func() {
		glog.Info("signalling stop at next EOF")
		close(fs.stopChan)
	})
}
<commit_msg>Fix golangci-lint issue<commit_after>// Copyright 2020 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package logstream

import (
	"bytes"
	"context"
	"errors"
	"expvar"
	"io"
	"os"
	"sync"
	"syscall"
	"time"

	"github.com/golang/glog"
	"github.com/google/mtail/internal/logline"
	"github.com/google/mtail/internal/waker"
)

// fileTruncates counts the truncations of a file stream.
var fileTruncates = expvar.NewMap("file_truncates_total")

// fileStream streams log lines from a regular file on the file system.  These
// log files are appended to by another process, and are either rotated or
// truncated by that (or yet another) process.  Rotation implies that a new
// inode with the same name has been created, the old file descriptor will be
// valid until EOF at which point it's considered completed.  A truncation means
// the same file descriptor is used but the file offset will be reset to 0.
// The latter is potentially lossy as far as mtail is concerned, if the last
// logs are not read before truncation occurs.  When an EOF is read, the
// goroutine tests for both truncation and inode change and resets or spins off
// a new goroutine and closes itself down.  The shared context is used for
// cancellation.
type fileStream struct {
	ctx   context.Context
	lines chan<- *logline.LogLine

	pathname string // Given name for the underlying file on the filesystem

	mu           sync.RWMutex // protects following fields.
	lastReadTime time.Time    // Last time a log line was read from this file
	completed    bool         // The filestream is completed and can no longer be used.

	stopOnce sync.Once     // Ensure stopChan only closed once.
	stopChan chan struct{} // Close to start graceful shutdown.
}

// newFileStream creates a new log stream from a regular file.
func newFileStream(ctx context.Context, wg *sync.WaitGroup, waker waker.Waker, pathname string, fi os.FileInfo, lines chan<- *logline.LogLine, streamFromStart bool) (LogStream, error) {
	fs := &fileStream{ctx: ctx, pathname: pathname, lastReadTime: time.Now(), lines: lines, stopChan: make(chan struct{})}
	if err := fs.stream(ctx, wg, waker, fi, streamFromStart); err != nil {
		return nil, err
	}
	return fs, nil
}

func (fs *fileStream) LastReadTime() time.Time {
	fs.mu.RLock()
	defer fs.mu.RUnlock()
	return fs.lastReadTime
}

func (fs *fileStream) stream(ctx context.Context, wg *sync.WaitGroup, waker waker.Waker, fi os.FileInfo, streamFromStart bool) error {
	fd, err := os.OpenFile(fs.pathname, os.O_RDONLY, 0o600)
	if err != nil {
		logErrors.Add(fs.pathname, 1)
		return err
	}
	logOpens.Add(fs.pathname, 1)
	glog.V(2).Infof("%v: opened new file", fd)
	if !streamFromStart {
		if _, err := fd.Seek(0, io.SeekEnd); err != nil {
			logErrors.Add(fs.pathname, 1)
			if err := fd.Close(); err != nil {
				logErrors.Add(fs.pathname, 1)
				glog.Info(err)
			}
			return err
		}
		glog.V(2).Infof("%v: seeked to end", fd)
	}
	b := make([]byte, defaultReadBufferSize)
	var lastBytes []byte
	partial := bytes.NewBufferString("")
	started := make(chan struct{})
	var total int
	wg.Add(1)
	go func() {
		defer wg.Done()
		defer func() {
			glog.V(2).Infof("%v: read total %d bytes from %s", fd, total, fs.pathname)
			glog.V(2).Infof("%v: closing file descriptor", fd)
			if err := fd.Close(); err != nil {
				logErrors.Add(fs.pathname, 1)
				glog.Info(err)
			}
			logCloses.Add(fs.pathname, 1)
		}()
		close(started)
		for {
			// Blocking read but regular files will return EOF straight away.
			count, err := fd.Read(b)
			glog.V(2).Infof("%v: read %d bytes, err is %v", fd, count, err)

			if count > 0 {
				total += count
				glog.V(2).Infof("%v: decode and send", fd)
				needSend := lastBytes
				needSend = append(needSend, b[:count]...)
				sendCount := decodeAndSend(ctx, fs.lines, fs.pathname, len(needSend), needSend, partial)
				if sendCount < len(needSend) {
					lastBytes = append([]byte{}, needSend[sendCount:]...)
				} else {
					lastBytes = []byte{}
				}
				fs.mu.Lock()
				fs.lastReadTime = time.Now()
				fs.mu.Unlock()
			}

			if err != nil && err != io.EOF {
				logErrors.Add(fs.pathname, 1)
				// TODO: This could be generalised to check for any retryable
				// errors, and end on unretriables; e.g. ESTALE looks
				// retryable.
				if errors.Is(err, syscall.ESTALE) {
					glog.Infof("%v: reopening stream due to %s", fd, err)
					if nerr := fs.stream(ctx, wg, waker, fi, true); nerr != nil {
						glog.Info(nerr)
					}
					// Close this stream.
					return
				}
				glog.Info(err)
			}

			// If we have read no bytes and are at EOF, check for truncation and rotation.
			if err == io.EOF && count == 0 {
				glog.V(2).Infof("%v: eof an no bytes", fd)
				// Both rotation and truncation need to stat, so check for
				// rotation first.  It is assumed that rotation is the more
				// common change pattern anyway.
				newfi, serr := os.Stat(fs.pathname)
				if serr != nil {
					glog.Info(serr)
					// If this is a NotExist error, then we should wrap up this
					// goroutine. The Tailer will create a new logstream if the
					// file is in the middle of a rotation and gets recreated
					// in the next moment.  We can't rely on the Tailer to tell
					// us we're deleted because the tailer can only tell us to
					// Stop, which ends up causing us to race here against
					// detection of IsCompleted.
					if os.IsNotExist(serr) {
						glog.V(2).Infof("%v: source no longer exists, exiting", fd)
						if partial.Len() > 0 {
							sendLine(ctx, fs.pathname, partial, fs.lines)
						}
						fs.mu.Lock()
						fs.completed = true
						fs.mu.Unlock()
						return
					}
					logErrors.Add(fs.pathname, 1)
					goto Sleep
				}
				if !os.SameFile(fi, newfi) {
					glog.V(2).Infof("%v: adding a new file routine", fd)
					if err := fs.stream(ctx, wg, waker, newfi, true); err != nil {
						glog.Info(err)
					}
					// We're at EOF so there's nothing left to read here.
					return
				}
				currentOffset, serr := fd.Seek(0, io.SeekCurrent)
				if serr != nil {
					logErrors.Add(fs.pathname, 1)
					glog.Info(serr)
					continue
				}
				glog.V(2).Infof("%v: current seek is %d", fd, currentOffset)
				glog.V(2).Infof("%v: new size is %d", fd, newfi.Size())
				// We know that newfi is from the current file.  Truncation can
				// only be detected if the new file is currently shorter than
				// the current seek offset.  In test this can be a race, but in
				// production it's unlikely that a new file writes more bytes
				// than the previous after rotation in the time it takes for
				// mtail to notice.
				if newfi.Size() < currentOffset {
					glog.V(2).Infof("%v: truncate? currentoffset is %d and size is %d", fd, currentOffset, newfi.Size())
					// About to lose all remaining data because of the truncate so flush the accumulator.
					if partial.Len() > 0 {
						sendLine(ctx, fs.pathname, partial, fs.lines)
					}
					p, serr := fd.Seek(0, io.SeekStart)
					if serr != nil {
						logErrors.Add(fs.pathname, 1)
						glog.Info(serr)
					}
					glog.V(2).Infof("%v: Seeked to %d", fd, p)
					fileTruncates.Add(fs.pathname, 1)
					continue
				}
			}

			// No error implies there is more to read in this file so go
			// straight back to read unless it looks like context is Done.
			if err == nil && ctx.Err() == nil {
				continue
			}

		Sleep:
			// If we get here it's because we've stalled.  First test to see if it's
			// time to exit.
			if err == io.EOF || ctx.Err() != nil {
				select {
				case <-fs.stopChan:
					glog.V(2).Infof("%v: stream has been stopped, exiting", fd)
					if partial.Len() > 0 {
						sendLine(ctx, fs.pathname, partial, fs.lines)
					}
					fs.mu.Lock()
					fs.completed = true
					fs.mu.Unlock()
					return
				case <-ctx.Done():
					glog.V(2).Infof("%v: stream has been cancelled, exiting", fd)
					if partial.Len() > 0 {
						sendLine(ctx, fs.pathname, partial, fs.lines)
					}
					fs.mu.Lock()
					fs.completed = true
					fs.mu.Unlock()
					return
				default:
					// keep going
				}
			}

			// Don't exit, instead yield and wait for a termination signal or
			// wakeup.
			glog.V(2).Infof("%v: waiting", fd)
			select {
			case <-fs.stopChan:
				// We may have started waiting here when the stop signal
				// arrives, but since that wait the file may have been
				// written to.  The file is not technically yet at EOF so
				// we need to go back and try one more read.  We'll exit
				// the stream in the select stanza above.
				glog.V(2).Infof("%v: Stopping after next read", fd)
			case <-ctx.Done():
				// Same for cancellation; this makes tests stable, but
				// could argue exiting immediately is less surprising.
				// Assumption is that this doesn't make a difference in
				// production.
				glog.V(2).Infof("%v: Cancelled after next read", fd)
			case <-waker.Wake():
				// sleep until next Wake()
				glog.V(2).Infof("%v: Wake received", fd)
			}
		}
	}()

	<-started
	return nil
}

func (fs *fileStream) IsComplete() bool {
	fs.mu.RLock()
	defer fs.mu.RUnlock()
	return fs.completed
}

// Stop implements the LogStream interface.
func (fs *fileStream) Stop() {
	fs.stopOnce.Do(func() {
		glog.Info("signalling stop at next EOF")
		close(fs.stopChan)
	})
}
<|endoftext|>


<commit_before>/* https://leetcode.com/problems/regular-expression-matching/#/description
Implement regular expression matching with support for '.' and '*'.

    '.' Matches any single character.
    '*' Matches zero or more of the preceding element.

    The matching should cover the entire input string (not partial).

    The function prototype should be:
    bool isMatch(const char *s, const char *p)

    Some examples:
    isMatch("aa","a") → false
    isMatch("aa","aa") → true
    isMatch("aaa","aa") → false
    isMatch("aa", "a*") → true
    isMatch("aa", ".*") → true
    isMatch("ab", ".*") → true
    isMatch("aab", "c*a*b") → true
*/

package leetcode

func isMatch(s string, p string) bool {
	return false
}
<commit_msg>add isMatch: recursion solution<commit_after>/* https://leetcode.com/problems/regular-expression-matching/#/description
Implement regular expression matching with support for '.' and '*'.

    '.' Matches any single character.
    '*' Matches zero or more of the preceding element.

    The matching should cover the entire input string (not partial).

    The function prototype should be:
    bool isMatch(const char *s, const char *p)

    Some examples:
    isMatch("aa","a") → false
    isMatch("aa","aa") → true
    isMatch("aaa","aa") → false
    isMatch("aa", "a*") → true
    isMatch("aa", ".*") → true
    isMatch("ab", ".*") → true
    isMatch("aab", "c*a*b") → true
*/

package leetcode

func isMatch(s string, p string) bool {
	// 递归实现
	lens, lenp := len(s), len(p)
	if lenp == 0 {
		return lens == 0
	}

	if lenp <= 1 || (lenp > 1 && p[1] != '*') {
		return lens != 0 && (p[0] == s[0] || p[0] == '.') && isMatch(s[1:], p[1:])
	}

	slice := 2
	if len(p) < 2 {
		slice = 1
	}

	for len(s) != 0 && (p[0] == s[0] || p[0] == '.') {
		if isMatch(s, p[slice:]) {
			return true
		}
		s = s[1:]
	}

	return isMatch(s, p[slice:])
}
<|endoftext|>


<commit_before>package couchbase

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"log"
	"net/http"
	"net/url"
	"runtime"
	"sort"
	"strings"
)

// The HTTP Client To Use
var HttpClient = http.DefaultClient

// Auth callback gets the auth username and password for the given
// bucket.
type AuthHandler interface {
	GetCredentials() (string, string)
}

type RestPool struct {
	Name         string `json:"name"`
	StreamingURI string `json:"streamingUri"`
	URI          string `json:"uri"`
}

type Pools struct {
	ComponentsVersion     map[string]string `json:"componentsVersion,omitempty"`
	ImplementationVersion string            `json:"implementationVersion"`
	IsAdmin               bool              `json:"isAdminCreds"`
	UUID                  string            `json:"uuid"`
	Pools                 []RestPool        `json:"pools"`
}

// A computer in a cluster running the couchbase software.
type Node struct {
	ClusterCompatibility int                `json:"clusterCompatibility"`
	ClusterMembership    string             `json:"clusterMembership"`
	CouchAPIBase         string             `json:"couchApiBase"`
	Hostname             string             `json:"hostname"`
	InterestingStats     map[string]float64 `json:"interestingStats,omitempty"`
	MCDMemoryAllocated   float64            `json:"mcdMemoryAllocated"`
	MCDMemoryReserved    float64            `json:"mcdMemoryReserved"`
	MemoryFree           float64            `json:"memoryFree"`
	MemoryTotal          float64            `json:"memoryTotal"`
	OS                   string             `json:"os"`
	Ports                map[string]int     `json:"ports"`
	Status               string             `json:"status"`
	Uptime               int                `json:"uptime,string"`
	Version              string             `json:"version"`
	ThisNode             bool               `json:"thisNode,omitempty"`
}

// A pool of nodes and buckets.
type Pool struct {
	BucketMap map[string]Bucket
	Nodes     []Node

	BucketURL map[string]string `json:"buckets"`

	client Client
}

// An individual bucket.  Herein lives the most useful stuff.
type Bucket struct {
	AuthType            string             `json:"authType"`
	Capabilities        []string           `json:"bucketCapabilities"`
	CapabilitiesVersion string             `json:"bucketCapabilitiesVer"`
	Type                string             `json:"bucketType"`
	Name                string             `json:"name"`
	NodeLocator         string             `json:"nodeLocator"`
	Nodes               []Node             `json:"nodes"`
	Quota               map[string]float64 `json:"quota,omitempty"`
	Replicas            int                `json:"replicaNumber"`
	Password            string             `json:"saslPassword"`
	URI                 string             `json:"uri"`
	StreamingURI        string             `json:"streamingUri"`
	LocalRandomKeyURI   string             `json:"localRandomKeyUri,omitempty"`
	UUID                string             `json:"uuid"`
	DDocs               struct {
		URI string `json:"uri"`
	} `json:"ddocs,omitempty"`
	VBucketServerMap struct {
		HashAlgorithm string   `json:"hashAlgorithm"`
		NumReplicas   int      `json:"numReplicas"`
		ServerList    []string `json:"serverList"`
		VBucketMap    [][]int  `json:"vBucketMap"`
	} `json:"vBucketServerMap"`
	BasicStats  map[string]interface{} `json:"basicStats,omitempty"`
	Controllers map[string]interface{} `json:"controllers,omitempty"`

	pool        *Pool
	connections []*connectionPool
	commonSufix string
	auth        AuthHandler
}

// Get the (sorted) list of memcached node addresses (hostname:port).
func (b Bucket) NodeAddresses() []string {
	rv := make([]string, len(b.VBucketServerMap.ServerList))
	copy(rv, b.VBucketServerMap.ServerList)
	sort.Strings(rv)
	return rv
}

// Get the longest common suffix of all host:port strings in the node list.
func (b Bucket) CommonAddressSuffix() string {
	input := []string{}
	for _, n := range b.Nodes {
		input = append(input, n.Hostname)
	}
	return FindCommonSuffix(input)
}

// The couchbase client gives access to all the things.
type Client struct {
	BaseURL  *url.URL
	Info     Pools
	Statuses [256]uint64
}

func maybeAddAuth(req *http.Request, ah AuthHandler) {
	if ah != nil {
		user, pass := ah.GetCredentials()
		req.Header.Set("Authorization", "Basic "+
			base64.StdEncoding.EncodeToString([]byte(user+":"+pass)))
	}
}

func (c *Client) parseURLResponse(path string, ah AuthHandler,
	out interface{}) error {

	u := *c.BaseURL
	u.User = nil
	if q := strings.Index(path, "?"); q > 0 {
		u.Path = path[:q]
		u.RawQuery = path[q+1:]
	} else {
		u.Path = path
	}

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return err
	}
	maybeAddAuth(req, ah)

	res, err := HttpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	d := json.NewDecoder(res.Body)
	if err = d.Decode(&out); err != nil {
		return err
	}
	return nil
}

type basicAuth struct {
	u, p string
}

func (b basicAuth) GetCredentials() (string, string) {
	return b.u, b.p
}

// Connect to a couchbase cluster.
func Connect(baseU string) (c Client, err error) {
	c.BaseURL, err = url.Parse(baseU)
	if err != nil {
		return
	}

	return c, c.parseURLResponse("/pools", nil, &c.Info)
}

func (b *Bucket) refresh() (err error) {
	pool := b.pool
	err = pool.client.parseURLResponse(b.URI, b.auth, b)
	if err != nil {
		return err
	}
	b.pool = pool
	for i := range b.connections {
		b.connections[i] = newConnectionPool(
			b.VBucketServerMap.ServerList[i], b.auth, 4)
	}
	return nil
}

func (p *Pool) refresh() (err error) {
	p.BucketMap = make(map[string]Bucket)

	var ah AuthHandler
	if user := p.client.BaseURL.User; user != nil {
		pw, _ := user.Password()
		ah = basicAuth{user.Username(), pw}
	}

	buckets := []Bucket{}
	err = p.client.parseURLResponse(p.BucketURL["uri"], ah, &buckets)
	if err != nil {
		return err
	}
	for _, b := range buckets {
		b.pool = p
		b.auth = p.getDefaultAuth(b.Name)
		b.connections = make([]*connectionPool, len(b.VBucketServerMap.ServerList))

		p.BucketMap[b.Name] = b
	}
	return nil
}

// Get a pool from within the couchbase cluster (usually "default").
func (c *Client) GetPool(name string) (p Pool, err error) {
	var poolURI string
	for _, p := range c.Info.Pools {
		if p.Name == name {
			poolURI = p.URI
		}
	}
	if poolURI == "" {
		return p, errors.New("No pool named " + name)
	}

	err = c.parseURLResponse(poolURI, nil, &p)

	p.client = *c

	p.refresh()
	return
}

// Mark this bucket as no longer needed, closing connections it may have open.
func (b *Bucket) Close() {
	if b.connections != nil {
		for _, c := range b.connections {
			if c != nil {
				c.Close()
			}
		}
		b.connections = nil
	}
}

func bucket_finalizer(b *Bucket) {
	if b.connections != nil {
		log.Printf("Warning: Finalizing a bucket with active connections.")
	}
}

func (p *Pool) getDefaultAuth(name string) AuthHandler {
	var pw string
	if p.client.BaseURL.User != nil {
		pw, _ = p.client.BaseURL.User.Password()
	}
	return &basicAuth{name, pw}
}

// Get a bucket from within this pool.
func (p *Pool) GetBucket(name string) (*Bucket, error) {
	rv, ok := p.BucketMap[name]
	if !ok {
		return nil, errors.New("No bucket named " + name)
	}
	runtime.SetFinalizer(&rv, bucket_finalizer)
	rv.refresh()
	rv.auth = p.getDefaultAuth(name)
	return &rv, nil
}

// Get the pool to which this bucket belongs.
func (b *Bucket) GetPool() *Pool {
	return b.pool
}

// Get the client from which we got this pool.
func (p *Pool) GetClient() *Client {
	return &p.client
}

// Convenience function for getting a named bucket from a URL
func GetBucket(endpoint, poolname, bucketname string) (*Bucket, error) {
	var err error
	client, err := Connect(endpoint)
	if err != nil {
		return nil, err
	}

	pool, err := client.GetPool(poolname)
	if err != nil {
		return nil, err
	}

	return pool.GetBucket(bucketname)
}
<commit_msg>Better HTTP error messages<commit_after>package couchbase

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"runtime"
	"sort"
	"strings"
)

// The HTTP Client To Use
var HttpClient = http.DefaultClient

// Auth callback gets the auth username and password for the given
// bucket.
type AuthHandler interface {
	GetCredentials() (string, string)
}

type RestPool struct {
	Name         string `json:"name"`
	StreamingURI string `json:"streamingUri"`
	URI          string `json:"uri"`
}

type Pools struct {
	ComponentsVersion     map[string]string `json:"componentsVersion,omitempty"`
	ImplementationVersion string            `json:"implementationVersion"`
	IsAdmin               bool              `json:"isAdminCreds"`
	UUID                  string            `json:"uuid"`
	Pools                 []RestPool        `json:"pools"`
}

// A computer in a cluster running the couchbase software.
type Node struct {
	ClusterCompatibility int                `json:"clusterCompatibility"`
	ClusterMembership    string             `json:"clusterMembership"`
	CouchAPIBase         string             `json:"couchApiBase"`
	Hostname             string             `json:"hostname"`
	InterestingStats     map[string]float64 `json:"interestingStats,omitempty"`
	MCDMemoryAllocated   float64            `json:"mcdMemoryAllocated"`
	MCDMemoryReserved    float64            `json:"mcdMemoryReserved"`
	MemoryFree           float64            `json:"memoryFree"`
	MemoryTotal          float64            `json:"memoryTotal"`
	OS                   string             `json:"os"`
	Ports                map[string]int     `json:"ports"`
	Status               string             `json:"status"`
	Uptime               int                `json:"uptime,string"`
	Version              string             `json:"version"`
	ThisNode             bool               `json:"thisNode,omitempty"`
}

// A pool of nodes and buckets.
type Pool struct {
	BucketMap map[string]Bucket
	Nodes     []Node

	BucketURL map[string]string `json:"buckets"`

	client Client
}

// An individual bucket.  Herein lives the most useful stuff.
type Bucket struct {
	AuthType            string             `json:"authType"`
	Capabilities        []string           `json:"bucketCapabilities"`
	CapabilitiesVersion string             `json:"bucketCapabilitiesVer"`
	Type                string             `json:"bucketType"`
	Name                string             `json:"name"`
	NodeLocator         string             `json:"nodeLocator"`
	Nodes               []Node             `json:"nodes"`
	Quota               map[string]float64 `json:"quota,omitempty"`
	Replicas            int                `json:"replicaNumber"`
	Password            string             `json:"saslPassword"`
	URI                 string             `json:"uri"`
	StreamingURI        string             `json:"streamingUri"`
	LocalRandomKeyURI   string             `json:"localRandomKeyUri,omitempty"`
	UUID                string             `json:"uuid"`
	DDocs               struct {
		URI string `json:"uri"`
	} `json:"ddocs,omitempty"`
	VBucketServerMap struct {
		HashAlgorithm string   `json:"hashAlgorithm"`
		NumReplicas   int      `json:"numReplicas"`
		ServerList    []string `json:"serverList"`
		VBucketMap    [][]int  `json:"vBucketMap"`
	} `json:"vBucketServerMap"`
	BasicStats  map[string]interface{} `json:"basicStats,omitempty"`
	Controllers map[string]interface{} `json:"controllers,omitempty"`

	pool        *Pool
	connections []*connectionPool
	commonSufix string
	auth        AuthHandler
}

// Get the (sorted) list of memcached node addresses (hostname:port).
func (b Bucket) NodeAddresses() []string {
	rv := make([]string, len(b.VBucketServerMap.ServerList))
	copy(rv, b.VBucketServerMap.ServerList)
	sort.Strings(rv)
	return rv
}

// Get the longest common suffix of all host:port strings in the node list.
func (b Bucket) CommonAddressSuffix() string {
	input := []string{}
	for _, n := range b.Nodes {
		input = append(input, n.Hostname)
	}
	return FindCommonSuffix(input)
}

// The couchbase client gives access to all the things.
type Client struct {
	BaseURL  *url.URL
	Info     Pools
	Statuses [256]uint64
}

func maybeAddAuth(req *http.Request, ah AuthHandler) {
	if ah != nil {
		user, pass := ah.GetCredentials()
		req.Header.Set("Authorization", "Basic "+
			base64.StdEncoding.EncodeToString([]byte(user+":"+pass)))
	}
}

func (c *Client) parseURLResponse(path string, ah AuthHandler,
	out interface{}) error {

	u := *c.BaseURL
	u.User = nil
	if q := strings.Index(path, "?"); q > 0 {
		u.Path = path[:q]
		u.RawQuery = path[q+1:]
	} else {
		u.Path = path
	}

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return err
	}
	maybeAddAuth(req, ah)

	res, err := HttpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		bod, _ := ioutil.ReadAll(io.LimitReader(res.Body, 512))
		return fmt.Errorf("HTTP error %v getting %q: %s",
			res.Status, u.String(), bod)
	}

	d := json.NewDecoder(res.Body)
	if err = d.Decode(&out); err != nil {
		return err
	}
	return nil
}

type basicAuth struct {
	u, p string
}

func (b basicAuth) GetCredentials() (string, string) {
	return b.u, b.p
}

// Connect to a couchbase cluster.
func Connect(baseU string) (c Client, err error) {
	c.BaseURL, err = url.Parse(baseU)
	if err != nil {
		return
	}

	return c, c.parseURLResponse("/pools", nil, &c.Info)
}

func (b *Bucket) refresh() (err error) {
	pool := b.pool
	err = pool.client.parseURLResponse(b.URI, b.auth, b)
	if err != nil {
		return err
	}
	b.pool = pool
	for i := range b.connections {
		b.connections[i] = newConnectionPool(
			b.VBucketServerMap.ServerList[i], b.auth, 4)
	}
	return nil
}

func (p *Pool) refresh() (err error) {
	p.BucketMap = make(map[string]Bucket)

	var ah AuthHandler
	if user := p.client.BaseURL.User; user != nil {
		pw, _ := user.Password()
		ah = basicAuth{user.Username(), pw}
	}

	buckets := []Bucket{}
	err = p.client.parseURLResponse(p.BucketURL["uri"], ah, &buckets)
	if err != nil {
		return err
	}
	for _, b := range buckets {
		b.pool = p
		b.auth = p.getDefaultAuth(b.Name)
		b.connections = make([]*connectionPool, len(b.VBucketServerMap.ServerList))

		p.BucketMap[b.Name] = b
	}
	return nil
}

// Get a pool from within the couchbase cluster (usually "default").
func (c *Client) GetPool(name string) (p Pool, err error) {
	var poolURI string
	for _, p := range c.Info.Pools {
		if p.Name == name {
			poolURI = p.URI
		}
	}
	if poolURI == "" {
		return p, errors.New("No pool named " + name)
	}

	err = c.parseURLResponse(poolURI, nil, &p)

	p.client = *c

	p.refresh()
	return
}

// Mark this bucket as no longer needed, closing connections it may have open.
func (b *Bucket) Close() {
	if b.connections != nil {
		for _, c := range b.connections {
			if c != nil {
				c.Close()
			}
		}
		b.connections = nil
	}
}

func bucket_finalizer(b *Bucket) {
	if b.connections != nil {
		log.Printf("Warning: Finalizing a bucket with active connections.")
	}
}

func (p *Pool) getDefaultAuth(name string) AuthHandler {
	var pw string
	if p.client.BaseURL.User != nil {
		pw, _ = p.client.BaseURL.User.Password()
	}
	return &basicAuth{name, pw}
}

// Get a bucket from within this pool.
func (p *Pool) GetBucket(name string) (*Bucket, error) {
	rv, ok := p.BucketMap[name]
	if !ok {
		return nil, errors.New("No bucket named " + name)
	}
	runtime.SetFinalizer(&rv, bucket_finalizer)
	rv.refresh()
	rv.auth = p.getDefaultAuth(name)
	return &rv, nil
}

// Get the pool to which this bucket belongs.
func (b *Bucket) GetPool() *Pool {
	return b.pool
}

// Get the client from which we got this pool.
func (p *Pool) GetClient() *Client {
	return &p.client
}

// Convenience function for getting a named bucket from a URL
func GetBucket(endpoint, poolname, bucketname string) (*Bucket, error) {
	var err error
	client, err := Connect(endpoint)
	if err != nil {
		return nil, err
	}

	pool, err := client.GetPool(poolname)
	if err != nil {
		return nil, err
	}

	return pool.GetBucket(bucketname)
}
<|endoftext|>


<commit_before>package liberty

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"time"

	"golang.scot/liberty/middleware"
)

type Config struct {
	Certs     []*Crt                     `yaml:"certs"`
	Proxies   []*ReverseProxy            `yaml:"proxies"`
	Whitelist []*middleware.ApiWhitelist `yaml:"whitelist"`
}

// Proxy is a reverse HTTP proxy
type Proxy struct {
	config   *Config
	group    *ServerGroup
	secure   map[string]*VHost
	insecure map[string]*VHost
}

// NewProxy returns a Proxy configured for use
func NewProxy(config *Config) *Proxy {
	p := &Proxy{
		config:   config,
		secure:   map[string]*VHost{},
		insecure: map[string]*VHost{},
	}

	servers := make([]*http.Server, 0)

	for _, proxy := range p.config.Proxies {
		host, _ := proxy.hostAndPath()
		if _, ok := p.secure[host]; !ok {
			p.secure[host] = &VHost{
				host:    host,
				handler: NewRouter(),
			}
		}
		if _, ok := p.insecure[host]; !ok {
			p.insecure[host] = &VHost{
				host:    host,
				handler: http.HandlerFunc(middleware.RedirectPerm),
			}
		}

		err := proxy.Configure(p.config.Whitelist, p.secure[host].handler, p.serveInsecure)
		if err != nil {
			fmt.Printf("the proxy for '%s' was not configured - %s\n", proxy.HostPath, err)
			continue
		}

		servers = append(servers, proxy.Servers...)
	}

	p.group = NewServerGroup(p, servers)
	p.group.setTLSConfig(p.vhostDomains())

	return p
}

func (b *Proxy) vhostDomains() []string {
	domains := make([]string, 0)
	for host, _ := range b.secure {
		domains = append(domains, host)
	}

	return domains
}

const gracePriod = 5 // seconds

// Serve incoming requests between a set of configured reverse proxies, uses
// kernel SO_REUSEPORT which conveniently maps incoming connections to the least
// used socket
func (b *Proxy) Serve() {
	sig := make(chan os.Signal)

	startServer := func(s *server, wg *sync.WaitGroup) {
		fmt.Println("server lisening: ", s.s.Addr)
		log.Println(s.s.Serve(s.Listener()))
	}

	var wg sync.WaitGroup
	wg.Add(len(b.group.servers))

	for _, s := range b.group.servers {
		go startServer(s, &wg)
	}

	signal.Notify(sig, os.Interrupt, os.Kill)
	<-sig
	log.Println("Draining server connections...")
	for _, s := range b.group.servers {
		go func() {
			ctx, _ := context.WithTimeout(context.Background(), gracePriod*time.Second)
			s.s.Shutdown(ctx)
			wg.Done()
		}()
	}

	wg.Wait()
	log.Println("Done, exiting")
}

func (b *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if vhost, ok := b.secure[r.Host]; ok {
		fmt.Println(vhost, r.URL.Path)
		vhost.handler.ServeHTTP(w, r)
		return
	}

	http.NotFound(w, r)
}

func (b *Proxy) serveInsecure(w http.ResponseWriter, r *http.Request) {
	if s, ok := b.insecure[r.Host]; ok {
		s.handler.ServeHTTP(w, r)
		return
	}

	http.NotFound(w, r)
}
<commit_msg>chore: misc tidy<commit_after>package liberty

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"time"

	"golang.scot/liberty/middleware"
)

type Config struct {
	Certs     []*Crt                     `yaml:"certs"`
	Proxies   []*ReverseProxy            `yaml:"proxies"`
	Whitelist []*middleware.ApiWhitelist `yaml:"whitelist"`
}

// Proxy is a reverse HTTP proxy
type Proxy struct {
	config   *Config
	group    *ServerGroup
	secure   map[string]*VHost
	insecure map[string]*VHost
}

// NewProxy returns a Proxy configured for use
func NewProxy(config *Config) *Proxy {
	p := &Proxy{
		config:   config,
		secure:   map[string]*VHost{},
		insecure: map[string]*VHost{},
	}

	servers := make([]*http.Server, 0)

	for _, proxy := range p.config.Proxies {
		host, _ := proxy.hostAndPath()
		if _, ok := p.secure[host]; !ok {
			p.secure[host] = &VHost{
				host:    host,
				handler: NewRouter(),
			}
		}
		if _, ok := p.insecure[host]; !ok {
			p.insecure[host] = &VHost{
				host:    host,
				handler: http.HandlerFunc(middleware.RedirectPerm),
			}
		}

		err := proxy.Configure(p.config.Whitelist, p.secure[host].handler, p.serveInsecure)
		if err != nil {
			fmt.Printf("the proxy for '%s' was not configured - %s\n", proxy.HostPath, err)
			continue
		}

		servers = append(servers, proxy.Servers...)
	}

	p.group = NewServerGroup(p, servers)
	p.group.setTLSConfig(p.vhostDomains())

	return p
}

func (b *Proxy) vhostDomains() []string {
	domains := make([]string, 0)
	for host, _ := range b.secure {
		domains = append(domains, host)
	}

	return domains
}

const gracePriod = 5 // seconds

// Serve incoming requests between a set of configured reverse proxies, uses
// kernel SO_REUSEPORT which conveniently maps incoming connections to the least
// used socket
func (b *Proxy) Serve() {

	startServer := func(s *server, wg *sync.WaitGroup) {
		fmt.Println("server lisening: ", s.s.Addr)
		log.Println(s.s.Serve(s.Listener()))
	}

	var wg sync.WaitGroup
	wg.Add(len(b.group.servers))

	for _, s := range b.group.servers {
		go startServer(s, &wg)
	}

	sig := make(chan os.Signal)
	signal.Notify(sig, os.Interrupt, os.Kill)
	<-sig

	log.Println("Draining server connections...")
	for _, s := range b.group.servers {
		go func() {
			ctx, _ := context.WithTimeout(context.Background(), gracePriod*time.Second)
			s.s.Shutdown(ctx)
			wg.Done()
		}()
	}

	wg.Wait()
	log.Println("Done, exiting")
}

func (b *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if vhost, ok := b.secure[r.Host]; ok {
		fmt.Println(vhost, r.URL.Path)
		vhost.handler.ServeHTTP(w, r)
		return
	}

	http.NotFound(w, r)
}

func (b *Proxy) serveInsecure(w http.ResponseWriter, r *http.Request) {
	if s, ok := b.insecure[r.Host]; ok {
		s.handler.ServeHTTP(w, r)
		return
	}

	http.NotFound(w, r)
}
<|endoftext|>


<commit_before>package proxy

import (
	"errors"
	"net/http"
	"net/http/httputil"
	"net/url"
	"path"
	"time"

	"github.com/gorilla/websocket"
	"github.com/koding/websocketproxy"
)

type ProxyOptions struct {
	// Number of times a request should be tried
	Retries int

	// Period to wait between retries
	Period time.Duration

	// Returns a url that we should proxy to for a given request
	Balancer func(req *http.Request) (string, error)

	// A static backend to route to
	Backend string

	// Rewriter is a function modifying the request before it's proxied
	// usually it should change the hostname of the request object, etc ...
	Rewriter func(req *http.Request, destUrl *url.URL)
}

type Proxy struct {
	*ProxyOptions

	// Http proxy
	httpProxy http.Handler

	// Websocket proxy
	websocketProxy http.Handler
}

// New returns a new Proxy instance based on the provided ProxyOptions
// either 'Backend' (static) or 'Balancer' must be provided
func New(opts ProxyOptions) (*Proxy, error) {
	// Validate Balancer and Backend options
	if opts.Balancer == nil {
		if opts.Backend == "" {
			return nil, errors.New("Please provide a Backend or a Balancer")
		} else if err := validateUrl(normalizeUrl(opts.Backend)); err != nil {
			return nil, err
		} else {
			// Normalize backend's url
			opts.Backend = normalizeUrl(opts.Backend)
		}
	}

	// Default for Retries
	if opts.Retries == 0 {
		opts.Retries = 1
	}

	// Default for Period
	if opts.Period == 0 {
		opts.Period = 100 * time.Millisecond
	}

	// Default for Rewriter
	if opts.Rewriter == nil {
		opts.Rewriter = DefaultRewriter
	}

	p := &Proxy{
		ProxyOptions: &opts,
	}

	return p.init(), nil
}

// ServeHTTP allows us to comply to the http.Handler interface
func (p *Proxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	if isWebsocket(req) {
		// we don't use https explicitly, ssl termination is done here
		req.URL.Scheme = "ws"
		p.websocketProxy.ServeHTTP(rw, req)
		return
	}

	p.httpProxy.ServeHTTP(rw, req)
}

// init sets up proxies and other stuff based on options
func (p *Proxy) init() *Proxy {
	// Setup http proxy
	p.httpProxy = &httputil.ReverseProxy{
		Director: p.director,
	}

	// Setup websocket proxy
	p.websocketProxy = &websocketproxy.WebsocketProxy{
		Backend: func(req *http.Request) *url.URL {
			url, _ := p.backend(req)
			p.Rewriter(req, url)
			return url
		},
		Upgrader: &websocket.Upgrader{
			ReadBufferSize:  4096,
			WriteBufferSize: 4096,
			CheckOrigin: func(r *http.Request) bool {
				return true
			},
		},
	}

	return p
}

// director rewrites a http.Request to route to the correct host
func (p *Proxy) director(req *http.Request) {
	url, err := p.backend(req)
	if url == nil || err != nil {
		return
	}

	p.Rewriter(req, url)
}

// backend return the full url object of the backend to route
// after fetching the host backend and resolving it relative
// to the current http request data
func (p *Proxy) backend(req *http.Request) (*url.URL, error) {
	hostUrl, err := p.backendHost(req)
	if err != nil {
		return nil, err
	}

	return fullBackendURL(req, hostUrl)
}

// backendHost parses the result of getBackend and ensures it's validity
func (p *Proxy) backendHost(req *http.Request) (*url.URL, error) {
	rawurl, err := p.getBackend(req)
	if err != nil {
		return nil, err
	}

	// Normalize URL
	backendUrl := normalizeUrl(rawurl)

	if err := validateUrl(backendUrl); err != nil {
		return nil, err
	}

	return url.Parse(backendUrl)
}

// getBackend gets the backend selected by the balancer or the static one set by the 'Backend' attribute
func (p *Proxy) getBackend(req *http.Request) (string, error) {
	if p.Balancer == nil && p.Backend != "" {
		return p.Backend, nil
	}
	return p.Balancer(req)
}

// fullBackendURL builds a url.URL object representing the destination's full URL
// resolving path's etc ... relative to the request
func fullBackendURL(req *http.Request, hostUrl *url.URL) (*url.URL, error) {
	// Pick scheme differently depending on if the request is http or websocket
	schemeFunc := httpScheme
	if isWebsocket(req) {
		schemeFunc = websocketScheme
	}

	newUrl, err := url.Parse(hostUrl.String())
	if err != nil {
		return nil, err
	}

	newUrl.Scheme = schemeFunc(hostUrl.Scheme)
	newUrl.Host = hostUrl.Host
	newUrl.Path = path.Join(hostUrl.Path, req.URL.Path)

	return newUrl, nil
}

// DefaultRewriter adjust the *http.Request object it's attributes
// match those of the destionation url (scheme, host, path, ...)
func DefaultRewriter(req *http.Request, url *url.URL) {
	req.URL.Scheme = url.Scheme
	req.URL.Host = url.Host
	req.URL.Path = url.Path

	// Rewrite outgoing request url
	req.Host = url.Host
}
<commit_msg>Remove unused code forcing "ws" for web sockets<commit_after>package proxy

import (
	"errors"
	"net/http"
	"net/http/httputil"
	"net/url"
	"path"
	"time"

	"github.com/gorilla/websocket"
	"github.com/koding/websocketproxy"
)

type ProxyOptions struct {
	// Number of times a request should be tried
	Retries int

	// Period to wait between retries
	Period time.Duration

	// Returns a url that we should proxy to for a given request
	Balancer func(req *http.Request) (string, error)

	// A static backend to route to
	Backend string

	// Rewriter is a function modifying the request before it's proxied
	// usually it should change the hostname of the request object, etc ...
	Rewriter func(req *http.Request, destUrl *url.URL)
}

type Proxy struct {
	*ProxyOptions

	// Http proxy
	httpProxy http.Handler

	// Websocket proxy
	websocketProxy http.Handler
}

// New returns a new Proxy instance based on the provided ProxyOptions
// either 'Backend' (static) or 'Balancer' must be provided
func New(opts ProxyOptions) (*Proxy, error) {
	// Validate Balancer and Backend options
	if opts.Balancer == nil {
		if opts.Backend == "" {
			return nil, errors.New("Please provide a Backend or a Balancer")
		} else if err := validateUrl(normalizeUrl(opts.Backend)); err != nil {
			return nil, err
		} else {
			// Normalize backend's url
			opts.Backend = normalizeUrl(opts.Backend)
		}
	}

	// Default for Retries
	if opts.Retries == 0 {
		opts.Retries = 1
	}

	// Default for Period
	if opts.Period == 0 {
		opts.Period = 100 * time.Millisecond
	}

	// Default for Rewriter
	if opts.Rewriter == nil {
		opts.Rewriter = DefaultRewriter
	}

	p := &Proxy{
		ProxyOptions: &opts,
	}

	return p.init(), nil
}

// ServeHTTP allows us to comply to the http.Handler interface
func (p *Proxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	if isWebsocket(req) {
		p.websocketProxy.ServeHTTP(rw, req)
		return
	}

	p.httpProxy.ServeHTTP(rw, req)
}

// init sets up proxies and other stuff based on options
func (p *Proxy) init() *Proxy {
	// Setup http proxy
	p.httpProxy = &httputil.ReverseProxy{
		Director: p.director,
	}

	// Setup websocket proxy
	p.websocketProxy = &websocketproxy.WebsocketProxy{
		Backend: func(req *http.Request) *url.URL {
			url, _ := p.backend(req)
			p.Rewriter(req, url)
			return url
		},
		Upgrader: &websocket.Upgrader{
			ReadBufferSize:  4096,
			WriteBufferSize: 4096,
			CheckOrigin: func(r *http.Request) bool {
				return true
			},
		},
	}

	return p
}

// director rewrites a http.Request to route to the correct host
func (p *Proxy) director(req *http.Request) {
	url, err := p.backend(req)
	if url == nil || err != nil {
		return
	}

	p.Rewriter(req, url)
}

// backend return the full url object of the backend to route
// after fetching the host backend and resolving it relative
// to the current http request data
func (p *Proxy) backend(req *http.Request) (*url.URL, error) {
	hostUrl, err := p.backendHost(req)
	if err != nil {
		return nil, err
	}

	return fullBackendURL(req, hostUrl)
}

// backendHost parses the result of getBackend and ensures it's validity
func (p *Proxy) backendHost(req *http.Request) (*url.URL, error) {
	rawurl, err := p.getBackend(req)
	if err != nil {
		return nil, err
	}

	// Normalize URL
	backendUrl := normalizeUrl(rawurl)

	if err := validateUrl(backendUrl); err != nil {
		return nil, err
	}

	return url.Parse(backendUrl)
}

// getBackend gets the backend selected by the balancer or the static one set by the 'Backend' attribute
func (p *Proxy) getBackend(req *http.Request) (string, error) {
	if p.Balancer == nil && p.Backend != "" {
		return p.Backend, nil
	}
	return p.Balancer(req)
}

// fullBackendURL builds a url.URL object representing the destination's full URL
// resolving path's etc ... relative to the request
func fullBackendURL(req *http.Request, hostUrl *url.URL) (*url.URL, error) {
	// Pick scheme differently depending on if the request is http or websocket
	schemeFunc := httpScheme
	if isWebsocket(req) {
		schemeFunc = websocketScheme
	}

	newUrl, err := url.Parse(hostUrl.String())
	if err != nil {
		return nil, err
	}

	newUrl.Scheme = schemeFunc(hostUrl.Scheme)
	newUrl.Host = hostUrl.Host
	newUrl.Path = path.Join(hostUrl.Path, req.URL.Path)

	return newUrl, nil
}

// DefaultRewriter adjust the *http.Request object it's attributes
// match those of the destionation url (scheme, host, path, ...)
func DefaultRewriter(req *http.Request, url *url.URL) {
	req.URL.Scheme = url.Scheme
	req.URL.Host = url.Host
	req.URL.Path = url.Path

	// Rewrite outgoing request url
	req.Host = url.Host
}
<|endoftext|>


<commit_before>// Copyright 2015 Igor Dolzhikov. All rights reserved.
// Use of this source code is governed by a license
// that can be found in the LICENSE file.

package spawn

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

const (
	cmdQueueCapacity = 100

	checkHealthTask = iota
	doJobTask
	repeatJobTask
)

// Queue data (queries, responses, etc)
type queue struct {
	id       string
	jobs     chan *queueJob
	task     chan int
	ask      chan struct{}
	response chan struct{}
	quit     chan struct{}
}

// queueJob produce a task which contains query/response and status (done)
type queueJob struct {
	done   chan struct{}
	query  chan []byte
	answer chan *http.Response
}

// queueBundle is the bundle for the queue data (queries, responses, etc)
type queueBundle struct {
	mutex   sync.Mutex
	records map[string]*queue
}

// check the queue, if it does not exist, create it
func (bundle *queueBundle) check(id string) (*queue, bool) {
	bundle.mutex.Lock()
	defer bundle.mutex.Unlock()

	// check for a new record
	_, ok := bundle.records[id]

	// if it is new
	if !ok {
		bundle.records[id] = &queue{
			id:       id,
			jobs:     make(chan *queueJob, MaxJobs),
			task:     make(chan int, MaxJobs),
			ask:      make(chan struct{}, cmdQueueCapacity),
			response: make(chan struct{}, cmdQueueCapacity),
			quit:     make(chan struct{}),
		}
		return bundle.records[id], false
	}

	// if it exists already
	return bundle.records[id], true
}

// remove the queue and stops the worker
func (bundle *queueBundle) remove(id string, timeout time.Duration) {
	bundle.mutex.Lock()
	defer bundle.mutex.Unlock()

	// if a queue exists, the worker must be stoped and a queue must be deleted
	if q, ok := bundle.records[id]; ok {

		// if the worker is alive
		if getResponse(q, timeout) {

			// send a 'quit' command to the worker
			q.quit <- struct{}{}

			// get a response from the worker
			<-q.response
		}

		// delete the queue
		delete(bundle.records, id)
	}
}

// getReponse method is waiting a response or get the false value if timeout
func getResponse(q *queue, timeout time.Duration) bool {
	ticker := time.NewTimer(time.Second * timeout)

	// sweeps unwanted responses if exist
	for {
		select {
		case <-q.response:
			continue
		default:
		}
		break
	}

	// Sends an ASK to the worker
	q.ask <- struct{}{}

	select {
	// Exit by timeout if the response does not get (worker is not alive)
	case <-ticker.C:
		// sweeps ask which sent before if exist
		select {
		case <-q.ask:
		default:
		}
		return false
	// Exit after the response has been received
	case <-q.response:
		return true
	}
}
<commit_msg>removed unused statuses<commit_after>// Copyright 2015 Igor Dolzhikov. All rights reserved.
// Use of this source code is governed by a license
// that can be found in the LICENSE file.

package spawn

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

const (
	cmdQueueCapacity = 100

	doJobTask = iota
)

// Queue data (queries, responses, etc)
type queue struct {
	id       string
	jobs     chan *queueJob
	task     chan int
	ask      chan struct{}
	response chan struct{}
	quit     chan struct{}
}

// queueJob produce a task which contains query/response and status (done)
type queueJob struct {
	done   chan struct{}
	query  chan []byte
	answer chan *http.Response
}

// queueBundle is the bundle for the queue data (queries, responses, etc)
type queueBundle struct {
	mutex   sync.Mutex
	records map[string]*queue
}

// check the queue, if it does not exist, create it
func (bundle *queueBundle) check(id string) (*queue, bool) {
	bundle.mutex.Lock()
	defer bundle.mutex.Unlock()

	// check for a new record
	_, ok := bundle.records[id]

	// if it is new
	if !ok {
		bundle.records[id] = &queue{
			id:       id,
			jobs:     make(chan *queueJob, MaxJobs),
			task:     make(chan int, MaxJobs),
			ask:      make(chan struct{}, cmdQueueCapacity),
			response: make(chan struct{}, cmdQueueCapacity),
			quit:     make(chan struct{}),
		}
		return bundle.records[id], false
	}

	// if it exists already
	return bundle.records[id], true
}

// remove the queue and stops the worker
func (bundle *queueBundle) remove(id string, timeout time.Duration) {
	bundle.mutex.Lock()
	defer bundle.mutex.Unlock()

	// if a queue exists, the worker must be stoped and a queue must be deleted
	if q, ok := bundle.records[id]; ok {

		// if the worker is alive
		if getResponse(q, timeout) {

			// send a 'quit' command to the worker
			q.quit <- struct{}{}

			// get a response from the worker
			<-q.response
		}

		// delete the queue
		delete(bundle.records, id)
	}
}

// getReponse method is waiting a response or get the false value if timeout
func getResponse(q *queue, timeout time.Duration) bool {
	ticker := time.NewTimer(time.Second * timeout)

	// sweeps unwanted responses if exist
	for {
		select {
		case <-q.response:
			continue
		default:
		}
		break
	}

	// Sends an ASK to the worker
	q.ask <- struct{}{}

	select {
	// Exit by timeout if the response does not get (worker is not alive)
	case <-ticker.C:
		// sweeps ask which sent before if exist
		select {
		case <-q.ask:
		default:
		}
		return false
	// Exit after the response has been received
	case <-q.response:
		return true
	}
}
<|endoftext|>


<commit_before>package nio

import "github.com/djherbis/buffer"

var empty struct{}

type Packet struct {
	data []byte
	ok   chan struct{}
}

func NewPacket() *Packet {
	return &Packet{
		ok: make(chan struct{}),
	}
}

type BufferQueue struct {
	top  *Packet
	data []byte
	buffer.Buffer
}

func NewBufferQueue(buf buffer.Buffer) *BufferQueue {
	return &BufferQueue{
		top:    NewPacket(),
		data:   make([]byte, 32*1024),
		Buffer: buf,
	}
}

func (q *BufferQueue) Len() (n int64) {
	return q.Buffer.Len()
}

func (q *BufferQueue) Peek() (b interface{}) {
	n, _ := q.Buffer.ReadAt(q.data, 0)
	q.top.data = q.data[:n]
	return q.top
}

func (q *BufferQueue) Pop() (b interface{}) {
	q.top.ok <- empty
	q.Buffer.FastForward(len(q.top.data))
	return q.top
}

func (q *BufferQueue) Push(b interface{}) {
	p := b.(*Packet)
	q.Write(p.data)
	p.ok <- empty
}
<commit_msg>Changed FastForward to FFwd<commit_after>package nio

import "github.com/djherbis/buffer"

var empty struct{}

type Packet struct {
	data []byte
	ok   chan struct{}
}

func NewPacket() *Packet {
	return &Packet{
		ok: make(chan struct{}),
	}
}

type BufferQueue struct {
	top  *Packet
	data []byte
	buffer.Buffer
}

func NewBufferQueue(buf buffer.Buffer) *BufferQueue {
	return &BufferQueue{
		top:    NewPacket(),
		data:   make([]byte, 32*1024),
		Buffer: buf,
	}
}

func (q *BufferQueue) Len() (n int64) {
	return q.Buffer.Len()
}

func (q *BufferQueue) Peek() (b interface{}) {
	n, _ := q.Buffer.ReadAt(q.data, 0)
	q.top.data = q.data[:n]
	return q.top
}

func (q *BufferQueue) Pop() (b interface{}) {
	q.top.ok <- empty
	q.Buffer.FFwd(len(q.top.data))
	return q.top
}

func (q *BufferQueue) Push(b interface{}) {
	p := b.(*Packet)
	q.Write(p.data)
	p.ok <- empty
}
<|endoftext|>


<commit_before>package shellquote

import (
	"bytes"
	"strings"
	"unicode/utf8"
)

// Join quotes each argument and joins them with a space.
// If passed to /bin/sh, the resulting string will be split back into the
// original arguments.
func Join(args ...string) string {
	var buf bytes.Buffer
	for i, arg := range args {
		if i != 0 {
			buf.WriteByte(' ')
		}
		quote(arg, &buf)
	}
	return buf.String()
}

const (
	specialChars      = "\\'\"`${[|&;<>()*?!"
	extraSpecialChars = " \t\r\n"
	prefixChars       = "~"
)

func quote(word string, buf *bytes.Buffer) {
	// We want to try to produce a "nice" output. As such, we will
	// backslash-escape most characters, but if we encounter a space, or if we
	// encounter an extra-special char (which doesn't work with
	// backslash-escaping) we switch over to quoting the whole word. We do this
	// with a space because it's typically easier for people to read multi-word
	// arguments when quoted with a space rather than with ugly backslashes
	// everywhere.
	origLen := buf.Len()

	cur, prev := word, word
	atStart := true
	for len(cur) > 0 {
		c, l := utf8.DecodeRuneInString(cur)
		cur = cur[l:]
		if strings.ContainsRune(specialChars, c) || (atStart && strings.ContainsRune(prefixChars, c)) {
			// copy the non-special chars up to this point
			if len(cur) < len(prev) {
				buf.WriteString(word[0 : len(prev)-len(cur)-l])
			}
			buf.WriteByte('\\')
			buf.WriteRune(c)
			prev = cur
		} else if strings.ContainsRune(extraSpecialChars, c) {
			// start over in quote mode
			buf.Truncate(origLen)
			goto quote
		}
		atStart = false
	}
	if len(prev) > 0 {
		buf.WriteString(prev)
	}
	return

quote:
	// quote mode
	// Use single-quotes, but if we find a single-quote in the word, we need
	// to terminate the string, emit an escaped quote, and start the string up
	// again
	inQuote := false
	for len(word) > 0 {
		i := strings.IndexRune(word, '\'')
		if i == -1 {
			break
		}
		if i > 0 {
			if !inQuote {
				buf.WriteByte('\'')
				inQuote = true
			}
			buf.WriteString(word[0:i])
			word = word[i+1:]
		}
		if inQuote {
			buf.WriteByte('\'')
			inQuote = false
		}
		buf.WriteString("\\'")
	}
	if len(word) > 0 {
		if !inQuote {
			buf.WriteByte('\'')
		}
		buf.WriteString(word)
		buf.WriteByte('\'')
	}
}
<commit_msg>Remove \r from the special chars list<commit_after>package shellquote

import (
	"bytes"
	"strings"
	"unicode/utf8"
)

// Join quotes each argument and joins them with a space.
// If passed to /bin/sh, the resulting string will be split back into the
// original arguments.
func Join(args ...string) string {
	var buf bytes.Buffer
	for i, arg := range args {
		if i != 0 {
			buf.WriteByte(' ')
		}
		quote(arg, &buf)
	}
	return buf.String()
}

const (
	specialChars      = "\\'\"`${[|&;<>()*?!"
	extraSpecialChars = " \t\n"
	prefixChars       = "~"
)

func quote(word string, buf *bytes.Buffer) {
	// We want to try to produce a "nice" output. As such, we will
	// backslash-escape most characters, but if we encounter a space, or if we
	// encounter an extra-special char (which doesn't work with
	// backslash-escaping) we switch over to quoting the whole word. We do this
	// with a space because it's typically easier for people to read multi-word
	// arguments when quoted with a space rather than with ugly backslashes
	// everywhere.
	origLen := buf.Len()

	cur, prev := word, word
	atStart := true
	for len(cur) > 0 {
		c, l := utf8.DecodeRuneInString(cur)
		cur = cur[l:]
		if strings.ContainsRune(specialChars, c) || (atStart && strings.ContainsRune(prefixChars, c)) {
			// copy the non-special chars up to this point
			if len(cur) < len(prev) {
				buf.WriteString(word[0 : len(prev)-len(cur)-l])
			}
			buf.WriteByte('\\')
			buf.WriteRune(c)
			prev = cur
		} else if strings.ContainsRune(extraSpecialChars, c) {
			// start over in quote mode
			buf.Truncate(origLen)
			goto quote
		}
		atStart = false
	}
	if len(prev) > 0 {
		buf.WriteString(prev)
	}
	return

quote:
	// quote mode
	// Use single-quotes, but if we find a single-quote in the word, we need
	// to terminate the string, emit an escaped quote, and start the string up
	// again
	inQuote := false
	for len(word) > 0 {
		i := strings.IndexRune(word, '\'')
		if i == -1 {
			break
		}
		if i > 0 {
			if !inQuote {
				buf.WriteByte('\'')
				inQuote = true
			}
			buf.WriteString(word[0:i])
			word = word[i+1:]
		}
		if inQuote {
			buf.WriteByte('\'')
			inQuote = false
		}
		buf.WriteString("\\'")
	}
	if len(word) > 0 {
		if !inQuote {
			buf.WriteByte('\'')
		}
		buf.WriteString(word)
		buf.WriteByte('\'')
	}
}
<|endoftext|>


<commit_before>package create

import (
	"fmt"

	"github.com/giantswarm/awstpr"
	microerror "github.com/giantswarm/microkit/error"

	awsutil "github.com/giantswarm/aws-operator/client/aws"
	"github.com/giantswarm/aws-operator/resources"
	awsresources "github.com/giantswarm/aws-operator/resources/aws"
)

type securityGroupInput struct {
	Clients   awsutil.Clients
	GroupName string
	VPCID     string
}

type rulesInput struct {
	Cluster                awstpr.CustomObject
	Rules                  []awsresources.SecurityGroupRule
	MastersSecurityGroupID string
	WorkersSecurityGroupID string
	IngressSecurityGroupID string
}

const (
	calicoBGPNetworkPort = 179
	httpPort             = 80
	httpsPort            = 443
	sshPort              = 22

	defaultCIDR = "0.0.0.0/0"
)

func (s *Service) createSecurityGroup(input securityGroupInput) (*awsresources.SecurityGroup, error) {
	securityGroup := &awsresources.SecurityGroup{
		Description: input.GroupName,
		GroupName:   input.GroupName,
		VpcID:       input.VPCID,
		AWSEntity:   awsresources.AWSEntity{Clients: input.Clients},
	}
	securityGroupCreated, err := securityGroup.CreateIfNotExists()
	if err != nil {
		return nil, microerror.MaskAny(err)
	}
	if securityGroupCreated {
		s.logger.Log("info", fmt.Sprintf("created security group '%s'", input.GroupName))
	} else {
		s.logger.Log("info", fmt.Sprintf("security group '%s' already exists, reusing", input.GroupName))
	}

	return securityGroup, nil
}

func (s *Service) deleteSecurityGroup(input securityGroupInput) error {
	var securityGroup resources.ResourceWithID
	securityGroup = &awsresources.SecurityGroup{
		Description: input.GroupName,
		GroupName:   input.GroupName,
		AWSEntity:   awsresources.AWSEntity{Clients: input.Clients},
	}
	if err := securityGroup.Delete(); err != nil {
		return microerror.MaskAny(err)
	} else {
		s.logger.Log("info", fmt.Sprintf("deleted security group '%s'", input.GroupName))
	}

	return nil
}

// masterRules returns the rules for the masters security group.
func (ri rulesInput) masterRules() []awsresources.SecurityGroupRule {
	return []awsresources.SecurityGroupRule{
		{
			Port:       ri.Cluster.Spec.Cluster.Kubernetes.API.SecurePort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       ri.Cluster.Spec.Cluster.Etcd.Port,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       sshPort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       calicoBGPNetworkPort,
			SourceCIDR: defaultCIDR,
		},
	}
}

// workerRules returns the rules for the workers security group.
func (ri rulesInput) workerRules() []awsresources.SecurityGroupRule {
	return []awsresources.SecurityGroupRule{
		{
			Port:       ri.Cluster.Spec.Cluster.Kubernetes.IngressController.InsecurePort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       ri.Cluster.Spec.Cluster.Kubernetes.Kubelet.Port,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       sshPort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       calicoBGPNetworkPort,
			SourceCIDR: defaultCIDR,
		},
	}
}

// ingressRules returns the rules for the ingress ELB security group.
func (ri rulesInput) ingressRules() []awsresources.SecurityGroupRule {
	return []awsresources.SecurityGroupRule{
		{
			Port:       httpPort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       httpsPort,
			SourceCIDR: defaultCIDR,
		},
	}
}

func securityGroupName(clusterName string, groupName string) string {
	return fmt.Sprintf("%s-%s", clusterName, groupName)
}
<commit_msg>service: add SecurePort to worker security group<commit_after>package create

import (
	"fmt"

	"github.com/giantswarm/awstpr"
	microerror "github.com/giantswarm/microkit/error"

	awsutil "github.com/giantswarm/aws-operator/client/aws"
	"github.com/giantswarm/aws-operator/resources"
	awsresources "github.com/giantswarm/aws-operator/resources/aws"
)

type securityGroupInput struct {
	Clients   awsutil.Clients
	GroupName string
	VPCID     string
}

type rulesInput struct {
	Cluster                awstpr.CustomObject
	Rules                  []awsresources.SecurityGroupRule
	MastersSecurityGroupID string
	WorkersSecurityGroupID string
	IngressSecurityGroupID string
}

const (
	calicoBGPNetworkPort = 179
	httpPort             = 80
	httpsPort            = 443
	sshPort              = 22

	defaultCIDR = "0.0.0.0/0"
)

func (s *Service) createSecurityGroup(input securityGroupInput) (*awsresources.SecurityGroup, error) {
	securityGroup := &awsresources.SecurityGroup{
		Description: input.GroupName,
		GroupName:   input.GroupName,
		VpcID:       input.VPCID,
		AWSEntity:   awsresources.AWSEntity{Clients: input.Clients},
	}
	securityGroupCreated, err := securityGroup.CreateIfNotExists()
	if err != nil {
		return nil, microerror.MaskAny(err)
	}
	if securityGroupCreated {
		s.logger.Log("info", fmt.Sprintf("created security group '%s'", input.GroupName))
	} else {
		s.logger.Log("info", fmt.Sprintf("security group '%s' already exists, reusing", input.GroupName))
	}

	return securityGroup, nil
}

func (s *Service) deleteSecurityGroup(input securityGroupInput) error {
	var securityGroup resources.ResourceWithID
	securityGroup = &awsresources.SecurityGroup{
		Description: input.GroupName,
		GroupName:   input.GroupName,
		AWSEntity:   awsresources.AWSEntity{Clients: input.Clients},
	}
	if err := securityGroup.Delete(); err != nil {
		return microerror.MaskAny(err)
	} else {
		s.logger.Log("info", fmt.Sprintf("deleted security group '%s'", input.GroupName))
	}

	return nil
}

// masterRules returns the rules for the masters security group.
func (ri rulesInput) masterRules() []awsresources.SecurityGroupRule {
	return []awsresources.SecurityGroupRule{
		{
			Port:       ri.Cluster.Spec.Cluster.Kubernetes.API.SecurePort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       ri.Cluster.Spec.Cluster.Etcd.Port,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       sshPort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       calicoBGPNetworkPort,
			SourceCIDR: defaultCIDR,
		},
	}
}

// workerRules returns the rules for the workers security group.
func (ri rulesInput) workerRules() []awsresources.SecurityGroupRule {
	return []awsresources.SecurityGroupRule{
		{
			Port:       ri.Cluster.Spec.Cluster.Kubernetes.IngressController.SecurePort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       ri.Cluster.Spec.Cluster.Kubernetes.IngressController.InsecurePort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       ri.Cluster.Spec.Cluster.Kubernetes.Kubelet.Port,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       sshPort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       calicoBGPNetworkPort,
			SourceCIDR: defaultCIDR,
		},
	}
}

// ingressRules returns the rules for the ingress ELB security group.
func (ri rulesInput) ingressRules() []awsresources.SecurityGroupRule {
	return []awsresources.SecurityGroupRule{
		{
			Port:       httpPort,
			SourceCIDR: defaultCIDR,
		},
		{
			Port:       httpsPort,
			SourceCIDR: defaultCIDR,
		},
	}
}

func securityGroupName(clusterName string, groupName string) string {
	return fmt.Sprintf("%s-%s", clusterName, groupName)
}
<|endoftext|>


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

package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"hash/fnv"
	"net"
	"regexp"
	"strconv"
	"strings"
	"unicode/utf8"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/model"
)

const (
	defaultHelp = "Metric autogenerated by statsd_exporter."
	regErrF     = "A change of configuration created inconsistent metrics for " +
		"%q. You have to restart the statsd_exporter, and you should " +
		"consider the effects on your monitoring setup. Error: %s"
)

var (
	illegalCharsRE = regexp.MustCompile(`[^a-zA-Z0-9_]`)

	hash   = fnv.New64a()
	strBuf bytes.Buffer // Used for hashing.
	intBuf = make([]byte, 8)
)

// hashNameAndLabels returns a hash value of the provided name string and all
// the label names and values in the provided labels map.
//
// Not safe for concurrent use! (Uses a shared buffer and hasher to save on
// allocations.)
func hashNameAndLabels(name string, labels prometheus.Labels) uint64 {
	hash.Reset()
	strBuf.Reset()
	strBuf.WriteString(name)
	hash.Write(strBuf.Bytes())
	binary.BigEndian.PutUint64(intBuf, model.LabelsToSignature(labels))
	hash.Write(intBuf)
	return hash.Sum64()
}

type CounterContainer struct {
	Elements map[uint64]prometheus.Counter
}

func NewCounterContainer() *CounterContainer {
	return &CounterContainer{
		Elements: make(map[uint64]prometheus.Counter),
	}
}

func (c *CounterContainer) Get(metricName string, labels prometheus.Labels) prometheus.Counter {
	hash := hashNameAndLabels(metricName, labels)
	counter, ok := c.Elements[hash]
	if !ok {
		counter = prometheus.NewCounter(prometheus.CounterOpts{
			Name:        metricName,
			Help:        defaultHelp,
			ConstLabels: labels,
		})
		c.Elements[hash] = counter
		if err := prometheus.Register(counter); err != nil {
			log.Fatalf(regErrF, metricName, err)
		}
	}
	return counter
}

type GaugeContainer struct {
	Elements map[uint64]prometheus.Gauge
}

func NewGaugeContainer() *GaugeContainer {
	return &GaugeContainer{
		Elements: make(map[uint64]prometheus.Gauge),
	}
}

func (c *GaugeContainer) Get(metricName string, labels prometheus.Labels) prometheus.Gauge {
	hash := hashNameAndLabels(metricName, labels)
	gauge, ok := c.Elements[hash]
	if !ok {
		gauge = prometheus.NewGauge(prometheus.GaugeOpts{
			Name:        metricName,
			Help:        defaultHelp,
			ConstLabels: labels,
		})
		c.Elements[hash] = gauge
		if err := prometheus.Register(gauge); err != nil {
			log.Fatalf(regErrF, metricName, err)
		}
	}
	return gauge
}

type SummaryContainer struct {
	Elements map[uint64]prometheus.Summary
}

func NewSummaryContainer() *SummaryContainer {
	return &SummaryContainer{
		Elements: make(map[uint64]prometheus.Summary),
	}
}

func (c *SummaryContainer) Get(metricName string, labels prometheus.Labels) prometheus.Summary {
	hash := hashNameAndLabels(metricName, labels)
	summary, ok := c.Elements[hash]
	if !ok {
		summary = prometheus.NewSummary(
			prometheus.SummaryOpts{
				Name:        metricName,
				Help:        defaultHelp,
				ConstLabels: labels,
			})
		c.Elements[hash] = summary
		if err := prometheus.Register(summary); err != nil {
			log.Fatalf(regErrF, metricName, err)
		}
	}
	return summary
}

type Event interface {
	MetricName() string
	Value() float64
	Labels() map[string]string
}

type CounterEvent struct {
	metricName string
	value      float64
	labels     map[string]string
}

func (c *CounterEvent) MetricName() string        { return c.metricName }
func (c *CounterEvent) Value() float64            { return c.value }
func (c *CounterEvent) Labels() map[string]string { return c.labels }

type GaugeEvent struct {
	metricName string
	value      float64
	relative   bool
	labels     map[string]string
}

func (g *GaugeEvent) MetricName() string        { return g.metricName }
func (g *GaugeEvent) Value() float64            { return g.value }
func (c *GaugeEvent) Labels() map[string]string { return c.labels }

type TimerEvent struct {
	metricName string
	value      float64
	labels     map[string]string
}

func (t *TimerEvent) MetricName() string        { return t.metricName }
func (t *TimerEvent) Value() float64            { return t.value }
func (c *TimerEvent) Labels() map[string]string { return c.labels }

type Events []Event

type Exporter struct {
	Counters  *CounterContainer
	Gauges    *GaugeContainer
	Summaries *SummaryContainer
	mapper    *metricMapper
	addSuffix bool
}

func escapeMetricName(metricName string) string {
	// If a metric starts with a digit, prepend an underscore.
	if metricName[0] >= '0' && metricName[0] <= '9' {
		metricName = "_" + metricName
	}

	// Replace all illegal metric chars with underscores.
	metricName = illegalCharsRE.ReplaceAllString(metricName, "_")
	return metricName
}

func (b *Exporter) suffix(metricName, suffix string) string {
	str := metricName
	if b.addSuffix {
		str += "_" + suffix
	}
	return str
}

func (b *Exporter) Listen(e <-chan Events) {
	for {
		events, ok := <-e
		if !ok {
			log.Debug("Channel is closed. Break out of Exporter.Listener.")
			return
		}
		for _, event := range events {
			metricName := ""
			prometheusLabels := event.Labels()

			labels, present := b.mapper.getMapping(event.MetricName())
			if present {
				metricName = labels["name"]
				for label, value := range labels {
					if label != "name" {
						prometheusLabels[label] = value
					}
				}
			} else {
				eventsUnmapped.Inc()
				metricName = escapeMetricName(event.MetricName())
			}

			switch ev := event.(type) {
			case *CounterEvent:
				// We don't accept negative values for counters. Incrementing the counter with a negative number
				// will cause the exporter to panic. Instead we will warn and continue to the next event.
				if event.Value() < 0.0 {
					log.Errorf("Counter %q is: '%f' (counter must be non-negative value)", metricName, event.Value())
					continue
				}

				counter := b.Counters.Get(
					b.suffix(metricName, "counter"),
					prometheusLabels,
				)

				counter.Add(event.Value())

				eventStats.WithLabelValues("counter").Inc()

			case *GaugeEvent:
				gauge := b.Gauges.Get(
					b.suffix(metricName, "gauge"),
					prometheusLabels,
				)

				if ev.relative {
					gauge.Add(event.Value())
				} else {
					gauge.Set(event.Value())
				}

				eventStats.WithLabelValues("gauge").Inc()

			case *TimerEvent:
				summary := b.Summaries.Get(
					b.suffix(metricName, "timer"),
					prometheusLabels,
				)
				summary.Observe(event.Value())

				eventStats.WithLabelValues("timer").Inc()

			default:
				log.Errorln("Unsupported event type")
				eventStats.WithLabelValues("illegal").Inc()
			}
		}
	}
}

func NewExporter(mapper *metricMapper, addSuffix bool) *Exporter {
	return &Exporter{
		addSuffix: addSuffix,
		Counters:  NewCounterContainer(),
		Gauges:    NewGaugeContainer(),
		Summaries: NewSummaryContainer(),
		mapper:    mapper,
	}
}

type StatsDListener struct {
	conn *net.UDPConn
}

func buildEvent(statType, metric string, value float64, relative bool, labels map[string]string) (Event, error) {
	switch statType {
	case "c":
		return &CounterEvent{
			metricName: metric,
			value:      float64(value),
			labels:     labels,
		}, nil
	case "g":
		return &GaugeEvent{
			metricName: metric,
			value:      float64(value),
			relative:   relative,
			labels:     labels,
		}, nil
	case "ms", "h":
		return &TimerEvent{
			metricName: metric,
			value:      float64(value),
			labels:     labels,
		}, nil
	case "s":
		return nil, fmt.Errorf("No support for StatsD sets")
	default:
		return nil, fmt.Errorf("Bad stat type %s", statType)
	}
}

func (l *StatsDListener) Listen(e chan<- Events) {
	buf := make([]byte, 65535)
	for {
		n, _, err := l.conn.ReadFromUDP(buf)
		if err != nil {
			log.Fatal(err)
		}
		l.handlePacket(buf[0:n], e)
	}
}

func parseDogStatsDTagsToLabels(component string) map[string]string {
	labels := map[string]string{}
	networkStats.WithLabelValues("dogstatsd_tags").Inc()
	tags := strings.Split(component, ",")
	for _, t := range tags {
		t = strings.TrimPrefix(t, "#")
		kv := strings.SplitN(t, ":", 2)

		if len(kv) < 2 || len(kv[1]) == 0 {
			networkStats.WithLabelValues("malformed_dogstatsd_tag").Inc()
			log.Errorf("Malformed or empty DogStatsD tag %s in component %s", t, component)
			continue
		}

		labels[escapeMetricName(kv[0])] = kv[1]
	}
	return labels
}

func (l *StatsDListener) handlePacket(packet []byte, e chan<- Events) {
	lines := strings.Split(string(packet), "\n")
	events := Events{}
	for _, line := range lines {
		if line == "" {
			continue
		}

		elements := strings.SplitN(line, ":", 2)
		if len(elements) < 2 || len(elements[0]) == 0 || !utf8.ValidString(line) {
			networkStats.WithLabelValues("malformed_line").Inc()
			log.Errorln("Bad line from StatsD:", line)
			continue
		}
		metric := elements[0]
		var samples []string
		if strings.Contains(elements[1], "|#") {
			// using datadog extensions, disable multi-metrics
			samples = elements[1:]
		} else {
			samples = strings.Split(elements[1], ":")
		}
	samples:
		for _, sample := range samples {
			components := strings.Split(sample, "|")
			samplingFactor := 1.0
			if len(components) < 2 || len(components) > 4 {
				networkStats.WithLabelValues("malformed_component").Inc()
				log.Errorln("Bad component on line:", line)
				continue
			}
			valueStr, statType := components[0], components[1]

			var relative = false
			if strings.Index(valueStr, "+") == 0 || strings.Index(valueStr, "-") == 0 {
				relative = true
			}

			value, err := strconv.ParseFloat(valueStr, 64)
			if err != nil {
				log.Errorf("Bad value %s on line: %s", valueStr, line)
				networkStats.WithLabelValues("malformed_value").Inc()
				continue
			}

			multiplyEvents := 1
			labels := map[string]string{}
			if len(components) >= 3 {
				for _, component := range components[2:] {
					if len(component) == 0 {
						log.Errorln("Empty component on line: ", line)
						networkStats.WithLabelValues("malformed_component").Inc()
						continue samples
					}
				}

				for _, component := range components[2:] {
					switch component[0] {
					case '@':
						if statType != "c" && statType != "ms" {
							log.Errorln("Illegal sampling factor for non-counter metric on line", line)
							networkStats.WithLabelValues("illegal_sample_factor").Inc()
							continue
						}
						samplingFactor, err = strconv.ParseFloat(component[1:], 64)
						if err != nil {
							log.Errorf("Invalid sampling factor %s on line %s", component[1:], line)
							networkStats.WithLabelValues("invalid_sample_factor").Inc()
						}
						if samplingFactor == 0 {
							samplingFactor = 1
						}

						if statType == "c" {
							value /= samplingFactor
						} else if statType == "ms" {
							multiplyEvents = int(1 / samplingFactor)
						}
					case '#':
						labels = parseDogStatsDTagsToLabels(component)
					default:
						log.Errorf("Invalid sampling factor or tag section %s on line %s", components[2], line)
						networkStats.WithLabelValues("invalid_sample_factor").Inc()
						continue
					}
				}
			}

			for i := 0; i < multiplyEvents; i++ {
				event, err := buildEvent(statType, metric, value, relative, labels)
				if err != nil {
					log.Errorf("Error building event on line %s: %s", line, err)
					networkStats.WithLabelValues("illegal_event").Inc()
					continue
				}
				events = append(events, event)
			}
			networkStats.WithLabelValues("legal").Inc()
		}
	}
	e <- events
}
<commit_msg>Don't crash on conflicting metric names<commit_after>// Copyright 2013 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"hash/fnv"
	"net"
	"regexp"
	"strconv"
	"strings"
	"unicode/utf8"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/model"
)

const (
	defaultHelp = "Metric autogenerated by statsd_exporter."
	regErrF     = "A change of configuration created inconsistent metrics for " +
		"%q. You have to restart the statsd_exporter, and you should " +
		"consider the effects on your monitoring setup. Error: %s"
)

var (
	illegalCharsRE = regexp.MustCompile(`[^a-zA-Z0-9_]`)

	hash   = fnv.New64a()
	strBuf bytes.Buffer // Used for hashing.
	intBuf = make([]byte, 8)
)

// hashNameAndLabels returns a hash value of the provided name string and all
// the label names and values in the provided labels map.
//
// Not safe for concurrent use! (Uses a shared buffer and hasher to save on
// allocations.)
func hashNameAndLabels(name string, labels prometheus.Labels) uint64 {
	hash.Reset()
	strBuf.Reset()
	strBuf.WriteString(name)
	hash.Write(strBuf.Bytes())
	binary.BigEndian.PutUint64(intBuf, model.LabelsToSignature(labels))
	hash.Write(intBuf)
	return hash.Sum64()
}

type CounterContainer struct {
	Elements map[uint64]prometheus.Counter
}

func NewCounterContainer() *CounterContainer {
	return &CounterContainer{
		Elements: make(map[uint64]prometheus.Counter),
	}
}

func (c *CounterContainer) Get(metricName string, labels prometheus.Labels) (prometheus.Counter, error) {
	hash := hashNameAndLabels(metricName, labels)
	counter, ok := c.Elements[hash]
	if !ok {
		counter = prometheus.NewCounter(prometheus.CounterOpts{
			Name:        metricName,
			Help:        defaultHelp,
			ConstLabels: labels,
		})
		c.Elements[hash] = counter
		if err := prometheus.Register(counter); err != nil {
			return nil, err
		}
	}
	return counter, nil
}

type GaugeContainer struct {
	Elements map[uint64]prometheus.Gauge
}

func NewGaugeContainer() *GaugeContainer {
	return &GaugeContainer{
		Elements: make(map[uint64]prometheus.Gauge),
	}
}

func (c *GaugeContainer) Get(metricName string, labels prometheus.Labels) (prometheus.Gauge, error) {
	hash := hashNameAndLabels(metricName, labels)
	gauge, ok := c.Elements[hash]
	if !ok {
		gauge = prometheus.NewGauge(prometheus.GaugeOpts{
			Name:        metricName,
			Help:        defaultHelp,
			ConstLabels: labels,
		})
		c.Elements[hash] = gauge
		if err := prometheus.Register(gauge); err != nil {
			return nil, err
		}
	}
	return gauge, nil
}

type SummaryContainer struct {
	Elements map[uint64]prometheus.Summary
}

func NewSummaryContainer() *SummaryContainer {
	return &SummaryContainer{
		Elements: make(map[uint64]prometheus.Summary),
	}
}

func (c *SummaryContainer) Get(metricName string, labels prometheus.Labels) (prometheus.Summary, error) {
	hash := hashNameAndLabels(metricName, labels)
	summary, ok := c.Elements[hash]
	if !ok {
		summary = prometheus.NewSummary(
			prometheus.SummaryOpts{
				Name:        metricName,
				Help:        defaultHelp,
				ConstLabels: labels,
			})
		c.Elements[hash] = summary
		if err := prometheus.Register(summary); err != nil {
			return nil, err
		}
	}
	return summary, nil
}

type Event interface {
	MetricName() string
	Value() float64
	Labels() map[string]string
}

type CounterEvent struct {
	metricName string
	value      float64
	labels     map[string]string
}

func (c *CounterEvent) MetricName() string        { return c.metricName }
func (c *CounterEvent) Value() float64            { return c.value }
func (c *CounterEvent) Labels() map[string]string { return c.labels }

type GaugeEvent struct {
	metricName string
	value      float64
	relative   bool
	labels     map[string]string
}

func (g *GaugeEvent) MetricName() string        { return g.metricName }
func (g *GaugeEvent) Value() float64            { return g.value }
func (c *GaugeEvent) Labels() map[string]string { return c.labels }

type TimerEvent struct {
	metricName string
	value      float64
	labels     map[string]string
}

func (t *TimerEvent) MetricName() string        { return t.metricName }
func (t *TimerEvent) Value() float64            { return t.value }
func (c *TimerEvent) Labels() map[string]string { return c.labels }

type Events []Event

type Exporter struct {
	Counters  *CounterContainer
	Gauges    *GaugeContainer
	Summaries *SummaryContainer
	mapper    *metricMapper
	addSuffix bool
}

func escapeMetricName(metricName string) string {
	// If a metric starts with a digit, prepend an underscore.
	if metricName[0] >= '0' && metricName[0] <= '9' {
		metricName = "_" + metricName
	}

	// Replace all illegal metric chars with underscores.
	metricName = illegalCharsRE.ReplaceAllString(metricName, "_")
	return metricName
}

func (b *Exporter) suffix(metricName, suffix string) string {
	str := metricName
	if b.addSuffix {
		str += "_" + suffix
	}
	return str
}

func (b *Exporter) Listen(e <-chan Events) {
	for {
		events, ok := <-e
		if !ok {
			log.Debug("Channel is closed. Break out of Exporter.Listener.")
			return
		}
		for _, event := range events {
			metricName := ""
			prometheusLabels := event.Labels()

			labels, present := b.mapper.getMapping(event.MetricName())
			if present {
				metricName = labels["name"]
				for label, value := range labels {
					if label != "name" {
						prometheusLabels[label] = value
					}
				}
			} else {
				eventsUnmapped.Inc()
				metricName = escapeMetricName(event.MetricName())
			}

			switch ev := event.(type) {
			case *CounterEvent:
				// We don't accept negative values for counters. Incrementing the counter with a negative number
				// will cause the exporter to panic. Instead we will warn and continue to the next event.
				if event.Value() < 0.0 {
					log.Errorf("Counter %q is: '%f' (counter must be non-negative value)", metricName, event.Value())
					continue
				}

				counter, err := b.Counters.Get(
					b.suffix(metricName, "counter"),
					prometheusLabels,
				)
				if err == nil {
					counter.Add(event.Value())

					eventStats.WithLabelValues("counter").Inc()
				} else {
					log.Errorf(regErrF, metricName, err)
				}

			case *GaugeEvent:
				gauge, err := b.Gauges.Get(
					b.suffix(metricName, "gauge"),
					prometheusLabels,
				)

				if err == nil {
					if ev.relative {
						gauge.Add(event.Value())
					} else {
						gauge.Set(event.Value())
					}

					eventStats.WithLabelValues("gauge").Inc()
				} else {
					log.Errorf(regErrF, metricName, err)
				}

			case *TimerEvent:
				summary, err := b.Summaries.Get(
					b.suffix(metricName, "timer"),
					prometheusLabels,
				)
				if err == nil {
					summary.Observe(event.Value())

					eventStats.WithLabelValues("timer").Inc()
				} else {
					log.Errorf(regErrF, metricName, err)
				}

			default:
				log.Errorln("Unsupported event type")
				eventStats.WithLabelValues("illegal").Inc()
			}
		}
	}
}

func NewExporter(mapper *metricMapper, addSuffix bool) *Exporter {
	return &Exporter{
		addSuffix: addSuffix,
		Counters:  NewCounterContainer(),
		Gauges:    NewGaugeContainer(),
		Summaries: NewSummaryContainer(),
		mapper:    mapper,
	}
}

type StatsDListener struct {
	conn *net.UDPConn
}

func buildEvent(statType, metric string, value float64, relative bool, labels map[string]string) (Event, error) {
	switch statType {
	case "c":
		return &CounterEvent{
			metricName: metric,
			value:      float64(value),
			labels:     labels,
		}, nil
	case "g":
		return &GaugeEvent{
			metricName: metric,
			value:      float64(value),
			relative:   relative,
			labels:     labels,
		}, nil
	case "ms", "h":
		return &TimerEvent{
			metricName: metric,
			value:      float64(value),
			labels:     labels,
		}, nil
	case "s":
		return nil, fmt.Errorf("No support for StatsD sets")
	default:
		return nil, fmt.Errorf("Bad stat type %s", statType)
	}
}

func (l *StatsDListener) Listen(e chan<- Events) {
	buf := make([]byte, 65535)
	for {
		n, _, err := l.conn.ReadFromUDP(buf)
		if err != nil {
			log.Fatal(err)
		}
		l.handlePacket(buf[0:n], e)
	}
}

func parseDogStatsDTagsToLabels(component string) map[string]string {
	labels := map[string]string{}
	networkStats.WithLabelValues("dogstatsd_tags").Inc()
	tags := strings.Split(component, ",")
	for _, t := range tags {
		t = strings.TrimPrefix(t, "#")
		kv := strings.SplitN(t, ":", 2)

		if len(kv) < 2 || len(kv[1]) == 0 {
			networkStats.WithLabelValues("malformed_dogstatsd_tag").Inc()
			log.Errorf("Malformed or empty DogStatsD tag %s in component %s", t, component)
			continue
		}

		labels[escapeMetricName(kv[0])] = kv[1]
	}
	return labels
}

func (l *StatsDListener) handlePacket(packet []byte, e chan<- Events) {
	lines := strings.Split(string(packet), "\n")
	events := Events{}
	for _, line := range lines {
		if line == "" {
			continue
		}

		elements := strings.SplitN(line, ":", 2)
		if len(elements) < 2 || len(elements[0]) == 0 || !utf8.ValidString(line) {
			networkStats.WithLabelValues("malformed_line").Inc()
			log.Errorln("Bad line from StatsD:", line)
			continue
		}
		metric := elements[0]
		var samples []string
		if strings.Contains(elements[1], "|#") {
			// using datadog extensions, disable multi-metrics
			samples = elements[1:]
		} else {
			samples = strings.Split(elements[1], ":")
		}
	samples:
		for _, sample := range samples {
			components := strings.Split(sample, "|")
			samplingFactor := 1.0
			if len(components) < 2 || len(components) > 4 {
				networkStats.WithLabelValues("malformed_component").Inc()
				log.Errorln("Bad component on line:", line)
				continue
			}
			valueStr, statType := components[0], components[1]

			var relative = false
			if strings.Index(valueStr, "+") == 0 || strings.Index(valueStr, "-") == 0 {
				relative = true
			}

			value, err := strconv.ParseFloat(valueStr, 64)
			if err != nil {
				log.Errorf("Bad value %s on line: %s", valueStr, line)
				networkStats.WithLabelValues("malformed_value").Inc()
				continue
			}

			multiplyEvents := 1
			labels := map[string]string{}
			if len(components) >= 3 {
				for _, component := range components[2:] {
					if len(component) == 0 {
						log.Errorln("Empty component on line: ", line)
						networkStats.WithLabelValues("malformed_component").Inc()
						continue samples
					}
				}

				for _, component := range components[2:] {
					switch component[0] {
					case '@':
						if statType != "c" && statType != "ms" {
							log.Errorln("Illegal sampling factor for non-counter metric on line", line)
							networkStats.WithLabelValues("illegal_sample_factor").Inc()
							continue
						}
						samplingFactor, err = strconv.ParseFloat(component[1:], 64)
						if err != nil {
							log.Errorf("Invalid sampling factor %s on line %s", component[1:], line)
							networkStats.WithLabelValues("invalid_sample_factor").Inc()
						}
						if samplingFactor == 0 {
							samplingFactor = 1
						}

						if statType == "c" {
							value /= samplingFactor
						} else if statType == "ms" {
							multiplyEvents = int(1 / samplingFactor)
						}
					case '#':
						labels = parseDogStatsDTagsToLabels(component)
					default:
						log.Errorf("Invalid sampling factor or tag section %s on line %s", components[2], line)
						networkStats.WithLabelValues("invalid_sample_factor").Inc()
						continue
					}
				}
			}

			for i := 0; i < multiplyEvents; i++ {
				event, err := buildEvent(statType, metric, value, relative, labels)
				if err != nil {
					log.Errorf("Error building event on line %s: %s", line, err)
					networkStats.WithLabelValues("illegal_event").Inc()
					continue
				}
				events = append(events, event)
			}
			networkStats.WithLabelValues("legal").Inc()
		}
	}
	e <- events
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"testing"
	"time"
)

type postgresSuite struct {
	postgres   *sql.DB
	user       UserRepository
	permission PermissionRepository
	group      GroupRepository
}

func setupPostgresSuite(t *testing.T) *postgresSuite {
	ps := postgresSuite{}
	ps.setup(t)

	return ps
}

func (ps *postgresSuite) setup(t *testing.T) {
	config.parse()
	var err error

	ps.postgres, err = sql.Open("postgres", config.postgres.connectionString)
	if err != nil {
		t.Errorf("connection to postgres failed with error: %s", err.Error())
		t.FailNow()
	}

	setupDatabase(ps.postgres)

	ps.user = newUserRepository(ps.postgres)
}

func (ps *postgresSuite) teardown(t *testing.T) {
	if err := tearDownDatabase(ps.postgres); err != nil {
		t.Errorf("unexpected error during database teardown: %s", err.Error())
	}

	ps.postgres.Close()
}

func assertf(t *testing.T, is bool, msg string, args ...interface{}) bool {
	if !is {
		t.Errorf(msg, args...)
	}

	return is
}

func assert(t *testing.T, is bool, msg string) bool {
	if !is {
		t.Errorf(msg)
	}

	return is
}

func assertfTime(t *testing.T, tm *time.Time, msg string, args ...interface{}) bool {
	if tm == nil || tm.IsZero() {
		t.Errorf(msg, args...)
		return false
	}

	return true
}
<commit_msg>postgres suite fix<commit_after>package main

import (
	"database/sql"
	"testing"
	"time"
)

type postgresSuite struct {
	postgres   *sql.DB
	user       UserRepository
	permission PermissionRepository
	group      GroupRepository
}

func setupPostgresSuite(t *testing.T) *postgresSuite {
	ps := postgresSuite{}
	ps.setup(t)

	return &ps
}

func (ps *postgresSuite) setup(t *testing.T) {
	config.parse()
	var err error

	ps.postgres, err = sql.Open("postgres", config.postgres.connectionString)
	if err != nil {
		t.Errorf("connection to postgres failed with error: %s", err.Error())
		t.FailNow()
	}

	setupDatabase(ps.postgres)

	ps.user = newUserRepository(ps.postgres)
}

func (ps *postgresSuite) teardown(t *testing.T) {
	if err := tearDownDatabase(ps.postgres); err != nil {
		t.Errorf("unexpected error during database teardown: %s", err.Error())
	}

	ps.postgres.Close()
}

func assertf(t *testing.T, is bool, msg string, args ...interface{}) bool {
	if !is {
		t.Errorf(msg, args...)
	}

	return is
}

func assert(t *testing.T, is bool, msg string) bool {
	if !is {
		t.Errorf(msg)
	}

	return is
}

func assertfTime(t *testing.T, tm *time.Time, msg string, args ...interface{}) bool {
	if tm == nil || tm.IsZero() {
		t.Errorf(msg, args...)
		return false
	}

	return true
}
<|endoftext|>