<commit_before>package nomad

import (
	"container/list"
	"fmt"
	"io"
	"net"
	"net/rpc"
	"sync"
	"sync/atomic"
	"time"

	"github.com/hashicorp/consul/tlsutil"
	"github.com/hashicorp/net-rpc-msgpackrpc"
	"github.com/hashicorp/yamux"
)

// streamClient is used to wrap a stream with an RPC client
type StreamClient struct {
	stream net.Conn
	codec  rpc.ClientCodec
}

func (sc *StreamClient) Close() {
	sc.stream.Close()
}

// Conn is a pooled connection to a Nomad server
type Conn struct {
	refCount    int32
	shouldClose int32

	addr     net.Addr
	session  *yamux.Session
	lastUsed time.Time
	version  int

	pool *ConnPool

	clients    *list.List
	clientLock sync.Mutex
}

// markForUse does all the bookkeeping required to ready a connection for use.
func (c *Conn) markForUse() {
	c.lastUsed = time.Now()
	atomic.AddInt32(&c.refCount, 1)
}

func (c *Conn) Close() error {
	return c.session.Close()
}

// getClient is used to get a cached or new client
func (c *Conn) getClient() (*StreamClient, error) {
	// Check for cached client
	c.clientLock.Lock()
	front := c.clients.Front()
	if front != nil {
		c.clients.Remove(front)
	}
	c.clientLock.Unlock()
	if front != nil {
		return front.Value.(*StreamClient), nil
	}

	// Open a new session
	stream, err := c.session.Open()
	if err != nil {
		return nil, err
	}

	// Create a client codec
	codec := NewClientCodec(stream)

	// Return a new stream client
	sc := &StreamClient{
		stream: stream,
		codec:  codec,
	}
	return sc, nil
}

// returnStream is used when done with a stream
// to allow re-use by a future RPC
func (c *Conn) returnClient(client *StreamClient) {
	didSave := false
	c.clientLock.Lock()
	if c.clients.Len() < c.pool.maxStreams && atomic.LoadInt32(&c.shouldClose) == 0 {
		c.clients.PushFront(client)
		didSave = true
	}
	c.clientLock.Unlock()
	if !didSave {
		client.Close()
	}
}

// ConnPool is used to maintain a connection pool to other
// Nomad servers. This is used to reduce the latency of
// RPC requests between servers. It is only used to pool
// connections in the rpcNomad mode. Raft connections
// are pooled separately.
type ConnPool struct {
	sync.Mutex

	// LogOutput is used to control logging
	logOutput io.Writer

	// The maximum time to keep a connection open
	maxTime time.Duration

	// The maximum number of open streams to keep
	maxStreams int

	// Pool maps an address to a open connection
	pool map[string]*Conn

	// limiter is used to throttle the number of connect attempts
	// to a given address. The first thread will attempt a connection
	// and put a channel in here, which all other threads will wait
	// on to close.
	limiter map[string]chan struct{}

	// TLS wrapper
	tlsWrap tlsutil.DCWrapper

	// Used to indicate the pool is shutdown
	shutdown   bool
	shutdownCh chan struct{}
}

// NewPool is used to make a new connection pool
// Maintain at most one connection per host, for up to maxTime.
// Set maxTime to 0 to disable reaping. maxStreams is used to control
// the number of idle streams allowed.
// If TLS settings are provided outgoing connections use TLS.
func NewPool(logOutput io.Writer, maxTime time.Duration, maxStreams int, tlsWrap tlsutil.DCWrapper) *ConnPool {
	pool := &ConnPool{
		logOutput:  logOutput,
		maxTime:    maxTime,
		maxStreams: maxStreams,
		pool:       make(map[string]*Conn),
		limiter:    make(map[string]chan struct{}),
		tlsWrap:    tlsWrap,
		shutdownCh: make(chan struct{}),
	}
	if maxTime > 0 {
		go pool.reap()
	}
	return pool
}

// Shutdown is used to close the connection pool
func (p *ConnPool) Shutdown() error {
	p.Lock()
	defer p.Unlock()

	for _, conn := range p.pool {
		conn.Close()
	}
	p.pool = make(map[string]*Conn)

	if p.shutdown {
		return nil
	}
	p.shutdown = true
	close(p.shutdownCh)
	return nil
}

// Acquire is used to get a connection that is
// pooled or to return a new connection
func (p *ConnPool) acquire(region string, addr net.Addr, version int) (*Conn, error) {
	// Check to see if there's a pooled connection available. This is up
	// here since it should the the vastly more common case than the rest
	// of the code here.
	p.Lock()
	c := p.pool[addr.String()]
	if c != nil {
		c.markForUse()
		p.Unlock()
		return c, nil
	}

	// If not (while we are still locked), set up the throttling structure
	// for this address, which will make everyone else wait until our
	// attempt is done.
	var wait chan struct{}
	var ok bool
	if wait, ok = p.limiter[addr.String()]; !ok {
		wait = make(chan struct{})
		p.limiter[addr.String()] = wait
	}
	isLeadThread := !ok
	p.Unlock()

	// If we are the lead thread, make the new connection and then wake
	// everybody else up to see if we got it.
	if isLeadThread {
		c, err := p.getNewConn(region, addr, version)
		p.Lock()
		delete(p.limiter, addr.String())
		close(wait)
		if err != nil {
			p.Unlock()
			return nil, err
		}

		p.pool[addr.String()] = c
		p.Unlock()
		return c, nil
	}

	// Otherwise, wait for the lead thread to attempt the connection
	// and use what's in the pool at that point.
	select {
	case <-p.shutdownCh:
		return nil, fmt.Errorf("rpc error: shutdown")
	case <-wait:
	}

	// See if the lead thread was able to get us a connection.
	p.Lock()
	if c := p.pool[addr.String()]; c != nil {
		c.markForUse()
		p.Unlock()
		return c, nil
	}

	p.Unlock()
	return nil, fmt.Errorf("rpc error: lead thread didn't get connection")
}

// getNewConn is used to return a new connection
func (p *ConnPool) getNewConn(region string, addr net.Addr, version int) (*Conn, error) {
	// Try to dial the conn
	conn, err := net.DialTimeout("tcp", addr.String(), 10*time.Second)
	if err != nil {
		return nil, err
	}

	// Cast to TCPConn
	if tcp, ok := conn.(*net.TCPConn); ok {
		tcp.SetKeepAlive(true)
		tcp.SetNoDelay(true)
	}

	// Check if TLS is enabled
	if p.tlsWrap != nil {
		// Switch the connection into TLS mode
		if _, err := conn.Write([]byte{byte(rpcTLS)}); err != nil {
			conn.Close()
			return nil, err
		}

		// Wrap the connection in a TLS client
		tlsConn, err := p.tlsWrap(region, conn)
		if err != nil {
			conn.Close()
			return nil, err
		}
		conn = tlsConn
	}

	// Write the multiplex byte to set the mode
	if _, err := conn.Write([]byte{byte(rpcMultiplex)}); err != nil {
		conn.Close()
		return nil, err
	}

	// Setup the logger
	conf := yamux.DefaultConfig()
	conf.LogOutput = p.logOutput

	// Create a multiplexed session
	session, err := yamux.Client(conn, conf)
	if err != nil {
		conn.Close()
		return nil, err
	}

	// Wrap the connection
	c := &Conn{
		refCount: 1,
		addr:     addr,
		session:  session,
		clients:  list.New(),
		lastUsed: time.Now(),
		version:  version,
		pool:     p,
	}
	return c, nil
}

// clearConn is used to clear any cached connection, potentially in response to an erro
func (p *ConnPool) clearConn(conn *Conn) {
	// Ensure returned streams are closed
	atomic.StoreInt32(&conn.shouldClose, 1)

	// Clear from the cache
	p.Lock()
	if c, ok := p.pool[conn.addr.String()]; ok && c == conn {
		delete(p.pool, conn.addr.String())
	}
	p.Unlock()

	// Close down immediately if idle
	if refCount := atomic.LoadInt32(&conn.refCount); refCount == 0 {
		conn.Close()
	}
}

// releaseConn is invoked when we are done with a conn to reduce the ref count
func (p *ConnPool) releaseConn(conn *Conn) {
	refCount := atomic.AddInt32(&conn.refCount, -1)
	if refCount == 0 && atomic.LoadInt32(&conn.shouldClose) == 1 {
		conn.Close()
	}
}

// getClient is used to get a usable client for an address and protocol version
func (p *ConnPool) getClient(region string, addr net.Addr, version int) (*Conn, *StreamClient, error) {
	retries := 0
START:
	// Try to get a conn first
	conn, err := p.acquire(region, addr, version)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to get conn: %v", err)
	}

	// Get a client
	client, err := conn.getClient()
	if err != nil {
		p.clearConn(conn)
		p.releaseConn(conn)

		// Try to redial, possible that the TCP session closed due to timeout
		if retries == 0 {
			retries++
			goto START
		}
		return nil, nil, fmt.Errorf("failed to start stream: %v", err)
	}
	return conn, client, nil
}

// RPC is used to make an RPC call to a remote host
func (p *ConnPool) RPC(region string, addr net.Addr, version int, method string, args interface{}, reply interface{}) error {
	// Get a usable client
	conn, sc, err := p.getClient(region, addr, version)
	if err != nil {
		return fmt.Errorf("rpc error: %v", err)
	}

	// Make the RPC call
	err = msgpackrpc.CallWithCodec(sc.codec, method, args, reply)
	if err != nil {
		sc.Close()
		p.releaseConn(conn)
		return fmt.Errorf("rpc error: %v", err)
	}

	// Done with the connection
	conn.returnClient(sc)
	p.releaseConn(conn)
	return nil
}

// Reap is used to close conns open over maxTime
func (p *ConnPool) reap() {
	for {
		// Sleep for a while
		select {
		case <-p.shutdownCh:
			return
		case <-time.After(time.Second):
		}

		// Reap all old conns
		p.Lock()
		var removed []string
		now := time.Now()
		for host, conn := range p.pool {
			// Skip recently used connections
			if now.Sub(conn.lastUsed) < p.maxTime {
				continue
			}

			// Skip connections with active streams
			if atomic.LoadInt32(&conn.refCount) > 0 {
				continue
			}

			// Close the conn
			conn.Close()

			// Remove from pool
			removed = append(removed, host)
		}
		for _, host := range removed {
			delete(p.pool, host)
		}
		p.Unlock()
	}
}
<commit_msg>nomad: shrink the recv buffer of idle streams<commit_after>package nomad

import (
	"container/list"
	"fmt"
	"io"
	"net"
	"net/rpc"
	"sync"
	"sync/atomic"
	"time"

	"github.com/hashicorp/consul/tlsutil"
	"github.com/hashicorp/net-rpc-msgpackrpc"
	"github.com/hashicorp/yamux"
)

// streamClient is used to wrap a stream with an RPC client
type StreamClient struct {
	stream net.Conn
	codec  rpc.ClientCodec
}

func (sc *StreamClient) Close() {
	sc.stream.Close()
}

// Conn is a pooled connection to a Nomad server
type Conn struct {
	refCount    int32
	shouldClose int32

	addr     net.Addr
	session  *yamux.Session
	lastUsed time.Time
	version  int

	pool *ConnPool

	clients    *list.List
	clientLock sync.Mutex
}

// markForUse does all the bookkeeping required to ready a connection for use.
func (c *Conn) markForUse() {
	c.lastUsed = time.Now()
	atomic.AddInt32(&c.refCount, 1)
}

func (c *Conn) Close() error {
	return c.session.Close()
}

// getClient is used to get a cached or new client
func (c *Conn) getClient() (*StreamClient, error) {
	// Check for cached client
	c.clientLock.Lock()
	front := c.clients.Front()
	if front != nil {
		c.clients.Remove(front)
	}
	c.clientLock.Unlock()
	if front != nil {
		return front.Value.(*StreamClient), nil
	}

	// Open a new session
	stream, err := c.session.Open()
	if err != nil {
		return nil, err
	}

	// Create a client codec
	codec := NewClientCodec(stream)

	// Return a new stream client
	sc := &StreamClient{
		stream: stream,
		codec:  codec,
	}
	return sc, nil
}

// returnStream is used when done with a stream
// to allow re-use by a future RPC
func (c *Conn) returnClient(client *StreamClient) {
	didSave := false
	c.clientLock.Lock()
	if c.clients.Len() < c.pool.maxStreams && atomic.LoadInt32(&c.shouldClose) == 0 {
		c.clients.PushFront(client)
		didSave = true

		// If this is a Yamux stream, shrink the internal buffers so that
		// we can GC the idle memory
		if ys, ok := client.stream.(*yamux.Stream); ok {
			ys.Shrink()
		}
	}
	c.clientLock.Unlock()
	if !didSave {
		client.Close()
	}
}

// ConnPool is used to maintain a connection pool to other
// Nomad servers. This is used to reduce the latency of
// RPC requests between servers. It is only used to pool
// connections in the rpcNomad mode. Raft connections
// are pooled separately.
type ConnPool struct {
	sync.Mutex

	// LogOutput is used to control logging
	logOutput io.Writer

	// The maximum time to keep a connection open
	maxTime time.Duration

	// The maximum number of open streams to keep
	maxStreams int

	// Pool maps an address to a open connection
	pool map[string]*Conn

	// limiter is used to throttle the number of connect attempts
	// to a given address. The first thread will attempt a connection
	// and put a channel in here, which all other threads will wait
	// on to close.
	limiter map[string]chan struct{}

	// TLS wrapper
	tlsWrap tlsutil.DCWrapper

	// Used to indicate the pool is shutdown
	shutdown   bool
	shutdownCh chan struct{}
}

