<commit_before>package cloudformation

import (
	"fmt"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/client"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/remind101/pkg/reporter"
	"golang.org/x/net/context"
)

var (
	defaultVisibilityHeartbeat = 1 * time.Minute
	defaultNumWorkers          = 10
)

// sqsClient duck types the sqs.SQS interface.
type sqsClient interface {
	ReceiveMessage(*sqs.ReceiveMessageInput) (*sqs.ReceiveMessageOutput, error)
	DeleteMessage(*sqs.DeleteMessageInput) (*sqs.DeleteMessageOutput, error)
	ChangeMessageVisibility(*sqs.ChangeMessageVisibilityInput) (*sqs.ChangeMessageVisibilityOutput, error)
}

// SQSDispatcher pulls messages from SQS, and dispatches them to a handler.
type SQSDispatcher struct {
	// Root context.Context to use. If a reporter.Reporter is embedded,
	// errors generated will be reporter there. If a logger.Logger is
	// embedded, logging will be logged there.
	Context context.Context

	// The SQS queue url to listen for CloudFormation Custom Resource
	// requests.
	QueueURL string

	// When a message is pulled off of sqs, the visibility timeout will be
	// extended by this much, and periodically extended while the handler
	// performs it's work. If this process crashes, then the sqs message
	// will be redelivered later.
	VisibilityHeartbeat time.Duration

	// Number of worker goroutines to start for receiving messages.
	NumWorkers int

	stopped chan struct{}
	after   func(time.Duration) <-chan time.Time
	sqs     sqsClient
}

func newSQSDispatcher(config client.ConfigProvider) *SQSDispatcher {
	return &SQSDispatcher{
		VisibilityHeartbeat: defaultVisibilityHeartbeat,
		NumWorkers:          defaultNumWorkers,
		after:               time.After,
		sqs:                 sqs.New(config),
	}
}

// Start starts multiple goroutines pulling messages off of the queue.
func (q *SQSDispatcher) Start(handle func(context.Context, *sqs.Message) error) {
	var wg sync.WaitGroup
	for i := 0; i < q.NumWorkers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			q.start(handle)
		}()
	}

	wg.Wait()
}

func (q *SQSDispatcher) Stop() {
	close(q.stopped)
}

// start starts a pulling messages off of the queue and passing them to the
// handler.
func (q *SQSDispatcher) start(handle func(context.Context, *sqs.Message) error) {
	var wg sync.WaitGroup
	for {
		select {
		case <-q.stopped:
			wg.Wait()
			return
		default:
			ctx := q.Context

			resp, err := q.sqs.ReceiveMessage(&sqs.ReceiveMessageInput{
				QueueUrl: aws.String(q.QueueURL),
			})
			if err != nil {
				reporter.Report(ctx, err)
				continue
			}

			for _, m := range resp.Messages {
				wg.Add(1)
				go func(m *sqs.Message) {
					defer wg.Done()
					if err := q.handle(ctx, handle, m); err != nil {
						reporter.Report(ctx, err)
					}
				}(m)
			}
		}
	}
}

func (q *SQSDispatcher) handle(ctx context.Context, handle func(context.Context, *sqs.Message) error, message *sqs.Message) (err error) {
	ctx, cancel := context.WithCancel(ctx)

	defer func() {
		if err == nil {
			_, err = q.sqs.DeleteMessage(&sqs.DeleteMessageInput{
				QueueUrl:      aws.String(q.QueueURL),
				ReceiptHandle: message.ReceiptHandle,
			})
		}
	}()

	var t <-chan time.Time
	t, err = q.extendMessageVisibilityTimeout(message.ReceiptHandle)

	errCh := make(chan error)
	go func() { errCh <- handle(ctx, message) }()

	for {
		select {
		case err = <-errCh:
			return
		case <-q.stopped:
			cancel()
		case <-t:
			t, err = q.extendMessageVisibilityTimeout(message.ReceiptHandle)
			if err != nil {
				return
			}
		}
	}
}

// extendMessageVisibilityTimeout extends the messages visibility timeout
// extends the timeout by VisibilityHeartbeat, and returns a channel that will
// receive after half of VisibilityTimeout has elapsed.
func (q *SQSDispatcher) extendMessageVisibilityTimeout(receiptHandle *string) (<-chan time.Time, error) {
	visibilityTimeout := int64(float64(q.VisibilityHeartbeat) / float64(time.Second))

	_, err := q.sqs.ChangeMessageVisibility(&sqs.ChangeMessageVisibilityInput{
		QueueUrl:          aws.String(q.QueueURL),
		ReceiptHandle:     receiptHandle,
		VisibilityTimeout: aws.Int64(visibilityTimeout),
	})
	if err != nil {
		return nil, fmt.Errorf("error extending message visibility timeout: %v", err)
	}

	return q.after(q.VisibilityHeartbeat / 2), nil
}
<commit_msg>Fix typo.<commit_after>package cloudformation

import (
	"fmt"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/client"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/remind101/pkg/reporter"
	"golang.org/x/net/context"
)

var (
	defaultVisibilityHeartbeat = 1 * time.Minute
	defaultNumWorkers          = 10
)

// sqsClient duck types the sqs.SQS interface.
type sqsClient interface {
	ReceiveMessage(*sqs.ReceiveMessageInput) (*sqs.ReceiveMessageOutput, error)
	DeleteMessage(*sqs.DeleteMessageInput) (*sqs.DeleteMessageOutput, error)
	ChangeMessageVisibility(*sqs.ChangeMessageVisibilityInput) (*sqs.ChangeMessageVisibilityOutput, error)
}

// SQSDispatcher pulls messages from SQS, and dispatches them to a handler.
type SQSDispatcher struct {
	// Root context.Context to use. If a reporter.Reporter is embedded,
	// errors generated will be reporter there. If a logger.Logger is
	// embedded, logging will be logged there.
	Context context.Context

	// The SQS queue url to listen for CloudFormation Custom Resource
	// requests.
	QueueURL string

	// When a message is pulled off of sqs, the visibility timeout will be
	// extended by this much, and periodically extended while the handler
	// performs it's work. If this process crashes, then the sqs message
	// will be redelivered later.
	VisibilityHeartbeat time.Duration

	// Number of worker goroutines to start for receiving messages.
	NumWorkers int

	stopped chan struct{}
	after   func(time.Duration) <-chan time.Time
	sqs     sqsClient
}

func newSQSDispatcher(config client.ConfigProvider) *SQSDispatcher {
	return &SQSDispatcher{
		VisibilityHeartbeat: defaultVisibilityHeartbeat,
		NumWorkers:          defaultNumWorkers,
		after:               time.After,
		sqs:                 sqs.New(config),
	}
}

// Start starts multiple goroutines pulling messages off of the queue.
func (q *SQSDispatcher) Start(handle func(context.Context, *sqs.Message) error) {
	var wg sync.WaitGroup
	for i := 0; i < q.NumWorkers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			q.start(handle)
		}()
	}

	wg.Wait()
}

func (q *SQSDispatcher) Stop() {
	close(q.stopped)
}

// start starts a pulling messages off of the queue and passing them to the
// handler.
func (q *SQSDispatcher) start(handle func(context.Context, *sqs.Message) error) {
	var wg sync.WaitGroup
	for {
		select {
		case <-q.stopped:
			wg.Wait()
			return
		default:
			ctx := q.Context

			resp, err := q.sqs.ReceiveMessage(&sqs.ReceiveMessageInput{
				QueueUrl: aws.String(q.QueueURL),
			})
			if err != nil {
				reporter.Report(ctx, err)
				continue
			}

			for _, m := range resp.Messages {
				wg.Add(1)
				go func(m *sqs.Message) {
					defer wg.Done()
					if err := q.handle(ctx, handle, m); err != nil {
						reporter.Report(ctx, err)
					}
				}(m)
			}
		}
	}
}

func (q *SQSDispatcher) handle(ctx context.Context, handle func(context.Context, *sqs.Message) error, message *sqs.Message) (err error) {
	ctx, cancel := context.WithCancel(ctx)

	defer func() {
		if err == nil {
			_, err = q.sqs.DeleteMessage(&sqs.DeleteMessageInput{
				QueueUrl:      aws.String(q.QueueURL),
				ReceiptHandle: message.ReceiptHandle,
			})
		}
	}()

	var t <-chan time.Time
	t, err = q.extendMessageVisibilityTimeout(message.ReceiptHandle)

	errCh := make(chan error)
	go func() { errCh <- handle(ctx, message) }()

	for {
		select {
		case err = <-errCh:
			return
		case <-q.stopped:
			cancel()
		case <-t:
			t, err = q.extendMessageVisibilityTimeout(message.ReceiptHandle)
			if err != nil {
				return
			}
		}
	}
}

// extendMessageVisibilityTimeout extends the messages visibility timeout by
// VisibilityHeartbeat, and returns a channel that will receive after half of
// VisibilityTimeout has elapsed.
func (q *SQSDispatcher) extendMessageVisibilityTimeout(receiptHandle *string) (<-chan time.Time, error) {
	visibilityTimeout := int64(float64(q.VisibilityHeartbeat) / float64(time.Second))

	_, err := q.sqs.ChangeMessageVisibility(&sqs.ChangeMessageVisibilityInput{
		QueueUrl:          aws.String(q.QueueURL),
		ReceiptHandle:     receiptHandle,
		VisibilityTimeout: aws.Int64(visibilityTimeout),
	})
	if err != nil {
		return nil, fmt.Errorf("error extending message visibility timeout: %v", err)
	}

	return q.after(q.VisibilityHeartbeat / 2), nil
}
<|endoftext|>


<commit_before>package apiproxy

import (
	"net/http"
)

// RevalidateFunc is the signature of functions used to determine whether the
// HTTP resource indicated by req should be revalidated by contacting the
// destination server.
type RevalidateFunc func(req *http.Request) bool

// RevalidationTransport is an implementation of net/http.RoundTripper that
// permits custom behavior with respect to cache entry revalidation for
// resources on the target server.
//
// If the request contains cache validators (an If-None-Match or
// If-Modified-Since header), then Revalidate is called to determine whether the
// cache entry should be revalidated (by being passed to the underlying
// transport). In this way, the Revalidate function can effectively extend or shorten cache
// age limits.
//
// If the request does not contain cache validators, then it is passed to the
// underlying transport.
type RevalidationTransport struct {
	// Revalidate is called on each request in RoundTrip. If it returns true,
	// RoundTrip synthesizes and returns an HTTP 304 Not Modified response.
	// Otherwise, the request is passed through to the underlying transport.
	Revalidate RevalidateFunc

	// Transport is the underlying transport. If nil, net/http.DefaultTransport is used.
	Transport http.RoundTripper
}

// RoundTrip takes a Request and returns a Response.
func (t *RevalidationTransport) RoundTrip(req *http.Request) (resp *http.Response, err error) {
	if t.Revalidate != nil && hasCacheValidator(req.Header) && !t.Revalidate(req) {
		resp = &http.Response{
			Request:          req,
			TransferEncoding: req.TransferEncoding,
			StatusCode:       http.StatusNotModified,
		}
		return
	}

	transport := t.Transport
	if transport == nil {
		transport = http.DefaultTransport
	}

	return transport.RoundTrip(req)
}

// hasCacheValidator returns true if the headers contain cache validators. See
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html#sec13.3 for more
// information.
func hasCacheValidator(headers http.Header) bool {
	return headers.Get("if-none-match") != "" || headers.Get("if-modified-since") != ""
}
<commit_msg>Set empty body on synthetized HTTP response to prevent a panic in net/http<commit_after>package apiproxy

import (
	"bytes"
	"io/ioutil"
	"net/http"
)

// RevalidateFunc is the signature of functions used to determine whether the
// HTTP resource indicated by req should be revalidated by contacting the
// destination server.
type RevalidateFunc func(req *http.Request) bool

// RevalidationTransport is an implementation of net/http.RoundTripper that
// permits custom behavior with respect to cache entry revalidation for
// resources on the target server.
//
// If the request contains cache validators (an If-None-Match or
// If-Modified-Since header), then Revalidate is called to determine whether the
// cache entry should be revalidated (by being passed to the underlying
// transport). In this way, the Revalidate function can effectively extend or shorten cache
// age limits.
//
// If the request does not contain cache validators, then it is passed to the
// underlying transport.
type RevalidationTransport struct {
	// Revalidate is called on each request in RoundTrip. If it returns true,
	// RoundTrip synthesizes and returns an HTTP 304 Not Modified response.
	// Otherwise, the request is passed through to the underlying transport.
	Revalidate RevalidateFunc

	// Transport is the underlying transport. If nil, net/http.DefaultTransport is used.
	Transport http.RoundTripper
}

