<commit_before>package gorpc

import (
	"bufio"
	"compress/flate"
	"encoding/gob"
	"fmt"
	"io"
	"sync"
	"sync/atomic"
	"time"
)

// Rpc client.
//
// The client must be started with Client.Start() before use.
//
// It is absolutely safe and encouraged using a single client across arbitrary
// number of concurrently running goroutines.
//
// Default client settings are optimized for high load, so don't override
// them without valid reason.
type Client struct {
	// Server TCP address to connect to.
	Addr string

	// The number of concurrent connections the client should establish
	// to the sever.
	// By default only one connection is established.
	Conns int

	// The maximum number of pending requests in the queue.
	// Default is 32768.
	PendingRequests int

	// Delay between request flushes.
	// Default value is 5ms.
	FlushDelay time.Duration

	// Maximum request time.
	// Default value is 30s.
	RequestTimeout time.Duration

	// Disable data compression.
	// By default data compression is enabled.
	DisableCompression bool

	// Size of send buffer per each TCP connection in bytes.
	// Default value is 1M.
	SendBufferSize int

	// Size of recv buffer per each TCP connection in bytes.
	// Default value is 1M.
	RecvBufferSize int

	// The client calls this callback when it needs new connection
	// to the server.
	// The client passed Client.Addr into Dial().
	//
	// Override this callback if you want custom underlying transport
	// for gorpc. For example, UDP-based, encrypted or SOAP-based :)
	// Don't forget overriding Client.Listener accordingly.
	//
	// If you need encrypted transport, then feel free using NewTLSDial()
	// on the client and NewTLSListener() helpers on the server.
	//
	// It is expected that the returned conn immediately
	// sends all the data passed via Write() to the server.
	// Otherwise gorpc may hang.
	// The conn implementation must call Flush() on underlying buffered
	// streams before returning from Write().
	//
	// By default it returns TCP connections established
	// to the Client.Addr.
	Dial func(addr string) (conn io.ReadWriteCloser, err error)

	// Connection statistics.
	//
	// The stats doesn't reset automatically. Feel free resetting it
	// any time you wish.
	Stats ConnStats

	requestsChan chan *clientMessage

	clientStopChan chan struct{}
	stopWg         sync.WaitGroup
}

// Starts rpc client. Establishes connection to the server on Client.Addr.
//
// All the response types the server may return must be registered
// via gorpc.RegisterType() before starting the client.
// There is no need in registering base Go types such as int, string, bool,
// float64, etc. or arrays, slices and maps containing base Go types.
func (c *Client) Start() {
	if c.clientStopChan != nil {
		panic("gorpc.Client: the given client is already started. Call Client.Stop() before calling Client.Start() again!")
	}

	if c.PendingRequests <= 0 {
		c.PendingRequests = 32768
	}
	if c.FlushDelay <= 0 {
		c.FlushDelay = 5 * time.Millisecond
	}
	if c.RequestTimeout <= 0 {
		c.RequestTimeout = 30 * time.Second
	}
	if c.SendBufferSize <= 0 {
		c.SendBufferSize = 1024 * 1024
	}
	if c.RecvBufferSize <= 0 {
		c.RecvBufferSize = 1024 * 1024
	}

	c.requestsChan = make(chan *clientMessage, c.PendingRequests)
	c.clientStopChan = make(chan struct{})

	if c.Conns <= 0 {
		c.Conns = 1
	}
	if c.Dial == nil {
		c.Dial = defaultDial
	}
	for i := 0; i < c.Conns; i++ {
		c.stopWg.Add(1)
		go clientHandler(c)
	}
}

// Stops rpc client. Stopped client can be started again.
func (c *Client) Stop() {
	close(c.clientStopChan)
	c.stopWg.Wait()
	c.clientStopChan = nil
}

// Sends the given request to the server and obtains response from the server.
// Returns non-nil error if the response cannot be obtained during
// Client.RequestTimeout or server connection problems occur.
//
// Request and response types may be arbitrary. All the response types
// the server may return must be registered via gorpc.RegisterType() before
// starting the client.
// There is no need in registering base Go types such as int, string, bool,
// float64, etc. or arrays, slices and maps containing base Go types.
//
// Don't forget starting the client with Client.Start() before calling Client.Call().
func (c *Client) Call(request interface{}) (response interface{}, err error) {
	return c.CallTimeout(request, c.RequestTimeout)
}

// Sends the given request to the server and obtains response from the server.
// Returns non-nil error if the response cannot be obtained during
// the given timeout or server connection problems occur.
//
// Request and response types may be arbitrary. All the response types
// the server may return must be registered via gorpc.RegisterType() before
// starting the client.
// There is no need in registering base Go types such as int, string, bool,
// float64, etc. or arrays, slices and maps containing base Go types.
//
// Don't forget starting the client with Client.Start() before calling Client.Call().
func (c *Client) CallTimeout(request interface{}, timeout time.Duration) (response interface{}, err error) {
	m := clientMessage{
		Request: request,
		Done:    make(chan struct{}),
	}
	tc := time.After(timeout)
	select {
	case c.requestsChan <- &m:
		select {
		case <-m.Done:
			return m.Response, m.Error
		case <-tc:
			err := fmt.Errorf("gorpc.Client: [%s]. Cannot obtain response during timeout=%s", c.Addr, timeout)
			logError("%s", err)
			return nil, err
		}
	case <-tc:
		err := fmt.Errorf("gorpc.Client: [%s]. Requests' queue with size=%d is overflown", c.Addr, cap(c.requestsChan))
		logError("%s", err)
		return nil, err
	}
}

func clientHandler(c *Client) {
	defer c.stopWg.Done()

	var conn io.ReadWriteCloser
	var err error

	for {
		dialChan := make(chan struct{})
		go func() {
			if conn, err = c.Dial(c.Addr); err != nil {
				logError("gorpc.Client: [%s]. Cannot establish rpc connection: [%s]", c.Addr, err)
				time.Sleep(time.Second)
			}
			close(dialChan)
		}()

		select {
		case <-c.clientStopChan:
			return
		case <-dialChan:
			atomic.AddUint64(&c.Stats.DialCalls, 1)
		}

		if err != nil {
			atomic.AddUint64(&c.Stats.DialErrors, 1)
			continue
		}
		clientHandleConnection(c, conn)
	}
}

func clientHandleConnection(c *Client, conn io.ReadWriteCloser) {
	var buf [1]byte
	if !c.DisableCompression {
		buf[0] = 1
	}
	_, err := conn.Write(buf[:])
	if err != nil {
		logError("gorpc.Client: [%s]. Error when writing handshake to server: [%s]", c.Addr, err)
		conn.Close()
		return
	}

	stopChan := make(chan struct{})

	pendingRequests := make(map[uint64]*clientMessage)
	var pendingRequestsLock sync.Mutex

	writerDone := make(chan error, 1)
	go clientWriter(c, conn, pendingRequests, &pendingRequestsLock, stopChan, writerDone)

	readerDone := make(chan error, 1)
	go clientReader(c, conn, pendingRequests, &pendingRequestsLock, readerDone)

	select {
	case err = <-writerDone:
		close(stopChan)
		conn.Close()
		<-readerDone
	case err = <-readerDone:
		close(stopChan)
		conn.Close()
		<-writerDone
	case <-c.clientStopChan:
		close(stopChan)
		conn.Close()
		<-readerDone
		<-writerDone
	}

	if err != nil {
		logError("%s", err)
	}
	for _, m := range pendingRequests {
		m.Error = err
		close(m.Done)
	}
}

type clientMessage struct {
	Request  interface{}
	Response interface{}
	Done     chan struct{}
	Error    error
}

func clientWriter(c *Client, w io.Writer, pendingRequests map[uint64]*clientMessage, pendingRequestsLock *sync.Mutex, stopChan <-chan struct{}, done chan<- error) {
	var err error
	defer func() { done <- err }()

	w = newWriterCounter(w, &c.Stats)
	bw := bufio.NewWriterSize(w, c.SendBufferSize)

	ww := bw
	var zw *flate.Writer
	if !c.DisableCompression {
		zw, _ = flate.NewWriter(bw, flate.BestSpeed)
		defer zw.Close()
		ww = bufio.NewWriterSize(zw, c.SendBufferSize)
	}
	e := gob.NewEncoder(ww)

	var flushChan <-chan time.Time

	var msgID uint64
	for {
		var rpcM *clientMessage

		select {
		case <-stopChan:
			return
		case rpcM = <-c.requestsChan:
			if flushChan == nil {
				flushChan = time.After(c.FlushDelay)
			}
		case <-flushChan:
			if !c.DisableCompression {
				if err := ww.Flush(); err != nil {
					err = fmt.Errorf("gorpc.Client: [%s]. Cannot flush data to compressed stream: [%s]", c.Addr, err)
					return
				}
				if err := zw.Flush(); err != nil {
					err = fmt.Errorf("gorpc.Client: [%s]. Cannot flush compressed data to wire: [%s]", c.Addr, err)
					return
				}
			}
			if err := bw.Flush(); err != nil {
				err = fmt.Errorf("gorpc.Client: [%s]. Cannot flush requests to wire: [%s]", c.Addr, err)
				return
			}
			flushChan = nil
			continue
		}

		msgID++
		pendingRequestsLock.Lock()
		n := len(pendingRequests)
		pendingRequests[msgID] = rpcM
		pendingRequestsLock.Unlock()

		if n > 10*c.PendingRequests {
			err = fmt.Errorf("gorpc.Client: [%s]. The server didn't return %d responses yet. Closing server connection in order to prevent client resource leaks", c.Addr, n)
			return
		}

		m := wireMessage{
			ID:   msgID,
			Data: rpcM.Request,
		}
		if err := e.Encode(&m); err != nil {
			err = fmt.Errorf("gorpc.Client: [%s]. Cannot send request to wire: [%s]", c.Addr, err)
			return
		}
	}
}

func clientReader(c *Client, r io.Reader, pendingRequests map[uint64]*clientMessage, pendingRequestsLock *sync.Mutex, done chan<- error) {
	var err error
	defer func() { done <- err }()

	r = newReaderCounter(r, &c.Stats)
	br := bufio.NewReaderSize(r, c.RecvBufferSize)

	rr := br
	if !c.DisableCompression {
		zr := flate.NewReader(br)
		defer zr.Close()
		rr = bufio.NewReaderSize(zr, c.RecvBufferSize)
	}
	d := gob.NewDecoder(rr)

	for {
		var m wireMessage
		if err := d.Decode(&m); err != nil {
			err = fmt.Errorf("gorpc.Client: [%s]. Cannot decode response: [%s]", c.Addr, err)
			return
		}

		pendingRequestsLock.Lock()
		rpcM, ok := pendingRequests[m.ID]
		delete(pendingRequests, m.ID)
		pendingRequestsLock.Unlock()
		if !ok {
			err = fmt.Errorf("gorpc.Client: [%s]. Unexpected msgID=[%d] obtained from server", c.Addr, m.ID)
			return
		}

		rpcM.Response = m.Data
		close(rpcM.Done)
	}
}
<commit_msg>fixed a typo<commit_after>package gorpc

import (
	"bufio"
	"compress/flate"
	"encoding/gob"
	"fmt"
	"io"
	"sync"
	"sync/atomic"
	"time"
)

// Rpc client.
//
// The client must be started with Client.Start() before use.
//
// It is absolutely safe and encouraged using a single client across arbitrary
// number of concurrently running goroutines.
//
// Default client settings are optimized for high load, so don't override
// them without valid reason.
type Client struct {
	// Server TCP address to connect to.
	Addr string

	// The number of concurrent connections the client should establish
	// to the sever.
	// By default only one connection is established.
	Conns int

	// The maximum number of pending requests in the queue.
	// Default is 32768.
	PendingRequests int

	// Delay between request flushes.
	// Default value is 5ms.
	FlushDelay time.Duration

	// Maximum request time.
	// Default value is 30s.
	RequestTimeout time.Duration

	// Disable data compression.
	// By default data compression is enabled.
	DisableCompression bool

	// Size of send buffer per each TCP connection in bytes.
	// Default value is 1M.
	SendBufferSize int

	// Size of recv buffer per each TCP connection in bytes.
	// Default value is 1M.
	RecvBufferSize int

	// The client calls this callback when it needs new connection
	// to the server.
	// The client passed Client.Addr into Dial().
	//
	// Override this callback if you want custom underlying transport
	// for gorpc. For example, UDP-based, encrypted or SOAP-based :)
	// Don't forget overriding Server.Listener accordingly.
	//
	// If you need encrypted transport, then feel free using NewTLSDial()
	// on the client and NewTLSListener() helpers on the server.
	//
	// It is expected that the returned conn immediately
	// sends all the data passed via Write() to the server.
	// Otherwise gorpc may hang.
	// The conn implementation must call Flush() on underlying buffered
	// streams before returning from Write().
	//
	// By default it returns TCP connections established
	// to the Client.Addr.
	Dial func(addr string) (conn io.ReadWriteCloser, err error)

	// Connection statistics.
	//
	// The stats doesn't reset automatically. Feel free resetting it
	// any time you wish.
	Stats ConnStats

	requestsChan chan *clientMessage

	clientStopChan chan struct{}
	stopWg         sync.WaitGroup
}

// Starts rpc client. Establishes connection to the server on Client.Addr.
//
// All the response types the server may return must be registered
// via gorpc.RegisterType() before starting the client.
// There is no need in registering base Go types such as int, string, bool,
// float64, etc. or arrays, slices and maps containing base Go types.
func (c *Client) Start() {
	if c.clientStopChan != nil {
		panic("gorpc.Client: the given client is already started. Call Client.Stop() before calling Client.Start() again!")
	}

	if c.PendingRequests <= 0 {
		c.PendingRequests = 32768
	}
	if c.FlushDelay <= 0 {
		c.FlushDelay = 5 * time.Millisecond
	}
	if c.RequestTimeout <= 0 {
		c.RequestTimeout = 30 * time.Second
	}
	if c.SendBufferSize <= 0 {
		c.SendBufferSize = 1024 * 1024
	}
	if c.RecvBufferSize <= 0 {
		c.RecvBufferSize = 1024 * 1024
	}

	c.requestsChan = make(chan *clientMessage, c.PendingRequests)
	c.clientStopChan = make(chan struct{})

	if c.Conns <= 0 {
		c.Conns = 1
	}
	if c.Dial == nil {
		c.Dial = defaultDial
	}
	for i := 0; i < c.Conns; i++ {
		c.stopWg.Add(1)
		go clientHandler(c)
	}
}