// NewPool is used to make a new connection pool
// Maintain at most one connection per host, for up to maxTime.
// Set maxTime to 0 to disable reaping. maxStreams is used to control
// the number of idle streams allowed.
// If TLS settings are provided outgoing connections use TLS.
func NewPool(logOutput io.Writer, maxTime time.Duration, maxStreams int, tlsWrap tlsutil.DCWrapper) *ConnPool {
	pool := &ConnPool{
		logOutput:  logOutput,
		maxTime:    maxTime,
		maxStreams: maxStreams,
		pool:       make(map[string]*Conn),
		limiter:    make(map[string]chan struct{}),
		tlsWrap:    tlsWrap,
		shutdownCh: make(chan struct{}),
	}
	if maxTime > 0 {
		go pool.reap()
	}
	return pool
}

// Shutdown is used to close the connection pool
func (p *ConnPool) Shutdown() error {
	p.Lock()
	defer p.Unlock()

	for _, conn := range p.pool {
		conn.Close()
	}
	p.pool = make(map[string]*Conn)

	if p.shutdown {
		return nil
	}
	p.shutdown = true
	close(p.shutdownCh)
	return nil
}

// Acquire is used to get a connection that is
// pooled or to return a new connection
func (p *ConnPool) acquire(region string, addr net.Addr, version int) (*Conn, error) {
	// Check to see if there's a pooled connection available. This is up
	// here since it should the the vastly more common case than the rest
	// of the code here.
	p.Lock()
	c := p.pool[addr.String()]
	if c != nil {
		c.markForUse()
		p.Unlock()
		return c, nil
	}

	// If not (while we are still locked), set up the throttling structure
	// for this address, which will make everyone else wait until our
	// attempt is done.
	var wait chan struct{}
	var ok bool
	if wait, ok = p.limiter[addr.String()]; !ok {
		wait = make(chan struct{})
		p.limiter[addr.String()] = wait
	}
	isLeadThread := !ok
	p.Unlock()

	// If we are the lead thread, make the new connection and then wake
	// everybody else up to see if we got it.
	if isLeadThread {
		c, err := p.getNewConn(region, addr, version)
		p.Lock()
		delete(p.limiter, addr.String())
		close(wait)
		if err != nil {
			p.Unlock()
			return nil, err
		}

		p.pool[addr.String()] = c
		p.Unlock()
		return c, nil
	}

	// Otherwise, wait for the lead thread to attempt the connection
	// and use what's in the pool at that point.
	select {
	case <-p.shutdownCh:
		return nil, fmt.Errorf("rpc error: shutdown")
	case <-wait:
	}

	// See if the lead thread was able to get us a connection.
	p.Lock()
	if c := p.pool[addr.String()]; c != nil {
		c.markForUse()
		p.Unlock()
		return c, nil
	}

	p.Unlock()
	return nil, fmt.Errorf("rpc error: lead thread didn't get connection")
}

// getNewConn is used to return a new connection
func (p *ConnPool) getNewConn(region string, addr net.Addr, version int) (*Conn, error) {
	// Try to dial the conn
	conn, err := net.DialTimeout("tcp", addr.String(), 10*time.Second)
	if err != nil {
		return nil, err
	}

	// Cast to TCPConn
	if tcp, ok := conn.(*net.TCPConn); ok {
		tcp.SetKeepAlive(true)
		tcp.SetNoDelay(true)
	}

	// Check if TLS is enabled
	if p.tlsWrap != nil {
		// Switch the connection into TLS mode
		if _, err := conn.Write([]byte{byte(rpcTLS)}); err != nil {
			conn.Close()
			return nil, err
		}

		// Wrap the connection in a TLS client
		tlsConn, err := p.tlsWrap(region, conn)
		if err != nil {
			conn.Close()
			return nil, err
		}
		conn = tlsConn
	}

	// Write the multiplex byte to set the mode
	if _, err := conn.Write([]byte{byte(rpcMultiplex)}); err != nil {
		conn.Close()
		return nil, err
	}

	// Setup the logger
	conf := yamux.DefaultConfig()
	conf.LogOutput = p.logOutput

	// Create a multiplexed session
	session, err := yamux.Client(conn, conf)
	if err != nil {
		conn.Close()
		return nil, err
	}

	// Wrap the connection
	c := &Conn{
		refCount: 1,
		addr:     addr,
		session:  session,
		clients:  list.New(),
		lastUsed: time.Now(),
		version:  version,
		pool:     p,
	}
	return c, nil
}

// clearConn is used to clear any cached connection, potentially in response to an erro
func (p *ConnPool) clearConn(conn *Conn) {
	// Ensure returned streams are closed
	atomic.StoreInt32(&conn.shouldClose, 1)

	// Clear from the cache
	p.Lock()
	if c, ok := p.pool[conn.addr.String()]; ok && c == conn {
		delete(p.pool, conn.addr.String())
	}
	p.Unlock()

	// Close down immediately if idle
	if refCount := atomic.LoadInt32(&conn.refCount); refCount == 0 {
		conn.Close()
	}
}

// releaseConn is invoked when we are done with a conn to reduce the ref count
func (p *ConnPool) releaseConn(conn *Conn) {
	refCount := atomic.AddInt32(&conn.refCount, -1)
	if refCount == 0 && atomic.LoadInt32(&conn.shouldClose) == 1 {
		conn.Close()
	}
}

// getClient is used to get a usable client for an address and protocol version
func (p *ConnPool) getClient(region string, addr net.Addr, version int) (*Conn, *StreamClient, error) {
	retries := 0
START:
	// Try to get a conn first
	conn, err := p.acquire(region, addr, version)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to get conn: %v", err)
	}

	// Get a client
	client, err := conn.getClient()
	if err != nil {
		p.clearConn(conn)
		p.releaseConn(conn)

		// Try to redial, possible that the TCP session closed due to timeout
		if retries == 0 {
			retries++
			goto START
		}
		return nil, nil, fmt.Errorf("failed to start stream: %v", err)
	}
	return conn, client, nil
}

// RPC is used to make an RPC call to a remote host
func (p *ConnPool) RPC(region string, addr net.Addr, version int, method string, args interface{}, reply interface{}) error {
	// Get a usable client
	conn, sc, err := p.getClient(region, addr, version)
	if err != nil {
		return fmt.Errorf("rpc error: %v", err)
	}

	// Make the RPC call
	err = msgpackrpc.CallWithCodec(sc.codec, method, args, reply)
	if err != nil {
		sc.Close()
		p.releaseConn(conn)
		return fmt.Errorf("rpc error: %v", err)
	}

	// Done with the connection
	conn.returnClient(sc)
	p.releaseConn(conn)
	return nil
}

// Reap is used to close conns open over maxTime
func (p *ConnPool) reap() {
	for {
		// Sleep for a while
		select {
		case <-p.shutdownCh:
			return
		case <-time.After(time.Second):
		}

		// Reap all old conns
		p.Lock()
		var removed []string
		now := time.Now()
		for host, conn := range p.pool {
			// Skip recently used connections
			if now.Sub(conn.lastUsed) < p.maxTime {
				continue
			}

			// Skip connections with active streams
			if atomic.LoadInt32(&conn.refCount) > 0 {
				continue
			}

			// Close the conn
			conn.Close()

			// Remove from pool
			removed = append(removed, host)
		}
		for _, host := range removed {
			delete(p.pool, host)
		}
		p.Unlock()
	}
}
<|endoftext|>


<commit_before>package controller

import (
	"github.com/gin-gonic/gin"

	"github.com/drone/drone/model"
	"github.com/drone/drone/store"
)

var (
	badgeSuccess = `<svg xmlns="http://www.w3.org/2000/svg" width="91" height="20"><linearGradient id="a" x2="0" y2="100%"><stop offset="0" stop-color="#bbb" stop-opacity=".1"/><stop offset="1" stop-opacity=".1"/></linearGradient><rect rx="3" width="91" height="20" fill="#555"/><rect rx="3" x="37" width="54" height="20" fill="#4c1"/><path fill="#4c1" d="M37 0h4v20h-4z"/><rect rx="3" width="91" height="20" fill="url(#a)"/><g fill="#fff" text-anchor="middle" font-family="DejaVu Sans,Verdana,Geneva,sans-serif" font-size="11"><text x="19.5" y="15" fill="#010101" fill-opacity=".3">build</text><text x="19.5" y="14">build</text><text x="63" y="15" fill="#010101" fill-opacity=".3">success</text><text x="63" y="14">success</text></g></svg>`
	badgeFailure = `<svg xmlns="http://www.w3.org/2000/svg" width="83" height="20"><linearGradient id="a" x2="0" y2="100%"><stop offset="0" stop-color="#bbb" stop-opacity=".1"/><stop offset="1" stop-opacity=".1"/></linearGradient><rect rx="3" width="83" height="20" fill="#555"/><rect rx="3" x="37" width="46" height="20" fill="#e05d44"/><path fill="#e05d44" d="M37 0h4v20h-4z"/><rect rx="3" width="83" height="20" fill="url(#a)"/><g fill="#fff" text-anchor="middle" font-family="DejaVu Sans,Verdana,Geneva,sans-serif" font-size="11"><text x="19.5" y="15" fill="#010101" fill-opacity=".3">build</text><text x="19.5" y="14">build</text><text x="59" y="15" fill="#010101" fill-opacity=".3">failure</text><text x="59" y="14">failure</text></g></svg>`
	badgeStarted = `<svg xmlns="http://www.w3.org/2000/svg" width="87" height="20"><linearGradient id="a" x2="0" y2="100%"><stop offset="0" stop-color="#bbb" stop-opacity=".1"/><stop offset="1" stop-opacity=".1"/></linearGradient><rect rx="3" width="87" height="20" fill="#555"/><rect rx="3" x="37" width="50" height="20" fill="#dfb317"/><path fill="#dfb317" d="M37 0h4v20h-4z"/><rect rx="3" width="87" height="20" fill="url(#a)"/><g fill="#fff" text-anchor="middle" font-family="DejaVu Sans,Verdana,Geneva,sans-serif" font-size="11"><text x="19.5" y="15" fill="#010101" fill-opacity=".3">build</text><text x="19.5" y="14">build</text><text x="61" y="15" fill="#010101" fill-opacity=".3">started</text><text x="61" y="14">started</text></g></svg>`
	badgeError   = `<svg xmlns="http://www.w3.org/2000/svg" width="76" height="20"><linearGradient id="a" x2="0" y2="100%"><stop offset="0" stop-color="#bbb" stop-opacity=".1"/><stop offset="1" stop-opacity=".1"/></linearGradient><rect rx="3" width="76" height="20" fill="#555"/><rect rx="3" x="37" width="39" height="20" fill="#9f9f9f"/><path fill="#9f9f9f" d="M37 0h4v20h-4z"/><rect rx="3" width="76" height="20" fill="url(#a)"/><g fill="#fff" text-anchor="middle" font-family="DejaVu Sans,Verdana,Geneva,sans-serif" font-size="11"><text x="19.5" y="15" fill="#010101" fill-opacity=".3">build</text><text x="19.5" y="14">build</text><text x="55.5" y="15" fill="#010101" fill-opacity=".3">error</text><text x="55.5" y="14">error</text></g></svg>`
	badgeNone    = `<svg xmlns="http://www.w3.org/2000/svg" width="75" height="20"><linearGradient id="a" x2="0" y2="100%"><stop offset="0" stop-color="#bbb" stop-opacity=".1"/><stop offset="1" stop-opacity=".1"/></linearGradient><rect rx="3" width="75" height="20" fill="#555"/><rect rx="3" x="37" width="38" height="20" fill="#9f9f9f"/><path fill="#9f9f9f" d="M37 0h4v20h-4z"/><rect rx="3" width="75" height="20" fill="url(#a)"/><g fill="#fff" text-anchor="middle" font-family="DejaVu Sans,Verdana,Geneva,sans-serif" font-size="11"><text x="19.5" y="15" fill="#010101" fill-opacity=".3">build</text><text x="19.5" y="14">build</text><text x="55" y="15" fill="#010101" fill-opacity=".3">none</text><text x="55" y="14">none</text></g></svg>`
)

func GetBadge(c *gin.Context) {
	repo, err := store.GetRepoOwnerName(c,
		c.Param("owner"),
		c.Param("name"),
	)
	if err != nil {
		c.AbortWithStatus(404)
		return
	}

	// an SVG response is always served, even when error, so
	// we can go ahead and set the content type appropriately.
	c.Writer.Header().Set("Content-Type", "image/svg+xml")

	// if no commit was found then display
	// the 'none' badge, instead of throwing
	// an error response
	branch := c.Query("branch")
	if len(branch) == 0 {
		branch = repo.Branch
	}

	build, err := store.GetBuildLast(c, repo, branch)
	if err != nil {
		c.String(404, badgeNone)
		return
	}

	switch build.Status {
	case model.StatusSuccess:
		c.String(200, badgeSuccess)
	case model.StatusFailure:
		c.String(200, badgeFailure)
	case model.StatusError, model.StatusKilled:
		c.String(200, badgeError)
	case model.StatusPending, model.StatusRunning:
		c.String(200, badgeStarted)
	default:
		c.String(404, badgeNone)
	}
}

