<commit_before>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package driver defines interfaces to be implemented by database
// drivers as used by package sql.
//
// Most code should use package sql.
package driver

import "errors"

// A driver Value is a value that drivers must be able to handle.
// A Value is either nil or an instance of one of these types:
//
//   int64
//   float64
//   bool
//   []byte
//   string   [*] everywhere except from Rows.Next.
//   time.Time
type Value interface{}

// Driver is the interface that must be implemented by a database
// driver.
type Driver interface {
	// Open returns a new connection to the database.
	// The name is a string in a driver-specific format.
	//
	// Open may return a cached connection (one previously
	// closed), but doing so is unnecessary; the sql package
	// maintains a pool of idle connections for efficient re-use.
	//
	// The returned connection is only used by one goroutine at a
	// time.
	Open(name string) (Conn, error)
}

// ErrSkip may be returned by some optional interfaces' methods to
// indicate at runtime that the fast path is unavailable and the sql
// package should continue as if the optional interface was not
// implemented. ErrSkip is only supported where explicitly
// documented.
var ErrSkip = errors.New("driver: skip fast-path; continue as if unimplemented")

// ErrBadConn should be returned by a driver to signal to the sql
// package that a driver.Conn is in a bad state (such as the server
// having earlier closed the connection) and the sql package should
// retry on a new connection.
//
// To prevent duplicate operations, ErrBadConn should NOT be returned
// if there's a possibility that the database server might have
// performed the operation. Even if the server sends back an error,
// you shouldn't return ErrBadConn.
var ErrBadConn = errors.New("driver: bad connection")

// Execer is an optional interface that may be implemented by a Conn.
//
// If a Conn does not implement Execer, the sql package's DB.Exec will
// first prepare a query, execute the statement, and then close the
// statement.
//
// Exec may return ErrSkip.
type Execer interface {
	Exec(query string, args []Value) (Result, error)
}

// Queryer is an optional interface that may be implemented by a Conn.
//
// If a Conn does not implement Queryer, the sql package's DB.Query will
// first prepare a query, execute the statement, and then close the
// statement.
//
// Query may return ErrSkip.
type Queryer interface {
	Query(query string, args []Value) (Rows, error)
}

// Conn is a connection to a database. It is not used concurrently
// by multiple goroutines.
//
// Conn is assumed to be stateful.
type Conn interface {
	// Prepare returns a prepared statement, bound to this connection.
	Prepare(query string) (Stmt, error)

	// Close invalidates and potentially stops any current
	// prepared statements and transactions, marking this
	// connection as no longer in use.
	//
	// Because the sql package maintains a free pool of
	// connections and only calls Close when there's a surplus of
	// idle connections, it shouldn't be necessary for drivers to
	// do their own connection caching.
	Close() error

	// Begin starts and returns a new transaction.
	Begin() (Tx, error)
}

// Result is the result of a query execution.
type Result interface {
	// LastInsertId returns the database's auto-generated ID
	// after, for example, an INSERT into a table with primary
	// key.
	LastInsertId() (int64, error)

	// RowsAffected returns the number of rows affected by the
	// query.
	RowsAffected() (int64, error)
}

// Stmt is a prepared statement. It is bound to a Conn and not
// used by multiple goroutines concurrently.
type Stmt interface {
	// Close closes the statement.
	//
	// As of Go 1.1, a Stmt will not be closed if it's in use
	// by any queries.
	Close() error

	// NumInput returns the number of placeholder parameters.
	//
	// If NumInput returns >= 0, the sql package will sanity check
	// argument counts from callers and return errors to the caller
	// before the statement's Exec or Query methods are called.
	//
	// NumInput may also return -1, if the driver doesn't know
	// its number of placeholders. In that case, the sql package
	// will not sanity check Exec or Query argument counts.
	NumInput() int

	// Exec executes a query that doesn't return rows, such
	// as an INSERT or UPDATE.
	Exec(args []Value) (Result, error)

	// Exec executes a query that may return rows, such as a
	// SELECT.
	Query(args []Value) (Rows, error)
}

// ColumnConverter may be optionally implemented by Stmt if the
// the statement is aware of its own columns' types and can
// convert from any type to a driver Value.
type ColumnConverter interface {
	// ColumnConverter returns a ValueConverter for the provided
	// column index.  If the type of a specific column isn't known
	// or shouldn't be handled specially, DefaultValueConverter
	// can be returned.
	ColumnConverter(idx int) ValueConverter
}

// Rows is an iterator over an executed query's results.
type Rows interface {
	// Columns returns the names of the columns. The number of
	// columns of the result is inferred from the length of the
	// slice.  If a particular column name isn't known, an empty
	// string should be returned for that entry.
	Columns() []string

	// Close closes the rows iterator.
	Close() error

	// Next is called to populate the next row of data into
	// the provided slice. The provided slice will be the same
	// size as the Columns() are wide.
	//
	// The dest slice may be populated only with
	// a driver Value type, but excluding string.
	// All string values must be converted to []byte.
	//
	// Next should return io.EOF when there are no more rows.
	Next(dest []Value) error
}

// Tx is a transaction.
type Tx interface {
	Commit() error
	Rollback() error
}

// RowsAffected implements Result for an INSERT or UPDATE operation
// which mutates a number of rows.
type RowsAffected int64

var _ Result = RowsAffected(0)

func (RowsAffected) LastInsertId() (int64, error) {
	return 0, errors.New("no LastInsertId available")
}

func (v RowsAffected) RowsAffected() (int64, error) {
	return int64(v), nil
}

// ResultNoRows is a pre-defined Result for drivers to return when a DDL
// command (such as a CREATE TABLE) succeeds. It returns an error for both
// LastInsertId and RowsAffected.
var ResultNoRows noRows

type noRows struct{}

var _ Result = noRows{}

func (noRows) LastInsertId() (int64, error) {
	return 0, errors.New("no LastInsertId available after DDL statement")
}

func (noRows) RowsAffected() (int64, error) {
	return 0, errors.New("no RowsAffected available after DDL statement")
}
<commit_msg>database/sql/driver: try to unstutter Value docs<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package driver defines interfaces to be implemented by database
// drivers as used by package sql.
//
// Most code should use package sql.
package driver

import "errors"

// Value is a value that drivers must be able to handle.
// It is either nil or an instance of one of these types:
//
//   int64
//   float64
//   bool
//   []byte
//   string   [*] everywhere except from Rows.Next.
//   time.Time
type Value interface{}

// Driver is the interface that must be implemented by a database
// driver.
type Driver interface {
	// Open returns a new connection to the database.
	// The name is a string in a driver-specific format.
	//
	// Open may return a cached connection (one previously
	// closed), but doing so is unnecessary; the sql package
	// maintains a pool of idle connections for efficient re-use.
	//
	// The returned connection is only used by one goroutine at a
	// time.
	Open(name string) (Conn, error)
}

// ErrSkip may be returned by some optional interfaces' methods to
// indicate at runtime that the fast path is unavailable and the sql
// package should continue as if the optional interface was not
// implemented. ErrSkip is only supported where explicitly
// documented.
var ErrSkip = errors.New("driver: skip fast-path; continue as if unimplemented")

// ErrBadConn should be returned by a driver to signal to the sql
// package that a driver.Conn is in a bad state (such as the server
// having earlier closed the connection) and the sql package should
// retry on a new connection.
//
// To prevent duplicate operations, ErrBadConn should NOT be returned
// if there's a possibility that the database server might have
// performed the operation. Even if the server sends back an error,
// you shouldn't return ErrBadConn.
var ErrBadConn = errors.New("driver: bad connection")

// Execer is an optional interface that may be implemented by a Conn.
//
// If a Conn does not implement Execer, the sql package's DB.Exec will
// first prepare a query, execute the statement, and then close the
// statement.
//
// Exec may return ErrSkip.
type Execer interface {
	Exec(query string, args []Value) (Result, error)
}

// Queryer is an optional interface that may be implemented by a Conn.
//
// If a Conn does not implement Queryer, the sql package's DB.Query will
// first prepare a query, execute the statement, and then close the
// statement.
//
// Query may return ErrSkip.
type Queryer interface {
	Query(query string, args []Value) (Rows, error)
}

// Conn is a connection to a database. It is not used concurrently
// by multiple goroutines.
//
// Conn is assumed to be stateful.
type Conn interface {
	// Prepare returns a prepared statement, bound to this connection.
	Prepare(query string) (Stmt, error)

	// Close invalidates and potentially stops any current
	// prepared statements and transactions, marking this
	// connection as no longer in use.
	//
	// Because the sql package maintains a free pool of
	// connections and only calls Close when there's a surplus of
	// idle connections, it shouldn't be necessary for drivers to
	// do their own connection caching.
	Close() error

	// Begin starts and returns a new transaction.
	Begin() (Tx, error)
}

// Result is the result of a query execution.
type Result interface {
	// LastInsertId returns the database's auto-generated ID
	// after, for example, an INSERT into a table with primary
	// key.
	LastInsertId() (int64, error)

	// RowsAffected returns the number of rows affected by the
	// query.
	RowsAffected() (int64, error)
}

// Stmt is a prepared statement. It is bound to a Conn and not
// used by multiple goroutines concurrently.
type Stmt interface {
	// Close closes the statement.
	//
	// As of Go 1.1, a Stmt will not be closed if it's in use
	// by any queries.
	Close() error

	// NumInput returns the number of placeholder parameters.
	//
	// If NumInput returns >= 0, the sql package will sanity check
	// argument counts from callers and return errors to the caller
	// before the statement's Exec or Query methods are called.
	//
	// NumInput may also return -1, if the driver doesn't know
	// its number of placeholders. In that case, the sql package
	// will not sanity check Exec or Query argument counts.
	NumInput() int

	// Exec executes a query that doesn't return rows, such
	// as an INSERT or UPDATE.
	Exec(args []Value) (Result, error)

	// Exec executes a query that may return rows, such as a
	// SELECT.
	Query(args []Value) (Rows, error)
}

// ColumnConverter may be optionally implemented by Stmt if the
// the statement is aware of its own columns' types and can
// convert from any type to a driver Value.
type ColumnConverter interface {
	// ColumnConverter returns a ValueConverter for the provided
	// column index.  If the type of a specific column isn't known
	// or shouldn't be handled specially, DefaultValueConverter
	// can be returned.
	ColumnConverter(idx int) ValueConverter
}

