<commit_before>// Package client implements a Go client for CFSSL API commands.
package client

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	stderr "errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/cloudflare/cfssl/api"
	"github.com/cloudflare/cfssl/auth"
	"github.com/cloudflare/cfssl/errors"
	"github.com/cloudflare/cfssl/info"
	"github.com/cloudflare/cfssl/log"
)

// A server points to a single remote CFSSL instance.
type server struct {
	URL            string
	TLSConfig      *tls.Config
	reqModifier    func(*http.Request, []byte)
	RequestTimeout time.Duration
	proxy          func(*http.Request) (*url.URL, error)
}

// A Remote points to at least one (but possibly multiple) remote
// CFSSL instances. It must be able to perform a authenticated and
// unauthenticated certificate signing requests, return information
// about the CA on the other end, and return a list of the hosts that
// are used by the remote.
type Remote interface {
	AuthSign(req, id []byte, provider auth.Provider) ([]byte, error)
	Sign(jsonData []byte) ([]byte, error)
	Info(jsonData []byte) (*info.Resp, error)
	Hosts() []string
	SetReqModifier(func(*http.Request, []byte))
	SetRequestTimeout(d time.Duration)
	SetProxy(func(*http.Request) (*url.URL, error))
}

// NewServer sets up a new server target. The address should be of
// The format [protocol:]name[:port] of the remote CFSSL instance.
// If no protocol is given http is default. If no port
// is specified, the CFSSL default port (8888) is used. If the name is
// a comma-separated list of hosts, an ordered group will be returned.
func NewServer(addr string) Remote {
	return NewServerTLS(addr, nil)
}

// NewServerTLS is the TLS version of NewServer
func NewServerTLS(addr string, tlsConfig *tls.Config) Remote {
	addrs := strings.Split(addr, ",")

	var remote Remote

	if len(addrs) > 1 {
		remote, _ = NewGroup(addrs, tlsConfig, StrategyOrderedList)
	} else {
		u, err := normalizeURL(addrs[0])
		if err != nil {
			log.Errorf("bad url: %v", err)
			return nil
		}
		srv := newServer(u, tlsConfig)
		if srv != nil {
			remote = srv
		}
	}
	return remote
}

func (srv *server) Hosts() []string {
	return []string{srv.URL}
}

func (srv *server) SetReqModifier(mod func(*http.Request, []byte)) {
	srv.reqModifier = mod
}

func (srv *server) SetRequestTimeout(timeout time.Duration) {
	srv.RequestTimeout = timeout
}

func (srv *server) SetProxy(proxy func(*http.Request) (*url.URL, error)) {
	srv.proxy = proxy
}

func newServer(u *url.URL, tlsConfig *tls.Config) *server {
	URL := u.String()
	return &server{
		URL:       URL,
		TLSConfig: tlsConfig,
	}
}

func (srv *server) getURL(endpoint string) string {
	return fmt.Sprintf("%s/api/v1/cfssl/%s", srv.URL, endpoint)
}

func (srv *server) createTransport() (transport *http.Transport) {
	transport = new(http.Transport)
	// Setup HTTPS client
	tlsConfig := srv.TLSConfig
	tlsConfig.BuildNameToCertificate()
	transport.TLSClientConfig = tlsConfig
	// Setup Proxy
	transport.Proxy = srv.proxy
	return transport
}

// post connects to the remote server and returns a Response struct
func (srv *server) post(url string, jsonData []byte) (*api.Response, error) {
	var resp *http.Response
	var err error
	client := &http.Client{}
	if srv.TLSConfig != nil {
		client.Transport = srv.createTransport()
	}
	if srv.RequestTimeout != 0 {
		client.Timeout = srv.RequestTimeout
	}
	req, err := http.NewRequest("POST", url, bytes.NewReader(jsonData))
	if err != nil {
		err = fmt.Errorf("failed POST to %s: %v", url, err)
		return nil, errors.Wrap(errors.APIClientError, errors.ClientHTTPError, err)
	}
	req.Close = true
	req.Header.Set("content-type", "application/json")
	if srv.reqModifier != nil {
		srv.reqModifier(req, jsonData)
	}
	resp, err = client.Do(req)
	if err != nil {
		err = fmt.Errorf("failed POST to %s: %v", url, err)
		return nil, errors.Wrap(errors.APIClientError, errors.ClientHTTPError, err)
	}
	defer req.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, errors.Wrap(errors.APIClientError, errors.IOError, err)
	}

	if resp.StatusCode != http.StatusOK {
		log.Errorf("http error with %s", url)
		return nil, errors.Wrap(errors.APIClientError, errors.ClientHTTPError, stderr.New(string(body)))
	}

	var response api.Response
	err = json.Unmarshal(body, &response)
	if err != nil {
		log.Debug("Unable to parse response body:", string(body))
		return nil, errors.Wrap(errors.APIClientError, errors.JSONError, err)
	}

	if !response.Success || response.Result == nil {
		if len(response.Errors) > 0 {
			return nil, errors.Wrap(errors.APIClientError, errors.ServerRequestFailed, stderr.New(response.Errors[0].Message))
		}
		return nil, errors.New(errors.APIClientError, errors.ServerRequestFailed)
	}

	return &response, nil
}

// AuthSign fills out an authenticated signing request to the server,
// receiving a certificate or error in response.
// It takes the serialized JSON request to send, remote address and
// authentication provider.
func (srv *server) AuthSign(req, id []byte, provider auth.Provider) ([]byte, error) {
	return srv.authReq(req, id, provider, "sign")
}

// AuthInfo fills out an authenticated info request to the server,
// receiving a certificate or error in response.
// It takes the serialized JSON request to send, remote address and
// authentication provider.
func (srv *server) AuthInfo(req, id []byte, provider auth.Provider) ([]byte, error) {
	return srv.authReq(req, id, provider, "info")
}

// authReq is the common logic for AuthSign and AuthInfo -- perform the given
// request, and return the resultant certificate.
// The target is either 'sign' or 'info'.
func (srv *server) authReq(req, ID []byte, provider auth.Provider, target string) ([]byte, error) {
	url := srv.getURL("auth" + target)

	token, err := provider.Token(req)
	if err != nil {
		return nil, errors.Wrap(errors.APIClientError, errors.AuthenticationFailure, err)
	}

	aReq := &auth.AuthenticatedRequest{
		Timestamp:     time.Now().Unix(),
		RemoteAddress: ID,
		Token:         token,
		Request:       req,
	}

	jsonData, err := json.Marshal(aReq)
	if err != nil {
		return nil, errors.Wrap(errors.APIClientError, errors.JSONError, err)
	}

	response, err := srv.post(url, jsonData)
	if err != nil {
		return nil, err
	}

	result, ok := response.Result.(map[string]interface{})
	if !ok {
		return nil, errors.New(errors.APIClientError, errors.JSONError)
	}

	cert, ok := result["certificate"].(string)
	if !ok {
		return nil, errors.New(errors.APIClientError, errors.JSONError)
	}

	return []byte(cert), nil
}

// Sign sends a signature request to the remote CFSSL server,
// receiving a signed certificate or an error in response.
// It takes the serialized JSON request to send.
func (srv *server) Sign(jsonData []byte) ([]byte, error) {
	return srv.request(jsonData, "sign")
}

// Info sends an info request to the remote CFSSL server, receiving a
// response or an error in response.
// It takes the serialized JSON request to send.
func (srv *server) Info(jsonData []byte) (*info.Resp, error) {
	res, err := srv.getResultMap(jsonData, "info")
	if err != nil {
		return nil, err
	}

	info := new(info.Resp)

	if val, ok := res["certificate"]; ok {
		info.Certificate = val.(string)
	}
	var usages []interface{}
	if val, ok := res["usages"]; ok && val != nil {
		usages = val.([]interface{})
	}
	if val, ok := res["expiry"]; ok && val != nil {
		info.ExpiryString = val.(string)
	}

	info.Usage = make([]string, len(usages))
	for i, s := range usages {
		info.Usage[i] = s.(string)
	}

	return info, nil
}

func (srv *server) getResultMap(jsonData []byte, target string) (result map[string]interface{}, err error) {
	url := srv.getURL(target)
	response, err := srv.post(url, jsonData)
	if err != nil {
		return
	}
	result, ok := response.Result.(map[string]interface{})
	if !ok {
		err = errors.Wrap(errors.APIClientError, errors.ClientHTTPError, stderr.New("response is formatted improperly"))
		return
	}
	return
}

// request performs the common logic for Sign and Info, performing the actual
// request and returning the resultant certificate.
func (srv *server) request(jsonData []byte, target string) ([]byte, error) {
	result, err := srv.getResultMap(jsonData, target)
	if err != nil {
		return nil, err
	}
	cert := result["certificate"].(string)
	if cert != "" {
		return []byte(cert), nil
	}

	return nil, errors.Wrap(errors.APIClientError, errors.ClientHTTPError, stderr.New("response doesn't contain certificate."))
}

// AuthRemote acts as a Remote with a default Provider for AuthSign.
type AuthRemote struct {
	Remote
	provider auth.Provider
}

// NewAuthServer sets up a new auth server target with an addr
// in the same format at NewServer and a default authentication provider to
// use for Sign requests.
func NewAuthServer(addr string, tlsConfig *tls.Config, provider auth.Provider) *AuthRemote {
	return &AuthRemote{
		Remote:   NewServerTLS(addr, tlsConfig),
		provider: provider,
	}
}

// Sign is overloaded to perform an AuthSign request using the default auth provider.
func (ar *AuthRemote) Sign(req []byte) ([]byte, error) {
	return ar.AuthSign(req, nil, ar.provider)
}

// nomalizeURL checks for http/https protocol, appends "http" as default protocol if not defiend in url
func normalizeURL(addr string) (*url.URL, error) {
	addr = strings.TrimSpace(addr)

	u, err := url.Parse(addr)
	if err != nil {
		return nil, err
	}

	if u.Opaque != "" {
		u.Host = net.JoinHostPort(u.Scheme, u.Opaque)
		u.Opaque = ""
	} else if u.Path != "" && !strings.Contains(u.Path, ":") {
		u.Host = net.JoinHostPort(u.Path, "8888")
		u.Path = ""
	} else if u.Scheme == "" {
		u.Host = u.Path
		u.Path = ""
	}

	if u.Scheme != "https" {
		u.Scheme = "http"
	}

	_, port, err := net.SplitHostPort(u.Host)
	if err != nil {
		_, port, err = net.SplitHostPort(u.Host + ":8888")
		if err != nil {
			return nil, err
		}
	}

	if port != "" {
		_, err = strconv.Atoi(port)
		if err != nil {
			return nil, err
		}
	}
	return u, nil
}
<commit_msg>Create transports with same defaults as http.DefaultTransport<commit_after>// Package client implements a Go client for CFSSL API commands.
package client

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	stderr "errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/cloudflare/cfssl/api"
	"github.com/cloudflare/cfssl/auth"
	"github.com/cloudflare/cfssl/errors"
	"github.com/cloudflare/cfssl/info"
	"github.com/cloudflare/cfssl/log"
)

// A server points to a single remote CFSSL instance.
type server struct {
	URL            string
	TLSConfig      *tls.Config
	reqModifier    func(*http.Request, []byte)
	RequestTimeout time.Duration
	proxy          func(*http.Request) (*url.URL, error)
}

// A Remote points to at least one (but possibly multiple) remote
// CFSSL instances. It must be able to perform a authenticated and
// unauthenticated certificate signing requests, return information
// about the CA on the other end, and return a list of the hosts that
// are used by the remote.
type Remote interface {
	AuthSign(req, id []byte, provider auth.Provider) ([]byte, error)
	Sign(jsonData []byte) ([]byte, error)
	Info(jsonData []byte) (*info.Resp, error)
	Hosts() []string
	SetReqModifier(func(*http.Request, []byte))
	SetRequestTimeout(d time.Duration)
	SetProxy(func(*http.Request) (*url.URL, error))
}

// NewServer sets up a new server target. The address should be of
// The format [protocol:]name[:port] of the remote CFSSL instance.
// If no protocol is given http is default. If no port
// is specified, the CFSSL default port (8888) is used. If the name is
// a comma-separated list of hosts, an ordered group will be returned.
func NewServer(addr string) Remote {
	return NewServerTLS(addr, nil)
}

// NewServerTLS is the TLS version of NewServer
func NewServerTLS(addr string, tlsConfig *tls.Config) Remote {
	addrs := strings.Split(addr, ",")

	var remote Remote

	if len(addrs) > 1 {
		remote, _ = NewGroup(addrs, tlsConfig, StrategyOrderedList)
	} else {
		u, err := normalizeURL(addrs[0])
		if err != nil {
			log.Errorf("bad url: %v", err)
			return nil
		}
		srv := newServer(u, tlsConfig)
		if srv != nil {
			remote = srv
		}
	}
	return remote
}

func (srv *server) Hosts() []string {
	return []string{srv.URL}
}

func (srv *server) SetReqModifier(mod func(*http.Request, []byte)) {
	srv.reqModifier = mod
}

func (srv *server) SetRequestTimeout(timeout time.Duration) {
	srv.RequestTimeout = timeout
}

func (srv *server) SetProxy(proxy func(*http.Request) (*url.URL, error)) {
	srv.proxy = proxy
}

func newServer(u *url.URL, tlsConfig *tls.Config) *server {
	URL := u.String()
	return &server{
		URL:       URL,
		TLSConfig: tlsConfig,
	}
}

