<commit_before>package mitm

import (
	"crypto/tls"
	"errors"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"sync"
	"time"
)

// Proxy is a forward proxy that substitutes its own certificate
// for incoming TLS connections in place of the upstream server's
// certificate.
type Proxy struct {
	// Wrap specifies a function for optionally wrapping upstream for
	// inspecting the decrypted HTTP request and response.
	Wrap func(upstream http.Handler) http.Handler

	// CA specifies the root CA for generating leaf certs for each incoming
	// TLS request.
	CA *tls.Certificate

	// TLSServerConfig specifies the tls.Config to use when generating leaf
	// cert using CA.
	TLSServerConfig *tls.Config

	// TLSServerConfig specifies the tls.Config to use when establishing a
	// downstream connection for proxying.
	TLSClientConfig *tls.Config

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

func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method == "CONNECT" {
		p.serveConnect(w, r)
		return
	}
	rp := &httputil.ReverseProxy{
		Director:      httpDirector,
		FlushInterval: p.FlushInterval,
	}
	p.Wrap(rp).ServeHTTP(w, r)
}

func (p *Proxy) serveConnect(w http.ResponseWriter, r *http.Request) {
	var (
		err   error
		sconn *tls.Conn
		name  = dnsName(r.Host)
	)

	if name == "" {
		log.Println("cannot determine cert name for " + r.Host)
		http.Error(w, "no upstream", 503)
		return
	}

	provisionalCert, err := p.cert(name)
	if err != nil {
		log.Println("cert", err)
		http.Error(w, "no upstream", 503)
		return
	}

	sConfig := new(tls.Config)
	if p.TLSServerConfig != nil {
		*sConfig = *p.TLSServerConfig
	}
	sConfig.Certificates = []tls.Certificate{*provisionalCert}
	sConfig.GetCertificate = func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
		cConfig := new(tls.Config)
		if p.TLSClientConfig != nil {
			*cConfig = *p.TLSClientConfig
		}
		cConfig.ServerName = hello.ServerName
		sconn, err = tls.Dial("tcp", r.Host, cConfig)
		if err != nil {
			log.Println("dial", r.Host, err)
			return nil, err
		}
		return p.cert(hello.ServerName)
	}

	cconn, err := handshake(w, sConfig)
	if err != nil {
		log.Println("handshake", r.Host, err)
		return
	}
	defer cconn.Close()
	if sconn == nil {
		log.Println("could not determine cert name for " + r.Host)
		return
	}
	defer sconn.Close()

	od := &oneShotDialer{c: sconn}
	rp := &httputil.ReverseProxy{
		Director:      httpsDirector,
		Transport:     &http.Transport{DialTLS: od.Dial},
		FlushInterval: p.FlushInterval,
	}

	ch := make(chan int)
	wc := &onCloseConn{cconn, func() { ch <- 0 }}
	http.Serve(&oneShotListener{wc}, p.Wrap(rp))
	<-ch
}

func (p *Proxy) cert(names ...string) (*tls.Certificate, error) {
	return genCert(p.CA, names)
}

var okHeader = []byte("HTTP/1.1 200 OK\r\n\r\n")

// handshake hijacks w's underlying net.Conn, responds to the CONNECT request
// and manually performs the TLS handshake. It returns the net.Conn or and
// error if any.
func handshake(w http.ResponseWriter, config *tls.Config) (net.Conn, error) {
	raw, _, err := w.(http.Hijacker).Hijack()
	if err != nil {
		http.Error(w, "no upstream", 503)
		return nil, err
	}
	if _, err = raw.Write(okHeader); err != nil {
		raw.Close()
		return nil, err
	}
	conn := tls.Server(raw, config)
	err = conn.Handshake()
	if err != nil {
		conn.Close()
		raw.Close()
		return nil, err
	}
	return conn, nil
}

func httpDirector(r *http.Request) {
	r.URL.Host = r.Host
	r.URL.Scheme = "http"
}

func httpsDirector(r *http.Request) {
	r.URL.Host = r.Host
	r.URL.Scheme = "https"
}

// dnsName returns the DNS name in addr, if any.
func dnsName(addr string) string {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return ""
	}
	return host
}

// namesOnCert returns the dns names
// in the peer's presented cert.
func namesOnCert(conn *tls.Conn) []string {
	// TODO(kr): handle IP addr SANs.
	c := conn.ConnectionState().PeerCertificates[0]
	if len(c.DNSNames) > 0 {
		// If Subject Alt Name is given,
		// we ignore the common name.
		// This matches behavior of crypto/x509.
		return c.DNSNames
	}
	return []string{c.Subject.CommonName}
}

// A oneShotDialer implements net.Dialer whos Accept only returns a
// net.Conn as specified by c followed by an error for each subsequent Dial.
type oneShotDialer struct {
	c  net.Conn
	mu sync.Mutex
}

func (d *oneShotDialer) Dial(network, addr string) (net.Conn, error) {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.c == nil {
		return nil, errors.New("closed")
	}
	c := d.c
	d.c = nil
	return c, nil
}

// A oneShotListener implements net.Listener whos Accept only returns a
// net.Conn as specified by c followed by an error for each subsequent Accept.
type oneShotListener struct {
	c net.Conn
}

func (l *oneShotListener) Accept() (net.Conn, error) {
	if l.c == nil {
		return nil, errors.New("closed")
	}
	c := l.c
	l.c = nil
	return c, nil
}

func (l *oneShotListener) Close() error {
	return nil
}

func (l *oneShotListener) Addr() net.Addr {
	return l.c.LocalAddr()
}

// A onCloseConn implements net.Conn and calls its f on Close.
type onCloseConn struct {
	net.Conn
	f func()
}

func (c *onCloseConn) Close() error {
	if c.f != nil {
		c.f()
		c.f = nil
	}
	return c.Conn.Close()
}
<commit_msg>mitm: fix doc<commit_after>package mitm

import (
	"crypto/tls"
	"errors"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"sync"
	"time"
)

// Proxy is a forward proxy that substitutes its own certificate
// for incoming TLS connections in place of the upstream server's
// certificate.
type Proxy struct {
	// Wrap specifies a function for optionally wrapping upstream for
	// inspecting the decrypted HTTP request and response.
	Wrap func(upstream http.Handler) http.Handler

	// CA specifies the root CA for generating leaf certs for each incoming
	// TLS request.
	CA *tls.Certificate

	// TLSServerConfig specifies the tls.Config to use when generating leaf
	// cert using CA.
	TLSServerConfig *tls.Config

	// TLSServerConfig specifies the tls.Config to use when establishing a
	// downstream connection for proxying.
	TLSClientConfig *tls.Config

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

func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method == "CONNECT" {
		p.serveConnect(w, r)
		return
	}
	rp := &httputil.ReverseProxy{
		Director:      httpDirector,
		FlushInterval: p.FlushInterval,
	}
	p.Wrap(rp).ServeHTTP(w, r)
}

func (p *Proxy) serveConnect(w http.ResponseWriter, r *http.Request) {
	var (
		err   error
		sconn *tls.Conn
		name  = dnsName(r.Host)
	)

	if name == "" {
		log.Println("cannot determine cert name for " + r.Host)
		http.Error(w, "no upstream", 503)
		return
	}

	provisionalCert, err := p.cert(name)
	if err != nil {
		log.Println("cert", err)
		http.Error(w, "no upstream", 503)
		return
	}

	sConfig := new(tls.Config)
	if p.TLSServerConfig != nil {
		*sConfig = *p.TLSServerConfig
	}
	sConfig.Certificates = []tls.Certificate{*provisionalCert}
	sConfig.GetCertificate = func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
		cConfig := new(tls.Config)
		if p.TLSClientConfig != nil {
			*cConfig = *p.TLSClientConfig
		}
		cConfig.ServerName = hello.ServerName
		sconn, err = tls.Dial("tcp", r.Host, cConfig)
		if err != nil {
			log.Println("dial", r.Host, err)
			return nil, err
		}
		return p.cert(hello.ServerName)
	}

	cconn, err := handshake(w, sConfig)
	if err != nil {
		log.Println("handshake", r.Host, err)
		return
	}
	defer cconn.Close()
	if sconn == nil {
		log.Println("could not determine cert name for " + r.Host)
		return
	}
	defer sconn.Close()

	od := &oneShotDialer{c: sconn}
	rp := &httputil.ReverseProxy{
		Director:      httpsDirector,
		Transport:     &http.Transport{DialTLS: od.Dial},
		FlushInterval: p.FlushInterval,
	}

	ch := make(chan int)
	wc := &onCloseConn{cconn, func() { ch <- 0 }}
	http.Serve(&oneShotListener{wc}, p.Wrap(rp))
	<-ch
}

func (p *Proxy) cert(names ...string) (*tls.Certificate, error) {
	return genCert(p.CA, names)
}

var okHeader = []byte("HTTP/1.1 200 OK\r\n\r\n")

// handshake hijacks w's underlying net.Conn, responds to the CONNECT request
// and manually performs the TLS handshake. It returns the net.Conn or and
// error if any.
func handshake(w http.ResponseWriter, config *tls.Config) (net.Conn, error) {
	raw, _, err := w.(http.Hijacker).Hijack()
	if err != nil {
		http.Error(w, "no upstream", 503)
		return nil, err
	}
	if _, err = raw.Write(okHeader); err != nil {
		raw.Close()
		return nil, err
	}
	conn := tls.Server(raw, config)
	err = conn.Handshake()
	if err != nil {
		conn.Close()
		raw.Close()
		return nil, err
	}
	return conn, nil
}

func httpDirector(r *http.Request) {
	r.URL.Host = r.Host
	r.URL.Scheme = "http"
}

func httpsDirector(r *http.Request) {
	r.URL.Host = r.Host
	r.URL.Scheme = "https"
}

// dnsName returns the DNS name in addr, if any.
func dnsName(addr string) string {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return ""
	}
	return host
}

// namesOnCert returns the dns names
// in the peer's presented cert.
func namesOnCert(conn *tls.Conn) []string {
	// TODO(kr): handle IP addr SANs.
	c := conn.ConnectionState().PeerCertificates[0]
	if len(c.DNSNames) > 0 {
		// If Subject Alt Name is given,
		// we ignore the common name.
		// This matches behavior of crypto/x509.
		return c.DNSNames
	}
	return []string{c.Subject.CommonName}
}

// A oneShotDialer implements net.Dialer whos Dial only returns a
// net.Conn as specified by c followed by an error for each subsequent Dial.
type oneShotDialer struct {
	c  net.Conn
	mu sync.Mutex
}

func (d *oneShotDialer) Dial(network, addr string) (net.Conn, error) {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.c == nil {
		return nil, errors.New("closed")
	}
	c := d.c
	d.c = nil
	return c, nil
}

// A oneShotListener implements net.Listener whos Accept only returns a
// net.Conn as specified by c followed by an error for each subsequent Accept.
type oneShotListener struct {
	c net.Conn
}

