<commit_before>package goproxy

import (
	"bufio"
	"crypto/tls"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync/atomic"
)

type ConnectActionLiteral int

const (
	ConnectAccept = iota
	ConnectReject
	ConnectMitm
	ConnectHijack
	ConnectHTTPMitm
)

var (
	OkConnect       = &ConnectAction{Action: ConnectAccept, TLSConfig: TLSConfigFromCA(&GoproxyCa)}
	MitmConnect     = &ConnectAction{Action: ConnectMitm, TLSConfig: TLSConfigFromCA(&GoproxyCa)}
	HTTPMitmConnect = &ConnectAction{Action: ConnectHTTPMitm, TLSConfig: TLSConfigFromCA(&GoproxyCa)}
	RejectConnect   = &ConnectAction{Action: ConnectReject, TLSConfig: TLSConfigFromCA(&GoproxyCa)}
)

type ConnectAction struct {
	Action    ConnectActionLiteral
	Hijack    func(req *http.Request, client net.Conn, ctx *ProxyCtx)
	TLSConfig func(host string, ctx *ProxyCtx) (*tls.Config, error)
}

func stripPort(s string) string {
	ix := strings.IndexRune(s, ':')
	if ix == -1 {
		return s
	}
	return s[:ix]
}

func (proxy *ProxyHttpServer) dial(network, addr string) (c net.Conn, err error) {
	if proxy.Tr.Dial != nil {
		return proxy.Tr.Dial(network, addr)
	}
	return net.Dial(network, addr)
}

func (proxy *ProxyHttpServer) connectDial(network, addr string) (c net.Conn, err error) {
	if proxy.ConnectDial == nil {
		return proxy.dial(network, addr)
	}
	return proxy.ConnectDial(network, addr)
}

func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request) {
	ctx := &ProxyCtx{Req: r, Session: atomic.AddInt64(&proxy.sess, 1), proxy: proxy}

	hij, ok := w.(http.Hijacker)
	if !ok {
		panic("httpserver does not support hijacking")
	}

	proxyClient, _, e := hij.Hijack()
	if e != nil {
		panic("Cannot hijack connection " + e.Error())
	}

	ctx.Logf("Running %d CONNECT handlers", len(proxy.httpsHandlers))
	todo, host := OkConnect, r.URL.Host
	for i, h := range proxy.httpsHandlers {
		newtodo, newhost := h.HandleConnect(host, ctx)

		// If found a result, break the loop immediately
		if newtodo != nil {
			todo, host = newtodo, newhost
			ctx.Logf("on %dth handler: %v %s", i, todo, host)
			break
		}
	}
	switch todo.Action {
	case ConnectAccept:
		if !hasPort.MatchString(host) {
			host += ":80"
		}
		targetSiteCon, err := proxy.connectDial("tcp", host)
		if err != nil {
			httpError(proxyClient, ctx, err)
			return
		}
		ctx.Logf("Accepting CONNECT to %s", host)
		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
		go copyAndClose(ctx, targetSiteCon, proxyClient)
		go copyAndClose(ctx, proxyClient, targetSiteCon)
	case ConnectHijack:
		ctx.Logf("Hijacking CONNECT to %s", host)
		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
		todo.Hijack(r, proxyClient, ctx)
	case ConnectHTTPMitm:
		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
		ctx.Logf("Assuming CONNECT is plain HTTP tunneling, mitm proxying it")
		targetSiteCon, err := proxy.connectDial("tcp", host)
		if err != nil {
			ctx.Warnf("Error dialing to %s: %s", host, err.Error())
			return
		}
		for {
			client := bufio.NewReader(proxyClient)
			remote := bufio.NewReader(targetSiteCon)
			req, err := http.ReadRequest(client)
			if err != nil && err != io.EOF {
				ctx.Warnf("cannot read request of MITM HTTP client: %+#v", err)
			}
			if err != nil {
				return
			}
			req, resp := proxy.filterRequest(req, ctx)
			if resp == nil {
				if err := req.Write(targetSiteCon); err != nil {
					httpError(proxyClient, ctx, err)
					return
				}
				resp, err = http.ReadResponse(remote, req)
				if err != nil {
					httpError(proxyClient, ctx, err)
					return
				}
			}
			resp = proxy.filterResponse(resp, ctx)
			if err := resp.Write(proxyClient); err != nil {
				httpError(proxyClient, ctx, err)
				return
			}
		}
	case ConnectMitm:
		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
		ctx.Logf("Assuming CONNECT is TLS, mitm proxying it")
		// this goes in a separate goroutine, so that the net/http server won't think we're
		// still handling the request even after hijacking the connection. Those HTTP CONNECT
		// request can take forever, and the server will be stuck when "closed".
		// TODO: Allow Server.Close() mechanism to shut down this connection as nicely as possible
		tlsConfig := defaultTLSConfig
		if todo.TLSConfig != nil {
			var err error
			tlsConfig, err = todo.TLSConfig(host, ctx)
			if err != nil {
				httpError(proxyClient, ctx, err)
				return
			}
		}
		go func() {
			//TODO: cache connections to the remote website
			rawClientTls := tls.Server(proxyClient, tlsConfig)
			if err := rawClientTls.Handshake(); err != nil {
				ctx.Warnf("Cannot handshake client %v %v", r.Host, err)
				return
			}
			defer rawClientTls.Close()
			clientTlsReader := bufio.NewReader(rawClientTls)
			for !isEof(clientTlsReader) {
				req, err := http.ReadRequest(clientTlsReader)
				if err != nil && err != io.EOF {
					return
				}
				if err != nil {
					ctx.Warnf("Cannot read TLS request from mitm'd client %v %v", r.Host, err)
					return
				}
				req.RemoteAddr = r.RemoteAddr // since we're converting the request, need to carry over the original connecting IP as well
				ctx.Logf("req %v", r.Host)
				req.URL, err = url.Parse("https://" + r.Host + req.URL.String())

				// Bug fix which goproxy fails to provide request
				// information URL in the context when does HTTPS MITM
				ctx.Req = req

				req, resp := proxy.filterRequest(req, ctx)
				if resp == nil {
					if err != nil {
						ctx.Warnf("Illegal URL %s", "https://"+r.Host+req.URL.Path)
						return
					}
					removeProxyHeaders(ctx, req)
					resp, err = ctx.RoundTrip(req)
					if err != nil {
						ctx.Warnf("Cannot read TLS response from mitm'd server %v", err)
						return
					}
					ctx.Logf("resp %v", resp.Status)
				}
				resp = proxy.filterResponse(resp, ctx)
				text := resp.Status
				statusCode := strconv.Itoa(resp.StatusCode) + " "
				if strings.HasPrefix(text, statusCode) {
					text = text[len(statusCode):]
				}
				// always use 1.1 to support chunked encoding
				if _, err := io.WriteString(rawClientTls, "HTTP/1.1"+" "+statusCode+text+"\r\n"); err != nil {
					ctx.Warnf("Cannot write TLS response HTTP status from mitm'd client: %v", err)
					return
				}
				// Since we don't know the length of resp, return chunked encoded response
				// TODO: use a more reasonable scheme
				resp.Header.Del("Content-Length")
				resp.Header.Set("Transfer-Encoding", "chunked")
				if err := resp.Header.Write(rawClientTls); err != nil {
					ctx.Warnf("Cannot write TLS response header from mitm'd client: %v", err)
					return
				}
				if _, err = io.WriteString(rawClientTls, "\r\n"); err != nil {
					ctx.Warnf("Cannot write TLS response header end from mitm'd client: %v", err)
					return
				}
				chunked := newChunkedWriter(rawClientTls)
				if _, err := io.Copy(chunked, resp.Body); err != nil {
					ctx.Warnf("Cannot write TLS response body from mitm'd client: %v", err)
					return
				}
				if err := chunked.Close(); err != nil {
					ctx.Warnf("Cannot write TLS chunked EOF from mitm'd client: %v", err)
					return
				}
				if _, err = io.WriteString(rawClientTls, "\r\n"); err != nil {
					ctx.Warnf("Cannot write TLS response chunked trailer from mitm'd client: %v", err)
					return
				}
			}
			ctx.Logf("Exiting on EOF")
		}()
	case ConnectReject:
		if ctx.Resp != nil {
			if err := ctx.Resp.Write(proxyClient); err != nil {
				ctx.Warnf("Cannot write response that reject http CONNECT: %v", err)
			}
		}
		proxyClient.Close()
	}
}

func httpError(w io.WriteCloser, ctx *ProxyCtx, err error) {
	if _, err := io.WriteString(w, "HTTP/1.1 502 Bad Gateway\r\n\r\n"); err != nil {
		ctx.Warnf("Error responding to client: %s", err)
	}
	if err := w.Close(); err != nil {
		ctx.Warnf("Error closing client connection: %s", err)
	}
}

func copyAndClose(ctx *ProxyCtx, w, r net.Conn) {
	connOk := true
	if _, err := io.Copy(w, r); err != nil {
		connOk = false
		ctx.Warnf("Error copying to client: %s", err)
	}
	if err := r.Close(); err != nil && connOk {
		ctx.Warnf("Error closing: %s", err)
	}
}

func dialerFromEnv(proxy *ProxyHttpServer) func(network, addr string) (net.Conn, error) {
	https_proxy := os.Getenv("HTTPS_PROXY")
	if https_proxy == "" {
		https_proxy = os.Getenv("https_proxy")
	}
	if https_proxy == "" {
		return nil
	}
	return proxy.NewConnectDialToProxy(https_proxy)
}

func (proxy *ProxyHttpServer) NewConnectDialToProxy(https_proxy string) func(network, addr string) (net.Conn, error) {
	u, err := url.Parse(https_proxy)
	if err != nil {
		return nil
	}
	if u.Scheme == "" || u.Scheme == "http" {
		if strings.IndexRune(u.Host, ':') == -1 {
			u.Host += ":80"
		}
		return func(network, addr string) (net.Conn, error) {
			connectReq := &http.Request{
				Method: "CONNECT",
				URL:    &url.URL{Opaque: addr},
				Host:   addr,
				Header: make(http.Header),
			}
			c, err := proxy.dial(network, u.Host)
			if err != nil {
				return nil, err
			}
			connectReq.Write(c)
			// Read response.
			// Okay to use and discard buffered reader here, because
			// TLS server will not speak until spoken to.
			br := bufio.NewReader(c)
			resp, err := http.ReadResponse(br, connectReq)
			if err != nil {
				c.Close()
				return nil, err
			}
			if resp.StatusCode != 200 {
				resp, _ := ioutil.ReadAll(resp.Body)
				c.Close()
				return nil, errors.New("proxy refused connection" + string(resp))
			}
			return c, nil
		}
	}
	if u.Scheme == "https" {
		if strings.IndexRune(u.Host, ':') == -1 {
			u.Host += ":443"
		}
		return func(network, addr string) (net.Conn, error) {
			c, err := proxy.dial(network, u.Host)
			if err != nil {
				return nil, err
			}
			c = tls.Client(c, proxy.Tr.TLSClientConfig)
			connectReq := &http.Request{
				Method: "CONNECT",
				URL:    &url.URL{Opaque: addr},
				Host:   addr,
				Header: make(http.Header),
			}
			connectReq.Write(c)
			// Read response.
			// Okay to use and discard buffered reader here, because
			// TLS server will not speak until spoken to.
			br := bufio.NewReader(c)
			resp, err := http.ReadResponse(br, connectReq)
			if err != nil {
				c.Close()
				return nil, err
			}
			if resp.StatusCode != 200 {
				body, _ := ioutil.ReadAll(io.LimitReader(resp.Body, 500))
				resp.Body.Close()
				c.Close()
				return nil, errors.New("proxy refused connection" + string(body))
			}
			return c, nil
		}
	}
	return nil
}

func TLSConfigFromCA(ca *tls.Certificate) func(host string, ctx *ProxyCtx) (*tls.Config, error) {
	return func(host string, ctx *ProxyCtx) (*tls.Config, error) {
		config := *defaultTLSConfig
		ctx.Logf("signing for %s", stripPort(host))
		cert, err := signHost(*ca, []string{stripPort(host)})
		if err != nil {
			ctx.Warnf("Cannot sign host certificate with provided CA: %s", err)
			return nil, err
		}
		config.Certificates = append(config.Certificates, cert)
		return &config, nil
	}
}
<commit_msg>Add customizable HTTPS proxy authentication<commit_after>package goproxy