// Stops rpc client. Stopped client can be started again.
func (c *Client) Stop() {
	close(c.clientStopChan)
	c.stopWg.Wait()
	c.clientStopChan = nil
}

// Sends the given request to the server and obtains response from the server.
// Returns non-nil error if the response cannot be obtained during
// Client.RequestTimeout or server connection problems occur.
//
// Request and response types may be arbitrary. All the response types
// the server may return must be registered via gorpc.RegisterType() before
// starting the client.
// There is no need in registering base Go types such as int, string, bool,
// float64, etc. or arrays, slices and maps containing base Go types.
//
// Don't forget starting the client with Client.Start() before calling Client.Call().
func (c *Client) Call(request interface{}) (response interface{}, err error) {
	return c.CallTimeout(request, c.RequestTimeout)
}

// Sends the given request to the server and obtains response from the server.
// Returns non-nil error if the response cannot be obtained during
// the given timeout or server connection problems occur.
//
// Request and response types may be arbitrary. All the response types
// the server may return must be registered via gorpc.RegisterType() before
// starting the client.
// There is no need in registering base Go types such as int, string, bool,
// float64, etc. or arrays, slices and maps containing base Go types.
//
// Don't forget starting the client with Client.Start() before calling Client.Call().
func (c *Client) CallTimeout(request interface{}, timeout time.Duration) (response interface{}, err error) {
	m := clientMessage{
		Request: request,
		Done:    make(chan struct{}),
	}
	tc := time.After(timeout)
	select {
	case c.requestsChan <- &m:
		select {
		case <-m.Done:
			return m.Response, m.Error
		case <-tc:
			err := fmt.Errorf("gorpc.Client: [%s]. Cannot obtain response during timeout=%s", c.Addr, timeout)
			logError("%s", err)
			return nil, err
		}
	case <-tc:
		err := fmt.Errorf("gorpc.Client: [%s]. Requests' queue with size=%d is overflown", c.Addr, cap(c.requestsChan))
		logError("%s", err)
		return nil, err
	}
}

func clientHandler(c *Client) {
	defer c.stopWg.Done()

	var conn io.ReadWriteCloser
	var err error

	for {
		dialChan := make(chan struct{})
		go func() {
			if conn, err = c.Dial(c.Addr); err != nil {
				logError("gorpc.Client: [%s]. Cannot establish rpc connection: [%s]", c.Addr, err)
				time.Sleep(time.Second)
			}
			close(dialChan)
		}()

		select {
		case <-c.clientStopChan:
			return
		case <-dialChan:
			atomic.AddUint64(&c.Stats.DialCalls, 1)
		}

		if err != nil {
			atomic.AddUint64(&c.Stats.DialErrors, 1)
			continue
		}
		clientHandleConnection(c, conn)
	}
}

func clientHandleConnection(c *Client, conn io.ReadWriteCloser) {
	var buf [1]byte
	if !c.DisableCompression {
		buf[0] = 1
	}
	_, err := conn.Write(buf[:])
	if err != nil {
		logError("gorpc.Client: [%s]. Error when writing handshake to server: [%s]", c.Addr, err)
		conn.Close()
		return
	}

	stopChan := make(chan struct{})

	pendingRequests := make(map[uint64]*clientMessage)
	var pendingRequestsLock sync.Mutex

	writerDone := make(chan error, 1)
	go clientWriter(c, conn, pendingRequests, &pendingRequestsLock, stopChan, writerDone)

	readerDone := make(chan error, 1)
	go clientReader(c, conn, pendingRequests, &pendingRequestsLock, readerDone)

	select {
	case err = <-writerDone:
		close(stopChan)
		conn.Close()
		<-readerDone
	case err = <-readerDone:
		close(stopChan)
		conn.Close()
		<-writerDone
	case <-c.clientStopChan:
		close(stopChan)
		conn.Close()
		<-readerDone
		<-writerDone
	}

	if err != nil {
		logError("%s", err)
	}
	for _, m := range pendingRequests {
		m.Error = err
		close(m.Done)
	}
}

type clientMessage struct {
	Request  interface{}
	Response interface{}
	Done     chan struct{}
	Error    error
}

func clientWriter(c *Client, w io.Writer, pendingRequests map[uint64]*clientMessage, pendingRequestsLock *sync.Mutex, stopChan <-chan struct{}, done chan<- error) {
	var err error
	defer func() { done <- err }()

	w = newWriterCounter(w, &c.Stats)
	bw := bufio.NewWriterSize(w, c.SendBufferSize)

	ww := bw
	var zw *flate.Writer
	if !c.DisableCompression {
		zw, _ = flate.NewWriter(bw, flate.BestSpeed)
		defer zw.Close()
		ww = bufio.NewWriterSize(zw, c.SendBufferSize)
	}
	e := gob.NewEncoder(ww)

	var flushChan <-chan time.Time

	var msgID uint64
	for {
		var rpcM *clientMessage

		select {
		case <-stopChan:
			return
		case rpcM = <-c.requestsChan:
			if flushChan == nil {
				flushChan = time.After(c.FlushDelay)
			}
		case <-flushChan:
			if !c.DisableCompression {
				if err := ww.Flush(); err != nil {
					err = fmt.Errorf("gorpc.Client: [%s]. Cannot flush data to compressed stream: [%s]", c.Addr, err)
					return
				}
				if err := zw.Flush(); err != nil {
					err = fmt.Errorf("gorpc.Client: [%s]. Cannot flush compressed data to wire: [%s]", c.Addr, err)
					return
				}
			}
			if err := bw.Flush(); err != nil {
				err = fmt.Errorf("gorpc.Client: [%s]. Cannot flush requests to wire: [%s]", c.Addr, err)
				return
			}
			flushChan = nil
			continue
		}

		msgID++
		pendingRequestsLock.Lock()
		n := len(pendingRequests)
		pendingRequests[msgID] = rpcM
		pendingRequestsLock.Unlock()

		if n > 10*c.PendingRequests {
			err = fmt.Errorf("gorpc.Client: [%s]. The server didn't return %d responses yet. Closing server connection in order to prevent client resource leaks", c.Addr, n)
			return
		}

		m := wireMessage{
			ID:   msgID,
			Data: rpcM.Request,
		}
		if err := e.Encode(&m); err != nil {
			err = fmt.Errorf("gorpc.Client: [%s]. Cannot send request to wire: [%s]", c.Addr, err)
			return
		}
	}
}

func clientReader(c *Client, r io.Reader, pendingRequests map[uint64]*clientMessage, pendingRequestsLock *sync.Mutex, done chan<- error) {
	var err error
	defer func() { done <- err }()

	r = newReaderCounter(r, &c.Stats)
	br := bufio.NewReaderSize(r, c.RecvBufferSize)

	rr := br
	if !c.DisableCompression {
		zr := flate.NewReader(br)
		defer zr.Close()
		rr = bufio.NewReaderSize(zr, c.RecvBufferSize)
	}
	d := gob.NewDecoder(rr)

	for {
		var m wireMessage
		if err := d.Decode(&m); err != nil {
			err = fmt.Errorf("gorpc.Client: [%s]. Cannot decode response: [%s]", c.Addr, err)
			return
		}

		pendingRequestsLock.Lock()
		rpcM, ok := pendingRequests[m.ID]
		delete(pendingRequests, m.ID)
		pendingRequestsLock.Unlock()
		if !ok {
			err = fmt.Errorf("gorpc.Client: [%s]. Unexpected msgID=[%d] obtained from server", c.Addr, m.ID)
			return
		}

		rpcM.Response = m.Data
		close(rpcM.Done)
	}
}
<|endoftext|>


<commit_before>/*
Package k8s implements a Kubernetes client.

	c, err := k8s.NewInClusterClient()
	if err != nil {
		// handle error
	}
	extensions := c.ExtensionsV1Beta1()

	ingresses, err := extensions.ListIngresses(ctx, c.Namespace)
	if err != nil {
		// handle error
	}

*/
package k8s

import (
	"bytes"
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/ericchiang/k8s/api/unversioned"
)

// String returns a pointer to a string. Useful for creating API objects
// that take pointers instead of literals.
//
//		cm := &v1.ConfigMap{
//			Metadata: &v1.ObjectMeta{
//				Name:      k8s.String("myconfigmap"),
//				Namespace: k8s.String("default"),
//			},
//			Data: map[string]string{
//				"foo": "bar",
//			},
//		}
//
func String(s string) *string { return &s }
func Int(i int) *int          { return &i }
func Bool(b bool) *bool       { return &b }

// Client is a Kuberntes client.
type Client struct {
	// The URL of the API server.
	Endpoint string

	// Default namespaces for objects that don't supply a namespace in
	// their object metadata.
	Namespace string

	// SetHeaders provides a hook for modifying the HTTP headers of all requests.
	//
	//		client, err := k8s.NewClient(config)
	//		if err != nil {
	//			// handle error
	//		}
	//		client.SetHeaders = func(h http.Header) error {
	//			h.Set("Authorization", "Bearer "+mytoken)
	//			return nil
	//		}
	//
	SetHeaders func(h http.Header) error

	Client *http.Client
}

func (c *Client) newRequest(verb, url string, body io.Reader) (*http.Request, error) {
	req, err := http.NewRequest(verb, url, body)
	if err != nil {
		return nil, err
	}
	if c.SetHeaders != nil {
		if err := c.SetHeaders(req.Header); err != nil {
			return nil, err
		}
	}
	return req, nil
}

// Option represents optional call parameters, such as label selectors.
type Option interface {
	queryParam() (key, val string)
}

type resourceVersionOption string

func (r resourceVersionOption) queryParam() (string, string) {
	return "resourceVersion", string(r)
}

// ResourceVersion causes watch operations to only show changes since
// a particular version of a resource.
func ResourceVersion(resourceVersion string) Option {
	return resourceVersionOption(resourceVersion)
}

type timeoutSeconds string

func (t timeoutSeconds) queryParam() (string, string) {
	return "timeoutSeconds", string(t)
}

// Timeout declares the timeout for list and watch operations. Timeout
// is only accurate to the second.
func Timeout(d time.Duration) Option {
	return timeoutSeconds(strconv.FormatInt(int64(d/time.Second), 10))
}

// NewClient initializes a client from a client config.
func NewClient(config *Config) (*Client, error) {
	if len(config.Contexts) == 0 {
		if config.CurrentContext != "" {
			return nil, fmt.Errorf("no contexts with name %q", config.CurrentContext)
		}

		if n := len(config.Clusters); n == 0 {
			return nil, errors.New("no clusters provided")
		} else if n > 1 {
			return nil, errors.New("multiple clusters but no current context")
		}
		if n := len(config.AuthInfos); n == 0 {
			return nil, errors.New("no users provided")
		} else if n > 1 {
			return nil, errors.New("multiple users but no current context")
		}

		return newClient(config.Clusters[0].Cluster, config.AuthInfos[0].AuthInfo, "")
	}

	var ctx Context
	if config.CurrentContext == "" {
		if n := len(config.Contexts); n == 0 {
			return nil, errors.New("no contexts provided")
		} else if n > 1 {
			return nil, errors.New("multiple contexts but no current context")
		}
		ctx = config.Contexts[0].Context
	} else {
		for _, c := range config.Contexts {
			if c.Name == config.CurrentContext {
				ctx = c.Context
				goto configFound
			}
		}
		return nil, fmt.Errorf("no config named %q", config.CurrentContext)
	configFound:
	}

	if ctx.Cluster == "" {
		return nil, fmt.Errorf("context doesn't have a cluster")
	}
	if ctx.AuthInfo == "" {
		return nil, fmt.Errorf("context doesn't have a user")
	}
	var (
		user    AuthInfo
		cluster Cluster
	)

	for _, u := range config.AuthInfos {
		if u.Name == ctx.AuthInfo {
			user = u.AuthInfo
			goto userFound
		}
	}
	return nil, fmt.Errorf("no user named %q", ctx.AuthInfo)
userFound:

	for _, c := range config.Clusters {
		if c.Name == ctx.Cluster {
			cluster = c.Cluster
			goto clusterFound
		}
	}
	return nil, fmt.Errorf("no cluster named %q", ctx.Cluster)
clusterFound:

	return newClient(cluster, user, ctx.Namespace)
}

// NewInClusterClient returns a client that uses the service account bearer token mounted
// into Kubernetes pods.
func NewInClusterClient() (*Client, error) {
	host, port := os.Getenv("KUBERNETES_SERVICE_HOST"), os.Getenv("KUBERNETES_SERVICE_PORT")
	if len(host) == 0 || len(port) == 0 {
		return nil, errors.New("unable to load in-cluster configuration, KUBERNETES_SERVICE_HOST and KUBERNETES_SERVICE_PORT must be defined")
	}
	namespace, err := ioutil.ReadFile("/var/run/secrets/kubernetes.io/serviceaccount/namespace")
	if err != nil {
		return nil, err
	}

	cluster := Cluster{
		Server:               "https://" + host + ":" + port,
		CertificateAuthority: "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt",
	}
	user := AuthInfo{TokenFile: "/var/run/secrets/kubernetes.io/serviceaccount/token"}
	return newClient(cluster, user, string(namespace))
}

func load(filepath string, data []byte) (out []byte, err error) {
	if filepath != "" {
		data, err = ioutil.ReadFile(filepath)
	}
	return data, err
}