func (srv *server) getURL(endpoint string) string {
	return fmt.Sprintf("%s/api/v1/cfssl/%s", srv.URL, endpoint)
}

func (srv *server) createTransport() *http.Transport {
	// Start with defaults from http.DefaultTransport
	transport := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		DialContext: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
			DualStack: true,
		}).DialContext,
		MaxIdleConns:          100,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	}
	// Setup HTTPS client
	tlsConfig := srv.TLSConfig
	tlsConfig.BuildNameToCertificate()
	transport.TLSClientConfig = tlsConfig
	// Setup Proxy
	if srv.proxy != nil {
		transport.Proxy = srv.proxy
	}
	return transport
}

// post connects to the remote server and returns a Response struct
func (srv *server) post(url string, jsonData []byte) (*api.Response, error) {
	var resp *http.Response
	var err error
	client := &http.Client{}
	if srv.TLSConfig != nil {
		client.Transport = srv.createTransport()
	}
	if srv.RequestTimeout != 0 {
		client.Timeout = srv.RequestTimeout
	}
	req, err := http.NewRequest("POST", url, bytes.NewReader(jsonData))
	if err != nil {
		err = fmt.Errorf("failed POST to %s: %v", url, err)
		return nil, errors.Wrap(errors.APIClientError, errors.ClientHTTPError, err)
	}
	req.Close = true
	req.Header.Set("content-type", "application/json")
	if srv.reqModifier != nil {
		srv.reqModifier(req, jsonData)
	}
	resp, err = client.Do(req)
	if err != nil {
		err = fmt.Errorf("failed POST to %s: %v", url, err)
		return nil, errors.Wrap(errors.APIClientError, errors.ClientHTTPError, err)
	}
	defer req.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, errors.Wrap(errors.APIClientError, errors.IOError, err)
	}

	if resp.StatusCode != http.StatusOK {
		log.Errorf("http error with %s", url)
		return nil, errors.Wrap(errors.APIClientError, errors.ClientHTTPError, stderr.New(string(body)))
	}

	var response api.Response
	err = json.Unmarshal(body, &response)
	if err != nil {
		log.Debug("Unable to parse response body:", string(body))
		return nil, errors.Wrap(errors.APIClientError, errors.JSONError, err)
	}

	if !response.Success || response.Result == nil {
		if len(response.Errors) > 0 {
			return nil, errors.Wrap(errors.APIClientError, errors.ServerRequestFailed, stderr.New(response.Errors[0].Message))
		}
		return nil, errors.New(errors.APIClientError, errors.ServerRequestFailed)
	}

	return &response, nil
}

// AuthSign fills out an authenticated signing request to the server,
// receiving a certificate or error in response.
// It takes the serialized JSON request to send, remote address and
// authentication provider.
func (srv *server) AuthSign(req, id []byte, provider auth.Provider) ([]byte, error) {
	return srv.authReq(req, id, provider, "sign")
}

// AuthInfo fills out an authenticated info request to the server,
// receiving a certificate or error in response.
// It takes the serialized JSON request to send, remote address and
// authentication provider.
func (srv *server) AuthInfo(req, id []byte, provider auth.Provider) ([]byte, error) {
	return srv.authReq(req, id, provider, "info")
}

// authReq is the common logic for AuthSign and AuthInfo -- perform the given
// request, and return the resultant certificate.
// The target is either 'sign' or 'info'.
func (srv *server) authReq(req, ID []byte, provider auth.Provider, target string) ([]byte, error) {
	url := srv.getURL("auth" + target)

	token, err := provider.Token(req)
	if err != nil {
		return nil, errors.Wrap(errors.APIClientError, errors.AuthenticationFailure, err)
	}

	aReq := &auth.AuthenticatedRequest{
		Timestamp:     time.Now().Unix(),
		RemoteAddress: ID,
		Token:         token,
		Request:       req,
	}

	jsonData, err := json.Marshal(aReq)
	if err != nil {
		return nil, errors.Wrap(errors.APIClientError, errors.JSONError, err)
	}

	response, err := srv.post(url, jsonData)
	if err != nil {
		return nil, err
	}

	result, ok := response.Result.(map[string]interface{})
	if !ok {
		return nil, errors.New(errors.APIClientError, errors.JSONError)
	}

	cert, ok := result["certificate"].(string)
	if !ok {
		return nil, errors.New(errors.APIClientError, errors.JSONError)
	}

	return []byte(cert), nil
}

// Sign sends a signature request to the remote CFSSL server,
// receiving a signed certificate or an error in response.
// It takes the serialized JSON request to send.
func (srv *server) Sign(jsonData []byte) ([]byte, error) {
	return srv.request(jsonData, "sign")
}

// Info sends an info request to the remote CFSSL server, receiving a
// response or an error in response.
// It takes the serialized JSON request to send.
func (srv *server) Info(jsonData []byte) (*info.Resp, error) {
	res, err := srv.getResultMap(jsonData, "info")
	if err != nil {
		return nil, err
	}

	info := new(info.Resp)

	if val, ok := res["certificate"]; ok {
		info.Certificate = val.(string)
	}
	var usages []interface{}
	if val, ok := res["usages"]; ok && val != nil {
		usages = val.([]interface{})
	}
	if val, ok := res["expiry"]; ok && val != nil {
		info.ExpiryString = val.(string)
	}

	info.Usage = make([]string, len(usages))
	for i, s := range usages {
		info.Usage[i] = s.(string)
	}

	return info, nil
}

func (srv *server) getResultMap(jsonData []byte, target string) (result map[string]interface{}, err error) {
	url := srv.getURL(target)
	response, err := srv.post(url, jsonData)
	if err != nil {
		return
	}
	result, ok := response.Result.(map[string]interface{})
	if !ok {
		err = errors.Wrap(errors.APIClientError, errors.ClientHTTPError, stderr.New("response is formatted improperly"))
		return
	}
	return
}

// request performs the common logic for Sign and Info, performing the actual
// request and returning the resultant certificate.
func (srv *server) request(jsonData []byte, target string) ([]byte, error) {
	result, err := srv.getResultMap(jsonData, target)
	if err != nil {
		return nil, err
	}
	cert := result["certificate"].(string)
	if cert != "" {
		return []byte(cert), nil
	}

	return nil, errors.Wrap(errors.APIClientError, errors.ClientHTTPError, stderr.New("response doesn't contain certificate."))
}

// AuthRemote acts as a Remote with a default Provider for AuthSign.
type AuthRemote struct {
	Remote
	provider auth.Provider
}

// NewAuthServer sets up a new auth server target with an addr
// in the same format at NewServer and a default authentication provider to
// use for Sign requests.
func NewAuthServer(addr string, tlsConfig *tls.Config, provider auth.Provider) *AuthRemote {
	return &AuthRemote{
		Remote:   NewServerTLS(addr, tlsConfig),
		provider: provider,
	}
}

// Sign is overloaded to perform an AuthSign request using the default auth provider.
func (ar *AuthRemote) Sign(req []byte) ([]byte, error) {
	return ar.AuthSign(req, nil, ar.provider)
}

// nomalizeURL checks for http/https protocol, appends "http" as default protocol if not defiend in url
func normalizeURL(addr string) (*url.URL, error) {
	addr = strings.TrimSpace(addr)

	u, err := url.Parse(addr)
	if err != nil {
		return nil, err
	}

	if u.Opaque != "" {
		u.Host = net.JoinHostPort(u.Scheme, u.Opaque)
		u.Opaque = ""
	} else if u.Path != "" && !strings.Contains(u.Path, ":") {
		u.Host = net.JoinHostPort(u.Path, "8888")
		u.Path = ""
	} else if u.Scheme == "" {
		u.Host = u.Path
		u.Path = ""
	}

	if u.Scheme != "https" {
		u.Scheme = "http"
	}

	_, port, err := net.SplitHostPort(u.Host)
	if err != nil {
		_, port, err = net.SplitHostPort(u.Host + ":8888")
		if err != nil {
			return nil, err
		}
	}

	if port != "" {
		_, err = strconv.Atoi(port)
		if err != nil {
			return nil, err
		}
	}
	return u, nil
}
<|endoftext|>


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

package filter

import (
	"bytes"
	"regexp"
)

// filterOutput filters output from juju.
//
// It removes all lines that does not represent useful output, like juju's
// logging and Python's deprecation warnings.
func FilterOutput(output []byte) []byte {
	var result [][]byte
	var ignore bool
	deprecation := []byte("DeprecationWarning")
	regexLog := regexp.MustCompile(`^\d{4}-\d{2}-\d{2}\s\d{2}:\d{2}:\d{2}`)
	regexSshWarning := regexp.MustCompile(`^Warning: Permanently added`)
	regexPythonWarning := regexp.MustCompile(`^.*warnings.warn`)
	regexUserWarning := regexp.MustCompile(`^.*UserWarning`)
	lines := bytes.Split(output, []byte{'\n'})
	for _, line := range lines {
		if ignore {
			ignore = false
			continue
		}
		if bytes.Contains(line, deprecation) {
			ignore = true
			continue
		}
		if !regexSshWarning.Match(line) && !regexLog.Match(line) && !regexPythonWarning.Match(line) && !regexUserWarning.Match(line) {
			result = append(result, line)
		}
	}
	return bytes.Join(result, []byte{'\n'})
}
<commit_msg>fixed the method name on docs<commit_after>// Copyright 2012 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package filter

import (
	"bytes"
	"regexp"
)

// FilterOutput filters output from juju.
//
// It removes all lines that does not represent useful output, like juju's
// logging and Python's deprecation warnings.
func FilterOutput(output []byte) []byte {
	var result [][]byte
	var ignore bool
	deprecation := []byte("DeprecationWarning")
	regexLog := regexp.MustCompile(`^\d{4}-\d{2}-\d{2}\s\d{2}:\d{2}:\d{2}`)
	regexSshWarning := regexp.MustCompile(`^Warning: Permanently added`)
	regexPythonWarning := regexp.MustCompile(`^.*warnings.warn`)
	regexUserWarning := regexp.MustCompile(`^.*UserWarning`)
	lines := bytes.Split(output, []byte{'\n'})
	for _, line := range lines {
		if ignore {
			ignore = false
			continue
		}
		if bytes.Contains(line, deprecation) {
			ignore = true
			continue
		}
		if !regexSshWarning.Match(line) && !regexLog.Match(line) && !regexPythonWarning.Match(line) && !regexUserWarning.Match(line) {
			result = append(result, line)
		}
	}
	return bytes.Join(result, []byte{'\n'})
}
<|endoftext|>


<commit_before>package api

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/CenturyLinkLabs/docker-reg-client/registry"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
)

type mockConnection struct {
	mock.Mock
}

func (m *mockConnection) Status() (Status, error) {
	args := m.Mock.Called()
	return args.Get(0).(Status), nil
}

func (m *mockConnection) Connect(repo string) error {
	args := m.Mock.Called(repo)
	return args.Error(0)
}

func (m *mockConnection) GetImageID(image string, tag string) (string, error) {
	args := m.Mock.Called(image, tag)
	return args.Get(0).(string), nil
}

func (m *mockConnection) GetAncestry(id string) ([]string, error) {
	args := m.Mock.Called(id)
	return args.Get(0).([]string), nil
}

func (m *mockConnection) GetMetadata(layer string) (*registry.ImageMetadata, error) {
	args := m.Mock.Called(layer)
	return args.Get(0).(*registry.ImageMetadata), nil
}

func TestMarshalStatus(t *testing.T) {
	status := Status{Message: "active", Service: "foo"}
	jsonTest, _ := json.Marshal(status)

	assert.Equal(t, `{"message":"active","service":"foo"}`, string(jsonTest))
}

func TestMarshalRequest(t *testing.T) {
	repo1 := Repo{Name: "foo", Tag: "latest"}
	repo2 := Repo{Name: "bar", Tag: "latest"}
	imageList := []Repo{repo1, repo2}
	repos := Request{Repos: imageList}
	jsonTest, _ := json.Marshal(repos)

	assert.Equal(t, `{"repos":[{"name":"foo","tag":"latest"},{"name":"bar","tag":"latest"}]}`, string(jsonTest))
}

func TestAnalyzeRequest(t *testing.T) {
	// setup
	fakeConn := new(mockConnection)
	api := newRegistryApi(fakeConn)
	layers := []string{"baz"}
	inBody, image := "{\"repos\":[{\"name\":\"foo\",\"tag\":\"latest\"}]}", "foo"

	// build request
	req, _ := http.NewRequest("POST", "http://localhost/analyze", strings.NewReader(inBody))
	w := httptest.NewRecorder()

	// build response
	metadata := new(registry.ImageMetadata)
	resp := make([]*registry.ImageMetadata, 1)
	resp[0] = metadata

	// test
	fakeConn.On("Connect", image).Return(nil)
	fakeConn.On("GetImageID", image, "latest").Return("fooID")
	fakeConn.On("GetAncestry", "fooID").Return(layers)
	fakeConn.On("GetMetadata", "baz").Return(metadata)
	api.handleAnalysis(w, req)

	// asserts
	fakeConn.AssertExpectations(t)
}

func TestStatusRequest(t *testing.T) {
	// setup
	fakeConn := new(mockConnection)
	api := newRegistryApi(fakeConn)

	// build request
	resp := Status{Message: "foo", Service: "bar"}
	req, _ := http.NewRequest("GET", "http://localhost/analyze", strings.NewReader("{}"))
	w := httptest.NewRecorder()

	// test
	fakeConn.On("Status").Return(resp, nil)
	api.handleStatus(w, req)

	// asserts
	fakeConn.AssertExpectations(t)
}
<commit_msg>fixed test<commit_after>package api

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/CenturyLinkLabs/docker-reg-client/registry"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
)