import (
	"bufio"
	"crypto/tls"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync/atomic"
)

type ConnectActionLiteral int

const (
	ConnectAccept = iota
	ConnectReject
	ConnectMitm
	ConnectHijack
	ConnectHTTPMitm
	ConnectProxyAuthHijack
)

var (
	OkConnect       = &ConnectAction{Action: ConnectAccept, TLSConfig: TLSConfigFromCA(&GoproxyCa)}
	MitmConnect     = &ConnectAction{Action: ConnectMitm, TLSConfig: TLSConfigFromCA(&GoproxyCa)}
	HTTPMitmConnect = &ConnectAction{Action: ConnectHTTPMitm, TLSConfig: TLSConfigFromCA(&GoproxyCa)}
	RejectConnect   = &ConnectAction{Action: ConnectReject, TLSConfig: TLSConfigFromCA(&GoproxyCa)}
)

type ConnectAction struct {
	Action    ConnectActionLiteral
	Hijack    func(req *http.Request, client net.Conn, ctx *ProxyCtx)
	TLSConfig func(host string, ctx *ProxyCtx) (*tls.Config, error)
}

func stripPort(s string) string {
	ix := strings.IndexRune(s, ':')
	if ix == -1 {
		return s
	}
	return s[:ix]
}

func (proxy *ProxyHttpServer) dial(network, addr string) (c net.Conn, err error) {
	if proxy.Tr.Dial != nil {
		return proxy.Tr.Dial(network, addr)
	}
	return net.Dial(network, addr)
}

func (proxy *ProxyHttpServer) connectDial(network, addr string) (c net.Conn, err error) {
	if proxy.ConnectDial == nil {
		return proxy.dial(network, addr)
	}
	return proxy.ConnectDial(network, addr)
}

func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request) {
	ctx := &ProxyCtx{Req: r, Session: atomic.AddInt64(&proxy.sess, 1), proxy: proxy}

	hij, ok := w.(http.Hijacker)
	if !ok {
		panic("httpserver does not support hijacking")
	}

	proxyClient, _, e := hij.Hijack()
	if e != nil {
		panic("Cannot hijack connection " + e.Error())
	}

	ctx.Logf("Running %d CONNECT handlers", len(proxy.httpsHandlers))
	todo, host := OkConnect, r.URL.Host
	for i, h := range proxy.httpsHandlers {
		newtodo, newhost := h.HandleConnect(host, ctx)

		// If found a result, break the loop immediately
		if newtodo != nil {
			todo, host = newtodo, newhost
			ctx.Logf("on %dth handler: %v %s", i, todo, host)
			break
		}
	}
	switch todo.Action {
	case ConnectAccept:
		if !hasPort.MatchString(host) {
			host += ":80"
		}
		targetSiteCon, err := proxy.connectDial("tcp", host)
		if err != nil {
			httpError(proxyClient, ctx, err)
			return
		}
		ctx.Logf("Accepting CONNECT to %s", host)
		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
		go copyAndClose(ctx, targetSiteCon, proxyClient)
		go copyAndClose(ctx, proxyClient, targetSiteCon)
	case ConnectHijack:
		ctx.Logf("Hijacking CONNECT to %s", host)
		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
		todo.Hijack(r, proxyClient, ctx)
	case ConnectHTTPMitm:
		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
		ctx.Logf("Assuming CONNECT is plain HTTP tunneling, mitm proxying it")
		targetSiteCon, err := proxy.connectDial("tcp", host)
		if err != nil {
			ctx.Warnf("Error dialing to %s: %s", host, err.Error())
			return
		}
		for {
			client := bufio.NewReader(proxyClient)
			remote := bufio.NewReader(targetSiteCon)
			req, err := http.ReadRequest(client)
			if err != nil && err != io.EOF {
				ctx.Warnf("cannot read request of MITM HTTP client: %+#v", err)
			}
			if err != nil {
				return
			}
			req, resp := proxy.filterRequest(req, ctx)
			if resp == nil {
				if err := req.Write(targetSiteCon); err != nil {
					httpError(proxyClient, ctx, err)
					return
				}
				resp, err = http.ReadResponse(remote, req)
				if err != nil {
					httpError(proxyClient, ctx, err)
					return
				}
			}
			resp = proxy.filterResponse(resp, ctx)
			if err := resp.Write(proxyClient); err != nil {
				httpError(proxyClient, ctx, err)
				return
			}
		}
	case ConnectMitm:
		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
		ctx.Logf("Assuming CONNECT is TLS, mitm proxying it")
		// this goes in a separate goroutine, so that the net/http server won't think we're
		// still handling the request even after hijacking the connection. Those HTTP CONNECT
		// request can take forever, and the server will be stuck when "closed".
		// TODO: Allow Server.Close() mechanism to shut down this connection as nicely as possible
		tlsConfig := defaultTLSConfig
		if todo.TLSConfig != nil {
			var err error
			tlsConfig, err = todo.TLSConfig(host, ctx)
			if err != nil {
				httpError(proxyClient, ctx, err)
				return
			}
		}
		go func() {
			//TODO: cache connections to the remote website
			rawClientTls := tls.Server(proxyClient, tlsConfig)
			if err := rawClientTls.Handshake(); err != nil {
				ctx.Warnf("Cannot handshake client %v %v", r.Host, err)
				return
			}
			defer rawClientTls.Close()
			clientTlsReader := bufio.NewReader(rawClientTls)
			for !isEof(clientTlsReader) {
				req, err := http.ReadRequest(clientTlsReader)
				if err != nil && err != io.EOF {
					return
				}
				if err != nil {
					ctx.Warnf("Cannot read TLS request from mitm'd client %v %v", r.Host, err)
					return
				}
				req.RemoteAddr = r.RemoteAddr // since we're converting the request, need to carry over the original connecting IP as well
				ctx.Logf("req %v", r.Host)
				req.URL, err = url.Parse("https://" + r.Host + req.URL.String())

				// Bug fix which goproxy fails to provide request
				// information URL in the context when does HTTPS MITM
				ctx.Req = req

				req, resp := proxy.filterRequest(req, ctx)
				if resp == nil {
					if err != nil {
						ctx.Warnf("Illegal URL %s", "https://"+r.Host+req.URL.Path)
						return
					}
					removeProxyHeaders(ctx, req)
					resp, err = ctx.RoundTrip(req)
					if err != nil {
						ctx.Warnf("Cannot read TLS response from mitm'd server %v", err)
						return
					}
					ctx.Logf("resp %v", resp.Status)
				}
				resp = proxy.filterResponse(resp, ctx)
				text := resp.Status
				statusCode := strconv.Itoa(resp.StatusCode) + " "
				if strings.HasPrefix(text, statusCode) {
					text = text[len(statusCode):]
				}
				// always use 1.1 to support chunked encoding
				if _, err := io.WriteString(rawClientTls, "HTTP/1.1"+" "+statusCode+text+"\r\n"); err != nil {
					ctx.Warnf("Cannot write TLS response HTTP status from mitm'd client: %v", err)
					return
				}
				// Since we don't know the length of resp, return chunked encoded response
				// TODO: use a more reasonable scheme
				resp.Header.Del("Content-Length")
				resp.Header.Set("Transfer-Encoding", "chunked")
				if err := resp.Header.Write(rawClientTls); err != nil {
					ctx.Warnf("Cannot write TLS response header from mitm'd client: %v", err)
					return
				}
				if _, err = io.WriteString(rawClientTls, "\r\n"); err != nil {
					ctx.Warnf("Cannot write TLS response header end from mitm'd client: %v", err)
					return
				}
				chunked := newChunkedWriter(rawClientTls)
				if _, err := io.Copy(chunked, resp.Body); err != nil {
					ctx.Warnf("Cannot write TLS response body from mitm'd client: %v", err)
					return
				}
				if err := chunked.Close(); err != nil {
					ctx.Warnf("Cannot write TLS chunked EOF from mitm'd client: %v", err)
					return
				}
				if _, err = io.WriteString(rawClientTls, "\r\n"); err != nil {
					ctx.Warnf("Cannot write TLS response chunked trailer from mitm'd client: %v", err)
					return
				}
			}
			ctx.Logf("Exiting on EOF")
		}()
	case ConnectProxyAuthHijack:
		proxyClient.Write([]byte("HTTP/1.1 407 Proxy Authentication Required\r\n"))
		todo.Hijack(r, proxyClient, ctx)
	case ConnectReject:
		if ctx.Resp != nil {
			if err := ctx.Resp.Write(proxyClient); err != nil {
				ctx.Warnf("Cannot write response that reject http CONNECT: %v", err)
			}
		}
		proxyClient.Close()
	}
}

func httpError(w io.WriteCloser, ctx *ProxyCtx, err error) {
	if _, err := io.WriteString(w, "HTTP/1.1 502 Bad Gateway\r\n\r\n"); err != nil {
		ctx.Warnf("Error responding to client: %s", err)
	}
	if err := w.Close(); err != nil {
		ctx.Warnf("Error closing client connection: %s", err)
	}
}

func copyAndClose(ctx *ProxyCtx, w, r net.Conn) {
	connOk := true
	if _, err := io.Copy(w, r); err != nil {
		connOk = false
		ctx.Warnf("Error copying to client: %s", err)
	}
	if err := r.Close(); err != nil && connOk {
		ctx.Warnf("Error closing: %s", err)
	}
}

func dialerFromEnv(proxy *ProxyHttpServer) func(network, addr string) (net.Conn, error) {
	https_proxy := os.Getenv("HTTPS_PROXY")
	if https_proxy == "" {
		https_proxy = os.Getenv("https_proxy")
	}
	if https_proxy == "" {
		return nil
	}
	return proxy.NewConnectDialToProxy(https_proxy)
}

func (proxy *ProxyHttpServer) NewConnectDialToProxy(https_proxy string) func(network, addr string) (net.Conn, error) {
	u, err := url.Parse(https_proxy)
	if err != nil {
		return nil
	}
	if u.Scheme == "" || u.Scheme == "http" {
		if strings.IndexRune(u.Host, ':') == -1 {
			u.Host += ":80"
		}
		return func(network, addr string) (net.Conn, error) {
			connectReq := &http.Request{
				Method: "CONNECT",
				URL:    &url.URL{Opaque: addr},
				Host:   addr,
				Header: make(http.Header),
			}
			c, err := proxy.dial(network, u.Host)
			if err != nil {
				return nil, err
			}
			connectReq.Write(c)
			// Read response.
			// Okay to use and discard buffered reader here, because
			// TLS server will not speak until spoken to.
			br := bufio.NewReader(c)
			resp, err := http.ReadResponse(br, connectReq)
			if err != nil {
				c.Close()
				return nil, err
			}
			if resp.StatusCode != 200 {
				resp, _ := ioutil.ReadAll(resp.Body)
				c.Close()
				return nil, errors.New("proxy refused connection" + string(resp))
			}
			return c, nil
		}
	}
	if u.Scheme == "https" {
		if strings.IndexRune(u.Host, ':') == -1 {
			u.Host += ":443"
		}
		return func(network, addr string) (net.Conn, error) {
			c, err := proxy.dial(network, u.Host)
			if err != nil {
				return nil, err
			}
			c = tls.Client(c, proxy.Tr.TLSClientConfig)
			connectReq := &http.Request{
				Method: "CONNECT",
				URL:    &url.URL{Opaque: addr},
				Host:   addr,
				Header: make(http.Header),
			}
			connectReq.Write(c)
			// Read response.
			// Okay to use and discard buffered reader here, because
			// TLS server will not speak until spoken to.
			br := bufio.NewReader(c)
			resp, err := http.ReadResponse(br, connectReq)
			if err != nil {
				c.Close()
				return nil, err
			}
			if resp.StatusCode != 200 {
				body, _ := ioutil.ReadAll(io.LimitReader(resp.Body, 500))
				resp.Body.Close()
				c.Close()
				return nil, errors.New("proxy refused connection" + string(body))
			}
			return c, nil
		}
	}
	return nil
}