// Rows is an iterator over an executed query's results.
type Rows interface {
	// Columns returns the names of the columns. The number of
	// columns of the result is inferred from the length of the
	// slice.  If a particular column name isn't known, an empty
	// string should be returned for that entry.
	Columns() []string

	// Close closes the rows iterator.
	Close() error

	// Next is called to populate the next row of data into
	// the provided slice. The provided slice will be the same
	// size as the Columns() are wide.
	//
	// The dest slice may be populated only with
	// a driver Value type, but excluding string.
	// All string values must be converted to []byte.
	//
	// Next should return io.EOF when there are no more rows.
	Next(dest []Value) error
}

// Tx is a transaction.
type Tx interface {
	Commit() error
	Rollback() error
}

// RowsAffected implements Result for an INSERT or UPDATE operation
// which mutates a number of rows.
type RowsAffected int64

var _ Result = RowsAffected(0)

func (RowsAffected) LastInsertId() (int64, error) {
	return 0, errors.New("no LastInsertId available")
}

func (v RowsAffected) RowsAffected() (int64, error) {
	return int64(v), nil
}

// ResultNoRows is a pre-defined Result for drivers to return when a DDL
// command (such as a CREATE TABLE) succeeds. It returns an error for both
// LastInsertId and RowsAffected.
var ResultNoRows noRows

type noRows struct{}

var _ Result = noRows{}

func (noRows) LastInsertId() (int64, error) {
	return 0, errors.New("no LastInsertId available after DDL statement")
}

func (noRows) RowsAffected() (int64, error) {
	return 0, errors.New("no RowsAffected available after DDL statement")
}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// HTTP reverse proxy handler

package httputil

import (
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"
)

// onExitFlushLoop is a callback set by tests to detect the state of the
// flushLoop() goroutine.
var onExitFlushLoop func()

var ErrHandled = errors.New("request handled")

// ReverseProxy is an HTTP Handler that takes an incoming request and
// sends it to another server, proxying the response back to the
// client.
type ReverseProxy struct {
	// Proxy is a function that either handles the request
	// by sending back data or modifies request to be
	// sent. In the former case, ErrHandled must be returned.
	Proxy func(http.ResponseWriter, *http.Request) error

	// The transport used to perform proxy requests.
	// If nil, http.DefaultTransport is used.
	Transport http.RoundTripper

	// FlushInterval specifies the flush interval
	// to flush to the client while copying the
	// response body.
	// If zero, no periodic flushing is done.
	FlushInterval time.Duration

	// ErrorLog specifies an optional logger for errors
	// that occur when attempting to proxy the request.
	// If nil, logging goes to os.Stderr via the log package's
	// standard logger.
	ErrorLog *log.Logger

	// BufferPool optionally specifies a buffer pool to
	// get byte slices for use by io.CopyBuffer when
	// copying HTTP response bodies.
	BufferPool BufferPool

	// Output additional debug logging
	Debug bool
}

// A BufferPool is an interface for getting and returning temporary
// byte slices for use by io.CopyBuffer.
type BufferPool interface {
	Get() []byte
	Put([]byte)
}

func singleJoiningSlash(a, b string) string {
	aslash := strings.HasSuffix(a, "/")
	bslash := strings.HasPrefix(b, "/")
	switch {
	case aslash && bslash:
		return a + b[1:]
	case !aslash && !bslash:
		return a + "/" + b
	}
	return a + b
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Proxy-Connection", // non-standard but still sent by libcurl and rejected by e.g. google
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te",      // canonicalized version of "TE"
	"Trailer", // not Trailers per URL above; http://www.rfc-editor.org/errata_search.php?eid=4522
	"Transfer-Encoding",
	"Upgrade",
}

type requestCanceler interface {
	CancelRequest(*http.Request)
}

type runOnFirstRead struct {
	io.Reader // optional; nil means empty body

	fn func() // Run before first Read, then set to nil
}

func (c *runOnFirstRead) Read(bs []byte) (int, error) {
	if c.fn != nil {
		c.fn()
		c.fn = nil
	}
	if c.Reader == nil {
		return 0, io.EOF
	}
	return c.Reader.Read(bs)
}

func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	transport := p.Transport
	if transport == nil {
		transport = http.DefaultTransport
	}

	outreq := new(http.Request)
	*outreq = *req // includes shallow copies of maps, but okay

	err := p.Proxy(rw, outreq)
	if err != nil {
		if err == ErrHandled {
			return
		}

		rw.WriteHeader(500)
		rw.Write([]byte(err.Error()))
		return
	}

	outreq.Proto = "HTTP/1.1"
	outreq.ProtoMajor = 1
	outreq.ProtoMinor = 1
	outreq.Close = false

	// Remove hop-by-hop headers to the backend.  Especially
	// important is "Connection" because we want a persistent
	// connection, regardless of what the client sent to us.  This
	// is modifying the same underlying map from req (shallow
	// copied above) so we only copy it if necessary.
	copiedHeaders := false
	for _, h := range hopHeaders {
		if outreq.Header.Get(h) != "" {
			if !copiedHeaders {
				outreq.Header = make(http.Header)
				copyHeader(outreq.Header, req.Header)
				copiedHeaders = true
			}
			outreq.Header.Del(h)
		}
	}

	if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
		// If we aren't the first proxy retain prior
		// X-Forwarded-For information as a comma+space
		// separated list and fold multiple headers into one.
		if prior, ok := outreq.Header["X-Forwarded-For"]; ok {
			clientIP = strings.Join(prior, ", ") + ", " + clientIP
		}
		outreq.Header.Set("X-Forwarded-For", clientIP)
	}

	if req.TLS != nil {
		outreq.Header.Set("X-Forwarded-Proto", "https")
	} else {
		outreq.Header.Set("X-Forwarded-Proto", "http")
	}

	if req.Header.Get("Connection") == "Upgrade" &&
		req.Header.Get("Upgrade") == "websocket" {
		p.websocketProxy(rw, outreq)
		return
	}

	if closeNotifier, ok := rw.(http.CloseNotifier); ok {
		if requestCanceler, ok := transport.(requestCanceler); ok {
			reqDone := make(chan struct{})
			defer close(reqDone)

			clientGone := closeNotifier.CloseNotify()

			outreq.Body = struct {
				io.Reader
				io.Closer
			}{
				Reader: &runOnFirstRead{
					Reader: outreq.Body,
					fn: func() {
						go func() {
							select {
							case <-clientGone:
								requestCanceler.CancelRequest(outreq)
							case <-reqDone:
							}
						}()
					},
				},
				Closer: outreq.Body,
			}
		}
	}

	res, err := transport.RoundTrip(outreq)
	if err != nil {
		p.logf("http: proxy error: %v", err)
		rw.WriteHeader(http.StatusInternalServerError)
		return
	}

	for _, h := range hopHeaders {
		res.Header.Del(h)
	}

	copyHeader(rw.Header(), res.Header)

	// The "Trailer" header isn't included in the Transport's response,
	// at least for *http.Transport. Build it up from Trailer.
	if len(res.Trailer) > 0 {
		var trailerKeys []string
		for k := range res.Trailer {
			trailerKeys = append(trailerKeys, k)
		}
		rw.Header().Add("Trailer", strings.Join(trailerKeys, ", "))
	}

	rw.WriteHeader(res.StatusCode)
	if len(res.Trailer) > 0 {
		// Force chunking if we saw a response trailer.
		// This prevents net/http from calculating the length for short
		// bodies and adding a Content-Length.
		if fl, ok := rw.(http.Flusher); ok {
			fl.Flush()
		}
	}
	p.copyResponse(rw, res.Body)
	res.Body.Close() // close now, instead of defer, to populate res.Trailer
	copyHeader(rw.Header(), res.Trailer)
}

func (p *ReverseProxy) websocketProxy(w http.ResponseWriter, r *http.Request) {
	if p.Debug {
		p.logf("http: websockets detected, using direct proxy")
	}

	var (
		conn net.Conn
		err  error
	)

	r.Header.Set("Connection", "Upgrade")
	r.Header.Set("Upgrade", "websocket")

	switch r.URL.Scheme {
	case "httpu":
		conn, err = net.Dial("unix", r.URL.Host)
	case "http":
		conn, err = net.Dial("tcp", r.URL.Host)
	case "https":
		conn, err = tls.Dial("tcp", r.URL.Host, nil)
	default:
		err = fmt.Errorf("Unsupport scheme: %s", r.URL.Scheme)
	}

	if err != nil {
		http.Error(w, "Error contacting backend server", 500)
		log.Printf("Error dialing websocket backend %s: %v", r.URL, err)
		return
	}

	defer conn.Close()

	hj, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "Not a hijacker?", 500)
		return
	}

	nc, buf, err := hj.Hijack()
	if err != nil {
		log.Printf("Hijack error: %v", err)
		return
	}

	buf.Flush()

	defer nc.Close()

	err = r.Write(conn)
	if err != nil {
		log.Printf("Error copying request to target: %v", err)
		return
	}

	errc := make(chan error, 2)
	cp := func(dst io.Writer, src io.Reader) {
		_, err := io.Copy(dst, src)
		errc <- err
	}

	go cp(conn, nc)
	go cp(nc, conn)

	if p.Debug {
		p.logf("http/websocket: closing connection to %s", nc.RemoteAddr())
	}

	<-errc
}

func (p *ReverseProxy) copyResponse(dst io.Writer, src io.Reader) {
	if p.FlushInterval != 0 {
		if wf, ok := dst.(writeFlusher); ok {
			mlw := &maxLatencyWriter{
				dst:     wf,
				latency: p.FlushInterval,
				done:    make(chan bool),
			}
			go mlw.flushLoop()
			defer mlw.stop()
			dst = mlw
		}
	}

	var buf []byte
	if p.BufferPool != nil {
		buf = p.BufferPool.Get()
	}
	io.CopyBuffer(dst, src, buf)
	if p.BufferPool != nil {
		p.BufferPool.Put(buf)
	}
}

func (p *ReverseProxy) logf(format string, args ...interface{}) {
	if p.ErrorLog != nil {
		p.ErrorLog.Printf(format, args...)
	} else {
		log.Printf(format, args...)
	}
}

type writeFlusher interface {
	io.Writer
	http.Flusher
}

type maxLatencyWriter struct {
	dst     writeFlusher
	latency time.Duration

	lk   sync.Mutex // protects Write + Flush
	done chan bool
}

func (m *maxLatencyWriter) Write(p []byte) (int, error) {
	m.lk.Lock()
	defer m.lk.Unlock()
	return m.dst.Write(p)
}