type mockConnection struct {
	mock.Mock
}

func (m *mockConnection) Status() (Status, error) {
	args := m.Mock.Called()
	return args.Get(0).(Status), nil
}

func (m *mockConnection) Connect(repo string) error {
	args := m.Mock.Called(repo)
	return args.Error(0)
}

func (m *mockConnection) GetImageID(image string, tag string) (string, error) {
	args := m.Mock.Called(image, tag)
	return args.Get(0).(string), nil
}

func (m *mockConnection) GetAncestry(id string) ([]string, error) {
	args := m.Mock.Called(id)
	return args.Get(0).([]string), nil
}

func (m *mockConnection) GetMetadata(layer string) (*registry.ImageMetadata, error) {
	args := m.Mock.Called(layer)
	return args.Get(0).(*registry.ImageMetadata), nil
}

func TestMarshalStatus(t *testing.T) {
	status := Status{Message: "active", Service: "foo"}
	jsonTest, _ := json.Marshal(status)

	assert.Equal(t, `{"message":"active","service":"foo"}`, string(jsonTest))
}

func TestMarshalRequest(t *testing.T) {
	repo1 := Repo{Name: "foo", Tag: "latest"}
	repo2 := Repo{Name: "bar", Tag: "latest"}
	imageList := []Repo{repo1, repo2}
	repos := Request{Repos: imageList}
	jsonTest, _ := json.Marshal(repos)

	assert.Equal(t, `{"repos":[{"name":"foo","tag":"latest","size":0},{"name":"bar","tag":"latest","size":0}]}`, string(jsonTest))
}

func TestAnalyzeRequest(t *testing.T) {
	// setup
	fakeConn := new(mockConnection)
	api := newRegistryApi(fakeConn)
	layers := []string{"baz"}
	inBody, image := "{\"repos\":[{\"name\":\"foo\",\"tag\":\"latest\"}]}", "foo"

	// build request
	req, _ := http.NewRequest("POST", "http://localhost/analyze", strings.NewReader(inBody))
	w := httptest.NewRecorder()

	// build response
	metadata := new(registry.ImageMetadata)
	resp := make([]*registry.ImageMetadata, 1)
	resp[0] = metadata

	// test
	fakeConn.On("Connect", image).Return(nil)
	fakeConn.On("GetImageID", image, "latest").Return("fooID")
	fakeConn.On("GetAncestry", "fooID").Return(layers)
	fakeConn.On("GetMetadata", "baz").Return(metadata)
	api.handleAnalysis(w, req)

	// asserts
	fakeConn.AssertExpectations(t)
}

func TestStatusRequest(t *testing.T) {
	// setup
	fakeConn := new(mockConnection)
	api := newRegistryApi(fakeConn)

	// build request
	resp := Status{Message: "foo", Service: "bar"}
	req, _ := http.NewRequest("GET", "http://localhost/analyze", strings.NewReader("{}"))
	w := httptest.NewRecorder()

	// test
	fakeConn.On("Status").Return(resp, nil)
	api.handleStatus(w, req)

	// asserts
	fakeConn.AssertExpectations(t)
}
<|endoftext|>


<commit_before>package main

import (
	"code.google.com/p/go.net/websocket"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"os/signal"
	"time"
)

const socksTimeout = 2
const bufSiz = 1500

func logDebug(format string, v ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", v...)
}

func proxy(local *net.TCPConn, ws *websocket.Conn) {
	var localToWs chan bool
	var wsToLocal chan bool

	// Local-to-WebSocket read loop.
	localToWs = make(chan bool, 1)
	go func() {
		buf := make([]byte, bufSiz)
		var err error
		for {
			n, er := local.Read(buf[:])
			if n > 0 {
				ew := websocket.Message.Send(ws, buf[:n])
				if ew != nil {
					err = ew
					break
				}
			}
			if er != nil {
				err = er
				break
			}
		}
		if err != nil && err != io.EOF {
			logDebug("%s", err)
		}
		local.CloseRead()
		ws.Close()

		localToWs <- true
	}()

	// WebSocket-to-local read loop.
	wsToLocal = make(chan bool, 1)
	go func() {
		var buf []byte
		var err error
		for {
			er := websocket.Message.Receive(ws, &buf)
			if er != nil {
				err = er
				break
			}
			n, ew := local.Write(buf)
			if ew != nil {
				err = ew
				break
			}
			if n != len(buf) {
				err = io.ErrShortWrite
				break
			}
		}
		if err != nil && err != io.EOF {
			logDebug("%s", err)
		}
		local.CloseWrite()
		ws.Close()

		wsToLocal <- true
	}()

	// Select twice, once for each read loop.
	select {
	case <-localToWs:
	case <-wsToLocal:
	}
	select {
	case <-localToWs:
	case <-wsToLocal:
	}
}

func handleConnection(conn *net.TCPConn) error {
	defer conn.Close()

	conn.SetDeadline(time.Now().Add(socksTimeout * time.Second))
	dest, err := readSocks4aConnect(conn)
	if err != nil {
		sendSocks4aResponseFailed(conn)
		return err
	}
	// Disable deadline.
	conn.SetDeadline(time.Time{})
	logDebug("SOCKS request for %s", dest)

	// We need the parsed IP and port for the SOCKS reply.
	destAddr, err := net.ResolveTCPAddr("tcp", dest)
	if err != nil {
		sendSocks4aResponseFailed(conn)
		return err
	}

	wsUrl := url.URL{Scheme: "ws", Host: dest}
	ws, err := websocket.Dial(wsUrl.String(), "", wsUrl.String())
	if err != nil {
		sendSocks4aResponseFailed(conn)
		return err
	}
	defer ws.Close()
	logDebug("WebSocket connection to %s", ws.Config().Location.String())

	sendSocks4aResponseGranted(conn, destAddr)

	proxy(conn, ws)

	return nil
}

func socksAcceptLoop(ln *net.TCPListener) error {
	for {
		socks, err := ln.AcceptTCP()
		if err != nil {
			return err
		}
		go func() {
			err := handleConnection(socks)
			if err != nil {
				logDebug("SOCKS from %s: %s", socks.RemoteAddr(), err)
			}
		}()
	}
	return nil
}

func startListener(addrStr string) (*net.TCPListener, error) {
	addr, err := net.ResolveTCPAddr("tcp", addrStr)
	if err != nil {
		return nil, err
	}
	ln, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return nil, err
	}
	go func() {
		err := socksAcceptLoop(ln)
		if err != nil {
			logDebug("accept: %s", err)
		}
	}()
	return ln, nil
}

func main() {
	const ptMethodName = "websocket"
	var socksAddrStrs = [...]string{"127.0.0.1:0", "[::1]:0"}

	ptClientSetup([]string{ptMethodName})

	for _, socksAddrStr := range socksAddrStrs {
		ln, err := startListener(socksAddrStr)
		if err != nil {
			ptCmethodError(ptMethodName, err.Error())
		}
		ptCmethod(ptMethodName, "socks4", ln.Addr())
	}
	ptCmethodsDone()

	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt)
	<-signalChan
}
<commit_msg>Do the proper thing on first and second SIGINT.<commit_after>package main

import (
	"code.google.com/p/go.net/websocket"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"os/signal"
	"time"
)

const socksTimeout = 2
const bufSiz = 1500

// When a connection handler starts, +1 is written to this channel; when it
// ends, -1 is written.
var handlerChan = make(chan int)

func logDebug(format string, v ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", v...)
}

func proxy(local *net.TCPConn, ws *websocket.Conn) {
	var localToWs chan bool
	var wsToLocal chan bool

	// Local-to-WebSocket read loop.
	localToWs = make(chan bool, 1)
	go func() {
		buf := make([]byte, bufSiz)
		var err error
		for {
			n, er := local.Read(buf[:])
			if n > 0 {
				ew := websocket.Message.Send(ws, buf[:n])
				if ew != nil {
					err = ew
					break
				}
			}
			if er != nil {
				err = er
				break
			}
		}
		if err != nil && err != io.EOF {
			logDebug("%s", err)
		}
		local.CloseRead()
		ws.Close()

		localToWs <- true
	}()

	// WebSocket-to-local read loop.
	wsToLocal = make(chan bool, 1)
	go func() {
		var buf []byte
		var err error
		for {
			er := websocket.Message.Receive(ws, &buf)
			if er != nil {
				err = er
				break
			}
			n, ew := local.Write(buf)
			if ew != nil {
				err = ew
				break
			}
			if n != len(buf) {
				err = io.ErrShortWrite
				break
			}
		}
		if err != nil && err != io.EOF {
			logDebug("%s", err)
		}
		local.CloseWrite()
		ws.Close()

		wsToLocal <- true
	}()

	// Select twice, once for each read loop.
	select {
	case <-localToWs:
	case <-wsToLocal:
	}
	select {
	case <-localToWs:
	case <-wsToLocal:
	}
}

func handleConnection(conn *net.TCPConn) error {
	defer conn.Close()

	handlerChan <- 1
	defer func() {
		handlerChan <- -1
	}()

	conn.SetDeadline(time.Now().Add(socksTimeout * time.Second))
	dest, err := readSocks4aConnect(conn)
	if err != nil {
		sendSocks4aResponseFailed(conn)
		return err
	}
	// Disable deadline.
	conn.SetDeadline(time.Time{})
	logDebug("SOCKS request for %s", dest)

	// We need the parsed IP and port for the SOCKS reply.
	destAddr, err := net.ResolveTCPAddr("tcp", dest)
	if err != nil {
		sendSocks4aResponseFailed(conn)
		return err
	}

	wsUrl := url.URL{Scheme: "ws", Host: dest}
	ws, err := websocket.Dial(wsUrl.String(), "", wsUrl.String())
	if err != nil {
		sendSocks4aResponseFailed(conn)
		return err
	}
	defer ws.Close()
	logDebug("WebSocket connection to %s", ws.Config().Location.String())

	sendSocks4aResponseGranted(conn, destAddr)

	proxy(conn, ws)

	return nil
}

func socksAcceptLoop(ln *net.TCPListener) error {
	for {
		socks, err := ln.AcceptTCP()
		if err != nil {
			return err
		}
		go func() {
			err := handleConnection(socks)
			if err != nil {
				logDebug("SOCKS from %s: %s", socks.RemoteAddr(), err)
			}
		}()
	}
	return nil
}

func startListener(addrStr string) (*net.TCPListener, error) {
	addr, err := net.ResolveTCPAddr("tcp", addrStr)
	if err != nil {
		return nil, err
	}
	ln, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return nil, err
	}
	go func() {
		err := socksAcceptLoop(ln)
		if err != nil {
			logDebug("accept: %s", err)
		}
	}()
	return ln, nil
}

func main() {
	const ptMethodName = "websocket"
	var socksAddrStrs = [...]string{"127.0.0.1:0", "[::1]:0"}

	ptClientSetup([]string{ptMethodName})

	listeners := make([]*net.TCPListener, 0)
	for _, socksAddrStr := range socksAddrStrs {
		ln, err := startListener(socksAddrStr)
		if err != nil {
			ptCmethodError(ptMethodName, err.Error())
		}
		ptCmethod(ptMethodName, "socks4", ln.Addr())
		listeners = append(listeners, ln)
	}
	ptCmethodsDone()

	var numHandlers int = 0

	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt)
	var sigint bool = false
	for !sigint {
		select {
		case n := <-handlerChan:
			numHandlers += n
		case <-signalChan:
			logDebug("SIGINT")
			sigint = true
		}
	}

	for _, ln := range listeners {
		ln.Close()
	}

	sigint = false
	for numHandlers != 0 && !sigint {
		select {
		case n := <-handlerChan:
			numHandlers += n
		case <-signalChan:
			logDebug("SIGINT")
			sigint = true
		}
	}
}
<|endoftext|>


<commit_before>package tcp

import (
	"encoding/binary"
	"errors"
	"io"
	"log"
	"net"
	"os"
	"sync"

	"github.com/influxdb/influxdb/data"
)

var (
	// ErrBindAddressRequired is returned when starting the Server
	// without a TCP or UDP listening address.
	ErrBindAddressRequired = errors.New("bind address required")

	// ErrServerClosed return when closing an already closed graphite server.
	ErrServerClosed = errors.New("server already closed")

	// ErrServerNotSpecified returned when Server is not specified.
	ErrServerNotSpecified = errors.New("server not present")
)

// Server processes data received over raw TCP connections.
type Server struct {
	writer   data.ShardWriter
	listener *net.Listener

	wg sync.WaitGroup

	Logger *log.Logger

	shutdown chan struct{}
}

// NewServer returns a new instance of a Server.
func NewServer(w data.ShardWriter) *Server {
	return &Server{
		writer:   w,
		Logger:   log.New(os.Stderr, "[tcp] ", log.LstdFlags),
		shutdown: make(chan struct{}),
	}
}