func TLSConfigFromCA(ca *tls.Certificate) func(host string, ctx *ProxyCtx) (*tls.Config, error) {
	return func(host string, ctx *ProxyCtx) (*tls.Config, error) {
		config := *defaultTLSConfig
		ctx.Logf("signing for %s", stripPort(host))
		cert, err := signHost(*ca, []string{stripPort(host)})
		if err != nil {
			ctx.Warnf("Cannot sign host certificate with provided CA: %s", err)
			return nil, err
		}
		config.Certificates = append(config.Certificates, cert)
		return &config, nil
	}
}
<|endoftext|>


<commit_before>package secstring

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"strings"
	"syscall"
)

type SecString struct {
	String  []byte
	Length  int
	cipher  cipher.Block
	iv      []byte
	Padding int
}

func memset(s []byte, c byte) {
	for i := 0; i < len(s); i++ {
		s[i] = c
	}
}

func NewSecString(str []byte) (*SecString, error) {
	ret := &SecString{Length: len(str)}
	var err error

	if padding := ret.Length % aes.BlockSize; padding != 0 {
		ret.Padding = aes.BlockSize - padding
	}

	ret.String, err = syscall.Mmap(-1, 0, ret.Length+ret.Padding, syscall.PROT_READ|syscall.PROT_WRITE, syscall.MAP_ANON|syscall.MAP_PRIVATE)
	if err != nil {
		return nil, err
	}

	if err := syscall.Mlock(ret.String); err != nil {
		syscall.Munmap(ret.String)
		return nil, err
	}

	for i := 0; i < ret.Length; i++ {
		ret.String[i] = str[i]
		str[i] = 0
	}

	key := make([]byte, 32)
	if _, err := rand.Read(key); err != nil {
		memset(ret.String, 0)
		syscall.Munmap(ret.String)
		return nil, err
	}

	if ret.cipher, err = aes.NewCipher(key); err != nil {
		memset(ret.String, 0)
		syscall.Munmap(ret.String)
		return nil, err
	}

	if err := ret.Encrypt(); err != nil {
		memset(ret.String, 0)
		syscall.Munmap(ret.String)
		return nil, err
	}

	return ret, nil
}

func (s *SecString) Encrypt() error {
	if err := syscall.Mprotect(s.String, syscall.PROT_READ|syscall.PROT_WRITE); err != nil {
		return err
	}

	s.iv = make([]byte, aes.BlockSize)
	if _, err := rand.Read(s.iv); err != nil {
		return err
	}

	encrypter := cipher.NewOFB(s.cipher, s.iv)
	encrypter.XORKeyStream(s.String, s.String)

	if err := syscall.Mprotect(s.String, syscall.PROT_READ); err != nil {
		return err
	}

	return nil
}

func (s *SecString) Decrypt() error {
	if err := syscall.Mprotect(s.String, syscall.PROT_READ|syscall.PROT_WRITE); err != nil {
		return err
	}

	decrypter := cipher.NewOFB(s.cipher, s.iv)
	decrypter.XORKeyStream(s.String, s.String)

	if err := syscall.Mprotect(s.String, syscall.PROT_READ); err != nil {
		return err
	}

	return nil
}

func FromString(str *string) (*SecString, error) {
	b := make([]byte, len(*str))
	for i := 0; i < len(*str); i++ {
		b[i] = (*str)[i]
	}
	*str = strings.Repeat("x", len(*str))
	return NewSecString(b)
}

func (s *SecString) Destroy() error {
	if err := syscall.Mprotect(s.String, syscall.PROT_READ|syscall.PROT_WRITE); err != nil {
		return err
	}

	memset(s.String, 0)

	if err := syscall.Munlock(s.String); err != nil {
		return err
	}

	if err := syscall.Munmap(s.String); err != nil {
		return err
	}

	s.String = nil
	return nil
}
<commit_msg>Added documentation<commit_after>package secstring

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"strings"
	"syscall"
)

type SecString struct {
	String  []byte // Encrypted/Decrypted string
	Length  int    // Length of the target string
	Padding int    // Length of padding added
	iv      []byte
	cipher  cipher.Block
}

func memset(s []byte, c byte) {
	for i := 0; i < len(s); i++ {
		s[i] = c
	}
}

// Takes a []byte and builds a SecString out of it, wiping str in the
// process. The SecString.String is encrypted after this function.
//
// A SecString should be destroyed when it's no longer needed to prevent memory leaks.
// It is probably a good idea to defer SecString.Destroy()
func NewSecString(str []byte) (*SecString, error) {
	ret := &SecString{Length: len(str)}
	var err error

	if padding := ret.Length % aes.BlockSize; padding != 0 {
		ret.Padding = aes.BlockSize - padding
	}

	ret.String, err = syscall.Mmap(-1, 0, ret.Length+ret.Padding, syscall.PROT_READ|syscall.PROT_WRITE, syscall.MAP_ANON|syscall.MAP_PRIVATE)
	if err != nil {
		return nil, err
	}

	if err := syscall.Mlock(ret.String); err != nil {
		syscall.Munmap(ret.String)
		return nil, err
	}

	for i := 0; i < ret.Length; i++ {
		ret.String[i] = str[i]
		str[i] = 0
	}

	key := make([]byte, 32)
	if _, err := rand.Read(key); err != nil {
		memset(ret.String, 0)
		syscall.Munmap(ret.String)
		return nil, err
	}

	if ret.cipher, err = aes.NewCipher(key); err != nil {
		memset(ret.String, 0)
		syscall.Munmap(ret.String)
		return nil, err
	}

	if err := ret.Encrypt(); err != nil {
		memset(ret.String, 0)
		syscall.Munmap(ret.String)
		return nil, err
	}

	return ret, nil
}

// Makes a new SecString from a string reference. Destroys str after creating
// the secstring
func FromString(str *string) (*SecString, error) {
	b := make([]byte, len(*str))
	for i := 0; i < len(*str); i++ {
		b[i] = (*str)[i]
	}
	*str = strings.Repeat("x", len(*str))
	return NewSecString(b)
}

// Encrypts SecString.String
func (s *SecString) Encrypt() error {
	if err := syscall.Mprotect(s.String, syscall.PROT_READ|syscall.PROT_WRITE); err != nil {
		return err
	}

	s.iv = make([]byte, aes.BlockSize)
	if _, err := rand.Read(s.iv); err != nil {
		return err
	}

	encrypter := cipher.NewOFB(s.cipher, s.iv)
	encrypter.XORKeyStream(s.String, s.String)

	if err := syscall.Mprotect(s.String, syscall.PROT_READ); err != nil {
		return err
	}

	return nil
}

// Decrypt SecString.String for use
func (s *SecString) Decrypt() error {
	if err := syscall.Mprotect(s.String, syscall.PROT_READ|syscall.PROT_WRITE); err != nil {
		return err
	}

	decrypter := cipher.NewOFB(s.cipher, s.iv)
	decrypter.XORKeyStream(s.String, s.String)

	if err := syscall.Mprotect(s.String, syscall.PROT_READ); err != nil {
		return err
	}

	return nil
}

// Destroys the s. *MUST* be called to prevent memory leaks. Probably best to
// be called in a defer
func (s *SecString) Destroy() error {
	if err := syscall.Mprotect(s.String, syscall.PROT_READ|syscall.PROT_WRITE); err != nil {
		return err
	}

	memset(s.String, 0)

	if err := syscall.Munlock(s.String); err != nil {
		return err
	}

	if err := syscall.Munmap(s.String); err != nil {
		return err
	}

	s.String = nil
	return nil
}
<|endoftext|>


<commit_before>package influxquery

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

	"github.com/funkygao/gafka/cmd/kguard/monitor"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/go-metrics"
	log "github.com/funkygao/log4go"
)

func init() {
	monitor.RegisterWatcher("ngx.err", func() monitor.Watcher {
		return &WatchNgxErr{
			Tick: time.Minute,
		}
	})
}

type WatchNgxErr struct {
	Zkzone *zk.ZkZone
	Stop   <-chan struct{}
	Tick   time.Duration
	Wg     *sync.WaitGroup

	addr string
	db   string
}

func (this *WatchNgxErr) Init(ctx monitor.Context) {
	this.Zkzone = ctx.ZkZone()
	this.Stop = ctx.StopChan()
	this.Wg = ctx.Inflight()

	this.addr = ctx.InfluxAddr()
	this.db = "kfk_prod"
}

func (this *WatchNgxErr) Run() {
	defer this.Wg.Done()

	if this.addr == "" || this.db == "" {
		log.Warn("empty addr or db, quit...")
		return
	}

	ticker := time.NewTicker(this.Tick)
	defer ticker.Stop()

	ngerr := metrics.NewRegisteredGauge("_ngx.err", nil)
	for {
		select {
		case <-this.Stop:
			log.Info("ngx.err stopped")
			return

		case <-ticker.C:
			n, err := this.ngixErrSum()
			if err != nil {
				log.Error("ngx.err: %v", err)
			} else {
				ngerr.Update(int64(n))
			}
		}
	}
}

func (this *WatchNgxErr) ngixErrSum() (int, error) {
	res, err := queryInfluxDB(this.addr, this.db,
		fmt.Sprintf(`SELECT mean("m5") FROM "pub.qps.meter" WHERE "appid" = 'logstash' AND "topic" = 'nginx_errlog_intra' AND time > now() - 1m GROUP BY time(1m) fill(0)`))
	if err != nil {
		return 0, err
	}

	total := 0
	for _, row := range res {
		for _, x := range row.Series {
			for _, val := range x.Values {
				// val[0] is time
				n, _ := val[1].(json.Number).Float64()
				total += int(n)
			}
		}
	}

	res, err = queryInfluxDB(this.addr, this.db,
		fmt.Sprintf(`SELECT mean("m5") FROM "pub.qps.meter" WHERE "appid" = 'logstash' AND "topic" = 'nginx_errlog_extra' AND time > now() - 1m GROUP BY time(1m) fill(0)`))
	if err != nil {
		return 0, err
	}

	for _, row := range res {
		for _, x := range row.Series {
			for _, val := range x.Values {
				// val[0] is time
				n, _ := val[1].(json.Number).Float64()
				total += int(n)
			}
		}
	}

	return total, nil
}
<commit_msg>use m1 instead of m5<commit_after>package influxquery

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

	"github.com/funkygao/gafka/cmd/kguard/monitor"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/go-metrics"
	log "github.com/funkygao/log4go"
)

func init() {
	monitor.RegisterWatcher("ngx.err", func() monitor.Watcher {
		return &WatchNgxErr{
			Tick: time.Minute,
		}
	})
}

type WatchNgxErr struct {
	Zkzone *zk.ZkZone
	Stop   <-chan struct{}
	Tick   time.Duration
	Wg     *sync.WaitGroup

	addr string
	db   string
}

func (this *WatchNgxErr) Init(ctx monitor.Context) {
	this.Zkzone = ctx.ZkZone()
	this.Stop = ctx.StopChan()
	this.Wg = ctx.Inflight()

	this.addr = ctx.InfluxAddr()
	this.db = "kfk_prod"
}

func (this *WatchNgxErr) Run() {
	defer this.Wg.Done()

	if this.addr == "" || this.db == "" {
		log.Warn("empty addr or db, quit...")
		return
	}

	ticker := time.NewTicker(this.Tick)
	defer ticker.Stop()

	ngerr := metrics.NewRegisteredGauge("_ngx.err", nil)
	for {
		select {
		case <-this.Stop:
			log.Info("ngx.err stopped")
			return

		case <-ticker.C:
			n, err := this.ngixErrSum()
			if err != nil {
				log.Error("ngx.err: %v", err)
			} else {
				ngerr.Update(int64(n))
			}
		}
	}
}

func (this *WatchNgxErr) ngixErrSum() (int, error) {
	res, err := queryInfluxDB(this.addr, this.db,
		fmt.Sprintf(`SELECT mean("m1") FROM "pub.qps.meter" WHERE "appid" = 'logstash' AND "topic" = 'nginx_errlog_intra' AND time > now() - 1m GROUP BY time(1m) fill(0)`))
	if err != nil {
		return 0, err
	}

	total := 0
	for _, row := range res {
		for _, x := range row.Series {
			for _, val := range x.Values {
				// val[0] is time
				n, _ := val[1].(json.Number).Float64()
				total += int(n)
			}
		}
	}

	res, err = queryInfluxDB(this.addr, this.db,
		fmt.Sprintf(`SELECT mean("m1") FROM "pub.qps.meter" WHERE "appid" = 'logstash' AND "topic" = 'nginx_errlog_extra' AND time > now() - 1m GROUP BY time(1m) fill(0)`))
	if err != nil {
		return 0, err
	}

	for _, row := range res {
		for _, x := range row.Series {
			for _, val := range x.Values {
				// val[0] is time
				n, _ := val[1].(json.Number).Float64()
				total += int(n)
			}
		}
	}

	return total, nil
}
<|endoftext|>