func (l *oneShotListener) Accept() (net.Conn, error) {
	if l.c == nil {
		return nil, errors.New("closed")
	}
	c := l.c
	l.c = nil
	return c, nil
}

func (l *oneShotListener) Close() error {
	return nil
}

func (l *oneShotListener) Addr() net.Addr {
	return l.c.LocalAddr()
}

// A onCloseConn implements net.Conn and calls its f on Close.
type onCloseConn struct {
	net.Conn
	f func()
}

func (c *onCloseConn) Close() error {
	if c.f != nil {
		c.f()
		c.f = nil
	}
	return c.Conn.Close()
}
<|endoftext|>


<commit_before>package models

import (
	"testing"
	"time"

	. "github.com/smartystreets/goconvey/convey"
)

func TestChannelparticipantNewChannelParticipant(t *testing.T) {
	Convey("while testing New Channel Participant", t, func() {
		Convey("Function call should return ChannelParticipant", func() {
			So(NewChannelParticipant(), ShouldNotBeNil)
		})
	})
}

func TestChannelparticipantGetId(t *testing.T) {
	Convey("While testing get id", t, func() {
		Convey("Initialized struct", func() {
			Convey("should return given id", func() {
				c := ChannelParticipant{Id: 45}
				So(c.GetId(), ShouldEqual, 45)
			})
		})
		Convey("Uninitialized struct", func() {
			Convey("should return 0", func() {
				So(NewChannelParticipant().GetId(), ShouldEqual, 0)
			})
			So(NewChannelParticipant, ShouldNotBeNil)
		})
	})
}

func TestChannelParticipantBeforeUpdate(t *testing.T) {
	Convey("While testing before update", t, func() {
		Convey("LastSeenAt should be updated", func() {
			lastSeenAt := time.Now().UTC()

			c := NewChannelParticipant()
			c.LastSeenAt = lastSeenAt

			// call before update
			err := c.BeforeUpdate()

			// make sure err is nil
			So(err, ShouldBeNil)
			// check preset last seen at is not same after calling
			// before update function
			So(c.LastSeenAt, ShouldNotEqual, lastSeenAt)
		})
	})
}
<commit_msg>social/test: test is added for channelparticipant<commit_after>package models

import (
	"testing"
	"time"

	. "github.com/smartystreets/goconvey/convey"
)

func TestChannelparticipantNewChannelParticipant(t *testing.T) {
	Convey("while testing New Channel Participant", t, func() {
		Convey("Function call should return ChannelParticipant", func() {
			So(NewChannelParticipant(), ShouldNotBeNil)
		})
	})
}

func TestChannelparticipantGetId(t *testing.T) {
	Convey("While testing get id", t, func() {
		Convey("Initialized struct", func() {
			Convey("should return given id", func() {
				c := ChannelParticipant{Id: 45}
				So(c.GetId(), ShouldEqual, 45)
			})
		})
		Convey("Uninitialized struct", func() {
			Convey("should return 0", func() {
				So(NewChannelParticipant().GetId(), ShouldEqual, 0)
			})
			So(NewChannelParticipant, ShouldNotBeNil)
		})
	})
}

func TestChannelParticipantTableName(t *testing.T) {
	Convey("While getting table name", t, func() {
		Convey("table names should match", func() {
			c := ChannelParticipant{}
			So(c.TableName(), ShouldEqual, "api.channel_participant")
		})
	})
}

func TestChannelParticipantBeforeUpdate(t *testing.T) {
	Convey("While testing before update", t, func() {
		Convey("LastSeenAt should be updated", func() {
			lastSeenAt := time.Now().UTC()

			c := NewChannelParticipant()
			c.LastSeenAt = lastSeenAt

			// call before update
			err := c.BeforeUpdate()

			// make sure err is nil
			So(err, ShouldBeNil)
			// check preset last seen at is not same after calling
			// before update function
			So(c.LastSeenAt, ShouldNotEqual, lastSeenAt)
		})
	})
}
<|endoftext|>


<commit_before>// Copyright (C) 2017 Google Inc.
//
// 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 service

import (
	"context"
	"fmt"

	"github.com/google/gapid/core/data/id"
	"github.com/google/gapid/core/data/protoutil"
	"github.com/google/gapid/core/image"
	"github.com/google/gapid/core/log"
	"github.com/google/gapid/core/os/device"
	"github.com/google/gapid/gapis/gfxapi"
	"github.com/google/gapid/gapis/memory"
	"github.com/google/gapid/gapis/service/box"
	"github.com/google/gapid/gapis/service/path"
	"github.com/google/gapid/gapis/stringtable"
)

type Service interface {
	// GetServerInfo returns information about the running server.
	GetServerInfo(ctx context.Context) (*ServerInfo, error)

	// GetAvailableStringTables returns list of available string table descriptions.
	GetAvailableStringTables(ctx context.Context) ([]*stringtable.Info, error)

	// GetStringTable returns the requested string table.
	GetStringTable(ctx context.Context, info *stringtable.Info) (*stringtable.StringTable, error)

	// ImportCapture imports capture data emitted by the graphics spy, returning
	// the new capture identifier.
	ImportCapture(ctx context.Context, name string, data []uint8) (*path.Capture, error)

	// ExportCapture returns a capture's data that can be consumed by
	// ImportCapture or LoadCapture.
	ExportCapture(ctx context.Context, c *path.Capture) ([]byte, error)

	// LoadCapture imports capture data from a local file, returning the new
	// capture identifier.
	LoadCapture(ctx context.Context, path string) (*path.Capture, error)

	// GetDevices returns the full list of replay devices avaliable to the server.
	// These include local replay devices and any connected Android devices.
	// This list may change over time, as devices are connected and disconnected.
	// If both connected Android and Local replay devices are found,
	// the local Android devices will be returned first.
	GetDevices(ctx context.Context) ([]*path.Device, error)

	// GetDevicesForReplay returns the list of replay devices avaliable to the
	// server that are capable of replaying the given capture.
	// These include local replay devices and any connected Android devices.
	// This list may change over time, as devices are connected and disconnected.
	// If both connected Android and Local replay devices are found,
	// the local Android devices will be returned first.
	GetDevicesForReplay(ctx context.Context, p *path.Capture) ([]*path.Device, error)

	// GetFramebufferAttachment returns the ImageInfo identifier describing the
	// given framebuffer attachment and device, immediately following the atom
	// after.
	// The provided RenderSettings structure can be used to adjust maximum desired
	// dimensions of the image, as well as applying debug visualizations.
	GetFramebufferAttachment(
		ctx context.Context,
		device *path.Device,
		after *path.Command,
		attachment gfxapi.FramebufferAttachment,
		settings *RenderSettings,
		hints *UsageHints) (*path.ImageInfo, error)

	// Get resolves and returns the object, value or memory at the path p.
	Get(ctx context.Context, p *path.Any) (interface{}, error)

	// Set creates a copy of the capture referenced by p, but with the object, value
	// or memory at p replaced with v. The path returned is identical to p, but with
	// the base changed to refer to the new capture.
	Set(ctx context.Context, p *path.Any, v interface{}) (*path.Any, error)

	// Follow returns the path to the object that the value at p links to.
	// If the value at p does not link to anything then nil is returned.
	Follow(ctx context.Context, p *path.Any) (*path.Any, error)

	// BeginCPUProfile starts CPU self-profiling of the server.
	// If the CPU is already being profiled then this function will return an
	// error.
	// This is a debug API, and may be removed in the future.
	BeginCPUProfile(ctx context.Context) error

	// EndCPUProfile ends the CPU profile, returning the pprof samples.
	// This is a debug API, and may be removed in the future.
	EndCPUProfile(ctx context.Context) ([]byte, error)

	// GetPerformanceCounters returns the values of all global counters as
	// a JSON blob.
	GetPerformanceCounters(ctx context.Context) ([]byte, error)

	// GetProfile returns the pprof profile with the given name.
	GetProfile(ctx context.Context, name string, debug int32) ([]byte, error)

	// GetLogStream calls the handler with each log record raised until the
	// context is cancelled.
	GetLogStream(context.Context, log.Handler) error

	// Find performs a search using req, streaming the results to h.
	Find(ctx context.Context, req *FindRequest, h FindHandler) error
}

// FindHandler is the handler of found items using Service.Find.
type FindHandler func(*FindResponse) error

// NewError attempts to box and return err into an Error.
// If err cannot be boxed into an Error then nil is returned.
func NewError(err error) *Error {
	switch err := err.(type) {
	case nil:
		return nil
	case *ErrDataUnavailable:
		return &Error{&Error_ErrDataUnavailable{err}}
	case *ErrInvalidPath:
		return &Error{&Error_ErrInvalidPath{err}}
	case *ErrInvalidArgument:
		return &Error{&Error_ErrInvalidArgument{err}}
	case *ErrPathNotFollowable:
		return &Error{&Error_ErrPathNotFollowable{err}}
	default:
		return &Error{&Error_ErrInternal{&ErrInternal{err.Error()}}}
	}
}

// Get returns the boxed error.
func (v *Error) Get() error { return protoutil.OneOf(v.Err).(error) }

// NewValue attempts to box and return v into a Value.
// If v cannot be boxed into a Value then nil is returned.
func NewValue(v interface{}) *Value {
	switch v := v.(type) {
	case nil:
		return &Value{}
	case *Capture:
		return &Value{&Value_Capture{v}}
	case *Contexts:
		return &Value{&Value_Contexts{v}}
	case *Command:
		return &Value{&Value_Command{v}}
	case *Commands:
		return &Value{&Value_Commands{v}}
	case *CommandTree:
		return &Value{&Value_CommandTree{v}}
	case *CommandTreeNode:
		return &Value{&Value_CommandTreeNode{v}}
	case *ConstantSet:
		return &Value{&Value_ConstantSet{v}}
	case *Event:
		return &Value{&Value_Event{v}}
	case *Events:
		return &Value{&Value_Events{v}}
	case *MemoryInfo:
		return &Value{&Value_MemoryInfo{v}}
	case *path.Any:
		return &Value{&Value_Path{v}}
	case path.Node:
		return &Value{&Value_Path{v.Path()}}
	case *Report:
		return &Value{&Value_Report{v}}
	case *Resources:
		return &Value{&Value_Resources{v}}
	case *StateTree:
		return &Value{&Value_StateTree{v}}
	case *StateTreeNode:
		return &Value{&Value_StateTreeNode{v}}
	case *gfxapi.Mesh:
		return &Value{&Value_Mesh{v}}
	case *gfxapi.Texture2D:
		return &Value{&Value_Texture_2D{v}}
	case *gfxapi.Cubemap:
		return &Value{&Value_Cubemap{v}}
	case *gfxapi.Shader:
		return &Value{&Value_Shader{v}}
	case *gfxapi.Program:
		return &Value{&Value_Program{v}}
	case *image.Info2D:
		return &Value{&Value_ImageInfo_2D{v}}
	case *device.Instance:
		return &Value{&Value_Device{v}}

	default:
		if v := box.NewValue(v); v != nil {
			return &Value{&Value_Box{v}}
		}
	}
	panic(fmt.Errorf("Cannot box value type %T", v))
}