// ListenAndServe instructs the Server to start processing connections
// on the given interface. iface must be in the form host:port
// If successful, it returns the host as the first argument
func (s *Server) ListenAndServe(laddr string) (string, error) {
	if laddr == "" { // Make sure we have an laddr
		return "", ErrBindAddressRequired
	}

	ln, err := net.Listen("tcp", laddr)
	if err != nil {
		return "", err
	}
	s.listener = &ln

	s.Logger.Println("listening on TCP connection", ln.Addr().String())
	s.wg.Add(1)
	go func() {
		defer s.wg.Done()

		for {
			// Are we shutting down? If so, exit
			select {
			case <-s.shutdown:
				return
			default:
			}

			conn, err := ln.Accept()
			if opErr, ok := err.(*net.OpError); ok && !opErr.Temporary() {
				s.Logger.Println("error temporarily accepting TCP connection", err.Error())
				continue
			}
			if err != nil {
				s.Logger.Println("TCP listener closed")
				return
			}

			s.wg.Add(1)
			go s.handleConnection(conn)
		}
	}()

	// Return the host we started up on. Mostly needed for testing
	return ln.Addr().String(), nil
}

// Close will close the listener
func (s *Server) Close() error {
	// Stop accepting client connections
	if s.listener != nil {
		err := (*s.listener).Close()
		if err != nil {
			return err
		}
	} else {
		return ErrServerClosed
	}
	// Shut down all handlers
	close(s.shutdown)
	s.wg.Wait()
	s.listener = nil

	return nil
}

// handleConnection services an individual TCP connection.
func (s *Server) handleConnection(conn net.Conn) {
	defer func() {
		conn.Close()
		s.wg.Done()
	}()

	messageChannel := make(chan byte)

	// Start our reader up in a go routine so we don't block checking our close channel
	go func() {
		var messageType byte
		err := binary.Read(conn, binary.LittleEndian, &messageType)
		if err != nil {
			s.Logger.Printf("unable to read message type %s", err)
			return
		}
		messageChannel <- messageType
	}()

	for {
		select {
		case <-s.shutdown:
			// Are we shutting down? If so, exit
			return
		case messageType := <-messageChannel:
			switch messageType {
			case writeShardRequestMessage:
				if err := s.WriteShardRequest(conn); err != nil {
					s.WriteShardResponse(conn, err)
					return
				}
				s.WriteShardResponse(conn, nil)
			}
		default:
		}
	}

}

func (s *Server) WriteShardRequest(conn net.Conn) error {
	messageChannel := make(chan data.WriteShardRequest)
	errChan := make(chan error)

	go func() {
		var size int64
		if err := binary.Read(conn, binary.LittleEndian, &size); err != nil {
			errChan <- err
			return
		}

		message := make([]byte, size)

		reader := io.LimitReader(conn, size)
		_, err := reader.Read(message)
		if err != nil {
			errChan <- err
			return
		}
		var wsr data.WriteShardRequest
		if err := wsr.UnmarshalBinary(message); err != nil {
			errChan <- err
			return
		}
		messageChannel <- wsr
	}()

	for {
		select {
		case <-s.shutdown:
			// Are we shutting down? If so, exit
			return nil
		case e := <-errChan:
			return e
		case wsr := <-messageChannel:
			if _, err := s.writer.WriteShard(wsr.ShardID(), wsr.Points()); err != nil {
				return err
			}
			return nil
		default:
		}
	}
}

func (s *Server) WriteShardResponse(conn net.Conn, e error) {
	s.Logger.Println("writing shard response")
	var mt byte = writeShardResponseMessage
	if err := binary.Write(conn, binary.LittleEndian, &mt); err != nil {
		s.Logger.Printf("error writing shard response message type: %s", err)
		return
	}

	var wsr data.WriteShardResponse
	if e != nil {
		wsr.SetCode(1)
		wsr.SetMessage(e.Error())
	} else {
		wsr.SetCode(0)
	}

	b, err := wsr.MarshalBinary()
	if err != nil {
		s.Logger.Printf("error marshalling shard response: %s", err)
		return
	}

	size := int64(len(b))

	if err := binary.Write(conn, binary.LittleEndian, &size); err != nil {
		s.Logger.Printf("error writing shard response length: %s", err)
		return
	}

	if _, err := conn.Write(b); err != nil {
		s.Logger.Printf("error writing shard response: %s", err)
		return
	}
}
<commit_msg>removing unused error<commit_after>package tcp

import (
	"encoding/binary"
	"errors"
	"io"
	"log"
	"net"
	"os"
	"sync"

	"github.com/influxdb/influxdb/data"
)

var (
	// ErrBindAddressRequired is returned when starting the Server
	// without a TCP or UDP listening address.
	ErrBindAddressRequired = errors.New("bind address required")

	// ErrServerClosed return when closing an already closed graphite server.
	ErrServerClosed = errors.New("server already closed")
)

// Server processes data received over raw TCP connections.
type Server struct {
	writer   data.ShardWriter
	listener *net.Listener

	wg sync.WaitGroup

	Logger *log.Logger

	shutdown chan struct{}
}

// NewServer returns a new instance of a Server.
func NewServer(w data.ShardWriter) *Server {
	return &Server{
		writer:   w,
		Logger:   log.New(os.Stderr, "[tcp] ", log.LstdFlags),
		shutdown: make(chan struct{}),
	}
}

// ListenAndServe instructs the Server to start processing connections
// on the given interface. iface must be in the form host:port
// If successful, it returns the host as the first argument
func (s *Server) ListenAndServe(laddr string) (string, error) {
	if laddr == "" { // Make sure we have an laddr
		return "", ErrBindAddressRequired
	}

	ln, err := net.Listen("tcp", laddr)
	if err != nil {
		return "", err
	}
	s.listener = &ln

	s.Logger.Println("listening on TCP connection", ln.Addr().String())
	s.wg.Add(1)
	go func() {
		defer s.wg.Done()

		for {
			// Are we shutting down? If so, exit
			select {
			case <-s.shutdown:
				return
			default:
			}

			conn, err := ln.Accept()
			if opErr, ok := err.(*net.OpError); ok && !opErr.Temporary() {
				s.Logger.Println("error temporarily accepting TCP connection", err.Error())
				continue
			}
			if err != nil {
				s.Logger.Println("TCP listener closed")
				return
			}

			s.wg.Add(1)
			go s.handleConnection(conn)
		}
	}()

	// Return the host we started up on. Mostly needed for testing
	return ln.Addr().String(), nil
}

// Close will close the listener
func (s *Server) Close() error {
	// Stop accepting client connections
	if s.listener != nil {
		err := (*s.listener).Close()
		if err != nil {
			return err
		}
	} else {
		return ErrServerClosed
	}
	// Shut down all handlers
	close(s.shutdown)
	s.wg.Wait()
	s.listener = nil

	return nil
}

// handleConnection services an individual TCP connection.
func (s *Server) handleConnection(conn net.Conn) {
	defer func() {
		conn.Close()
		s.wg.Done()
	}()

	messageChannel := make(chan byte)

	// Start our reader up in a go routine so we don't block checking our close channel
	go func() {
		var messageType byte
		err := binary.Read(conn, binary.LittleEndian, &messageType)
		if err != nil {
			s.Logger.Printf("unable to read message type %s", err)
			return
		}
		messageChannel <- messageType
	}()

	for {
		select {
		case <-s.shutdown:
			// Are we shutting down? If so, exit
			return
		case messageType := <-messageChannel:
			switch messageType {
			case writeShardRequestMessage:
				if err := s.WriteShardRequest(conn); err != nil {
					s.WriteShardResponse(conn, err)
					return
				}
				s.WriteShardResponse(conn, nil)
			}
		default:
		}
	}

}

func (s *Server) WriteShardRequest(conn net.Conn) error {
	messageChannel := make(chan data.WriteShardRequest)
	errChan := make(chan error)

	go func() {
		var size int64
		if err := binary.Read(conn, binary.LittleEndian, &size); err != nil {
			errChan <- err
			return
		}

		message := make([]byte, size)

		reader := io.LimitReader(conn, size)
		_, err := reader.Read(message)
		if err != nil {
			errChan <- err
			return
		}
		var wsr data.WriteShardRequest
		if err := wsr.UnmarshalBinary(message); err != nil {
			errChan <- err
			return
		}
		messageChannel <- wsr
	}()

	for {
		select {
		case <-s.shutdown:
			// Are we shutting down? If so, exit
			return nil
		case e := <-errChan:
			return e
		case wsr := <-messageChannel:
			if _, err := s.writer.WriteShard(wsr.ShardID(), wsr.Points()); err != nil {
				return err
			}
			return nil
		default:
		}
	}
}