<commit_before>// Copyright (C) 2017 Damon Revoe. All rights reserved.
// Use of this source code is governed by the MIT
// license, which can be found in the LICENSE file.

package main

import (
	"errors"
	"log"
	"os"
	"os/exec"
	"path/filepath"

	"github.com/spf13/cobra"
)

func generateAndBootstrapPackage(pkgSelection []string) error {
	packageIndex, err := buildPackageIndex()
	if err != nil {
		return err
	}

	pkgRootDir := filepath.Join(getWorkspaceDir(), "packages")

	type packageAndGenerator struct {
		pd         *packageDefinition
		packageDir string
		generator  func() error
	}

	var packagesAndGenerators []packageAndGenerator

	for _, packageName := range pkgSelection {
		pd, ok := packageIndex.packageByName[packageName]
		if !ok {
			return errors.New("no such package: " + packageName)
		}

		packageDir := filepath.Join(pkgRootDir, pd.packageName)

		generator, err := pd.getPackageGeneratorFunc(packageDir)
		if err != nil {
			return err
		}

		packagesAndGenerators = append(packagesAndGenerators,
			packageAndGenerator{pd, packageDir, generator})
	}

	for _, pg := range packagesAndGenerators {
		// Generate autoconf and automake sources for the package.
		if err = pg.generator(); err != nil {
			return err
		}

		// Bootstrap the package if the 'configure' script
		// does not exist.
		_, err = os.Lstat(filepath.Join(pg.packageDir, "configure"))
		if os.IsNotExist(err) {
			cmd := exec.Command("./autogen.sh")
			cmd.Dir = pg.packageDir
			if err = cmd.Run(); err != nil {
				return errors.New(filepath.Join(pg.packageDir,
					"autogen.sh") + ": " + err.Error())
			}
		}
	}

	return nil
}

// SelectCmd represents the select command
var selectCmd = &cobra.Command{
	Use:   "select package_range...",
	Short: "Choose one or more packages to work on",
	Args:  cobra.MinimumNArgs(1),
	Run: func(_ *cobra.Command, args []string) {
		if err := generateAndBootstrapPackage(args); err != nil {
			log.Fatal(err)
		}
	},
}

func init() {
	RootCmd.AddCommand(selectCmd)

	selectCmd.Flags().SortFlags = false
	addQuietFlag(selectCmd)
	addWorkspaceDirFlag(selectCmd)
	addPkgPathFlag(selectCmd)
}
<commit_msg>Implement draft of 'configure --help' parsing<commit_after>// Copyright (C) 2017 Damon Revoe. All rights reserved.
// Use of this source code is governed by the MIT
// license, which can be found in the LICENSE file.

package main

import (
	"bufio"
	"errors"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/spf13/cobra"
)

func generateAndBootstrapPackage(pkgSelection []string) error {
	packageIndex, err := buildPackageIndex()
	if err != nil {
		return err
	}

	pkgRootDir := filepath.Join(getWorkspaceDir(), "packages")

	type packageAndGenerator struct {
		pd         *packageDefinition
		packageDir string
		generator  func() error
	}

	var packagesAndGenerators []packageAndGenerator

	for _, packageName := range pkgSelection {
		pd, ok := packageIndex.packageByName[packageName]
		if !ok {
			return errors.New("no such package: " + packageName)
		}

		packageDir := filepath.Join(pkgRootDir, pd.packageName)

		generator, err := pd.getPackageGeneratorFunc(packageDir)
		if err != nil {
			return err
		}

		packagesAndGenerators = append(packagesAndGenerators,
			packageAndGenerator{pd, packageDir, generator})
	}

	optRegexp := regexp.MustCompile(`^--(enable|disable|with|without)` +
		`-([^\s\[=]+)([^\s]*)\s*(.*)$`)

	autoconfOptions := map[string]bool{
		"FEATURE":             true,
		"PACKAGE":             true,
		"aix-soname":          true,
		"dependency-tracking": true,
		"fast-install":        true,
		"gnu-ld":              true,
		"libtool-lock":        true,
		"option-checking":     true,
		"pic":                 true,
		"pkgconfigdir":        true,
		"shared":              true,
		"silent-rules":        true,
		"static":              true,
		"sysroot":             true,
	}

	for _, pg := range packagesAndGenerators {
		// Generate autoconf and automake sources for the package.
		if err = pg.generator(); err != nil {
			return err
		}

		configurePathname := filepath.Join(pg.packageDir, "configure")

		// Bootstrap the package if 'configure' does not exist.
		_, err = os.Lstat(configurePathname)
		if os.IsNotExist(err) {
			bootstrapCmd := exec.Command("./autogen.sh")
			bootstrapCmd.Dir = pg.packageDir
			if err = bootstrapCmd.Run(); err != nil {
				return errors.New(filepath.Join(pg.packageDir,
					"autogen.sh") + ": " + err.Error())
			}
		}

		configureHelpCmd := exec.Command(configurePathname, "--help")
		configureHelpStdout, err := configureHelpCmd.StdoutPipe()
		if err != nil {
			return err
		}
		if err = configureHelpCmd.Start(); err != nil {
			return err
		}
		helpScanner := bufio.NewScanner(configureHelpStdout)
		type optionOrFeature = struct {
			keyword     string
			option      string
			arg         string
			description string
		}
		var options []optionOrFeature
		var currentOption *optionOrFeature

		for helpScanner.Scan() {
			helpLine := strings.TrimRight(helpScanner.Text(), " ")

			if helpLine == "" ||
				!strings.HasPrefix(helpLine, " ") {
				if currentOption != nil {
					options = append(options, *currentOption)
					currentOption = nil
				}
				continue
			}

			helpLine = strings.TrimLeft(helpLine, " ")

			if strings.HasPrefix(helpLine, "--") {
				if currentOption != nil {
					options = append(options, *currentOption)
					currentOption = nil
				}
			} else {
				if currentOption != nil {
					if currentOption.description != "" {
						currentOption.description += " "
					}
					currentOption.description += helpLine
				}
				continue
			}

			matches := optRegexp.FindStringSubmatch(helpLine)

			if len(matches) > 4 {
				option := matches[2]
				if _, present := autoconfOptions[option]; present {
					continue
				}
				currentOption = &optionOrFeature{
					matches[1],
					option,
					matches[3],
					matches[4]}
			}
		}
		if err := helpScanner.Err(); err != nil {
			return err
		}
		if err = configureHelpCmd.Wait(); err != nil {
			return err
		}
		for _, opt := range options {
			fmt.Println(opt.keyword)
			fmt.Println(opt.option)
			fmt.Println(opt.arg)
			fmt.Println(opt.description)
		}
	}

	return nil
}

// SelectCmd represents the select command
var selectCmd = &cobra.Command{
	Use:   "select package_range...",
	Short: "Choose one or more packages to work on",
	Args:  cobra.MinimumNArgs(1),
	Run: func(_ *cobra.Command, args []string) {
		if err := generateAndBootstrapPackage(args); err != nil {
			log.Fatal(err)
		}
	},
}

func init() {
	RootCmd.AddCommand(selectCmd)

	selectCmd.Flags().SortFlags = false
	addQuietFlag(selectCmd)
	addWorkspaceDirFlag(selectCmd)
	addPkgPathFlag(selectCmd)
}
<|endoftext|>


<commit_before>package deploy

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	kube "k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/meta"
	kubecli "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	"k8s.io/kubernetes/pkg/kubectl/cmd/config"
	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/util/strategicpatch"
)

const DefaultContext = ""

// KubeCluster is able to deploy to Kubernetes clusters. This is a very simple implementation with no error recovery.
type KubeCluster struct {
	client  *kubecli.Client
	context string
}

// NewKubeClusterFromContext creates a KubeCluster using a Kubernetes client with the configuration of the given context.
// If the context name is empty, the default context will be used.
func NewKubeClusterFromContext(name string) (*KubeCluster, error) {
	rules := defaultLoadingRules()

	overrides := &clientcmd.ConfigOverrides{
		CurrentContext: name,
	}

	config := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(rules, overrides)

	clientConfig, err := config.ClientConfig()
	if err != nil {
		if len(name) == 0 {
			return nil, fmt.Errorf("could not use default context: %v", err)
		}
		return nil, fmt.Errorf("could not use context `%s`: %v", name, err)
	}

	client, err := kubecli.New(clientConfig)
	if err != nil {
		return nil, fmt.Errorf("could not create Kubernetes client: %v", err)
	}

	return &KubeCluster{
		client:  client,
		context: name,
	}, nil
}

// Context returns the kubectl context being used
func (c *KubeCluster) Context() string {
	return c.context
}

// Deploy creates/updates the Deployment's objects on the Kubernetes cluster. If update is not set, will error if objects exist.
// Currently no error recovery is implemented; if there is an error the deployment process will immediately halt and return the error.
func (c *KubeCluster) Deploy(dep *Deployment, update bool) error {
	if c.client == nil {
		return errors.New("client not setup (was nil)")
	}

	// create namespaces before everything else
	for _, ns := range dep.namespaces {
		_, err := c.client.Namespaces().Create(ns)
		if err != nil && !alreadyExists(err) {
			return err
		}
	}

	// TODO: add continue on error and error lists
	for _, obj := range dep.Objects() {
		// don't create namespaces again
		if _, isNamespace := obj.(*kube.Namespace); isNamespace {
			continue
		}

		err := c.deploy(obj, update)
		if err != nil {
			return err
		}
	}

	// deployed successfully
	return nil
}

// deploy creates the object on the connected Kubernetes instance. Errors if object exists and not updating.
func (c *KubeCluster) deploy(obj KubeObject, update bool) error {
	if obj == nil {
		return errors.New("tried to deploy nil object")
	}

	mapping, err := mapping(obj)
	if err != nil {
		return err
	}

	if update {
		_, err := c.update(obj, true, mapping)
		if err != nil {
			return err
		}
		return nil
	}

	_, err = c.create(obj, mapping)
	return err
}

// update replaces the currently deployed version with a new one. If the objects already match then nothing is done.
func (c *KubeCluster) update(obj KubeObject, create bool, mapping *meta.RESTMapping) (KubeObject, error) {
	meta := obj.GetObjectMeta()

	deployed, err := c.get(meta.GetNamespace(), meta.GetName(), true, mapping)
	if doesNotExist(err) && create {
		return c.create(obj, mapping)
	} else if err != nil {
		return nil, err
	}

	// TODO: need a better way to handle resource versioning
	// set resource version on local to same as remote
	deployedVersion := deployed.GetObjectMeta().GetResourceVersion()
	meta.SetResourceVersion(deployedVersion)

	copyImmutables(deployed, obj)

	// if local matches deployed, do nothing
	if kube.Semantic.DeepEqual(obj, deployed) {
		return deployed, nil
	}

	patch, err := diff(deployed, obj)
	if err != nil {
		return nil, fmt.Errorf("could not create diff: %v", err)
	}

	req := c.client.RESTClient.Patch(kube.StrategicMergePatchType).
		Name(meta.GetName()).
		Body(patch)

	setRequestObjectInfo(req, meta.GetNamespace(), mapping)

	runtimeObj, err := req.Do().Get()
	if err != nil {
		return nil, resourceError("update", meta.GetNamespace(), meta.GetName(), mapping, err)
	}

	return asKubeObject(runtimeObj)
}

// get retrieves the object from the cluster.
func (c *KubeCluster) get(namespace, name string, export bool, mapping *meta.RESTMapping) (KubeObject, error) {
	req := c.client.RESTClient.Get().Name(name)
	setRequestObjectInfo(req, namespace, mapping)

	if export {
		req.Param("export", "true")
	}

	runtimeObj, err := req.Do().Get()
	if err != nil {
		return nil, resourceError("get", namespace, name, mapping, err)
	}

	return asKubeObject(runtimeObj)
}