// RoundTrip takes a Request and returns a Response.
func (t *RevalidationTransport) RoundTrip(req *http.Request) (resp *http.Response, err error) {
	if t.Revalidate != nil && hasCacheValidator(req.Header) && !t.Revalidate(req) {
		resp = &http.Response{
			Request:          req,
			TransferEncoding: req.TransferEncoding,
			StatusCode:       http.StatusNotModified,
			Body:             ioutil.NopCloser(bytes.NewReader([]byte(""))),
		}
		return
	}

	transport := t.Transport
	if transport == nil {
		transport = http.DefaultTransport
	}

	return transport.RoundTrip(req)
}

// hasCacheValidator returns true if the headers contain cache validators. See
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html#sec13.3 for more
// information.
func hasCacheValidator(headers http.Header) bool {
	return headers.Get("if-none-match") != "" || headers.Get("if-modified-since") != ""
}
<|endoftext|>


<commit_before>package irma

import (
	"bytes"
	"context"
	"encoding/hex"
	"encoding/json"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/go-errors/errors"
	"github.com/hashicorp/go-retryablehttp"
	"github.com/privacybydesign/gabi"
	"github.com/privacybydesign/gabi/revocation"
	sseclient "github.com/sietseringers/go-sse"
	"github.com/sirupsen/logrus"
	prefixed "github.com/x-cray/logrus-prefixed-formatter"

	"github.com/privacybydesign/irmago/internal/common"
	"github.com/privacybydesign/irmago/internal/disable_sigpipe"
)

// HTTPTransport sends and receives JSON messages to a HTTP server.
type HTTPTransport struct {
	Server     string
	Binary     bool
	ForceHTTPS bool
	client     *retryablehttp.Client
	headers    http.Header
}

var HTTPHeaders = map[string]http.Header{}

// Logger is used for logging. If not set, init() will initialize it to logrus.StandardLogger().
var Logger *logrus.Logger

var transportlogger *log.Logger

func init() {
	logger := logrus.New()
	logger.SetFormatter(&prefixed.TextFormatter{
		DisableColors:   true,
		FullTimestamp:   true,
		TimestampFormat: "15:04:05.000000",
	})
	SetLogger(logger)
}

func SetLogger(logger *logrus.Logger) {
	Logger = logger
	gabi.Logger = Logger
	common.Logger = Logger
	revocation.Logger = Logger
	sseclient.Logger = log.New(Logger.WithField("type", "sseclient").WriterLevel(logrus.TraceLevel), "", 0)
}

// NewHTTPTransport returns a new HTTPTransport.
func NewHTTPTransport(serverURL string, forceHTTPS bool) *HTTPTransport {
	if Logger.IsLevelEnabled(logrus.TraceLevel) {
		transportlogger = log.New(Logger.WriterLevel(logrus.TraceLevel), "transport: ", 0)
	} else {
		transportlogger = log.New(ioutil.Discard, "", 0)
	}

	if serverURL != "" && !strings.HasSuffix(serverURL, "/") {
		serverURL += "/"
	}

	// Create a transport that dials with a SIGPIPE handler (which is only active on iOS)
	var innerTransport http.Transport

	innerTransport.Dial = func(network, addr string) (c net.Conn, err error) {
		c, err = net.Dial(network, addr)
		if err != nil {
			return c, err
		}
		if err = disable_sigpipe.DisableSigPipe(c); err != nil {
			return c, err
		}
		return c, nil
	}

	client := &retryablehttp.Client{
		Logger:       transportlogger,
		RetryWaitMin: 100 * time.Millisecond,
		RetryWaitMax: 200 * time.Millisecond,
		RetryMax:     2,
		Backoff:      retryablehttp.DefaultBackoff,
		CheckRetry: func(ctx context.Context, resp *http.Response, err error) (bool, error) {
			// Don't retry on 5xx (which retryablehttp does by default)
			return err != nil || resp.StatusCode == 0, err
		},
		HTTPClient: &http.Client{
			Timeout:   time.Second * 3,
			Transport: &innerTransport,
		},
	}

	var host string
	u, err := url.Parse(serverURL)
	if err != nil {
		Logger.Warnf("failed to parse URL %s: %s", serverURL, err.Error())
	} else {
		host = u.Host
	}
	headers := HTTPHeaders[host].Clone()
	if headers == nil {
		headers = http.Header{}
	}
	return &HTTPTransport{
		Server:     serverURL,
		ForceHTTPS: forceHTTPS,
		headers:    headers,
		client:     client,
	}
}

func (transport *HTTPTransport) marshal(o interface{}) ([]byte, error) {
	if transport.Binary {
		return MarshalBinary(o)
	}
	return json.Marshal(o)
}

func (transport *HTTPTransport) unmarshal(data []byte, dst interface{}) error {
	if transport.Binary {
		return UnmarshalBinary(data, dst)
	}
	return json.Unmarshal(data, dst)
}

func (transport *HTTPTransport) unmarshalValidate(data []byte, dst interface{}) error {
	if transport.Binary {
		return UnmarshalValidateBinary(data, dst)
	}
	return UnmarshalValidate(data, dst)
}

func (transport *HTTPTransport) log(prefix string, message interface{}, binary bool) {
	if !Logger.IsLevelEnabled(logrus.TraceLevel) {
		return // do nothing if nothing would be printed anyway
	}
	var str string
	switch s := message.(type) {
	case []byte:
		str = string(s)
	case string:
		str = s
	default:
		tmp, _ := json.Marshal(message)
		str = string(tmp)
		binary = false
	}
	if !binary {
		Logger.Tracef("transport: %s: %s", prefix, str)
	} else {
		Logger.Tracef("transport: %s (hex): %s", prefix, hex.EncodeToString([]byte(str)))
	}
}

// SetHeader sets a header to be sent in requests.
func (transport *HTTPTransport) SetHeader(name, val string) {
	transport.headers.Set(name, val)
}

func (transport *HTTPTransport) request(
	url string, method string, reader io.Reader, contenttype string,
) (response *http.Response, err error) {
	var req retryablehttp.Request
	u := transport.Server + url
	if common.ForceHTTPS && transport.ForceHTTPS && !strings.HasPrefix(u, "https") {
		return nil, &SessionError{ErrorType: ErrorHTTPS, Err: errors.New("remote server does not use https")}
	}
	req.Request, err = http.NewRequest(method, u, reader)
	if err != nil {
		return nil, &SessionError{ErrorType: ErrorTransport, Err: err}
	}
	req.Header = transport.headers.Clone()
	if req.Header.Get("User-agent") == "" {
		req.Header.Set("User-Agent", "irmago")
	}
	if reader != nil && contenttype != "" {
		req.Header.Set("Content-Type", contenttype)
	}
	res, err := transport.client.Do(&req)
	if err != nil {
		return nil, &SessionError{ErrorType: ErrorTransport, Err: err}
	}
	return res, nil
}

func (transport *HTTPTransport) jsonRequest(url string, method string, result interface{}, object interface{}) error {
	if method != http.MethodPost && method != http.MethodGet && method != http.MethodDelete {
		panic("Unsupported HTTP method " + method)
	}
	if method == http.MethodGet && object != nil {
		panic("Cannot GET and also post an object")
	}

	var reader io.Reader
	var contenttype string
	if object != nil {
		switch o := object.(type) {
		case []byte:
			transport.log("body", o, true)
			contenttype = "application/octet-stream"
			reader = bytes.NewBuffer(o)
		case string:
			transport.log("body", o, false)
			contenttype = "text/plain; charset=UTF-8"
			reader = bytes.NewBuffer([]byte(o))
		default:
			marshaled, err := transport.marshal(object)
			if err != nil {
				return &SessionError{ErrorType: ErrorSerialization, Err: err}
			}
			transport.log("body", string(marshaled), transport.Binary)
			if transport.Binary {
				contenttype = "application/octet-stream"
			} else {
				contenttype = "application/json; charset=UTF-8"
			}
			reader = bytes.NewBuffer(marshaled)
		}
	}

	res, err := transport.request(url, method, reader, contenttype)
	if err != nil {
		return err
	}
	if method == http.MethodDelete {
		return nil
	}

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return &SessionError{ErrorType: ErrorServerResponse, Err: err, RemoteStatus: res.StatusCode}
	}
	if res.StatusCode == http.StatusNoContent {
		if result != nil {
			return &SessionError{
				ErrorType:    ErrorServerResponse,
				Err:          errors.New("'204 No Content' received, but result was expected"),
				RemoteStatus: res.StatusCode,
			}
		}
	} else if res.StatusCode != http.StatusOK {
		apierr := &RemoteError{}
		err = transport.unmarshal(body, apierr)
		if err != nil || apierr.ErrorName == "" { // Not an ApiErrorMessage
			return &SessionError{ErrorType: ErrorServerResponse, Err: err, RemoteStatus: res.StatusCode}
		}
		transport.log("error", apierr, false)
		return &SessionError{ErrorType: ErrorApi, RemoteStatus: res.StatusCode, RemoteError: apierr}
	}

	transport.log("response", body, transport.Binary)
	if result == nil { // caller doesn't care about server response
		return nil
	}
	if _, resultstr := result.(*string); resultstr {
		*result.(*string) = string(body)
	} else {
		err = transport.unmarshalValidate(body, result)
		if err != nil {
			return &SessionError{ErrorType: ErrorServerResponse, Err: err, RemoteStatus: res.StatusCode}
		}
	}

	return nil
}

func (transport *HTTPTransport) GetBytes(url string) ([]byte, error) {
	res, err := transport.request(url, http.MethodGet, nil, "")
	if err != nil {
		return nil, &SessionError{ErrorType: ErrorTransport, Err: err}
	}

	if res.StatusCode != 200 {
		return nil, &SessionError{ErrorType: ErrorServerResponse, RemoteStatus: res.StatusCode}
	}
	b, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, &SessionError{ErrorType: ErrorServerResponse, Err: err, RemoteStatus: res.StatusCode}
	}
	return b, nil
}

// Post sends the object to the server and parses its response into result.
func (transport *HTTPTransport) Post(url string, result interface{}, object interface{}) error {
	return transport.jsonRequest(url, http.MethodPost, result, object)
}

// Get performs a GET request and parses the server's response into result.
func (transport *HTTPTransport) Get(url string, result interface{}) error {
	return transport.jsonRequest(url, http.MethodGet, result, nil)
}

// Delete performs a DELETE.
func (transport *HTTPTransport) Delete() {
	_ = transport.jsonRequest("", http.MethodDelete, nil, nil)
}
<commit_msg>Feat: added option to specify which TLS configuration to use for outbound connections<commit_after>package irma

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/hex"
	"encoding/json"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/go-errors/errors"
	"github.com/hashicorp/go-retryablehttp"
	"github.com/privacybydesign/gabi"
	"github.com/privacybydesign/gabi/revocation"
	sseclient "github.com/sietseringers/go-sse"
	"github.com/sirupsen/logrus"
	prefixed "github.com/x-cray/logrus-prefixed-formatter"

	"github.com/privacybydesign/irmago/internal/common"
	"github.com/privacybydesign/irmago/internal/disable_sigpipe"
)

// HTTPTransport sends and receives JSON messages to a HTTP server.
type HTTPTransport struct {
	Server     string
	Binary     bool
	ForceHTTPS bool
	client     *retryablehttp.Client
	headers    http.Header
}

var HTTPHeaders = map[string]http.Header{}

// Logger is used for logging. If not set, init() will initialize it to logrus.StandardLogger().
var Logger *logrus.Logger

var transportlogger *log.Logger

var tlsClientConfig *tls.Config

func init() {
	logger := logrus.New()
	logger.SetFormatter(&prefixed.TextFormatter{
		DisableColors:   true,
		FullTimestamp:   true,
		TimestampFormat: "15:04:05.000000",
	})
	SetLogger(logger)
}

func SetLogger(logger *logrus.Logger) {
	Logger = logger
	gabi.Logger = Logger
	common.Logger = Logger
	revocation.Logger = Logger
	sseclient.Logger = log.New(Logger.WithField("type", "sseclient").WriterLevel(logrus.TraceLevel), "", 0)
}

// SetTLSClientConfig sets the TLS configuration being used for future outbound connections.
// A TLS configuration instance should not be modified after being set.
func SetTLSClientConfig(config *tls.Config) {
	tlsClientConfig = config
}