func (m *maxLatencyWriter) flushLoop() {
	t := time.NewTicker(m.latency)
	defer t.Stop()
	for {
		select {
		case <-m.done:
			if onExitFlushLoop != nil {
				onExitFlushLoop()
			}
			return
		case <-t.C:
			m.lk.Lock()
			m.dst.Flush()
			m.lk.Unlock()
		}
	}
}

func (m *maxLatencyWriter) stop() { m.done <- true }
<commit_msg>Better handling of websocket headers<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// HTTP reverse proxy handler

package httputil

import (
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"
)

// onExitFlushLoop is a callback set by tests to detect the state of the
// flushLoop() goroutine.
var onExitFlushLoop func()

var ErrHandled = errors.New("request handled")

// ReverseProxy is an HTTP Handler that takes an incoming request and
// sends it to another server, proxying the response back to the
// client.
type ReverseProxy struct {
	// Proxy is a function that either handles the request
	// by sending back data or modifies request to be
	// sent. In the former case, ErrHandled must be returned.
	Proxy func(http.ResponseWriter, *http.Request) error

	// The transport used to perform proxy requests.
	// If nil, http.DefaultTransport is used.
	Transport http.RoundTripper

	// FlushInterval specifies the flush interval
	// to flush to the client while copying the
	// response body.
	// If zero, no periodic flushing is done.
	FlushInterval time.Duration

	// ErrorLog specifies an optional logger for errors
	// that occur when attempting to proxy the request.
	// If nil, logging goes to os.Stderr via the log package's
	// standard logger.
	ErrorLog *log.Logger

	// BufferPool optionally specifies a buffer pool to
	// get byte slices for use by io.CopyBuffer when
	// copying HTTP response bodies.
	BufferPool BufferPool

	// Output additional debug logging
	Debug bool
}

// A BufferPool is an interface for getting and returning temporary
// byte slices for use by io.CopyBuffer.
type BufferPool interface {
	Get() []byte
	Put([]byte)
}

func singleJoiningSlash(a, b string) string {
	aslash := strings.HasSuffix(a, "/")
	bslash := strings.HasPrefix(b, "/")
	switch {
	case aslash && bslash:
		return a + b[1:]
	case !aslash && !bslash:
		return a + "/" + b
	}
	return a + b
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Proxy-Connection", // non-standard but still sent by libcurl and rejected by e.g. google
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te",      // canonicalized version of "TE"
	"Trailer", // not Trailers per URL above; http://www.rfc-editor.org/errata_search.php?eid=4522
	"Transfer-Encoding",
	"Upgrade",
}

type requestCanceler interface {
	CancelRequest(*http.Request)
}

type runOnFirstRead struct {
	io.Reader // optional; nil means empty body

	fn func() // Run before first Read, then set to nil
}

func (c *runOnFirstRead) Read(bs []byte) (int, error) {
	if c.fn != nil {
		c.fn()
		c.fn = nil
	}
	if c.Reader == nil {
		return 0, io.EOF
	}
	return c.Reader.Read(bs)
}

func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	transport := p.Transport
	if transport == nil {
		transport = http.DefaultTransport
	}

	outreq := new(http.Request)
	*outreq = *req // includes shallow copies of maps, but okay

	err := p.Proxy(rw, outreq)
	if err != nil {
		if err == ErrHandled {
			return
		}

		rw.WriteHeader(500)
		rw.Write([]byte(err.Error()))
		return
	}

	outreq.Proto = "HTTP/1.1"
	outreq.ProtoMajor = 1
	outreq.ProtoMinor = 1
	outreq.Close = false

	// Remove hop-by-hop headers to the backend.  Especially
	// important is "Connection" because we want a persistent
	// connection, regardless of what the client sent to us.  This
	// is modifying the same underlying map from req (shallow
	// copied above) so we only copy it if necessary.
	copiedHeaders := false
	for _, h := range hopHeaders {
		if outreq.Header.Get(h) != "" {
			if !copiedHeaders {
				outreq.Header = make(http.Header)
				copyHeader(outreq.Header, req.Header)
				copiedHeaders = true
			}
			outreq.Header.Del(h)
		}
	}

	if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
		// If we aren't the first proxy retain prior
		// X-Forwarded-For information as a comma+space
		// separated list and fold multiple headers into one.
		if prior, ok := outreq.Header["X-Forwarded-For"]; ok {
			clientIP = strings.Join(prior, ", ") + ", " + clientIP
		}
		outreq.Header.Set("X-Forwarded-For", clientIP)
	}

	if req.TLS != nil {
		outreq.Header.Set("X-Forwarded-Proto", "https")
	} else {
		outreq.Header.Set("X-Forwarded-Proto", "http")
	}

	if strings.Contains(req.Header.Get("Connection"), "Upgrade") &&
		req.Header.Get("Upgrade") == "websocket" {
		p.websocketProxy(rw, outreq)
		return
	}

	if closeNotifier, ok := rw.(http.CloseNotifier); ok {
		if requestCanceler, ok := transport.(requestCanceler); ok {
			reqDone := make(chan struct{})
			defer close(reqDone)

			clientGone := closeNotifier.CloseNotify()

			outreq.Body = struct {
				io.Reader
				io.Closer
			}{
				Reader: &runOnFirstRead{
					Reader: outreq.Body,
					fn: func() {
						go func() {
							select {
							case <-clientGone:
								requestCanceler.CancelRequest(outreq)
							case <-reqDone:
							}
						}()
					},
				},
				Closer: outreq.Body,
			}
		}
	}

	res, err := transport.RoundTrip(outreq)
	if err != nil {
		p.logf("http: proxy error: %v", err)
		rw.WriteHeader(http.StatusInternalServerError)
		return
	}

	for _, h := range hopHeaders {
		res.Header.Del(h)
	}

	copyHeader(rw.Header(), res.Header)

	// The "Trailer" header isn't included in the Transport's response,
	// at least for *http.Transport. Build it up from Trailer.
	if len(res.Trailer) > 0 {
		var trailerKeys []string
		for k := range res.Trailer {
			trailerKeys = append(trailerKeys, k)
		}
		rw.Header().Add("Trailer", strings.Join(trailerKeys, ", "))
	}

	rw.WriteHeader(res.StatusCode)
	if len(res.Trailer) > 0 {
		// Force chunking if we saw a response trailer.
		// This prevents net/http from calculating the length for short
		// bodies and adding a Content-Length.
		if fl, ok := rw.(http.Flusher); ok {
			fl.Flush()
		}
	}
	p.copyResponse(rw, res.Body)
	res.Body.Close() // close now, instead of defer, to populate res.Trailer
	copyHeader(rw.Header(), res.Trailer)
}

func (p *ReverseProxy) websocketProxy(w http.ResponseWriter, r *http.Request) {
	if p.Debug {
		p.logf("http: websockets detected, using direct proxy")
	}

	var (
		conn net.Conn
		err  error
	)

	r.Header.Set("Connection", "Upgrade")
	r.Header.Set("Upgrade", "websocket")

	switch r.URL.Scheme {
	case "httpu":
		conn, err = net.Dial("unix", r.URL.Host)
	case "http":
		conn, err = net.Dial("tcp", r.URL.Host)
	case "https":
		conn, err = tls.Dial("tcp", r.URL.Host, nil)
	default:
		err = fmt.Errorf("Unsupport scheme: %s", r.URL.Scheme)
	}

	if err != nil {
		http.Error(w, "Error contacting backend server", 500)
		log.Printf("Error dialing websocket backend %s: %v", r.URL, err)
		return
	}

	defer conn.Close()

	hj, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "Not a hijacker?", 500)
		return
	}

	nc, buf, err := hj.Hijack()
	if err != nil {
		log.Printf("Hijack error: %v", err)
		return
	}

	buf.Flush()

	defer nc.Close()

	err = r.Write(conn)
	if err != nil {
		log.Printf("Error copying request to target: %v", err)
		return
	}

	errc := make(chan error, 2)
	cp := func(dst io.Writer, src io.Reader) {
		_, err := io.Copy(dst, src)
		errc <- err
	}

	go cp(conn, nc)
	go cp(nc, conn)

	if p.Debug {
		p.logf("http/websocket: closing connection to %s", nc.RemoteAddr())
	}

	<-errc
}

func (p *ReverseProxy) copyResponse(dst io.Writer, src io.Reader) {
	if p.FlushInterval != 0 {
		if wf, ok := dst.(writeFlusher); ok {
			mlw := &maxLatencyWriter{
				dst:     wf,
				latency: p.FlushInterval,
				done:    make(chan bool),
			}
			go mlw.flushLoop()
			defer mlw.stop()
			dst = mlw
		}
	}

	var buf []byte
	if p.BufferPool != nil {
		buf = p.BufferPool.Get()
	}
	io.CopyBuffer(dst, src, buf)
	if p.BufferPool != nil {
		p.BufferPool.Put(buf)
	}
}

func (p *ReverseProxy) logf(format string, args ...interface{}) {
	if p.ErrorLog != nil {
		p.ErrorLog.Printf(format, args...)
	} else {
		log.Printf(format, args...)
	}
}

type writeFlusher interface {
	io.Writer
	http.Flusher
}

type maxLatencyWriter struct {
	dst     writeFlusher
	latency time.Duration

	lk   sync.Mutex // protects Write + Flush
	done chan bool
}

func (m *maxLatencyWriter) Write(p []byte) (int, error) {
	m.lk.Lock()
	defer m.lk.Unlock()
	return m.dst.Write(p)
}

func (m *maxLatencyWriter) flushLoop() {
	t := time.NewTicker(m.latency)
	defer t.Stop()
	for {
		select {
		case <-m.done:
			if onExitFlushLoop != nil {
				onExitFlushLoop()
			}
			return
		case <-t.C:
			m.lk.Lock()
			m.dst.Flush()
			m.lk.Unlock()
		}
	}
}

func (m *maxLatencyWriter) stop() { m.done <- true }
<|endoftext|>


<commit_before>// Package feed describes RSS format and provides parsing
package feed

// based on http://siongui.github.io/2015/03/03/go-parse-web-feed-rss-atom/

import (
	"bufio"
	"bytes"
	"encoding/xml"
	"fmt"
	"html/template"
	"io"
	"net/http"
	"strings"
	"time"

	log "github.com/go-pkgz/lgr"
	"github.com/pkg/errors"
)

// Rss2 feed
type Rss2 struct {
	XMLName       xml.Name `xml:"rss"`
	Version       string   `xml:"version,attr"`
	Title         string   `xml:"channel>title"`
	Language      string   `xml:"channel>lang"`
	Link          string   `xml:"channel>link"`
	Description   string   `xml:"channel>description"`
	PubDate       string   `xml:"channel>pubDate"`
	LastBuildDate string   `xml:"channel>lastBuildDate"`

	ItemList []Item `xml:"channel>item"`
}