// create adds the object to the cluster.
func (c *KubeCluster) create(obj KubeObject, mapping *meta.RESTMapping) (KubeObject, error) {
	meta := obj.GetObjectMeta()
	req := c.client.RESTClient.Post().Body(obj)

	setRequestObjectInfo(req, meta.GetNamespace(), mapping)

	runtimeObj, err := req.Do().Get()
	if err != nil {
		return nil, resourceError("create", meta.GetName(), meta.GetNamespace(), mapping, err)
	}

	return asKubeObject(runtimeObj)
}

// setRequestObjectInfo adds necessary type information to requests.
func setRequestObjectInfo(req *kubecli.Request, namespace string, mapping *meta.RESTMapping) {
	// if namespace scoped resource, set namespace
	req.NamespaceIfScoped(namespace, isNamespaceScoped(mapping))

	// set resource name
	req.Resource(mapping.Resource)
}

// alreadyExists checks if the error is for a resource already existing.
func alreadyExists(err error) bool {
	if err == nil {
		return false
	}
	return strings.HasSuffix(err.Error(), "already exists")
}

// doesNotExist checks if the error is for a non-existent resource.
func doesNotExist(err error) bool {
	if err == nil {
		return false
	}
	return strings.HasSuffix(err.Error(), "not found")
}

// mapping returns the appropriate RESTMapping for the object.
func mapping(obj KubeObject) (*meta.RESTMapping, error) {
	gvk, err := kube.Scheme.ObjectKind(obj)
	if err != nil {
		return nil, err
	}

	mapping, err := kube.RESTMapper.RESTMapping(gvk.GroupKind(), gvk.Version)
	if err != nil {
		return nil, fmt.Errorf("could not create RESTMapping for %s: %v", gvk, err)
	}
	return mapping, nil
}

// isNamespaceScoped returns if the mapping is scoped by Namespace.
func isNamespaceScoped(mapping *meta.RESTMapping) bool {
	return mapping.Scope.Name() == meta.RESTScopeNameNamespace
}

// defaultLoadingRules use the same rules (as of 2/17/16) as kubectl.
func defaultLoadingRules() *clientcmd.ClientConfigLoadingRules {
	opts := config.NewDefaultPathOptions()

	loadingRules := opts.LoadingRules
	loadingRules.Precedence = opts.GetLoadingPrecedence()
	return loadingRules
}

// diff creates a patch.
func diff(original, modified runtime.Object) (patch []byte, err error) {
	origBytes, err := json.Marshal(original)
	if err != nil {
		return nil, err
	}

	modBytes, err := json.Marshal(modified)
	if err != nil {
		return nil, err
	}

	return strategicpatch.CreateTwoWayMergePatch(origBytes, modBytes, original)
}

// asKubeObject attempts use the object as a KubeObject. It will return an error if not possible.
func asKubeObject(runtimeObj runtime.Object) (KubeObject, error) {
	kubeObj, ok := runtimeObj.(KubeObject)
	if !ok {
		return nil, errors.New("was unable to use runtime.Object as deploy.KubeObject")
	}
	return kubeObj, nil
}

func resourceError(action, namespace, name string, mapping *meta.RESTMapping, err error) error {
	if mapping == nil || mapping.GroupVersionKind.IsEmpty() {
		return fmt.Errorf("could not %s '%s/%s': %v", action, namespace, name, err)
	}
	gvk := mapping.GroupVersionKind
	return fmt.Errorf("could not %s '%s/%s' (%s): %v", action, namespace, name, gvk.Kind, err)
}

// copyImmutables sets any immutable fields from src on dst. Will panic if objects not of same type.
func copyImmutables(src, dst KubeObject) {
	if src == nil || dst == nil {
		return
	}

	// each type has specific fields that must be copied
	switch src := src.(type) {
	case *kube.Service:
		dst := dst.(*kube.Service)
		dst.Spec.ClusterIP = src.Spec.ClusterIP
	}
}
<commit_msg>now displays endpoints while deploying<commit_after>package deploy

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	kube "k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/meta"
	kubecli "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	"k8s.io/kubernetes/pkg/kubectl/cmd/config"
	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/util/strategicpatch"
)

const DefaultContext = ""

// KubeCluster is able to deploy to Kubernetes clusters. This is a very simple implementation with no error recovery.
type KubeCluster struct {
	client  *kubecli.Client
	context string
}

// NewKubeClusterFromContext creates a KubeCluster using a Kubernetes client with the configuration of the given context.
// If the context name is empty, the default context will be used.
func NewKubeClusterFromContext(name string) (*KubeCluster, error) {
	rules := defaultLoadingRules()

	overrides := &clientcmd.ConfigOverrides{
		CurrentContext: name,
	}

	config := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(rules, overrides)

	clientConfig, err := config.ClientConfig()
	if err != nil {
		if len(name) == 0 {
			return nil, fmt.Errorf("could not use default context: %v", err)
		}
		return nil, fmt.Errorf("could not use context `%s`: %v", name, err)
	}

	client, err := kubecli.New(clientConfig)
	if err != nil {
		return nil, fmt.Errorf("could not create Kubernetes client: %v", err)
	}

	return &KubeCluster{
		client:  client,
		context: name,
	}, nil
}

// Context returns the kubectl context being used
func (c *KubeCluster) Context() string {
	return c.context
}

// Deploy creates/updates the Deployment's objects on the Kubernetes cluster. If update is not set, will error if objects exist.
// Currently no error recovery is implemented; if there is an error the deployment process will immediately halt and return the error.
func (c *KubeCluster) Deploy(dep *Deployment, update bool) error {
	if c.client == nil {
		return errors.New("client not setup (was nil)")
	}

	// create namespaces before everything else
	for _, ns := range dep.namespaces {
		_, err := c.client.Namespaces().Create(ns)
		if err != nil && !alreadyExists(err) {
			return err
		}
	}

	// TODO: add continue on error and error lists
	for _, obj := range dep.Objects() {
		// don't create namespaces again
		if _, isNamespace := obj.(*kube.Namespace); isNamespace {
			continue
		}

		err := c.deploy(obj, update)
		if err != nil {
			return err
		}
	}

	printLoadBalancers(c.client, dep.services)

	// deployed successfully
	return nil
}

// deploy creates the object on the connected Kubernetes instance. Errors if object exists and not updating.
func (c *KubeCluster) deploy(obj KubeObject, update bool) error {
	if obj == nil {
		return errors.New("tried to deploy nil object")
	}

	mapping, err := mapping(obj)
	if err != nil {
		return err
	}

	if update {
		_, err := c.update(obj, true, mapping)
		if err != nil {
			return err
		}
		return nil
	}

	_, err = c.create(obj, mapping)
	return err
}

// update replaces the currently deployed version with a new one. If the objects already match then nothing is done.
func (c *KubeCluster) update(obj KubeObject, create bool, mapping *meta.RESTMapping) (KubeObject, error) {
	meta := obj.GetObjectMeta()

	deployed, err := c.get(meta.GetNamespace(), meta.GetName(), true, mapping)
	if doesNotExist(err) && create {
		return c.create(obj, mapping)
	} else if err != nil {
		return nil, err
	}

	// TODO: need a better way to handle resource versioning
	// set resource version on local to same as remote
	deployedVersion := deployed.GetObjectMeta().GetResourceVersion()
	meta.SetResourceVersion(deployedVersion)

	copyImmutables(deployed, obj)

	// if local matches deployed, do nothing
	if kube.Semantic.DeepEqual(obj, deployed) {
		return deployed, nil
	}

	patch, err := diff(deployed, obj)
	if err != nil {
		return nil, fmt.Errorf("could not create diff: %v", err)
	}

	req := c.client.RESTClient.Patch(kube.StrategicMergePatchType).
		Name(meta.GetName()).
		Body(patch)

	setRequestObjectInfo(req, meta.GetNamespace(), mapping)

	runtimeObj, err := req.Do().Get()
	if err != nil {
		return nil, resourceError("update", meta.GetNamespace(), meta.GetName(), mapping, err)
	}

	return asKubeObject(runtimeObj)
}

// get retrieves the object from the cluster.
func (c *KubeCluster) get(namespace, name string, export bool, mapping *meta.RESTMapping) (KubeObject, error) {
	req := c.client.RESTClient.Get().Name(name)
	setRequestObjectInfo(req, namespace, mapping)

	if export {
		req.Param("export", "true")
	}

	runtimeObj, err := req.Do().Get()
	if err != nil {
		return nil, resourceError("get", namespace, name, mapping, err)
	}

	return asKubeObject(runtimeObj)
}

// create adds the object to the cluster.
func (c *KubeCluster) create(obj KubeObject, mapping *meta.RESTMapping) (KubeObject, error) {
	meta := obj.GetObjectMeta()
	req := c.client.RESTClient.Post().Body(obj)

	setRequestObjectInfo(req, meta.GetNamespace(), mapping)

	runtimeObj, err := req.Do().Get()
	if err != nil {
		return nil, resourceError("create", meta.GetName(), meta.GetNamespace(), mapping, err)
	}

	return asKubeObject(runtimeObj)
}

// setRequestObjectInfo adds necessary type information to requests.
func setRequestObjectInfo(req *kubecli.Request, namespace string, mapping *meta.RESTMapping) {
	// if namespace scoped resource, set namespace
	req.NamespaceIfScoped(namespace, isNamespaceScoped(mapping))

	// set resource name
	req.Resource(mapping.Resource)
}

// alreadyExists checks if the error is for a resource already existing.
func alreadyExists(err error) bool {
	if err == nil {
		return false
	}
	return strings.HasSuffix(err.Error(), "already exists")
}

// doesNotExist checks if the error is for a non-existent resource.
func doesNotExist(err error) bool {
	if err == nil {
		return false
	}
	return strings.HasSuffix(err.Error(), "not found")
}

// mapping returns the appropriate RESTMapping for the object.
func mapping(obj KubeObject) (*meta.RESTMapping, error) {
	gvk, err := kube.Scheme.ObjectKind(obj)
	if err != nil {
		return nil, err
	}

	mapping, err := kube.RESTMapper.RESTMapping(gvk.GroupKind(), gvk.Version)
	if err != nil {
		return nil, fmt.Errorf("could not create RESTMapping for %s: %v", gvk, err)
	}
	return mapping, nil
}

// isNamespaceScoped returns if the mapping is scoped by Namespace.
func isNamespaceScoped(mapping *meta.RESTMapping) bool {
	return mapping.Scope.Name() == meta.RESTScopeNameNamespace
}

// defaultLoadingRules use the same rules (as of 2/17/16) as kubectl.
func defaultLoadingRules() *clientcmd.ClientConfigLoadingRules {
	opts := config.NewDefaultPathOptions()

	loadingRules := opts.LoadingRules
	loadingRules.Precedence = opts.GetLoadingPrecedence()
	return loadingRules
}

// diff creates a patch.
func diff(original, modified runtime.Object) (patch []byte, err error) {
	origBytes, err := json.Marshal(original)
	if err != nil {
		return nil, err
	}

	modBytes, err := json.Marshal(modified)
	if err != nil {
		return nil, err
	}

	return strategicpatch.CreateTwoWayMergePatch(origBytes, modBytes, original)
}

// asKubeObject attempts use the object as a KubeObject. It will return an error if not possible.
func asKubeObject(runtimeObj runtime.Object) (KubeObject, error) {
	kubeObj, ok := runtimeObj.(KubeObject)
	if !ok {
		return nil, errors.New("was unable to use runtime.Object as deploy.KubeObject")
	}
	return kubeObj, nil
}

func resourceError(action, namespace, name string, mapping *meta.RESTMapping, err error) error {
	if mapping == nil || mapping.GroupVersionKind.IsEmpty() {
		return fmt.Errorf("could not %s '%s/%s': %v", action, namespace, name, err)
	}
	gvk := mapping.GroupVersionKind
	return fmt.Errorf("could not %s '%s/%s' (%s): %v", action, namespace, name, gvk.Kind, err)
}

// copyImmutables sets any immutable fields from src on dst. Will panic if objects not of same type.
func copyImmutables(src, dst KubeObject) {
	if src == nil || dst == nil {
		return
	}

	// each type has specific fields that must be copied
	switch src := src.(type) {
	case *kube.Service:
		dst := dst.(*kube.Service)
		dst.Spec.ClusterIP = src.Spec.ClusterIP
	}
}