// NewHTTPTransport returns a new HTTPTransport.
func NewHTTPTransport(serverURL string, forceHTTPS bool) *HTTPTransport {
	if Logger.IsLevelEnabled(logrus.TraceLevel) {
		transportlogger = log.New(Logger.WriterLevel(logrus.TraceLevel), "transport: ", 0)
	} else {
		transportlogger = log.New(ioutil.Discard, "", 0)
	}

	if serverURL != "" && !strings.HasSuffix(serverURL, "/") {
		serverURL += "/"
	}

	// Create a transport that dials with a SIGPIPE handler (which is only active on iOS)
	innerTransport := &http.Transport{
		TLSClientConfig: tlsClientConfig,
		Dial: func(network, addr string) (c net.Conn, err error) {
			c, err = net.Dial(network, addr)
			if err != nil {
				return c, err
			}
			if err = disable_sigpipe.DisableSigPipe(c); err != nil {
				return c, err
			}
			return c, nil
		},
	}

	client := &retryablehttp.Client{
		Logger:       transportlogger,
		RetryWaitMin: 100 * time.Millisecond,
		RetryWaitMax: 200 * time.Millisecond,
		RetryMax:     2,
		Backoff:      retryablehttp.DefaultBackoff,
		CheckRetry: func(ctx context.Context, resp *http.Response, err error) (bool, error) {
			// Don't retry on 5xx (which retryablehttp does by default)
			return err != nil || resp.StatusCode == 0, err
		},
		HTTPClient: &http.Client{
			Timeout:   time.Second * 3,
			Transport: innerTransport,
		},
	}

	var host string
	u, err := url.Parse(serverURL)
	if err != nil {
		Logger.Warnf("failed to parse URL %s: %s", serverURL, err.Error())
	} else {
		host = u.Host
	}
	headers := HTTPHeaders[host].Clone()
	if headers == nil {
		headers = http.Header{}
	}
	return &HTTPTransport{
		Server:     serverURL,
		ForceHTTPS: forceHTTPS,
		headers:    headers,
		client:     client,
	}
}

func (transport *HTTPTransport) marshal(o interface{}) ([]byte, error) {
	if transport.Binary {
		return MarshalBinary(o)
	}
	return json.Marshal(o)
}

func (transport *HTTPTransport) unmarshal(data []byte, dst interface{}) error {
	if transport.Binary {
		return UnmarshalBinary(data, dst)
	}
	return json.Unmarshal(data, dst)
}

func (transport *HTTPTransport) unmarshalValidate(data []byte, dst interface{}) error {
	if transport.Binary {
		return UnmarshalValidateBinary(data, dst)
	}
	return UnmarshalValidate(data, dst)
}

func (transport *HTTPTransport) log(prefix string, message interface{}, binary bool) {
	if !Logger.IsLevelEnabled(logrus.TraceLevel) {
		return // do nothing if nothing would be printed anyway
	}
	var str string
	switch s := message.(type) {
	case []byte:
		str = string(s)
	case string:
		str = s
	default:
		tmp, _ := json.Marshal(message)
		str = string(tmp)
		binary = false
	}
	if !binary {
		Logger.Tracef("transport: %s: %s", prefix, str)
	} else {
		Logger.Tracef("transport: %s (hex): %s", prefix, hex.EncodeToString([]byte(str)))
	}
}

// SetHeader sets a header to be sent in requests.
func (transport *HTTPTransport) SetHeader(name, val string) {
	transport.headers.Set(name, val)
}

func (transport *HTTPTransport) request(
	url string, method string, reader io.Reader, contenttype string,
) (response *http.Response, err error) {
	var req retryablehttp.Request
	u := transport.Server + url
	if common.ForceHTTPS && transport.ForceHTTPS && !strings.HasPrefix(u, "https") {
		return nil, &SessionError{ErrorType: ErrorHTTPS, Err: errors.New("remote server does not use https")}
	}
	req.Request, err = http.NewRequest(method, u, reader)
	if err != nil {
		return nil, &SessionError{ErrorType: ErrorTransport, Err: err}
	}
	req.Header = transport.headers.Clone()
	if req.Header.Get("User-agent") == "" {
		req.Header.Set("User-Agent", "irmago")
	}
	if reader != nil && contenttype != "" {
		req.Header.Set("Content-Type", contenttype)
	}
	res, err := transport.client.Do(&req)
	if err != nil {
		return nil, &SessionError{ErrorType: ErrorTransport, Err: err}
	}
	return res, nil
}

func (transport *HTTPTransport) jsonRequest(url string, method string, result interface{}, object interface{}) error {
	if method != http.MethodPost && method != http.MethodGet && method != http.MethodDelete {
		panic("Unsupported HTTP method " + method)
	}
	if method == http.MethodGet && object != nil {
		panic("Cannot GET and also post an object")
	}

	var reader io.Reader
	var contenttype string
	if object != nil {
		switch o := object.(type) {
		case []byte:
			transport.log("body", o, true)
			contenttype = "application/octet-stream"
			reader = bytes.NewBuffer(o)
		case string:
			transport.log("body", o, false)
			contenttype = "text/plain; charset=UTF-8"
			reader = bytes.NewBuffer([]byte(o))
		default:
			marshaled, err := transport.marshal(object)
			if err != nil {
				return &SessionError{ErrorType: ErrorSerialization, Err: err}
			}
			transport.log("body", string(marshaled), transport.Binary)
			if transport.Binary {
				contenttype = "application/octet-stream"
			} else {
				contenttype = "application/json; charset=UTF-8"
			}
			reader = bytes.NewBuffer(marshaled)
		}
	}

	res, err := transport.request(url, method, reader, contenttype)
	if err != nil {
		return err
	}
	if method == http.MethodDelete {
		return nil
	}

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return &SessionError{ErrorType: ErrorServerResponse, Err: err, RemoteStatus: res.StatusCode}
	}
	if res.StatusCode == http.StatusNoContent {
		if result != nil {
			return &SessionError{
				ErrorType:    ErrorServerResponse,
				Err:          errors.New("'204 No Content' received, but result was expected"),
				RemoteStatus: res.StatusCode,
			}
		}
	} else if res.StatusCode != http.StatusOK {
		apierr := &RemoteError{}
		err = transport.unmarshal(body, apierr)
		if err != nil || apierr.ErrorName == "" { // Not an ApiErrorMessage
			return &SessionError{ErrorType: ErrorServerResponse, Err: err, RemoteStatus: res.StatusCode}
		}
		transport.log("error", apierr, false)
		return &SessionError{ErrorType: ErrorApi, RemoteStatus: res.StatusCode, RemoteError: apierr}
	}

	transport.log("response", body, transport.Binary)
	if result == nil { // caller doesn't care about server response
		return nil
	}
	if _, resultstr := result.(*string); resultstr {
		*result.(*string) = string(body)
	} else {
		err = transport.unmarshalValidate(body, result)
		if err != nil {
			return &SessionError{ErrorType: ErrorServerResponse, Err: err, RemoteStatus: res.StatusCode}
		}
	}

	return nil
}

func (transport *HTTPTransport) GetBytes(url string) ([]byte, error) {
	res, err := transport.request(url, http.MethodGet, nil, "")
	if err != nil {
		return nil, &SessionError{ErrorType: ErrorTransport, Err: err}
	}

	if res.StatusCode != 200 {
		return nil, &SessionError{ErrorType: ErrorServerResponse, RemoteStatus: res.StatusCode}
	}
	b, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, &SessionError{ErrorType: ErrorServerResponse, Err: err, RemoteStatus: res.StatusCode}
	}
	return b, nil
}

// Post sends the object to the server and parses its response into result.
func (transport *HTTPTransport) Post(url string, result interface{}, object interface{}) error {
	return transport.jsonRequest(url, http.MethodPost, result, object)
}

// Get performs a GET request and parses the server's response into result.
func (transport *HTTPTransport) Get(url string, result interface{}) error {
	return transport.jsonRequest(url, http.MethodGet, result, nil)
}

// Delete performs a DELETE.
func (transport *HTTPTransport) Delete() {
	_ = transport.jsonRequest("", http.MethodDelete, nil, nil)
}
<|endoftext|>


<commit_before>package momos

import (
	"bytes"
	"errors"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/nats-io/nuid"

	"github.com/PuerkitoBio/goquery"
	clientCache "github.com/gregjones/httpcache"
)

var (
	ErrRequest            = errors.New("Request error")
	ErrTimeout            = errors.New("Timeout error")
	ErrInvalidStatusCode  = errors.New("Invalid status code")
	ErrInvalidContentType = errors.New("Invalid content type")
)

var cache = clientCache.NewMemoryCacheTransport()

type ssiResult struct {
	name    string
	payload []byte
	error   error
}

type proxyTransport struct {
	http.RoundTripper
}

func (t *proxyTransport) RoundTrip(req *http.Request) (resp *http.Response, err error) {

	Log.Debugf("start processing request %q", req.URL)
	timeStart := time.Now()

	resp, err = t.RoundTripper.RoundTrip(req)
	if err != nil {
		Log.Errorf("could not create RoundTripper from %q", req.URL)
		return nil, err
	}

	// Only html files are scanned
	contentType := resp.Header.Get("Content-Type")
	if !strings.HasPrefix(contentType, "text/html") {
		return resp, nil
	}

	doc, err := goquery.NewDocumentFromReader(resp.Body)

	if err != nil {
		Log.Errorf("illegal response body from %q", req.URL)
		return nil, err
	}

	ssiCount := 0
	ssiElements := map[string]SSIElement{}

	doc.Find("ssi").Each(func(i int, element *goquery.Selection) {
		se := SSIElement{Element: element}

		se.SetTimeout(element.AttrOr("timeout", "2000"))
		se.SetSrc(element.AttrOr("src", ""))
		se.SetName(element.AttrOr("name", nuid.Next()))
		se.SetTemplate(element.AttrOr("template", "false"))
		se.SetFilterIncludes(element.AttrOr("no-scripts", "true"))

		se.GetErrorTag()
		se.GetTimeoutTag()

		se.templateContext = TemplateContext{
			DateLocal: time.Now().Local().Format("2006-01-02"),
			Date:      time.Now().Format(time.RFC3339),
			RequestId: req.Header.Get("X-Request-Id"),
			Name:      se.name,
		}

		ssiElements[se.name] = se
		ssiCount++
	})

	ch := make(chan ssiResult)

	timeStartRequest := time.Now()

	for _, el := range ssiElements {
		go makeRequest(el.name, el.src, ch, el.timeout)
	}

	for i := 0; i < ssiCount; i++ {
		res := <-ch
		el := ssiElements[res.name]
		if res.error == nil {
			Log.Debugf("fragment (%v) - Request to %v took %v", el.name, el.src, time.Since(timeStartRequest))
			el.SetupSuccess(res.payload)
		} else {
			el.SetupFallback(res.error)
			Log.Debugf("Fragment (%v) - Request to %v error: %q", el.name, el.src, res.error)
		}
	}

	close(ch)

	htmlDoc, err := doc.Html()

	if err != nil {
		Log.Errorf("Could not get html from document %q", req.URL)
		return nil, err
	}

	// assign new reader with content
	content := []byte(htmlDoc)
	body := ioutil.NopCloser(bytes.NewReader(content))
	resp.Body = body
	resp.ContentLength = int64(len(content)) // update content length
	resp.Header.Set("Content-Length", strconv.Itoa(len(content)))

	Log.Debugf("Processing complete %q took %q", req.URL, time.Since(timeStart))

	return resp, nil
}

func makeRequest(name string, url string, ch chan<- ssiResult, timeout time.Duration) {
	// @TODO don't create a new client per request
	var Client = &http.Client{
		Transport: cache,
		Timeout:   timeout,
	}

	resp, err := Client.Get(url)

	if err != nil {
		ch <- ssiResult{name: name, error: err}
	} else {
		contentType := resp.Header.Get("Content-Type")
		if !strings.HasPrefix(contentType, "text/html") {
			ch <- ssiResult{name: name, error: ErrInvalidContentType}
		} else if resp.StatusCode > 199 && resp.StatusCode < 300 {
			body, _ := ioutil.ReadAll(resp.Body)
			defer resp.Body.Close()

			// https://github.com/gregjones/httpcache
			if resp.Header.Get("X-From-Cache") == "1" {
				Log.Debugf("Fragment (%v) - Response was cached", name)
			} else {
				Log.Debugf("Fragment (%v) - Response was refreshed", name)
			}

			ch <- ssiResult{name: name, payload: body}
		} else {
			ch <- ssiResult{name: name, error: ErrInvalidStatusCode}
		}
	}
}
<commit_msg>log as error<commit_after>package momos