func newClient(cluster Cluster, user AuthInfo, namespace string) (*Client, error) {
	if cluster.Server == "" {
		// NOTE: kubectl defaults to localhost:8080, but it's probably better to just
		// be strict.
		return nil, fmt.Errorf("no cluster endpoint provided")
	}

	ca, err := load(cluster.CertificateAuthority, cluster.CertificateAuthorityData)
	if err != nil {
		return nil, fmt.Errorf("loading certificate authority: %v", err)
	}

	clientCert, err := load(user.ClientCertificate, user.ClientCertificateData)
	if err != nil {
		return nil, fmt.Errorf("load client cert: %v", err)
	}
	clientKey, err := load(user.ClientKey, user.ClientKeyData)
	if err != nil {
		return nil, fmt.Errorf("load client cert: %v", err)
	}

	// See https://github.com/gtank/cryptopasta
	tlsConfig := &tls.Config{MinVersion: tls.VersionTLS12}

	if len(ca) != 0 {
		tlsConfig.RootCAs = x509.NewCertPool()
		if !tlsConfig.RootCAs.AppendCertsFromPEM(ca) {
			return nil, errors.New("certificate authority doesn't contain any certificates")
		}
	}
	if len(clientCert) != 0 {
		cert, err := tls.X509KeyPair(clientCert, clientKey)
		if err != nil {
			return nil, fmt.Errorf("invalid client cert and key pair: %v", err)
		}
		tlsConfig.Certificates = []tls.Certificate{cert}
	}

	token := user.Token
	if user.TokenFile != "" {
		data, err := ioutil.ReadFile(user.TokenFile)
		if err != nil {
			return nil, fmt.Errorf("load token file: %v", err)
		}
		token = string(data)
	}

	client := &Client{
		Endpoint:  cluster.Server,
		Namespace: namespace,
		Client: &http.Client{
			Transport: &http.Transport{
				Proxy: http.ProxyFromEnvironment,
				DialContext: (&net.Dialer{
					Timeout:   30 * time.Second,
					KeepAlive: 30 * time.Second,
				}).DialContext,
				TLSClientConfig:       tlsConfig,
				MaxIdleConns:          100,
				IdleConnTimeout:       90 * time.Second,
				TLSHandshakeTimeout:   10 * time.Second,
				ExpectContinueTimeout: 1 * time.Second,
			},
		},
	}

	if token != "" {
		client.SetHeaders = func(h http.Header) error {
			h.Set("Authorization", "Bearer "+token)
			return nil
		}
	}
	if user.Username != "" && user.Password != "" {
		auth := user.Username + ":" + user.Password
		auth = "Basic " + base64.StdEncoding.EncodeToString([]byte(auth))
		client.SetHeaders = func(h http.Header) error {
			h.Set("Authorization", auth)
			return nil
		}
	}
	return client, nil
}

// APIError is an error from a unexpected status code.
type APIError struct {
	// The status object returned by the Kubernetes API,
	Status *unversioned.Status
}

func (e *APIError) Error() string { return fmt.Sprintf("kubernetes api: %s", e.Status.Message) }

func checkStatusCode(c *codec, statusCode int, body []byte) error {
	if statusCode/100 == 2 {
		return nil
	}

	status := new(unversioned.Status)
	if err := c.unmarshal(body, status); err != nil {
		return fmt.Errorf("decode error status: %v", err)
	}
	return &APIError{status}
}

func (c *Client) client() *http.Client {
	if c.Client == nil {
		return http.DefaultClient
	}
	return c.Client
}

func (c *Client) namespaceFor(namespace string) string {
	if namespace != "" {
		return namespace
	}
	return c.Namespace
}

// The following methods hold the logic for interacting with the Kubernetes API. Generated
// clients are thin wrappers on top of these methods.
//
// This client implements specs in the "API Conventions" developer document, which can be
// found here:
//
//   https://github.com/kubernetes/kubernetes/blob/master/docs/devel/api-conventions.md

func (c *Client) urlFor(apiGroup, apiVersion, namespace, resource, name string, options ...Option) string {
	basePath := "apis/"
	if apiGroup == "" {
		basePath = "api/"
	}

	var p string
	if namespace != "" {
		p = path.Join(basePath, apiGroup, apiVersion, "namespaces", namespace, resource, name)
	} else {
		p = path.Join(basePath, apiGroup, apiVersion, resource, name)
	}
	endpoint := ""
	if strings.HasSuffix(c.Endpoint, "/") {
		endpoint = c.Endpoint + p
	} else {
		endpoint = c.Endpoint + "/" + p
	}
	if len(options) == 0 {
		return endpoint
	}

	v := url.Values{}
	for _, option := range options {
		key, val := option.queryParam()
		v.Set(key, val)
	}
	return endpoint + "?" + v.Encode()
}

func (c *Client) create(ctx context.Context, codec *codec, verb, url string, req, resp interface{}) error {
	body, err := codec.marshal(req)
	if err != nil {
		return err
	}

	r, err := c.newRequest(verb, url, bytes.NewReader(body))
	if err != nil {
		return err
	}
	r.Header.Set("Content-Type", codec.contentType)
	r.Header.Set("Accept", codec.contentType)

	re, err := c.client().Do(r)
	if err != nil {
		return err
	}
	defer re.Body.Close()

	respBody, err := ioutil.ReadAll(re.Body)
	if err != nil {
		return fmt.Errorf("read body: %v", err)
	}

	if err := checkStatusCode(codec, re.StatusCode, respBody); err != nil {
		return err
	}
	return codec.unmarshal(respBody, resp)
}

func (c *Client) delete(ctx context.Context, codec *codec, url string) error {
	r, err := c.newRequest("DELETE", url, nil)
	if err != nil {
		return err
	}
	r.Header.Set("Accept", codec.contentType)

	re, err := c.client().Do(r)
	if err != nil {
		return err
	}
	defer re.Body.Close()

	respBody, err := ioutil.ReadAll(re.Body)
	if err != nil {
		return fmt.Errorf("read body: %v", err)
	}

	if err := checkStatusCode(codec, re.StatusCode, respBody); err != nil {
		return err
	}
	return nil
}

// get can be used to either get or list a given resource.
func (c *Client) get(ctx context.Context, codec *codec, url string, resp interface{}) error {
	r, err := c.newRequest("GET", url, nil)
	if err != nil {
		return err
	}
	r.Header.Set("Accept", codec.contentType)
	re, err := c.client().Do(r)
	if err != nil {
		return err
	}
	defer re.Body.Close()

	respBody, err := ioutil.ReadAll(re.Body)
	if err != nil {
		return fmt.Errorf("read body: %v", err)
	}

	if err := checkStatusCode(codec, re.StatusCode, respBody); err != nil {
		return err
	}
	return codec.unmarshal(respBody, resp)
}
<commit_msg>*: add status code to API error<commit_after>/*
Package k8s implements a Kubernetes client.

	c, err := k8s.NewInClusterClient()
	if err != nil {
		// handle error
	}
	extensions := c.ExtensionsV1Beta1()

	ingresses, err := extensions.ListIngresses(ctx, c.Namespace)
	if err != nil {
		// handle error
	}

*/
package k8s

import (
	"bytes"
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/ericchiang/k8s/api/unversioned"
)

// String returns a pointer to a string. Useful for creating API objects
// that take pointers instead of literals.
//
//		cm := &v1.ConfigMap{
//			Metadata: &v1.ObjectMeta{
//				Name:      k8s.String("myconfigmap"),
//				Namespace: k8s.String("default"),
//			},
//			Data: map[string]string{
//				"foo": "bar",
//			},
//		}
//
func String(s string) *string { return &s }
func Int(i int) *int          { return &i }
func Bool(b bool) *bool       { return &b }

// Client is a Kuberntes client.
type Client struct {
	// The URL of the API server.
	Endpoint string

	// Default namespaces for objects that don't supply a namespace in
	// their object metadata.
	Namespace string

	// SetHeaders provides a hook for modifying the HTTP headers of all requests.
	//
	//		client, err := k8s.NewClient(config)
	//		if err != nil {
	//			// handle error
	//		}
	//		client.SetHeaders = func(h http.Header) error {
	//			h.Set("Authorization", "Bearer "+mytoken)
	//			return nil
	//		}
	//
	SetHeaders func(h http.Header) error

	Client *http.Client
}

func (c *Client) newRequest(verb, url string, body io.Reader) (*http.Request, error) {
	req, err := http.NewRequest(verb, url, body)
	if err != nil {
		return nil, err
	}
	if c.SetHeaders != nil {
		if err := c.SetHeaders(req.Header); err != nil {
			return nil, err
		}
	}
	return req, nil
}

// Option represents optional call parameters, such as label selectors.
type Option interface {
	queryParam() (key, val string)
}

type resourceVersionOption string

func (r resourceVersionOption) queryParam() (string, string) {
	return "resourceVersion", string(r)
}

// ResourceVersion causes watch operations to only show changes since
// a particular version of a resource.
func ResourceVersion(resourceVersion string) Option {
	return resourceVersionOption(resourceVersion)
}

type timeoutSeconds string

func (t timeoutSeconds) queryParam() (string, string) {
	return "timeoutSeconds", string(t)
}

// Timeout declares the timeout for list and watch operations. Timeout
// is only accurate to the second.
func Timeout(d time.Duration) Option {
	return timeoutSeconds(strconv.FormatInt(int64(d/time.Second), 10))
}

// NewClient initializes a client from a client config.
func NewClient(config *Config) (*Client, error) {
	if len(config.Contexts) == 0 {
		if config.CurrentContext != "" {
			return nil, fmt.Errorf("no contexts with name %q", config.CurrentContext)
		}

		if n := len(config.Clusters); n == 0 {
			return nil, errors.New("no clusters provided")
		} else if n > 1 {
			return nil, errors.New("multiple clusters but no current context")
		}
		if n := len(config.AuthInfos); n == 0 {
			return nil, errors.New("no users provided")
		} else if n > 1 {
			return nil, errors.New("multiple users but no current context")
		}

		return newClient(config.Clusters[0].Cluster, config.AuthInfos[0].AuthInfo, "")
	}

	var ctx Context
	if config.CurrentContext == "" {
		if n := len(config.Contexts); n == 0 {
			return nil, errors.New("no contexts provided")
		} else if n > 1 {
			return nil, errors.New("multiple contexts but no current context")
		}
		ctx = config.Contexts[0].Context
	} else {
		for _, c := range config.Contexts {
			if c.Name == config.CurrentContext {
				ctx = c.Context
				goto configFound
			}
		}
		return nil, fmt.Errorf("no config named %q", config.CurrentContext)
	configFound:
	}

	if ctx.Cluster == "" {
		return nil, fmt.Errorf("context doesn't have a cluster")
	}
	if ctx.AuthInfo == "" {
		return nil, fmt.Errorf("context doesn't have a user")
	}
	var (
		user    AuthInfo
		cluster Cluster
	)

	for _, u := range config.AuthInfos {
		if u.Name == ctx.AuthInfo {
			user = u.AuthInfo
			goto userFound
		}
	}
	return nil, fmt.Errorf("no user named %q", ctx.AuthInfo)
userFound:

	for _, c := range config.Clusters {
		if c.Name == ctx.Cluster {
			cluster = c.Cluster
			goto clusterFound
		}
	}
	return nil, fmt.Errorf("no cluster named %q", ctx.Cluster)
clusterFound:

	return newClient(cluster, user, ctx.Namespace)
}

// NewInClusterClient returns a client that uses the service account bearer token mounted
// into Kubernetes pods.
func NewInClusterClient() (*Client, error) {
	host, port := os.Getenv("KUBERNETES_SERVICE_HOST"), os.Getenv("KUBERNETES_SERVICE_PORT")
	if len(host) == 0 || len(port) == 0 {
		return nil, errors.New("unable to load in-cluster configuration, KUBERNETES_SERVICE_HOST and KUBERNETES_SERVICE_PORT must be defined")
	}
	namespace, err := ioutil.ReadFile("/var/run/secrets/kubernetes.io/serviceaccount/namespace")
	if err != nil {
		return nil, err
	}

	cluster := Cluster{
		Server:               "https://" + host + ":" + port,
		CertificateAuthority: "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt",
	}
	user := AuthInfo{TokenFile: "/var/run/secrets/kubernetes.io/serviceaccount/token"}
	return newClient(cluster, user, string(namespace))
}

func load(filepath string, data []byte) (out []byte, err error) {
	if filepath != "" {
		data, err = ioutil.ReadFile(filepath)
	}
	return data, err
}

func newClient(cluster Cluster, user AuthInfo, namespace string) (*Client, error) {
	if cluster.Server == "" {
		// NOTE: kubectl defaults to localhost:8080, but it's probably better to just
		// be strict.
		return nil, fmt.Errorf("no cluster endpoint provided")
	}

	ca, err := load(cluster.CertificateAuthority, cluster.CertificateAuthorityData)
	if err != nil {
		return nil, fmt.Errorf("loading certificate authority: %v", err)
	}

	clientCert, err := load(user.ClientCertificate, user.ClientCertificateData)
	if err != nil {
		return nil, fmt.Errorf("load client cert: %v", err)
	}
	clientKey, err := load(user.ClientKey, user.ClientKeyData)
	if err != nil {
		return nil, fmt.Errorf("load client cert: %v", err)
	}

	// See https://github.com/gtank/cryptopasta
	tlsConfig := &tls.Config{MinVersion: tls.VersionTLS12}

	if len(ca) != 0 {
		tlsConfig.RootCAs = x509.NewCertPool()
		if !tlsConfig.RootCAs.AppendCertsFromPEM(ca) {
			return nil, errors.New("certificate authority doesn't contain any certificates")
		}
	}
	if len(clientCert) != 0 {
		cert, err := tls.X509KeyPair(clientCert, clientKey)
		if err != nil {
			return nil, fmt.Errorf("invalid client cert and key pair: %v", err)
		}
		tlsConfig.Certificates = []tls.Certificate{cert}
	}

	token := user.Token
	if user.TokenFile != "" {
		data, err := ioutil.ReadFile(user.TokenFile)
		if err != nil {
			return nil, fmt.Errorf("load token file: %v", err)
		}
		token = string(data)
	}

	client := &Client{
		Endpoint:  cluster.Server,
		Namespace: namespace,
		Client: &http.Client{
			Transport: &http.Transport{
				Proxy: http.ProxyFromEnvironment,
				DialContext: (&net.Dialer{
					Timeout:   30 * time.Second,
					KeepAlive: 30 * time.Second,
				}).DialContext,
				TLSClientConfig:       tlsConfig,
				MaxIdleConns:          100,
				IdleConnTimeout:       90 * time.Second,
				TLSHandshakeTimeout:   10 * time.Second,
				ExpectContinueTimeout: 1 * time.Second,
			},
		},
	}

	if token != "" {
		client.SetHeaders = func(h http.Header) error {
			h.Set("Authorization", "Bearer "+token)
			return nil
		}
	}
	if user.Username != "" && user.Password != "" {
		auth := user.Username + ":" + user.Password
		auth = "Basic " + base64.StdEncoding.EncodeToString([]byte(auth))
		client.SetHeaders = func(h http.Header) error {
			h.Set("Authorization", auth)
			return nil
		}
	}
	return client, nil
}