func printLoadBalancers(client *kubecli.Client, services []*kube.Service) {
	if len(services) == 0 {
		return
	}

	first := true
	completed := map[string]bool{}

	// checks when we've seen every service
	done := func () bool {
		for _, s := range services {
			if s.Spec.Type == kube.ServiceTypeLoadBalancer && !completed[s.Name] {
				return false
			}
		}
		return true
	}

	for {
		if done() {
			return
		}

		if first {
			fmt.Println("Waiting for load balancer deployment...")
			first = false
		}

		for _, s := range services {
			if s.Spec.Type == kube.ServiceTypeLoadBalancer && !completed[s.Name] {
				clusterVers, err := client.Services(s.Namespace).Get(s.Name)
				if err != nil {
					fmt.Printf("Error getting service `%s`: %v\n", s.Name, err)
				}

				loadBalancers := clusterVers.Status.LoadBalancer.Ingress
				if len (loadBalancers) == 1 {
					completed[s.Name] = true
					fmt.Printf("Service '%s/%s' available at: \t%s\n", s.Namespace, s.Name,loadBalancers[0].IP)
				}
			}
		}
	}
}<|endoftext|>


<commit_before>/*
Copyright 2017 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.
*/

package util

import (
	"testing"
)

func TestValidateExactArgNumber(t *testing.T) {
	var tests = []struct {
		name                string
		args, supportedArgs []string
		expectedErr         bool
	}{
		{
			name:          "one arg given and one arg expected",
			args:          []string{"my-node-1234"},
			supportedArgs: []string{"node-name"},
			expectedErr:   false,
		},
		{
			name:          "two args given and two args expected",
			args:          []string{"my-node-1234", "foo"},
			supportedArgs: []string{"node-name", "second-toplevel-arg"},
			expectedErr:   false,
		},
		{
			name:          "too few supplied args",
			args:          []string{},
			supportedArgs: []string{"node-name"},
			expectedErr:   true,
		},
		{
			name:          "too few non-empty args",
			args:          []string{""},
			supportedArgs: []string{"node-name"},
			expectedErr:   true,
		},
		{
			name:          "too many args",
			args:          []string{"my-node-1234", "foo"},
			supportedArgs: []string{"node-name"},
			expectedErr:   true,
		},
	}
	for _, rt := range tests {
		t.Run(rt.name, func(t *testing.T) {
			actual := ValidateExactArgNumber(rt.args, rt.supportedArgs)
			if (actual != nil) != rt.expectedErr {
				t.Errorf(
					"failed ValidateExactArgNumber:\n\texpected error: %t\n\t  actual error: %t",
					rt.expectedErr,
					(actual != nil),
				)
			}
		})
	}
}
<commit_msg>Add test for cmd/kubeadm/app/cmd/util/cmdutil.go<commit_after>/*
Copyright 2017 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.
*/

package util

import (
	"k8s.io/client-go/tools/clientcmd"
	"testing"
)

func TestValidateExactArgNumber(t *testing.T) {
	var tests = []struct {
		name                string
		args, supportedArgs []string
		expectedErr         bool
	}{
		{
			name:          "one arg given and one arg expected",
			args:          []string{"my-node-1234"},
			supportedArgs: []string{"node-name"},
			expectedErr:   false,
		},
		{
			name:          "two args given and two args expected",
			args:          []string{"my-node-1234", "foo"},
			supportedArgs: []string{"node-name", "second-toplevel-arg"},
			expectedErr:   false,
		},
		{
			name:          "too few supplied args",
			args:          []string{},
			supportedArgs: []string{"node-name"},
			expectedErr:   true,
		},
		{
			name:          "too few non-empty args",
			args:          []string{""},
			supportedArgs: []string{"node-name"},
			expectedErr:   true,
		},
		{
			name:          "too many args",
			args:          []string{"my-node-1234", "foo"},
			supportedArgs: []string{"node-name"},
			expectedErr:   true,
		},
	}
	for _, rt := range tests {
		t.Run(rt.name, func(t *testing.T) {
			actual := ValidateExactArgNumber(rt.args, rt.supportedArgs)
			if (actual != nil) != rt.expectedErr {
				t.Errorf(
					"failed ValidateExactArgNumber:\n\texpected error: %t\n\t  actual error: %t",
					rt.expectedErr,
					(actual != nil),
				)
			}
		})
	}
}

func TestGetKubeConfigPath(t *testing.T) {
	var tests = []struct {
		name     string
		file     string
		expected string
	}{
		{
			name:     "provide an empty value",
			file:     "",
			expected: clientcmd.NewDefaultClientConfigLoadingRules().GetDefaultFilename(),
		},
		{
			name:     "provide a non-empty value",
			file:     "kubelet.kubeconfig",
			expected: "kubelet.kubeconfig",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if GetKubeConfigPath(tt.file) != tt.expected {
				t.Error("unexpected result")
			}
		})
	}
}
<|endoftext|>


<commit_before>package etwlogrus

import (
	"sort"

	"github.com/Microsoft/go-winio/pkg/etw"
	"github.com/sirupsen/logrus"
)

// Hook is a Logrus hook which logs received events to ETW.
type Hook struct {
	provider      *etw.Provider
	closeProvider bool
}

// NewHook registers a new ETW provider and returns a hook to log from it. The
// provider will be closed when the hook is closed.
func NewHook(providerName string) (*Hook, error) {
	provider, err := etw.NewProvider(providerName, nil)
	if err != nil {
		return nil, err
	}

	return &Hook{provider, true}, nil
}

// NewHookFromProvider creates a new hook based on an existing ETW provider. The
// provider will not be closed when the hook is closed.
func NewHookFromProvider(provider *etw.Provider) (*Hook, error) {
	return &Hook{provider, false}, nil
}

// Levels returns the set of levels that this hook wants to receive log entries
// for.
func (h *Hook) Levels() []logrus.Level {
	return []logrus.Level{
		logrus.TraceLevel,
		logrus.DebugLevel,
		logrus.InfoLevel,
		logrus.WarnLevel,
		logrus.ErrorLevel,
		logrus.FatalLevel,
		logrus.PanicLevel,
	}
}

var logrusToETWLevelMap = map[logrus.Level]etw.Level{
	logrus.PanicLevel: etw.LevelAlways,
	logrus.FatalLevel: etw.LevelCritical,
	logrus.ErrorLevel: etw.LevelError,
	logrus.WarnLevel:  etw.LevelWarning,
	logrus.InfoLevel:  etw.LevelInfo,
	logrus.DebugLevel: etw.LevelVerbose,
	logrus.TraceLevel: etw.LevelVerbose,
}

// Fire receives each Logrus entry as it is logged, and logs it to ETW.
func (h *Hook) Fire(e *logrus.Entry) error {
	// Logrus defines more levels than ETW typically uses, but analysis is
	// easiest when using a consistent set of levels across ETW providers, so we
	// map the Logrus levels to ETW levels.
	level := logrusToETWLevelMap[e.Level]
	if !h.provider.IsEnabledForLevel(level) {
		return nil
	}

	// Sort the fields by name so they are consistent in each instance
	// of an event. Otherwise, the fields don't line up in WPA.
	names := make([]string, 0, len(e.Data))
	hasError := false
	for k := range e.Data {
		names := make([]string, 0, len(e.Data))
		if k == logrus.ErrorKey {
			// Always put the error last because it is optional in some events.
			hasError = true
		} else {
			names = append(names, k)
		}
	}
	sort.Strings(names)

	// Reserve extra space for the message field.
	fields := make([]etw.FieldOpt, 0, len(e.Data)+1)
	fields = append(fields, etw.StringField("Message", e.Message))
	for _, k := range names {
		fields = append(fields, etw.SmartField(k, e.Data[k]))
	}
	if hasError {
		fields = append(fields, etw.SmartField(logrus.ErrorKey, e.Data[logrus.ErrorKey]))
	}

	return h.provider.WriteEvent(
		"LogrusEntry",
		etw.WithEventOpts(etw.WithLevel(level)),
		fields)
}

// Close cleans up the hook and closes the ETW provider. If the provder was
// registered by etwlogrus, it will be closed as part of `Close`. If the
// provider was passed in, it will not be closed.
func (h *Hook) Close() error {
	if h.closeProvider {
		return h.provider.Close()
	}
	return nil
}
<commit_msg>etwlogrus: Ignore error from ETW WriteEvent<commit_after>package etwlogrus

import (
	"sort"

	"github.com/Microsoft/go-winio/pkg/etw"
	"github.com/sirupsen/logrus"
)

// Hook is a Logrus hook which logs received events to ETW.
type Hook struct {
	provider      *etw.Provider
	closeProvider bool
}

// NewHook registers a new ETW provider and returns a hook to log from it. The
// provider will be closed when the hook is closed.
func NewHook(providerName string) (*Hook, error) {
	provider, err := etw.NewProvider(providerName, nil)
	if err != nil {
		return nil, err
	}

	return &Hook{provider, true}, nil
}

// NewHookFromProvider creates a new hook based on an existing ETW provider. The
// provider will not be closed when the hook is closed.
func NewHookFromProvider(provider *etw.Provider) (*Hook, error) {
	return &Hook{provider, false}, nil
}

// Levels returns the set of levels that this hook wants to receive log entries
// for.
func (h *Hook) Levels() []logrus.Level {
	return []logrus.Level{
		logrus.TraceLevel,
		logrus.DebugLevel,
		logrus.InfoLevel,
		logrus.WarnLevel,
		logrus.ErrorLevel,
		logrus.FatalLevel,
		logrus.PanicLevel,
	}
}

var logrusToETWLevelMap = map[logrus.Level]etw.Level{
	logrus.PanicLevel: etw.LevelAlways,
	logrus.FatalLevel: etw.LevelCritical,
	logrus.ErrorLevel: etw.LevelError,
	logrus.WarnLevel:  etw.LevelWarning,
	logrus.InfoLevel:  etw.LevelInfo,
	logrus.DebugLevel: etw.LevelVerbose,
	logrus.TraceLevel: etw.LevelVerbose,
}

// Fire receives each Logrus entry as it is logged, and logs it to ETW.
func (h *Hook) Fire(e *logrus.Entry) error {
	// Logrus defines more levels than ETW typically uses, but analysis is
	// easiest when using a consistent set of levels across ETW providers, so we
	// map the Logrus levels to ETW levels.
	level := logrusToETWLevelMap[e.Level]
	if !h.provider.IsEnabledForLevel(level) {
		return nil
	}

	// Sort the fields by name so they are consistent in each instance
	// of an event. Otherwise, the fields don't line up in WPA.
	names := make([]string, 0, len(e.Data))
	hasError := false
	for k := range e.Data {
		names := make([]string, 0, len(e.Data))
		if k == logrus.ErrorKey {
			// Always put the error last because it is optional in some events.
			hasError = true
		} else {
			names = append(names, k)
		}
	}
	sort.Strings(names)

	// Reserve extra space for the message field.
	fields := make([]etw.FieldOpt, 0, len(e.Data)+1)
	fields = append(fields, etw.StringField("Message", e.Message))
	for _, k := range names {
		fields = append(fields, etw.SmartField(k, e.Data[k]))
	}
	if hasError {
		fields = append(fields, etw.SmartField(logrus.ErrorKey, e.Data[logrus.ErrorKey]))
	}

	// Firing an ETW event is essentially best effort, as the event write can
	// fail for reasons completely out of the control of the event writer (such
	// as a session listening for the event having no available space in its
	// buffers). Therefore, we don't return the error from WriteEvent, as it is
	// just noise in many cases.
	h.provider.WriteEvent(
		"LogrusEntry",
		etw.WithEventOpts(etw.WithLevel(level)),
		fields)

	return nil
}

// Close cleans up the hook and closes the ETW provider. If the provder was
// registered by etwlogrus, it will be closed as part of `Close`. If the
// provider was passed in, it will not be closed.
func (h *Hook) Close() error {
	if h.closeProvider {
		return h.provider.Close()
	}
	return nil
}
<|endoftext|>


<commit_before>package flagparser

import (
	"fmt"
	"strings"

	"github.com/docker/swarmkit/api"
	"github.com/spf13/pflag"
)