import (
	"bytes"
	"errors"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/nats-io/nuid"

	"github.com/PuerkitoBio/goquery"
	clientCache "github.com/gregjones/httpcache"
)

var (
	ErrRequest            = errors.New("Request error")
	ErrTimeout            = errors.New("Timeout error")
	ErrInvalidStatusCode  = errors.New("Invalid status code")
	ErrInvalidContentType = errors.New("Invalid content type")
)

var cache = clientCache.NewMemoryCacheTransport()

type ssiResult struct {
	name    string
	payload []byte
	error   error
}

type proxyTransport struct {
	http.RoundTripper
}

func (t *proxyTransport) RoundTrip(req *http.Request) (resp *http.Response, err error) {

	Log.Debugf("start processing request %q", req.URL)
	timeStart := time.Now()

	resp, err = t.RoundTripper.RoundTrip(req)
	if err != nil {
		Log.Errorf("could not create RoundTripper from %q", req.URL)
		return nil, err
	}

	// Only html files are scanned
	contentType := resp.Header.Get("Content-Type")
	if !strings.HasPrefix(contentType, "text/html") {
		return resp, nil
	}

	doc, err := goquery.NewDocumentFromReader(resp.Body)

	if err != nil {
		Log.Errorf("illegal response body from %q", req.URL)
		return nil, err
	}

	ssiCount := 0
	ssiElements := map[string]SSIElement{}

	doc.Find("ssi").Each(func(i int, element *goquery.Selection) {
		se := SSIElement{Element: element}

		se.SetTimeout(element.AttrOr("timeout", "2000"))
		se.SetSrc(element.AttrOr("src", ""))
		se.SetName(element.AttrOr("name", nuid.Next()))
		se.SetTemplate(element.AttrOr("template", "false"))
		se.SetFilterIncludes(element.AttrOr("no-scripts", "true"))

		se.GetErrorTag()
		se.GetTimeoutTag()

		se.templateContext = TemplateContext{
			DateLocal: time.Now().Local().Format("2006-01-02"),
			Date:      time.Now().Format(time.RFC3339),
			RequestId: req.Header.Get("X-Request-Id"),
			Name:      se.name,
		}

		ssiElements[se.name] = se
		ssiCount++
	})

	ch := make(chan ssiResult)

	timeStartRequest := time.Now()

	for _, el := range ssiElements {
		go makeRequest(el.name, el.src, ch, el.timeout)
	}

	for i := 0; i < ssiCount; i++ {
		res := <-ch
		el := ssiElements[res.name]
		if res.error == nil {
			Log.Debugf("fragment (%v) - Request to %v took %v", el.name, el.src, time.Since(timeStartRequest))
			el.SetupSuccess(res.payload)
		} else {
			el.SetupFallback(res.error)
			Log.Errorf("Fragment (%v) - Request to %v \nerror: %q", el.name, el.src, res.error)
		}
	}

	close(ch)

	htmlDoc, err := doc.Html()

	if err != nil {
		Log.Errorf("Could not get html from document %q", req.URL)
		return nil, err
	}

	// assign new reader with content
	content := []byte(htmlDoc)
	body := ioutil.NopCloser(bytes.NewReader(content))
	resp.Body = body
	resp.ContentLength = int64(len(content)) // update content length
	resp.Header.Set("Content-Length", strconv.Itoa(len(content)))

	Log.Debugf("Processing complete %q took %q", req.URL, time.Since(timeStart))

	return resp, nil
}