// Item for rss
type Item struct {
	// Required
	Title       string        `xml:"title"`
	Link        string        `xml:"link"`
	Description template.HTML `xml:"description"`
	// Optional
	Content   template.HTML `xml:"encoded"`
	PubDate   string        `xml:"pubDate"`
	Comments  string        `xml:"comments"`
	Enclosure Enclosure     `xml:"enclosure"`
	GUID      string        `xml:"guid"`

	// Internal
	DT time.Time `xml:"-"`
}

// Enclosure element from item
type Enclosure struct {
	URL    string `xml:"url,attr"`
	Length int    `xml:"length,attr"`
	Type   string `xml:"type,attr"`
}

// Atom1 is atom feed
type Atom1 struct {
	XMLName   xml.Name `xml:"http://www.w3.org/2005/Atom feed"`
	Title     string   `xml:"title"`
	Subtitle  string   `xml:"subtitle"`
	ID        string   `xml:"id"`
	Updated   string   `xml:"updated"`
	Rights    string   `xml:"rights"`
	Link      Link     `xml:"link"`
	Author    Author   `xml:"author"`
	EntryList []Entry  `xml:"entry"`
}

// Link element for xml
type Link struct {
	Href string `xml:"href,attr"`
}

// Author element for xml
type Author struct {
	Name  string `xml:"name"`
	Email string `xml:"email"`
}

// Entry from atom
type Entry struct {
	Title     string    `xml:"title"`
	Summary   string    `xml:"summary"`
	Content   string    `xml:"content"`
	ID        string    `xml:"id"`
	Updated   string    `xml:"updated"`
	Link      Link      `xml:"link"`
	Author    Author    `xml:"author"`
	Enclosure Enclosure `xml:"enclosure"`
}

// Parse gets url to rss feed and returns Rss2 items
func Parse(uri string) (result Rss2, err error) {
	resp, err := http.Get(uri) // nolint
	if err != nil {
		return result, err
	}
	defer func() {
		if e := resp.Body.Close(); e != nil {
			log.Printf("[WARN] failed to close body, %s", e)
		}
	}()

	var b bytes.Buffer
	bwriter := bufio.NewWriter(&b)
	if _, err = io.Copy(bwriter, resp.Body); err != nil {
		return result, err
	}

	result, e := parseFeedContent(b.Bytes())
	if e != nil {
		return Rss2{}, errors.Wrap(err, "parsing error")
	}

	return result.Normalize()
}

func atom1ToRss2(a Atom1) Rss2 {
	r := Rss2{
		Title:       a.Title,
		Link:        a.Link.Href,
		Description: a.Subtitle,
		PubDate:     a.Updated,
	}
	r.ItemList = make([]Item, len(a.EntryList))
	for i, entry := range a.EntryList {
		r.ItemList[i].Title = entry.Title
		r.ItemList[i].Link = entry.Link.Href
		if entry.Content == "" {
			r.ItemList[i].Description = template.HTML(entry.Summary) // nolint
		} else {
			r.ItemList[i].Description = template.HTML(entry.Content) // nolint
		}
	}
	return r
}

const atomErrStr = "expected element type <rss> but have <feed>"

func parseAtom(content []byte) (Rss2, error) {
	a := Atom1{}
	err := xml.Unmarshal(content, &a)
	if err != nil {
		return Rss2{}, errors.Wrap(err, "can't parse atom1")
	}
	return atom1ToRss2(a), nil
}

func parseFeedContent(content []byte) (Rss2, error) {
	v := Rss2{}
	err := xml.Unmarshal(content, &v)
	if err != nil {
		if err.Error() == atomErrStr {
			// try Atom 1.0
			return parseAtom(content)
		}
		return v, errors.Wrap(err, "can't parse feed content")
	}

	if v.Version == "2.0" {
		// RSS 2.0
		for i := range v.ItemList {
			if v.ItemList[i].Content != "" {
				v.ItemList[i].Description = v.ItemList[i].Content
			}
		}
		return v, nil
	}

	return v, errors.New("not RSS 2.0")
}

// Normalize converts to RFC822 = "02 Jan 06 15:04 MST"
func (rss *Rss2) Normalize() (Rss2, error) {

	dt, err := rss.normalizeDate(rss.LastBuildDate)
	if err != nil {
		dt, err = rss.normalizeDate(rss.PubDate)
	}

	if err == nil {
		rss.PubDate = dt.Format(time.RFC1123Z)
	}

	for i, item := range rss.ItemList {
		if dt, err := rss.normalizeDate(item.PubDate); err == nil {
			rss.ItemList[i].DT = dt
			rss.ItemList[i].PubDate = dt.Format(time.RFC1123Z)
		}
		rss.ItemList[i].Title = strings.Replace(item.Title, "\n", "", -1)
		rss.ItemList[i].Title = strings.TrimSpace(rss.ItemList[i].Title)
	}
	return *rss, nil
}

func (rss *Rss2) normalizeDate(dt string) (time.Time, error) {
	if dt == "" {
		return time.Now(), fmt.Errorf("can't normalize empty pubDate")
	}
	if ts, err := time.Parse(time.RFC822, dt); err == nil {
		return ts, nil
	}
	if ts, err := time.Parse(time.RFC822Z, dt); err == nil {
		return ts, nil
	}
	if ts, err := time.Parse(time.RFC1123Z, dt); err == nil {
		return ts, nil
	}
	if ts, err := time.Parse(time.RFC1123, dt); err == nil {
		return ts, nil
	}
	if ts, err := time.Parse("2006-01-02 15:04:05 -0700", dt); err == nil {
		return ts, nil
	}
	log.Printf("[DEBUG] can't normalize %s", dt)
	return time.Now(), fmt.Errorf("can't normalize %s", dt)
}
<commit_msg>change lang tag to language<commit_after>// Package feed describes RSS format and provides parsing
package feed

// based on http://siongui.github.io/2015/03/03/go-parse-web-feed-rss-atom/

import (
	"bufio"
	"bytes"
	"encoding/xml"
	"fmt"
	"html/template"
	"io"
	"net/http"
	"strings"
	"time"

	log "github.com/go-pkgz/lgr"
	"github.com/pkg/errors"
)

// Rss2 feed
type Rss2 struct {
	XMLName       xml.Name `xml:"rss"`
	Version       string   `xml:"version,attr"`
	Title         string   `xml:"channel>title"`
	Language      string   `xml:"channel>language"`
	Link          string   `xml:"channel>link"`
	Description   string   `xml:"channel>description"`
	PubDate       string   `xml:"channel>pubDate"`
	LastBuildDate string   `xml:"channel>lastBuildDate"`

	ItemList []Item `xml:"channel>item"`
}

// Item for rss
type Item struct {
	// Required
	Title       string        `xml:"title"`
	Link        string        `xml:"link"`
	Description template.HTML `xml:"description"`
	// Optional
	Content   template.HTML `xml:"encoded"`
	PubDate   string        `xml:"pubDate"`
	Comments  string        `xml:"comments"`
	Enclosure Enclosure     `xml:"enclosure"`
	GUID      string        `xml:"guid"`

	// Internal
	DT time.Time `xml:"-"`
}

// Enclosure element from item
type Enclosure struct {
	URL    string `xml:"url,attr"`
	Length int    `xml:"length,attr"`
	Type   string `xml:"type,attr"`
}

// Atom1 is atom feed
type Atom1 struct {
	XMLName   xml.Name `xml:"http://www.w3.org/2005/Atom feed"`
	Title     string   `xml:"title"`
	Subtitle  string   `xml:"subtitle"`
	ID        string   `xml:"id"`
	Updated   string   `xml:"updated"`
	Rights    string   `xml:"rights"`
	Link      Link     `xml:"link"`
	Author    Author   `xml:"author"`
	EntryList []Entry  `xml:"entry"`
}

// Link element for xml
type Link struct {
	Href string `xml:"href,attr"`
}

// Author element for xml
type Author struct {
	Name  string `xml:"name"`
	Email string `xml:"email"`
}

// Entry from atom
type Entry struct {
	Title     string    `xml:"title"`
	Summary   string    `xml:"summary"`
	Content   string    `xml:"content"`
	ID        string    `xml:"id"`
	Updated   string    `xml:"updated"`
	Link      Link      `xml:"link"`
	Author    Author    `xml:"author"`
	Enclosure Enclosure `xml:"enclosure"`
}

// Parse gets url to rss feed and returns Rss2 items
func Parse(uri string) (result Rss2, err error) {
	resp, err := http.Get(uri) // nolint
	if err != nil {
		return result, err
	}
	defer func() {
		if e := resp.Body.Close(); e != nil {
			log.Printf("[WARN] failed to close body, %s", e)
		}
	}()

	var b bytes.Buffer
	bwriter := bufio.NewWriter(&b)
	if _, err = io.Copy(bwriter, resp.Body); err != nil {
		return result, err
	}

	result, e := parseFeedContent(b.Bytes())
	if e != nil {
		return Rss2{}, errors.Wrap(err, "parsing error")
	}

	return result.Normalize()
}

func atom1ToRss2(a Atom1) Rss2 {
	r := Rss2{
		Title:       a.Title,
		Link:        a.Link.Href,
		Description: a.Subtitle,
		PubDate:     a.Updated,
	}
	r.ItemList = make([]Item, len(a.EntryList))
	for i, entry := range a.EntryList {
		r.ItemList[i].Title = entry.Title
		r.ItemList[i].Link = entry.Link.Href
		if entry.Content == "" {
			r.ItemList[i].Description = template.HTML(entry.Summary) // nolint
		} else {
			r.ItemList[i].Description = template.HTML(entry.Content) // nolint
		}
	}
	return r
}

const atomErrStr = "expected element type <rss> but have <feed>"

func parseAtom(content []byte) (Rss2, error) {
	a := Atom1{}
	err := xml.Unmarshal(content, &a)
	if err != nil {
		return Rss2{}, errors.Wrap(err, "can't parse atom1")
	}
	return atom1ToRss2(a), nil
}

func parseFeedContent(content []byte) (Rss2, error) {
	v := Rss2{}
	err := xml.Unmarshal(content, &v)
	if err != nil {
		if err.Error() == atomErrStr {
			// try Atom 1.0
			return parseAtom(content)
		}
		return v, errors.Wrap(err, "can't parse feed content")
	}

	if v.Version == "2.0" {
		// RSS 2.0
		for i := range v.ItemList {
			if v.ItemList[i].Content != "" {
				v.ItemList[i].Description = v.ItemList[i].Content
			}
		}
		return v, nil
	}

	return v, errors.New("not RSS 2.0")
}