func GetCC(c *gin.Context) {
	repo, err := store.GetRepoOwnerName(c,
		c.Param("owner"),
		c.Param("name"),
	)
	if err != nil {
		c.AbortWithStatus(404)
		return
	}

	builds, err := store.GetBuildList(c, repo)
	if err != nil || len(builds) == 0 {
		c.AbortWithStatus(404)
		return
	}

	cc := model.NewCC(repo, builds[0], "")
	c.XML(200, cc)
}
<commit_msg>Corrects CC Menu url<commit_after>package controller

import (
	"fmt"

	"github.com/gin-gonic/gin"

	"github.com/drone/drone/model"
	"github.com/drone/drone/shared/httputil"
	"github.com/drone/drone/store"
)

var (
	badgeSuccess = `<svg xmlns="http://www.w3.org/2000/svg" width="91" height="20"><linearGradient id="a" x2="0" y2="100%"><stop offset="0" stop-color="#bbb" stop-opacity=".1"/><stop offset="1" stop-opacity=".1"/></linearGradient><rect rx="3" width="91" height="20" fill="#555"/><rect rx="3" x="37" width="54" height="20" fill="#4c1"/><path fill="#4c1" d="M37 0h4v20h-4z"/><rect rx="3" width="91" height="20" fill="url(#a)"/><g fill="#fff" text-anchor="middle" font-family="DejaVu Sans,Verdana,Geneva,sans-serif" font-size="11"><text x="19.5" y="15" fill="#010101" fill-opacity=".3">build</text><text x="19.5" y="14">build</text><text x="63" y="15" fill="#010101" fill-opacity=".3">success</text><text x="63" y="14">success</text></g></svg>`
	badgeFailure = `<svg xmlns="http://www.w3.org/2000/svg" width="83" height="20"><linearGradient id="a" x2="0" y2="100%"><stop offset="0" stop-color="#bbb" stop-opacity=".1"/><stop offset="1" stop-opacity=".1"/></linearGradient><rect rx="3" width="83" height="20" fill="#555"/><rect rx="3" x="37" width="46" height="20" fill="#e05d44"/><path fill="#e05d44" d="M37 0h4v20h-4z"/><rect rx="3" width="83" height="20" fill="url(#a)"/><g fill="#fff" text-anchor="middle" font-family="DejaVu Sans,Verdana,Geneva,sans-serif" font-size="11"><text x="19.5" y="15" fill="#010101" fill-opacity=".3">build</text><text x="19.5" y="14">build</text><text x="59" y="15" fill="#010101" fill-opacity=".3">failure</text><text x="59" y="14">failure</text></g></svg>`
	badgeStarted = `<svg xmlns="http://www.w3.org/2000/svg" width="87" height="20"><linearGradient id="a" x2="0" y2="100%"><stop offset="0" stop-color="#bbb" stop-opacity=".1"/><stop offset="1" stop-opacity=".1"/></linearGradient><rect rx="3" width="87" height="20" fill="#555"/><rect rx="3" x="37" width="50" height="20" fill="#dfb317"/><path fill="#dfb317" d="M37 0h4v20h-4z"/><rect rx="3" width="87" height="20" fill="url(#a)"/><g fill="#fff" text-anchor="middle" font-family="DejaVu Sans,Verdana,Geneva,sans-serif" font-size="11"><text x="19.5" y="15" fill="#010101" fill-opacity=".3">build</text><text x="19.5" y="14">build</text><text x="61" y="15" fill="#010101" fill-opacity=".3">started</text><text x="61" y="14">started</text></g></svg>`
	badgeError   = `<svg xmlns="http://www.w3.org/2000/svg" width="76" height="20"><linearGradient id="a" x2="0" y2="100%"><stop offset="0" stop-color="#bbb" stop-opacity=".1"/><stop offset="1" stop-opacity=".1"/></linearGradient><rect rx="3" width="76" height="20" fill="#555"/><rect rx="3" x="37" width="39" height="20" fill="#9f9f9f"/><path fill="#9f9f9f" d="M37 0h4v20h-4z"/><rect rx="3" width="76" height="20" fill="url(#a)"/><g fill="#fff" text-anchor="middle" font-family="DejaVu Sans,Verdana,Geneva,sans-serif" font-size="11"><text x="19.5" y="15" fill="#010101" fill-opacity=".3">build</text><text x="19.5" y="14">build</text><text x="55.5" y="15" fill="#010101" fill-opacity=".3">error</text><text x="55.5" y="14">error</text></g></svg>`
	badgeNone    = `<svg xmlns="http://www.w3.org/2000/svg" width="75" height="20"><linearGradient id="a" x2="0" y2="100%"><stop offset="0" stop-color="#bbb" stop-opacity=".1"/><stop offset="1" stop-opacity=".1"/></linearGradient><rect rx="3" width="75" height="20" fill="#555"/><rect rx="3" x="37" width="38" height="20" fill="#9f9f9f"/><path fill="#9f9f9f" d="M37 0h4v20h-4z"/><rect rx="3" width="75" height="20" fill="url(#a)"/><g fill="#fff" text-anchor="middle" font-family="DejaVu Sans,Verdana,Geneva,sans-serif" font-size="11"><text x="19.5" y="15" fill="#010101" fill-opacity=".3">build</text><text x="19.5" y="14">build</text><text x="55" y="15" fill="#010101" fill-opacity=".3">none</text><text x="55" y="14">none</text></g></svg>`
)

func GetBadge(c *gin.Context) {
	repo, err := store.GetRepoOwnerName(c,
		c.Param("owner"),
		c.Param("name"),
	)
	if err != nil {
		c.AbortWithStatus(404)
		return
	}

	// an SVG response is always served, even when error, so
	// we can go ahead and set the content type appropriately.
	c.Writer.Header().Set("Content-Type", "image/svg+xml")

	// if no commit was found then display
	// the 'none' badge, instead of throwing
	// an error response
	branch := c.Query("branch")
	if len(branch) == 0 {
		branch = repo.Branch
	}

	build, err := store.GetBuildLast(c, repo, branch)
	if err != nil {
		c.String(404, badgeNone)
		return
	}

	switch build.Status {
	case model.StatusSuccess:
		c.String(200, badgeSuccess)
	case model.StatusFailure:
		c.String(200, badgeFailure)
	case model.StatusError, model.StatusKilled:
		c.String(200, badgeError)
	case model.StatusPending, model.StatusRunning:
		c.String(200, badgeStarted)
	default:
		c.String(404, badgeNone)
	}
}

func GetCC(c *gin.Context) {
	repo, err := store.GetRepoOwnerName(c,
		c.Param("owner"),
		c.Param("name"),
	)
	if err != nil {
		c.AbortWithStatus(404)
		return
	}

	builds, err := store.GetBuildList(c, repo)
	if err != nil || len(builds) == 0 {
		c.AbortWithStatus(404)
		return
	}

	url := fmt.Sprintf("%s/%s/%d", httputil.GetURL(c.Request), repo.FullName, builds[0].Number)
	cc := model.NewCC(repo, builds[0], url)
	c.XML(200, cc)
}
<|endoftext|>


<commit_before>// This file contains functions for transpiling common branching and control
// flow, such as "if", "while", "do" and "for". The more complicated control
// flows like "switch" will be put into their own file of the same or sensible
// name.

package transpiler

import (
	"fmt"
	"go/token"

	"github.com/elliotchance/c2go/ast"
	"github.com/elliotchance/c2go/program"
	"github.com/elliotchance/c2go/types"
	"github.com/elliotchance/c2go/util"

	goast "go/ast"
)

func transpileIfStmt(n *ast.IfStmt, p *program.Program) (
	*goast.IfStmt, []goast.Stmt, []goast.Stmt, error) {
	preStmts := []goast.Stmt{}
	postStmts := []goast.Stmt{}
	children := n.Children

	// There is always 4 or 5 children in an IfStmt. For example:
	//
	//     if (i == 0) {
	//         return 0;
	//     } else {
	//         return 1;
	//     }
	//
	// 1. Not sure what this is for. This gets removed.
	// 2. Not sure what this is for.
	// 3. conditional = BinaryOperator: i == 0
	// 4. body = CompoundStmt: { return 0; }
	// 5. elseBody = CompoundStmt: { return 1; }
	//
	// elseBody will be nil if there is no else clause.

	// On linux I have seen only 4 children for an IfStmt with the same
	// definitions above, but missing the first argument. Since we don't
	// know what the first argument is for anyway we will just remove it on
	// Mac if necessary.
	if len(children) == 5 && children[0] != nil {
		panic("non-nil child 0 in IfStmt")
	}
	if len(children) == 5 {
		children = children[1:]
	}

	// From here on there must be 4 children.
	if len(children) != 4 {
		panic(fmt.Sprintf("Expected 4 children in IfStmt, got %#v", children))
	}

	// Maybe we will discover what the nil value is?
	if children[0] != nil {
		panic("non-nil child 0 in IfStmt")
	}

	conditional, conditionalType, newPre, newPost, err := transpileToExpr(children[1], p)
	if err != nil {
		return nil, nil, nil, err
	}

	// The condition in Go must always be a bool.
	boolCondition, err := types.CastExpr(p, conditional, conditionalType, "bool")
	p.AddMessage(ast.GenerateWarningOrErrorMessage(err, n, boolCondition == nil))

	if boolCondition == nil {
		boolCondition = util.NewNil()
	}

	body, newPre, newPost, err := transpileToBlockStmt(children[2], p)
	if err != nil {
		return nil, nil, nil, err
	}

	preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

	r := &goast.IfStmt{
		Cond: boolCondition,
		Body: body,
	}

	if children[3] != nil {
		elseBody, newPre, newPost, err := transpileToBlockStmt(children[3], p)
		if err != nil {
			return nil, nil, nil, err
		}

		preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

		r.Else = elseBody
	}

	return r, newPre, newPost, nil
}

func transpileForStmt(n *ast.ForStmt, p *program.Program) (
	*goast.ForStmt, []goast.Stmt, []goast.Stmt, error) {
	preStmts := []goast.Stmt{}
	postStmts := []goast.Stmt{}

	children := n.Children

	// There are always 5 children in a ForStmt, for example:
	//
	//     for ( c = 0 ; c < n ; c++ ) {
	//         doSomething();
	//     }
	//
	// 1. initExpression = BinaryStmt: c = 0
	// 2. Not sure what this is for, but it's always nil. There is a panic
	//    below in case we discover what it is used for (pun intended).
	// 3. conditionalExpression = BinaryStmt: c < n
	// 4. stepExpression = BinaryStmt: c++
	// 5. body = CompoundStmt: { CallExpr }

	if len(children) != 5 {
		panic(fmt.Sprintf("Expected 5 children in ForStmt, got %#v", children))
	}

	// TODO: The second child of a ForStmt appears to always be null.
	// Are there any cases where it is used?
	if children[1] != nil {
		panic("non-nil child 1 in ForStmt")
	}

	// If we have 2 and more initializations like
	// in operator for
	// for( a = 0, b = 0, c = 0; a < 5; a ++)
	switch children[0].(type) {
	case *ast.BinaryOperator:
		c := children[0].(*ast.BinaryOperator)
		if c.Operator == "," {
			// recursive action to code like that:
			// a = 0;
			// b = 0;
			// for(c = 0 ; a < 5 ; a++)
			//
			// Add after in preStmts
			beforeFor, newPre, newPost, err := transpileToStmt(c.Children[0], p)
			if err != nil {
				return nil, nil, nil, err
			}
			preStmts = append(preStmts, beforeFor)
			preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)
			// Now, in init only initialization
			children[0] = c.Children[len(c.Children)-1]
		}
	}

	init, newPre, newPost, err := transpileToStmt(children[0], p)
	if err != nil {
		return nil, nil, nil, err
	}

	preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

	post, newPre, newPost, err := transpileToStmt(children[3], p)
	if err != nil {
		return nil, nil, nil, err
	}

	preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

	body, newPre, newPost, err := transpileToBlockStmt(children[4], p)
	if err != nil {
		return nil, nil, nil, err
	}

	preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

	// The condition can be nil. This means an infinite loop and will be
	// rendered in Go as "for {".
	var condition goast.Expr
	if children[2] != nil {
		var conditionType string
		var newPre, newPost []goast.Stmt
		condition, conditionType, newPre, newPost, err = transpileToExpr(children[2], p)
		if err != nil {
			return nil, nil, nil, err
		}

		preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

		condition, err = types.CastExpr(p, condition, conditionType, "bool")
		p.AddMessage(ast.GenerateWarningOrErrorMessage(err, n, condition == nil))

		if condition == nil {
			condition = util.NewNil()
		}
	}

	return &goast.ForStmt{
		Init: init,
		Cond: condition,
		Post: post,
		Body: body,
	}, preStmts, postStmts, nil
}