// parseNpipe only supports a very simple version of anonymous npipes for
// testing the most basic of data flows. Replace with a --mount flag, similar
// to what we have in docker service.
func parseNpipe(flags *pflag.FlagSet, spec *api.ServiceSpec) error {
	if flags.Changed("npipe") {
		npipes, err := flags.GetStringSlice("npipe")
		if err != nil {
			return err
		}

		container := spec.Task.GetContainer()

		for _, npipe := range npipes {
			if strings.Contains(npipe, ":") {
				return fmt.Errorf("npipe format %q not supported", npipe)
			}
			container.Mounts = append(container.Mounts, api.Mount{
				Type:   api.MountTypeNamedPipe,
				Target: npipe,
			})
		}
	}

	return nil
}
<commit_msg>Corrected --npipe parameter parsing<commit_after>package flagparser

import (
	"fmt"
	"strings"

	"github.com/docker/swarmkit/api"
	"github.com/spf13/pflag"
)

// parseNpipe only supports a very simple version of anonymous npipes for
// testing the most basic of data flows. Replace with a --mount flag, similar
// to what we have in docker service.
func parseNpipe(flags *pflag.FlagSet, spec *api.ServiceSpec) error {
	if flags.Changed("npipe") {
		npipes, err := flags.GetStringSlice("npipe")
		if err != nil {
			return err
		}

		container := spec.Task.GetContainer()

		for _, npipe := range npipes {
			parts := strings.SplitN(npipe, ":", 2)
			if len(parts) != 2 {
				return fmt.Errorf("npipe format %q not supported", npipe)
			}
			container.Mounts = append(container.Mounts, api.Mount{
				Type:   api.MountTypeNamedPipe,
				Source: parts[0],
				Target: parts[1],
			})
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package glob

import (
	"os"
	"reflect"
	"sort"
	"strings"
	"testing"

	"src.elv.sh/pkg/testutil"
)

var (
	mkdirs = []string{"a", "b", "c", "d1", "d1/e", "d1/e/f", "d1/e/f/g",
		"d2", "d2/e", "d2/e/f", "d2/e/f/g"}
	mkdirDots = []string{".el"}
	creates   = []string{"a/X", "a/Y", "b/X", "c/Y",
		"dX", "dXY",
		"lorem", "ipsum",
		"d1/e/f/g/X", "d2/e/f/g/X"}
	createDots = []string{".x", ".el/x"}
)

type globCase struct {
	pattern string
	want    []string
}

var globCases = []globCase{
	{"*", []string{"a", "b", "c", "d1", "d2", "dX", "dXY", "lorem", "ipsum"}},
	{".", []string{"."}},
	{"./*", []string{"./a", "./b", "./c", "./d1", "./d2", "./dX", "./dXY", "./lorem", "./ipsum"}},
	{"..", []string{".."}},
	{"a/..", []string{"a/.."}},
	{"a/../*", []string{"a/../a", "a/../b", "a/../c", "a/../d1", "a/../d2", "a/../dX", "a/../dXY", "a/../lorem", "a/../ipsum"}},
	{"*/", []string{"a/", "b/", "c/", "d1/", "d2/"}},
	{"**", append(mkdirs, creates...)},
	{"*/X", []string{"a/X", "b/X"}},
	{"**X", []string{"a/X", "b/X", "dX", "d1/e/f/g/X", "d2/e/f/g/X"}},
	{"*/*/*", []string{"d1/e/f", "d2/e/f"}},
	{"l*m", []string{"lorem"}},
	{"d*", []string{"d1", "d2", "dX", "dXY"}},
	{"d*/", []string{"d1/", "d2/"}},
	{"d**", []string{"d1", "d1/e", "d1/e/f", "d1/e/f/g", "d1/e/f/g/X",
		"d2", "d2/e", "d2/e/f", "d2/e/f/g", "d2/e/f/g/X", "dX", "dXY"}},
	{"?", []string{"a", "b", "c"}},
	{"??", []string{"d1", "d2", "dX"}},

	// Nonexistent paths.
	{"xxxx", []string{}},
	{"xxxx/*", []string{}},
	{"a/*/", []string{}},

	// TODO Test cases against dotfiles.
}

func TestGlob_Relative(t *testing.T) {
	testGlob(t, false)
}

func TestGlob_Absolute(t *testing.T) {
	testGlob(t, true)
}

func testGlob(t *testing.T, abs bool) {
	dir, cleanup := testutil.InTestDir()
	dir = strings.ReplaceAll(dir, string(os.PathSeparator), "/")
	defer cleanup()

	for _, dir := range append(mkdirs, mkdirDots...) {
		err := os.Mkdir(dir, 0755)
		if err != nil {
			panic(err)
		}
	}
	for _, file := range append(creates, createDots...) {
		f, err := os.Create(file)
		if err != nil {
			panic(err)
		}
		f.Close()
	}
	for _, tc := range globCases {
		pattern := tc.pattern
		if abs {
			pattern = dir + "/" + pattern
		}
		wantResults := make([]string, len(tc.want))
		for i, result := range tc.want {
			if abs {
				wantResults[i] = dir + "/" + result
			} else {
				wantResults[i] = result
			}
		}
		sort.Strings(wantResults)

		results := globPaths(pattern)

		if !reflect.DeepEqual(results, wantResults) {
			t.Errorf(`Glob(%q) => %v, want %v`, pattern, results, wantResults)
		}
	}
}

// Regression test for b.elv.sh/1220
func TestGlob_InvalidUTF8InFilename(t *testing.T) {
	_, cleanup := testutil.InTestDir()
	defer cleanup()

	name := string([]byte{255}) + ".c"
	f, err := os.Create(name)
	if err != nil {
		// The system may refuse to create a file whose name is not UTF-8. This
		// happens on macOS 11 with an APFS filesystem.
		t.Skip("create: ", err)
	}
	f.Close()

	_, err = os.Stat(name)
	if err != nil {
		// The system may pretend to have created the file successfully,
		// but substitute the invalid sequences with U+FFFD. This happens on
		// Windows 10 with an NTFS filesystem.
		t.Skip("stat: ", err)
	}

	paths := globPaths("*.c")
	wantPaths := []string{name}
	if !reflect.DeepEqual(paths, wantPaths) {
		t.Errorf("got %v, want %v", paths, wantPaths)
	}
}

func globPaths(pattern string) []string {
	paths := []string{}
	Glob(pattern, func(pathInfo PathInfo) bool {
		paths = append(paths, pathInfo.Path)
		return true
	})
	sort.Strings(paths)
	return paths
}
<commit_msg>Actually make the test reproduce the crash described in #1220.<commit_after>package glob

import (
	"os"
	"reflect"
	"sort"
	"strings"
	"testing"

	"src.elv.sh/pkg/testutil"
)

var (
	mkdirs = []string{"a", "b", "c", "d1", "d1/e", "d1/e/f", "d1/e/f/g",
		"d2", "d2/e", "d2/e/f", "d2/e/f/g"}
	mkdirDots = []string{".el"}
	creates   = []string{"a/X", "a/Y", "b/X", "c/Y",
		"dX", "dXY",
		"lorem", "ipsum",
		"d1/e/f/g/X", "d2/e/f/g/X"}
	createDots = []string{".x", ".el/x"}
)

type globCase struct {
	pattern string
	want    []string
}

var globCases = []globCase{
	{"*", []string{"a", "b", "c", "d1", "d2", "dX", "dXY", "lorem", "ipsum"}},
	{".", []string{"."}},
	{"./*", []string{"./a", "./b", "./c", "./d1", "./d2", "./dX", "./dXY", "./lorem", "./ipsum"}},
	{"..", []string{".."}},
	{"a/..", []string{"a/.."}},
	{"a/../*", []string{"a/../a", "a/../b", "a/../c", "a/../d1", "a/../d2", "a/../dX", "a/../dXY", "a/../lorem", "a/../ipsum"}},
	{"*/", []string{"a/", "b/", "c/", "d1/", "d2/"}},
	{"**", append(mkdirs, creates...)},
	{"*/X", []string{"a/X", "b/X"}},
	{"**X", []string{"a/X", "b/X", "dX", "d1/e/f/g/X", "d2/e/f/g/X"}},
	{"*/*/*", []string{"d1/e/f", "d2/e/f"}},
	{"l*m", []string{"lorem"}},
	{"d*", []string{"d1", "d2", "dX", "dXY"}},
	{"d*/", []string{"d1/", "d2/"}},
	{"d**", []string{"d1", "d1/e", "d1/e/f", "d1/e/f/g", "d1/e/f/g/X",
		"d2", "d2/e", "d2/e/f", "d2/e/f/g", "d2/e/f/g/X", "dX", "dXY"}},
	{"?", []string{"a", "b", "c"}},
	{"??", []string{"d1", "d2", "dX"}},

	// Nonexistent paths.
	{"xxxx", []string{}},
	{"xxxx/*", []string{}},
	{"a/*/", []string{}},

	// TODO Test cases against dotfiles.
}

func TestGlob_Relative(t *testing.T) {
	testGlob(t, false)
}

func TestGlob_Absolute(t *testing.T) {
	testGlob(t, true)
}

func testGlob(t *testing.T, abs bool) {
	dir, cleanup := testutil.InTestDir()
	dir = strings.ReplaceAll(dir, string(os.PathSeparator), "/")
	defer cleanup()

	for _, dir := range append(mkdirs, mkdirDots...) {
		err := os.Mkdir(dir, 0755)
		if err != nil {
			panic(err)
		}
	}
	for _, file := range append(creates, createDots...) {
		f, err := os.Create(file)
		if err != nil {
			panic(err)
		}
		f.Close()
	}
	for _, tc := range globCases {
		pattern := tc.pattern
		if abs {
			pattern = dir + "/" + pattern
		}
		wantResults := make([]string, len(tc.want))
		for i, result := range tc.want {
			if abs {
				wantResults[i] = dir + "/" + result
			} else {
				wantResults[i] = result
			}
		}
		sort.Strings(wantResults)

		results := globPaths(pattern)

		if !reflect.DeepEqual(results, wantResults) {
			t.Errorf(`Glob(%q) => %v, want %v`, pattern, results, wantResults)
		}
	}
}

// Regression test for b.elv.sh/1220
func TestGlob_InvalidUTF8InFilename(t *testing.T) {
	_, cleanup := testutil.InTestDir()
	defer cleanup()

	name := string([]byte{255}) + "x"
	f, err := os.Create(name)
	if err != nil {
		// The system may refuse to create a file whose name is not UTF-8. This
		// happens on macOS 11 with an APFS filesystem.
		t.Skip("create: ", err)
	}
	f.Close()

	_, err = os.Stat(name)
	if err != nil {
		// The system may pretend to have created the file successfully,
		// but substitute the invalid sequences with U+FFFD. This happens on
		// Windows 10 with an NTFS filesystem.
		t.Skip("stat: ", err)
	}

	paths := globPaths("*x")
	wantPaths := []string{name}
	if !reflect.DeepEqual(paths, wantPaths) {
		t.Errorf("got %v, want %v", paths, wantPaths)
	}
}

func globPaths(pattern string) []string {
	paths := []string{}
	Glob(pattern, func(pathInfo PathInfo) bool {
		paths = append(paths, pathInfo.Path)
		return true
	})
	sort.Strings(paths)
	return paths
}
<|endoftext|>


<commit_before>package api

import (
	"fmt"
	"os"
	"reflect"
	"strconv"
	"testing"
)

func TestFetchBrokerByID(t *testing.T) {
	if os.Getenv("CIRCONUS_API_TOKEN") == "" {
		t.Skip("skipping test; $CIRCONUS_API_TOKEN not set")
	}
	if os.Getenv("CIRC_API_TEST_BROKER_ID") == "" {
		t.Skip("skipping test; $CIRC_API_TEST_BROKER_ID not set")
	}

	t.Log("Testing correct return from API call")

	ac := &Config{}
	ac.TokenKey = os.Getenv("CIRCONUS_API_TOKEN")
	apih, err := NewAPI(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	bid := os.Getenv("CIRC_API_TEST_BROKER_ID")
	if bid == "" {
		t.Fatal("Invalid broker id (empty)")
	}

	id, err := strconv.Atoi(bid)
	if err != nil {
		t.Fatalf("Unable to convert broker id %s to int", bid)
	}

	brokerID := BrokerIDType(id)

	broker, err := apih.FetchBrokerByID(brokerID)
	if err != nil {
		t.Fatalf("Expected no error, got '%v'", err)
	}

	actualType := reflect.TypeOf(broker)
	expectedType := "*api.Broker"
	if actualType.String() != expectedType {
		t.Errorf("Expected %s, got %s", expectedType, actualType.String())
	}

	expectedCid := fmt.Sprintf("/broker/%s", strconv.Itoa(int(brokerID)))
	if broker.Cid != expectedCid {
		t.Fatalf("%s != %s", broker.Cid, expectedCid)
	}

	t.Logf("Broker returned %s %s", broker.Name, broker.Cid)
}

func TestFetchBrokerByCID(t *testing.T) {
	if os.Getenv("CIRCONUS_API_TOKEN") == "" {
		t.Skip("skipping test; $CIRCONUS_API_TOKEN not set")
	}
	if os.Getenv("CIRC_API_TEST_BROKER_CID") == "" {
		t.Skip("skipping test; $CIRC_API_TEST_BROKER_CID not set")
	}

	t.Log("Testing correct return from API call")

	ac := &Config{}
	ac.TokenKey = os.Getenv("CIRCONUS_API_TOKEN")
	apih, err := NewAPI(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	cid := os.Getenv("CIRC_API_TEST_BROKER_CID")
	if cid == "" {
		t.Fatal("Invalid broker cid (empty)")
	}

	brokerCid := BrokerCIDType(cid)

	broker, err := apih.FetchBrokerByCID(brokerCid)
	if err != nil {
		t.Fatalf("Expected no error, got '%v'", err)
	}

	actualType := reflect.TypeOf(broker)
	expectedType := "*api.Broker"
	if actualType.String() != expectedType {
		t.Errorf("Expected %s, got %s", expectedType, actualType.String())
	}

	expectedCid := string(brokerCid)
	if broker.Cid != expectedCid {
		t.Fatalf("%s != %s", broker.Cid, expectedCid)
	}

	t.Logf("Broker returned %s %s", broker.Name, broker.Cid)
}

func TestFetchBrokerListByTag(t *testing.T) {
	if os.Getenv("CIRCONUS_API_TOKEN") == "" {
		t.Skip("skipping test; $CIRCONUS_API_TOKEN not set")
	}
	if os.Getenv("CIRC_API_TEST_BROKER_TAG") == "" {
		t.Skip("skipping test; $CIRC_API_TEST_BROKER_TAG not set")
	}

	t.Log("Testing correct return from API call")

	ac := &Config{}
	ac.TokenKey = os.Getenv("CIRCONUS_API_TOKEN")
	apih, err := NewAPI(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	tag := os.Getenv("CIRC_API_TEST_BROKER_TAG")
	if tag == "" {
		t.Fatal("Invalid broker tag (empty)")
	}

	selectTag := BrokerSearchTagType(tag)

	brokers, err := apih.FetchBrokerListByTag(selectTag)
	if err != nil {
		t.Fatalf("Expected no error, got '%v'", err)
	}

	actualType := reflect.TypeOf(brokers)
	expectedType := "[]api.Broker"
	if actualType.String() != expectedType {
		t.Errorf("Expected %s, got %s", expectedType, actualType.String())
	}

	t.Logf("%d brokers returned", len(brokers))
}

func TestFetchBrokerList(t *testing.T) {
	if os.Getenv("CIRCONUS_API_TOKEN") == "" {
		t.Skip("skipping test; $CIRCONUS_API_TOKEN not set")
	}

	t.Log("Testing correct return from API call")

	ac := &Config{}
	ac.TokenKey = os.Getenv("CIRCONUS_API_TOKEN")
	apih, err := NewAPI(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	brokers, err := apih.FetchBrokerList()
	if err != nil {
		t.Fatalf("Expected no error, got '%v'", err)
	}

	actualType := reflect.TypeOf(brokers)
	expectedType := "[]api.Broker"
	if actualType.String() != expectedType {
		t.Errorf("Expected %s, got %s", expectedType, actualType.String())
	}

	t.Logf("%d brokers returned", len(brokers))
}
<commit_msg>use the generic types for ID, CID, Search query and tag.<commit_after>package api

import (
	"fmt"
	"os"
	"reflect"
	"strconv"
	"testing"
)

// Implicit tests:
//
// FetchBrokerByCID is called by FetchBrokerByID
// BrokerSearch is called by FetchBrokerListByTag

func TestFetchBrokerByID(t *testing.T) {
	if os.Getenv("CIRCONUS_API_TOKEN") == "" {
		t.Skip("skipping test; $CIRCONUS_API_TOKEN not set")
	}
	if os.Getenv("CIRC_API_TEST_BROKER_ID") == "" {
		t.Skip("skipping test; $CIRC_API_TEST_BROKER_ID not set")
	}

	t.Log("Testing correct return from API call")

	ac := &Config{}
	ac.TokenKey = os.Getenv("CIRCONUS_API_TOKEN")
	apih, err := NewAPI(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	bid := os.Getenv("CIRC_API_TEST_BROKER_ID")
	if bid == "" {
		t.Fatal("Invalid broker id (empty)")
	}

	id, err := strconv.Atoi(bid)
	if err != nil {
		t.Fatalf("Unable to convert broker id %s to int", bid)
	}

	brokerID := IDType(id)

	broker, err := apih.FetchBrokerByID(brokerID)
	if err != nil {
		t.Fatalf("Expected no error, got '%v'", err)
	}

	actualType := reflect.TypeOf(broker)
	expectedType := "*api.Broker"
	if actualType.String() != expectedType {
		t.Errorf("Expected %s, got %s", expectedType, actualType.String())
	}

	expectedCid := fmt.Sprintf("/broker/%s", strconv.Itoa(int(brokerID)))
	if broker.Cid != expectedCid {
		t.Fatalf("%s != %s", broker.Cid, expectedCid)
	}

	t.Logf("Broker returned %s %s", broker.Name, broker.Cid)
}

func TestFetchBrokerListByTag(t *testing.T) {
	if os.Getenv("CIRCONUS_API_TOKEN") == "" {
		t.Skip("skipping test; $CIRCONUS_API_TOKEN not set")
	}
	if os.Getenv("CIRC_API_TEST_BROKER_TAG") == "" {
		t.Skip("skipping test; $CIRC_API_TEST_BROKER_TAG not set")
	}

	t.Log("Testing correct return from API call")

	ac := &Config{}
	ac.TokenKey = os.Getenv("CIRCONUS_API_TOKEN")
	apih, err := NewAPI(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	tag := os.Getenv("CIRC_API_TEST_BROKER_TAG")
	if tag == "" {
		t.Fatal("Invalid broker tag (empty)")
	}

	selectTag := SearchTagType(tag)

	brokers, err := apih.FetchBrokerListByTag(selectTag)
	if err != nil {
		t.Fatalf("Expected no error, got '%v'", err)
	}

	actualType := reflect.TypeOf(brokers)
	expectedType := "[]api.Broker"
	if actualType.String() != expectedType {
		t.Errorf("Expected %s, got %s", expectedType, actualType.String())
	}

	t.Logf("%d brokers returned", len(brokers))
}

func TestFetchBrokerList(t *testing.T) {
	if os.Getenv("CIRCONUS_API_TOKEN") == "" {
		t.Skip("skipping test; $CIRCONUS_API_TOKEN not set")
	}

	t.Log("Testing correct return from API call")

	ac := &Config{}
	ac.TokenKey = os.Getenv("CIRCONUS_API_TOKEN")
	apih, err := NewAPI(ac)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	brokers, err := apih.FetchBrokerList()
	if err != nil {
		t.Fatalf("Expected no error, got '%v'", err)
	}

	actualType := reflect.TypeOf(brokers)
	expectedType := "[]api.Broker"
	if actualType.String() != expectedType {
		t.Errorf("Expected %s, got %s", expectedType, actualType.String())
	}

	t.Logf("%d brokers returned", len(brokers))
}
<|endoftext|>


<commit_before>package semver

import (
	"reflect"
	"testing"
)

func TestNew(t *testing.T) {
	parts := []string{"1", "2", "3", "4", "5"}

	ver, err := New(parts[0], parts[1], parts[2], parts[3], parts[4])
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	if !reflect.DeepEqual(parts, ver.parts()) {
		t.Fatalf("bad: %v", parts)
	}
}

func TestNewFromString(t *testing.T) {
	// All components present
	ver, err := NewFromString("1.2.3-4.0a+5.1b")
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	expected := &SemVer{
		Major:  "1",
		Minor:  "2",
		Patch:  "3",
		PreRel: "4.0a",
		Build:  "5.1b",
	}

	if !reflect.DeepEqual(ver, expected) {
		t.Fatalf("bad: %#v", ver)
	}

	// Build number present, pre-release absent
	ver, err = NewFromString("1.2.3+4.0a")
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	expected = &SemVer{
		Major:  "1",
		Minor:  "2",
		Patch:  "3",
		PreRel: "",
		Build:  "4.0a",
	}

	if !reflect.DeepEqual(ver, expected) {
		t.Fatalf("bad: %#v", ver)
	}

	// Pre-release present, build number absent
	ver, err = NewFromString("1.2.3-4.0a")
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	expected = &SemVer{
		Major:  "1",
		Minor:  "2",
		Patch:  "3",
		PreRel: "4.0a",
		Build:  "",
	}

	if !reflect.DeepEqual(ver, expected) {
		t.Fatalf("bad: %#v", ver)
	}
}
<commit_msg>Adding tests for basic semver functionality<commit_after>package semver

import (
	"reflect"
	"testing"
)

func TestNew(t *testing.T) {
	parts := []string{"1", "2", "3", "4", "5"}

	ver, err := New(parts[0], parts[1], parts[2], parts[3], parts[4])
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	if !reflect.DeepEqual(parts, ver.parts()) {
		t.Fatalf("bad: %v", parts)
	}
}

func TestNewFromString(t *testing.T) {
	ver, err := NewFromString("1.2.3")
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	expected := &SemVer{
		Major:  "1",
		Minor:  "2",
		Patch:  "3",
		PreRel: "",
		Build:  "",
	}

	if !reflect.DeepEqual(ver, expected) {
		t.Fatalf("bad: %#v", ver)
	}
}

func TestNewFromString_PrereleaseOnly(t *testing.T) {
	ver, err := NewFromString("1.2.3-4.0a")
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	expected := &SemVer{
		Major:  "1",
		Minor:  "2",
		Patch:  "3",
		PreRel: "4.0a",
		Build:  "",
	}

	if !reflect.DeepEqual(ver, expected) {
		t.Fatalf("bad: %#v", ver)
	}
}

func TestNewFromString_BuildOnly(t *testing.T) {
	ver, err := NewFromString("1.2.3+4.0a")
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	expected := &SemVer{
		Major:  "1",
		Minor:  "2",
		Patch:  "3",
		PreRel: "",
		Build:  "4.0a",
	}

	if !reflect.DeepEqual(ver, expected) {
		t.Fatalf("bad: %#v", ver)
	}
}

func TestNewFromString_PrereleaseAndBuild(t *testing.T) {
	ver, err := NewFromString("1.2.3-alpha5+4.0a")
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	expected := &SemVer{
		Major:  "1",
		Minor:  "2",
		Patch:  "3",
		PreRel: "alpha5",
		Build:  "4.0a",
	}

	if !reflect.DeepEqual(ver, expected) {
		t.Fatalf("bad: %#v", ver)
	}
}

func TestNewFromString_PrerelaseWithPoints(t *testing.T) {
	ver, err := NewFromString("1.2.3-alpha5.3.1")
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	expected := &SemVer{
		Major:  "1",
		Minor:  "2",
		Patch:  "3",
		PreRel: "alpha5.3.1",
		Build:  "",
	}

	if !reflect.DeepEqual(ver, expected) {
		t.Fatalf("bad: %#v", ver)
	}
}

func TestNewFromString_BuildWithPoints(t *testing.T) {
	ver, err := NewFromString("1.2.3+5.3.1a")
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	expected := &SemVer{
		Major:  "1",
		Minor:  "2",
		Patch:  "3",
		PreRel: "",
		Build:  "5.3.1a",
	}

	if !reflect.DeepEqual(ver, expected) {
		t.Fatalf("bad: %#v", ver)
	}
}
<|endoftext|>