// Normalize converts to RFC822 = "02 Jan 06 15:04 MST"
func (rss *Rss2) Normalize() (Rss2, error) {

	dt, err := rss.normalizeDate(rss.LastBuildDate)
	if err != nil {
		dt, err = rss.normalizeDate(rss.PubDate)
	}

	if err == nil {
		rss.PubDate = dt.Format(time.RFC1123Z)
	}

	for i, item := range rss.ItemList {
		if dt, err := rss.normalizeDate(item.PubDate); err == nil {
			rss.ItemList[i].DT = dt
			rss.ItemList[i].PubDate = dt.Format(time.RFC1123Z)
		}
		rss.ItemList[i].Title = strings.Replace(item.Title, "\n", "", -1)
		rss.ItemList[i].Title = strings.TrimSpace(rss.ItemList[i].Title)
	}
	return *rss, nil
}

func (rss *Rss2) normalizeDate(dt string) (time.Time, error) {
	if dt == "" {
		return time.Now(), fmt.Errorf("can't normalize empty pubDate")
	}
	if ts, err := time.Parse(time.RFC822, dt); err == nil {
		return ts, nil
	}
	if ts, err := time.Parse(time.RFC822Z, dt); err == nil {
		return ts, nil
	}
	if ts, err := time.Parse(time.RFC1123Z, dt); err == nil {
		return ts, nil
	}
	if ts, err := time.Parse(time.RFC1123, dt); err == nil {
		return ts, nil
	}
	if ts, err := time.Parse("2006-01-02 15:04:05 -0700", dt); err == nil {
		return ts, nil
	}
	log.Printf("[DEBUG] can't normalize %s", dt)
	return time.Now(), fmt.Errorf("can't normalize %s", dt)
}
<|endoftext|>


<commit_before>// Copyright 2013, 2015 Canonical Ltd.
// Copyright 2015 Cloudbase Solutions SRL
// Licensed under the AGPLv3, see LICENCE file for details.

package containerinit

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	"github.com/juju/utils/proxy"

	"github.com/juju/juju/cloudconfig"
	"github.com/juju/juju/cloudconfig/cloudinit"
	"github.com/juju/juju/cloudconfig/instancecfg"
	"github.com/juju/juju/container"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/network"
	"github.com/juju/juju/service"
	"github.com/juju/juju/service/common"
)

var (
	logger = loggo.GetLogger("juju.cloudconfig.containerinit")
)

// WriteUserData generates the cloud-init user-data using the
// specified machine and network config for a container, and writes
// the serialized form out to a cloud-init file in the directory
// specified.
func WriteUserData(
	instanceConfig *instancecfg.InstanceConfig,
	networkConfig *container.NetworkConfig,
	directory string,
) (string, error) {
	userData, err := CloudInitUserData(instanceConfig, networkConfig)
	if err != nil {
		logger.Errorf("failed to create user data: %v", err)
		return "", err
	}
	return WriteCloudInitFile(directory, userData)
}

// WriteCloudInitFile writes the data out to a cloud-init file in the
// directory specified, and returns the filename.
func WriteCloudInitFile(directory string, userData []byte) (string, error) {
	userDataFilename := filepath.Join(directory, "cloud-init")
	if err := ioutil.WriteFile(userDataFilename, userData, 0644); err != nil {
		logger.Errorf("failed to write user data: %v", err)
		return "", err
	}
	return userDataFilename, nil
}

// networkConfigTemplate defines how to render /etc/network/interfaces
// file for a container with one or more NICs.
const networkConfigTemplate = `
# loopback interface
auto lo
iface lo inet loopback{{define "static"}}
{{.InterfaceName | printf "# interface %q"}}{{if not .NoAutoStart}}
auto {{.InterfaceName}}{{end}}
iface {{.InterfaceName}} inet manual{{if gt (len .DNSServers) 0}}
    dns-nameservers{{range $dns := .DNSServers}} {{$dns.Value}}{{end}}{{end}}{{if gt (len .DNSSearch) 0}}
    dns-search {{.DNSSearch}}{{end}}
    pre-up ip address add {{.Address.Value}}/32 dev {{.InterfaceName}} &> /dev/null || true
    up ip route replace {{.GatewayAddress.Value}} dev {{.InterfaceName}}
    up ip route replace default via {{.GatewayAddress.Value}}
    down ip route del default via {{.GatewayAddress.Value}} &> /dev/null || true
    down ip route del {{.GatewayAddress.Value}} dev {{.InterfaceName}} &> /dev/null || true
    post-down ip address del {{.Address.Value}}/32 dev {{.InterfaceName}} &> /dev/null || true
{{end}}{{define "dhcp"}}
{{.InterfaceName | printf "# interface %q"}}{{if not .NoAutoStart}}
auto {{.InterfaceName}}{{end}}
iface {{.InterfaceName}} inet dhcp
{{end}}{{range $nic := . }}{{if eq $nic.ConfigType "static"}}
{{template "static" $nic}}{{else}}{{template "dhcp" $nic}}{{end}}{{end}}`

// multiBridgeNetworkConfigTemplate defines how to render /etc/network/interfaces
// file for a multi-NIC container.
const multiBridgeNetworkConfigTemplate = `
auto lo
iface lo inet loopback
{{range $nic := .}}{{template "single" $nic}}{{end}}
{{define "single"}}{{if not .NoAutoStart}}
auto {{.InterfaceName}}{{end}}
iface {{.InterfaceName}} inet manual{{if .DNSServers}}
  dns-nameservers{{range $srv := .DNSServers}} {{$srv.Value}}{{end}}{{end}}{{if .DNSSearchDomains}}
  dns-search{{range $dom := .DNSSearchDomains}} {{$dom}}{{end}}{{end}}
  pre-up ip address add {{.CIDRAddress}} dev {{.InterfaceName}} || true
  up ip route replace {{.CIDR}} dev {{.InterfaceName}} || true
  down ip route del {{.CIDR}} dev {{.InterfaceName}} || true
  post-down address del {{.CIDRAddress}} dev {{.InterfaceName}} || true{{if .GatewayAddress.Value}}
  up ip route replace default via {{.GatewayAddress.Value}} || true
  down ip route del default via {{.GatewayAddress.Value}} || true{{end}}
{{end}}`

var networkInterfacesFile = "/etc/network/interfaces.d/00-juju.cfg"

// GenerateNetworkConfig renders a network config for one or more
// network interfaces, using the given non-nil networkConfig
// containing a non-empty Interfaces field.
func GenerateNetworkConfig(networkConfig *container.NetworkConfig) (string, error) {
	if networkConfig == nil || len(networkConfig.Interfaces) == 0 {
		// Don't generate networking config.
		logger.Tracef("no network config to generate")
		return "", nil
	}
	logger.Debugf("generating network config from %#v", *networkConfig)

	// Copy the InterfaceInfo before modifying the original.
	interfacesCopy := make([]network.InterfaceInfo, len(networkConfig.Interfaces))
	copy(interfacesCopy, networkConfig.Interfaces)
	for i, info := range interfacesCopy {
		if info.MACAddress != "" {
			info.MACAddress = ""
		}
		if info.InterfaceName != "eth0" {
			info.GatewayAddress = network.Address{}
		}
		interfacesCopy[i] = info
	}

	// Render the config first.
	tmpl, err := template.New("config").Parse(multiBridgeNetworkConfigTemplate)
	if err != nil {
		return "", errors.Annotate(err, "cannot parse network config template")
	}

	var buf bytes.Buffer
	if err := tmpl.Execute(&buf, interfacesCopy); err != nil {
		return "", errors.Annotate(err, "cannot render network config")
	}

	generatedConfig := buf.String()
	logger.Debugf("generated network config from %#v\nusing%#v:\n%s", interfacesCopy, networkConfig.Interfaces, generatedConfig)

	return generatedConfig, nil
}

// newCloudInitConfigWithNetworks creates a cloud-init config which
// might include per-interface networking config if both networkConfig
// is not nil and its Interfaces field is not empty.
func newCloudInitConfigWithNetworks(series string, networkConfig *container.NetworkConfig) (cloudinit.CloudConfig, error) {
	cloudConfig, err := cloudinit.New(series)
	if err != nil {
		return nil, errors.Trace(err)
	}
	config, err := GenerateNetworkConfig(networkConfig)
	if err != nil || len(config) == 0 {
		return cloudConfig, errors.Trace(err)
	}

	cloudConfig.AddBootTextFile(networkInterfacesFile, config, 0644)
	cloudConfig.AddRunCmd("ifup -a || true")
	return cloudConfig, nil
}

func CloudInitUserData(
	instanceConfig *instancecfg.InstanceConfig,
	networkConfig *container.NetworkConfig,
) ([]byte, error) {
	cloudConfig, err := newCloudInitConfigWithNetworks(instanceConfig.Series, networkConfig)
	if err != nil {
		return nil, errors.Trace(err)
	}
	udata, err := cloudconfig.NewUserdataConfig(instanceConfig, cloudConfig)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if err = udata.Configure(); err != nil {
		return nil, errors.Trace(err)
	}
	// Run ifconfig to get the addresses of the internal container at least
	// logged in the host.
	cloudConfig.AddRunCmd("ifconfig")

	if instanceConfig.MachineContainerHostname != "" {
		cloudConfig.SetAttr("hostname", instanceConfig.MachineContainerHostname)
	}

	data, err := cloudConfig.RenderYAML()
	if err != nil {
		return nil, errors.Trace(err)
	}
	return data, nil
}