func transpileWhileStmt(n *ast.WhileStmt, p *program.Program) (
	*goast.ForStmt, []goast.Stmt, []goast.Stmt, error) {
	preStmts := []goast.Stmt{}
	postStmts := []goast.Stmt{}

	// TODO: The first child of a WhileStmt appears to always be null.
	// Are there any cases where it is used?
	children := n.Children[1:]

	body, newPre, newPost, err := transpileToBlockStmt(children[1], p)
	if err != nil {
		return nil, nil, nil, err
	}

	preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

	condition, conditionType, newPre, newPost, err := transpileToExpr(children[0], p)
	if err != nil {
		return nil, nil, nil, err
	}

	preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

	cond, err := types.CastExpr(p, condition, conditionType, "bool")
	p.AddMessage(ast.GenerateWarningOrErrorMessage(err, n, cond == nil))

	if cond == nil {
		cond = util.NewNil()
	}

	return &goast.ForStmt{
		Cond: cond,
		Body: body,
	}, preStmts, postStmts, nil
}

func transpileDoStmt(n *ast.DoStmt, p *program.Program) (
	*goast.ForStmt, []goast.Stmt, []goast.Stmt, error) {
	preStmts := []goast.Stmt{}
	postStmts := []goast.Stmt{}
	children := n.Children

	body, newPre, newPost, err := transpileToBlockStmt(children[0], p)
	if err != nil {
		return nil, nil, nil, err
	}

	preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

	condition, conditionType, newPre, newPost, err := transpileToExpr(children[1], p)
	if err != nil {
		return nil, nil, nil, err
	}

	preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

	// Add IfStmt to the end of the loop to check the condition.
	x, err := types.CastExpr(p, condition, conditionType, "bool")
	p.AddMessage(ast.GenerateWarningOrErrorMessage(err, n, x == nil))

	if x == nil {
		x = util.NewNil()
	}

	body.List = append(body.List, &goast.IfStmt{
		Cond: &goast.UnaryExpr{
			Op: token.NOT,
			X:  x,
		},
		Body: &goast.BlockStmt{
			List: []goast.Stmt{&goast.BranchStmt{Tok: token.BREAK}},
		},
	})

	return &goast.ForStmt{
		Body: body,
	}, preStmts, postStmts, nil
}

func transpileContinueStmt(n *ast.ContinueStmt, p *program.Program) (*goast.BranchStmt, error) {
	return &goast.BranchStmt{
		Tok: token.CONTINUE,
	}, nil
}
<commit_msg>Infinite increment in for comma<commit_after>// This file contains functions for transpiling common branching and control
// flow, such as "if", "while", "do" and "for". The more complicated control
// flows like "switch" will be put into their own file of the same or sensible
// name.

package transpiler

import (
	"fmt"
	"go/token"

	"github.com/elliotchance/c2go/ast"
	"github.com/elliotchance/c2go/program"
	"github.com/elliotchance/c2go/types"
	"github.com/elliotchance/c2go/util"

	goast "go/ast"
)

func transpileIfStmt(n *ast.IfStmt, p *program.Program) (
	*goast.IfStmt, []goast.Stmt, []goast.Stmt, error) {
	preStmts := []goast.Stmt{}
	postStmts := []goast.Stmt{}
	children := n.Children

	// There is always 4 or 5 children in an IfStmt. For example:
	//
	//     if (i == 0) {
	//         return 0;
	//     } else {
	//         return 1;
	//     }
	//
	// 1. Not sure what this is for. This gets removed.
	// 2. Not sure what this is for.
	// 3. conditional = BinaryOperator: i == 0
	// 4. body = CompoundStmt: { return 0; }
	// 5. elseBody = CompoundStmt: { return 1; }
	//
	// elseBody will be nil if there is no else clause.

	// On linux I have seen only 4 children for an IfStmt with the same
	// definitions above, but missing the first argument. Since we don't
	// know what the first argument is for anyway we will just remove it on
	// Mac if necessary.
	if len(children) == 5 && children[0] != nil {
		panic("non-nil child 0 in IfStmt")
	}
	if len(children) == 5 {
		children = children[1:]
	}

	// From here on there must be 4 children.
	if len(children) != 4 {
		panic(fmt.Sprintf("Expected 4 children in IfStmt, got %#v", children))
	}

	// Maybe we will discover what the nil value is?
	if children[0] != nil {
		panic("non-nil child 0 in IfStmt")
	}

	conditional, conditionalType, newPre, newPost, err := transpileToExpr(children[1], p)
	if err != nil {
		return nil, nil, nil, err
	}

	// The condition in Go must always be a bool.
	boolCondition, err := types.CastExpr(p, conditional, conditionalType, "bool")
	p.AddMessage(ast.GenerateWarningOrErrorMessage(err, n, boolCondition == nil))

	if boolCondition == nil {
		boolCondition = util.NewNil()
	}

	body, newPre, newPost, err := transpileToBlockStmt(children[2], p)
	if err != nil {
		return nil, nil, nil, err
	}

	preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

	r := &goast.IfStmt{
		Cond: boolCondition,
		Body: body,
	}

	if children[3] != nil {
		elseBody, newPre, newPost, err := transpileToBlockStmt(children[3], p)
		if err != nil {
			return nil, nil, nil, err
		}

		preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

		r.Else = elseBody
	}

	return r, newPre, newPost, nil
}

func transpileForStmt(n *ast.ForStmt, p *program.Program) (
	*goast.ForStmt, []goast.Stmt, []goast.Stmt, error) {
	preStmts := []goast.Stmt{}
	postStmts := []goast.Stmt{}

	children := n.Children

	// There are always 5 children in a ForStmt, for example:
	//
	//     for ( c = 0 ; c < n ; c++ ) {
	//         doSomething();
	//     }
	//
	// 1. initExpression = BinaryStmt: c = 0
	// 2. Not sure what this is for, but it's always nil. There is a panic
	//    below in case we discover what it is used for (pun intended).
	// 3. conditionalExpression = BinaryStmt: c < n
	// 4. stepExpression = BinaryStmt: c++
	// 5. body = CompoundStmt: { CallExpr }

	if len(children) != 5 {
		panic(fmt.Sprintf("Expected 5 children in ForStmt, got %#v", children))
	}

	// TODO: The second child of a ForStmt appears to always be null.
	// Are there any cases where it is used?
	if children[1] != nil {
		panic("non-nil child 1 in ForStmt")
	}

	// If we have 2 and more initializations like
	// in operator for
	// for( a = 0, b = 0, c = 0; a < 5; a ++)
	switch children[0].(type) {
	case *ast.BinaryOperator:
		c := children[0].(*ast.BinaryOperator)
		if c.Operator == "," {
			// recursive action to code like that:
			// a = 0;
			// b = 0;
			// for(c = 0 ; a < 5 ; a++)
			//
			// Add after in preStmts
			beforeFor, newPre, newPost, err := transpileToStmt(c.Children[0], p)
			if err != nil {
				return nil, nil, nil, err
			}
			preStmts = append(preStmts, beforeFor)
			preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)
			// Now, in init only initialization
			children[0] = c.Children[len(c.Children)-1]
		}
	}

	init, newPre, newPost, err := transpileToStmt(children[0], p)
	if err != nil {
		return nil, nil, nil, err
	}

	preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

	// If we have 2 and more increments
	// in operator for
	// for( a = 0; a < 5; a ++, b++, c+=2)
	switch children[3].(type) {
	case *ast.BinaryOperator:
		c := children[3].(*ast.BinaryOperator)
		if c.Operator == "," {
			// recursive action to code like that:
			// a = 0;
			// b = 0;
			// for(a = 0 ; a < 5 ; c+=2){
			// body
			// a++;
			// b++
			// }
			//
			// Add after in preStmts
			var compound ast.CompoundStmt
			compound.Children = append(compound.Children, children[4])
			compound.Children = append(compound.Children, c.Children[0:len(c.Children)-1]...)
			children[4] = &compound
			children[3] = c.Children[len(c.Children)-1]
		}
	}

	post, newPre, newPost, err := transpileToStmt(children[3], p)
	if err != nil {
		return nil, nil, nil, err
	}

	preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

	body, newPre, newPost, err := transpileToBlockStmt(children[4], p)
	if err != nil {
		return nil, nil, nil, err
	}

	preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

	// The condition can be nil. This means an infinite loop and will be
	// rendered in Go as "for {".
	var condition goast.Expr
	if children[2] != nil {
		var conditionType string
		var newPre, newPost []goast.Stmt
		condition, conditionType, newPre, newPost, err = transpileToExpr(children[2], p)
		if err != nil {
			return nil, nil, nil, err
		}

		preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

		condition, err = types.CastExpr(p, condition, conditionType, "bool")
		p.AddMessage(ast.GenerateWarningOrErrorMessage(err, n, condition == nil))

		if condition == nil {
			condition = util.NewNil()
		}
	}

	return &goast.ForStmt{
		Init: init,
		Cond: condition,
		Post: post,
		Body: body,
	}, preStmts, postStmts, nil
}

func transpileWhileStmt(n *ast.WhileStmt, p *program.Program) (
	*goast.ForStmt, []goast.Stmt, []goast.Stmt, error) {
	preStmts := []goast.Stmt{}
	postStmts := []goast.Stmt{}

	// TODO: The first child of a WhileStmt appears to always be null.
	// Are there any cases where it is used?
	children := n.Children[1:]

	body, newPre, newPost, err := transpileToBlockStmt(children[1], p)
	if err != nil {
		return nil, nil, nil, err
	}

	preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

	condition, conditionType, newPre, newPost, err := transpileToExpr(children[0], p)
	if err != nil {
		return nil, nil, nil, err
	}

	preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

	cond, err := types.CastExpr(p, condition, conditionType, "bool")
	p.AddMessage(ast.GenerateWarningOrErrorMessage(err, n, cond == nil))

	if cond == nil {
		cond = util.NewNil()
	}

	return &goast.ForStmt{
		Cond: cond,
		Body: body,
	}, preStmts, postStmts, nil
}

func transpileDoStmt(n *ast.DoStmt, p *program.Program) (
	*goast.ForStmt, []goast.Stmt, []goast.Stmt, error) {
	preStmts := []goast.Stmt{}
	postStmts := []goast.Stmt{}
	children := n.Children

	body, newPre, newPost, err := transpileToBlockStmt(children[0], p)
	if err != nil {
		return nil, nil, nil, err
	}

	preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

	condition, conditionType, newPre, newPost, err := transpileToExpr(children[1], p)
	if err != nil {
		return nil, nil, nil, err
	}

	preStmts, postStmts = combinePreAndPostStmts(preStmts, postStmts, newPre, newPost)

	// Add IfStmt to the end of the loop to check the condition.
	x, err := types.CastExpr(p, condition, conditionType, "bool")
	p.AddMessage(ast.GenerateWarningOrErrorMessage(err, n, x == nil))

	if x == nil {
		x = util.NewNil()
	}

	body.List = append(body.List, &goast.IfStmt{
		Cond: &goast.UnaryExpr{
			Op: token.NOT,
			X:  x,
		},
		Body: &goast.BlockStmt{
			List: []goast.Stmt{&goast.BranchStmt{Tok: token.BREAK}},
		},
	})

	return &goast.ForStmt{
		Body: body,
	}, preStmts, postStmts, nil
}

func transpileContinueStmt(n *ast.ContinueStmt, p *program.Program) (*goast.BranchStmt, error) {
	return &goast.BranchStmt{
		Tok: token.CONTINUE,
	}, nil
}
<|endoftext|>


<commit_before>package controllers

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	jwt "github.com/dgrijalva/jwt-go"
	"github.com/prest/prest/config"
)

// Response representation
type Response struct {
	LoggedUser interface{} `json:"user_info"`
	Token      string      `json:"token"`
}

// RavensRequest representation
type RavensRequest struct {
	Type       string   `json:"type_of"`
	Subject    string   `json:"subject"`
	Recipients []string `json:"recipients"`
	Sender     string   `json:"sender"`
	SenderName string   `json:"sender_name"`
	Content    string   `json:"content"`
}

// AuthClaims JWT
type AuthClaims struct {
	jwt.StandardClaims
}

// Token for user
func Token(u User) (t string, err error) {
	// add expiry time in configuration (in minute format, so we support the maximum need)
	expireToken := time.Now().Add(time.Hour * 6).Unix()
	claims := AuthClaims{
		StandardClaims: jwt.StandardClaims{
			ExpiresAt: expireToken,
			Id:        strconv.Itoa(u.ID),
			IssuedAt:  time.Now().Unix(),
			Issuer:    strconv.Itoa(u.CustomerID),
		},
	}
	tok := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return tok.SignedString([]byte(config.PrestConf.JWTKey))
}

const unf = "user not found"