// APIError is an error from a unexpected status code.
type APIError struct {
	// The status object returned by the Kubernetes API,
	Status *unversioned.Status

	// Status code returned by the HTTP request.
	Code int
}

func (e *APIError) Error() string {
	if e.Status != nil && e.Status.Message != nil && e.Status.Status != nil {
		return fmt.Sprintf("kubernetes api: %s %d %s", *e.Status.Status, e.Code, *e.Status.Message)
	}
	return fmt.Sprintf("%#v", e)
}

func checkStatusCode(c *codec, statusCode int, body []byte) error {
	if statusCode/100 == 2 {
		return nil
	}

	status := new(unversioned.Status)
	if err := c.unmarshal(body, status); err != nil {
		return fmt.Errorf("decode error status: %v", err)
	}
	return &APIError{status, statusCode}
}

func (c *Client) client() *http.Client {
	if c.Client == nil {
		return http.DefaultClient
	}
	return c.Client
}

func (c *Client) namespaceFor(namespace string) string {
	if namespace != "" {
		return namespace
	}
	return c.Namespace
}

// The following methods hold the logic for interacting with the Kubernetes API. Generated
// clients are thin wrappers on top of these methods.
//
// This client implements specs in the "API Conventions" developer document, which can be
// found here:
//
//   https://github.com/kubernetes/kubernetes/blob/master/docs/devel/api-conventions.md

func (c *Client) urlFor(apiGroup, apiVersion, namespace, resource, name string, options ...Option) string {
	basePath := "apis/"
	if apiGroup == "" {
		basePath = "api/"
	}

	var p string
	if namespace != "" {
		p = path.Join(basePath, apiGroup, apiVersion, "namespaces", namespace, resource, name)
	} else {
		p = path.Join(basePath, apiGroup, apiVersion, resource, name)
	}
	endpoint := ""
	if strings.HasSuffix(c.Endpoint, "/") {
		endpoint = c.Endpoint + p
	} else {
		endpoint = c.Endpoint + "/" + p
	}
	if len(options) == 0 {
		return endpoint
	}

	v := url.Values{}
	for _, option := range options {
		key, val := option.queryParam()
		v.Set(key, val)
	}
	return endpoint + "?" + v.Encode()
}

func (c *Client) create(ctx context.Context, codec *codec, verb, url string, req, resp interface{}) error {
	body, err := codec.marshal(req)
	if err != nil {
		return err
	}

	r, err := c.newRequest(verb, url, bytes.NewReader(body))
	if err != nil {
		return err
	}
	r.Header.Set("Content-Type", codec.contentType)
	r.Header.Set("Accept", codec.contentType)

	re, err := c.client().Do(r)
	if err != nil {
		return err
	}
	defer re.Body.Close()

	respBody, err := ioutil.ReadAll(re.Body)
	if err != nil {
		return fmt.Errorf("read body: %v", err)
	}

	if err := checkStatusCode(codec, re.StatusCode, respBody); err != nil {
		return err
	}
	return codec.unmarshal(respBody, resp)
}

func (c *Client) delete(ctx context.Context, codec *codec, url string) error {
	r, err := c.newRequest("DELETE", url, nil)
	if err != nil {
		return err
	}
	r.Header.Set("Accept", codec.contentType)

	re, err := c.client().Do(r)
	if err != nil {
		return err
	}
	defer re.Body.Close()

	respBody, err := ioutil.ReadAll(re.Body)
	if err != nil {
		return fmt.Errorf("read body: %v", err)
	}

	if err := checkStatusCode(codec, re.StatusCode, respBody); err != nil {
		return err
	}
	return nil
}

// get can be used to either get or list a given resource.
func (c *Client) get(ctx context.Context, codec *codec, url string, resp interface{}) error {
	r, err := c.newRequest("GET", url, nil)
	if err != nil {
		return err
	}
	r.Header.Set("Accept", codec.contentType)
	re, err := c.client().Do(r)
	if err != nil {
		return err
	}
	defer re.Body.Close()

	respBody, err := ioutil.ReadAll(re.Body)
	if err != nil {
		return fmt.Errorf("read body: %v", err)
	}

	if err := checkStatusCode(codec, re.StatusCode, respBody); err != nil {
		return err
	}
	return codec.unmarshal(respBody, resp)
}
<|endoftext|>


<commit_before>package doubleratchet

import (
	"github.com/stretchr/testify/require"
	"testing"
)

var sk = [32]byte{0xeb, 0x8, 0x10, 0x7c, 0x33, 0x54, 0x0, 0x20, 0xe9, 0x4f, 0x6c, 0x84, 0xe4, 0x39, 0x50, 0x5a, 0x2f, 0x60, 0xbe, 0x81, 0xa, 0x78, 0x8b, 0xeb, 0x1e, 0x2c, 0x9, 0x8d, 0x4b, 0x4d, 0xc1, 0x40}

func TestNew_Basic(t *testing.T) {
	// Act.
	var (
		si, err = New(sk)
		s       = si.(*state)
	)

	// Assert.
	require.Nil(t, err)
	require.Equal(t, sk, s.RK)
	require.Equal(t, [32]byte{}, s.DHr)
	require.NotEqual(t, [32]byte{}, s.DHs.PrivateKey())
	require.NotEqual(t, [32]byte{}, s.DHs.PublicKey())
	require.Equal(t, [32]byte{}, s.CKs)
	require.Equal(t, [32]byte{}, s.CKr)
	require.EqualValues(t, 0, s.Ns)
	require.EqualValues(t, 0, s.Nr)
	require.EqualValues(t, 0, s.PN)
	require.NotNil(t, s.MkSkipped)
	require.NotEqual(t, 0, s.MaxSkip)
	require.NotNil(t, s.Crypto)
}

func TestNew_WithMaxSkip_OK(t *testing.T) {
	// Act.
	var (
		si, err = New(sk, WithMaxSkip(100))
		s       = si.(*state)
	)

	// Assert.
	require.Nil(t, err)
	require.EqualValues(t, 100, s.MaxSkip)
}

func TestNew_WithMaxSkip_Negative(t *testing.T) {
	// Act.
	var (
		_, err = New(sk, WithMaxSkip(-10))
	)

	// Assert.
	require.NotNil(t, err)
}

//func TestNew_WithRemoteKey(t *testing.T) {
//	// Act.
//	var (
//		si, err = New(sk, WithRemoteKey(100))
//		s       = si.(*state)
//	)
//
//	// Assert.
//
//}
<commit_msg>Implemented WithRemoteKey option tests<commit_after>package doubleratchet

import (
	"github.com/stretchr/testify/require"
	"testing"
)

var (
	sk        = [32]byte{0xeb, 0x8, 0x10, 0x7c, 0x33, 0x54, 0x0, 0x20, 0xe9, 0x4f, 0x6c, 0x84, 0xe4, 0x39, 0x50, 0x5a, 0x2f, 0x60, 0xbe, 0x81, 0xa, 0x78, 0x8b, 0xeb, 0x1e, 0x2c, 0x9, 0x8d, 0x4b, 0x4d, 0xc1, 0x40}
	alicePair = dhPair{
		privateKey: [32]byte{0x88, 0x6a, 0x50, 0xbf, 0x65, 0x63, 0x4b, 0xfb, 0xf5, 0xcf, 0x7d, 0xe0, 0x79, 0xec, 0x3c, 0x70, 0x56, 0x96, 0xa5, 0x4a, 0x8a, 0xfb, 0xfa, 0x97, 0xde, 0x4a, 0x2f, 0xdd, 0xdd, 0x22, 0x40, 0x51},
		publicKey:  [32]byte{0x6, 0x45, 0x36, 0xa5, 0xed, 0xa0, 0xae, 0xaf, 0x62, 0x4f, 0x20, 0x63, 0x3b, 0x8e, 0xc1, 0x7, 0xe8, 0xe7, 0x45, 0x1, 0x8d, 0x14, 0xdb, 0xf8, 0x9, 0x51, 0x3c, 0x5f, 0xbd, 0x33, 0x7, 0x44},
	}
	bobPair = dhPair{
		privateKey: [32]byte{0xf0, 0x22, 0x54, 0xf4, 0xcb, 0xa2, 0x60, 0xc8, 0xeb, 0xe, 0x83, 0xb, 0xc8, 0xb2, 0xfb, 0x18, 0x6f, 0x1b, 0xa4, 0xa2, 0x6e, 0x45, 0xc, 0xeb, 0xff, 0x74, 0xce, 0x65, 0x8b, 0x6e, 0x4c, 0x5d},
		publicKey:  [32]byte{0xe3, 0xbe, 0xb9, 0x4e, 0x70, 0x17, 0x37, 0xc, 0x1, 0x8f, 0xa9, 0x7e, 0xef, 0x4, 0xfb, 0x23, 0xac, 0xea, 0x28, 0xf7, 0xa9, 0x56, 0xcc, 0x1d, 0x46, 0xf3, 0xb5, 0x1d, 0x7d, 0x7d, 0x5e, 0x2c},
	}
)

func TestNew_Basic(t *testing.T) {
	// Act.
	var (
		si, err = New(sk)
		s       = si.(*state)
	)

	// Assert.
	require.Nil(t, err)
	require.Equal(t, sk, s.RK)
	require.Equal(t, [32]byte{}, s.DHr)
	require.NotEqual(t, [32]byte{}, s.DHs.PrivateKey())
	require.NotEqual(t, [32]byte{}, s.DHs.PublicKey())
	require.Equal(t, [32]byte{}, s.CKs)
	require.Equal(t, [32]byte{}, s.CKr)
	require.EqualValues(t, 0, s.Ns)
	require.EqualValues(t, 0, s.Nr)
	require.EqualValues(t, 0, s.PN)
	require.NotNil(t, s.MkSkipped)
	require.NotEqual(t, 0, s.MaxSkip)
	require.NotNil(t, s.Crypto)
}

func TestNew_WithMaxSkip_OK(t *testing.T) {
	// Act.
	var (
		si, err = New(sk, WithMaxSkip(100))
		s       = si.(*state)
	)

	// Assert.
	require.Nil(t, err)
	require.EqualValues(t, 100, s.MaxSkip)
}

func TestNew_WithMaxSkip_Negative(t *testing.T) {
	// Act.
	var (
		_, err = New(sk, WithMaxSkip(-10))
	)

	// Assert.
	require.NotNil(t, err)
}

func TestNew_WithRemoteKey(t *testing.T) {
	// Act.
	var (
		si, err = New(sk, WithRemoteKey(bobPair.PublicKey()))
		s       = si.(*state)
	)

	// Assert.
	require.Nil(t, err)
	require.Equal(t, bobPair.PublicKey(), s.DHr)
	require.NotEqual(t, [32]byte{}, s.RK)
	require.NotEqual(t, sk, s.RK)
	require.NotEqual(t, [32]byte{}, s.CKs)
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"log"
	"time"

	"github.com/phil-mansfield/shellfish/cmd/catalog"
	"github.com/phil-mansfield/shellfish/cmd/env"
	"github.com/phil-mansfield/shellfish/cmd/halo"
	"github.com/phil-mansfield/shellfish/cmd/memo"
	"github.com/phil-mansfield/shellfish/io"
	"github.com/phil-mansfield/shellfish/logging"
	"github.com/phil-mansfield/shellfish/parse"
)

const finderCells = 150

// IDConfig contains the configuration fileds for the 'id' mode of the shellfish
// tool.
type IDConfig struct {
	idType                     string
	ids                        []int64
	idStart, idEnd, snap, mult int64

	exclusionStrategy   string
	exclusionRadiusMult float64
}

var _ Mode = &IDConfig{}

// ExampleConfig creates an example id.config file.
func (config *IDConfig) ExampleConfig() string {
	return `[id.config]
#####################
## Required Fields ##
#####################

# Index of the snapshot to be analyzed.
Snap = 100

IDs = 10, 11, 12, 13, 14

#####################
## Optional Fields ##
#####################

# IDType indicates what the input IDs correspond to. It can be set to the
# following modes:
# halo-id - The numeric IDs given in the halo catalog.
# m200m   - The rank of the halos when sorted by M200m.
#
# Defaults to m200m if not set.
# IDType = m200m

# An alternative way of specifying IDs is to select start and end (inclusive)
# ID values. If the IDs variable is not set, both of these values must be set.
#
# IDStart = 10
# IDEnd = 15

# ExclusionStrategy determines how to exclude IDs from the given set. This is
# useful because splashback shells are not particularly meaningful for
# subhalos. It can be set to the following modes:
# none    - No halos are removed
# subhalo - Halos flagged as subhalos in the catalog are removed
# overlap - Halos which have an R200m shell that overlaps with a larger halo's
#           R200m shell are removed
#
# ExclusionStrategy defaults to overlap if not set.
#
# ExclusionStrategy = overlap

# ExclusionRadiusMult is a multiplier of R200m applied for the sake of
# determining exclusions.
#
# ExclusionRadiusMult defaults to 1 if not set.
#
# ExclustionRadiusMult = 1

# Mult is the number of times a given ID should be repeated. This is most useful
# if you want to estimate the scatter in shell measurements for halos with a
# given set of shell parameters.
#
# Mult defaults to 1 if not set.
#
# Mult = 1`
}