// TemplateUserData returns a minimal user data necessary for the template.
// This should have the authorized keys, base packages, the cloud archive if
// necessary,  initial apt proxy config, and it should do the apt-get
// update/upgrade initially.
func TemplateUserData(
	series string,
	authorizedKeys string,
	aptProxy proxy.Settings,
	aptMirror string,
	enablePackageUpdates bool,
	enableOSUpgrades bool,
	networkConfig *container.NetworkConfig,
) ([]byte, error) {
	var config cloudinit.CloudConfig
	var err error
	if networkConfig != nil {
		config, err = newCloudInitConfigWithNetworks(series, networkConfig)
		if err != nil {
			return nil, errors.Trace(err)
		}
	} else {
		config, err = cloudinit.New(series)
		if err != nil {
			return nil, errors.Trace(err)
		}
	}
	cloudconfig.SetUbuntuUser(config, authorizedKeys)
	config.AddScripts(
		"set -xe", // ensure we run all the scripts or abort.
	)
	// For LTS series which need support for the cloud-tools archive,
	// we need to enable apt-get update regardless of the environ
	// setting, otherwise provisioning will fail.
	if series == "precise" && !enablePackageUpdates {
		logger.Infof("series %q requires cloud-tools archive: enabling updates", series)
		enablePackageUpdates = true
	}

	if enablePackageUpdates && config.RequiresCloudArchiveCloudTools() {
		config.AddCloudArchiveCloudTools()
	}
	config.AddPackageCommands(aptProxy, aptMirror, enablePackageUpdates, enableOSUpgrades)

	initSystem, err := service.VersionInitSystem(series)
	if err != nil {
		return nil, errors.Trace(err)
	}
	cmds, err := shutdownInitCommands(initSystem, series)
	if err != nil {
		return nil, errors.Trace(err)
	}
	config.AddScripts(strings.Join(cmds, "\n"))

	data, err := config.RenderYAML()
	if err != nil {
		return nil, err
	}
	return data, nil
}

// defaultEtcNetworkInterfaces is the contents of
// /etc/network/interfaces file which is left on the template LXC
// container on shutdown. This is needed to allow cloned containers to
// start in case no network config is provided during cloud-init, e.g.
// when AUFS is used.
const defaultEtcNetworkInterfaces = `
# loopback interface
auto lo
iface lo inet loopback

# primary interface
auto eth0
iface eth0 inet dhcp
`

func shutdownInitCommands(initSystem, series string) ([]string, error) {
	// These files are removed just before the template shuts down.
	cleanupOnShutdown := []string{
		// We remove any dhclient lease files so there's no chance a
		// clone to reuse a lease from the template it was cloned
		// from.
		"/var/lib/dhcp/dhclient*",
		// Both of these sets of files below are recreated on boot and
		// if we leave them in the template's rootfs boot logs coming
		// from cloned containers will be appended. It's better to
		// keep clean logs for diagnosing issues / debugging.
		"/var/log/cloud-init*.log",
	}

	// Using EOC below as the template shutdown script is itself
	// passed through cat > ... < EOF.
	replaceNetConfCmd := fmt.Sprintf(
		"/bin/cat > /etc/network/interfaces << EOC%sEOC\n  ",
		defaultEtcNetworkInterfaces,
	)
	paths := strings.Join(cleanupOnShutdown, " ")
	removeCmd := fmt.Sprintf("/bin/rm -fr %s\n  ", paths)
	shutdownCmd := "/sbin/shutdown -h now"
	name := "juju-template-restart"
	desc := "juju shutdown job"

	execStart := shutdownCmd
	if environs.AddressAllocationEnabled("") { // we only care the provider is not MAAS here.
		// Only do the cleanup and replacement of /e/n/i when address
		// allocation feature flag is enabled.
		execStart = replaceNetConfCmd + removeCmd + shutdownCmd
	}

	conf := common.Conf{
		Desc:         desc,
		Transient:    true,
		AfterStopped: "cloud-final",
		ExecStart:    execStart,
	}
	// systemd uses targets for synchronization of services
	if initSystem == service.InitSystemSystemd {
		conf.AfterStopped = "cloud-config.target"
	}

	svc, err := service.NewService(name, conf, series)
	if err != nil {
		return nil, errors.Trace(err)
	}

	cmds, err := svc.InstallCommands()
	if err != nil {
		return nil, errors.Trace(err)
	}

	startCommands, err := svc.StartCommands()
	if err != nil {
		return nil, errors.Trace(err)
	}
	cmds = append(cmds, startCommands...)

	return cmds, nil
}
<commit_msg>cloudconfig/containerinit: Refactored network config generation<commit_after>// Copyright 2013, 2015 Canonical Ltd.
// Copyright 2015 Cloudbase Solutions SRL
// Licensed under the AGPLv3, see LICENCE file for details.

package containerinit

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	"github.com/juju/utils/proxy"
	"github.com/juju/utils/set"

	"github.com/juju/juju/cloudconfig"
	"github.com/juju/juju/cloudconfig/cloudinit"
	"github.com/juju/juju/cloudconfig/instancecfg"
	"github.com/juju/juju/container"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/network"
	"github.com/juju/juju/service"
	"github.com/juju/juju/service/common"
)

var (
	logger = loggo.GetLogger("juju.cloudconfig.containerinit")
)

// WriteUserData generates the cloud-init user-data using the
// specified machine and network config for a container, and writes
// the serialized form out to a cloud-init file in the directory
// specified.
func WriteUserData(
	instanceConfig *instancecfg.InstanceConfig,
	networkConfig *container.NetworkConfig,
	directory string,
) (string, error) {
	userData, err := CloudInitUserData(instanceConfig, networkConfig)
	if err != nil {
		logger.Errorf("failed to create user data: %v", err)
		return "", err
	}
	return WriteCloudInitFile(directory, userData)
}

// WriteCloudInitFile writes the data out to a cloud-init file in the
// directory specified, and returns the filename.
func WriteCloudInitFile(directory string, userData []byte) (string, error) {
	userDataFilename := filepath.Join(directory, "cloud-init")
	if err := ioutil.WriteFile(userDataFilename, userData, 0644); err != nil {
		logger.Errorf("failed to write user data: %v", err)
		return "", err
	}
	return userDataFilename, nil
}

var networkInterfacesFile = "/etc/network/interfaces"

// GenerateNetworkConfig renders a network config for one or more network
// interfaces, using the given non-nil networkConfig containing a non-empty
// Interfaces field.
func GenerateNetworkConfig(networkConfig *container.NetworkConfig) (string, error) {
	if networkConfig == nil || len(networkConfig.Interfaces) == 0 {
		logger.Tracef("no network config to generate")
		return "", nil
	}
	logger.Debugf("generating network config from %#v", *networkConfig)

	prepared := PrepareNetworkConfigFromInterfaces(networkConfig.Interfaces)

	var output bytes.Buffer
	gatewayWritten := false
	for _, name := range prepared.InterfaceNames {
		output.WriteString("\n")
		if name == "lo" {
			output.WriteString("auto ")
			autoStarted := strings.Join(prepared.AutoStarted, " ")
			output.WriteString(autoStarted + "\n\n")
			output.WriteString("iface lo inet loopback\n")

			dnsServers := strings.Join(prepared.DNSServers, " ")
			if dnsServers != "" {
				output.WriteString("  dns-nameservers ")
				output.WriteString(dnsServers + "\n")
			}

			dnsSearchDomains := strings.Join(prepared.DNSSearchDomains, " ")
			if dnsSearchDomains != "" {
				output.WriteString("  dns-search ")
				output.WriteString(dnsSearchDomains + "\n")
			}
			continue
		}

		address, hasAddress := prepared.NameToAddress[name]
		if !hasAddress {
			output.WriteString("iface " + name + " inet manual\n")
			continue
		}

		output.WriteString("iface " + name + " inet static\n")
		output.WriteString("  address " + address + "\n")
		if !gatewayWritten && prepared.GatewayAddress != "" {
			output.WriteString("  gateway " + prepared.GatewayAddress + "\n")
			gatewayWritten = true // write it only once
		}
	}

	generatedConfig := output.String()
	logger.Debugf("generated network config:\n%s", generatedConfig)

	return generatedConfig, nil
}

// PreparedConfig holds all the necessary information to render a persistent
// network config to a file.
type PreparedConfig struct {
	InterfaceNames   []string
	AutoStarted      []string
	DNSServers       []string
	DNSSearchDomains []string
	NameToAddress    map[string]string
	GatewayAddress   string
}

// PrepareNetworkConfigFromInterfaces collects the necessary information to
// render a persistent network config from the given slice of
// network.InterfaceInfo. The result always includes the loopback interface.
func PrepareNetworkConfigFromInterfaces(interfaces []network.InterfaceInfo) *PreparedConfig {
	// Collect those two as they need all to rendered once.
	dnsServers := set.NewStrings()
	dnsSearchDomains := set.NewStrings()
	gatewayAddress := ""

	// Collect those for later rendering.
	namesInOrder := make([]string, 1, len(interfaces)+1)
	nameToAddress := make(map[string]string)

	// Always include the loopback.
	namesInOrder[0] = "lo"
	autoStarted := set.NewStrings("lo")

	for _, info := range interfaces {
		if !info.NoAutoStart {
			autoStarted.Add(info.InterfaceName)
		}

		if cidr := info.CIDRAddress(); cidr != "" {
			nameToAddress[info.InterfaceName] = cidr
		}

		for _, dns := range info.DNSServers {
			dnsServers.Add(dns.Value)
		}

		dnsSearchDomains = dnsSearchDomains.Union(set.NewStrings(info.DNSSearchDomains...))

		if info.InterfaceName == "eth0" && gatewayAddress == "" {
			// Only set gateway once for the primary NIC.
			gatewayAddress = info.GatewayAddress.Value
		}

		namesInOrder = append(namesInOrder, info.InterfaceName)
	}

	prepared := &PreparedConfig{
		InterfaceNames:   namesInOrder,
		NameToAddress:    nameToAddress,
		AutoStarted:      autoStarted.SortedValues(),
		DNSServers:       dnsServers.SortedValues(),
		DNSSearchDomains: dnsSearchDomains.SortedValues(),
		GatewayAddress:   gatewayAddress,
	}

	logger.Debugf("prepared network config for rendering: %+v", prepared)
	return prepared
}

// newCloudInitConfigWithNetworks creates a cloud-init config which
// might include per-interface networking config if both networkConfig
// is not nil and its Interfaces field is not empty.
func newCloudInitConfigWithNetworks(series string, networkConfig *container.NetworkConfig) (cloudinit.CloudConfig, error) {
	cloudConfig, err := cloudinit.New(series)
	if err != nil {
		return nil, errors.Trace(err)
	}
	config, err := GenerateNetworkConfig(networkConfig)
	if err != nil || len(config) == 0 {
		return cloudConfig, errors.Trace(err)
	}

	cloudConfig.AddBootTextFile(networkInterfacesFile, config, 0644)
	cloudConfig.AddRunCmd("ifup -a || true")
	return cloudConfig, nil
}