func (s *Server) WriteShardResponse(conn net.Conn, e error) {
	s.Logger.Println("writing shard response")
	var mt byte = writeShardResponseMessage
	if err := binary.Write(conn, binary.LittleEndian, &mt); err != nil {
		s.Logger.Printf("error writing shard response message type: %s", err)
		return
	}

	var wsr data.WriteShardResponse
	if e != nil {
		wsr.SetCode(1)
		wsr.SetMessage(e.Error())
	} else {
		wsr.SetCode(0)
	}

	b, err := wsr.MarshalBinary()
	if err != nil {
		s.Logger.Printf("error marshalling shard response: %s", err)
		return
	}

	size := int64(len(b))

	if err := binary.Write(conn, binary.LittleEndian, &size); err != nil {
		s.Logger.Printf("error writing shard response length: %s", err)
		return
	}

	if _, err := conn.Write(b); err != nil {
		s.Logger.Printf("error writing shard response: %s", err)
		return
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// +k8s:deepcopy-gen=package
// +k8s:conversion-gen=k8s.io/kube-aggregator/pkg/apis/apiregistration
// +k8s:openapi-gen=true
// +groupName=apiregistration.k8s.io

// Package v1beta1 contains the API Registration API, which is responsible for
// registering an API `Group`/`Version` with another kubernetes like API server.
// The `APIService` holds information about the other API server in
// `APIServiceSpec` type as well as general `TypeMeta` and `ObjectMeta`. The
// `APIServiceSpec` type have the main configuration needed to do the
// aggregation. Any request coming for specified `Group`/`Version` will be
// directed to the service defined by `ServiceReference` (on port 443) after
// validating the target using provided `CABundle` or skipping validation
// if development flag `InsecureSkipTLSVerify` is set. `Priority` is controlling
// the order of this API group in the overall discovery document.
// The return status is a set of conditions for this aggregation. Currently
// there is only one condition named "Available", if true, it means the
// api/server requests will be redirected to specified API server.
package v1 // import "k8s.io/kube-aggregator/pkg/apis/apiregistration/v1"
<commit_msg>update comments of doc.go in stagging/src/k8s.io/kube-aggregator<commit_after>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// +k8s:deepcopy-gen=package
// +k8s:conversion-gen=k8s.io/kube-aggregator/pkg/apis/apiregistration
// +k8s:openapi-gen=true
// +groupName=apiregistration.k8s.io

// Package v1 contains the API Registration API, which is responsible for
// registering an API `Group`/`Version` with another kubernetes like API server.
// The `APIService` holds information about the other API server in
// `APIServiceSpec` type as well as general `TypeMeta` and `ObjectMeta`. The
// `APIServiceSpec` type have the main configuration needed to do the
// aggregation. Any request coming for specified `Group`/`Version` will be
// directed to the service defined by `ServiceReference` (on port 443) after
// validating the target using provided `CABundle` or skipping validation
// if development flag `InsecureSkipTLSVerify` is set. `Priority` is controlling
// the order of this API group in the overall discovery document.
// The return status is a set of conditions for this aggregation. Currently
// there is only one condition named "Available", if true, it means the
// api/server requests will be redirected to specified API server.
package v1 // import "k8s.io/kube-aggregator/pkg/apis/apiregistration/v1"
<|endoftext|>


<commit_before>package ddevapp_test

import (
	"fmt"
	"github.com/drud/ddev/pkg/exec"
	"github.com/drud/ddev/pkg/globalconfig"
	"github.com/drud/ddev/pkg/nodeps"
	"github.com/stretchr/testify/require"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"testing"

	. "github.com/drud/ddev/pkg/ddevapp"
	"github.com/drud/ddev/pkg/testcommon"
	asrt "github.com/stretchr/testify/assert"
)

/**
 * A valid site (with backups) must be present which matches the test site and environment name
 * defined in the constants below.
 */
const acquiaPullTestSite = "eeamoreno.dev"
const acquiaPushTestSite = "eeamoreno.stg"

const acquiaPullSiteURL = "http://eeamorenodev.prod.acquia-sites.com/"
const acquiaSiteExpectation = "Super easy vegetarian pasta"

// TestAcquiaPull ensures we can pull backups from Acquia
func TestAcquiaPull(t *testing.T) {
	acquiaKey := ""
	acquiaSecret := ""
	sshkey := ""
	if acquiaKey = os.Getenv("DDEV_ACQUIA_API_KEY"); acquiaKey == "" {
		t.Skipf("No DDEV_ACQUIA_API_KEY env var has been set. Skipping %v", t.Name())
	}
	if acquiaSecret = os.Getenv("DDEV_ACQUIA_API_SECRET"); acquiaSecret == "" {
		t.Skipf("No DDEV_ACQUIA_SECRET env var has been set. Skipping %v", t.Name())
	}
	if sshkey = os.Getenv("DDEV_ACQUIA_SSH_KEY"); sshkey == "" {
		t.Skipf("No DDEV_ACQUIA_SSH_KEY env var has been set. Skipping %v", t.Name())
	}
	sshkey = strings.Replace(sshkey, "<SPLIT>", "\n", -1)

	require.True(t, isPullSiteValid(acquiaPullSiteURL, acquiaSiteExpectation), "acquiaPullSiteURL %s isn't working right", acquiaPullSiteURL)
	// Set up tests and give ourselves a working directory.
	assert := asrt.New(t)
	origDir, _ := os.Getwd()

	webEnvSave := globalconfig.DdevGlobalConfig.WebEnvironment

	globalconfig.DdevGlobalConfig.WebEnvironment = []string{"ACQUIA_API_KEY=" + acquiaKey, "ACQUIA_API_SECRET=" + acquiaSecret}
	err := globalconfig.WriteGlobalConfig(globalconfig.DdevGlobalConfig)
	assert.NoError(err)

	siteDir := testcommon.CreateTmpDir(t.Name())
	err = os.MkdirAll(filepath.Join(siteDir, "docroot/sites/default"), 0777)
	assert.NoError(err)
	err = os.Chdir(siteDir)
	assert.NoError(err)

	err = setupSSHKey(t, sshkey, filepath.Join(origDir, "testdata", t.Name()))
	require.NoError(t, err)

	app, err := NewApp(siteDir, true)
	assert.NoError(err)

	t.Cleanup(func() {
		err = app.Stop(true, false)
		assert.NoError(err)

		globalconfig.DdevGlobalConfig.WebEnvironment = webEnvSave
		err = globalconfig.WriteGlobalConfig(globalconfig.DdevGlobalConfig)
		assert.NoError(err)

		_ = os.Chdir(origDir)
		err = os.RemoveAll(siteDir)
		assert.NoError(err)
	})

	app.Name = t.Name()
	app.Type = nodeps.AppTypeDrupal9

	_ = app.Stop(true, false)
	err = app.WriteConfig()
	assert.NoError(err)

	testcommon.ClearDockerEnv()

	err = PopulateExamplesCommandsHomeadditions(app.Name)
	require.NoError(t, err)

	err = app.Start()
	require.NoError(t, err)

	// Make sure we have drush
	_, _, err = app.Exec(&ExecOpts{
		Cmd: "composer require --no-interaction drush/drush:* >/dev/null 2>/dev/null",
	})
	require.NoError(t, err)

	// Build our acquia.yaml from the example file
	s, err := os.ReadFile(app.GetConfigPath("providers/acquia.yaml.example"))
	require.NoError(t, err)
	x := strings.Replace(string(s), "project_id:", fmt.Sprintf("project_id: %s\n#project_id:", acquiaPullTestSite), -1)
	err = os.WriteFile(app.GetConfigPath("providers/acquia.yaml"), []byte(x), 0666)
	assert.NoError(err)
	err = app.WriteConfig()
	require.NoError(t, err)
	err = app.MutagenSyncFlush()
	require.NoError(t, err)

	provider, err := app.GetProvider("acquia")
	require.NoError(t, err)
	err = app.Pull(provider, false, false, false)
	require.NoError(t, err)

	assert.FileExists(filepath.Join(app.GetHostUploadDirFullPath(), "chocolate-brownie-umami.jpg"))
	out, err := exec.RunCommand("bash", []string{"-c", fmt.Sprintf(`echo 'select COUNT(*) from users_field_data where mail="randy@example.com";' | %s mysql -N`, DdevBin)})
	assert.NoError(err)
	assert.True(strings.HasPrefix(out, "1\n"))
}

// TestAcquiaPush ensures we can push to acquia for a configured environment.
func TestAcquiaPush(t *testing.T) {
	acquiaKey := ""
	acquiaSecret := ""
	sshkey := ""
	if os.Getenv("DDEV_ALLOW_ACQUIA_PUSH") != "true" {
		t.Skip("TestAcquiaPush is currently embargoed by DDEV_ALLOW_ACQUIA_PUSH not set to true")
	}
	if acquiaKey = os.Getenv("DDEV_ACQUIA_API_KEY"); acquiaKey == "" {
		t.Skipf("No DDEV_ACQUIA_API_KEY env var has been set. Skipping %v", t.Name())
	}
	if acquiaSecret = os.Getenv("DDEV_ACQUIA_API_SECRET"); acquiaSecret == "" {
		t.Skipf("No DDEV_ACQUIA_API_SECRET env var has been set. Skipping %v", t.Name())
	}
	if sshkey = os.Getenv("DDEV_ACQUIA_SSH_KEY"); sshkey == "" {
		t.Skipf("No DDEV_ACQUIA_SSH_KEY env var has been set. Skipping %v", t.Name())
	}
	sshkey = strings.Replace(sshkey, "<SPLIT>", "\n", -1)

	// Set up tests and give ourselves a working directory.
	assert := asrt.New(t)
	origDir, _ := os.Getwd()

	webEnvSave := globalconfig.DdevGlobalConfig.WebEnvironment

	globalconfig.DdevGlobalConfig.WebEnvironment = []string{"ACQUIA_API_KEY=" + acquiaKey, "ACQUIA_API_SECRET=" + acquiaSecret}
	err := globalconfig.WriteGlobalConfig(globalconfig.DdevGlobalConfig)
	assert.NoError(err)

	// Use a D9 codebase for drush to work right
	d9code := FullTestSites[8]
	d9code.Name = t.Name()
	err = globalconfig.RemoveProjectInfo(t.Name())
	require.NoError(t, err)
	err = d9code.Prepare()
	require.NoError(t, err)
	app, err := NewApp(d9code.Dir, false)
	require.NoError(t, err)
	_ = app.Stop(true, false)

	err = os.Chdir(d9code.Dir)
	require.NoError(t, err)

	err = setupSSHKey(t, sshkey, filepath.Join(origDir, "testdata", t.Name()))
	require.NoError(t, err)

	t.Cleanup(func() {
		err = app.Stop(true, false)
		assert.NoError(err)

		globalconfig.DdevGlobalConfig.WebEnvironment = webEnvSave
		err = globalconfig.WriteGlobalConfig(globalconfig.DdevGlobalConfig)
		assert.NoError(err)

		_ = os.Chdir(origDir)
		err = os.RemoveAll(d9code.Dir)
		assert.NoError(err)
	})

	app.Hooks = map[string][]YAMLTask{"post-push": {{"exec-host": "touch hello-post-push-" + app.Name}}, "pre-push": {{"exec-host": "touch hello-pre-push-" + app.Name}}}
	_ = app.Stop(true, false)

	err = app.WriteConfig()
	require.NoError(t, err)

	testcommon.ClearDockerEnv()

	err = PopulateExamplesCommandsHomeadditions(app.Name)
	require.NoError(t, err)

	err = app.Start()
	require.NoError(t, err)

	// Make sure we have drush
	_, _, err = app.Exec(&ExecOpts{
		Cmd: "composer require --no-interaction drush/drush:* >/dev/null 2>/dev/null",
	})
	require.NoError(t, err)

	_, _, err = app.Exec(&ExecOpts{
		Cmd: "time drush si -y minimal",
	})
	require.NoError(t, err)

	// Create database and files entries that we can verify after push
	tval := nodeps.RandomString(10)
	writeQuery := fmt.Sprintf(`mysql -e 'CREATE TABLE IF NOT EXISTS %s ( title VARCHAR(255) NOT NULL ); INSERT INTO %s VALUES("%s");'`, t.Name(), t.Name(), tval)
	_, _, err = app.Exec(&ExecOpts{
		Cmd: writeQuery,
	})
	require.NoError(t, err)

	fName := tval + ".txt"
	fContent := []byte(tval)
	err = os.WriteFile(filepath.Join(d9code.Dir, "sites/default/files", fName), fContent, 0644)
	assert.NoError(err)

	// Build our PUSH acquia.yaml from the example file
	s, err := os.ReadFile(app.GetConfigPath("providers/acquia.yaml.example"))
	require.NoError(t, err)
	x := strings.Replace(string(s), "project_id:", fmt.Sprintf("project_id: %s\n#project_id:", acquiaPushTestSite), -1)
	err = os.WriteFile(app.GetConfigPath("providers/acquia.yaml"), []byte(x), 0666)
	assert.NoError(err)
	err = app.WriteConfig()
	require.NoError(t, err)

	provider, err := app.GetProvider("acquia")
	require.NoError(t, err)

	err = app.Push(provider, false, false)
	require.NoError(t, err)

	// Test that the database row was added
	readQuery := fmt.Sprintf(`echo 'SELECT title FROM %s WHERE title="%s"' | drush @%s --alias-path=~/.drush sql-cli --extra=-N`, t.Name(), tval, acquiaPushTestSite)
	out, _, err := app.Exec(&ExecOpts{
		Cmd: readQuery,
	})
	require.NoError(t, err)
	assert.Contains(out, tval)

	// Test that the file arrived there (by rsyncing it back)
	out, _, err = app.Exec(&ExecOpts{
		Cmd: fmt.Sprintf("drush --alias-path=~/.drush rsync -y @%s:%%files/%s /tmp && cat /tmp/%s", acquiaPushTestSite, fName, fName),
	})
	require.NoError(t, err)
	assert.Contains(out, tval)

	err = app.MutagenSyncFlush()
	assert.NoError(err)

	assert.FileExists("hello-pre-push-" + app.Name)
	assert.FileExists("hello-post-push-" + app.Name)
	err = os.Remove("hello-pre-push-" + app.Name)
	assert.NoError(err)
	err = os.Remove("hello-post-push-" + app.Name)
	assert.NoError(err)
}

// isPullSiteValid just checks to make sure the site we're testing against is OK
func isPullSiteValid(siteURL string, siteExpectation string) bool {
	resp, err := http.Get(siteURL)
	if err != nil {
		return false
	}
	//nolint: errcheck
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return false
	}
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return false
	}
	return strings.Contains(string(body), siteExpectation)
}
<commit_msg>Make Acquia tests work again (#3662)<commit_after>package ddevapp_test

import (
	"fmt"
	"github.com/drud/ddev/pkg/exec"
	"github.com/drud/ddev/pkg/globalconfig"
	"github.com/drud/ddev/pkg/nodeps"
	"github.com/stretchr/testify/require"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"testing"

	. "github.com/drud/ddev/pkg/ddevapp"
	"github.com/drud/ddev/pkg/testcommon"
	asrt "github.com/stretchr/testify/assert"
)

/**
 * A valid site (with backups) must be present which matches the test site and environment name
 * defined in the constants below.
 */
const acquiaPullTestSite = "ddevdemo.dev"
const acquiaPushTestSite = "ddevdemo.test"

const acquiaPullSiteURL = "http://ddevdemodev.prod.acquia-sites.com/"
const acquiaSiteExpectation = "Super easy vegetarian pasta"

// TestAcquiaPull ensures we can pull backups from Acquia
func TestAcquiaPull(t *testing.T) {
	acquiaKey := ""
	acquiaSecret := ""
	sshkey := ""
	if acquiaKey = os.Getenv("DDEV_ACQUIA_API_KEY"); acquiaKey == "" {
		t.Skipf("No DDEV_ACQUIA_API_KEY env var has been set. Skipping %v", t.Name())
	}
	if acquiaSecret = os.Getenv("DDEV_ACQUIA_API_SECRET"); acquiaSecret == "" {
		t.Skipf("No DDEV_ACQUIA_SECRET env var has been set. Skipping %v", t.Name())
	}
	if sshkey = os.Getenv("DDEV_ACQUIA_SSH_KEY"); sshkey == "" {
		t.Skipf("No DDEV_ACQUIA_SSH_KEY env var has been set. Skipping %v", t.Name())
	}
	sshkey = strings.Replace(sshkey, "<SPLIT>", "\n", -1)

	require.True(t, isPullSiteValid(acquiaPullSiteURL, acquiaSiteExpectation), "acquiaPullSiteURL %s isn't working right", acquiaPullSiteURL)
	// Set up tests and give ourselves a working directory.
	assert := asrt.New(t)
	origDir, _ := os.Getwd()

	webEnvSave := globalconfig.DdevGlobalConfig.WebEnvironment

	globalconfig.DdevGlobalConfig.WebEnvironment = []string{"ACQUIA_API_KEY=" + acquiaKey, "ACQUIA_API_SECRET=" + acquiaSecret}
	err := globalconfig.WriteGlobalConfig(globalconfig.DdevGlobalConfig)
	assert.NoError(err)

	siteDir := testcommon.CreateTmpDir(t.Name())
	err = os.MkdirAll(filepath.Join(siteDir, "docroot/sites/default"), 0777)
	assert.NoError(err)
	err = os.Chdir(siteDir)
	assert.NoError(err)

	err = setupSSHKey(t, sshkey, filepath.Join(origDir, "testdata", t.Name()))
	require.NoError(t, err)

	app, err := NewApp(siteDir, true)
	assert.NoError(err)

	t.Cleanup(func() {
		err = app.Stop(true, false)
		assert.NoError(err)

		globalconfig.DdevGlobalConfig.WebEnvironment = webEnvSave
		err = globalconfig.WriteGlobalConfig(globalconfig.DdevGlobalConfig)
		assert.NoError(err)

		_ = os.Chdir(origDir)
		err = os.RemoveAll(siteDir)
		assert.NoError(err)
	})

	app.Name = t.Name()
	app.Type = nodeps.AppTypeDrupal9

	_ = app.Stop(true, false)
	err = app.WriteConfig()
	assert.NoError(err)

	testcommon.ClearDockerEnv()

	err = PopulateExamplesCommandsHomeadditions(app.Name)
	require.NoError(t, err)

	err = app.Start()
	require.NoError(t, err)

	// Make sure we have drush
	_, _, err = app.Exec(&ExecOpts{
		Cmd: "composer require --no-interaction drush/drush:* >/dev/null 2>/dev/null",
	})
	require.NoError(t, err)

	// Build our acquia.yaml from the example file
	s, err := os.ReadFile(app.GetConfigPath("providers/acquia.yaml.example"))
	require.NoError(t, err)
	x := strings.Replace(string(s), "project_id:", fmt.Sprintf("project_id: %s\n#project_id:", acquiaPullTestSite), -1)
	err = os.WriteFile(app.GetConfigPath("providers/acquia.yaml"), []byte(x), 0666)
	assert.NoError(err)
	err = app.WriteConfig()
	require.NoError(t, err)
	err = app.MutagenSyncFlush()
	require.NoError(t, err)

	provider, err := app.GetProvider("acquia")
	require.NoError(t, err)
	err = app.Pull(provider, false, false, false)
	require.NoError(t, err)

	assert.FileExists(filepath.Join(app.GetHostUploadDirFullPath(), "chocolate-brownie-umami.jpg"))
	out, err := exec.RunCommand("bash", []string{"-c", fmt.Sprintf(`echo 'select COUNT(*) from users_field_data where mail="randy@example.com";' | %s mysql -N`, DdevBin)})
	assert.NoError(err)
	assert.True(strings.HasPrefix(out, "1\n"))
}

// TestAcquiaPush ensures we can push to acquia for a configured environment.
func TestAcquiaPush(t *testing.T) {
	acquiaKey := ""
	acquiaSecret := ""
	sshkey := ""
	if os.Getenv("DDEV_ALLOW_ACQUIA_PUSH") != "true" {
		t.Skip("TestAcquiaPush is currently embargoed by DDEV_ALLOW_ACQUIA_PUSH not set to true")
	}
	if acquiaKey = os.Getenv("DDEV_ACQUIA_API_KEY"); acquiaKey == "" {
		t.Skipf("No DDEV_ACQUIA_API_KEY env var has been set. Skipping %v", t.Name())
	}
	if acquiaSecret = os.Getenv("DDEV_ACQUIA_API_SECRET"); acquiaSecret == "" {
		t.Skipf("No DDEV_ACQUIA_API_SECRET env var has been set. Skipping %v", t.Name())
	}
	if sshkey = os.Getenv("DDEV_ACQUIA_SSH_KEY"); sshkey == "" {
		t.Skipf("No DDEV_ACQUIA_SSH_KEY env var has been set. Skipping %v", t.Name())
	}
	sshkey = strings.Replace(sshkey, "<SPLIT>", "\n", -1)

	// Set up tests and give ourselves a working directory.
	assert := asrt.New(t)
	origDir, _ := os.Getwd()

	webEnvSave := globalconfig.DdevGlobalConfig.WebEnvironment

	globalconfig.DdevGlobalConfig.WebEnvironment = []string{"ACQUIA_API_KEY=" + acquiaKey, "ACQUIA_API_SECRET=" + acquiaSecret}
	err := globalconfig.WriteGlobalConfig(globalconfig.DdevGlobalConfig)
	assert.NoError(err)

	// Use a D9 codebase for drush to work right
	d9code := FullTestSites[8]
	d9code.Name = t.Name()
	err = globalconfig.RemoveProjectInfo(t.Name())
	require.NoError(t, err)
	err = d9code.Prepare()
	require.NoError(t, err)
	app, err := NewApp(d9code.Dir, false)
	require.NoError(t, err)
	_ = app.Stop(true, false)

	err = os.Chdir(d9code.Dir)
	require.NoError(t, err)

	err = setupSSHKey(t, sshkey, filepath.Join(origDir, "testdata", t.Name()))
	require.NoError(t, err)

	t.Cleanup(func() {
		err = app.Stop(true, false)
		assert.NoError(err)

		globalconfig.DdevGlobalConfig.WebEnvironment = webEnvSave
		err = globalconfig.WriteGlobalConfig(globalconfig.DdevGlobalConfig)
		assert.NoError(err)

		_ = os.Chdir(origDir)
		err = os.RemoveAll(d9code.Dir)
		assert.NoError(err)
	})

	app.Hooks = map[string][]YAMLTask{"post-push": {{"exec-host": "touch hello-post-push-" + app.Name}}, "pre-push": {{"exec-host": "touch hello-pre-push-" + app.Name}}}
	_ = app.Stop(true, false)

	err = app.WriteConfig()
	require.NoError(t, err)

	testcommon.ClearDockerEnv()

	err = PopulateExamplesCommandsHomeadditions(app.Name)
	require.NoError(t, err)

	err = app.Start()
	require.NoError(t, err)

	// Make sure we have drush
	_, _, err = app.Exec(&ExecOpts{
		Cmd: "composer require --no-interaction drush/drush:* >/dev/null 2>/dev/null",
	})
	require.NoError(t, err)

	_, _, err = app.Exec(&ExecOpts{
		Cmd: "time drush si -y minimal",
	})
	require.NoError(t, err)

	// Create database and files entries that we can verify after push
	tval := nodeps.RandomString(10)
	writeQuery := fmt.Sprintf(`mysql -e 'CREATE TABLE IF NOT EXISTS %s ( title VARCHAR(255) NOT NULL ); INSERT INTO %s VALUES("%s");'`, t.Name(), t.Name(), tval)
	_, _, err = app.Exec(&ExecOpts{
		Cmd: writeQuery,
	})
	require.NoError(t, err)

	fName := tval + ".txt"
	fContent := []byte(tval)
	err = os.WriteFile(filepath.Join(d9code.Dir, "sites/default/files", fName), fContent, 0644)
	assert.NoError(err)

	// Build our PUSH acquia.yaml from the example file
	s, err := os.ReadFile(app.GetConfigPath("providers/acquia.yaml.example"))
	require.NoError(t, err)
	x := strings.Replace(string(s), "project_id:", fmt.Sprintf("project_id: %s\n#project_id:", acquiaPushTestSite), -1)
	err = os.WriteFile(app.GetConfigPath("providers/acquia.yaml"), []byte(x), 0666)
	assert.NoError(err)
	err = app.WriteConfig()
	require.NoError(t, err)

	provider, err := app.GetProvider("acquia")
	require.NoError(t, err)

	err = app.Push(provider, false, false)
	require.NoError(t, err)

	// Test that the database row was added
	readQuery := fmt.Sprintf(`echo 'SELECT title FROM %s WHERE title="%s"' | drush @%s --alias-path=~/.drush sql-cli --extra=-N`, t.Name(), tval, acquiaPushTestSite)
	out, _, err := app.Exec(&ExecOpts{
		Cmd: readQuery,
	})
	require.NoError(t, err)
	assert.Contains(out, tval)

	// Test that the file arrived there (by rsyncing it back)
	out, _, err = app.Exec(&ExecOpts{
		Cmd: fmt.Sprintf("drush --alias-path=~/.drush rsync -y @%s:%%files/%s /tmp && cat /tmp/%s", acquiaPushTestSite, fName, fName),
	})
	require.NoError(t, err)
	assert.Contains(out, tval)

	err = app.MutagenSyncFlush()
	assert.NoError(err)

	assert.FileExists("hello-pre-push-" + app.Name)
	assert.FileExists("hello-post-push-" + app.Name)
	err = os.Remove("hello-pre-push-" + app.Name)
	assert.NoError(err)
	err = os.Remove("hello-post-push-" + app.Name)
	assert.NoError(err)
}

// isPullSiteValid just checks to make sure the site we're testing against is OK
func isPullSiteValid(siteURL string, siteExpectation string) bool {
	resp, err := http.Get(siteURL)
	if err != nil {
		return false
	}
	//nolint: errcheck
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return false
	}
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return false
	}
	return strings.Contains(string(body), siteExpectation)
}
<|endoftext|>