// Auth controller
func Auth(w http.ResponseWriter, r *http.Request) {
	user, password, ok := r.BasicAuth()
	if !ok {
		http.Error(w, unf, http.StatusBadRequest)
		return
	}
	loggedUser, err := basicPasswordCheck(strings.ToLower(user), password)
	if err != nil {
		http.Error(w, err.Error(), http.StatusUnauthorized)
		return
	}
	token, err := Token(loggedUser)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	resp := Response{
		LoggedUser: loggedUser,
		Token:      token,
	}
	err = json.NewEncoder(w).Encode(resp)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

// basicPasswordCheck
func basicPasswordCheck(user, password string) (obj interface{}, err error) {
	/**
	table name, fields (user and password) and encryption must be defined in
	the configuration file (toml)
	by default this endpoint will not be available, it is necessary to activate
	in the configuration file
	*/
	query := `SELECT * FROM users WHERE user=$1 AND password=$2 LIMIT 1`
	sc := config.Get.DBAdapter.Query(query, user, password)
	if sc.Err() != nil {
		err = sc.Err()
		return
	}
	n, err := sc.Scan(&obj)
	if err != nil {
		return
	}
	if n != 1 {
		err = fmt.Errorf(unf)
		return
	}
	return
}
<commit_msg>feat(controlers): implement basic authentication<commit_after>package controllers

import (
	"crypto/md5"
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	jwt "github.com/dgrijalva/jwt-go"
	"github.com/prest/prest/config"
)

// Response representation
type Response struct {
	LoggedUser interface{} `json:"user_info"`
	Token      string      `json:"token"`
}

// RavensRequest representation
type RavensRequest struct {
	Type       string   `json:"type_of"`
	Subject    string   `json:"subject"`
	Recipients []string `json:"recipients"`
	Sender     string   `json:"sender"`
	SenderName string   `json:"sender_name"`
	Content    string   `json:"content"`
}

// AuthClaims JWT
type AuthClaims struct {
	jwt.StandardClaims
}

type User struct {
	ID         int
	CustomerID int
}

// Token for user
func Token(u User) (t string, err error) {
	// add expiry time in configuration (in minute format, so we support the maximum need)
	expireToken := time.Now().Add(time.Hour * 6).Unix()
	claims := AuthClaims{
		StandardClaims: jwt.StandardClaims{
			ExpiresAt: expireToken,
			Id:        strconv.Itoa(u.ID),
			IssuedAt:  time.Now().Unix(),
			Issuer:    strconv.Itoa(u.CustomerID),
		},
	}
	tok := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return tok.SignedString([]byte(config.PrestConf.JWTKey))
}

const unf = "user not found"

// Auth controller
func Auth(w http.ResponseWriter, r *http.Request) {
	user, password, ok := r.BasicAuth()
	if !ok {
		http.Error(w, unf, http.StatusBadRequest)
		return
	}
	loggedUser, err := basicPasswordCheck(strings.ToLower(user), password)
	if err != nil {
		http.Error(w, err.Error(), http.StatusUnauthorized)
		return
	}
	token, err := Token(loggedUser)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	resp := Response{
		LoggedUser: loggedUser,
		Token:      token,
	}
	err = json.NewEncoder(w).Encode(resp)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

// basicPasswordCheck
func basicPasswordCheck(user, password string) (obj User, err error) {
	/**
	table name, fields (user and password) and encryption must be defined in
	the configuration file (toml)
	by default this endpoint will not be available, it is necessary to activate
	in the configuration file
	*/
	sc := config.PrestConf.Adapter.Query(getSelectQuery(), user, encrypt(password))
	if sc.Err() != nil {
		err = sc.Err()
		return
	}
	n, err := sc.Scan(&obj)
	if err != nil {
		return
	}
	if n != 1 {
		err = fmt.Errorf(unf)
	}

	return
}

// getSelectQuery create the query to authenticate the user
func getSelectQuery() (query string) {
	query = fmt.Sprintf(`SELECT * FROM %s WHERE %s=$1 AND %s=$2 LIMIT 1`, config.PrestConf.AuthTable, config.PrestConf.AuthUsername, config.PrestConf.AuthPassword)

	return
}

// encrypt will apply the encryption algorithm to the password
func encrypt(password string) (encrypted string) {
	switch config.PrestConf.AuthEncrypt {
	case "MD5":
		encrypted = fmt.Sprintf("%x", md5.Sum([]byte(password)))
		break

	case "SHA1":
		encrypted = fmt.Sprintf("%x", sha1.Sum([]byte(password)))
		break

	}
	return
}
<|endoftext|>


<commit_before>/* This file is part of VoltDB.
 * Copyright (C) 2008-2016 VoltDB Inc.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with VoltDB.  If not, see <http://www.gnu.org/licenses/>.
 */

package voltdbclient

import (
	"bytes"
	"database/sql/driver"
	"errors"
	"io"
	"net"
	"runtime"
	"sync"
	"time"
)

// start back pressure when this many bytes are queued for write
const maxQueuedBytes = 20

// connectionData are the values returned by a successful login.
type connectionData struct {
	hostId      int32
	connId      int64
	leaderAddr  int32
	buildString string
}

type nodeConn struct {
	connInfo      string
	reader        io.Reader
	connData      connectionData
	asyncsChannel chan voltResponse
	asyncs        map[int64]*voltAsyncResponse
	asyncsMutex   sync.Mutex
	nl            *networkListener
	nlwg          *sync.WaitGroup
	// used to write procedure invocation out to network.
	nw   *networkWriter
	nwCh chan<- *procedureInvocation
	// used to wait for writer to return on close.
	nwwg *sync.WaitGroup

	// queued bytes will be read/written by the main client thread and also
	// by the network writer thread.
	queuedBytes int
	qbMutex     sync.Mutex

	open bool
}

func newNodeConn(vc *VoltConn, ci string, reader io.Reader, writer io.Writer, connData connectionData) *nodeConn {
	var nc = new(nodeConn)

	nc.connInfo = ci
	nc.reader = reader
	nc.connData = connData
	nc.asyncsChannel = make(chan voltResponse)
	nc.asyncs = make(map[int64]*voltAsyncResponse)
	nc.asyncsMutex = sync.Mutex{}
	wg := sync.WaitGroup{}
	nc.nlwg = &wg
	nc.nl = newListener(vc, ci, reader, nc.nlwg)
	// The buffer won't be allocated up front, so it's ok to make this big.
	// In practice the buffer will be limited by back pressure
	ch := make(chan *procedureInvocation, 1000)
	nc.nwCh = ch
	wg = sync.WaitGroup{}
	nc.nwwg = &wg
	nc.nw = newNetworkWriter(writer, ch, nc.nwwg)

	nc.open = true
	nc.queuedBytes = 0
	nc.qbMutex = sync.Mutex{}

	nc.nl.start()
	nc.nw.start()
	go nc.processAsyncs()
	return nc
}

func (nc *nodeConn) close() (err error) {
	if !nc.open {
		return
	}

	// stop the network writer and reader
	nc.nw.stop()
	nc.nl.stop()

	// close the tcp conn, will unblock the listener.
	if nc.reader != nil {
		tcpConn := nc.reader.(*net.TCPConn)
		err = tcpConn.Close()
	}

	// network reader and writer should return.
	nc.nwwg.Wait()
	nc.nlwg.Wait()

	nc.open = false
	return err
}

func (nc *nodeConn) exec(pi *procedureInvocation) (driver.Result, error) {
	if !nc.open {
		return nil, errors.New("Connection is closed")
	}
	c := nc.nl.registerRequest(nc, pi)
	nc.incrementQueuedBytes(pi.getLen())
	nc.nwCh <- pi

	select {
	case resp := <-c:
		rslt := resp.(VoltResult)
		if err := rslt.getError(); err != nil {
			return nil, err
		}
		return rslt, nil
	case <-time.After(pi.timeout):
		return nil, errors.New("timeout")
	}
}

func (nc *nodeConn) execAsync(resCons AsyncResponseConsumer, pi *procedureInvocation) error {
	if !nc.open {
		return errors.New("Connection is closed")
	}
	c := nc.nl.registerRequest(nc, pi)
	vasr := newVoltAsyncResponse(nc, pi.handle, c, false, resCons)
	nc.registerAsync(pi.handle, vasr)
	nc.incrementQueuedBytes(pi.getLen())
	nc.nwCh <- pi
	return nil
}

func (nc *nodeConn) query(pi *procedureInvocation) (driver.Rows, error) {
	if !nc.open {
		return nil, errors.New("Connection is closed")
	}
	c := nc.nl.registerRequest(nc, pi)
	nc.incrementQueuedBytes(pi.getLen())
	nc.nwCh <- pi

	select {
	case resp := <-c:
		rows := resp.(VoltRows)
		if err := rows.getError(); err != nil {
			return nil, err
		}
		return rows, nil
	case <-time.After(pi.timeout):
		// TODO: make an error type for timeout
		return nil, errors.New("timeout")
	}
}

// QueryAsync executes a query asynchronously.  The invoking thread will block
// until the query is sent over the network to the server.  The eventual
// response will be handled by the given AsyncResponseConsumer, this processing
// happens in the 'response' thread.
func (nc *nodeConn) queryAsync(rowsCons AsyncResponseConsumer, pi *procedureInvocation) error {
	if !nc.open {
		return errors.New("Connection is closed")
	}
	c := nc.nl.registerRequest(nc, pi)
	vasr := newVoltAsyncResponse(nc, pi.handle, c, true, rowsCons)
	nc.registerAsync(pi.handle, vasr)
	nc.incrementQueuedBytes(pi.getLen())
	nc.nwCh <- pi
	return nil
}

func (nc *nodeConn) incrementQueuedBytes(numBytes int) {
	nc.qbMutex.Lock()
	nc.queuedBytes += numBytes
	nc.qbMutex.Unlock()
}

func (nc *nodeConn) decrementQueuedBytes(numBytes int) {
	nc.qbMutex.Lock()
	nc.queuedBytes -= numBytes
	nc.qbMutex.Unlock()
}

func (nc *nodeConn) getQueuedBytes() int {
	var qb int
	nc.qbMutex.Lock()
	qb = nc.queuedBytes
	nc.qbMutex.Unlock()
	return qb
}

func (nc *nodeConn) drain() {
	var numAsyncs int
	for {
		nc.asyncsMutex.Lock()
		numAsyncs = len(nc.asyncs)
		nc.asyncsMutex.Unlock()
		if numAsyncs == 0 {
			break
		}
		runtime.Gosched()
	}
}

func (nc *nodeConn) processAsyncs() {
	for {
		resp := <-nc.asyncsChannel
		handle := resp.getHandle()
		nc.asyncsMutex.Lock()
		async := nc.asyncs[handle]
		nc.asyncsMutex.Unlock()

		if async.isQuery() {
			vrows := resp.(VoltRows)
			if err := vrows.getError(); err != nil {
				async.getArc().ConsumeError(err)
			} else {
				async.getArc().ConsumeRows(vrows)

			}
			nc.removeAsync(handle)
			continue
		} else {
			vrslt := resp.(VoltResult)
			if err := vrslt.getError(); err != nil {
				async.getArc().ConsumeError(err)
			} else {
				async.getArc().ConsumeResult(vrslt)
			}
			nc.removeAsync(handle)
			continue
		}
	}
}

func (nc *nodeConn) hasBP() bool {
	return nc.getQueuedBytes() >= maxQueuedBytes
}

func (nc *nodeConn) registerAsync(handle int64, vasr *voltAsyncResponse) {
	nc.asyncsMutex.Lock()
	nc.asyncs[handle] = vasr
	nc.asyncsMutex.Unlock()
	go func() {
		nc.asyncsChannel <- <-vasr.channel()
	}()
}

func (nc *nodeConn) removeAsync(han int64) {
	nc.asyncsMutex.Lock()
	delete(nc.asyncs, han)
	nc.asyncsMutex.Unlock()
}

func writeLoginMessage(writer io.Writer, buf *bytes.Buffer) {
	// length includes protocol version.
	length := buf.Len() + 2
	var netmsg bytes.Buffer
	writeInt(&netmsg, int32(length))
	writeProtoVersion(&netmsg)
	writePasswordHashVersion(&netmsg)
	// 1 copy + 1 n/w write benchmarks faster than 2 n/w writes.
	io.Copy(&netmsg, buf)
	io.Copy(writer, &netmsg)
}

func readLoginResponse(reader io.Reader) (*connectionData, error) {
	buf, err := readMessage(reader)
	if err != nil {
		return nil, err
	}
	connData, err := deserializeLoginResponse(buf)
	return connData, err
}

// AsyncResponseConsumer is a type that consumes responses from asynchronous
// Queries and Execs.
// In the VoltDB go client, asynchronous requests are continuously processed by
// one or more goroutines executing in the background.  When a response from
// the server is received for an asynchronous request, one of the methods in
// this interface is invoked.  An instance of AyncResponseConsumer is passed
// when an asynchronous request is made, this instance will process the
// response for that request.
type AsyncResponseConsumer interface {

	// This method is invoked when an error is returned by an async Query
	// or an Exec.
	ConsumeError(error)
	// This method is invoked when a Result is returned by an async Exec.
	ConsumeResult(driver.Result)
	// This method is invoked when Rows is returned by an async Query.
	ConsumeRows(driver.Rows)
}

type voltAsyncResponse struct {
	conn *nodeConn
	han  int64
	ch   <-chan voltResponse
	isQ  bool
	arc  AsyncResponseConsumer
}

func newVoltAsyncResponse(conn *nodeConn, han int64, ch <-chan voltResponse, isQuery bool, arc AsyncResponseConsumer) *voltAsyncResponse {
	var vasr = new(voltAsyncResponse)
	vasr.conn = conn
	vasr.han = han
	vasr.ch = ch
	vasr.isQ = isQuery
	vasr.arc = arc
	return vasr
}

func (vasr *voltAsyncResponse) getArc() AsyncResponseConsumer {
	return vasr.arc
}

func (vasr *voltAsyncResponse) channel() <-chan voltResponse {
	return vasr.ch
}

func (vasr *voltAsyncResponse) handle() int64 {
	return vasr.han
}

func (vasr *voltAsyncResponse) isQuery() bool {
	return vasr.isQ
}

// Null Value type
type nullValue struct {
	colType int8
}

func (nv *nullValue) getColType() int8 {
	return nv.colType
}
<commit_msg>ENG-10558 minor update<commit_after>/* This file is part of VoltDB.
 * Copyright (C) 2008-2016 VoltDB Inc.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with VoltDB.  If not, see <http://www.gnu.org/licenses/>.
 */

package voltdbclient

import (
	"bytes"
	"database/sql/driver"
	"errors"
	"io"
	"net"
	"runtime"
	"sync"
	"time"
)

// start back pressure when this many bytes are queued for write
const maxQueuedBytes = 262144

// connectionData are the values returned by a successful login.
type connectionData struct {
	hostId      int32
	connId      int64
	leaderAddr  int32
	buildString string
}

type nodeConn struct {
	connInfo      string
	reader        io.Reader
	connData      connectionData
	asyncsChannel chan voltResponse
	asyncs        map[int64]*voltAsyncResponse
	asyncsMutex   sync.Mutex
	nl            *networkListener
	nlwg          *sync.WaitGroup
	// used to write procedure invocation out to network.
	nw   *networkWriter
	nwCh chan<- *procedureInvocation
	// used to wait for writer to return on close.
	nwwg *sync.WaitGroup

	// queued bytes will be read/written by the main client thread and also
	// by the network writer thread.
	queuedBytes int
	qbMutex     sync.Mutex

	open bool
}

func newNodeConn(vc *VoltConn, ci string, reader io.Reader, writer io.Writer, connData connectionData) *nodeConn {
	var nc = new(nodeConn)

	nc.connInfo = ci
	nc.reader = reader
	nc.connData = connData
	nc.asyncsChannel = make(chan voltResponse)
	nc.asyncs = make(map[int64]*voltAsyncResponse)
	nc.asyncsMutex = sync.Mutex{}
	wg := sync.WaitGroup{}
	nc.nlwg = &wg
	nc.nl = newListener(vc, ci, reader, nc.nlwg)
	// The buffer won't be allocated up front, so it's ok to make this big.
	// In practice the buffer will be limited by back pressure
	ch := make(chan *procedureInvocation, 1000)
	nc.nwCh = ch
	wg = sync.WaitGroup{}
	nc.nwwg = &wg
	nc.nw = newNetworkWriter(writer, ch, nc.nwwg)

	nc.open = true
	nc.queuedBytes = 0
	nc.qbMutex = sync.Mutex{}

	nc.nl.start()
	nc.nw.start()
	go nc.processAsyncs()
	return nc
}

func (nc *nodeConn) close() (err error) {
	if !nc.open {
		return
	}

	// stop the network writer and reader
	nc.nw.stop()
	nc.nl.stop()

	// close the tcp conn, will unblock the listener.
	if nc.reader != nil {
		tcpConn := nc.reader.(*net.TCPConn)
		err = tcpConn.Close()
	}

	// network reader and writer should return.
	nc.nwwg.Wait()
	nc.nlwg.Wait()

	nc.open = false
	return err
}

func (nc *nodeConn) exec(pi *procedureInvocation) (driver.Result, error) {
	if !nc.open {
		return nil, errors.New("Connection is closed")
	}
	c := nc.nl.registerRequest(nc, pi)
	nc.incrementQueuedBytes(pi.getLen())
	nc.nwCh <- pi

	select {
	case resp := <-c:
		rslt := resp.(VoltResult)
		if err := rslt.getError(); err != nil {
			return nil, err
		}
		return rslt, nil
	case <-time.After(pi.timeout):
		return nil, errors.New("timeout")
	}
}

func (nc *nodeConn) execAsync(resCons AsyncResponseConsumer, pi *procedureInvocation) error {
	if !nc.open {
		return errors.New("Connection is closed")
	}
	c := nc.nl.registerRequest(nc, pi)
	vasr := newVoltAsyncResponse(nc, pi.handle, c, false, resCons)
	nc.registerAsync(pi.handle, vasr)
	nc.incrementQueuedBytes(pi.getLen())
	nc.nwCh <- pi
	return nil
}

func (nc *nodeConn) query(pi *procedureInvocation) (driver.Rows, error) {
	if !nc.open {
		return nil, errors.New("Connection is closed")
	}
	c := nc.nl.registerRequest(nc, pi)
	nc.incrementQueuedBytes(pi.getLen())
	nc.nwCh <- pi

	select {
	case resp := <-c:
		rows := resp.(VoltRows)
		if err := rows.getError(); err != nil {
			return nil, err
		}
		return rows, nil
	case <-time.After(pi.timeout):
		// TODO: make an error type for timeout
		return nil, errors.New("timeout")
	}
}

// QueryAsync executes a query asynchronously.  The invoking thread will block
// until the query is sent over the network to the server.  The eventual
// response will be handled by the given AsyncResponseConsumer, this processing
// happens in the 'response' thread.
func (nc *nodeConn) queryAsync(rowsCons AsyncResponseConsumer, pi *procedureInvocation) error {
	if !nc.open {
		return errors.New("Connection is closed")
	}
	c := nc.nl.registerRequest(nc, pi)
	vasr := newVoltAsyncResponse(nc, pi.handle, c, true, rowsCons)
	nc.registerAsync(pi.handle, vasr)
	nc.incrementQueuedBytes(pi.getLen())
	nc.nwCh <- pi
	return nil
}

func (nc *nodeConn) incrementQueuedBytes(numBytes int) {
	nc.qbMutex.Lock()
	nc.queuedBytes += numBytes
	nc.qbMutex.Unlock()
}

func (nc *nodeConn) decrementQueuedBytes(numBytes int) {
	nc.qbMutex.Lock()
	nc.queuedBytes -= numBytes
	nc.qbMutex.Unlock()
}

func (nc *nodeConn) getQueuedBytes() int {
	var qb int
	nc.qbMutex.Lock()
	qb = nc.queuedBytes
	nc.qbMutex.Unlock()
	return qb
}

func (nc *nodeConn) drain() {
	var numAsyncs int
	for {
		nc.asyncsMutex.Lock()
		numAsyncs = len(nc.asyncs)
		nc.asyncsMutex.Unlock()
		if numAsyncs == 0 {
			break
		}
		runtime.Gosched()
	}
}

func (nc *nodeConn) processAsyncs() {
	for {
		resp := <-nc.asyncsChannel
		handle := resp.getHandle()
		nc.asyncsMutex.Lock()
		async := nc.asyncs[handle]
		nc.asyncsMutex.Unlock()

		if async.isQuery() {
			vrows := resp.(VoltRows)
			if err := vrows.getError(); err != nil {
				async.getArc().ConsumeError(err)
			} else {
				async.getArc().ConsumeRows(vrows)

			}
			nc.removeAsync(handle)
			continue
		} else {
			vrslt := resp.(VoltResult)
			if err := vrslt.getError(); err != nil {
				async.getArc().ConsumeError(err)
			} else {
				async.getArc().ConsumeResult(vrslt)
			}
			nc.removeAsync(handle)
			continue
		}
	}
}

func (nc *nodeConn) hasBP() bool {
	return nc.getQueuedBytes() >= maxQueuedBytes
}

func (nc *nodeConn) registerAsync(handle int64, vasr *voltAsyncResponse) {
	nc.asyncsMutex.Lock()
	nc.asyncs[handle] = vasr
	nc.asyncsMutex.Unlock()
	go func() {
		nc.asyncsChannel <- <-vasr.channel()
	}()
}

func (nc *nodeConn) removeAsync(han int64) {
	nc.asyncsMutex.Lock()
	delete(nc.asyncs, han)
	nc.asyncsMutex.Unlock()
}

func writeLoginMessage(writer io.Writer, buf *bytes.Buffer) {
	// length includes protocol version.
	length := buf.Len() + 2
	var netmsg bytes.Buffer
	writeInt(&netmsg, int32(length))
	writeProtoVersion(&netmsg)
	writePasswordHashVersion(&netmsg)
	// 1 copy + 1 n/w write benchmarks faster than 2 n/w writes.
	io.Copy(&netmsg, buf)
	io.Copy(writer, &netmsg)
}

func readLoginResponse(reader io.Reader) (*connectionData, error) {
	buf, err := readMessage(reader)
	if err != nil {
		return nil, err
	}
	connData, err := deserializeLoginResponse(buf)
	return connData, err
}

// AsyncResponseConsumer is a type that consumes responses from asynchronous
// Queries and Execs.
// In the VoltDB go client, asynchronous requests are continuously processed by
// one or more goroutines executing in the background.  When a response from
// the server is received for an asynchronous request, one of the methods in
// this interface is invoked.  An instance of AyncResponseConsumer is passed
// when an asynchronous request is made, this instance will process the
// response for that request.
type AsyncResponseConsumer interface {

	// This method is invoked when an error is returned by an async Query
	// or an Exec.
	ConsumeError(error)
	// This method is invoked when a Result is returned by an async Exec.
	ConsumeResult(driver.Result)
	// This method is invoked when Rows is returned by an async Query.
	ConsumeRows(driver.Rows)
}

type voltAsyncResponse struct {
	conn *nodeConn
	han  int64
	ch   <-chan voltResponse
	isQ  bool
	arc  AsyncResponseConsumer
}

func newVoltAsyncResponse(conn *nodeConn, han int64, ch <-chan voltResponse, isQuery bool, arc AsyncResponseConsumer) *voltAsyncResponse {
	var vasr = new(voltAsyncResponse)
	vasr.conn = conn
	vasr.han = han
	vasr.ch = ch
	vasr.isQ = isQuery
	vasr.arc = arc
	return vasr
}

func (vasr *voltAsyncResponse) getArc() AsyncResponseConsumer {
	return vasr.arc
}

func (vasr *voltAsyncResponse) channel() <-chan voltResponse {
	return vasr.ch
}

func (vasr *voltAsyncResponse) handle() int64 {
	return vasr.han
}

func (vasr *voltAsyncResponse) isQuery() bool {
	return vasr.isQ
}

// Null Value type
type nullValue struct {
	colType int8
}

func (nv *nullValue) getColType() int8 {
	return nv.colType
}
<|endoftext|>


<commit_before>/* This file is part of VoltDB.
 * Copyright (C) 2008-2018 VoltDB Inc.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with VoltDB.  If not, see <http://www.gnu.org/licenses/>.
 */

package voltdbclient

import (
	"database/sql/driver"
	"errors"
	"fmt"
	"time"
)

var errTimeoutExecutingQuery = errors.New("voltdbclient: Timed out executing query")

// Exec executes a query that doesn't return rows, such as an INSERT or UPDATE.
// Exec is available on both VoltConn and on VoltStatement.
// Uses DefaultQueryTimeout.
func (c *Conn) Exec(query string, args []driver.Value) (driver.Result, error) {
	return c.ExecTimeout(query, args, DefaultQueryTimeout)
}

// ExecTimeout executes a query that doesn't return rows, such as an INSERT or
// UPDATE. ExecTimeout is available on both VoltConn and on VoltStatement.
// Specifies a duration for timeout.
func (c *Conn) ExecTimeout(query string, args []driver.Value, timeout time.Duration) (driver.Result, error) {
	responseCh := make(chan voltResponse, 1)
	pi := newSyncProcedureInvocation(c.getNextHandle(), false, query, args, responseCh, timeout)
	_, err := c.submit(pi)
	if err != nil {
		return nil, err
	}
	tm := time.NewTimer(pi.timeout)
	defer tm.Stop()
	sec := time.NewTicker(time.Second)
	defer sec.Stop()
	for {
		if pi.conn.isClosed() {
			return nil, VoltError{voltResponse: voltResponseInfo{status: ConnectionTimeout, clusterRoundTripTime: -1},
				error: fmt.Errorf("%s: writing on a closed node connection",
					pi.conn.connInfo)}
		}
		select {
		case <-sec.C:
			if err := pi.conn.Ping(); err != nil {
				pi.conn.markClosed()
				return nil, VoltError{voltResponse: voltResponseInfo{status: ConnectionTimeout,
					clusterRoundTripTime: -1}, error: err}
			}
		case resp := <-pi.responseCh:
			switch resp.(type) {
			case VoltResult:
				return resp.(VoltResult), nil
			case VoltError:
				return nil, resp.(VoltError)
			default:
				panic("unexpected response type")
			}
		case <-tm.C:
			return nil, VoltError{voltResponse: voltResponseInfo{status: ConnectionTimeout, clusterRoundTripTime: -1}, error: errors.New("timeout")}
		}
	}
}

// ExecAsync is analogous to Exec but is run asynchronously.  That is, an
// invocation of this method blocks only until a request is sent to the VoltDB
// server.  Uses DefaultQueryTimeout.
func (c *Conn) ExecAsync(resCons AsyncResponseConsumer, query string, args []driver.Value) {
	c.ExecAsyncTimeout(resCons, query, args, DefaultQueryTimeout)
}

// ExecAsyncTimeout is analogous to Exec but is run asynchronously.  That is, an
// invocation of this method blocks only until a request is sent to the VoltDB
// server.  Specifies a duration for timeout.
func (c *Conn) ExecAsyncTimeout(resCons AsyncResponseConsumer, query string, args []driver.Value, timeout time.Duration) error {
	pi := newAsyncProcedureInvocation(c.getNextHandle(), false, query, args, timeout, resCons)
	_, err := c.submit(pi)
	return err
}

// Prepare creates a prepared statement for later queries or executions.
// The Statement returned by Prepare is bound to this VoltConn.
func (c *Conn) Prepare(query string) (driver.Stmt, error) {
	stmt := newVoltStatement(c, query)
	return *stmt, nil
}

// Query executes a query that returns rows, typically a SELECT. The args are
// for any placeholder parameters in the query.
// Uses DefaultQueryTimeout.
func (c *Conn) Query(query string, args []driver.Value) (driver.Rows, error) {
	return c.QueryTimeout(query, args, DefaultQueryTimeout)
}

// QueryTimeout executes a query that returns rows, typically a SELECT. The args
// are for any placeholder parameters in the query.
// Specifies a duration for timeout.
func (c *Conn) QueryTimeout(query string, args []driver.Value, timeout time.Duration) (driver.Rows, error) {
	responseCh := make(chan voltResponse, 1)
	pi := newSyncProcedureInvocation(c.getNextHandle(), true, query, args, responseCh, timeout)
	_, err := c.submit(pi)
	if err != nil {
		return nil, err
	}
	tm := time.NewTimer(pi.timeout)
	defer tm.Stop()
	sec := time.NewTicker(time.Second)
	defer sec.Stop()
	for {
		if pi.conn.isClosed() {
			return nil, VoltError{voltResponse: voltResponseInfo{status: ConnectionTimeout, clusterRoundTripTime: -1},
				error: fmt.Errorf("%s: writing on a closed node connection",
					pi.conn.connInfo)}
		}
		select {
		case <-sec.C:
			if err := pi.conn.Ping(); err != nil {
				pi.conn.markClosed()
				return nil, VoltError{voltResponse: voltResponseInfo{status: ConnectionTimeout, clusterRoundTripTime: -1},
					error: err}
			}
		case resp := <-pi.responseCh:
			switch e := resp.(type) {
			case VoltRows:
				return e, nil
			case VoltError:
				return nil, e
			default:
				panic("unexpected response type")
			}
		case <-tm.C:
			return nil, VoltError{voltResponse: voltResponseInfo{status: ConnectionTimeout, clusterRoundTripTime: -1}, error: errTimeoutExecutingQuery}
		}
	}

}

// QueryAsync executes a query asynchronously.  The invoking thread will block
// until the query is sent over the network to the server.  The eventual
// response will be handled by the given AsyncResponseConsumer, this processing
// happens in the 'response' thread.  Uses DefaultQueryTimeout.
func (c *Conn) QueryAsync(rowsCons AsyncResponseConsumer, query string, args []driver.Value) {
	c.QueryAsyncTimeout(rowsCons, query, args, DefaultQueryTimeout)
}

// QueryAsyncTimeout executes a query asynchronously.  The invoking thread will
// block until the query is sent over the network to the server.  The eventual
// response will be handled by the given AsyncResponseConsumer, this processing
// happens in the 'response' thread.  Specifies a duration for timeout.
func (c *Conn) QueryAsyncTimeout(rowsCons AsyncResponseConsumer, query string, args []driver.Value, timeout time.Duration) error {
	pi := newAsyncProcedureInvocation(c.getNextHandle(), true, query, args, timeout, rowsCons)
	_, err := c.submit(pi)
	return err
}
<commit_msg>update documentation<commit_after>/* This file is part of VoltDB.
 * Copyright (C) 2008-2018 VoltDB Inc.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with VoltDB.  If not, see <http://www.gnu.org/licenses/>.
 */

package voltdbclient

import (
	"database/sql/driver"
	"errors"
	"fmt"
	"time"
)

var errTimeoutExecutingQuery = errors.New("voltdbclient: Timed out executing query")

// Exec executes a query that doesn't return rows, such as an INSERT or UPDATE.
// Exec is available on both VoltConn and on VoltStatement.
// Uses DefaultQueryTimeout.
func (c *Conn) Exec(query string, args []driver.Value) (driver.Result, error) {
	return c.ExecTimeout(query, args, DefaultQueryTimeout)
}

// ExecTimeout executes a query that doesn't return rows, such as an INSERT or
// UPDATE. ExecTimeout is available on both VoltConn and on VoltStatement.
// Specifies a duration for timeout.
func (c *Conn) ExecTimeout(query string, args []driver.Value, timeout time.Duration) (driver.Result, error) {
	// Works similar to QueryTimeout but send exec queries instead.
	responseCh := make(chan voltResponse, 1)
	pi := newSyncProcedureInvocation(c.getNextHandle(), false, query, args, responseCh, timeout)
	_, err := c.submit(pi)
	if err != nil {
		return nil, err
	}
	tm := time.NewTimer(pi.timeout)
	defer tm.Stop()
	sec := time.NewTicker(time.Second)
	defer sec.Stop()
	for {
		if pi.conn.isClosed() {
			return nil, VoltError{voltResponse: voltResponseInfo{status: ConnectionTimeout, clusterRoundTripTime: -1},
				error: fmt.Errorf("%s: writing on a closed node connection",
					pi.conn.connInfo)}
		}
		select {
		case <-sec.C:
			if err := pi.conn.Ping(); err != nil {
				pi.conn.markClosed()
				return nil, VoltError{voltResponse: voltResponseInfo{status: ConnectionTimeout,
					clusterRoundTripTime: -1}, error: err}
			}
		case resp := <-pi.responseCh:
			switch resp.(type) {
			case VoltResult:
				return resp.(VoltResult), nil
			case VoltError:
				return nil, resp.(VoltError)
			default:
				panic("unexpected response type")
			}
		case <-tm.C:
			return nil, VoltError{voltResponse: voltResponseInfo{status: ConnectionTimeout, clusterRoundTripTime: -1}, error: errors.New("timeout")}
		}
	}
}

// ExecAsync is analogous to Exec but is run asynchronously.  That is, an
// invocation of this method blocks only until a request is sent to the VoltDB
// server.  Uses DefaultQueryTimeout.
func (c *Conn) ExecAsync(resCons AsyncResponseConsumer, query string, args []driver.Value) {
	c.ExecAsyncTimeout(resCons, query, args, DefaultQueryTimeout)
}

// ExecAsyncTimeout is analogous to Exec but is run asynchronously.  That is, an
// invocation of this method blocks only until a request is sent to the VoltDB
// server.  Specifies a duration for timeout.
func (c *Conn) ExecAsyncTimeout(resCons AsyncResponseConsumer, query string, args []driver.Value, timeout time.Duration) error {
	pi := newAsyncProcedureInvocation(c.getNextHandle(), false, query, args, timeout, resCons)
	_, err := c.submit(pi)
	return err
}

// Prepare creates a prepared statement for later queries or executions.
// The Statement returned by Prepare is bound to this VoltConn.
func (c *Conn) Prepare(query string) (driver.Stmt, error) {
	stmt := newVoltStatement(c, query)
	return *stmt, nil
}

// Query executes a query that returns rows, typically a SELECT. The args are
// for any placeholder parameters in the query.
// Uses DefaultQueryTimeout.
func (c *Conn) Query(query string, args []driver.Value) (driver.Rows, error) {
	return c.QueryTimeout(query, args, DefaultQueryTimeout)
}

// QueryTimeout executes a query that returns rows, typically a SELECT. The args
// are for any placeholder parameters in the query.
// Specifies a duration for timeout.
func (c *Conn) QueryTimeout(query string, args []driver.Value, timeout time.Duration) (driver.Rows, error) {
	// This is sync sql query to voltdb. In reality there is no sync implementation
	// of querying voltdb in this client.
	//
	// What happens here is, we invoke the procedure invocation and wait until we
	// receive a response on the procedureInvocation channel.
	// The call to c.submit ensures that the procedure invocation request did hit a
	// voltdb host.
	//
	// To account for the case the current connection that is serving the procedure
	// invocation is lost. We send ping request every 1 second.If we detect the
	// connection is down, we mark the connection as closed so that it will no
	// longer be used(NOTE: This is useful in cluster mode).
	responseCh := make(chan voltResponse, 1)
	pi := newSyncProcedureInvocation(c.getNextHandle(), true, query, args, responseCh, timeout)
	_, err := c.submit(pi)
	if err != nil {
		return nil, err
	}
	tm := time.NewTimer(pi.timeout)
	defer tm.Stop()
	sec := time.NewTicker(time.Second)
	defer sec.Stop()
	for {
		if pi.conn.isClosed() {
			return nil, VoltError{voltResponse: voltResponseInfo{status: ConnectionTimeout, clusterRoundTripTime: -1},
				error: fmt.Errorf("%s: writing on a closed node connection",
					pi.conn.connInfo)}
		}
		select {
		case <-sec.C:
			if err := pi.conn.Ping(); err != nil {
				pi.conn.markClosed()
				return nil, VoltError{voltResponse: voltResponseInfo{status: ConnectionTimeout, clusterRoundTripTime: -1},
					error: err}
			}
		case resp := <-pi.responseCh:
			switch e := resp.(type) {
			case VoltRows:
				return e, nil
			case VoltError:
				return nil, e
			default:
				panic("unexpected response type")
			}
		case <-tm.C:
			return nil, VoltError{voltResponse: voltResponseInfo{status: ConnectionTimeout, clusterRoundTripTime: -1}, error: errTimeoutExecutingQuery}
		}
	}

}

// QueryAsync executes a query asynchronously.  The invoking thread will block
// until the query is sent over the network to the server.  The eventual
// response will be handled by the given AsyncResponseConsumer, this processing
// happens in the 'response' thread.  Uses DefaultQueryTimeout.
func (c *Conn) QueryAsync(rowsCons AsyncResponseConsumer, query string, args []driver.Value) {
	c.QueryAsyncTimeout(rowsCons, query, args, DefaultQueryTimeout)
}

// QueryAsyncTimeout executes a query asynchronously.  The invoking thread will
// block until the query is sent over the network to the server.  The eventual
// response will be handled by the given AsyncResponseConsumer, this processing
// happens in the 'response' thread.  Specifies a duration for timeout.
func (c *Conn) QueryAsyncTimeout(rowsCons AsyncResponseConsumer, query string, args []driver.Value, timeout time.Duration) error {
	pi := newAsyncProcedureInvocation(c.getNextHandle(), true, query, args, timeout, rowsCons)
	_, err := c.submit(pi)
	return err
}
<|endoftext|>


<commit_before>package vxlan

import (
	//"fmt"
	//"strings"
	//"time"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/go-plugins-helpers/network"
	//"github.com/samalba/dockerclient"
	"github.com/vishvananda/netlink"
)

type Driver struct {
	network.Driver
	networks map[string]*NetworkState
}

// NetworkState is filled in at network creation time
// it contains state that we wish to keep for each network
type NetworkState struct {
	Bridge *netlink.Bridge
	VXLan  *netlink.Vxlan
}

func NewDriver() (*Driver, error) {
	d := &Driver{
		networks: make(map[string]*NetworkState),
	}
	return d, nil
}

func (d *Driver) CreateNetwork(r *network.CreateNetworkRequest) error {
	log.Debugf("Create network request: %+v", r)

	//if r.Options == nil {
	//	return "", fmt.Errorf("No options provided")
	//}

	vxlanName := "vxlan42"
	vxlanID := 42
	//if r.Options["vxlanID"] != nil {
	//	vxlanID = r.Options["vxlanID"]
	//}

	bridgeName := "br_vxlan42"

	bridge := &netlink.Bridge{
		LinkAttrs: netlink.LinkAttrs{Name: bridgeName},
	}
	netlink.LinkAdd(bridge)

	vxlan := &netlink.Vxlan{
		LinkAttrs: netlink.LinkAttrs{Name: vxlanName},
		VxlanId:   vxlanID,
	}
	netlink.LinkAdd(vxlan)

	ns := &NetworkState{
		VXLan:  vxlan,
		Bridge: bridge,
	}
	d.networks[r.NetworkID] = ns

	netlink.LinkSetMaster(vxlan, bridge)

	return nil
}
<commit_msg>add delete<commit_after>package vxlan

import (
	//"fmt"
	//"strings"
	//"time"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/go-plugins-helpers/network"
	//"github.com/samalba/dockerclient"
	"github.com/vishvananda/netlink"
)

type Driver struct {
	network.Driver
	networks map[string]*NetworkState
}

// NetworkState is filled in at network creation time
// it contains state that we wish to keep for each network
type NetworkState struct {
	Bridge *netlink.Bridge
	VXLan  *netlink.Vxlan
}

func NewDriver() (*Driver, error) {
	d := &Driver{
		networks: make(map[string]*NetworkState),
	}
	return d, nil
}

func (d *Driver) CreateNetwork(r *network.CreateNetworkRequest) error {
	log.Debugf("Create network request: %+v", r)

	name := r.NetworkID

	//if r.Options == nil {
	//	return "", fmt.Errorf("No options provided")
	//}

	vxlanName := "vx_42" // + name
	vxlanID := 42
	//if r.Options["vxlanID"] != nil {
	//	vxlanID = r.Options["vxlanID"]
	//}

	bridgeName := "br_42" // + name

	bridge := &netlink.Bridge{
		LinkAttrs: netlink.LinkAttrs{Name: bridgeName},
	}
	netlink.LinkAdd(bridge)

	vxlan := &netlink.Vxlan{
		LinkAttrs: netlink.LinkAttrs{Name: vxlanName},
		VxlanId:   vxlanID,
	}
	netlink.LinkAdd(vxlan)

	ns := &NetworkState{
		VXLan:  vxlan,
		Bridge: bridge,
	}
	d.networks[name] = ns

	netlink.LinkSetMaster(vxlan, bridge)

	return nil
}

func (d *Driver) DeleteNetwork(r *network.DeleteNetworkRequest) error {
	name := r.NetworkID

	netlink.LinkDel(d.networks[name].VXLan)
	netlink.LinkDel(d.networks[name].Bridge)

	return nil
}
<|endoftext|>


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

import (
	"fmt"

	"github.com/xwb1989/sqlparser/dependency/sqltypes"

	"github.com/xwb1989/sqlparser/dependency/querypb"
)

// Normalize changes the statement to use bind values, and
// updates the bind vars to those values. The supplied prefix
// is used to generate the bind var names. The function ensures
// that there are no collisions with existing bind vars.
// Within Select constructs, bind vars are deduped. This allows
// us to identify vindex equality. Otherwise, every value is
// treated as distinct.
func Normalize(stmt Statement, bindVars map[string]*querypb.BindVariable, prefix string) {
	nz := newNormalizer(stmt, bindVars, prefix)
	_ = Walk(nz.WalkStatement, stmt)
}

type normalizer struct {
	stmt     Statement
	bindVars map[string]*querypb.BindVariable
	prefix   string
	reserved map[string]struct{}
	counter  int
	vals     map[string]string
}

func newNormalizer(stmt Statement, bindVars map[string]*querypb.BindVariable, prefix string) *normalizer {
	return &normalizer{
		stmt:     stmt,
		bindVars: bindVars,
		prefix:   prefix,
		reserved: GetBindvars(stmt),
		counter:  1,
		vals:     make(map[string]string),
	}
}

// WalkStatement is the top level walk function.
// If it encounters a Select, it switches to a mode
// where variables are deduped.
func (nz *normalizer) WalkStatement(node SQLNode) (bool, error) {
	switch node := node.(type) {
	case *Select:
		_ = Walk(nz.WalkSelect, node)
		// Don't continue
		return false, nil
	case *SQLVal:
		nz.convertSQLVal(node)
	case *ComparisonExpr:
		nz.convertComparison(node)
	}
	return true, nil
}

// WalkSelect normalizes the AST in Select mode.
func (nz *normalizer) WalkSelect(node SQLNode) (bool, error) {
	switch node := node.(type) {
	case *SQLVal:
		nz.convertSQLValDedup(node)
	case *ComparisonExpr:
		nz.convertComparison(node)
	}
	return true, nil
}

func (nz *normalizer) convertSQLValDedup(node *SQLVal) {
	// If value is too long, don't dedup.
	if len(node.Val) > 256 {
		nz.convertSQLVal(node)
		return
	}

	// Make the bindvar
	bval := nz.sqlToBindvar(node)
	if bval == nil {
		return
	}

	// Check if there's a bindvar for that value already.
	var key string
	if bval.Type == sqltypes.VarBinary {
		// Prefixing strings with "'" ensures that a string
		// and number that have the same representation don't
		// collide.
		key = "'" + string(node.Val)
	} else {
		key = string(node.Val)
	}
	bvname, ok := nz.vals[key]
	if !ok {
		// If there's no such bindvar, make a new one.
		bvname = nz.newName()
		nz.vals[key] = bvname
		nz.bindVars[bvname] = bval
	}

	// Modify the AST node to a bindvar.
	node.Type = ValArg
	node.Val = append([]byte(":"), bvname...)
}

// convertSQLVal converts an SQLVal without the dedup.
func (nz *normalizer) convertSQLVal(node *SQLVal) {
	bval := nz.sqlToBindvar(node)
	if bval == nil {
		return
	}

	bvname := nz.newName()
	nz.bindVars[bvname] = bval

	node.Type = ValArg
	node.Val = append([]byte(":"), bvname...)
}

// convertComparison attempts to convert IN clauses to
// use the list bind var construct. If it fails, it returns
// with no change made. The walk function will then continue
// and iterate on converting each individual value into separate
// bind vars.
func (nz *normalizer) convertComparison(node *ComparisonExpr) {
	switch node.Operator {
	case InStr, NotInStr:
	default:
		return
	}
	// It's either IN or NOT IN.
	tupleVals, ok := node.Right.(ValTuple)
	if !ok {
		return
	}
	// The RHS is a tuple of values.
	// Make a list bindvar.
	bvals := &querypb.BindVariable{
		Type: querypb.Type_TUPLE,
	}
	for _, val := range tupleVals {
		bval := nz.sqlToBindvar(val)
		if bval == nil {
			return
		}
		bvals.Values = append(bvals.Values, &querypb.Value{
			Type:  bval.Type,
			Value: bval.Value,
		})
	}
	bvname := nz.newName()
	nz.bindVars[bvname] = bvals
	// Modify RHS to be a list bindvar.
	node.Right = ListArg(append([]byte("::"), bvname...))
}

func (nz *normalizer) sqlToBindvar(node SQLNode) *querypb.BindVariable {
	if node, ok := node.(*SQLVal); ok {
		var v sqltypes.Value
		var err error
		switch node.Type {
		case StrVal:
			v, err = sqltypes.NewValue(sqltypes.VarBinary, node.Val)
		case IntVal:
			v, err = sqltypes.NewValue(sqltypes.Int64, node.Val)
		case FloatVal:
			v, err = sqltypes.NewValue(sqltypes.Float64, node.Val)
		default:
			return nil
		}
		if err != nil {
			return nil
		}
		return sqltypes.ValueBindVariable(v)
	}
	return nil
}

func (nz *normalizer) newName() string {
	for {
		newName := fmt.Sprintf("%s%d", nz.prefix, nz.counter)
		if _, ok := nz.reserved[newName]; !ok {
			nz.reserved[newName] = struct{}{}
			return newName
		}
		nz.counter++
	}
}

// GetBindvars returns a map of the bind vars referenced in the statement.
// TODO(sougou); This function gets called again from vtgate/planbuilder.
// Ideally, this should be done only once.
func GetBindvars(stmt Statement) map[string]struct{} {
	bindvars := make(map[string]struct{})
	_ = Walk(func(node SQLNode) (kontinue bool, err error) {
		switch node := node.(type) {
		case *SQLVal:
			if node.Type == ValArg {
				bindvars[string(node.Val[1:])] = struct{}{}
			}
		case ListArg:
			bindvars[string(node[2:])] = struct{}{}
		}
		return true, nil
	}, stmt)
	return bindvars
}
<commit_msg>v3: normalizer: address review comments<commit_after>/*
Copyright 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 sqlparser

import (
	"fmt"

	"github.com/xwb1989/sqlparser/dependency/sqltypes"

	"github.com/xwb1989/sqlparser/dependency/querypb"
)

// Normalize changes the statement to use bind values, and
// updates the bind vars to those values. The supplied prefix
// is used to generate the bind var names. The function ensures
// that there are no collisions with existing bind vars.
// Within Select constructs, bind vars are deduped. This allows
// us to identify vindex equality. Otherwise, every value is
// treated as distinct.
func Normalize(stmt Statement, bindVars map[string]*querypb.BindVariable, prefix string) {
	nz := newNormalizer(stmt, bindVars, prefix)
	_ = Walk(nz.WalkStatement, stmt)
}

type normalizer struct {
	stmt     Statement
	bindVars map[string]*querypb.BindVariable
	prefix   string
	reserved map[string]struct{}
	counter  int
	vals     map[string]string
}

func newNormalizer(stmt Statement, bindVars map[string]*querypb.BindVariable, prefix string) *normalizer {
	return &normalizer{
		stmt:     stmt,
		bindVars: bindVars,
		prefix:   prefix,
		reserved: GetBindvars(stmt),
		counter:  1,
		vals:     make(map[string]string),
	}
}

// WalkStatement is the top level walk function.
// If it encounters a Select, it switches to a mode
// where variables are deduped.
func (nz *normalizer) WalkStatement(node SQLNode) (bool, error) {
	switch node := node.(type) {
	case *Select:
		_ = Walk(nz.WalkSelect, node)
		// Don't continue
		return false, nil
	case *SQLVal:
		nz.convertSQLVal(node)
	case *ComparisonExpr:
		nz.convertComparison(node)
	}
	return true, nil
}

// WalkSelect normalizes the AST in Select mode.
func (nz *normalizer) WalkSelect(node SQLNode) (bool, error) {
	switch node := node.(type) {
	case *SQLVal:
		nz.convertSQLValDedup(node)
	case *ComparisonExpr:
		nz.convertComparison(node)
	}
	return true, nil
}

func (nz *normalizer) convertSQLValDedup(node *SQLVal) {
	// If value is too long, don't dedup.
	// Such values are most likely not for vindexes.
	// We save a lot of CPU because we avoid building
	// the key for them.
	if len(node.Val) > 256 {
		nz.convertSQLVal(node)
		return
	}

	// Make the bindvar
	bval := nz.sqlToBindvar(node)
	if bval == nil {
		return
	}

	// Check if there's a bindvar for that value already.
	var key string
	if bval.Type == sqltypes.VarBinary {
		// Prefixing strings with "'" ensures that a string
		// and number that have the same representation don't
		// collide.
		key = "'" + string(node.Val)
	} else {
		key = string(node.Val)
	}
	bvname, ok := nz.vals[key]
	if !ok {
		// If there's no such bindvar, make a new one.
		bvname = nz.newName()
		nz.vals[key] = bvname
		nz.bindVars[bvname] = bval
	}

	// Modify the AST node to a bindvar.
	node.Type = ValArg
	node.Val = append([]byte(":"), bvname...)
}

// convertSQLVal converts an SQLVal without the dedup.
func (nz *normalizer) convertSQLVal(node *SQLVal) {
	bval := nz.sqlToBindvar(node)
	if bval == nil {
		return
	}

	bvname := nz.newName()
	nz.bindVars[bvname] = bval

	node.Type = ValArg
	node.Val = append([]byte(":"), bvname...)
}

// convertComparison attempts to convert IN clauses to
// use the list bind var construct. If it fails, it returns
// with no change made. The walk function will then continue
// and iterate on converting each individual value into separate
// bind vars.
func (nz *normalizer) convertComparison(node *ComparisonExpr) {
	if node.Operator != InStr && node.Operator != NotInStr {
		return
	}
	tupleVals, ok := node.Right.(ValTuple)
	if !ok {
		return
	}
	// The RHS is a tuple of values.
	// Make a list bindvar.
	bvals := &querypb.BindVariable{
		Type: querypb.Type_TUPLE,
	}
	for _, val := range tupleVals {
		bval := nz.sqlToBindvar(val)
		if bval == nil {
			return
		}
		bvals.Values = append(bvals.Values, &querypb.Value{
			Type:  bval.Type,
			Value: bval.Value,
		})
	}
	bvname := nz.newName()
	nz.bindVars[bvname] = bvals
	// Modify RHS to be a list bindvar.
	node.Right = ListArg(append([]byte("::"), bvname...))
}

func (nz *normalizer) sqlToBindvar(node SQLNode) *querypb.BindVariable {
	if node, ok := node.(*SQLVal); ok {
		var v sqltypes.Value
		var err error
		switch node.Type {
		case StrVal:
			v, err = sqltypes.NewValue(sqltypes.VarBinary, node.Val)
		case IntVal:
			v, err = sqltypes.NewValue(sqltypes.Int64, node.Val)
		case FloatVal:
			v, err = sqltypes.NewValue(sqltypes.Float64, node.Val)
		default:
			return nil
		}
		if err != nil {
			return nil
		}
		return sqltypes.ValueBindVariable(v)
	}
	return nil
}

func (nz *normalizer) newName() string {
	for {
		newName := fmt.Sprintf("%s%d", nz.prefix, nz.counter)
		if _, ok := nz.reserved[newName]; !ok {
			nz.reserved[newName] = struct{}{}
			return newName
		}
		nz.counter++
	}
}

// GetBindvars returns a map of the bind vars referenced in the statement.
// TODO(sougou); This function gets called again from vtgate/planbuilder.
// Ideally, this should be done only once.
func GetBindvars(stmt Statement) map[string]struct{} {
	bindvars := make(map[string]struct{})
	_ = Walk(func(node SQLNode) (kontinue bool, err error) {
		switch node := node.(type) {
		case *SQLVal:
			if node.Type == ValArg {
				bindvars[string(node.Val[1:])] = struct{}{}
			}
		case ListArg:
			bindvars[string(node[2:])] = struct{}{}
		}
		return true, nil
	}, stmt)
	return bindvars
}
<|endoftext|>