// Get returns the boxed value.
func (v *Value) Get() interface{} {
	switch v := v.Val.(type) {
	case nil:
		return nil
	case *Value_Box:
		return v.Box.Get()
	default:
		return protoutil.OneOf(v)
	}
}

// NewMemoryRange constructs and returns a new MemoryRange from the
// memory.Range.
func NewMemoryRange(rng memory.Range) *MemoryRange {
	return &MemoryRange{Base: rng.Base, Size: rng.Size}
}

// NewMemoryRanges constructs and returns a new slice of MemoryRanges from the
// memory.RangeList.
func NewMemoryRanges(l memory.RangeList) []*MemoryRange {
	out := make([]*MemoryRange, len(l))
	for i, r := range l {
		out[i] = NewMemoryRange(r)
	}
	return out
}

// Link returns the link to the atom pointed by a report item.
// If nil, nil is returned then the path cannot be followed.
func (i ReportItem) Link(ctx context.Context, p path.Node) (path.Node, error) {
	if capture := path.FindCapture(p); capture != nil {
		return capture.Command(i.Command), nil
	}
	return nil, nil
}

// Find looks up a resource by type and identifier.
func (r *Resources) Find(ty gfxapi.ResourceType, id id.ID) *Resource {
	for _, t := range r.Types {
		if t.Type == ty {
			for _, r := range t.Resources {
				if r.Path.Id.ID() == id {
					return r
				}
			}
			break
		}
	}
	return nil
}
<commit_msg>gapis/service: Add case for *Context in NewValue()<commit_after>// Copyright (C) 2017 Google Inc.
//
// 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 service

import (
	"context"
	"fmt"

	"github.com/google/gapid/core/data/id"
	"github.com/google/gapid/core/data/protoutil"
	"github.com/google/gapid/core/image"
	"github.com/google/gapid/core/log"
	"github.com/google/gapid/core/os/device"
	"github.com/google/gapid/gapis/gfxapi"
	"github.com/google/gapid/gapis/memory"
	"github.com/google/gapid/gapis/service/box"
	"github.com/google/gapid/gapis/service/path"
	"github.com/google/gapid/gapis/stringtable"
)

type Service interface {
	// GetServerInfo returns information about the running server.
	GetServerInfo(ctx context.Context) (*ServerInfo, error)

	// GetAvailableStringTables returns list of available string table descriptions.
	GetAvailableStringTables(ctx context.Context) ([]*stringtable.Info, error)

	// GetStringTable returns the requested string table.
	GetStringTable(ctx context.Context, info *stringtable.Info) (*stringtable.StringTable, error)

	// ImportCapture imports capture data emitted by the graphics spy, returning
	// the new capture identifier.
	ImportCapture(ctx context.Context, name string, data []uint8) (*path.Capture, error)

	// ExportCapture returns a capture's data that can be consumed by
	// ImportCapture or LoadCapture.
	ExportCapture(ctx context.Context, c *path.Capture) ([]byte, error)

	// LoadCapture imports capture data from a local file, returning the new
	// capture identifier.
	LoadCapture(ctx context.Context, path string) (*path.Capture, error)

	// GetDevices returns the full list of replay devices avaliable to the server.
	// These include local replay devices and any connected Android devices.
	// This list may change over time, as devices are connected and disconnected.
	// If both connected Android and Local replay devices are found,
	// the local Android devices will be returned first.
	GetDevices(ctx context.Context) ([]*path.Device, error)

	// GetDevicesForReplay returns the list of replay devices avaliable to the
	// server that are capable of replaying the given capture.
	// These include local replay devices and any connected Android devices.
	// This list may change over time, as devices are connected and disconnected.
	// If both connected Android and Local replay devices are found,
	// the local Android devices will be returned first.
	GetDevicesForReplay(ctx context.Context, p *path.Capture) ([]*path.Device, error)

	// GetFramebufferAttachment returns the ImageInfo identifier describing the
	// given framebuffer attachment and device, immediately following the atom
	// after.
	// The provided RenderSettings structure can be used to adjust maximum desired
	// dimensions of the image, as well as applying debug visualizations.
	GetFramebufferAttachment(
		ctx context.Context,
		device *path.Device,
		after *path.Command,
		attachment gfxapi.FramebufferAttachment,
		settings *RenderSettings,
		hints *UsageHints) (*path.ImageInfo, error)

	// Get resolves and returns the object, value or memory at the path p.
	Get(ctx context.Context, p *path.Any) (interface{}, error)

	// Set creates a copy of the capture referenced by p, but with the object, value
	// or memory at p replaced with v. The path returned is identical to p, but with
	// the base changed to refer to the new capture.
	Set(ctx context.Context, p *path.Any, v interface{}) (*path.Any, error)

	// Follow returns the path to the object that the value at p links to.
	// If the value at p does not link to anything then nil is returned.
	Follow(ctx context.Context, p *path.Any) (*path.Any, error)

	// BeginCPUProfile starts CPU self-profiling of the server.
	// If the CPU is already being profiled then this function will return an
	// error.
	// This is a debug API, and may be removed in the future.
	BeginCPUProfile(ctx context.Context) error

	// EndCPUProfile ends the CPU profile, returning the pprof samples.
	// This is a debug API, and may be removed in the future.
	EndCPUProfile(ctx context.Context) ([]byte, error)

	// GetPerformanceCounters returns the values of all global counters as
	// a JSON blob.
	GetPerformanceCounters(ctx context.Context) ([]byte, error)

	// GetProfile returns the pprof profile with the given name.
	GetProfile(ctx context.Context, name string, debug int32) ([]byte, error)

	// GetLogStream calls the handler with each log record raised until the
	// context is cancelled.
	GetLogStream(context.Context, log.Handler) error

	// Find performs a search using req, streaming the results to h.
	Find(ctx context.Context, req *FindRequest, h FindHandler) error
}

// FindHandler is the handler of found items using Service.Find.
type FindHandler func(*FindResponse) error

// NewError attempts to box and return err into an Error.
// If err cannot be boxed into an Error then nil is returned.
func NewError(err error) *Error {
	switch err := err.(type) {
	case nil:
		return nil
	case *ErrDataUnavailable:
		return &Error{&Error_ErrDataUnavailable{err}}
	case *ErrInvalidPath:
		return &Error{&Error_ErrInvalidPath{err}}
	case *ErrInvalidArgument:
		return &Error{&Error_ErrInvalidArgument{err}}
	case *ErrPathNotFollowable:
		return &Error{&Error_ErrPathNotFollowable{err}}
	default:
		return &Error{&Error_ErrInternal{&ErrInternal{err.Error()}}}
	}
}

// Get returns the boxed error.
func (v *Error) Get() error { return protoutil.OneOf(v.Err).(error) }

// NewValue attempts to box and return v into a Value.
// If v cannot be boxed into a Value then nil is returned.
func NewValue(v interface{}) *Value {
	switch v := v.(type) {
	case nil:
		return &Value{}
	case *Capture:
		return &Value{&Value_Capture{v}}
	case *Context:
		return &Value{&Value_Context{v}}
	case *Contexts:
		return &Value{&Value_Contexts{v}}
	case *Command:
		return &Value{&Value_Command{v}}
	case *Commands:
		return &Value{&Value_Commands{v}}
	case *CommandTree:
		return &Value{&Value_CommandTree{v}}
	case *CommandTreeNode:
		return &Value{&Value_CommandTreeNode{v}}
	case *ConstantSet:
		return &Value{&Value_ConstantSet{v}}
	case *Event:
		return &Value{&Value_Event{v}}
	case *Events:
		return &Value{&Value_Events{v}}
	case *MemoryInfo:
		return &Value{&Value_MemoryInfo{v}}
	case *path.Any:
		return &Value{&Value_Path{v}}
	case path.Node:
		return &Value{&Value_Path{v.Path()}}
	case *Report:
		return &Value{&Value_Report{v}}
	case *Resources:
		return &Value{&Value_Resources{v}}
	case *StateTree:
		return &Value{&Value_StateTree{v}}
	case *StateTreeNode:
		return &Value{&Value_StateTreeNode{v}}
	case *gfxapi.Mesh:
		return &Value{&Value_Mesh{v}}
	case *gfxapi.Texture2D:
		return &Value{&Value_Texture_2D{v}}
	case *gfxapi.Cubemap:
		return &Value{&Value_Cubemap{v}}
	case *gfxapi.Shader:
		return &Value{&Value_Shader{v}}
	case *gfxapi.Program:
		return &Value{&Value_Program{v}}
	case *image.Info2D:
		return &Value{&Value_ImageInfo_2D{v}}
	case *device.Instance:
		return &Value{&Value_Device{v}}

	default:
		if v := box.NewValue(v); v != nil {
			return &Value{&Value_Box{v}}
		}
	}
	panic(fmt.Errorf("Cannot box value type %T", v))
}

// Get returns the boxed value.
func (v *Value) Get() interface{} {
	switch v := v.Val.(type) {
	case nil:
		return nil
	case *Value_Box:
		return v.Box.Get()
	default:
		return protoutil.OneOf(v)
	}
}

// NewMemoryRange constructs and returns a new MemoryRange from the
// memory.Range.
func NewMemoryRange(rng memory.Range) *MemoryRange {
	return &MemoryRange{Base: rng.Base, Size: rng.Size}
}

// NewMemoryRanges constructs and returns a new slice of MemoryRanges from the
// memory.RangeList.
func NewMemoryRanges(l memory.RangeList) []*MemoryRange {
	out := make([]*MemoryRange, len(l))
	for i, r := range l {
		out[i] = NewMemoryRange(r)
	}
	return out
}

// Link returns the link to the atom pointed by a report item.
// If nil, nil is returned then the path cannot be followed.
func (i ReportItem) Link(ctx context.Context, p path.Node) (path.Node, error) {
	if capture := path.FindCapture(p); capture != nil {
		return capture.Command(i.Command), nil
	}
	return nil, nil
}