<commit_before>package builder

import (
	"github.com/Aptomi/aptomi/pkg/external"
	"github.com/Aptomi/aptomi/pkg/external/secrets"
	"github.com/Aptomi/aptomi/pkg/external/users"
	"github.com/Aptomi/aptomi/pkg/lang"
	"github.com/Aptomi/aptomi/pkg/object"
	"github.com/Aptomi/aptomi/pkg/util"
	"math/rand"
)

var randSeed = int64(239)
var idLength = 16

// PolicyBuilder is a utility struct to help build policy objects
// It is primarily used in unit tests
type PolicyBuilder struct {
	random    *rand.Rand
	namespace string
	policy    *lang.Policy
	users     *users.UserLoaderMock
	secrets   *secrets.SecretLoaderMock

	domainAdmin     *lang.User
	domainAdminView *lang.PolicyView
}

// NewPolicyBuilder creates a new PolicyBuilder with a default "main" namespace
func NewPolicyBuilder() *PolicyBuilder {
	return NewPolicyBuilderWithNS("main")
}

// NewPolicyBuilderWithNS creates a new PolicyBuilder
func NewPolicyBuilderWithNS(namespace string) *PolicyBuilder {
	result := &PolicyBuilder{
		random:    rand.New(rand.NewSource(randSeed)),
		namespace: namespace,
		policy:    lang.NewPolicy(),
		users:     users.NewUserLoaderMock(),
		secrets:   secrets.NewSecretLoaderMock(),
	}

	for _, rule := range lang.ACLRulesBootstrap {
		err := result.policy.AddObject(rule)
		if err != nil {
			panic(err)
		}
	}

	result.domainAdmin = result.AddUserDomainAdmin()
	result.domainAdminView = result.policy.View(result.domainAdmin)

	return result
}

// SwitchNamespace switches the current namespace where objects will be generated
func (builder *PolicyBuilder) SwitchNamespace(namespace string) {
	builder.namespace = namespace
}

// AddDependency creates a new dependency and adds it to the policy
func (builder *PolicyBuilder) AddDependency(user *lang.User, contract *lang.Contract) *lang.Dependency {
	result := &lang.Dependency{
		Metadata: lang.Metadata{
			Kind:      lang.DependencyObject.Kind,
			Namespace: builder.namespace,
			Name:      util.RandomID(builder.random, idLength),
		},
		UserID:   user.ID,
		Contract: contract.Namespace + "/" + contract.Name,
		Labels:   make(map[string]string),
	}

	builder.addObject(builder.domainAdminView, result)
	return result
}

// AddUser creates a new user who can consume services from the 'main' namespace and adds it to the policy
func (builder *PolicyBuilder) AddUser() *lang.User {
	result := &lang.User{
		ID:     util.RandomID(builder.random, idLength),
		Name:   util.RandomID(builder.random, idLength),
		Labels: map[string]string{"role": "aptomi_main_ns_consumer"},
	}
	builder.users.AddUser(result)
	return result
}

// AddUserDomainAdmin creates a new user who is a domain admin and adds it to the policy
func (builder *PolicyBuilder) AddUserDomainAdmin() *lang.User {
	result := builder.AddUser()
	result.Labels["role"] = "aptomi_domain_admin"
	return result
}

// AddService creates a new service and adds it to the policy
func (builder *PolicyBuilder) AddService(owner *lang.User) *lang.Service {
	var ownerID = ""
	if owner != nil {
		ownerID = owner.ID
	}
	result := &lang.Service{
		Metadata: lang.Metadata{
			Kind:      lang.ServiceObject.Kind,
			Namespace: builder.namespace,
			Name:      util.RandomID(builder.random, idLength),
		},
		Owner: ownerID,
	}
	builder.addObject(builder.domainAdminView, result)
	return result
}

// AddContract creates a new contract for a given service and adds it to the policy
func (builder *PolicyBuilder) AddContract(service *lang.Service, criteria *lang.Criteria) *lang.Contract {
	result := &lang.Contract{
		Metadata: lang.Metadata{
			Kind:      lang.ContractObject.Kind,
			Namespace: builder.namespace,
			Name:      util.RandomID(builder.random, idLength),
		},
		Contexts: []*lang.Context{{
			Name:     util.RandomID(builder.random, idLength),
			Criteria: criteria,
			Allocation: &lang.Allocation{
				Service: service.Name,
			},
		}},
	}
	builder.addObject(builder.domainAdminView, result)
	return result
}

// AddContractMultipleContexts creates contract with multiple contexts for a given service and adds it to the policy
func (builder *PolicyBuilder) AddContractMultipleContexts(service *lang.Service, criteriaArray ...*lang.Criteria) *lang.Contract {
	result := &lang.Contract{
		Metadata: lang.Metadata{
			Kind:      lang.ContractObject.Kind,
			Namespace: builder.namespace,
			Name:      util.RandomID(builder.random, idLength),
		},
	}
	for _, criteria := range criteriaArray {
		result.Contexts = append(result.Contexts,
			&lang.Context{
				Name:     util.RandomID(builder.random, idLength),
				Criteria: criteria,
				Allocation: &lang.Allocation{
					Service: service.Name,
				},
			},
		)
	}

	builder.addObject(builder.domainAdminView, result)
	return result
}

// AddRule creates a new rule and adds it to the policy
func (builder *PolicyBuilder) AddRule(criteria *lang.Criteria, actions *lang.RuleActions) *lang.Rule {
	result := &lang.Rule{
		Metadata: lang.Metadata{
			Kind:      lang.RuleObject.Kind,
			Namespace: builder.namespace,
			Name:      util.RandomID(builder.random, idLength),
		},
		Weight:   len(builder.policy.GetObjectsByKind(lang.RuleObject.Kind)),
		Criteria: criteria,
		Actions:  actions,
	}
	builder.addObject(builder.domainAdminView, result)
	return result
}