// ReadConfig reads in an id.config file into config.
func (config *IDConfig) ReadConfig(fname string) error {

	vars := parse.NewConfigVars("id.config")
	vars.String(&config.idType, "IDType", "m200m")
	vars.Ints(&config.ids, "IDs", []int64{})
	vars.Int(&config.idStart, "IDStart", -1)
	vars.Int(&config.idEnd, "IDEnd", -1)
	vars.Int(&config.mult, "Mult", 1)
	vars.Int(&config.snap, "Snap", -1)
	vars.String(&config.exclusionStrategy, "ExclusionStrategy", "overlap")
	vars.Float(&config.exclusionRadiusMult, "ExclusionRadiusMult", 1)

	if fname == "" {
		return nil
	}
	if err := parse.ReadConfig(fname, vars); err != nil {
		return err
	}

	return config.validate()
}

// validate checks whether all the fields of config are valid.
func (config *IDConfig) validate() error {
	switch config.idType {
	case "halo-id", "m200m":
	default:
		return fmt.Errorf("The 'IDType' variable is set to '%s', which I "+
			"don't recognize.", config.idType)
	}

	switch config.exclusionStrategy {
	case "none", "subhalo":
	case "overlap":
		if config.exclusionRadiusMult <= 0 {
			return fmt.Errorf("The 'ExclusionRadiusMult' varaible is set to "+
				"%g, but it needs to be positive.", config.exclusionRadiusMult)
		}
	default:
		return fmt.Errorf("The 'ExclusionStrategy' variable is set to '%s', "+
			"which I don't recognize.", config.exclusionStrategy)
	}

	// TODO: Check the ranges of the IDs as well as IDStart and IDEnd
	if len(config.ids) == 0 {
		switch {
		case config.idStart == -1 && config.idEnd == -1:
			return fmt.Errorf("'IDs' variable not set.")
		case config.idStart == -1:
			return fmt.Errorf("'IDStart variable not set.")
		case config.idEnd == -1:
			return fmt.Errorf("'IDEnd' variable not set.")
		case config.idEnd < config.idStart:
			return fmt.Errorf("'IDEnd' variable set to %d, but 'IDStart' "+
				"variable set to %d.", config.idEnd, config.idStart)
		}
	}

	switch {
	case config.snap == -1:
		return fmt.Errorf("'Snap' variable not set.")
	case config.snap < 0:
		return fmt.Errorf("'Snap' variable set to %d.", config.snap)
	}

	if config.mult <= 0 {
		return fmt.Errorf("'Mult' variable set to %d", config.mult)
	}

	return nil
}

// Run executes the ID mode of shellfish tool.
func (config *IDConfig) Run(
	flags []string, gConfig *GlobalConfig, e *env.Environment, stdin []string,
) ([]string, error) {

	if logging.Mode != logging.Nil {
		log.Println(`
##################
## shellfish id ##
##################`,
		)
	}
	var t time.Time
	if logging.Mode == logging.Performance {
		t = time.Now()
	}

	if config.snap == -1 {
		return nil, fmt.Errorf("Either no id.config file was provided or " +
			"the 'Snap' variable wasn't set.")
	}

	if config.snap < gConfig.SnapMin || config.snap > gConfig.SnapMax {
		return nil, fmt.Errorf("'Snap' = %d, but 'SnapMin' = %d and "+
			"'SnapMax = %d'", config.snap, gConfig.SnapMin, gConfig.SnapMax)
	}

	// Get IDs and snapshots

	rawIds := getIDs(config.idStart, config.idEnd, config.ids)

	vars := &halo.VarColumns{
		ID:    int(gConfig.HaloIDColumn),
		X:     int(gConfig.HaloPositionColumns[0]),
		Y:     int(gConfig.HaloPositionColumns[1]),
		Z:     int(gConfig.HaloPositionColumns[2]),
		M200m: int(gConfig.HaloM200mColumn),
	}

	var (
		ids, snaps []int
		buf        io.VectorBuffer
	)

	switch config.idType {
	case "halo-id":
		snaps = make([]int, len(rawIds))
		for i := range snaps {
			snaps[i] = int(config.snap)
		}
		ids = rawIds

		var err error
		buf, err = getVectorBuffer(
			e.ParticleCatalog(snaps[0], 0),
			gConfig.SnapshotType, gConfig.Endianness,
		)
		if err != nil {
			return nil, err
		}
	case "m200m":
		snaps = make([]int, len(rawIds))
		for i := range snaps {
			snaps[i] = int(config.snap)
		}

		var err error
		buf, err = getVectorBuffer(
			e.ParticleCatalog(snaps[0], 0),
			gConfig.SnapshotType, gConfig.Endianness,
		)
		if err != nil {
			return nil, err
		}

		ids, err = convertSortedIDs(rawIds, int(config.snap), vars, buf, e)
		if err != nil {
			return nil, err
		}

	default:
		panic("Impossible")
	}

	// Tag subhalos, if neccessary.
	exclude := make([]bool, len(ids))
	switch config.exclusionStrategy {
	case "none":
	case "subhalo":
		panic("subhalo is not implemented")
	case "overlap":
		var err error
		exclude, err = findOverlapSubs(ids, snaps, vars, buf, e, config)
		if err != nil {
			return nil, err
		}
	}

	// Generate lines
	intCols := [][]int{ids, snaps}
	floatCols := [][]float64{}
	colOrder := []int{0, 1}
	lines := catalog.FormatCols(intCols, floatCols, colOrder)

	// Filter
	fLines := []string{}
	for i := range lines {
		if !exclude[i] {
			fLines = append(fLines, lines[i])
		}
	}

	// Multiply
	mLines := []string{}
	for i := range fLines {
		for j := 0; j < int(config.mult); j++ {
			mLines = append(mLines, fLines[i])
		}
	}

	cString := catalog.CommentString(
		[]string{"ID", "Snapshot"}, []string{}, []int{0, 1}, []int{1, 1},
	)
	mLines = append([]string{cString}, mLines...)

	if logging.Mode == logging.Performance {
		log.Printf("Time: %s", time.Since(t).String())
		log.Printf("Memory:\n%s", logging.MemString())
	}

	return mLines, nil
}

func getIDs(idStart, idEnd int64, ids []int64) []int {
	if idStart != -1 {
		out := make([]int, idEnd-idStart)
		for i := range out {
			out[i] = int(idStart) + i
		}
		return out
	} else {
		out := make([]int, len(ids))
		for i := range out {
			out[i] = int(ids[i])
		}
		return out
	}
}

func convertSortedIDs(
	rawIDs []int, snap int, vars *halo.VarColumns,
	buf io.VectorBuffer, e *env.Environment,
) ([]int, error) {
	maxID := 0
	for _, id := range rawIDs {
		if id > maxID {
			maxID = id
		}
	}

	rids, err := memo.ReadSortedRockstarIDs(snap, maxID, vars, buf, e)
	if err != nil {
		return nil, err
	}

	ids := make([]int, len(rawIDs))
	for i := range ids {
		ids[i] = rids[rawIDs[i]]
	}
	return ids, nil
}

func findOverlapSubs(
	rawIDs, snaps []int, vars *halo.VarColumns,
	buf io.VectorBuffer, e *env.Environment, config *IDConfig,
) ([]bool, error) {
	isSub := make([]bool, len(rawIDs))

	// Group by snapshot.
	snapGroups := make(map[int][]int)
	groupIdxs := make(map[int][]int)
	for i, id := range rawIDs {
		snap := snaps[i]
		snapGroups[snap] = append(snapGroups[snap], id)
		groupIdxs[snap] = append(groupIdxs[snap], i)
	}

	// Load each snapshot.
	hds, _, err := memo.ReadHeaders(snaps[0], buf, e)
	if err != nil {
		return nil, err
	}
	hd := hds[0]

	for snap, group := range snapGroups {
		rids, err := memo.ReadSortedRockstarIDs(snap, -1, vars, buf, e)
		if err != nil {
			return nil, err
		}
		_, xs, ys, zs, _, rs, err := memo.ReadRockstar(snap, rids, vars, buf, e)

		g := halo.NewGrid(finderCells, hd.TotalWidth, len(xs))
		g.Insert(xs, ys, zs)
		sf := halo.NewSubhaloFinder(g)
		sf.FindSubhalos(xs, ys, zs, rs, config.exclusionRadiusMult)

		for i, id := range group {
			origIdx := groupIdxs[snap][i]
			// TODO: Holy linear search, batman! Fix this.
			for j, checkID := range rids {
				if checkID == id {
					isSub[origIdx] = sf.HostCount(j) > 0
					break
				} else if j == len(rids)-1 {
					return nil, fmt.Errorf("ID %d not in halo list.", id)
				}
			}
		}
	}
	return isSub, nil
}
<commit_msg>Added a subhalo-finding mode.<commit_after>package cmd

import (
	"fmt"
	"log"
	"time"

	"github.com/phil-mansfield/shellfish/cmd/catalog"
	"github.com/phil-mansfield/shellfish/cmd/env"
	"github.com/phil-mansfield/shellfish/cmd/halo"
	"github.com/phil-mansfield/shellfish/cmd/memo"
	"github.com/phil-mansfield/shellfish/io"
	"github.com/phil-mansfield/shellfish/logging"
	"github.com/phil-mansfield/shellfish/parse"
)

const finderCells = 150

// IDConfig contains the configuration fileds for the 'id' mode of the shellfish
// tool.
type IDConfig struct {
	idType                     string
	ids                        []int64
	idStart, idEnd, snap, mult int64

	exclusionStrategy   string
	exclusionRadiusMult float64
}

var _ Mode = &IDConfig{}

// ExampleConfig creates an example id.config file.
func (config *IDConfig) ExampleConfig() string {
	return `[id.config]
#####################
## Required Fields ##
#####################

# Index of the snapshot to be analyzed.
Snap = 100

IDs = 10, 11, 12, 13, 14

#####################
## Optional Fields ##
#####################

# IDType indicates what the input IDs correspond to. It can be set to the
# following modes:
# halo-id - The numeric IDs given in the halo catalog.
# m200m   - The rank of the halos when sorted by M200m.
#
# Defaults to m200m if not set.
# IDType = m200m

# An alternative way of specifying IDs is to select start and end (inclusive)
# ID values. If the IDs variable is not set, both of these values must be set.
#
# IDStart = 10
# IDEnd = 15

# ExclusionStrategy determines how to exclude IDs from the given set. This is
# useful because splashback shells are not particularly meaningful for
# subhalos. It can be set to the following modes:
# none    - No halos are removed
# subhalo - Halos flagged as subhalos in the catalog are removed
# overlap - Halos which have an R200m shell that overlaps with a larger halo's
#           R200m shell are removed
#
# ExclusionStrategy defaults to overlap if not set.
#
# ExclusionStrategy = overlap

# ExclusionRadiusMult is a multiplier of R200m applied for the sake of
# determining exclusions.
#
# ExclusionRadiusMult defaults to 1 if not set.
#
# ExclustionRadiusMult = 1

# Mult is the number of times a given ID should be repeated. This is most useful
# if you want to estimate the scatter in shell measurements for halos with a
# given set of shell parameters.
#
# Mult defaults to 1 if not set.
#
# Mult = 1`
}

// ReadConfig reads in an id.config file into config.
func (config *IDConfig) ReadConfig(fname string) error {

	vars := parse.NewConfigVars("id.config")
	vars.String(&config.idType, "IDType", "m200m")
	vars.Ints(&config.ids, "IDs", []int64{})
	vars.Int(&config.idStart, "IDStart", -1)
	vars.Int(&config.idEnd, "IDEnd", -1)
	vars.Int(&config.mult, "Mult", 1)
	vars.Int(&config.snap, "Snap", -1)
	vars.String(&config.exclusionStrategy, "ExclusionStrategy", "overlap")
	vars.Float(&config.exclusionRadiusMult, "ExclusionRadiusMult", 1)

	if fname == "" {
		return nil
	}
	if err := parse.ReadConfig(fname, vars); err != nil {
		return err
	}

	return config.validate()
}

// validate checks whether all the fields of config are valid.
func (config *IDConfig) validate() error {
	switch config.idType {
	case "halo-id", "m200m":
	default:
		return fmt.Errorf("The 'IDType' variable is set to '%s', which I "+
			"don't recognize.", config.idType)
	}

	switch config.exclusionStrategy {
	case "none", "subhalo":
	case "overlap":
		if config.exclusionRadiusMult <= 0 {
			return fmt.Errorf("The 'ExclusionRadiusMult' varaible is set to "+
				"%g, but it needs to be positive.", config.exclusionRadiusMult)
		}
	default:
		return fmt.Errorf("The 'ExclusionStrategy' variable is set to '%s', "+
			"which I don't recognize.", config.exclusionStrategy)
	}

	// TODO: Check the ranges of the IDs as well as IDStart and IDEnd
	if len(config.ids) == 0 {
		switch {
		case config.idStart == -1 && config.idEnd == -1:
			return fmt.Errorf("'IDs' variable not set.")
		case config.idStart == -1:
			return fmt.Errorf("'IDStart variable not set.")
		case config.idEnd == -1:
			return fmt.Errorf("'IDEnd' variable not set.")
		case config.idEnd < config.idStart:
			return fmt.Errorf("'IDEnd' variable set to %d, but 'IDStart' "+
				"variable set to %d.", config.idEnd, config.idStart)
		}
	}

	switch {
	case config.snap == -1:
		return fmt.Errorf("'Snap' variable not set.")
	case config.snap < 0:
		return fmt.Errorf("'Snap' variable set to %d.", config.snap)
	}

	if config.mult <= 0 {
		return fmt.Errorf("'Mult' variable set to %d", config.mult)
	}

	return nil
}