// Find looks up a resource by type and identifier.
func (r *Resources) Find(ty gfxapi.ResourceType, id id.ID) *Resource {
	for _, t := range r.Types {
		if t.Type == ty {
			for _, r := range t.Resources {
				if r.Path.Id.ID() == id {
					return r
				}
			}
			break
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package gorethink

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

	p "gopkg.in/dancannon/gorethink.v2/ql2"
)

// Mocking is based on the amazing package github.com/stretchr/testify

// testingT is an interface wrapper around *testing.T
type testingT interface {
	Logf(format string, args ...interface{})
	Errorf(format string, args ...interface{})
	FailNow()
}

// MockAnything can be used in place of any term, this is useful when you want
// mock similar queries or queries that you don't quite know the exact structure
// of.
func MockAnything() Term {
	t := constructRootTerm("MockAnything", p.Term_DATUM, nil, nil)
	t.isMockAnything = true

	return t
}

func (t Term) MockAnything() Term {
	t = constructMethodTerm(t, "MockAnything", p.Term_DATUM, nil, nil)
	t.isMockAnything = true

	return t
}

// MockQuery represents a mocked query and is used for setting expectations,
// as well as recording activity.
type MockQuery struct {
	parent *Mock

	// Holds the query and term
	Query Query

	// Holds the JSON representation of query
	BuiltQuery []byte

	// Holds the response that should be returned when this method is executed.
	Response interface{}

	// Holds the error that should be returned when this method is executed.
	Error error

	// The number of times to return the return arguments when setting
	// expectations. 0 means to always return the value.
	Repeatability int

	// Holds a channel that will be used to block the Return until it either
	// recieves a message or is closed. nil means it returns immediately.
	WaitFor <-chan time.Time

	// Amount of times this query has been executed
	executed int
}

func newMockQuery(parent *Mock, q Query) *MockQuery {
	// Build and marshal term
	builtQuery, err := json.Marshal(q.Build())
	if err != nil {
		panic(fmt.Sprintf("Failed to build query: %s", err))
	}

	return &MockQuery{
		parent:        parent,
		Query:         q,
		BuiltQuery:    builtQuery,
		Response:      make([]interface{}, 0),
		Repeatability: 0,
		WaitFor:       nil,
	}
}

func newMockQueryFromTerm(parent *Mock, t Term, opts map[string]interface{}) *MockQuery {
	q, err := parent.newQuery(t, opts)
	if err != nil {
		panic(fmt.Sprintf("Failed to build query: %s", err))
	}

	return newMockQuery(parent, q)
}

func (mq *MockQuery) lock() {
	mq.parent.mu.Lock()
}

func (mq *MockQuery) unlock() {
	mq.parent.mu.Unlock()
}

// Return specifies the return arguments for the expectation.
//
//    mock.On(r.Table("test")).Return(nil, errors.New("failed"))
func (mq *MockQuery) Return(response interface{}, err error) *MockQuery {
	mq.lock()
	defer mq.unlock()

	mq.Response = response
	mq.Error = err

	return mq
}

// Once indicates that that the mock should only return the value once.
//
//    mock.On(r.Table("test")).Return(result, nil).Once()
func (mq *MockQuery) Once() *MockQuery {
	return mq.Times(1)
}

// Twice indicates that that the mock should only return the value twice.
//
//    mock.On(r.Table("test")).Return(result, nil).Twice()
func (mq *MockQuery) Twice() *MockQuery {
	return mq.Times(2)
}

// Times indicates that that the mock should only return the indicated number
// of times.
//
//    mock.On(r.Table("test")).Return(result, nil).Times(5)
func (mq *MockQuery) Times(i int) *MockQuery {
	mq.lock()
	defer mq.unlock()
	mq.Repeatability = i
	return mq
}

// WaitUntil sets the channel that will block the mock's return until its closed
// or a message is received.
//
//    mock.On(r.Table("test")).WaitUntil(time.After(time.Second))
func (mq *MockQuery) WaitUntil(w <-chan time.Time) *MockQuery {
	mq.lock()
	defer mq.unlock()
	mq.WaitFor = w
	return mq
}

// After sets how long to block until the query returns
//
//    mock.On(r.Table("test")).After(time.Second)
func (mq *MockQuery) After(d time.Duration) *MockQuery {
	return mq.WaitUntil(time.After(d))
}

// On chains a new expectation description onto the mocked interface. This
// allows syntax like.
//
//    Mock.
//       On(r.Table("test")).Return(result, nil).
//       On(r.Table("test2")).Return(nil, errors.New("Some Error"))
func (mq *MockQuery) On(t Term) *MockQuery {
	return mq.parent.On(t)
}

// Mock is used to mock query execution and verify that the expected queries are
// being executed. Mocks are used by creating an instance using NewMock and then
// passing this when running your queries instead of a session. For example:
//
//     mock := r.NewMock()
//     mock.on(r.Table("test")).Return([]interface{}{data}, nil)
//
//     cursor, err := r.Table("test").Run(mock)
//
//     mock.AssertExpectations(t)
type Mock struct {
	mu   sync.Mutex
	opts ConnectOpts

	ExpectedQueries []*MockQuery
	Queries         []MockQuery
}

// NewMock creates an instance of Mock, you can optionally pass ConnectOpts to
// the function, if passed any mocked query will be generated using those
// options.
func NewMock(opts ...ConnectOpts) *Mock {
	m := &Mock{
		ExpectedQueries: make([]*MockQuery, 0),
		Queries:         make([]MockQuery, 0),
	}

	if len(opts) > 0 {
		m.opts = opts[0]
	}

	return m
}

// On starts a description of an expectation of the specified query
// being executed.
//
//     mock.On(r.Table("test"))
func (m *Mock) On(t Term, opts ...map[string]interface{}) *MockQuery {
	var qopts map[string]interface{}
	if len(opts) > 0 {
		qopts = opts[0]
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	mq := newMockQueryFromTerm(m, t, qopts)
	m.ExpectedQueries = append(m.ExpectedQueries, mq)
	return mq
}

// AssertExpectations asserts that everything specified with On and Return was
// in fact executed as expected. Queries may have been executed in any order.
func (m *Mock) AssertExpectations(t testingT) bool {
	var somethingMissing bool
	var failedExpectations int

	// iterate through each expectation
	expectedQueries := m.expectedQueries()
	for _, expectedQuery := range expectedQueries {
		if !m.queryWasExecuted(expectedQuery) && expectedQuery.executed == 0 {
			somethingMissing = true
			failedExpectations++
			t.Logf("❌\t%s", expectedQuery.Query.Term.String())
		} else {
			m.mu.Lock()
			if expectedQuery.Repeatability > 0 {
				somethingMissing = true
				failedExpectations++
			} else {
				t.Logf("✅\t%s", expectedQuery.Query.Term.String())
			}
			m.mu.Unlock()
		}
	}

	if somethingMissing {
		t.Errorf("FAIL: %d out of %d expectation(s) were met.\n\tThe query you are testing needs to be executed %d more times(s).", len(expectedQueries)-failedExpectations, len(expectedQueries), failedExpectations)
	}

	return !somethingMissing
}

// AssertNumberOfExecutions asserts that the query was executed expectedExecutions times.
func (m *Mock) AssertNumberOfExecutions(t testingT, expectedQuery *MockQuery, expectedExecutions int) bool {
	var actualExecutions int
	for _, query := range m.queries() {
		if query.Query.Term.compare(*expectedQuery.Query.Term, map[int64]int64{}) && query.Repeatability > -1 {
			// if bytes.Equal(query.BuiltQuery, expectedQuery.BuiltQuery) {
			actualExecutions++
		}
	}

	if expectedExecutions != actualExecutions {
		t.Errorf("Expected number of executions (%d) does not match the actual number of executions (%d).", expectedExecutions, actualExecutions)
		return false
	}

	return true
}

// AssertExecuted asserts that the method was executed.
// It can produce a false result when an argument is a pointer type and the underlying value changed after executing the mocked method.
func (m *Mock) AssertExecuted(t testingT, expectedQuery *MockQuery) bool {
	if !m.queryWasExecuted(expectedQuery) {
		t.Errorf("The query \"%s\" should have been executed, but was not.", expectedQuery.Query.Term.String())
		return false
	}
	return true
}

// AssertNotExecuted asserts that the method was not executed.
// It can produce a false result when an argument is a pointer type and the underlying value changed after executing the mocked method.
func (m *Mock) AssertNotExecuted(t testingT, expectedQuery *MockQuery) bool {
	if m.queryWasExecuted(expectedQuery) {
		t.Errorf("The query \"%s\" was executed, but should NOT have been.", expectedQuery.Query.Term.String())
		return false
	}
	return true
}

func (m *Mock) IsConnected() bool {
	return true
}

func (m *Mock) Query(q Query) (*Cursor, error) {
	found, query := m.findExpectedQuery(q)

	if found < 0 {
		panic(fmt.Sprintf("gorethink: mock: This query was unexpected:\n\t\t%s", q.Term.String()))
	} else {
		m.mu.Lock()
		switch {
		case query.Repeatability == 1:
			query.Repeatability = -1
			query.executed++

		case query.Repeatability > 1:
			query.Repeatability--
			query.executed++

		case query.Repeatability == 0:
			query.executed++
		}
		m.mu.Unlock()
	}

	// add the query
	m.mu.Lock()
	m.Queries = append(m.Queries, *newMockQuery(m, q))
	m.mu.Unlock()

	// block if specified
	if query.WaitFor != nil {
		<-query.WaitFor
	}

	// Return error without building cursor if non-nil
	if query.Error != nil {
		return nil, query.Error
	}

	// Build cursor and return
	c := newCursor(nil, "", query.Query.Token, query.Query.Term, query.Query.Opts)
	c.finished = true
	c.fetching = false
	c.isAtom = true

	responseVal := reflect.ValueOf(query.Response)
	if responseVal.Kind() == reflect.Slice || responseVal.Kind() == reflect.Array {
		for i := 0; i < responseVal.Len(); i++ {
			c.buffer = append(c.buffer, responseVal.Index(i).Interface())
		}
	} else {
		c.buffer = append(c.buffer, query.Response)
	}

	return c, nil
}

func (m *Mock) Exec(q Query) error {
	_, err := m.Query(q)

	return err
}

func (m *Mock) newQuery(t Term, opts map[string]interface{}) (Query, error) {
	return newQuery(t, opts, &m.opts)
}

func (m *Mock) findExpectedQuery(q Query) (int, *MockQuery) {
	m.mu.Lock()
	defer m.mu.Unlock()

	for i, query := range m.ExpectedQueries {
		// if bytes.Equal(query.BuiltQuery, builtQuery) && query.Repeatability > -1 {
		if query.Query.Term.compare(*q.Term, map[int64]int64{}) && query.Repeatability > -1 {
			return i, query
		}
	}

	return -1, nil
}

func (m *Mock) queryWasExecuted(expectedQuery *MockQuery) bool {
	for _, query := range m.queries() {
		if query.Query.Term.compare(*expectedQuery.Query.Term, map[int64]int64{}) {
			// if bytes.Equal(query.BuiltQuery, expectedQuery.BuiltQuery) {
			return true
		}
	}

	// we didn't find the expected query
	return false
}

func (m *Mock) expectedQueries() []*MockQuery {
	m.mu.Lock()
	defer m.mu.Unlock()
	return append([]*MockQuery{}, m.ExpectedQueries...)
}

func (m *Mock) queries() []MockQuery {
	m.mu.Lock()
	defer m.mu.Unlock()
	return append([]MockQuery{}, m.Queries...)
}
<commit_msg>mock: Correct typo in example<commit_after>package gorethink

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

	p "gopkg.in/dancannon/gorethink.v2/ql2"
)

// Mocking is based on the amazing package github.com/stretchr/testify

// testingT is an interface wrapper around *testing.T
type testingT interface {
	Logf(format string, args ...interface{})
	Errorf(format string, args ...interface{})
	FailNow()
}

// MockAnything can be used in place of any term, this is useful when you want
// mock similar queries or queries that you don't quite know the exact structure
// of.
func MockAnything() Term {
	t := constructRootTerm("MockAnything", p.Term_DATUM, nil, nil)
	t.isMockAnything = true

	return t
}

func (t Term) MockAnything() Term {
	t = constructMethodTerm(t, "MockAnything", p.Term_DATUM, nil, nil)
	t.isMockAnything = true

	return t
}

// MockQuery represents a mocked query and is used for setting expectations,
// as well as recording activity.
type MockQuery struct {
	parent *Mock

	// Holds the query and term
	Query Query

	// Holds the JSON representation of query
	BuiltQuery []byte

	// Holds the response that should be returned when this method is executed.
	Response interface{}

	// Holds the error that should be returned when this method is executed.
	Error error

	// The number of times to return the return arguments when setting
	// expectations. 0 means to always return the value.
	Repeatability int

	// Holds a channel that will be used to block the Return until it either
	// recieves a message or is closed. nil means it returns immediately.
	WaitFor <-chan time.Time

	// Amount of times this query has been executed
	executed int
}

func newMockQuery(parent *Mock, q Query) *MockQuery {
	// Build and marshal term
	builtQuery, err := json.Marshal(q.Build())
	if err != nil {
		panic(fmt.Sprintf("Failed to build query: %s", err))
	}

	return &MockQuery{
		parent:        parent,
		Query:         q,
		BuiltQuery:    builtQuery,
		Response:      make([]interface{}, 0),
		Repeatability: 0,
		WaitFor:       nil,
	}
}

func newMockQueryFromTerm(parent *Mock, t Term, opts map[string]interface{}) *MockQuery {
	q, err := parent.newQuery(t, opts)
	if err != nil {
		panic(fmt.Sprintf("Failed to build query: %s", err))
	}

	return newMockQuery(parent, q)
}

func (mq *MockQuery) lock() {
	mq.parent.mu.Lock()
}

func (mq *MockQuery) unlock() {
	mq.parent.mu.Unlock()
}

// Return specifies the return arguments for the expectation.
//
//    mock.On(r.Table("test")).Return(nil, errors.New("failed"))
func (mq *MockQuery) Return(response interface{}, err error) *MockQuery {
	mq.lock()
	defer mq.unlock()

	mq.Response = response
	mq.Error = err

	return mq
}

// Once indicates that that the mock should only return the value once.
//
//    mock.On(r.Table("test")).Return(result, nil).Once()
func (mq *MockQuery) Once() *MockQuery {
	return mq.Times(1)
}

// Twice indicates that that the mock should only return the value twice.
//
//    mock.On(r.Table("test")).Return(result, nil).Twice()
func (mq *MockQuery) Twice() *MockQuery {
	return mq.Times(2)
}

// Times indicates that that the mock should only return the indicated number
// of times.
//
//    mock.On(r.Table("test")).Return(result, nil).Times(5)
func (mq *MockQuery) Times(i int) *MockQuery {
	mq.lock()
	defer mq.unlock()
	mq.Repeatability = i
	return mq
}

// WaitUntil sets the channel that will block the mock's return until its closed
// or a message is received.
//
//    mock.On(r.Table("test")).WaitUntil(time.After(time.Second))
func (mq *MockQuery) WaitUntil(w <-chan time.Time) *MockQuery {
	mq.lock()
	defer mq.unlock()
	mq.WaitFor = w
	return mq
}

// After sets how long to block until the query returns
//
//    mock.On(r.Table("test")).After(time.Second)
func (mq *MockQuery) After(d time.Duration) *MockQuery {
	return mq.WaitUntil(time.After(d))
}

// On chains a new expectation description onto the mocked interface. This
// allows syntax like.
//
//    Mock.
//       On(r.Table("test")).Return(result, nil).
//       On(r.Table("test2")).Return(nil, errors.New("Some Error"))
func (mq *MockQuery) On(t Term) *MockQuery {
	return mq.parent.On(t)
}

// Mock is used to mock query execution and verify that the expected queries are
// being executed. Mocks are used by creating an instance using NewMock and then
// passing this when running your queries instead of a session. For example:
//
//     mock := r.NewMock()
//     mock.On(r.Table("test")).Return([]interface{}{data}, nil)
//
//     cursor, err := r.Table("test").Run(mock)
//
//     mock.AssertExpectations(t)
type Mock struct {
	mu   sync.Mutex
	opts ConnectOpts

	ExpectedQueries []*MockQuery
	Queries         []MockQuery
}

// NewMock creates an instance of Mock, you can optionally pass ConnectOpts to
// the function, if passed any mocked query will be generated using those
// options.
func NewMock(opts ...ConnectOpts) *Mock {
	m := &Mock{
		ExpectedQueries: make([]*MockQuery, 0),
		Queries:         make([]MockQuery, 0),
	}

	if len(opts) > 0 {
		m.opts = opts[0]
	}

	return m
}

// On starts a description of an expectation of the specified query
// being executed.
//
//     mock.On(r.Table("test"))
func (m *Mock) On(t Term, opts ...map[string]interface{}) *MockQuery {
	var qopts map[string]interface{}
	if len(opts) > 0 {
		qopts = opts[0]
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	mq := newMockQueryFromTerm(m, t, qopts)
	m.ExpectedQueries = append(m.ExpectedQueries, mq)
	return mq
}

// AssertExpectations asserts that everything specified with On and Return was
// in fact executed as expected. Queries may have been executed in any order.
func (m *Mock) AssertExpectations(t testingT) bool {
	var somethingMissing bool
	var failedExpectations int

	// iterate through each expectation
	expectedQueries := m.expectedQueries()
	for _, expectedQuery := range expectedQueries {
		if !m.queryWasExecuted(expectedQuery) && expectedQuery.executed == 0 {
			somethingMissing = true
			failedExpectations++
			t.Logf("❌\t%s", expectedQuery.Query.Term.String())
		} else {
			m.mu.Lock()
			if expectedQuery.Repeatability > 0 {
				somethingMissing = true
				failedExpectations++
			} else {
				t.Logf("✅\t%s", expectedQuery.Query.Term.String())
			}
			m.mu.Unlock()
		}
	}

	if somethingMissing {
		t.Errorf("FAIL: %d out of %d expectation(s) were met.\n\tThe query you are testing needs to be executed %d more times(s).", len(expectedQueries)-failedExpectations, len(expectedQueries), failedExpectations)
	}

	return !somethingMissing
}

// AssertNumberOfExecutions asserts that the query was executed expectedExecutions times.
func (m *Mock) AssertNumberOfExecutions(t testingT, expectedQuery *MockQuery, expectedExecutions int) bool {
	var actualExecutions int
	for _, query := range m.queries() {
		if query.Query.Term.compare(*expectedQuery.Query.Term, map[int64]int64{}) && query.Repeatability > -1 {
			// if bytes.Equal(query.BuiltQuery, expectedQuery.BuiltQuery) {
			actualExecutions++
		}
	}

	if expectedExecutions != actualExecutions {
		t.Errorf("Expected number of executions (%d) does not match the actual number of executions (%d).", expectedExecutions, actualExecutions)
		return false
	}

	return true
}

// AssertExecuted asserts that the method was executed.
// It can produce a false result when an argument is a pointer type and the underlying value changed after executing the mocked method.
func (m *Mock) AssertExecuted(t testingT, expectedQuery *MockQuery) bool {
	if !m.queryWasExecuted(expectedQuery) {
		t.Errorf("The query \"%s\" should have been executed, but was not.", expectedQuery.Query.Term.String())
		return false
	}
	return true
}

// AssertNotExecuted asserts that the method was not executed.
// It can produce a false result when an argument is a pointer type and the underlying value changed after executing the mocked method.
func (m *Mock) AssertNotExecuted(t testingT, expectedQuery *MockQuery) bool {
	if m.queryWasExecuted(expectedQuery) {
		t.Errorf("The query \"%s\" was executed, but should NOT have been.", expectedQuery.Query.Term.String())
		return false
	}
	return true
}

func (m *Mock) IsConnected() bool {
	return true
}

func (m *Mock) Query(q Query) (*Cursor, error) {
	found, query := m.findExpectedQuery(q)

	if found < 0 {
		panic(fmt.Sprintf("gorethink: mock: This query was unexpected:\n\t\t%s", q.Term.String()))
	} else {
		m.mu.Lock()
		switch {
		case query.Repeatability == 1:
			query.Repeatability = -1
			query.executed++

		case query.Repeatability > 1:
			query.Repeatability--
			query.executed++

		case query.Repeatability == 0:
			query.executed++
		}
		m.mu.Unlock()
	}

	// add the query
	m.mu.Lock()
	m.Queries = append(m.Queries, *newMockQuery(m, q))
	m.mu.Unlock()

	// block if specified
	if query.WaitFor != nil {
		<-query.WaitFor
	}

	// Return error without building cursor if non-nil
	if query.Error != nil {
		return nil, query.Error
	}

	// Build cursor and return
	c := newCursor(nil, "", query.Query.Token, query.Query.Term, query.Query.Opts)
	c.finished = true
	c.fetching = false
	c.isAtom = true

	responseVal := reflect.ValueOf(query.Response)
	if responseVal.Kind() == reflect.Slice || responseVal.Kind() == reflect.Array {
		for i := 0; i < responseVal.Len(); i++ {
			c.buffer = append(c.buffer, responseVal.Index(i).Interface())
		}
	} else {
		c.buffer = append(c.buffer, query.Response)
	}

	return c, nil
}

func (m *Mock) Exec(q Query) error {
	_, err := m.Query(q)

	return err
}

func (m *Mock) newQuery(t Term, opts map[string]interface{}) (Query, error) {
	return newQuery(t, opts, &m.opts)
}

func (m *Mock) findExpectedQuery(q Query) (int, *MockQuery) {
	m.mu.Lock()
	defer m.mu.Unlock()

	for i, query := range m.ExpectedQueries {
		// if bytes.Equal(query.BuiltQuery, builtQuery) && query.Repeatability > -1 {
		if query.Query.Term.compare(*q.Term, map[int64]int64{}) && query.Repeatability > -1 {
			return i, query
		}
	}

	return -1, nil
}

func (m *Mock) queryWasExecuted(expectedQuery *MockQuery) bool {
	for _, query := range m.queries() {
		if query.Query.Term.compare(*expectedQuery.Query.Term, map[int64]int64{}) {
			// if bytes.Equal(query.BuiltQuery, expectedQuery.BuiltQuery) {
			return true
		}
	}

	// we didn't find the expected query
	return false
}

func (m *Mock) expectedQueries() []*MockQuery {
	m.mu.Lock()
	defer m.mu.Unlock()
	return append([]*MockQuery{}, m.ExpectedQueries...)
}

func (m *Mock) queries() []MockQuery {
	m.mu.Lock()
	defer m.mu.Unlock()
	return append([]MockQuery{}, m.Queries...)
}
<|endoftext|>


<commit_before>package main

import (
	"code.google.com/p/go.crypto/ssh"
	"code.google.com/p/goauth2/oauth"
	"crypto"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"github.com/google/go-github/github"
	"github.com/gorilla/mux"
	"github.com/kylelemons/go-gypsy/yaml"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/user"
	"path"
	"reflect"
	"strings"
	"sync"
)

var (
	port       = "8080"
	sshPort    = "22"
	configFile = "config.yml"
)

type Host struct {
	URI          string
	LatestCommit string
}

type Environment struct {
	Name               string
	RepoPath           string
	Hosts              []Host
	Branch             string
	LatestGitHubCommit string
}

type Project struct {
	Name         string
	GitHubURL    string
	RepoName     string
	RepoOwner    string
	Environments []Environment
}

/*
Thanks to Russ Cox: https://groups.google.com/d/msg/golang-nuts/OEdSDgEC7js/iyhU9DW_IKcJ
eq reports whether the first argument is equal to
any of the remaining arguments. (for use in templates)
*/
func eq(args ...interface{}) bool {
	if len(args) == 0 {
		return false
	}
	x := args[0]
	switch x := x.(type) {
	case string, int, int64, byte, float32, float64:
		for _, y := range args[1:] {
			if x == y {
				return true
			}
		}
		return false
	}

	for _, y := range args[1:] {
		if reflect.DeepEqual(x, y) {
			return true
		}
	}
	return false
}

func (h *Host) ShortCommitHash() string {
	return h.LatestCommit[:7]
}

func getPrivateKey(filename string) []byte {
	content, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Panic("Failed to open private key file: " + err.Error())
	}
	return content
}

type keychain struct {
	key *rsa.PrivateKey
}

func (k *keychain) Key(i int) (interface{}, error) {
	if i != 0 {
		return nil, nil
	}
	return &k.key.PublicKey, nil
}

func (k *keychain) Sign(i int, rand io.Reader, data []byte) (sig []byte, err error) {
	hashFunc := crypto.SHA1
	h := hashFunc.New()
	h.Write(data)
	digest := h.Sum(nil)
	return rsa.SignPKCS1v15(rand, k.key, hashFunc, digest)
}

//  Get the most recent commit hash on a given branch from GitHub
func latestGitHubCommit(c *github.Client, repoOwner, repoName, branchName string) string {
	opts := &github.CommitsListOptions{SHA: branchName}
	commits, _, err := c.Repositories.ListCommits(repoOwner, repoName, opts)
	if err != nil {
		log.Panic(err)
	}

	return *commits[0].SHA
}

func remoteCmdOutput(username, hostname, privateKey, cmd string) []byte {
	block, _ := pem.Decode([]byte(privateKey))
	rsakey, _ := x509.ParsePKCS1PrivateKey(block.Bytes)
	clientKey := &keychain{rsakey}
	clientConfig := &ssh.ClientConfig{
		User: username,
		Auth: []ssh.ClientAuth{
			ssh.ClientAuthKeyring(clientKey),
		},
	}
	client, err := ssh.Dial("tcp", hostname, clientConfig)
	if err != nil {
		log.Panic("Failed to dial: " + err.Error())
	}
	session, err := client.NewSession()
	if err != nil {
		log.Panic("Failed to create session: " + err.Error())
	}
	defer session.Close()
	output, err := session.Output(cmd)
	if err != nil {
		log.Panic("Failed to run cmd: " + err.Error())
	}
	return output
}

func latestDeployedCommit(username, hostname string, e Environment) []byte {
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	privateKey := string(getPrivateKey(path.Join(usr.HomeDir, "/.ssh/id_rsa")))
	output := remoteCmdOutput(username, hostname, privateKey, fmt.Sprintf("git --git-dir=%s rev-parse %s", e.RepoPath, e.Branch))

	return output
}

func getYAMLString(n yaml.Node, key string) string {
	return n.(yaml.Map)[key].(yaml.Scalar).String()
}

func parseYAMLEnvironment(m yaml.Node) Environment {
	e := Environment{}
	for k, v := range m.(yaml.Map) {
		e.Name = k
		e.Branch = getYAMLString(v, "branch")
		e.RepoPath = getYAMLString(v, "repo_path")
		for _, v := range v.(yaml.Map)["hosts"].(yaml.List) {
			h := Host{URI: v.(yaml.Scalar).String()}
			e.Hosts = append(e.Hosts, h)
		}
	}
	return e
}

func parseYAML() (allProjects []Project, deployUser string) {
	config, err := yaml.ReadFile(configFile)
	if err != nil {
		log.Fatal(err)
	}
	deployUser, err = config.Get("deploy_user")
	if err != nil {
		log.Fatal("config.yml is missing deploy_user: " + err.Error())
	}
	configRoot, _ := config.Root.(yaml.Map)
	projects, _ := configRoot["projects"].(yaml.List)
	allProjects = []Project{}
	for _, p := range projects {
		for _, v := range p.(yaml.Map) {
			proj := Project{Name: getYAMLString(v, "project_name"),
				GitHubURL: getYAMLString(v, "github_url"),
				RepoName:  getYAMLString(v, "repo_name"),
				RepoOwner: getYAMLString(v, "repo_owner")}
			for _, v := range v.(yaml.Map)["environments"].(yaml.List) {
				proj.Environments = append(proj.Environments, parseYAMLEnvironment(v))
			}
			allProjects = append(allProjects, proj)
		}
	}
	return allProjects, deployUser
}

func getCommit(wg *sync.WaitGroup, projects []Project, env Environment, host Host, deployUser string, i, j, k int) {
	defer wg.Done()
	host.LatestCommit = string(latestDeployedCommit(deployUser, host.URI+":"+sshPort, env))
	host.LatestCommit = strings.Trim(host.LatestCommit, "\n\r")
	projects[i].Environments[j].Hosts[k] = host
}

func retrieveCommits(projects []Project, deployUser string) []Project {
	// define a wait group to wait for all goroutines to finish
	var wg sync.WaitGroup
	githubToken := os.Getenv("GITHUB_API_TOKEN")
	t := &oauth.Transport{
		Token: &oauth.Token{AccessToken: githubToken},
	}
	client := github.NewClient(t.Client())
	for i, project := range projects {
		for j, environment := range project.Environments {
			for k, host := range environment.Hosts {
				// start a goroutine for SSH-ing on to the machine
				wg.Add(1)
				go getCommit(&wg, projects, environment, host, deployUser, i, j, k)
			}
			environment.LatestGitHubCommit = latestGitHubCommit(client, project.RepoOwner, project.RepoName, environment.Branch)
			projects[i].Environments[j] = environment
		}
	}
	// wait for goroutines to finish
	wg.Wait()
	return projects
}

func HomeHandler(w http.ResponseWriter, r *http.Request) {
	projects, deployUser := parseYAML()
	projects = retrieveCommits(projects, deployUser)
	t, err := template.New("index.html").Funcs(template.FuncMap{"eq": eq}).ParseFiles("templates/index.html")
	if err != nil {
		log.Panic(err)
	}
	err = t.Execute(w, map[string]interface{}{"Projects": projects})
	if err != nil {
		log.Panic(err)
	}
}

func main() {
	r := mux.NewRouter()
	r.HandleFunc("/", HomeHandler)
	fmt.Println("Running on localhost:" + port)
	http.ListenAndServe(":"+port, r)
}
<commit_msg>use goroutines to get latest github commits<commit_after>package main

import (
	"code.google.com/p/go.crypto/ssh"
	"code.google.com/p/goauth2/oauth"
	"crypto"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"github.com/google/go-github/github"
	"github.com/gorilla/mux"
	"github.com/kylelemons/go-gypsy/yaml"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/user"
	"path"
	"reflect"
	"strings"
	"sync"
)

var (
	port       = "8080"
	sshPort    = "22"
	configFile = "config.yml"
)

type Host struct {
	URI          string
	LatestCommit string
}

type Environment struct {
	Name               string
	RepoPath           string
	Hosts              []Host
	Branch             string
	LatestGitHubCommit string
}

type Project struct {
	Name         string
	GitHubURL    string
	RepoName     string
	RepoOwner    string
	Environments []Environment
}

/*
Thanks to Russ Cox: https://groups.google.com/d/msg/golang-nuts/OEdSDgEC7js/iyhU9DW_IKcJ
eq reports whether the first argument is equal to
any of the remaining arguments. (for use in templates)
*/
func eq(args ...interface{}) bool {
	if len(args) == 0 {
		return false
	}
	x := args[0]
	switch x := x.(type) {
	case string, int, int64, byte, float32, float64:
		for _, y := range args[1:] {
			if x == y {
				return true
			}
		}
		return false
	}

	for _, y := range args[1:] {
		if reflect.DeepEqual(x, y) {
			return true
		}
	}
	return false
}

func (h *Host) ShortCommitHash() string {
	return h.LatestCommit[:7]
}

func getPrivateKey(filename string) []byte {
	content, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Panic("Failed to open private key file: " + err.Error())
	}
	return content
}

type keychain struct {
	key *rsa.PrivateKey
}

func (k *keychain) Key(i int) (interface{}, error) {
	if i != 0 {
		return nil, nil
	}
	return &k.key.PublicKey, nil
}

func (k *keychain) Sign(i int, rand io.Reader, data []byte) (sig []byte, err error) {
	hashFunc := crypto.SHA1
	h := hashFunc.New()
	h.Write(data)
	digest := h.Sum(nil)
	return rsa.SignPKCS1v15(rand, k.key, hashFunc, digest)
}

func remoteCmdOutput(username, hostname, privateKey, cmd string) []byte {
	block, _ := pem.Decode([]byte(privateKey))
	rsakey, _ := x509.ParsePKCS1PrivateKey(block.Bytes)
	clientKey := &keychain{rsakey}
	clientConfig := &ssh.ClientConfig{
		User: username,
		Auth: []ssh.ClientAuth{
			ssh.ClientAuthKeyring(clientKey),
		},
	}
	client, err := ssh.Dial("tcp", hostname, clientConfig)
	if err != nil {
		log.Panic("Failed to dial: " + err.Error())
	}
	session, err := client.NewSession()
	if err != nil {
		log.Panic("Failed to create session: " + err.Error())
	}
	defer session.Close()
	output, err := session.Output(cmd)
	if err != nil {
		log.Panic("Failed to run cmd: " + err.Error())
	}
	return output
}

func latestDeployedCommit(username, hostname string, e Environment) []byte {
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	privateKey := string(getPrivateKey(path.Join(usr.HomeDir, "/.ssh/id_rsa")))
	output := remoteCmdOutput(username, hostname, privateKey, fmt.Sprintf("git --git-dir=%s rev-parse %s", e.RepoPath, e.Branch))

	return output
}

func getYAMLString(n yaml.Node, key string) string {
	return n.(yaml.Map)[key].(yaml.Scalar).String()
}

func parseYAMLEnvironment(m yaml.Node) Environment {
	e := Environment{}
	for k, v := range m.(yaml.Map) {
		e.Name = k
		e.Branch = getYAMLString(v, "branch")
		e.RepoPath = getYAMLString(v, "repo_path")
		for _, v := range v.(yaml.Map)["hosts"].(yaml.List) {
			h := Host{URI: v.(yaml.Scalar).String()}
			e.Hosts = append(e.Hosts, h)
		}
	}
	return e
}

func parseYAML() (allProjects []Project, deployUser string) {
	config, err := yaml.ReadFile(configFile)
	if err != nil {
		log.Fatal(err)
	}
	deployUser, err = config.Get("deploy_user")
	if err != nil {
		log.Fatal("config.yml is missing deploy_user: " + err.Error())
	}
	configRoot, _ := config.Root.(yaml.Map)
	projects, _ := configRoot["projects"].(yaml.List)
	allProjects = []Project{}
	for _, p := range projects {
		for _, v := range p.(yaml.Map) {
			proj := Project{Name: getYAMLString(v, "project_name"),
				GitHubURL: getYAMLString(v, "github_url"),
				RepoName:  getYAMLString(v, "repo_name"),
				RepoOwner: getYAMLString(v, "repo_owner")}
			for _, v := range v.(yaml.Map)["environments"].(yaml.List) {
				proj.Environments = append(proj.Environments, parseYAMLEnvironment(v))
			}
			allProjects = append(allProjects, proj)
		}
	}
	return allProjects, deployUser
}

func getCommit(wg *sync.WaitGroup, projects []Project, env Environment, host Host, deployUser string, i, j, k int) {
	defer wg.Done()
	host.LatestCommit = string(latestDeployedCommit(deployUser, host.URI+":"+sshPort, env))
	host.LatestCommit = strings.Trim(host.LatestCommit, "\n\r")
	projects[i].Environments[j].Hosts[k] = host
}

//  Get the most recent commit hash on a given branch from GitHub
func getLatestGitHubCommit(wg *sync.WaitGroup, projects []Project, environment Environment, c *github.Client, repoOwner, repoName string, i, j int) {
	defer wg.Done()
	opts := &github.CommitsListOptions{SHA: environment.Branch}
	commits, _, err := c.Repositories.ListCommits(repoOwner, repoName, opts)
	if err != nil {
		log.Panic(err)
	}
	environment.LatestGitHubCommit = *commits[0].SHA
	projects[i].Environments[j] = environment
}

func retrieveCommits(projects []Project, deployUser string) []Project {
	// define a wait group to wait for all goroutines to finish
	var wg sync.WaitGroup
	githubToken := os.Getenv("GITHUB_API_TOKEN")
	t := &oauth.Transport{
		Token: &oauth.Token{AccessToken: githubToken},
	}
	client := github.NewClient(t.Client())
	for i, project := range projects {
		for j, environment := range project.Environments {
			for k, host := range environment.Hosts {
				// start a goroutine for SSH-ing on to the machine
				wg.Add(1)
				go getCommit(&wg, projects, environment, host, deployUser, i, j, k)
			}
			wg.Add(1)
			go getLatestGitHubCommit(&wg, projects, environment, client, project.RepoOwner, project.RepoName, i, j)
		}
	}
	// wait for goroutines to finish
	wg.Wait()
	return projects
}

func HomeHandler(w http.ResponseWriter, r *http.Request) {
	projects, deployUser := parseYAML()
	projects = retrieveCommits(projects, deployUser)
	t, err := template.New("index.html").Funcs(template.FuncMap{"eq": eq}).ParseFiles("templates/index.html")
	if err != nil {
		log.Panic(err)
	}
	err = t.Execute(w, map[string]interface{}{"Projects": projects})
	if err != nil {
		log.Panic(err)
	}
}

func main() {
	r := mux.NewRouter()
	r.HandleFunc("/", HomeHandler)
	fmt.Println("Running on localhost:" + port)
	http.ListenAndServe(":"+port, r)
}
<|endoftext|>


<commit_before>// Copyright 2015 PingCAP, Inc.
//
// Copyright 2015 Wenbin Xiao
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package kv

import (
	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/tidb/kv/memkv"
	"github.com/pingcap/tidb/util/types"
)

type btreeBuffer struct {
	tree *memkv.Tree
}

// NewBTreeBuffer returns a breeBuffer.
func NewBTreeBuffer() MemBuffer {
	return &btreeBuffer{
		tree: memkv.NewTree(types.Collators[true]),
	}
}

// Get returns the value associated with the key; ErrNotExist error if the key does not exist.
func (b *btreeBuffer) Get(k Key) ([]byte, error) {
	v, ok := b.tree.Get(toIfaces(k))
	if !ok {
		return nil, ErrNotExist
	}
	return fromIfaces(v), nil
}

// Set associates the key with the value.
func (b *btreeBuffer) Set(k []byte, v []byte) error {
	b.tree.Set(toIfaces(k), toIfaces(v))
	return nil
}

// Release clear the whole buffer.
func (b *btreeBuffer) Release() {
	b.tree.Clear()
}

type btreeIter struct {
	e  *memkv.Enumerator
	k  string
	v  []byte
	ok bool
}

// NewIterator creates a new Iterator based on the provided param
func (b *btreeBuffer) NewIterator(param interface{}) Iterator {
	var e *memkv.Enumerator
	var err error
	if param == nil {
		e, err = b.tree.SeekFirst()
		if err != nil {
			return &btreeIter{ok: false}
		}
	} else {
		key := toIfaces(param.([]byte))
		e, _ = b.tree.Seek(key)
	}
	iter := &btreeIter{e: e}
	// the initial push...
	err = iter.Next()
	if err != nil {
		log.Error(err)
		return &btreeIter{ok: false}
	}
	return iter
}

// Close implements Iterator Close.
func (i *btreeIter) Close() {
	//noop
}

// Key implements Iterator Key.
func (i *btreeIter) Key() string {
	return i.k
}

// Value implements Iterator Value.
func (i *btreeIter) Value() []byte {
	return i.v
}

// Next implements Iterator Next.
func (i *btreeIter) Next() error {
	k, v, err := i.e.Next()
	i.k, i.v, i.ok = string(fromIfaces(k)), fromIfaces(v), err == nil
	return errors.Trace(err)
}

// Valid implements Iterator Valid.
func (i *btreeIter) Valid() bool {
	return i.ok
}

func toIfaces(v []byte) []interface{} {
	return []interface{}{v}
}

func fromIfaces(v []interface{}) []byte {
	if v == nil {
		return nil
	}
	return v[0].([]byte)
}
<commit_msg>*: address review comments<commit_after>// Copyright 2015 PingCAP, Inc.
//
// Copyright 2015 Wenbin Xiao
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package kv

import (
	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/tidb/kv/memkv"
	"github.com/pingcap/tidb/util/types"
)

type btreeBuffer struct {
	tree *memkv.Tree
}

// NewBTreeBuffer returns a breeBuffer.
func NewBTreeBuffer() MemBuffer {
	return &btreeBuffer{
		tree: memkv.NewTree(types.Collators[true]),
	}
}

// Get returns the value associated with the key; ErrNotExist error if the key does not exist.
func (b *btreeBuffer) Get(k Key) ([]byte, error) {
	v, ok := b.tree.Get(toIfaces(k))
	if !ok {
		return nil, ErrNotExist
	}
	return fromIfaces(v), nil
}

// Set associates the key with the value.
func (b *btreeBuffer) Set(k []byte, v []byte) error {
	b.tree.Set(toIfaces(k), toIfaces(v))
	return nil
}

// Release clear the whole buffer.
func (b *btreeBuffer) Release() {
	b.tree.Clear()
}

type btreeIter struct {
	e  *memkv.Enumerator
	k  string
	v  []byte
	ok bool
}

// NewIterator creates a new Iterator based on the provided param
func (b *btreeBuffer) NewIterator(param interface{}) Iterator {
	var e *memkv.Enumerator
	var err error
	if param == nil {
		e, err = b.tree.SeekFirst()
		if err != nil {
			return &btreeIter{ok: false}
		}
	} else {
		key := toIfaces(param.([]byte))
		e, _ = b.tree.Seek(key)
	}
	iter := &btreeIter{e: e}
	// the initial push...
	err = iter.Next()
	if err != nil {
		log.Error(err)
		return &btreeIter{ok: false}
	}
	return iter
}

// Close implements Iterator Close.
func (i *btreeIter) Close() {
	//noop
}

// Key implements Iterator Key.
func (i *btreeIter) Key() string {
	return i.k
}

// Value implements Iterator Value.
func (i *btreeIter) Value() []byte {
	return i.v
}

// Next implements Iterator Next.
func (i *btreeIter) Next() error {
	k, v, err := i.e.Next()
	if err != nil {
		i.ok = false
		return errors.Trace(err)
	}
	i.k, i.v, i.ok = string(fromIfaces(k)), fromIfaces(v), true
	return nil
}

// Valid implements Iterator Valid.
func (i *btreeIter) Valid() bool {
	return i.ok
}

func toIfaces(v []byte) []interface{} {
	return []interface{}{v}
}

func fromIfaces(v []interface{}) []byte {
	if v == nil {
		return nil
	}
	return v[0].([]byte)
}
<|endoftext|>


<commit_before>package alicloud

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"strings"

	"github.com/aliyun/aliyun-oss-go-sdk/oss"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/mitchellh/go-homedir"
	"time"
)