// AddCluster creates a new cluster and adds it to the policy
func (builder *PolicyBuilder) AddCluster() *lang.Cluster {
	result := &lang.Cluster{
		Metadata: lang.Metadata{
			Kind:      lang.ClusterObject.Kind,
			Namespace: object.SystemNS,
			Name:      util.RandomID(builder.random, idLength),
		},
		Type: "kubernetes",
	}
	builder.addObject(builder.domainAdminView, result)
	return result
}

// Criteria creates a criteria with one require-all, one require-any, and one require-none
func (builder *PolicyBuilder) Criteria(all string, any string, none string) *lang.Criteria {
	return &lang.Criteria{
		RequireAll:  []string{all},
		RequireAny:  []string{any},
		RequireNone: []string{none},
	}
}

// CriteriaTrue creates a criteria which always evaluates to true
func (builder *PolicyBuilder) CriteriaTrue() *lang.Criteria {
	return &lang.Criteria{
		RequireAny: []string{"true"},
	}
}

// AllocationKeys creates allocation keys
func (builder *PolicyBuilder) AllocationKeys(key string) []string {
	return []string{key}
}

// UnknownComponent creates an unknown component for a service (not code and not contract)
func (builder *PolicyBuilder) UnknownComponent() *lang.ServiceComponent {
	return &lang.ServiceComponent{
		Name: util.RandomID(builder.random, idLength),
	}
}

// CodeComponent creates a new code component for a service
func (builder *PolicyBuilder) CodeComponent(codeParams util.NestedParameterMap, discoveryParams util.NestedParameterMap) *lang.ServiceComponent {
	return &lang.ServiceComponent{
		Name: util.RandomID(builder.random, idLength),
		Code: &lang.Code{
			Type:   "aptomi/code/unittests",
			Params: codeParams,
		},
		Discovery: discoveryParams,
	}
}

// ContractComponent creates a new contract component for a service
func (builder *PolicyBuilder) ContractComponent(contract *lang.Contract) *lang.ServiceComponent {
	return &lang.ServiceComponent{
		Name:     util.RandomID(builder.random, idLength),
		Contract: contract.Namespace + "/" + contract.Name,
	}
}

// AddServiceComponent adds a given service component to the service
func (builder *PolicyBuilder) AddServiceComponent(service *lang.Service, component *lang.ServiceComponent) *lang.ServiceComponent {
	service.Components = append(service.Components, component)
	return component
}

// AddComponentDependency adds a component dependency on another component
func (builder *PolicyBuilder) AddComponentDependency(component *lang.ServiceComponent, dependsOn *lang.ServiceComponent) {
	component.Dependencies = append(component.Dependencies, dependsOn.Name)
}

// RuleActions creates a new RuleActions object
func (builder *PolicyBuilder) RuleActions(labelOps lang.LabelOperations) *lang.RuleActions {
	result := &lang.RuleActions{}
	if labelOps != nil {
		result.ChangeLabels = lang.ChangeLabelsAction(labelOps)
	}
	return result
}

// Policy returns the generated policy
func (builder *PolicyBuilder) Policy() *lang.Policy {
	return builder.policy
}

// External returns the generated external data
func (builder *PolicyBuilder) External() *external.Data {
	return external.NewData(
		builder.users,
		builder.secrets,
	)
}

// Namespace returns the current namespace
func (builder *PolicyBuilder) Namespace() string {
	return builder.namespace
}

// Internal function to add objects to the policy
func (builder *PolicyBuilder) addObject(view *lang.PolicyView, obj object.Base) {
	err := view.AddObject(obj)
	if err != nil {
		panic(err)
	}
}
<commit_msg>got rid of bootstrap rules in policy builder<commit_after>package builder

import (
	"github.com/Aptomi/aptomi/pkg/external"
	"github.com/Aptomi/aptomi/pkg/external/secrets"
	"github.com/Aptomi/aptomi/pkg/external/users"
	"github.com/Aptomi/aptomi/pkg/lang"
	"github.com/Aptomi/aptomi/pkg/object"
	"github.com/Aptomi/aptomi/pkg/util"
	"math/rand"
)

var randSeed = int64(239)
var idLength = 16

// PolicyBuilder is a utility struct to help build policy objects
// It is primarily used in unit tests
type PolicyBuilder struct {
	random    *rand.Rand
	namespace string
	policy    *lang.Policy
	users     *users.UserLoaderMock
	secrets   *secrets.SecretLoaderMock

	domainAdmin     *lang.User
	domainAdminView *lang.PolicyView
}

// NewPolicyBuilder creates a new PolicyBuilder with a default "main" namespace
func NewPolicyBuilder() *PolicyBuilder {
	return NewPolicyBuilderWithNS("main")
}

// NewPolicyBuilderWithNS creates a new PolicyBuilder
func NewPolicyBuilderWithNS(namespace string) *PolicyBuilder {
	result := &PolicyBuilder{
		random:    rand.New(rand.NewSource(randSeed)),
		namespace: namespace,
		policy:    lang.NewPolicy(),
		users:     users.NewUserLoaderMock(),
		secrets:   secrets.NewSecretLoaderMock(),
	}

	result.domainAdmin = result.AddUserDomainAdmin()
	result.domainAdminView = result.policy.View(result.domainAdmin)

	return result
}

// SwitchNamespace switches the current namespace where objects will be generated
func (builder *PolicyBuilder) SwitchNamespace(namespace string) {
	builder.namespace = namespace
}

// AddDependency creates a new dependency and adds it to the policy
func (builder *PolicyBuilder) AddDependency(user *lang.User, contract *lang.Contract) *lang.Dependency {
	result := &lang.Dependency{
		Metadata: lang.Metadata{
			Kind:      lang.DependencyObject.Kind,
			Namespace: builder.namespace,
			Name:      util.RandomID(builder.random, idLength),
		},
		UserID:   user.ID,
		Contract: contract.Namespace + "/" + contract.Name,
		Labels:   make(map[string]string),
	}

	builder.addObject(builder.domainAdminView, result)
	return result
}

// AddUser creates a new user who can consume services from the 'main' namespace and adds it to the policy
func (builder *PolicyBuilder) AddUser() *lang.User {
	result := &lang.User{
		ID:     util.RandomID(builder.random, idLength),
		Name:   util.RandomID(builder.random, idLength),
		Labels: map[string]string{},
		Admin:  true, // this will ensure that this user can consume services
	}
	builder.users.AddUser(result)
	return result
}

// AddUserDomainAdmin creates a new user who is a domain admin and adds it to the policy
func (builder *PolicyBuilder) AddUserDomainAdmin() *lang.User {
	return builder.AddUser()
}

// AddService creates a new service and adds it to the policy
func (builder *PolicyBuilder) AddService(owner *lang.User) *lang.Service {
	var ownerID = ""
	if owner != nil {
		ownerID = owner.ID
	}
	result := &lang.Service{
		Metadata: lang.Metadata{
			Kind:      lang.ServiceObject.Kind,
			Namespace: builder.namespace,
			Name:      util.RandomID(builder.random, idLength),
		},
		Owner: ownerID,
	}
	builder.addObject(builder.domainAdminView, result)
	return result
}

// AddContract creates a new contract for a given service and adds it to the policy
func (builder *PolicyBuilder) AddContract(service *lang.Service, criteria *lang.Criteria) *lang.Contract {
	result := &lang.Contract{
		Metadata: lang.Metadata{
			Kind:      lang.ContractObject.Kind,
			Namespace: builder.namespace,
			Name:      util.RandomID(builder.random, idLength),
		},
		Contexts: []*lang.Context{{
			Name:     util.RandomID(builder.random, idLength),
			Criteria: criteria,
			Allocation: &lang.Allocation{
				Service: service.Name,
			},
		}},
	}
	builder.addObject(builder.domainAdminView, result)
	return result
}

// AddContractMultipleContexts creates contract with multiple contexts for a given service and adds it to the policy
func (builder *PolicyBuilder) AddContractMultipleContexts(service *lang.Service, criteriaArray ...*lang.Criteria) *lang.Contract {
	result := &lang.Contract{
		Metadata: lang.Metadata{
			Kind:      lang.ContractObject.Kind,
			Namespace: builder.namespace,
			Name:      util.RandomID(builder.random, idLength),
		},
	}
	for _, criteria := range criteriaArray {
		result.Contexts = append(result.Contexts,
			&lang.Context{
				Name:     util.RandomID(builder.random, idLength),
				Criteria: criteria,
				Allocation: &lang.Allocation{
					Service: service.Name,
				},
			},
		)
	}

	builder.addObject(builder.domainAdminView, result)
	return result
}

// AddRule creates a new rule and adds it to the policy
func (builder *PolicyBuilder) AddRule(criteria *lang.Criteria, actions *lang.RuleActions) *lang.Rule {
	result := &lang.Rule{
		Metadata: lang.Metadata{
			Kind:      lang.RuleObject.Kind,
			Namespace: builder.namespace,
			Name:      util.RandomID(builder.random, idLength),
		},
		Weight:   len(builder.policy.GetObjectsByKind(lang.RuleObject.Kind)),
		Criteria: criteria,
		Actions:  actions,
	}
	builder.addObject(builder.domainAdminView, result)
	return result
}

// AddCluster creates a new cluster and adds it to the policy
func (builder *PolicyBuilder) AddCluster() *lang.Cluster {
	result := &lang.Cluster{
		Metadata: lang.Metadata{
			Kind:      lang.ClusterObject.Kind,
			Namespace: object.SystemNS,
			Name:      util.RandomID(builder.random, idLength),
		},
		Type: "kubernetes",
	}
	builder.addObject(builder.domainAdminView, result)
	return result
}

// Criteria creates a criteria with one require-all, one require-any, and one require-none
func (builder *PolicyBuilder) Criteria(all string, any string, none string) *lang.Criteria {
	return &lang.Criteria{
		RequireAll:  []string{all},
		RequireAny:  []string{any},
		RequireNone: []string{none},
	}
}

// CriteriaTrue creates a criteria which always evaluates to true
func (builder *PolicyBuilder) CriteriaTrue() *lang.Criteria {
	return &lang.Criteria{
		RequireAny: []string{"true"},
	}
}

// AllocationKeys creates allocation keys
func (builder *PolicyBuilder) AllocationKeys(key string) []string {
	return []string{key}
}

// UnknownComponent creates an unknown component for a service (not code and not contract)
func (builder *PolicyBuilder) UnknownComponent() *lang.ServiceComponent {
	return &lang.ServiceComponent{
		Name: util.RandomID(builder.random, idLength),
	}
}

// CodeComponent creates a new code component for a service
func (builder *PolicyBuilder) CodeComponent(codeParams util.NestedParameterMap, discoveryParams util.NestedParameterMap) *lang.ServiceComponent {
	return &lang.ServiceComponent{
		Name: util.RandomID(builder.random, idLength),
		Code: &lang.Code{
			Type:   "aptomi/code/unittests",
			Params: codeParams,
		},
		Discovery: discoveryParams,
	}
}

// ContractComponent creates a new contract component for a service
func (builder *PolicyBuilder) ContractComponent(contract *lang.Contract) *lang.ServiceComponent {
	return &lang.ServiceComponent{
		Name:     util.RandomID(builder.random, idLength),
		Contract: contract.Namespace + "/" + contract.Name,
	}
}

// AddServiceComponent adds a given service component to the service
func (builder *PolicyBuilder) AddServiceComponent(service *lang.Service, component *lang.ServiceComponent) *lang.ServiceComponent {
	service.Components = append(service.Components, component)
	return component
}

// AddComponentDependency adds a component dependency on another component
func (builder *PolicyBuilder) AddComponentDependency(component *lang.ServiceComponent, dependsOn *lang.ServiceComponent) {
	component.Dependencies = append(component.Dependencies, dependsOn.Name)
}

// RuleActions creates a new RuleActions object
func (builder *PolicyBuilder) RuleActions(labelOps lang.LabelOperations) *lang.RuleActions {
	result := &lang.RuleActions{}
	if labelOps != nil {
		result.ChangeLabels = lang.ChangeLabelsAction(labelOps)
	}
	return result
}

// Policy returns the generated policy
func (builder *PolicyBuilder) Policy() *lang.Policy {
	return builder.policy
}

// External returns the generated external data
func (builder *PolicyBuilder) External() *external.Data {
	return external.NewData(
		builder.users,
		builder.secrets,
	)
}

// Namespace returns the current namespace
func (builder *PolicyBuilder) Namespace() string {
	return builder.namespace
}