// Run executes the ID mode of shellfish tool.
func (config *IDConfig) Run(
	flags []string, gConfig *GlobalConfig, e *env.Environment, stdin []string,
) ([]string, error) {

	if logging.Mode != logging.Nil {
		log.Println(`
##################
## shellfish id ##
##################`,
		)
	}
	var t time.Time
	if logging.Mode == logging.Performance {
		t = time.Now()
	}

	if config.snap == -1 {
		return nil, fmt.Errorf("Either no id.config file was provided or " +
			"the 'Snap' variable wasn't set.")
	}

	if config.snap < gConfig.SnapMin || config.snap > gConfig.SnapMax {
		return nil, fmt.Errorf("'Snap' = %d, but 'SnapMin' = %d and "+
			"'SnapMax = %d'", config.snap, gConfig.SnapMin, gConfig.SnapMax)
	}

	// Get IDs and snapshots

	rawIds, err := getIDs(config.idStart, config.idEnd, config.ids, stdin)
	if err != nil {
		return nil, err
	}

	vars := &halo.VarColumns{
		ID:    int(gConfig.HaloIDColumn),
		X:     int(gConfig.HaloPositionColumns[0]),
		Y:     int(gConfig.HaloPositionColumns[1]),
		Z:     int(gConfig.HaloPositionColumns[2]),
		M200m: int(gConfig.HaloM200mColumn),
	}

	var (
		ids, snaps []int
		buf        io.VectorBuffer
	)

	switch config.idType {
	case "halo-id":
		snaps = make([]int, len(rawIds))
		for i := range snaps {
			snaps[i] = int(config.snap)
		}
		ids = rawIds

		var err error
		buf, err = getVectorBuffer(
			e.ParticleCatalog(snaps[0], 0),
			gConfig.SnapshotType, gConfig.Endianness,
		)
		if err != nil {
			return nil, err
		}
	case "m200m":
		snaps = make([]int, len(rawIds))
		for i := range snaps {
			snaps[i] = int(config.snap)
		}

		var err error
		buf, err = getVectorBuffer(
			e.ParticleCatalog(snaps[0], 0),
			gConfig.SnapshotType, gConfig.Endianness,
		)
		if err != nil {
			return nil, err
		}

		ids, err = convertSortedIDs(rawIds, int(config.snap), vars, buf, e)
		if err != nil {
			return nil, err
		}

	default:
		panic("Impossible")
	}

	// Tag subhalos, if neccessary.
	exclude := make([]bool, len(ids))
	switch config.exclusionStrategy {
	case "none":
	case "subhalo":
		panic("subhalo is not implemented")
	case "neighbors":
		buf, err := getVectorBuffer(
			e.ParticleCatalog(snaps[0], 0),
			gConfig.SnapshotType, gConfig.Endianness,
		)
		if err != nil {
			return nil, err
		}

		ids, snaps, err = readSubIDs(ids, snaps, vars, buf, e, config)
		if err != nil {
			return nil, err
		}

		exclude = make([]bool, len(ids))
	case "overlap":
		var err error
		exclude, err = findOverlapSubs(ids, snaps, vars, buf, e, config)
		if err != nil {
			return nil, err
		}
	}

	// Generate lines
	intCols := [][]int{ids, snaps}
	floatCols := [][]float64{}
	colOrder := []int{0, 1}
	lines := catalog.FormatCols(intCols, floatCols, colOrder)

	// Filter
	fLines := []string{}
	for i := range lines {
		if !exclude[i] {
			fLines = append(fLines, lines[i])
		}
	}

	// Multiply
	mLines := []string{}
	for i := range fLines {
		for j := 0; j < int(config.mult); j++ {
			mLines = append(mLines, fLines[i])
		}
	}

	cString := catalog.CommentString(
		[]string{"ID", "Snapshot"}, []string{}, []int{0, 1}, []int{1, 1},
	)
	mLines = append([]string{cString}, mLines...)

	if logging.Mode == logging.Performance {
		log.Printf("Time: %s", time.Since(t).String())
		log.Printf("Memory:\n%s", logging.MemString())
	}

	return mLines, nil
}

func getIDs(idStart, idEnd int64, ids []int64, stdin []string) ([]int, error) {
	if idStart != -1 {
		out := make([]int, idEnd-idStart)
		for i := range out {
			out[i] = int(idStart) + i
		}
		return out, nil
	} else if len(ids) > 0 {
		out := make([]int, len(ids))
		for i := range out {
			out[i] = int(ids[i])
		}
		return out, nil
	} else {
		intCols, _, err := catalog.ParseCols(stdin, []int{0}, []int{})
		if err != nil {
			return nil, err
		}
		return intCols[0], nil
	}
}

func convertSortedIDs(
	rawIDs []int, snap int, vars *halo.VarColumns,
	buf io.VectorBuffer, e *env.Environment,
) ([]int, error) {
	maxID := 0
	for _, id := range rawIDs {
		if id > maxID {
			maxID = id
		}
	}

	rids, err := memo.ReadSortedRockstarIDs(snap, maxID, vars, buf, e)
	if err != nil {
		return nil, err
	}

	ids := make([]int, len(rawIDs))
	for i := range ids {
		ids[i] = rids[rawIDs[i]]
	}
	return ids, nil
}

func findOverlapSubs(
	rawIDs, snaps []int, vars *halo.VarColumns,
	buf io.VectorBuffer, e *env.Environment, config *IDConfig,
) ([]bool, error) {
	isSub := make([]bool, len(rawIDs))

	// Group by snapshot.
	snapGroups := make(map[int][]int)
	groupIdxs := make(map[int][]int)
	for i, id := range rawIDs {
		snap := snaps[i]
		snapGroups[snap] = append(snapGroups[snap], id)
		groupIdxs[snap] = append(groupIdxs[snap], i)
	}

	// Load each snapshot.
	hds, _, err := memo.ReadHeaders(snaps[0], buf, e)
	if err != nil {
		return nil, err
	}
	hd := hds[0]

	for snap, group := range snapGroups {
		rids, err := memo.ReadSortedRockstarIDs(snap, -1, vars, buf, e)
		if err != nil {
			return nil, err
		}
		_, xs, ys, zs, _, rs, err := memo.ReadRockstar(snap, rids, vars, buf, e)

		g := halo.NewGrid(finderCells, hd.TotalWidth, len(xs))
		g.Insert(xs, ys, zs)
		sf := halo.NewSubhaloFinder(g)
		sf.FindSubhalos(xs, ys, zs, rs, config.exclusionRadiusMult)

		for i, id := range group {
			origIdx := groupIdxs[snap][i]
			// TODO: Holy linear search, batman! Fix this.
			for j, checkID := range rids {
				if checkID == id {
					isSub[origIdx] = sf.HostCount(j) > 0
					break
				} else if j == len(rids)-1 {
					return nil, fmt.Errorf("ID %d not in halo list.", id)
				}
			}
		}
	}
	return isSub, nil
}

func readSubIDs(
	ids, snaps []int, vars *halo.VarColumns,
	buf io.VectorBuffer, e *env.Environment,
	config *IDConfig,
) (
	sIDs, sSnaps []int, err error,
) {
	snapBins, idxBins := binBySnap(snaps, ids)
	subIDs := make([][]int, len(ids))

	hds, _, err := memo.ReadHeaders(snaps[0], buf, e)
	if err != nil {
		return nil, nil, err
	}
	hd := hds[0]

	for snap, _ := range snapBins {
		if snap == -1 {
			continue
		}
		snapIDs := snapBins[snap]
		idxs := idxBins[snap]

		allIDs, err := memo.ReadSortedRockstarIDs(snap, -1, vars, buf, e)
		if err != nil { return nil, nil, err }
		_, xs, ys, zs, _, rs, err := memo.ReadRockstar(
			snap, allIDs, vars, buf, e,
		)

		g := halo.NewGrid(finderCells, hd.TotalWidth, len(xs))
		g.Insert(xs, ys, xs)
		sf := halo.NewSubhaloFinder(g)
		sf.FindSubhalos(xs, ys, zs, rs, config.exclusionRadiusMult)

		f := newIntFinder(allIDs)
		for i, id := range snapIDs {
			idx := idxs[i]
			sIdx, ok := f.find(id)
			if !ok {
				return nil, nil, fmt.Errorf("Could not find ID %d.", id)
			}

			subMassIDs := sf.Subhalos(sIdx)
			for i := range subMassIDs {
				subIDs[idx][i] = allIDs[subMassIDs[i]]
			}
		}
	}
	sIDs, sSnaps = []int{}, []int{}
	for i := range subIDs {
		sIDs = append(sIDs, ids[i])
		sSnaps = append(sSnaps, snaps[i])
		for _, id := range subIDs[i] {
			sIDs = append(sIDs, id)
			sSnaps = append(sSnaps, snaps[i])
		}
	}

	return sIDs, sSnaps, nil
}

// A quick generic wrapper for doing those one-to-one mappings I need to do so
// often. Written like this so the backend can be swapped out easily.
type intFinder struct {
	m map[int]int
}

// NewIntFinder creates a new IntFinder struct for a given slice of Rockstar
// IDs.
func newIntFinder(rids []int) intFinder {
	f := intFinder{}
	f.m = make(map[int]int)
	for i, rid := range rids { f.m[rid] = i }
	return f
}

// Find returns the index which the given ID corresponds to and true if the
// ID is in the finder. Otherwise, false is returned.
func (f intFinder) find(rid int) (int, bool) {
	line, ok := f.m[rid]
	return line, ok
}<|endoftext|>


<commit_before>package libnetwork

import (
	"fmt"
	"io/ioutil"
	"os"
	"testing"

	"github.com/docker/libkv/store"
	"github.com/docker/libnetwork/config"
	"github.com/docker/libnetwork/datastore"
	"github.com/docker/libnetwork/netlabel"
	"github.com/docker/libnetwork/options"
)

func TestZooKeeperBackend(t *testing.T) {
	c, err := testNewController(t, "zk", "127.0.0.1:2181")
	if err != nil {
		t.Fatal(err)
	}
	c.Stop()
}

func testNewController(t *testing.T, provider, url string) (NetworkController, error) {
	cfgOptions, err := OptionBoltdbWithRandomDBFile()
	if err != nil {
		return nil, err
	}
	cfgOptions = append(cfgOptions, config.OptionKVProvider(provider))
	cfgOptions = append(cfgOptions, config.OptionKVProviderURL(url))
	return New(cfgOptions...)
}

func TestBoltdbBackend(t *testing.T) {
	defer os.Remove(datastore.DefaultScopes("")[datastore.LocalScope].Client.Address)
	testLocalBackend(t, "", "", nil)
	defer os.Remove("/tmp/boltdb.db")
	config := &store.Config{Bucket: "testBackend"}
	testLocalBackend(t, "boltdb", "/tmp/boltdb.db", config)

}

func testLocalBackend(t *testing.T, provider, url string, storeConfig *store.Config) {
	cfgOptions := []config.Option{}
	cfgOptions = append(cfgOptions, config.OptionLocalKVProvider(provider))
	cfgOptions = append(cfgOptions, config.OptionLocalKVProviderURL(url))
	cfgOptions = append(cfgOptions, config.OptionLocalKVProviderConfig(storeConfig))

	driverOptions := options.Generic{}
	genericOption := make(map[string]interface{})
	genericOption[netlabel.GenericData] = driverOptions
	cfgOptions = append(cfgOptions, config.OptionDriverConfig("host", genericOption))

	ctrl, err := New(cfgOptions...)
	if err != nil {
		t.Fatalf("Error new controller: %v", err)
	}
	nw, err := ctrl.NewNetwork("host", "host")
	if err != nil {
		t.Fatalf("Error creating default \"host\" network: %v", err)
	}
	ep, err := nw.CreateEndpoint("newendpoint", []EndpointOption{}...)
	if err != nil {
		t.Fatalf("Error creating endpoint: %v", err)
	}
	store := ctrl.(*controller).getStore(datastore.LocalScope).KVStore()
	if exists, err := store.Exists(datastore.Key(datastore.NetworkKeyPrefix, string(nw.ID()))); !exists || err != nil {
		t.Fatalf("Network key should have been created.")
	}
	if exists, err := store.Exists(datastore.Key([]string{datastore.EndpointKeyPrefix, string(nw.ID()), string(ep.ID())}...)); !exists || err != nil {
		t.Fatalf("Endpoint key should have been created.")
	}
	store.Close()

	// test restore of local store
	ctrl, err = New(cfgOptions...)
	if err != nil {
		t.Fatalf("Error creating controller: %v", err)
	}
	if _, err = ctrl.NetworkByID(nw.ID()); err != nil {
		t.Fatalf("Error getting network %v", err)
	}
}

func TestNoPersist(t *testing.T) {
	cfgOptions, err := OptionBoltdbWithRandomDBFile()
	if err != nil {
		t.Fatalf("Error creating random boltdb file : %v", err)
	}
	ctrl, err := New(cfgOptions...)
	if err != nil {
		t.Fatalf("Error new controller: %v", err)
	}
	nw, err := ctrl.NewNetwork("host", "host", NetworkOptionPersist(false))
	if err != nil {
		t.Fatalf("Error creating default \"host\" network: %v", err)
	}
	ep, err := nw.CreateEndpoint("newendpoint", []EndpointOption{}...)
	if err != nil {
		t.Fatalf("Error creating endpoint: %v", err)
	}
	store := ctrl.(*controller).getStore(datastore.LocalScope).KVStore()
	if exists, _ := store.Exists(datastore.Key(datastore.NetworkKeyPrefix, string(nw.ID()))); exists {
		t.Fatalf("Network with persist=false should not be stored in KV Store")
	}
	if exists, _ := store.Exists(datastore.Key([]string{datastore.EndpointKeyPrefix, string(nw.ID()), string(ep.ID())}...)); exists {
		t.Fatalf("Endpoint in Network with persist=false should not be stored in KV Store")
	}
	store.Close()
}