func resourceAlicloudOssBucketObject() *schema.Resource {
	return &schema.Resource{
		Create: resourceAlicloudOssBucketObjectPut,
		Read:   resourceAlicloudOssBucketObjectRead,
		Update: resourceAlicloudOssBucketObjectPut,
		Delete: resourceAlicloudOssBucketObjectDelete,

		Schema: map[string]*schema.Schema{
			"bucket": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"key": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"source": {
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"content"},
			},

			"content": {
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"source"},
			},

			"acl": {
				Type:         schema.TypeString,
				Default:      oss.ACLPrivate,
				Optional:     true,
				ValidateFunc: validateOssBucketAcl,
			},

			"content_type": {
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"content_length": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"cache_control": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"content_disposition": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"content_encoding": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"content_md5": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"expires": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"server_side_encryption": {
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateOssBucketObjectServerSideEncryption,
				Computed:     true,
			},

			"etag": {
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAlicloudOssBucketObjectPut(d *schema.ResourceData, meta interface{}) error {

	bucket, err := meta.(*AliyunClient).ossconn.Bucket(d.Get("bucket").(string))
	if err != nil {
		return fmt.Errorf("Error getting bucket: %#v", err)
	}

	var filePath string
	var body io.Reader

	if v, ok := d.GetOk("source"); ok {
		source := v.(string)
		path, err := homedir.Expand(source)
		if err != nil {
			return fmt.Errorf("Error expanding homedir in source (%s): %s", source, err)
		}

		filePath = path
	} else if v, ok := d.GetOk("content"); ok {
		content := v.(string)
		body = bytes.NewReader([]byte(content))
	} else {
		return fmt.Errorf("[ERROR] Must specify \"source\" or \"content\" field")
	}

	key := d.Get("key").(string)
	options, err := buildObjectHeaderOptions(d)
	if err != nil {
		return err
	}
	if filePath != "" {
		err = bucket.PutObjectFromFile(key, filePath, options...)
	}

	if body != nil {
		err = bucket.PutObject(key, body, options...)
	}

	if err != nil {
		return fmt.Errorf("Error putting object in Oss bucket (%#v): %s", bucket, err)
	}

	d.SetId(key)
	return resourceAlicloudOssBucketObjectRead(d, meta)
}

func resourceAlicloudOssBucketObjectRead(d *schema.ResourceData, meta interface{}) error {
	bucket, err := meta.(*AliyunClient).ossconn.Bucket(d.Get("bucket").(string))
	if err != nil {
		return fmt.Errorf("Error getting bucket: %#v", err)
	}

	options, err := buildObjectHeaderOptions(d)
	if err != nil {
		return fmt.Errorf("Error building object header options: %#v", err)
	}

	object, err := bucket.GetObjectDetailedMeta(d.Get("key").(string), options...)
	if err != nil {
		return fmt.Errorf("Error Reading Object: %#v", err)
	}

	if object == nil {
		log.Printf("[WARN] Reading Object: %#v, object does not exist", object)
		d.SetId("")
		return nil
	}

	log.Printf("[DEBUG] Reading Oss Bucket Object meta: %s", object)

	d.Set("content_type", object.Get("Content-Type"))
	d.Set("content_length", object.Get("Content-Length"))
	d.Set("cache_control", object.Get("Cache-Control"))
	d.Set("content_disposition", object.Get("Content-Disposition"))
	d.Set("content_encoding", object.Get("Content-Encoding"))
	d.Set("expires", object.Get("Expires"))
	d.Set("server_side_encryption", object.Get("ServerSideEncryption"))
	d.Set("etag", strings.Trim(object.Get("ETag"), `"`))

	return nil
}

func resourceAlicloudOssBucketObjectDelete(d *schema.ResourceData, meta interface{}) error {
	bucket, err := meta.(*AliyunClient).ossconn.Bucket(d.Get("bucket").(string))
	if err != nil {
		return fmt.Errorf("Error getting bucket: %#v", err)
	}
	return resource.Retry(5*time.Minute, func() *resource.RetryError {
		exist, err := bucket.IsObjectExist(d.Id())
		if err != nil {
			return resource.NonRetryableError(fmt.Errorf("OSS delete object got an error: %#v", err))
		}

		if !exist {
			return nil
		}

		if err := bucket.DeleteObject(d.Id()); err != nil {
			return resource.RetryableError(fmt.Errorf("OSS object %#v is in use - trying again while it is deleted.", d.Id()))
		}

		return nil
	})

	return nil
}

func buildObjectHeaderOptions(d *schema.ResourceData) (options []oss.Option, err error) {

	if v, ok := d.GetOk("acl"); ok {
		options = append(options, oss.ACL(oss.ACLType(v.(string))))
	}

	if v, ok := d.GetOk("content_type"); ok {
		options = append(options, oss.ContentType(v.(string)))
	}

	if v, ok := d.GetOk("cache_control"); ok {
		options = append(options, oss.CacheControl(v.(string)))
	}

	if v, ok := d.GetOk("content_disposition"); ok {
		options = append(options, oss.ContentDisposition(v.(string)))
	}

	if v, ok := d.GetOk("content_encoding"); ok {
		options = append(options, oss.ContentEncoding(v.(string)))
	}

	if v, ok := d.GetOk("content_md5"); ok {
		options = append(options, oss.ContentMD5(v.(string)))
	}

	if v, ok := d.GetOk("expires"); ok {
		options = append(options, oss.Expires(v.(time.Time)))
	}

	if v, ok := d.GetOk("server_side_encryption"); ok {
		options = append(options, oss.ServerSideEncryption(v.(string)))
	}
	if options == nil || len(options) == 0 {
		log.Printf("[WARN] Object header options is nil.")
	}
	return options, nil
}
<commit_msg>remove useless code<commit_after>package alicloud

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"strings"

	"github.com/aliyun/aliyun-oss-go-sdk/oss"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/mitchellh/go-homedir"
	"time"
)

func resourceAlicloudOssBucketObject() *schema.Resource {
	return &schema.Resource{
		Create: resourceAlicloudOssBucketObjectPut,
		Read:   resourceAlicloudOssBucketObjectRead,
		Update: resourceAlicloudOssBucketObjectPut,
		Delete: resourceAlicloudOssBucketObjectDelete,

		Schema: map[string]*schema.Schema{
			"bucket": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"key": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"source": {
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"content"},
			},

			"content": {
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"source"},
			},

			"acl": {
				Type:         schema.TypeString,
				Default:      oss.ACLPrivate,
				Optional:     true,
				ValidateFunc: validateOssBucketAcl,
			},

			"content_type": {
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"content_length": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"cache_control": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"content_disposition": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"content_encoding": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"content_md5": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"expires": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"server_side_encryption": {
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateOssBucketObjectServerSideEncryption,
				Computed:     true,
			},

			"etag": {
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAlicloudOssBucketObjectPut(d *schema.ResourceData, meta interface{}) error {

	bucket, err := meta.(*AliyunClient).ossconn.Bucket(d.Get("bucket").(string))
	if err != nil {
		return fmt.Errorf("Error getting bucket: %#v", err)
	}

	var filePath string
	var body io.Reader

	if v, ok := d.GetOk("source"); ok {
		source := v.(string)
		path, err := homedir.Expand(source)
		if err != nil {
			return fmt.Errorf("Error expanding homedir in source (%s): %s", source, err)
		}

		filePath = path
	} else if v, ok := d.GetOk("content"); ok {
		content := v.(string)
		body = bytes.NewReader([]byte(content))
	} else {
		return fmt.Errorf("[ERROR] Must specify \"source\" or \"content\" field")
	}

	key := d.Get("key").(string)
	options, err := buildObjectHeaderOptions(d)
	if err != nil {
		return err
	}
	if filePath != "" {
		err = bucket.PutObjectFromFile(key, filePath, options...)
	}

	if body != nil {
		err = bucket.PutObject(key, body, options...)
	}

	if err != nil {
		return fmt.Errorf("Error putting object in Oss bucket (%#v): %s", bucket, err)
	}

	d.SetId(key)
	return resourceAlicloudOssBucketObjectRead(d, meta)
}

func resourceAlicloudOssBucketObjectRead(d *schema.ResourceData, meta interface{}) error {
	bucket, err := meta.(*AliyunClient).ossconn.Bucket(d.Get("bucket").(string))
	if err != nil {
		return fmt.Errorf("Error getting bucket: %#v", err)
	}

	options, err := buildObjectHeaderOptions(d)
	if err != nil {
		return fmt.Errorf("Error building object header options: %#v", err)
	}

	object, err := bucket.GetObjectDetailedMeta(d.Get("key").(string), options...)
	if err != nil {
		return fmt.Errorf("Error Reading Object: %#v", err)
	}

	if object == nil {
		log.Printf("[WARN] Reading Object: %#v, object does not exist", object)
		d.SetId("")
		return nil
	}

	log.Printf("[DEBUG] Reading Oss Bucket Object meta: %s", object)

	d.Set("content_type", object.Get("Content-Type"))
	d.Set("content_length", object.Get("Content-Length"))
	d.Set("cache_control", object.Get("Cache-Control"))
	d.Set("content_disposition", object.Get("Content-Disposition"))
	d.Set("content_encoding", object.Get("Content-Encoding"))
	d.Set("expires", object.Get("Expires"))
	d.Set("server_side_encryption", object.Get("ServerSideEncryption"))
	d.Set("etag", strings.Trim(object.Get("ETag"), `"`))

	return nil
}

func resourceAlicloudOssBucketObjectDelete(d *schema.ResourceData, meta interface{}) error {
	bucket, err := meta.(*AliyunClient).ossconn.Bucket(d.Get("bucket").(string))
	if err != nil {
		return fmt.Errorf("Error getting bucket: %#v", err)
	}
	return resource.Retry(5*time.Minute, func() *resource.RetryError {
		exist, err := bucket.IsObjectExist(d.Id())
		if err != nil {
			return resource.NonRetryableError(fmt.Errorf("OSS delete object got an error: %#v", err))
		}

		if !exist {
			return nil
		}

		if err := bucket.DeleteObject(d.Id()); err != nil {
			return resource.RetryableError(fmt.Errorf("OSS object %#v is in use - trying again while it is deleted.", d.Id()))
		}

		return nil
	})

}

func buildObjectHeaderOptions(d *schema.ResourceData) (options []oss.Option, err error) {

	if v, ok := d.GetOk("acl"); ok {
		options = append(options, oss.ACL(oss.ACLType(v.(string))))
	}

	if v, ok := d.GetOk("content_type"); ok {
		options = append(options, oss.ContentType(v.(string)))
	}

	if v, ok := d.GetOk("cache_control"); ok {
		options = append(options, oss.CacheControl(v.(string)))
	}

	if v, ok := d.GetOk("content_disposition"); ok {
		options = append(options, oss.ContentDisposition(v.(string)))
	}

	if v, ok := d.GetOk("content_encoding"); ok {
		options = append(options, oss.ContentEncoding(v.(string)))
	}

	if v, ok := d.GetOk("content_md5"); ok {
		options = append(options, oss.ContentMD5(v.(string)))
	}

	if v, ok := d.GetOk("expires"); ok {
		options = append(options, oss.Expires(v.(time.Time)))
	}

	if v, ok := d.GetOk("server_side_encryption"); ok {
		options = append(options, oss.ServerSideEncryption(v.(string)))
	}
	if options == nil || len(options) == 0 {
		log.Printf("[WARN] Object header options is nil.")
	}
	return options, nil
}
<|endoftext|>