// Internal function to add objects to the policy
func (builder *PolicyBuilder) addObject(view *lang.PolicyView, obj object.Base) {
	err := view.AddObject(obj)
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Rook Authors. All rights reserved.

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 mds

import (
	"context"
	"fmt"

	"github.com/pkg/errors"
	cephv1 "github.com/rook/rook/pkg/apis/ceph.rook.io/v1"
	"github.com/rook/rook/pkg/clusterd"
	"github.com/rook/rook/pkg/operator/ceph/config"
	"github.com/rook/rook/pkg/operator/ceph/controller"
	"github.com/rook/rook/pkg/operator/k8sutil"
	apps "k8s.io/api/apps/v1"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
	// MDS cache memory limit should be set to 50-60% of RAM reserved for the MDS container
	// MDS uses approximately 125% of the value of mds_cache_memory_limit in RAM.
	// Eventually we will tune this automatically: http://tracker.ceph.com/issues/36663
	mdsCacheMemoryLimitFactor = 0.5
)

func (c *Cluster) makeDeployment(mdsConfig *mdsConfig) (*apps.Deployment, error) {

	mdsContainer := c.makeMdsDaemonContainer(mdsConfig)
	config.ConfigureLivenessProbe(cephv1.KeyMds, mdsContainer, c.clusterSpec.HealthCheck)

	podSpec := v1.PodTemplateSpec{
		ObjectMeta: metav1.ObjectMeta{
			Name:   mdsConfig.ResourceName,
			Labels: c.podLabels(mdsConfig, true),
		},
		Spec: v1.PodSpec{
			InitContainers: []v1.Container{
				c.makeChownInitContainer(mdsConfig),
			},
			Containers: []v1.Container{
				mdsContainer,
			},
			RestartPolicy:     v1.RestartPolicyAlways,
			Volumes:           controller.DaemonVolumes(mdsConfig.DataPathMap, mdsConfig.ResourceName),
			HostNetwork:       c.clusterSpec.Network.IsHost(),
			PriorityClassName: c.fs.Spec.MetadataServer.PriorityClassName,
		},
	}

	// Replace default unreachable node toleration
	k8sutil.AddUnreachableNodeToleration(&podSpec.Spec)

	// If the log collector is enabled we add the side-car container
	if c.clusterSpec.LogCollector.Enabled {
		shareProcessNamespace := true
		podSpec.Spec.ShareProcessNamespace = &shareProcessNamespace
		podSpec.Spec.Containers = append(podSpec.Spec.Containers, *controller.LogCollectorContainer(fmt.Sprintf("ceph-mds.%s", mdsConfig.DaemonID), c.clusterInfo.Namespace, *c.clusterSpec))
	}

	c.fs.Spec.MetadataServer.Annotations.ApplyToObjectMeta(&podSpec.ObjectMeta)
	c.fs.Spec.MetadataServer.Labels.ApplyToObjectMeta(&podSpec.ObjectMeta)
	c.fs.Spec.MetadataServer.Placement.ApplyToPodSpec(&podSpec.Spec)

	replicas := int32(1)
	d := &apps.Deployment{
		ObjectMeta: metav1.ObjectMeta{
			Name:      mdsConfig.ResourceName,
			Namespace: c.fs.Namespace,
			Labels:    c.podLabels(mdsConfig, true),
		},
		Spec: apps.DeploymentSpec{
			Selector: &metav1.LabelSelector{
				MatchLabels: c.podLabels(mdsConfig, false),
			},
			Template: podSpec,
			Replicas: &replicas,
			Strategy: apps.DeploymentStrategy{
				Type: apps.RecreateDeploymentStrategyType,
			},
		},
	}

	if c.clusterSpec.Network.IsHost() {
		d.Spec.Template.Spec.DNSPolicy = v1.DNSClusterFirstWithHostNet
	} else if c.clusterSpec.Network.NetworkSpec.IsMultus() {
		if err := k8sutil.ApplyMultus(c.clusterSpec.Network.NetworkSpec, &podSpec.ObjectMeta); err != nil {
			return nil, err
		}
	}

	k8sutil.AddRookVersionLabelToDeployment(d)
	c.fs.Spec.MetadataServer.Annotations.ApplyToObjectMeta(&d.ObjectMeta)
	c.fs.Spec.MetadataServer.Labels.ApplyToObjectMeta(&d.ObjectMeta)
	controller.AddCephVersionLabelToDeployment(c.clusterInfo.CephVersion, d)

	return d, nil
}

func (c *Cluster) makeChownInitContainer(mdsConfig *mdsConfig) v1.Container {
	return controller.ChownCephDataDirsInitContainer(
		*mdsConfig.DataPathMap,
		c.clusterSpec.CephVersion.Image,
		controller.DaemonVolumeMounts(mdsConfig.DataPathMap, mdsConfig.ResourceName),
		c.fs.Spec.MetadataServer.Resources,
		controller.PodSecurityContext(),
	)
}

func (c *Cluster) makeMdsDaemonContainer(mdsConfig *mdsConfig) v1.Container {
	args := append(
		controller.DaemonFlags(c.clusterInfo, c.clusterSpec, mdsConfig.DaemonID),
		"--foreground",
	)

	container := v1.Container{
		Name: "mds",
		Command: []string{
			"ceph-mds",
		},
		Args:            args,
		Image:           c.clusterSpec.CephVersion.Image,
		VolumeMounts:    controller.DaemonVolumeMounts(mdsConfig.DataPathMap, mdsConfig.ResourceName),
		Env:             controller.DaemonEnvVars(c.clusterSpec.CephVersion.Image),
		Resources:       c.fs.Spec.MetadataServer.Resources,
		SecurityContext: controller.PodSecurityContext(),
		LivenessProbe:   controller.GenerateLivenessProbeExecDaemon(config.MdsType, mdsConfig.DaemonID),
		WorkingDir:      config.VarLogCephDir,
	}

	return container
}

func (c *Cluster) podLabels(mdsConfig *mdsConfig, includeNewLabels bool) map[string]string {
	labels := controller.CephDaemonAppLabels(AppName, c.fs.Namespace, "mds", mdsConfig.DaemonID, includeNewLabels)
	labels["rook_file_system"] = c.fs.Name
	return labels
}

func getMdsDeployments(context *clusterd.Context, namespace, fsName string) (*apps.DeploymentList, error) {
	fsLabelSelector := fmt.Sprintf("rook_file_system=%s", fsName)
	deps, err := k8sutil.GetDeployments(context.Clientset, namespace, fsLabelSelector)
	if err != nil {
		return nil, errors.Wrapf(err, "could not get deployments for filesystem %s (matching label selector %q)", fsName, fsLabelSelector)
	}
	return deps, nil
}

func deleteMdsDeployment(clusterdContext *clusterd.Context, namespace string, deployment *apps.Deployment) error {
	ctx := context.TODO()
	// Delete the mds deployment
	logger.Infof("deleting mds deployment %s", deployment.Name)
	var gracePeriod int64
	propagation := metav1.DeletePropagationForeground
	options := &metav1.DeleteOptions{GracePeriodSeconds: &gracePeriod, PropagationPolicy: &propagation}
	if err := clusterdContext.Clientset.AppsV1().Deployments(namespace).Delete(ctx, deployment.GetName(), *options); err != nil {
		return errors.Wrapf(err, "failed to delete mds deployment %s", deployment.GetName())
	}
	return nil
}
<commit_msg>ceph: fix mds liveness probe<commit_after>/*
Copyright 2016 The Rook Authors. All rights reserved.

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 mds

import (
	"context"
	"fmt"

	"github.com/pkg/errors"
	cephv1 "github.com/rook/rook/pkg/apis/ceph.rook.io/v1"
	"github.com/rook/rook/pkg/clusterd"
	"github.com/rook/rook/pkg/operator/ceph/config"
	"github.com/rook/rook/pkg/operator/ceph/controller"
	"github.com/rook/rook/pkg/operator/k8sutil"
	apps "k8s.io/api/apps/v1"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
	// MDS cache memory limit should be set to 50-60% of RAM reserved for the MDS container
	// MDS uses approximately 125% of the value of mds_cache_memory_limit in RAM.
	// Eventually we will tune this automatically: http://tracker.ceph.com/issues/36663
	mdsCacheMemoryLimitFactor = 0.5
)

func (c *Cluster) makeDeployment(mdsConfig *mdsConfig) (*apps.Deployment, error) {

	mdsContainer := c.makeMdsDaemonContainer(mdsConfig)
	mdsContainer = config.ConfigureLivenessProbe(cephv1.KeyMds, mdsContainer, c.clusterSpec.HealthCheck)

	podSpec := v1.PodTemplateSpec{
		ObjectMeta: metav1.ObjectMeta{
			Name:   mdsConfig.ResourceName,
			Labels: c.podLabels(mdsConfig, true),
		},
		Spec: v1.PodSpec{
			InitContainers: []v1.Container{
				c.makeChownInitContainer(mdsConfig),
			},
			Containers: []v1.Container{
				mdsContainer,
			},
			RestartPolicy:     v1.RestartPolicyAlways,
			Volumes:           controller.DaemonVolumes(mdsConfig.DataPathMap, mdsConfig.ResourceName),
			HostNetwork:       c.clusterSpec.Network.IsHost(),
			PriorityClassName: c.fs.Spec.MetadataServer.PriorityClassName,
		},
	}

	// Replace default unreachable node toleration
	k8sutil.AddUnreachableNodeToleration(&podSpec.Spec)

	// If the log collector is enabled we add the side-car container
	if c.clusterSpec.LogCollector.Enabled {
		shareProcessNamespace := true
		podSpec.Spec.ShareProcessNamespace = &shareProcessNamespace
		podSpec.Spec.Containers = append(podSpec.Spec.Containers, *controller.LogCollectorContainer(fmt.Sprintf("ceph-mds.%s", mdsConfig.DaemonID), c.clusterInfo.Namespace, *c.clusterSpec))
	}

	c.fs.Spec.MetadataServer.Annotations.ApplyToObjectMeta(&podSpec.ObjectMeta)
	c.fs.Spec.MetadataServer.Labels.ApplyToObjectMeta(&podSpec.ObjectMeta)
	c.fs.Spec.MetadataServer.Placement.ApplyToPodSpec(&podSpec.Spec)

	replicas := int32(1)
	d := &apps.Deployment{
		ObjectMeta: metav1.ObjectMeta{
			Name:      mdsConfig.ResourceName,
			Namespace: c.fs.Namespace,
			Labels:    c.podLabels(mdsConfig, true),
		},
		Spec: apps.DeploymentSpec{
			Selector: &metav1.LabelSelector{
				MatchLabels: c.podLabels(mdsConfig, false),
			},
			Template: podSpec,
			Replicas: &replicas,
			Strategy: apps.DeploymentStrategy{
				Type: apps.RecreateDeploymentStrategyType,
			},
		},
	}

	if c.clusterSpec.Network.IsHost() {
		d.Spec.Template.Spec.DNSPolicy = v1.DNSClusterFirstWithHostNet
	} else if c.clusterSpec.Network.NetworkSpec.IsMultus() {
		if err := k8sutil.ApplyMultus(c.clusterSpec.Network.NetworkSpec, &podSpec.ObjectMeta); err != nil {
			return nil, err
		}
	}

	k8sutil.AddRookVersionLabelToDeployment(d)
	c.fs.Spec.MetadataServer.Annotations.ApplyToObjectMeta(&d.ObjectMeta)
	c.fs.Spec.MetadataServer.Labels.ApplyToObjectMeta(&d.ObjectMeta)
	controller.AddCephVersionLabelToDeployment(c.clusterInfo.CephVersion, d)

	return d, nil
}

func (c *Cluster) makeChownInitContainer(mdsConfig *mdsConfig) v1.Container {
	return controller.ChownCephDataDirsInitContainer(
		*mdsConfig.DataPathMap,
		c.clusterSpec.CephVersion.Image,
		controller.DaemonVolumeMounts(mdsConfig.DataPathMap, mdsConfig.ResourceName),
		c.fs.Spec.MetadataServer.Resources,
		controller.PodSecurityContext(),
	)
}

func (c *Cluster) makeMdsDaemonContainer(mdsConfig *mdsConfig) v1.Container {
	args := append(
		controller.DaemonFlags(c.clusterInfo, c.clusterSpec, mdsConfig.DaemonID),
		"--foreground",
	)

	container := v1.Container{
		Name: "mds",
		Command: []string{
			"ceph-mds",
		},
		Args:            args,
		Image:           c.clusterSpec.CephVersion.Image,
		VolumeMounts:    controller.DaemonVolumeMounts(mdsConfig.DataPathMap, mdsConfig.ResourceName),
		Env:             controller.DaemonEnvVars(c.clusterSpec.CephVersion.Image),
		Resources:       c.fs.Spec.MetadataServer.Resources,
		SecurityContext: controller.PodSecurityContext(),
		LivenessProbe:   controller.GenerateLivenessProbeExecDaemon(config.MdsType, mdsConfig.DaemonID),
		WorkingDir:      config.VarLogCephDir,
	}

	return container
}

func (c *Cluster) podLabels(mdsConfig *mdsConfig, includeNewLabels bool) map[string]string {
	labels := controller.CephDaemonAppLabels(AppName, c.fs.Namespace, "mds", mdsConfig.DaemonID, includeNewLabels)
	labels["rook_file_system"] = c.fs.Name
	return labels
}

func getMdsDeployments(context *clusterd.Context, namespace, fsName string) (*apps.DeploymentList, error) {
	fsLabelSelector := fmt.Sprintf("rook_file_system=%s", fsName)
	deps, err := k8sutil.GetDeployments(context.Clientset, namespace, fsLabelSelector)
	if err != nil {
		return nil, errors.Wrapf(err, "could not get deployments for filesystem %s (matching label selector %q)", fsName, fsLabelSelector)
	}
	return deps, nil
}

func deleteMdsDeployment(clusterdContext *clusterd.Context, namespace string, deployment *apps.Deployment) error {
	ctx := context.TODO()
	// Delete the mds deployment
	logger.Infof("deleting mds deployment %s", deployment.Name)
	var gracePeriod int64
	propagation := metav1.DeletePropagationForeground
	options := &metav1.DeleteOptions{GracePeriodSeconds: &gracePeriod, PropagationPolicy: &propagation}
	if err := clusterdContext.Clientset.AppsV1().Deployments(namespace).Delete(ctx, deployment.GetName(), *options); err != nil {
		return errors.Wrapf(err, "failed to delete mds deployment %s", deployment.GetName())
	}
	return nil
}
<|endoftext|>