// OptionBoltdbWithRandomDBFile function returns a random dir for local store backend
func OptionBoltdbWithRandomDBFile() ([]config.Option, error) {
	tmp, err := ioutil.TempFile("", "libnetwork-")
	if err != nil {
		return nil, fmt.Errorf("Error creating temp file: %v", err)
	}
	if err := tmp.Close(); err != nil {
		return nil, fmt.Errorf("Error closing temp file: %v", err)
	}
	cfgOptions := []config.Option{}
	cfgOptions = append(cfgOptions, config.OptionLocalKVProvider("boltdb"))
	cfgOptions = append(cfgOptions, config.OptionLocalKVProviderURL(tmp.Name()))
	sCfg := &store.Config{Bucket: "testBackend"}
	cfgOptions = append(cfgOptions, config.OptionLocalKVProviderConfig(sCfg))
	return cfgOptions, nil
}

func TestMultipleControllersWithSameStore(t *testing.T) {
	cfgOptions, err := OptionBoltdbWithRandomDBFile()
	if err != nil {
		t.Fatalf("Error getting random boltdb configs %v", err)
	}
	ctrl1, err := New(cfgOptions...)
	if err != nil {
		t.Fatalf("Error new controller: %v", err)
	}
	defer ctrl1.Stop()
	// Use the same boltdb file without closing the previous controller
	_, err = New(cfgOptions...)
	if err != nil {
		t.Fatalf("Local store must support concurrent controllers")
	}
}
<commit_msg>Add customer_prefix to the store URL for integration test<commit_after>package libnetwork

import (
	"fmt"
	"io/ioutil"
	"os"
	"testing"

	"github.com/docker/libkv/store"
	"github.com/docker/libnetwork/config"
	"github.com/docker/libnetwork/datastore"
	"github.com/docker/libnetwork/netlabel"
	"github.com/docker/libnetwork/options"
)

func TestZooKeeperBackend(t *testing.T) {
	c, err := testNewController(t, "zk", "127.0.0.1:2181/custom_prefix")
	if err != nil {
		t.Fatal(err)
	}
	c.Stop()
}

func testNewController(t *testing.T, provider, url string) (NetworkController, error) {
	cfgOptions, err := OptionBoltdbWithRandomDBFile()
	if err != nil {
		return nil, err
	}
	cfgOptions = append(cfgOptions, config.OptionKVProvider(provider))
	cfgOptions = append(cfgOptions, config.OptionKVProviderURL(url))
	return New(cfgOptions...)
}

func TestBoltdbBackend(t *testing.T) {
	defer os.Remove(datastore.DefaultScopes("")[datastore.LocalScope].Client.Address)
	testLocalBackend(t, "", "", nil)
	defer os.Remove("/tmp/boltdb.db")
	config := &store.Config{Bucket: "testBackend"}
	testLocalBackend(t, "boltdb", "/tmp/boltdb.db", config)

}

func testLocalBackend(t *testing.T, provider, url string, storeConfig *store.Config) {
	cfgOptions := []config.Option{}
	cfgOptions = append(cfgOptions, config.OptionLocalKVProvider(provider))
	cfgOptions = append(cfgOptions, config.OptionLocalKVProviderURL(url))
	cfgOptions = append(cfgOptions, config.OptionLocalKVProviderConfig(storeConfig))

	driverOptions := options.Generic{}
	genericOption := make(map[string]interface{})
	genericOption[netlabel.GenericData] = driverOptions
	cfgOptions = append(cfgOptions, config.OptionDriverConfig("host", genericOption))

	ctrl, err := New(cfgOptions...)
	if err != nil {
		t.Fatalf("Error new controller: %v", err)
	}
	nw, err := ctrl.NewNetwork("host", "host")
	if err != nil {
		t.Fatalf("Error creating default \"host\" network: %v", err)
	}
	ep, err := nw.CreateEndpoint("newendpoint", []EndpointOption{}...)
	if err != nil {
		t.Fatalf("Error creating endpoint: %v", err)
	}
	store := ctrl.(*controller).getStore(datastore.LocalScope).KVStore()
	if exists, err := store.Exists(datastore.Key(datastore.NetworkKeyPrefix, string(nw.ID()))); !exists || err != nil {
		t.Fatalf("Network key should have been created.")
	}
	if exists, err := store.Exists(datastore.Key([]string{datastore.EndpointKeyPrefix, string(nw.ID()), string(ep.ID())}...)); !exists || err != nil {
		t.Fatalf("Endpoint key should have been created.")
	}
	store.Close()

	// test restore of local store
	ctrl, err = New(cfgOptions...)
	if err != nil {
		t.Fatalf("Error creating controller: %v", err)
	}
	if _, err = ctrl.NetworkByID(nw.ID()); err != nil {
		t.Fatalf("Error getting network %v", err)
	}
}

func TestNoPersist(t *testing.T) {
	cfgOptions, err := OptionBoltdbWithRandomDBFile()
	if err != nil {
		t.Fatalf("Error creating random boltdb file : %v", err)
	}
	ctrl, err := New(cfgOptions...)
	if err != nil {
		t.Fatalf("Error new controller: %v", err)
	}
	nw, err := ctrl.NewNetwork("host", "host", NetworkOptionPersist(false))
	if err != nil {
		t.Fatalf("Error creating default \"host\" network: %v", err)
	}
	ep, err := nw.CreateEndpoint("newendpoint", []EndpointOption{}...)
	if err != nil {
		t.Fatalf("Error creating endpoint: %v", err)
	}
	store := ctrl.(*controller).getStore(datastore.LocalScope).KVStore()
	if exists, _ := store.Exists(datastore.Key(datastore.NetworkKeyPrefix, string(nw.ID()))); exists {
		t.Fatalf("Network with persist=false should not be stored in KV Store")
	}
	if exists, _ := store.Exists(datastore.Key([]string{datastore.EndpointKeyPrefix, string(nw.ID()), string(ep.ID())}...)); exists {
		t.Fatalf("Endpoint in Network with persist=false should not be stored in KV Store")
	}
	store.Close()
}

// OptionBoltdbWithRandomDBFile function returns a random dir for local store backend
func OptionBoltdbWithRandomDBFile() ([]config.Option, error) {
	tmp, err := ioutil.TempFile("", "libnetwork-")
	if err != nil {
		return nil, fmt.Errorf("Error creating temp file: %v", err)
	}
	if err := tmp.Close(); err != nil {
		return nil, fmt.Errorf("Error closing temp file: %v", err)
	}
	cfgOptions := []config.Option{}
	cfgOptions = append(cfgOptions, config.OptionLocalKVProvider("boltdb"))
	cfgOptions = append(cfgOptions, config.OptionLocalKVProviderURL(tmp.Name()))
	sCfg := &store.Config{Bucket: "testBackend"}
	cfgOptions = append(cfgOptions, config.OptionLocalKVProviderConfig(sCfg))
	return cfgOptions, nil
}