func makeRequest(name string, url string, ch chan<- ssiResult, timeout time.Duration) {
	// @TODO don't create a new client per request
	var Client = &http.Client{
		Transport: cache,
		Timeout:   timeout,
	}

	resp, err := Client.Get(url)

	if err != nil {
		ch <- ssiResult{name: name, error: err}
	} else {
		contentType := resp.Header.Get("Content-Type")
		if !strings.HasPrefix(contentType, "text/html") {
			ch <- ssiResult{name: name, error: ErrInvalidContentType}
		} else if resp.StatusCode > 199 && resp.StatusCode < 300 {
			body, _ := ioutil.ReadAll(resp.Body)
			defer resp.Body.Close()

			// https://github.com/gregjones/httpcache
			if resp.Header.Get("X-From-Cache") == "1" {
				Log.Debugf("Fragment (%v) - Response was cached", name)
			} else {
				Log.Debugf("Fragment (%v) - Response was refreshed", name)
			}

			ch <- ssiResult{name: name, payload: body}
		} else {
			ch <- ssiResult{name: name, error: ErrInvalidStatusCode}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	gitmedia ".."
	"fmt"
	"os"
	"path/filepath"
)

var (
	allCommands map[string]map[string]string
	gitMediaBin string
)

func main() {
	for wd, commands := range allCommands {
		fmt.Println("Integration tests for", wd)
		for cmd, expected := range commands {
			if err := os.Chdir(wd); err != nil {
				fmt.Println("Cannot chdir to", wd)
				os.Exit(1)
			}
			fmt.Println("$ git-media", cmd)
			actual := gitmedia.SimpleExec(gitMediaBin, cmd)
			if actual != expected {
				fmt.Printf("expected:\n%s\n\n", expected)
				fmt.Printf("actual:\n%s\n", actual)
			}
		}
	}
}

func init() {
	wd, err := os.Getwd()
	if err != nil {
		panic(err)
	}

	gitMediaBin = filepath.Join(wd, "bin", "git-media")

	allCommands = make(map[string]map[string]string)

	// tests on the git-media repository, which has no actual git-media assets :)
	allCommands[wd] = map[string]string{
		"version": "git-media v" + gitmedia.Version,
		"config": "Endpoint=\n" +
			"LocalMediaDir=" + filepath.Join(wd, ".git", "media") + "\n" +
			"TempDir=" + filepath.Join(os.TempDir(), "git-media"),
	}
}
<commit_msg>ンンー ンンンン ンーンン<commit_after>package main

import (
	gitmedia ".."
	"fmt"
	"os"
	"path/filepath"
)

var (
	allCommands map[string]map[string]string
	gitMediaBin string
)

func main() {
	for wd, commands := range allCommands {
		fmt.Println("Integration tests for", wd)
		for cmd, expected := range commands {
			if err := os.Chdir(wd); err != nil {
				fmt.Println("Cannot chdir to", wd)
				os.Exit(1)
			}
			fmt.Println("$ git-media", cmd)
			actual := gitmedia.SimpleExec(gitMediaBin, cmd)
			if actual != expected {
				fmt.Printf("- expected\n%s\n\n", expected)
				fmt.Printf("- actual\n%s\n", actual)
			}
		}
		fmt.Println("")
	}
}

func init() {
	wd, err := os.Getwd()
	if err != nil {
		panic(err)
	}

	gitMediaBin = filepath.Join(wd, "bin", "git-media")

	allCommands = make(map[string]map[string]string)

	// tests on the git-media repository, which has no actual git-media assets :)
	allCommands[wd] = map[string]string{
		"version": "git-media v" + gitmedia.Version,
		"config": "Endpoint=\n" +
			"LocalMediaDir=" + filepath.Join(wd, ".git", "media") + "\n" +
			"TempDir=" + filepath.Join(os.TempDir(), "git-media"),
	}

	// tests on the git-media .git dir
	allCommands[filepath.Join(wd, ".git")] = allCommands[wd]
}
<|endoftext|>


<commit_before>package internal

import (
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/kms"
	"github.com/stretchr/testify/assert"
)

func TestKms(t *testing.T) {
	mockKMS := &MockKMSAPI{}
	defer mockKMS.AssertExpectations(t)
	kmsClients["myregion:myprofile"] = mockKMS
	kmsCrypter := KMSCrypter{}

	mockKMS.On("Encrypt",
		&kms.EncryptInput{
			KeyId:     aws.String("mykey"),
			Plaintext: []byte("mypass"),
		},
	).Return(
		&kms.EncryptOutput{
			CiphertextBlob: []byte("myciphertextblob"),
		},
		nil,
	)
	secret, myDecryptParams, err := kmsCrypter.Encrypt("mypass", map[string]string{
		"region": "myregion",
		"profile": "myprofile",
		"keyID":  "mykey",
	})
	assert.Equal(t, myDecryptParams, DecryptParams{
		"region": "myregion",
	})
	assert.Nil(t, err)

	mockKMS.On("Decrypt",
		&kms.DecryptInput{
			CiphertextBlob: []byte("myciphertextblob"),
		},
	).Return(
		&kms.DecryptOutput{
			Plaintext: []byte("mypass"),
		},
		nil,
	)
	myDecryptParams["profile"] = "myprofile"

	plaintext, err := kmsCrypter.Decrypt(secret, myDecryptParams)
	assert.Nil(t, err)
	assert.Equal(t, "mypass", plaintext)
}

func TestKmsDefaultProfile(t *testing.T) {
	mockKMS := &MockKMSAPI{}
	defer mockKMS.AssertExpectations(t)
	kmsClients["myregion:default"] = mockKMS
	kmsCrypter := KMSCrypter{}

	mockKMS.On("Encrypt",
		&kms.EncryptInput{
			KeyId:     aws.String("mykey"),
			Plaintext: []byte("mypass"),
		},
	).Return(
		&kms.EncryptOutput{
			CiphertextBlob: []byte("myciphertextblob"),
		},
		nil,
	)
	secret, myDecryptParams, err := kmsCrypter.Encrypt("mypass", map[string]string{
		"region": "myregion",
		"keyID":  "mykey",
	})
	assert.Equal(t, myDecryptParams, DecryptParams{
		"region": "myregion",
	})
	assert.Nil(t, err)

	mockKMS.On("Decrypt",
		&kms.DecryptInput{
			CiphertextBlob: []byte("myciphertextblob"),
		},
	).Return(
		&kms.DecryptOutput{
			Plaintext: []byte("mypass"),
		},
		nil,
	)

	plaintext, err := kmsCrypter.Decrypt(secret, myDecryptParams)
	assert.Nil(t, err)
	assert.Equal(t, "mypass", plaintext)
}
<commit_msg>go fmt (#11)<commit_after>package internal

import (
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/kms"
	"github.com/stretchr/testify/assert"
)

func TestKms(t *testing.T) {
	mockKMS := &MockKMSAPI{}
	defer mockKMS.AssertExpectations(t)
	kmsClients["myregion:myprofile"] = mockKMS
	kmsCrypter := KMSCrypter{}

	mockKMS.On("Encrypt",
		&kms.EncryptInput{
			KeyId:     aws.String("mykey"),
			Plaintext: []byte("mypass"),
		},
	).Return(
		&kms.EncryptOutput{
			CiphertextBlob: []byte("myciphertextblob"),
		},
		nil,
	)
	secret, myDecryptParams, err := kmsCrypter.Encrypt("mypass", map[string]string{
		"region":  "myregion",
		"profile": "myprofile",
		"keyID":   "mykey",
	})
	assert.Equal(t, myDecryptParams, DecryptParams{
		"region": "myregion",
	})
	assert.Nil(t, err)

	mockKMS.On("Decrypt",
		&kms.DecryptInput{
			CiphertextBlob: []byte("myciphertextblob"),
		},
	).Return(
		&kms.DecryptOutput{
			Plaintext: []byte("mypass"),
		},
		nil,
	)
	myDecryptParams["profile"] = "myprofile"

	plaintext, err := kmsCrypter.Decrypt(secret, myDecryptParams)
	assert.Nil(t, err)
	assert.Equal(t, "mypass", plaintext)
}

func TestKmsDefaultProfile(t *testing.T) {
	mockKMS := &MockKMSAPI{}
	defer mockKMS.AssertExpectations(t)
	kmsClients["myregion:default"] = mockKMS
	kmsCrypter := KMSCrypter{}

	mockKMS.On("Encrypt",
		&kms.EncryptInput{
			KeyId:     aws.String("mykey"),
			Plaintext: []byte("mypass"),
		},
	).Return(
		&kms.EncryptOutput{
			CiphertextBlob: []byte("myciphertextblob"),
		},
		nil,
	)
	secret, myDecryptParams, err := kmsCrypter.Encrypt("mypass", map[string]string{
		"region": "myregion",
		"keyID":  "mykey",
	})
	assert.Equal(t, myDecryptParams, DecryptParams{
		"region": "myregion",
	})
	assert.Nil(t, err)

	mockKMS.On("Decrypt",
		&kms.DecryptInput{
			CiphertextBlob: []byte("myciphertextblob"),
		},
	).Return(
		&kms.DecryptOutput{
			Plaintext: []byte("mypass"),
		},
		nil,
	)

	plaintext, err := kmsCrypter.Decrypt(secret, myDecryptParams)
	assert.Nil(t, err)
	assert.Equal(t, "mypass", plaintext)
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Jetstack cert-manager contributors.

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 certificates

import (
	"context"
	"crypto"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"math/big"
	"reflect"
	"time"

	"github.com/kr/pretty"
	k8sErrors "k8s.io/apimachinery/pkg/api/errors"
	corelisters "k8s.io/client-go/listers/core/v1"
	"k8s.io/client-go/tools/cache"

	apiutil "github.com/jetstack/cert-manager/pkg/api/util"
	"github.com/jetstack/cert-manager/pkg/apis/certmanager/v1alpha2"
	cmclient "github.com/jetstack/cert-manager/pkg/client/clientset/versioned"
	cmlisters "github.com/jetstack/cert-manager/pkg/client/listers/certmanager/v1alpha2"
	controllerpkg "github.com/jetstack/cert-manager/pkg/controller"
	logf "github.com/jetstack/cert-manager/pkg/logs"
	"github.com/jetstack/cert-manager/pkg/metrics"
	"github.com/jetstack/cert-manager/pkg/util"
	"github.com/jetstack/cert-manager/pkg/util/errors"
	"github.com/jetstack/cert-manager/pkg/util/kube"
	"github.com/jetstack/cert-manager/pkg/util/pki"
)

var (
	certificateGvk = v1alpha2.SchemeGroupVersion.WithKind("Certificate")
)

type calculateDurationUntilRenewFn func(context.Context, *x509.Certificate, *v1alpha2.Certificate) time.Duration

func getCertificateForKey(ctx context.Context, key string, lister cmlisters.CertificateLister) (*v1alpha2.Certificate, error) {
	namespace, name, err := cache.SplitMetaNamespaceKey(key)
	if err != nil {
		return nil, nil
	}

	crt, err := lister.Certificates(namespace).Get(name)
	if k8sErrors.IsNotFound(err) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}

	return crt, nil
}

func certificateGetter(lister cmlisters.CertificateLister) func(namespace, name string) (interface{}, error) {
	return func(namespace, name string) (interface{}, error) {
		return lister.Certificates(namespace).Get(name)
	}
}

var keyFunc = controllerpkg.KeyFunc

func certificateMatchesSpec(crt *v1alpha2.Certificate, key crypto.Signer, cert *x509.Certificate, secretLister corelisters.SecretLister) (bool, []string) {
	var errs []string

	// TODO: add checks for KeySize, KeyAlgorithm fields
	// TODO: add checks for Organization field
	// TODO: add checks for IsCA field

	// check if the private key is the corresponding pair to the certificate

	matches, err := pki.PublicKeyMatchesCertificate(key.Public(), cert)
	if err != nil {
		errs = append(errs, err.Error())
	} else if !matches {
		errs = append(errs, fmt.Sprintf("Certificate private key does not match certificate"))
	}

	// If CN is set on the resource then it should exist on the certificate as
	// the Common Name or a DNS Name
	expectedCN := crt.Spec.CommonName
	gotCN := append(cert.DNSNames, cert.Subject.CommonName)
	if !util.Contains(gotCN, expectedCN) {
		errs = append(errs, fmt.Sprintf("Common Name on TLS certificate not up to date (%q): %s",
			expectedCN, gotCN))
	}

	// validate the dns names are correct
	expectedDNSNames := crt.Spec.DNSNames
	if !util.Subset(cert.DNSNames, expectedDNSNames) {
		errs = append(errs, fmt.Sprintf("DNS names on TLS certificate not up to date: %q", cert.DNSNames))
	}

	expectedURIs := crt.Spec.URISANs
	if !util.EqualUnsorted(pki.URLsToString(cert.URIs), expectedURIs) {
		errs = append(errs, fmt.Sprintf("URI SANs on TLS certificate not up to date: %q", cert.URIs))
	}

	// validate the ip addresses are correct
	if !util.EqualUnsorted(pki.IPAddressesToString(cert.IPAddresses), crt.Spec.IPAddresses) {
		errs = append(errs, fmt.Sprintf("IP addresses on TLS certificate not up to date: %q", pki.IPAddressesToString(cert.IPAddresses)))
	}

	// get a copy of the current secret resource
	// Note that we already know that it exists, no need to check for errors
	// TODO: Refactor so that the secret is passed as argument?
	secret, err := secretLister.Secrets(crt.Namespace).Get(crt.Spec.SecretName)

	// validate that the issuer is correct
	if crt.Spec.IssuerRef.Name != secret.Annotations[v1alpha2.IssuerNameAnnotationKey] {
		errs = append(errs, fmt.Sprintf("Issuer of the certificate is not up to date: %q", secret.Annotations[v1alpha2.IssuerNameAnnotationKey]))
	}

	// validate that the issuer kind is correct
	if apiutil.IssuerKind(crt.Spec.IssuerRef) != secret.Annotations[v1alpha2.IssuerKindAnnotationKey] {
		errs = append(errs, fmt.Sprintf("Issuer kind of the certificate is not up to date: %q", secret.Annotations[v1alpha2.IssuerKindAnnotationKey]))
	}

	return len(errs) == 0, errs
}

func scheduleRenewal(ctx context.Context, lister corelisters.SecretLister, calc calculateDurationUntilRenewFn, queueFn func(interface{}, time.Duration), crt *v1alpha2.Certificate) {
	log := logf.FromContext(ctx)
	log = log.WithValues(
		logf.RelatedResourceNameKey, crt.Spec.SecretName,
		logf.RelatedResourceNamespaceKey, crt.Namespace,
		logf.RelatedResourceKindKey, "Secret",
	)

	key, err := keyFunc(crt)
	if err != nil {
		log.Error(err, "error getting key for certificate resource")
		return
	}

	cert, err := kube.SecretTLSCert(ctx, lister, crt.Namespace, crt.Spec.SecretName)
	if err != nil {
		if !errors.IsInvalidData(err) {
			log.Error(err, "error getting secret for certificate resource")
		}
		return
	}

	renewIn := calc(ctx, cert, crt)
	queueFn(key, renewIn)

	log.WithValues("duration_until_renewal", renewIn.String()).Info("certificate scheduled for renewal")
}

// staticTemporarySerialNumber is a fixed serial number we check for when
// updating the status of a certificate.
// It is used to identify temporarily generated certificates, so that friendly
// status messages can be displayed to users.
const staticTemporarySerialNumber = 0x1234567890

func isTemporaryCertificate(cert *x509.Certificate) bool {
	if cert == nil {
		return false
	}
	return cert.SerialNumber.Int64() == staticTemporarySerialNumber
}

// generateLocallySignedTemporaryCertificate signs a temporary certificate for
// the given certificate resource using a one-use temporary CA that is then
// discarded afterwards.
// This is to mitigate a potential attack against x509 certificates that use a
// predictable serial number and weak MD5 hashing algorithms.
// In practice, this shouldn't really be a concern anyway.
func generateLocallySignedTemporaryCertificate(crt *v1alpha2.Certificate, pk []byte) ([]byte, error) {
	// generate a throwaway self-signed root CA
	caPk, err := pki.GenerateECPrivateKey(pki.ECCurve521)
	if err != nil {
		return nil, err
	}
	caCertTemplate, err := pki.GenerateTemplate(&v1alpha2.Certificate{
		Spec: v1alpha2.CertificateSpec{
			CommonName: "cert-manager.local",
			IsCA:       true,
		},
	})
	if err != nil {
		return nil, err
	}
	_, caCert, err := pki.SignCertificate(caCertTemplate, caCertTemplate, caPk.Public(), caPk)
	if err != nil {
		return nil, err
	}

	// sign a temporary certificate using the root CA
	template, err := pki.GenerateTemplate(crt)
	if err != nil {
		return nil, err
	}
	template.SerialNumber = big.NewInt(staticTemporarySerialNumber)

	signeeKey, err := pki.DecodePrivateKeyBytes(pk)
	if err != nil {
		return nil, err
	}

	b, _, err := pki.SignCertificate(template, caCert, signeeKey.Public(), caPk)
	if err != nil {
		return nil, err
	}

	return b, nil
}

func updateCertificateStatus(ctx context.Context, m *metrics.Metrics, cmClient cmclient.Interface, old, new *v1alpha2.Certificate) (*v1alpha2.Certificate, error) {
	defer m.UpdateCertificateStatus(new)

	log := logf.FromContext(ctx, "updateStatus")
	oldBytes, _ := json.Marshal(old.Status)
	newBytes, _ := json.Marshal(new.Status)
	if reflect.DeepEqual(oldBytes, newBytes) {
		return nil, nil
	}
	log.V(logf.DebugLevel).Info("updating resource due to change in status", "diff", pretty.Diff(string(oldBytes), string(newBytes)))
	return cmClient.CertmanagerV1alpha2().Certificates(new.Namespace).UpdateStatus(new)
}

func certificateHasTemporaryCertificateAnnotation(crt *v1alpha2.Certificate) bool {
	if crt.Annotations == nil {
		return false
	}

	if val, ok := crt.Annotations[v1alpha2.IssueTemporaryCertificateAnnotation]; ok && val == "true" {
		return true
	}

	return false
}
<commit_msg>Make annotation of secret if nil to prevent panic<commit_after>/*
Copyright 2019 The Jetstack cert-manager contributors.

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 certificates

import (
	"context"
	"crypto"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"math/big"
	"reflect"
	"time"

	"github.com/kr/pretty"
	k8sErrors "k8s.io/apimachinery/pkg/api/errors"
	corelisters "k8s.io/client-go/listers/core/v1"
	"k8s.io/client-go/tools/cache"

	apiutil "github.com/jetstack/cert-manager/pkg/api/util"
	"github.com/jetstack/cert-manager/pkg/apis/certmanager/v1alpha2"
	cmclient "github.com/jetstack/cert-manager/pkg/client/clientset/versioned"
	cmlisters "github.com/jetstack/cert-manager/pkg/client/listers/certmanager/v1alpha2"
	controllerpkg "github.com/jetstack/cert-manager/pkg/controller"
	logf "github.com/jetstack/cert-manager/pkg/logs"
	"github.com/jetstack/cert-manager/pkg/metrics"
	"github.com/jetstack/cert-manager/pkg/util"
	"github.com/jetstack/cert-manager/pkg/util/errors"
	"github.com/jetstack/cert-manager/pkg/util/kube"
	"github.com/jetstack/cert-manager/pkg/util/pki"
)

var (
	certificateGvk = v1alpha2.SchemeGroupVersion.WithKind("Certificate")
)

type calculateDurationUntilRenewFn func(context.Context, *x509.Certificate, *v1alpha2.Certificate) time.Duration

func getCertificateForKey(ctx context.Context, key string, lister cmlisters.CertificateLister) (*v1alpha2.Certificate, error) {
	namespace, name, err := cache.SplitMetaNamespaceKey(key)
	if err != nil {
		return nil, nil
	}

	crt, err := lister.Certificates(namespace).Get(name)
	if k8sErrors.IsNotFound(err) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}

	return crt, nil
}

func certificateGetter(lister cmlisters.CertificateLister) func(namespace, name string) (interface{}, error) {
	return func(namespace, name string) (interface{}, error) {
		return lister.Certificates(namespace).Get(name)
	}
}

var keyFunc = controllerpkg.KeyFunc

func certificateMatchesSpec(crt *v1alpha2.Certificate, key crypto.Signer, cert *x509.Certificate, secretLister corelisters.SecretLister) (bool, []string) {
	var errs []string

	// TODO: add checks for KeySize, KeyAlgorithm fields
	// TODO: add checks for Organization field
	// TODO: add checks for IsCA field

	// check if the private key is the corresponding pair to the certificate

	matches, err := pki.PublicKeyMatchesCertificate(key.Public(), cert)
	if err != nil {
		errs = append(errs, err.Error())
	} else if !matches {
		errs = append(errs, fmt.Sprintf("Certificate private key does not match certificate"))
	}

	// If CN is set on the resource then it should exist on the certificate as
	// the Common Name or a DNS Name
	expectedCN := crt.Spec.CommonName
	gotCN := append(cert.DNSNames, cert.Subject.CommonName)
	if !util.Contains(gotCN, expectedCN) {
		errs = append(errs, fmt.Sprintf("Common Name on TLS certificate not up to date (%q): %s",
			expectedCN, gotCN))
	}

	// validate the dns names are correct
	expectedDNSNames := crt.Spec.DNSNames
	if !util.Subset(cert.DNSNames, expectedDNSNames) {
		errs = append(errs, fmt.Sprintf("DNS names on TLS certificate not up to date: %q", cert.DNSNames))
	}

	expectedURIs := crt.Spec.URISANs
	if !util.EqualUnsorted(pki.URLsToString(cert.URIs), expectedURIs) {
		errs = append(errs, fmt.Sprintf("URI SANs on TLS certificate not up to date: %q", cert.URIs))
	}

	// validate the ip addresses are correct
	if !util.EqualUnsorted(pki.IPAddressesToString(cert.IPAddresses), crt.Spec.IPAddresses) {
		errs = append(errs, fmt.Sprintf("IP addresses on TLS certificate not up to date: %q", pki.IPAddressesToString(cert.IPAddresses)))
	}

	// get a copy of the current secret resource
	// Note that we already know that it exists, no need to check for errors
	// TODO: Refactor so that the secret is passed as argument?
	secret, err := secretLister.Secrets(crt.Namespace).Get(crt.Spec.SecretName)

	if secret.Annotations == nil {
		secret.Annotations = make(map[string]string)
	}

	// validate that the issuer is correct
	if crt.Spec.IssuerRef.Name != secret.Annotations[v1alpha2.IssuerNameAnnotationKey] {
		errs = append(errs, fmt.Sprintf("Issuer of the certificate is not up to date: %q", secret.Annotations[v1alpha2.IssuerNameAnnotationKey]))
	}

	// validate that the issuer kind is correct
	if apiutil.IssuerKind(crt.Spec.IssuerRef) != secret.Annotations[v1alpha2.IssuerKindAnnotationKey] {
		errs = append(errs, fmt.Sprintf("Issuer kind of the certificate is not up to date: %q", secret.Annotations[v1alpha2.IssuerKindAnnotationKey]))
	}

	return len(errs) == 0, errs
}

func scheduleRenewal(ctx context.Context, lister corelisters.SecretLister, calc calculateDurationUntilRenewFn, queueFn func(interface{}, time.Duration), crt *v1alpha2.Certificate) {
	log := logf.FromContext(ctx)
	log = log.WithValues(
		logf.RelatedResourceNameKey, crt.Spec.SecretName,
		logf.RelatedResourceNamespaceKey, crt.Namespace,
		logf.RelatedResourceKindKey, "Secret",
	)

	key, err := keyFunc(crt)
	if err != nil {
		log.Error(err, "error getting key for certificate resource")
		return
	}

	cert, err := kube.SecretTLSCert(ctx, lister, crt.Namespace, crt.Spec.SecretName)
	if err != nil {
		if !errors.IsInvalidData(err) {
			log.Error(err, "error getting secret for certificate resource")
		}
		return
	}

	renewIn := calc(ctx, cert, crt)
	queueFn(key, renewIn)

	log.WithValues("duration_until_renewal", renewIn.String()).Info("certificate scheduled for renewal")
}

// staticTemporarySerialNumber is a fixed serial number we check for when
// updating the status of a certificate.
// It is used to identify temporarily generated certificates, so that friendly
// status messages can be displayed to users.
const staticTemporarySerialNumber = 0x1234567890

func isTemporaryCertificate(cert *x509.Certificate) bool {
	if cert == nil {
		return false
	}
	return cert.SerialNumber.Int64() == staticTemporarySerialNumber
}

// generateLocallySignedTemporaryCertificate signs a temporary certificate for
// the given certificate resource using a one-use temporary CA that is then
// discarded afterwards.
// This is to mitigate a potential attack against x509 certificates that use a
// predictable serial number and weak MD5 hashing algorithms.
// In practice, this shouldn't really be a concern anyway.
func generateLocallySignedTemporaryCertificate(crt *v1alpha2.Certificate, pk []byte) ([]byte, error) {
	// generate a throwaway self-signed root CA
	caPk, err := pki.GenerateECPrivateKey(pki.ECCurve521)
	if err != nil {
		return nil, err
	}
	caCertTemplate, err := pki.GenerateTemplate(&v1alpha2.Certificate{
		Spec: v1alpha2.CertificateSpec{
			CommonName: "cert-manager.local",
			IsCA:       true,
		},
	})
	if err != nil {
		return nil, err
	}
	_, caCert, err := pki.SignCertificate(caCertTemplate, caCertTemplate, caPk.Public(), caPk)
	if err != nil {
		return nil, err
	}

	// sign a temporary certificate using the root CA
	template, err := pki.GenerateTemplate(crt)
	if err != nil {
		return nil, err
	}
	template.SerialNumber = big.NewInt(staticTemporarySerialNumber)

	signeeKey, err := pki.DecodePrivateKeyBytes(pk)
	if err != nil {
		return nil, err
	}

	b, _, err := pki.SignCertificate(template, caCert, signeeKey.Public(), caPk)
	if err != nil {
		return nil, err
	}

	return b, nil
}

func updateCertificateStatus(ctx context.Context, m *metrics.Metrics, cmClient cmclient.Interface, old, new *v1alpha2.Certificate) (*v1alpha2.Certificate, error) {
	defer m.UpdateCertificateStatus(new)

	log := logf.FromContext(ctx, "updateStatus")
	oldBytes, _ := json.Marshal(old.Status)
	newBytes, _ := json.Marshal(new.Status)
	if reflect.DeepEqual(oldBytes, newBytes) {
		return nil, nil
	}
	log.V(logf.DebugLevel).Info("updating resource due to change in status", "diff", pretty.Diff(string(oldBytes), string(newBytes)))
	return cmClient.CertmanagerV1alpha2().Certificates(new.Namespace).UpdateStatus(new)
}

func certificateHasTemporaryCertificateAnnotation(crt *v1alpha2.Certificate) bool {
	if crt.Annotations == nil {
		return false
	}

	if val, ok := crt.Annotations[v1alpha2.IssueTemporaryCertificateAnnotation]; ok && val == "true" {
		return true
	}

	return false
}
<|endoftext|>


<commit_before>package cluster

import (
	"bufio"
	"fmt"
	"reflect"
	"strings"

	kapi "k8s.io/kubernetes/pkg/api"
	kerrs "k8s.io/kubernetes/pkg/api/errors"
	kclient "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/labels"

	authorizationapi "github.com/openshift/origin/pkg/authorization/api"
	osclient "github.com/openshift/origin/pkg/client"
	"github.com/openshift/origin/pkg/diagnostics/types"
	osapi "github.com/openshift/origin/pkg/image/api"
)

// ClusterRegistry is a Diagnostic to check that there is a working Docker registry.
type ClusterRegistry struct {
	KubeClient *kclient.Client
	OsClient   *osclient.Client
}

const (
	ClusterRegistryName = "ClusterRegistry"

	registryName   = "docker-registry"
	registryVolume = "registry-storage"
	clGetRegNone   = `
There is no "%s" service in project "%s". This is not strictly required to
be present; however, it is required for builds, and its absence probably
indicates an incomplete installation.

Please consult the documentation and use the 'oadm registry' command
to create a Docker registry.`

	clGetRegFailed = `
Client error while retrieving registry service. Client retrieved records
during discovery, so this is likely to be a transient error. Try running
diagnostics again. If this message persists, there may be a permissions
problem with getting records. The error was:

(%T) %[1]v `

	clRegNoPods = `
The "%s" service exists but has no associated pods, so it
is not available. Builds and deployments that use the registry will fail.`

	clRegNoRunningPods = `
The "%s" service exists but no pods currently running, so it
is not available. Builds and deployments that use the registry will fail.`

	clRegMultiPods = `
The "%s" service has multiple associated pods each using
ephemeral storage. These are likely to have inconsistent stores of
images. Builds and deployments that use images from the registry may
fail sporadically. Use a single registry or add a shared storage volume
to the registries.`

	clRegPodDown = `
The "%s" pod for the "%s" service is not running.
This may be transient, a scheduling error, or something else.`
	clRegPodLog = `
Failed to read the logs for the "%s" pod belonging to
the "%s" service. This is not a problem by itself but
prevents diagnostics from looking for errors in those logs. The
error encountered was:
%s`

	clRegPodConn = `
The pod logs for the "%s" pod belonging to
the "%s" service indicated a problem connecting to the
master to notify it about a new image. This typically results in builds
succeeding but not triggering deployments (as they wait on notifications
to the ImageStream from the build).

There are many reasons for this step to fail, including invalid
credentials, master outages, DNS failures, network errors, and so on. It
can be temporary or ongoing. Check the most recent error message from the
registry pod logs to determine the nature of the problem:

%s`

	clRegPodErr = `
The pod logs for the "%s" pod belonging to
the "%s" service indicated unknown errors.
This could result in problems with builds or deployments.
Please examine the log entries to determine if there might be
any related problems:
%s`

	clRegNoEP = `
The "%[1]s" service exists with %d associated pod(s), but there
are %d endpoints in the "%[1]s" service.
This mismatch likely indicates a system bug, and builds and
deployments that require the registry may fail sporadically.`

	clRegISDelFail = `
The diagnostics created an ImageStream named "%[1]s"
for test purposes and then attempted to delete it, which failed. This
should be an unusual, transient occurrence. The error encountered in
deleting it was:

%s

This message is just to notify you that this object exists.
You ought to be able to delete this object with:

oc delete imagestream/%[1]s -n default
`

	clRegISMismatch = `
Diagnostics created a test ImageStream and compared the registry IP
it received to the registry IP available via the %[1]s service.

%[1]s      : %[2]s
ImageStream registry : %[3]s

They do not match, which probably means that an administrator re-created
the %[1]s service but the master has cached the old service
IP address. Builds or deployments that use ImageStreams with the wrong
%[1]s IP will fail under this condition.

To resolve this issue, restarting the master (to clear the cache) should
be sufficient. Existing ImageStreams may need to be re-created.`
)

func (d *ClusterRegistry) Name() string {
	return ClusterRegistryName
}

func (d *ClusterRegistry) Description() string {
	return "Check that there is a working Docker registry"
}

func (d *ClusterRegistry) CanRun() (bool, error) {
	if d.OsClient == nil || d.KubeClient == nil {
		return false, fmt.Errorf("must have kube and os clients")
	}
	return userCan(d.OsClient, authorizationapi.AuthorizationAttributes{
		Namespace:    kapi.NamespaceDefault,
		Verb:         "get",
		Resource:     "services",
		ResourceName: registryName,
	})
}

func (d *ClusterRegistry) Check() types.DiagnosticResult {
	r := types.NewDiagnosticResult(ClusterRegistryName)
	if service := d.getRegistryService(r); service != nil {
		// Check that it actually has pod(s) selected and running
		if runningPods := d.getRegistryPods(service, r); len(runningPods) == 0 {
			r.Error("DClu1001", nil, fmt.Sprintf(clRegNoRunningPods, registryName))
			return r
		} else if d.checkRegistryEndpoints(runningPods, r) { // Check that matching endpoint exists on the service
			// attempt to create an imagestream and see if it gets the same registry service IP from the service cache
			d.verifyRegistryImageStream(service, r)
		}
	}
	return r
}

func (d *ClusterRegistry) getRegistryService(r types.DiagnosticResult) *kapi.Service {
	service, err := d.KubeClient.Services(kapi.NamespaceDefault).Get(registryName)
	if err != nil && reflect.TypeOf(err) == reflect.TypeOf(&kerrs.StatusError{}) {
		r.Warn("DClu1002", err, fmt.Sprintf(clGetRegNone, registryName, kapi.NamespaceDefault))
		return nil
	} else if err != nil {
		r.Error("DClu1003", err, fmt.Sprintf(clGetRegFailed, err))
		return nil
	}
	r.Debug("DClu1004", fmt.Sprintf("Found %s service with ports %v", registryName, service.Spec.Ports))
	return service
}

func (d *ClusterRegistry) getRegistryPods(service *kapi.Service, r types.DiagnosticResult) []*kapi.Pod {
	runningPods := []*kapi.Pod{}
	pods, err := d.KubeClient.Pods(kapi.NamespaceDefault).List(labels.SelectorFromSet(service.Spec.Selector), fields.Everything())
	if err != nil {
		r.Error("DClu1005", err, fmt.Sprintf("Finding pods for '%s' service failed. This should never happen. Error: (%T) %[2]v", registryName, err))
		return runningPods
	} else if len(pods.Items) < 1 {
		r.Error("DClu1006", nil, fmt.Sprintf(clRegNoPods, registryName))
		return runningPods
	} else if len(pods.Items) > 1 {
		// multiple registry pods using EmptyDir will be inconsistent
		for _, volume := range pods.Items[0].Spec.Volumes {
			if volume.Name == registryVolume && volume.EmptyDir != nil {
				r.Error("DClu1007", nil, fmt.Sprintf(clRegMultiPods, registryName))
				break
			}
		}
	}
	for _, pod := range pods.Items {
		r.Debug("DClu1008", fmt.Sprintf("Found %s pod with name %s", registryName, pod.ObjectMeta.Name))
		if pod.Status.Phase != kapi.PodRunning {
			r.Warn("DClu1009", nil, fmt.Sprintf(clRegPodDown, pod.ObjectMeta.Name, registryName))
		} else {
			runningPods = append(runningPods, &pod)
			// Check the logs for that pod for common issues (credentials, DNS resolution failure)
			d.checkRegistryLogs(&pod, r)
		}
	}
	return runningPods
}

func (d *ClusterRegistry) checkRegistryLogs(pod *kapi.Pod, r types.DiagnosticResult) {
	// pull out logs from the pod
	readCloser, err := d.KubeClient.RESTClient.Get().
		Namespace("default").Name(pod.ObjectMeta.Name).
		Resource("pods").SubResource("log").
		Param("follow", "false").
		Param("container", pod.Spec.Containers[0].Name).
		Stream()
	if err != nil {
		r.Warn("DClu1010", nil, fmt.Sprintf(clRegPodLog, pod.ObjectMeta.Name, registryName, fmt.Sprintf("(%T) %[1]v", err)))
		return
	}
	defer readCloser.Close()

	clientError := ""
	registryError := ""
	scanner := bufio.NewScanner(readCloser)
	for scanner.Scan() {
		logLine := scanner.Text()
		// TODO: once the logging API gets "since" and "tail" and "limit", limit to more recent log entries
		// https://github.com/kubernetes/kubernetes/issues/12447
		if strings.Contains(logLine, `level=error msg="client error:`) {
			clientError = logLine // end up showing only the most recent client error
		} else if strings.Contains(logLine, "level=error msg=") {
			registryError += "\n" + logLine // gather generic errors
		}
	}
	if clientError != "" {
		r.Error("DClu1011", nil, fmt.Sprintf(clRegPodConn, pod.ObjectMeta.Name, registryName, clientError))
	}
	if registryError != "" {
		r.Warn("DClu1012", nil, fmt.Sprintf(clRegPodErr, pod.ObjectMeta.Name, registryName, registryError))
	}

}

func (d *ClusterRegistry) checkRegistryEndpoints(pods []*kapi.Pod, r types.DiagnosticResult) bool {
	endPoint, err := d.KubeClient.Endpoints(kapi.NamespaceDefault).Get(registryName)
	if err != nil {
		r.Error("DClu1013", err, fmt.Sprintf(`Finding endpoints for "%s" service failed. This should never happen. Error: (%[2]T) %[2]v`, registryName, err))
		return false
	}
	numEP := 0
	for _, subs := range endPoint.Subsets {
		numEP += len(subs.Addresses)
	}
	if numEP != len(pods) {
		r.Warn("DClu1014", nil, fmt.Sprintf(clRegNoEP, registryName, len(pods), numEP))
		return false
	}
	return true
}

func (d *ClusterRegistry) verifyRegistryImageStream(service *kapi.Service, r types.DiagnosticResult) {
	imgStream, err := d.OsClient.ImageStreams(kapi.NamespaceDefault).Create(&osapi.ImageStream{ObjectMeta: kapi.ObjectMeta{GenerateName: "diagnostic-test"}})
	if err != nil {
		r.Error("DClu1015", err, fmt.Sprintf("Creating test ImageStream failed. Error: (%T) %[1]v", err))
		return
	}
	defer func() { // delete what we created, or notify that we couldn't
		if err := d.OsClient.ImageStreams(kapi.NamespaceDefault).Delete(imgStream.ObjectMeta.Name); err != nil {
			r.Warn("DClu1016", err, fmt.Sprintf(clRegISDelFail, imgStream.ObjectMeta.Name, fmt.Sprintf("(%T) %[1]s", err)))
		}
	}()
	imgStream, err = d.OsClient.ImageStreams(kapi.NamespaceDefault).Get(imgStream.ObjectMeta.Name) // status is filled in post-create
	if err != nil {
		r.Error("DClu1017", err, fmt.Sprintf("Getting created test ImageStream failed. Error: (%T) %[1]v", err))
		return
	}
	r.Debug("DClu1018", fmt.Sprintf("Created test ImageStream: %[1]v", imgStream))
	cacheHost := strings.SplitN(imgStream.Status.DockerImageRepository, "/", 2)[0]
	serviceHost := fmt.Sprintf("%s:%d", service.Spec.ClusterIP, service.Spec.Ports[0].Port)
	if cacheHost != serviceHost {
		r.Error("DClu1019", nil, fmt.Sprintf(clRegISMismatch, registryName, serviceHost, cacheHost))
	}
}
<commit_msg>Scan for selinux write error in registry diagnostic.<commit_after>package cluster

import (
	"bufio"
	"fmt"
	"reflect"
	"regexp"
	"strings"

	kapi "k8s.io/kubernetes/pkg/api"
	kerrs "k8s.io/kubernetes/pkg/api/errors"
	kclient "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/labels"

	authorizationapi "github.com/openshift/origin/pkg/authorization/api"
	osclient "github.com/openshift/origin/pkg/client"
	"github.com/openshift/origin/pkg/diagnostics/types"
	osapi "github.com/openshift/origin/pkg/image/api"
)

// ClusterRegistry is a Diagnostic to check that there is a working Docker registry.
type ClusterRegistry struct {
	KubeClient *kclient.Client
	OsClient   *osclient.Client
}

const (
	ClusterRegistryName = "ClusterRegistry"

	registryName   = "docker-registry"
	registryVolume = "registry-storage"
	clGetRegNone   = `
There is no "%s" service in project "%s". This is not strictly required to
be present; however, it is required for builds, and its absence probably
indicates an incomplete installation.

Please consult the documentation and use the 'oadm registry' command
to create a Docker registry.`

	clGetRegFailed = `
Client error while retrieving registry service. Client retrieved records
during discovery, so this is likely to be a transient error. Try running
diagnostics again. If this message persists, there may be a permissions
problem with getting records. The error was:

(%T) %[1]v `

	clRegNoPods = `
The "%s" service exists but has no associated pods, so it
is not available. Builds and deployments that use the registry will fail.`

	clRegNoRunningPods = `
The "%s" service exists but no pods currently running, so it
is not available. Builds and deployments that use the registry will fail.`

	clRegMultiPods = `
The "%s" service has multiple associated pods each using
ephemeral storage. These are likely to have inconsistent stores of
images. Builds and deployments that use images from the registry may
fail sporadically. Use a single registry or add a shared storage volume
to the registries.`

	clRegPodDown = `
The "%s" pod for the "%s" service is not running.
This may be transient, a scheduling error, or something else.`
	clRegPodLog = `
Failed to read the logs for the "%s" pod belonging to
the "%s" service. This is not a problem by itself but
prevents diagnostics from looking for errors in those logs. The
error encountered was:
%s`

	clRegPodConn = `
The pod logs for the "%s" pod belonging to
the "%s" service indicated a problem connecting to the
master to notify it about a new image. This typically results in builds
succeeding but not triggering deployments (as they wait on notifications
to the ImageStream from the build).

There are many reasons for this step to fail, including invalid
credentials, master outages, DNS failures, network errors, and so on. It
can be temporary or ongoing. Check the most recent error message from the
registry pod logs to determine the nature of the problem:

%s`

	clRegPodErr = `
The pod logs for the "%s" pod belonging to
the "%s" service indicated unknown errors.
This could result in problems with builds or deployments.
Please examine the log entries to determine if there might be
any related problems:
%s`

	clRegSelinuxErr = `
The pod logs for the "%s" pod belonging to
the "%s" service indicated the registry is unable to write to disk.
This may indicate an SELinux denial, or problems with volume
ownership/permissions.

For volume permission problems please consult the Persistent Storage section
of the Administrator's Guide.

In the case of SELinux this may be resolved on the node by running:

    sudo chcon -R -t svirt_sandbox_file_t [PATH_TO]/openshift.local.volumes

%s`

	clRegNoEP = `
The "%[1]s" service exists with %d associated pod(s), but there
are %d endpoints in the "%[1]s" service.
This mismatch likely indicates a system bug, and builds and
deployments that require the registry may fail sporadically.`

	clRegISDelFail = `
The diagnostics created an ImageStream named "%[1]s"
for test purposes and then attempted to delete it, which failed. This
should be an unusual, transient occurrence. The error encountered in
deleting it was:

%s

This message is just to notify you that this object exists.
You ought to be able to delete this object with:

oc delete imagestream/%[1]s -n default
`

	clRegISMismatch = `
Diagnostics created a test ImageStream and compared the registry IP
it received to the registry IP available via the %[1]s service.

%[1]s      : %[2]s
ImageStream registry : %[3]s

They do not match, which probably means that an administrator re-created
the %[1]s service but the master has cached the old service
IP address. Builds or deployments that use ImageStreams with the wrong
%[1]s IP will fail under this condition.

To resolve this issue, restarting the master (to clear the cache) should
be sufficient. Existing ImageStreams may need to be re-created.`
)

func (d *ClusterRegistry) Name() string {
	return ClusterRegistryName
}

func (d *ClusterRegistry) Description() string {
	return "Check that there is a working Docker registry"
}

func (d *ClusterRegistry) CanRun() (bool, error) {
	if d.OsClient == nil || d.KubeClient == nil {
		return false, fmt.Errorf("must have kube and os clients")
	}
	return userCan(d.OsClient, authorizationapi.AuthorizationAttributes{
		Namespace:    kapi.NamespaceDefault,
		Verb:         "get",
		Resource:     "services",
		ResourceName: registryName,
	})
}

func (d *ClusterRegistry) Check() types.DiagnosticResult {
	r := types.NewDiagnosticResult(ClusterRegistryName)
	if service := d.getRegistryService(r); service != nil {
		// Check that it actually has pod(s) selected and running
		if runningPods := d.getRegistryPods(service, r); len(runningPods) == 0 {
			r.Error("DClu1001", nil, fmt.Sprintf(clRegNoRunningPods, registryName))
			return r
		} else if d.checkRegistryEndpoints(runningPods, r) { // Check that matching endpoint exists on the service
			// attempt to create an imagestream and see if it gets the same registry service IP from the service cache
			d.verifyRegistryImageStream(service, r)
		}
	}
	return r
}

func (d *ClusterRegistry) getRegistryService(r types.DiagnosticResult) *kapi.Service {
	service, err := d.KubeClient.Services(kapi.NamespaceDefault).Get(registryName)
	if err != nil && reflect.TypeOf(err) == reflect.TypeOf(&kerrs.StatusError{}) {
		r.Warn("DClu1002", err, fmt.Sprintf(clGetRegNone, registryName, kapi.NamespaceDefault))
		return nil
	} else if err != nil {
		r.Error("DClu1003", err, fmt.Sprintf(clGetRegFailed, err))
		return nil
	}
	r.Debug("DClu1004", fmt.Sprintf("Found %s service with ports %v", registryName, service.Spec.Ports))
	return service
}

func (d *ClusterRegistry) getRegistryPods(service *kapi.Service, r types.DiagnosticResult) []*kapi.Pod {
	runningPods := []*kapi.Pod{}
	pods, err := d.KubeClient.Pods(kapi.NamespaceDefault).List(labels.SelectorFromSet(service.Spec.Selector), fields.Everything())
	if err != nil {
		r.Error("DClu1005", err, fmt.Sprintf("Finding pods for '%s' service failed. This should never happen. Error: (%T) %[2]v", registryName, err))
		return runningPods
	} else if len(pods.Items) < 1 {
		r.Error("DClu1006", nil, fmt.Sprintf(clRegNoPods, registryName))
		return runningPods
	} else if len(pods.Items) > 1 {
		// multiple registry pods using EmptyDir will be inconsistent
		for _, volume := range pods.Items[0].Spec.Volumes {
			if volume.Name == registryVolume && volume.EmptyDir != nil {
				r.Error("DClu1007", nil, fmt.Sprintf(clRegMultiPods, registryName))
				break
			}
		}
	}
	for _, pod := range pods.Items {
		r.Debug("DClu1008", fmt.Sprintf("Found %s pod with name %s", registryName, pod.ObjectMeta.Name))
		if pod.Status.Phase != kapi.PodRunning {
			r.Warn("DClu1009", nil, fmt.Sprintf(clRegPodDown, pod.ObjectMeta.Name, registryName))
		} else {
			runningPods = append(runningPods, &pod)
			// Check the logs for that pod for common issues (credentials, DNS resolution failure)
			d.checkRegistryLogs(&pod, r)
		}
	}
	return runningPods
}

func (d *ClusterRegistry) checkRegistryLogs(pod *kapi.Pod, r types.DiagnosticResult) {
	// pull out logs from the pod
	readCloser, err := d.KubeClient.RESTClient.Get().
		Namespace("default").Name(pod.ObjectMeta.Name).
		Resource("pods").SubResource("log").
		Param("follow", "false").
		Param("container", pod.Spec.Containers[0].Name).
		Stream()
	if err != nil {
		r.Warn("DClu1010", nil, fmt.Sprintf(clRegPodLog, pod.ObjectMeta.Name, registryName, fmt.Sprintf("(%T) %[1]v", err)))
		return
	}
	defer readCloser.Close()

	// Indicator that selinux is blocking the registry from writing to disk:
	selinuxErrorRegex, _ := regexp.Compile(".*level=error.*mkdir.*permission denied.*")
	// If seen after the above error regex, we know the problem has since been fixed:
	selinuxSuccessRegex, _ := regexp.Compile(".*level=info.*response completed.*http.request.method=PUT.*")

	clientError := ""
	registryError := ""
	selinuxError := ""

	scanner := bufio.NewScanner(readCloser)
	for scanner.Scan() {
		logLine := scanner.Text()
		// TODO: once the logging API gets "since" and "tail" and "limit", limit to more recent log entries
		// https://github.com/kubernetes/kubernetes/issues/12447
		if strings.Contains(logLine, `level=error msg="client error:`) {
			clientError = logLine // end up showing only the most recent client error
		} else if selinuxErrorRegex.MatchString(logLine) {
			selinuxError = logLine
		} else if selinuxSuccessRegex.MatchString(logLine) {
			// Check for a successful registry push, if this occurs after a selinux error
			// we can safely clear it, the problem has already been fixed.
			selinuxError = ""
		} else if strings.Contains(logLine, "level=error msg=") {
			registryError += "\n" + logLine // gather generic errors
		}
	}
	if clientError != "" {
		r.Error("DClu1011", nil, fmt.Sprintf(clRegPodConn, pod.ObjectMeta.Name, registryName, clientError))
	}
	if selinuxError != "" {
		r.Error("DClu1020", nil, fmt.Sprintf(clRegSelinuxErr, pod.ObjectMeta.Name, registryName, selinuxError))
	}
	if registryError != "" {
		r.Warn("DClu1012", nil, fmt.Sprintf(clRegPodErr, pod.ObjectMeta.Name, registryName, registryError))
	}
}

func (d *ClusterRegistry) checkRegistryEndpoints(pods []*kapi.Pod, r types.DiagnosticResult) bool {
	endPoint, err := d.KubeClient.Endpoints(kapi.NamespaceDefault).Get(registryName)
	if err != nil {
		r.Error("DClu1013", err, fmt.Sprintf(`Finding endpoints for "%s" service failed. This should never happen. Error: (%[2]T) %[2]v`, registryName, err))
		return false
	}
	numEP := 0
	for _, subs := range endPoint.Subsets {
		numEP += len(subs.Addresses)
	}
	if numEP != len(pods) {
		r.Warn("DClu1014", nil, fmt.Sprintf(clRegNoEP, registryName, len(pods), numEP))
		return false
	}
	return true
}

func (d *ClusterRegistry) verifyRegistryImageStream(service *kapi.Service, r types.DiagnosticResult) {
	imgStream, err := d.OsClient.ImageStreams(kapi.NamespaceDefault).Create(&osapi.ImageStream{ObjectMeta: kapi.ObjectMeta{GenerateName: "diagnostic-test"}})
	if err != nil {
		r.Error("DClu1015", err, fmt.Sprintf("Creating test ImageStream failed. Error: (%T) %[1]v", err))
		return
	}
	defer func() { // delete what we created, or notify that we couldn't
		if err := d.OsClient.ImageStreams(kapi.NamespaceDefault).Delete(imgStream.ObjectMeta.Name); err != nil {
			r.Warn("DClu1016", err, fmt.Sprintf(clRegISDelFail, imgStream.ObjectMeta.Name, fmt.Sprintf("(%T) %[1]s", err)))
		}
	}()
	imgStream, err = d.OsClient.ImageStreams(kapi.NamespaceDefault).Get(imgStream.ObjectMeta.Name) // status is filled in post-create
	if err != nil {
		r.Error("DClu1017", err, fmt.Sprintf("Getting created test ImageStream failed. Error: (%T) %[1]v", err))
		return
	}
	r.Debug("DClu1018", fmt.Sprintf("Created test ImageStream: %[1]v", imgStream))
	cacheHost := strings.SplitN(imgStream.Status.DockerImageRepository, "/", 2)[0]
	serviceHost := fmt.Sprintf("%s:%d", service.Spec.ClusterIP, service.Spec.Ports[0].Port)
	if cacheHost != serviceHost {
		r.Error("DClu1019", nil, fmt.Sprintf(clRegISMismatch, registryName, serviceHost, cacheHost))
	}
}
<|endoftext|>


<commit_before>package integration

import (
	"os"
	"testing"

	"github.com/pkg/errors"

	"google.golang.org/grpc"

	"github.com/havoc-io/mutagen/cmd"
	"github.com/havoc-io/mutagen/pkg/agent"
	"github.com/havoc-io/mutagen/pkg/daemon"
	mgrpc "github.com/havoc-io/mutagen/pkg/grpc"
	daemonsvc "github.com/havoc-io/mutagen/pkg/service/daemon"
	promptsvc "github.com/havoc-io/mutagen/pkg/service/prompt"
	sessionsvc "github.com/havoc-io/mutagen/pkg/service/session"
	"github.com/havoc-io/mutagen/pkg/session"

	// Explicitly import packages that need to register protocol handlers.
	_ "github.com/havoc-io/mutagen/pkg/integration/protocols/netpipe"
	_ "github.com/havoc-io/mutagen/pkg/protocols/docker"
	_ "github.com/havoc-io/mutagen/pkg/protocols/local"
	_ "github.com/havoc-io/mutagen/pkg/protocols/ssh"
)

// sessionManager is the session manager for the integration testing daemon. It
// is exposed for integration tests that operate at the API level (as opposed to
// the gRPC or command line level).
var sessionManager *session.Manager

// testMainInternal is the internal testing entry point, needed so that shutdown
// operations can be deferred (since TestMain will invoke os.Exit). It copies
// the mutagen executable to a well-known path, sets up the agent bundle to work
// during testing, sets up a functionally complete daemon instance for testing,
// runs integration tests, and finally tears down all of the aforementioned
// infrastructure.
func testMainInternal(m *testing.M) (int, error) {
	// Override the expected agent bundle location.
	agent.ExpectedBundleLocation = agent.BundleLocationBuildDirectory

	// Acquire the daemon lock and defer its release.
	lock, err := daemon.AcquireLock()
	if err != nil {
		return -1, errors.Wrap(err, "unable to acquire daemon lock")
	}
	defer lock.Unlock()

	// Create a session manager and defer its shutdown. Note that we assign to
	// the global instance here.
	sessionManager, err = session.NewManager()
	if err != nil {
		return -1, errors.Wrap(err, "unable to create session manager")
	}
	defer sessionManager.Shutdown()

	// Create the gRPC server and defer its stoppage. We use a hard stop rather
	// than a graceful stop so that it doesn't hang on open requests.
	server := grpc.NewServer(
		grpc.MaxSendMsgSize(mgrpc.MaximumIPCMessageSize),
		grpc.MaxRecvMsgSize(mgrpc.MaximumIPCMessageSize),
	)
	defer server.Stop()

	// Create and register the daemon service and defer its shutdown.
	daemonServer := daemonsvc.NewServer()
	daemonsvc.RegisterDaemonServer(server, daemonServer)
	defer daemonServer.Shutdown()

	// Create and register the prompt service.
	promptsvc.RegisterPromptingServer(server, promptsvc.NewServer())

	// Create and register the session service.
	sessionsvc.RegisterSessionsServer(server, sessionsvc.NewServer(sessionManager))

	// Create the daemon listener and defer its closure.
	listener, err := daemon.NewListener()
	if err != nil {
		return -1, errors.Wrap(err, "unable to create daemon listener")
	}
	defer listener.Close()

	// Serve incoming connections in a separate Goroutine. We don't monitor for
	// errors since there's nothing that we can do about them and because
	// they'll likely show up in the test output anyway.
	go server.Serve(listener)

	// Run tests.
	return m.Run(), nil
}

// TestMain is the entry point for integration tests (overriding the default
// generated entry point).
func TestMain(m *testing.M) {
	// Invoke the internal entry point. If there's an error, print it out before
	// exiting.
	result, err := testMainInternal(m)
	if err != nil {
		cmd.Error(err)
	}

	// Exit with the result.
	os.Exit(result)
}
<commit_msg>Added fix for integration test imports.<commit_after>package integration

import (
	"os"
	"testing"

	"github.com/pkg/errors"

	"google.golang.org/grpc"

	"github.com/havoc-io/mutagen/cmd"
	"github.com/havoc-io/mutagen/pkg/agent"
	"github.com/havoc-io/mutagen/pkg/daemon"
	daemonsvc "github.com/havoc-io/mutagen/pkg/service/daemon"
	promptsvc "github.com/havoc-io/mutagen/pkg/service/prompt"
	sessionsvc "github.com/havoc-io/mutagen/pkg/service/session"
	"github.com/havoc-io/mutagen/pkg/session"

	// Explicitly import packages that need to register protocol handlers.
	_ "github.com/havoc-io/mutagen/pkg/integration/protocols/netpipe"
	_ "github.com/havoc-io/mutagen/pkg/protocols/docker"
	_ "github.com/havoc-io/mutagen/pkg/protocols/local"
	_ "github.com/havoc-io/mutagen/pkg/protocols/ssh"
)

// sessionManager is the session manager for the integration testing daemon. It
// is exposed for integration tests that operate at the API level (as opposed to
// the gRPC or command line level).
var sessionManager *session.Manager

// testMainInternal is the internal testing entry point, needed so that shutdown
// operations can be deferred (since TestMain will invoke os.Exit). It copies
// the mutagen executable to a well-known path, sets up the agent bundle to work
// during testing, sets up a functionally complete daemon instance for testing,
// runs integration tests, and finally tears down all of the aforementioned
// infrastructure.
func testMainInternal(m *testing.M) (int, error) {
	// Override the expected agent bundle location.
	agent.ExpectedBundleLocation = agent.BundleLocationBuildDirectory

	// Acquire the daemon lock and defer its release.
	lock, err := daemon.AcquireLock()
	if err != nil {
		return -1, errors.Wrap(err, "unable to acquire daemon lock")
	}
	defer lock.Unlock()

	// Create a session manager and defer its shutdown. Note that we assign to
	// the global instance here.
	sessionManager, err = session.NewManager()
	if err != nil {
		return -1, errors.Wrap(err, "unable to create session manager")
	}
	defer sessionManager.Shutdown()

	// Create the gRPC server and defer its stoppage. We use a hard stop rather
	// than a graceful stop so that it doesn't hang on open requests.
	server := grpc.NewServer(
		grpc.MaxSendMsgSize(daemon.MaximumIPCMessageSize),
		grpc.MaxRecvMsgSize(daemon.MaximumIPCMessageSize),
	)
	defer server.Stop()

	// Create and register the daemon service and defer its shutdown.
	daemonServer := daemonsvc.NewServer()
	daemonsvc.RegisterDaemonServer(server, daemonServer)
	defer daemonServer.Shutdown()

	// Create and register the prompt service.
	promptsvc.RegisterPromptingServer(server, promptsvc.NewServer())

	// Create and register the session service.
	sessionsvc.RegisterSessionsServer(server, sessionsvc.NewServer(sessionManager))

	// Create the daemon listener and defer its closure.
	listener, err := daemon.NewListener()
	if err != nil {
		return -1, errors.Wrap(err, "unable to create daemon listener")
	}
	defer listener.Close()

	// Serve incoming connections in a separate Goroutine. We don't monitor for
	// errors since there's nothing that we can do about them and because
	// they'll likely show up in the test output anyway.
	go server.Serve(listener)

	// Run tests.
	return m.Run(), nil
}

// TestMain is the entry point for integration tests (overriding the default
// generated entry point).
func TestMain(m *testing.M) {
	// Invoke the internal entry point. If there's an error, print it out before
	// exiting.
	result, err := testMainInternal(m)
	if err != nil {
		cmd.Error(err)
	}

	// Exit with the result.
	os.Exit(result)
}
<|endoftext|>