func CloudInitUserData(
	instanceConfig *instancecfg.InstanceConfig,
	networkConfig *container.NetworkConfig,
) ([]byte, error) {
	cloudConfig, err := newCloudInitConfigWithNetworks(instanceConfig.Series, networkConfig)
	if err != nil {
		return nil, errors.Trace(err)
	}
	udata, err := cloudconfig.NewUserdataConfig(instanceConfig, cloudConfig)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if err = udata.Configure(); err != nil {
		return nil, errors.Trace(err)
	}
	// Run ifconfig to get the addresses of the internal container at least
	// logged in the host.
	cloudConfig.AddRunCmd("ifconfig")

	if instanceConfig.MachineContainerHostname != "" {
		cloudConfig.SetAttr("hostname", instanceConfig.MachineContainerHostname)
	}

	data, err := cloudConfig.RenderYAML()
	if err != nil {
		return nil, errors.Trace(err)
	}
	return data, nil
}

// TemplateUserData returns a minimal user data necessary for the template.
// This should have the authorized keys, base packages, the cloud archive if
// necessary,  initial apt proxy config, and it should do the apt-get
// update/upgrade initially.
func TemplateUserData(
	series string,
	authorizedKeys string,
	aptProxy proxy.Settings,
	aptMirror string,
	enablePackageUpdates bool,
	enableOSUpgrades bool,
	networkConfig *container.NetworkConfig,
) ([]byte, error) {
	var config cloudinit.CloudConfig
	var err error
	if networkConfig != nil {
		config, err = newCloudInitConfigWithNetworks(series, networkConfig)
		if err != nil {
			return nil, errors.Trace(err)
		}
	} else {
		config, err = cloudinit.New(series)
		if err != nil {
			return nil, errors.Trace(err)
		}
	}
	cloudconfig.SetUbuntuUser(config, authorizedKeys)
	config.AddScripts(
		"set -xe", // ensure we run all the scripts or abort.
	)
	// For LTS series which need support for the cloud-tools archive,
	// we need to enable apt-get update regardless of the environ
	// setting, otherwise provisioning will fail.
	if series == "precise" && !enablePackageUpdates {
		logger.Infof("series %q requires cloud-tools archive: enabling updates", series)
		enablePackageUpdates = true
	}

	if enablePackageUpdates && config.RequiresCloudArchiveCloudTools() {
		config.AddCloudArchiveCloudTools()
	}
	config.AddPackageCommands(aptProxy, aptMirror, enablePackageUpdates, enableOSUpgrades)

	initSystem, err := service.VersionInitSystem(series)
	if err != nil {
		return nil, errors.Trace(err)
	}
	cmds, err := shutdownInitCommands(initSystem, series)
	if err != nil {
		return nil, errors.Trace(err)
	}
	config.AddScripts(strings.Join(cmds, "\n"))

	data, err := config.RenderYAML()
	if err != nil {
		return nil, err
	}
	return data, nil
}

// defaultEtcNetworkInterfaces is the contents of
// /etc/network/interfaces file which is left on the template LXC
// container on shutdown. This is needed to allow cloned containers to
// start in case no network config is provided during cloud-init, e.g.
// when AUFS is used.
const defaultEtcNetworkInterfaces = `
# loopback interface
auto lo
iface lo inet loopback

# primary interface
auto eth0
iface eth0 inet dhcp
`

func shutdownInitCommands(initSystem, series string) ([]string, error) {
	// These files are removed just before the template shuts down.
	cleanupOnShutdown := []string{
		// We remove any dhclient lease files so there's no chance a
		// clone to reuse a lease from the template it was cloned
		// from.
		"/var/lib/dhcp/dhclient*",
		// Both of these sets of files below are recreated on boot and
		// if we leave them in the template's rootfs boot logs coming
		// from cloned containers will be appended. It's better to
		// keep clean logs for diagnosing issues / debugging.
		"/var/log/cloud-init*.log",
	}

	// Using EOC below as the template shutdown script is itself
	// passed through cat > ... < EOF.
	replaceNetConfCmd := fmt.Sprintf(
		"/bin/cat > /etc/network/interfaces << EOC%sEOC\n  ",
		defaultEtcNetworkInterfaces,
	)
	paths := strings.Join(cleanupOnShutdown, " ")
	removeCmd := fmt.Sprintf("/bin/rm -fr %s\n  ", paths)
	shutdownCmd := "/sbin/shutdown -h now"
	name := "juju-template-restart"
	desc := "juju shutdown job"

	execStart := shutdownCmd
	if environs.AddressAllocationEnabled("") { // we only care the provider is not MAAS here.
		// Only do the cleanup and replacement of /e/n/i when address
		// allocation feature flag is enabled.
		execStart = replaceNetConfCmd + removeCmd + shutdownCmd
	}

	conf := common.Conf{
		Desc:         desc,
		Transient:    true,
		AfterStopped: "cloud-final",
		ExecStart:    execStart,
	}
	// systemd uses targets for synchronization of services
	if initSystem == service.InitSystemSystemd {
		conf.AfterStopped = "cloud-config.target"
	}

	svc, err := service.NewService(name, conf, series)
	if err != nil {
		return nil, errors.Trace(err)
	}

	cmds, err := svc.InstallCommands()
	if err != nil {
		return nil, errors.Trace(err)
	}

	startCommands, err := svc.StartCommands()
	if err != nil {
		return nil, errors.Trace(err)
	}
	cmds = append(cmds, startCommands...)

	return cmds, nil
}
<|endoftext|>


<commit_before>// Copyright 2017 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"
	"errors"
	"net/http"
	"os"
	"os/exec"
	"testing"
	"time"

	"github.com/Stackdriver/stackdriver-prometheus-sidecar/metadata"
	"github.com/Stackdriver/stackdriver-prometheus-sidecar/retrieval"
	"github.com/go-kit/kit/log"
	"github.com/google/go-cmp/cmp"
	"github.com/prometheus/prometheus/pkg/labels"
	"github.com/prometheus/prometheus/pkg/textparse"
	"github.com/prometheus/prometheus/promql"
	metric_pb "google.golang.org/genproto/googleapis/api/metric"
)

var promPath string

func TestMain(m *testing.M) {
	if os.Getenv("RUN_MAIN") == "" {
		// Run the test directly.
		os.Exit(m.Run())
	}

	main()
}

// As soon as prometheus starts responding to http request should be able to accept Interrupt signals for a gracefull shutdown.
func TestStartupInterrupt(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}

	prom := exec.Command(os.Args[0], "--stackdriver.project-id=1234", "--prometheus.wal-directory=testdata/wal")
	prom.Env = append(os.Environ(), "RUN_MAIN=1")
	var bout, berr bytes.Buffer
	prom.Stdout = &bout
	prom.Stderr = &berr
	err := prom.Start()
	if err != nil {
		t.Errorf("execution error: %v", err)
		return
	}

	done := make(chan error)
	go func() {
		done <- prom.Wait()
	}()

	var startedOk bool
	var stoppedErr error

Loop:
	for x := 0; x < 10; x++ {
		// error=nil means the sidecar has started so can send the interrupt signal and wait for the grace shutdown.
		if _, err := http.Get("http://localhost:9091/metrics"); err == nil {
			startedOk = true
			prom.Process.Signal(os.Interrupt)
			select {
			case stoppedErr = <-done:
				break Loop
			case <-time.After(10 * time.Second):
			}
			break Loop
		} else {
			select {
			case stoppedErr = <-done:
				break Loop
			default: // try again
			}
		}
		time.Sleep(500 * time.Millisecond)
	}

	t.Logf("stodut: %v\n", bout.String())
	t.Logf("stderr: %v\n", berr.String())
	if !startedOk {
		t.Errorf("prometheus didn't start in the specified timeout")
		return
	}
	if err := prom.Process.Kill(); err == nil {
		t.Errorf("prometheus didn't shutdown gracefully after sending the Interrupt signal")
	} else if stoppedErr != nil && stoppedErr.Error() != "signal: interrupt" { // TODO - find a better way to detect when the process didn't exit as expected!
		t.Errorf("prometheus exited with an unexpected error:%v", stoppedErr)
	}
}

func TestParseWhitelists(t *testing.T) {
	logger := log.NewLogfmtLogger(log.NewSyncWriter(os.Stderr))
	for _, tt := range []struct {
		name         string
		filtersets   []string
		filters      []string
		wantMatchers int
	}{
		{"both filters and filtersets defined",
			[]string{
				`metric_name`,
				`metric_name{label="value"}`,
			},
			[]string{
				`__name__="test1"`,
				`a1=~"test2.+"`,
				`a2!="test3"`,
				`a3!~"test4.*"`,
			}, 3},
		{"just filtersets", []string{"metric_name"}, []string{}, 1},
		{"just filters", []string{}, []string{`__name__="foo"`}, 1},
		{"neither filtersets nor filters", []string{}, []string{}, 0},
	} {
		t.Run(tt.name, func(t *testing.T) {
			// Test success cases.
			parsed, err := parseFiltersets(logger, tt.filtersets, tt.filters)
			if err != nil {
				t.Fatal(err)
			}
			if len(parsed) != tt.wantMatchers {
				t.Fatalf("expected %d matchers; got %d", tt.wantMatchers, len(parsed))
			}
		})
	}

	// Test failure cases.
	for _, tt := range []struct {
		name       string
		filtersets []string
		filters    []string
	}{
		{"Invalid character in key", []string{}, []string{`a-b="1"`}},
		{"Missing trailing quote", []string{}, []string{`a="1`}},
		{"Missing leading quote", []string{}, []string{`a=1"`}},
		{"Invalid operator", []string{}, []string{`a!=="1"`}},
		{"Invalid operator in filterset", []string{`{a!=="1"}`}, []string{}},
		{"Empty filterset", []string{""}, []string{}},
	} {
		t.Run(tt.name, func(t *testing.T) {
			if _, err := parseFiltersets(logger, tt.filtersets, tt.filters); err == nil {
				t.Fatalf("expected error, but got none")
			}
		})
	}
}