func TestMultipleControllersWithSameStore(t *testing.T) {
	cfgOptions, err := OptionBoltdbWithRandomDBFile()
	if err != nil {
		t.Fatalf("Error getting random boltdb configs %v", err)
	}
	ctrl1, err := New(cfgOptions...)
	if err != nil {
		t.Fatalf("Error new controller: %v", err)
	}
	defer ctrl1.Stop()
	// Use the same boltdb file without closing the previous controller
	_, err = New(cfgOptions...)
	if err != nil {
		t.Fatalf("Local store must support concurrent controllers")
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2020 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 deployer

import (
	"errors"
	"fmt"
	"os"
	"path"
	"strings"

	"k8s.io/kops/tests/e2e/pkg/util"
	"sigs.k8s.io/kubetest2/pkg/exec"
)

const (
	defaultJobName = "pull-kops-e2e-kubernetes-aws"
	defaultGCSPath = "gs://kops-ci/pulls/%v/pull-%v"
)

func (d *deployer) Build() error {
	if err := d.init(); err != nil {
		return err
	}
	if err := d.BuildOptions.Build(); err != nil {
		return err
	}
	// Copy the kops binary into the test's RunDir to be included in the tester's PATH
	if d.KopsBinaryPath != "" {
		return util.Copy(d.KopsBinaryPath, path.Join(d.commonOptions.RunDir(), "kops"))
	} else {
		return nil
	}
}

func (d *deployer) verifyBuildFlags() error {
	if d.KopsRoot == "" {
		if goPath := os.Getenv("GOPATH"); goPath != "" {
			d.KopsRoot = path.Join(goPath, "src", "k8s.io", "kops")
		} else {
			return errors.New("required --kops-root when building from source")
		}
	}
	if d.StageLocation != "" {
		if !strings.HasPrefix(d.StageLocation, "gs://") {
			return errors.New("stage-location must be a gs:// path")
		}
	} else {
		stageLocation, err := defaultStageLocation(d.KopsRoot)
		if err != nil {
			return err
		}
		d.StageLocation = stageLocation
	}
	if d.KopsBaseURL == "" && os.Getenv("KOPS_BASE_URL") == "" {
		d.KopsBaseURL = strings.Replace(d.StageLocation, "gs://", "https://storage.googleapis.com/", 1)
	}

	fi, err := os.Stat(d.KopsRoot)
	if err != nil {
		return err
	}
	if !fi.Mode().IsDir() {
		return errors.New("--kops-root must be a directory")
	}
	if d.KopsVersionMarker != "" {
		return errors.New("cannot use --kops-version-marker with --build")
	}

	d.BuildOptions.KopsRoot = d.KopsRoot
	d.BuildOptions.StageLocation = d.StageLocation
	return nil
}

func defaultStageLocation(kopsRoot string) (string, error) {
	jobName := os.Getenv("JOB_NAME")
	if jobName == "" {
		jobName = defaultJobName
	}

	cmd := exec.Command("git", "describe", "--always")
	cmd.SetDir(kopsRoot)
	output, err := exec.CombinedOutputLines(cmd)
	if err != nil {
		return "", err
	} else if len(output) != 1 {
		return "", fmt.Errorf("unexpected output from git describe: %v", output)
	}
	sha := strings.TrimSpace(output[0])
	return fmt.Sprintf(defaultGCSPath, jobName, sha), nil
}
<commit_msg>Use short commit sha for default stage location instead of git-describe<commit_after>/*
Copyright 2020 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 deployer

import (
	"errors"
	"fmt"
	"os"
	"path"
	"strings"

	"k8s.io/kops/tests/e2e/pkg/util"
	"sigs.k8s.io/kubetest2/pkg/exec"
)

const (
	defaultJobName = "pull-kops-e2e-kubernetes-aws"
	defaultGCSPath = "gs://kops-ci/pulls/%v/pull-%v"
)

func (d *deployer) Build() error {
	if err := d.init(); err != nil {
		return err
	}
	if err := d.BuildOptions.Build(); err != nil {
		return err
	}
	// Copy the kops binary into the test's RunDir to be included in the tester's PATH
	if d.KopsBinaryPath != "" {
		return util.Copy(d.KopsBinaryPath, path.Join(d.commonOptions.RunDir(), "kops"))
	} else {
		return nil
	}
}

func (d *deployer) verifyBuildFlags() error {
	if d.KopsRoot == "" {
		if goPath := os.Getenv("GOPATH"); goPath != "" {
			d.KopsRoot = path.Join(goPath, "src", "k8s.io", "kops")
		} else {
			return errors.New("required --kops-root when building from source")
		}
	}
	if d.StageLocation != "" {
		if !strings.HasPrefix(d.StageLocation, "gs://") {
			return errors.New("stage-location must be a gs:// path")
		}
	} else {
		stageLocation, err := defaultStageLocation(d.KopsRoot)
		if err != nil {
			return err
		}
		d.StageLocation = stageLocation
	}
	if d.KopsBaseURL == "" && os.Getenv("KOPS_BASE_URL") == "" {
		d.KopsBaseURL = strings.Replace(d.StageLocation, "gs://", "https://storage.googleapis.com/", 1)
	}

	fi, err := os.Stat(d.KopsRoot)
	if err != nil {
		return err
	}
	if !fi.Mode().IsDir() {
		return errors.New("--kops-root must be a directory")
	}
	if d.KopsVersionMarker != "" {
		return errors.New("cannot use --kops-version-marker with --build")
	}

	d.BuildOptions.KopsRoot = d.KopsRoot
	d.BuildOptions.StageLocation = d.StageLocation
	return nil
}

func defaultStageLocation(kopsRoot string) (string, error) {
	jobName := os.Getenv("JOB_NAME")
	if jobName == "" {
		jobName = defaultJobName
	}

	cmd := exec.Command("git", "rev-parse", "--short", "HEAD")
	cmd.SetDir(kopsRoot)
	output, err := exec.CombinedOutputLines(cmd)
	if err != nil {
		return "", err
	} else if len(output) != 1 {
		return "", fmt.Errorf("unexpected output from git describe: %v", output)
	}
	sha := strings.TrimSpace(output[0])
	return fmt.Sprintf(defaultGCSPath, jobName, sha), nil
}
<|endoftext|>


<commit_before>package utils

import (
	"database/sql"
	"reflect"
	"strings"
	"sync"
	"time"
)

// SQLScanHelper helper class for reading sql to Struct
// Columns must be named the same
// Ex: in sql a column name is col1, in struct must be Col1
type SQLScanHelper struct {
	sync.RWMutex
	columnNames []string
}

// Clear - clears the columns array.
// Used to be able to reuse the scan helper for another SQL
func (s *SQLScanHelper) Clear() {
	s.RLock()
	defer s.RUnlock()

	s.columnNames = nil
}

// Scan - reads sql statement into a struct
func (s *SQLScanHelper) Scan(u *DbUtils, rows *sql.Rows, dest interface{}) error {
	s.RLock()
	defer s.RUnlock()

	if s.columnNames == nil || len(s.columnNames) == 0 {
		cols, err := rows.Columns()
		if err != nil {
			return err
		}

		s.columnNames = cols
	}

	nrCols := len(s.columnNames)
	pointers := make([]interface{}, nrCols)
	fieldTypes := make([]reflect.Type, nrCols)

	structVal := reflect.ValueOf(dest).Elem()
	nFields := structVal.NumField()

	for i, colName := range s.columnNames {
		loweredColName := strings.ToLower(colName)
		for j := 0; j < nFields; j++ {
			typeField := structVal.Type().Field(j)
			tag := typeField.Tag

			if tag.Get("sql") == loweredColName {
				pointers[i] = structVal.Field(j).Addr().Interface()
				fieldTypes[i] = typeField.Type
				break
			}
		}
	}

	err := rows.Scan(pointers...)
	if err != nil {
		return err
	}

	if u.dbType == Oracle {
		// in oci, the timestamp is comming up as local time zone
		// even if you ask for the UTC
		dt := time.Now()
		dtType := reflect.TypeOf(dt)

		for i := 0; i < nrCols; i++ {
			if fieldTypes[i] == dtType {
				dtval := pointers[i].(*time.Time)
				strdt := Date2string(*dtval, ISODateTimestamp)
				*dtval = String2dateNoErr(strdt, UTCDateTimestamp)
			}
		}
	}

	return nil
}
<commit_msg>sql col tag<commit_after>package utils

import (
	"database/sql"
	"reflect"
	"strings"
	"sync"
	"time"
)

// SQLScanHelper helper class for reading sql to Struct
// Columns in struct must be marked with a `sql:"col_name"` tag
// Ex: in sql a column name is col1, in struct the col tag must be `sql:"col1"`
type SQLScanHelper struct {
	sync.RWMutex
	columnNames []string
}

// Clear - clears the columns array.
// Used to be able to reuse the scan helper for another SQL
func (s *SQLScanHelper) Clear() {
	s.RLock()
	defer s.RUnlock()

	s.columnNames = nil
}

// Scan - reads sql statement into a struct
func (s *SQLScanHelper) Scan(u *DbUtils, rows *sql.Rows, dest interface{}) error {
	s.RLock()
	defer s.RUnlock()

	if s.columnNames == nil || len(s.columnNames) == 0 {
		cols, err := rows.Columns()
		if err != nil {
			return err
		}

		s.columnNames = cols
	}

	nrCols := len(s.columnNames)
	pointers := make([]interface{}, nrCols)
	fieldTypes := make([]reflect.Type, nrCols)

	structVal := reflect.ValueOf(dest).Elem()
	nFields := structVal.NumField()

	for i, colName := range s.columnNames {
		loweredColName := strings.ToLower(colName)
		for j := 0; j < nFields; j++ {
			typeField := structVal.Type().Field(j)
			tag := typeField.Tag

			if tag.Get("sql") == loweredColName {
				pointers[i] = structVal.Field(j).Addr().Interface()
				fieldTypes[i] = typeField.Type
				break
			}
		}
	}

	err := rows.Scan(pointers...)
	if err != nil {
		return err
	}

	if u.dbType == Oracle {
		// in oci, the timestamp is comming up as local time zone
		// even if you ask for the UTC
		dt := time.Now()
		dtType := reflect.TypeOf(dt)

		for i := 0; i < nrCols; i++ {
			if fieldTypes[i] == dtType {
				dtval := pointers[i].(*time.Time)
				strdt := Date2string(*dtval, ISODateTimestamp)
				*dtval = String2dateNoErr(strdt, UTCDateTimestamp)
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes 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 mungers

import (
	"bytes"
	"fmt"
	"strings"

	"k8s.io/contrib/mungegithub/features"
	"k8s.io/contrib/mungegithub/github"

	"github.com/golang/glog"
	"github.com/spf13/cobra"
)

const (
	clearAfterMergeName = "cherrypick-clear-after-merge"
)

// ClearPickAfterMerge will remove the the cherrypick-candidate label from
// any PR that does not have a 'release' milestone set.
type ClearPickAfterMerge struct {
	features *features.Features
}

func init() {
	RegisterMungerOrDie(&ClearPickAfterMerge{})
}

// Name is the name usable in --pr-mungers
func (c *ClearPickAfterMerge) Name() string { return clearAfterMergeName }

// RequiredFeatures is a slice of 'features' that must be provided
func (c *ClearPickAfterMerge) RequiredFeatures() []string { return []string{features.RepoFeatureName} }

// Initialize will initialize the munger
func (c *ClearPickAfterMerge) Initialize(config *github.Config, features *features.Features) error {
	c.features = features
	return nil
}

// EachLoop is called at the start of every munge loop
func (c *ClearPickAfterMerge) EachLoop() error { return nil }

// AddFlags will add any request flags to the cobra `cmd`
func (c *ClearPickAfterMerge) AddFlags(cmd *cobra.Command, config *github.Config) {}

func handleFound(obj *github.MungeObject, gitMsg []byte, branch string) error {
	msg := string(gitMsg)
	o := strings.SplitN(msg, "\n", 2)
	sha := o[0]
	msg = fmt.Sprintf("Commit %s found in the %q branch appears to be this PR. Removing the %q label. If this s an error find help to get your PR picked.", sha, branch, cpCandidateLabel)
	obj.WriteComment(msg)
	obj.RemoveLabel(cpCandidateLabel)
	return nil
}

// Munge is the workhorse the will actually make updates to the PR
func (c *ClearPickAfterMerge) Munge(obj *github.MungeObject) {
	if !obj.IsPR() {
		return
	}
	if !obj.HasLabel(cpCandidateLabel) {
		return
	}

	if merged, err := obj.IsMerged(); !merged || err != nil {
		return
	}

	releaseMilestone := obj.ReleaseMilestone()
	if releaseMilestone == "" || len(releaseMilestone) != 4 {
		glog.Errorf("Found invalid milestone: %q", releaseMilestone)
		return
	}
	rel := releaseMilestone[1:]
	branch := "release-" + rel

	sha := obj.MergeCommit()
	if sha == nil {
		glog.Errorf("Unable to get SHA of merged %d", sha)
		return
	}

	logMsg := fmt.Sprintf("Merge pull request #%d from ", *obj.Issue.Number)
	bLogMsg := []byte(logMsg)

	cherrypickMsg := fmt.Sprintf("(cherry picked from commit %s)", *sha)
	args := []string{"log", "--pretty=tformat:%H%n%s%n%b", "--grep", cherrypickMsg, "origin/" + branch}
	out, err := c.features.Repos.GitCommand(args)
	if err != nil {
		glog.Errorf("Error grepping for cherrypick -x message out=%q: %v", string(out), err)
		return
	}
	if bytes.Contains(out, bLogMsg) {
		glog.Infof("Found cherry-pick using -x information")
		handleFound(obj, out, branch)
		return
	}

	args = []string{"log", "--pretty=tformat:%H%n%s%n%b", "--grep", logMsg, "origin/" + branch}
	out, err = c.features.Repos.GitCommand(args)
	if err != nil {
		glog.Errorf("Error grepping for log message out=%q: %v", string(out), err)
		return
	}
	if bytes.Contains(out, bLogMsg) {
		glog.Infof("Found cherry-pick using log matching")
		handleFound(obj, out, branch)
		return
	}

	return
}
<commit_msg>mungegithub: Remove cp-candidate label if the words "Automated cherry pick of #NNNNN" are in the branch<commit_after>/*
Copyright 2016 The Kubernetes 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 mungers

import (
	"fmt"
	"strings"

	"k8s.io/contrib/mungegithub/features"
	"k8s.io/contrib/mungegithub/github"

	"github.com/golang/glog"
	"github.com/spf13/cobra"
)

const (
	clearAfterMergeName = "cherrypick-clear-after-merge"
)

// ClearPickAfterMerge will remove the the cherrypick-candidate label from
// any PR that does not have a 'release' milestone set.
type ClearPickAfterMerge struct {
	features *features.Features
}

func init() {
	RegisterMungerOrDie(&ClearPickAfterMerge{})
}

// Name is the name usable in --pr-mungers
func (c *ClearPickAfterMerge) Name() string { return clearAfterMergeName }

// RequiredFeatures is a slice of 'features' that must be provided
func (c *ClearPickAfterMerge) RequiredFeatures() []string { return []string{features.RepoFeatureName} }

// Initialize will initialize the munger
func (c *ClearPickAfterMerge) Initialize(config *github.Config, features *features.Features) error {
	c.features = features
	return nil
}

// EachLoop is called at the start of every munge loop
func (c *ClearPickAfterMerge) EachLoop() error { return nil }

// AddFlags will add any request flags to the cobra `cmd`
func (c *ClearPickAfterMerge) AddFlags(cmd *cobra.Command, config *github.Config) {}

func handleFound(obj *github.MungeObject, branch string) error {
	msg := fmt.Sprintf("Commit found in the %q branch appears to be this PR. Removing the %q label. If this is an error find help to get your PR picked.", branch, cpCandidateLabel)
	obj.WriteComment(msg)
	obj.RemoveLabel(cpCandidateLabel)
	return nil
}

// foundLog will return if the given `logString` exists on the branch in question.
// it will also return the actual logs for further processing
func (c *ClearPickAfterMerge) foundLog(branch string, logString string) (bool, string) {
	args := []string{"merge-base", "--fork-point", "origin/master", "origin/" + branch}
	out, err := c.features.Repos.GitCommand(args)
	base := string(out)
	if err != nil {
		glog.Errorf("Unable to find the fork point for branch %s. %s:%v", branch, base, err)
		return false, ""
	}
	base = strings.TrimSuffix(base, "\n")

	// if release-1.2 branched from master at abcdef123 this should result in:
	// abcdef123..origin/release-1.2
	logRefs := fmt.Sprintf("%s..origin/%s", base, branch)

	args = []string{"log", "--pretty=tformat:%H%n%s%n%b", "--grep", logString, logRefs}
	out, err = c.features.Repos.GitCommand(args)
	logs := string(out)
	if err != nil {
		glog.Errorf("Error grepping logs out=%q: %v", logs, err)
		return false, ""
	}
	glog.V(10).Infof("args:%v", args)

	glog.V(10).Infof("Searching for %q in %q", logString, logs)
	if !strings.Contains(logs, logString) {
		return false, ""
	}
	return true, logs
}

// Can we find a commit in the changelog that looks like it was done using git cherry-pick -m1 -x ?
func (c *ClearPickAfterMerge) foundByPickDashX(obj *github.MungeObject, branch string) bool {
	sha := obj.MergeCommit()
	if sha == nil {
		glog.Errorf("Unable to get SHA of merged PR %d", *obj.Issue.Number)
		return false
	}

	cherrypickMsg := fmt.Sprintf("(cherry picked from commit %s)", *sha)
	found, logs := c.foundLog(branch, cherrypickMsg)
	if !found {
		return false
	}

	// double check for the 'non -x' message
	logMsg := fmt.Sprintf("Merge pull request #%d from ", *obj.Issue.Number)
	if !strings.Contains(logs, logMsg) {
		return false
	}
	glog.Infof("Found cherry-pick for %d using -x information in branch %q", *obj.Issue.Number, branch)
	return true
}

// Can we find a commit in the changelog that looks like it was done using git cherry-pick -m1 ?
func (c *ClearPickAfterMerge) foundByPickWithoutDashX(obj *github.MungeObject, branch string) bool {
	logMsg := fmt.Sprintf("Merge pull request #%d from ", *obj.Issue.Number)

	found, _ := c.foundLog(branch, logMsg)
	if found {
		glog.Infof("Found cherry-pick for %d using log matching for `git cherry-pick` in branch %q", *obj.Issue.Number, branch)
	}
	return found
}

// Check that the commit messages for all commits in the PR are on the branch
func (c *ClearPickAfterMerge) foundByAllCommits(obj *github.MungeObject, branch string) bool {
	commits, err := obj.GetCommits()
	if err != nil {
		glog.Infof("unable to get commits")
		return false
	}
	for _, commit := range commits {
		if commit.Commit == nil {
			return false
		}
		if commit.Commit.Message == nil {
			return false
		}
		found, _ := c.foundLog(branch, *commit.Commit.Message)
		if !found {
			glog.Infof("found == false")
			return false
		}
	}
	return true
}

// Can we find a commit in the changelog that looks like it was done using the hack/cherry_pick_pull.sh script ?
func (c *ClearPickAfterMerge) foundByScript(obj *github.MungeObject, branch string) bool {
	logMsg := fmt.Sprintf("Cherry pick of #%d on %s.", *obj.Issue.Number, branch)

	found, _ := c.foundLog(branch, logMsg)
	if found {
		glog.Infof("Found cherry-pick for %d using log matching for `hack/cherry_pick_pull.sh` in branch %q", *obj.Issue.Number, branch)
	}
	return found
}

// Munge is the workhorse the will actually make updates to the PR
func (c *ClearPickAfterMerge) Munge(obj *github.MungeObject) {
	if !obj.IsPR() {
		return
	}
	if !obj.HasLabel(cpCandidateLabel) {
		return
	}

	if merged, err := obj.IsMerged(); !merged || err != nil {
		return
	}

	releaseMilestone := obj.ReleaseMilestone()
	if releaseMilestone == "" || len(releaseMilestone) != 4 {
		glog.Errorf("Found invalid milestone: %q", releaseMilestone)
		return
	}
	rel := releaseMilestone[1:]
	branch := "release-" + rel

	if c.foundByPickDashX(obj, branch) {
		handleFound(obj, branch)
		return
	}

	if c.foundByAllCommits(obj, branch) {
		handleFound(obj, branch)
		return
	}

	if c.foundByPickWithoutDashX(obj, branch) {
		handleFound(obj, branch)
		return
	}

	if c.foundByScript(obj, branch) {
		handleFound(obj, branch)
		return
	}

	return
}
<|endoftext|>