func TestProcessFileConfig(t *testing.T) {
	mustParseMetricSelector := func(input string) []*labels.Matcher {
		m, err := promql.ParseMetricSelector(input)
		if err != nil {
			t.Fatalf("bad test input %v: %v", input, err)
		}
		return m
	}
	for _, tt := range []struct {
		name           string
		config         fileConfig
		renameMappings map[string]string
		staticMetadata []*metadata.Entry
		aggregations   retrieval.CounterAggregatorConfig
		err            error
	}{
		{
			"empty",
			fileConfig{},
			map[string]string{},
			[]*metadata.Entry{},
			retrieval.CounterAggregatorConfig{},
			nil,
		},
		{
			"smoke",
			fileConfig{
				MetricRenames: []metricRenamesConfig{
					{From: "from", To: "to"},
				},
				StaticMetadata: []staticMetadataConfig{
					{Metric: "int64_counter", Type: "counter", ValueType: "int64", Help: "help1"},
					{Metric: "double_gauge", Type: "gauge", ValueType: "double", Help: "help2"},
					{Metric: "default_gauge", Type: "gauge"},
				},
				AggregatedCounters: []aggregatedCountersConfig{
					{
						Metric:  "network_transmit_bytes",
						Help:    "total number of bytes sent over eth0",
						Filters: []string{"filter1", "filter2"},
					},
				},
			},
			map[string]string{"from": "to"},
			[]*metadata.Entry{
				&metadata.Entry{Metric: "int64_counter", MetricType: textparse.MetricTypeCounter, ValueType: metric_pb.MetricDescriptor_INT64, Help: "help1"},
				&metadata.Entry{Metric: "double_gauge", MetricType: textparse.MetricTypeGauge, ValueType: metric_pb.MetricDescriptor_DOUBLE, Help: "help2"},
				&metadata.Entry{Metric: "default_gauge", MetricType: textparse.MetricTypeGauge},
			},
			retrieval.CounterAggregatorConfig{
				"network_transmit_bytes": &retrieval.CounterAggregatorMetricConfig{
					Matchers: [][]*labels.Matcher{
						mustParseMetricSelector("filter1"),
						mustParseMetricSelector("filter2"),
					},
					Help: "total number of bytes sent over eth0",
				},
			},
			nil,
		},
		{
			"missing_metric_type",
			fileConfig{
				StaticMetadata: []staticMetadataConfig{{Metric: "int64_default", ValueType: "int64"}},
			},
			nil, nil, nil,
			errors.New("invalid metric type \"\""),
		},
	} {
		t.Run(tt.name, func(t *testing.T) {
			renameMappings, staticMetadata, aggregations, err := processFileConfig(tt.config)
			if diff := cmp.Diff(tt.renameMappings, renameMappings); diff != "" {
				t.Errorf("renameMappings mismatch: %v", diff)
			}
			if diff := cmp.Diff(tt.staticMetadata, staticMetadata); diff != "" {
				t.Errorf("staticMetadata mismatch: %v", diff)
			}
			if diff := cmp.Diff(tt.aggregations, aggregations); diff != "" {
				t.Errorf("aggregations mismatch: %v", diff)
			}
			if (tt.err != nil && err != nil && tt.err.Error() != err.Error()) ||
				(tt.err == nil && err != nil) || (tt.err != nil && err == nil) {
				t.Errorf("error mismatch: got %v, expected %v", err, tt.err)
			}
		})
	}
}
<commit_msg>Make main_test.go more precise and clear. (#202)<commit_after>// Copyright 2017 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"
	"errors"
	"net/http"
	"os"
	"os/exec"
	"testing"
	"time"

	"github.com/Stackdriver/stackdriver-prometheus-sidecar/metadata"
	"github.com/Stackdriver/stackdriver-prometheus-sidecar/retrieval"
	"github.com/go-kit/kit/log"
	"github.com/google/go-cmp/cmp"
	"github.com/prometheus/prometheus/pkg/labels"
	"github.com/prometheus/prometheus/pkg/textparse"
	"github.com/prometheus/prometheus/promql"
	metric_pb "google.golang.org/genproto/googleapis/api/metric"
)

func TestMain(m *testing.M) {
	if os.Getenv("RUN_MAIN") == "" {
		// Run the test directly.
		os.Exit(m.Run())
	}

	main()
}

// As soon as prometheus starts responding to http request should be able to accept Interrupt signals for a gracefull shutdown.
func TestStartupInterrupt(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}

	cmd := exec.Command(os.Args[0], "--stackdriver.project-id=1234", "--prometheus.wal-directory=testdata/wal")
	cmd.Env = append(os.Environ(), "RUN_MAIN=1")
	var bout, berr bytes.Buffer
	cmd.Stdout = &bout
	cmd.Stderr = &berr
	err := cmd.Start()
	if err != nil {
		t.Errorf("execution error: %v", err)
		return
	}

	done := make(chan error)
	go func() {
		done <- cmd.Wait()
	}()

	var startedOk bool
	var stoppedErr error

Loop:
	for x := 0; x < 10; x++ {
		// error=nil means the sidecar has started so can send the interrupt signal and wait for the grace shutdown.
		if _, err := http.Get("http://localhost:9091/metrics"); err == nil {
			startedOk = true
			cmd.Process.Signal(os.Interrupt)
			select {
			case stoppedErr = <-done:
				break Loop
			case <-time.After(10 * time.Second):
			}
			break Loop
		} else {
			select {
			case stoppedErr = <-done:
				break Loop
			default: // try again
			}
		}
		time.Sleep(500 * time.Millisecond)
	}

	t.Logf("stodut: %v\n", bout.String())
	t.Logf("stderr: %v\n", berr.String())
	if !startedOk {
		t.Errorf("prometheus-stackdriver-sidecar didn't start in the specified timeout")
		return
	}
	if err := cmd.Process.Kill(); err == nil {
		t.Errorf("prometheus-stackdriver-sidecar didn't shutdown gracefully after sending the Interrupt signal")
	} else if stoppedErr != nil && stoppedErr.Error() != "signal: interrupt" { // TODO - find a better way to detect when the process didn't exit as expected!
		t.Errorf("prometheus-stackdriver-sidecar exited with an unexpected error:%v", stoppedErr)
	}
}

func TestParseWhitelists(t *testing.T) {
	logger := log.NewLogfmtLogger(log.NewSyncWriter(os.Stderr))
	for _, tt := range []struct {
		name         string
		filtersets   []string
		filters      []string
		wantMatchers int
	}{
		{"both filters and filtersets defined",
			[]string{
				`metric_name`,
				`metric_name{label="value"}`,
			},
			[]string{
				`__name__="test1"`,
				`a1=~"test2.+"`,
				`a2!="test3"`,
				`a3!~"test4.*"`,
			}, 3},
		{"just filtersets", []string{"metric_name"}, []string{}, 1},
		{"just filters", []string{}, []string{`__name__="foo"`}, 1},
		{"neither filtersets nor filters", []string{}, []string{}, 0},
	} {
		t.Run(tt.name, func(t *testing.T) {
			// Test success cases.
			parsed, err := parseFiltersets(logger, tt.filtersets, tt.filters)
			if err != nil {
				t.Fatal(err)
			}
			if len(parsed) != tt.wantMatchers {
				t.Fatalf("expected %d matchers; got %d", tt.wantMatchers, len(parsed))
			}
		})
	}

	// Test failure cases.
	for _, tt := range []struct {
		name       string
		filtersets []string
		filters    []string
	}{
		{"Invalid character in key", []string{}, []string{`a-b="1"`}},
		{"Missing trailing quote", []string{}, []string{`a="1`}},
		{"Missing leading quote", []string{}, []string{`a=1"`}},
		{"Invalid operator", []string{}, []string{`a!=="1"`}},
		{"Invalid operator in filterset", []string{`{a!=="1"}`}, []string{}},
		{"Empty filterset", []string{""}, []string{}},
	} {
		t.Run(tt.name, func(t *testing.T) {
			if _, err := parseFiltersets(logger, tt.filtersets, tt.filters); err == nil {
				t.Fatalf("expected error, but got none")
			}
		})
	}
}

func TestProcessFileConfig(t *testing.T) {
	mustParseMetricSelector := func(input string) []*labels.Matcher {
		m, err := promql.ParseMetricSelector(input)
		if err != nil {
			t.Fatalf("bad test input %v: %v", input, err)
		}
		return m
	}
	for _, tt := range []struct {
		name           string
		config         fileConfig
		renameMappings map[string]string
		staticMetadata []*metadata.Entry
		aggregations   retrieval.CounterAggregatorConfig
		err            error
	}{
		{
			"empty",
			fileConfig{},
			map[string]string{},
			[]*metadata.Entry{},
			retrieval.CounterAggregatorConfig{},
			nil,
		},
		{
			"smoke",
			fileConfig{
				MetricRenames: []metricRenamesConfig{
					{From: "from", To: "to"},
				},
				StaticMetadata: []staticMetadataConfig{
					{Metric: "int64_counter", Type: "counter", ValueType: "int64", Help: "help1"},
					{Metric: "double_gauge", Type: "gauge", ValueType: "double", Help: "help2"},
					{Metric: "default_gauge", Type: "gauge"},
				},
				AggregatedCounters: []aggregatedCountersConfig{
					{
						Metric:  "network_transmit_bytes",
						Help:    "total number of bytes sent over eth0",
						Filters: []string{"filter1", "filter2"},
					},
				},
			},
			map[string]string{"from": "to"},
			[]*metadata.Entry{
				&metadata.Entry{Metric: "int64_counter", MetricType: textparse.MetricTypeCounter, ValueType: metric_pb.MetricDescriptor_INT64, Help: "help1"},
				&metadata.Entry{Metric: "double_gauge", MetricType: textparse.MetricTypeGauge, ValueType: metric_pb.MetricDescriptor_DOUBLE, Help: "help2"},
				&metadata.Entry{Metric: "default_gauge", MetricType: textparse.MetricTypeGauge},
			},
			retrieval.CounterAggregatorConfig{
				"network_transmit_bytes": &retrieval.CounterAggregatorMetricConfig{
					Matchers: [][]*labels.Matcher{
						mustParseMetricSelector("filter1"),
						mustParseMetricSelector("filter2"),
					},
					Help: "total number of bytes sent over eth0",
				},
			},
			nil,
		},
		{
			"missing_metric_type",
			fileConfig{
				StaticMetadata: []staticMetadataConfig{{Metric: "int64_default", ValueType: "int64"}},
			},
			nil, nil, nil,
			errors.New("invalid metric type \"\""),
		},
	} {
		t.Run(tt.name, func(t *testing.T) {
			renameMappings, staticMetadata, aggregations, err := processFileConfig(tt.config)
			if diff := cmp.Diff(tt.renameMappings, renameMappings); diff != "" {
				t.Errorf("renameMappings mismatch: %v", diff)
			}
			if diff := cmp.Diff(tt.staticMetadata, staticMetadata); diff != "" {
				t.Errorf("staticMetadata mismatch: %v", diff)
			}
			if diff := cmp.Diff(tt.aggregations, aggregations); diff != "" {
				t.Errorf("aggregations mismatch: %v", diff)
			}
			if (tt.err != nil && err != nil && tt.err.Error() != err.Error()) ||
				(tt.err == nil && err != nil) || (tt.err != nil && err == nil) {
				t.Errorf("error mismatch: got %v, expected %v", err, tt.err)
			}
		})
	}
}
<|endoftext|>