<commit_before>//
// Copyright (c) 2016 Intel Corporation
//
// 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 libsnnet

import (
	"net"

	"github.com/vishvananda/netlink"
)

// NewCnciVnic is used to initialize the CnciVnic properties
// This has to be called prior to Create() or GetDevice()
func newCnciVnic(id string) (*CnciVnic, error) {
	CnciVnic := &CnciVnic{}
	CnciVnic.Link = &netlink.Macvtap{}
	CnciVnic.Link.Mode = netlink.MACVLAN_MODE_VEPA
	//TODO: Use bridge mode as default for now for easier integration
	//CnciVnic.Link.Mode = netlink.MACVLAN_MODE_BRIDGE
	//TODO: Figure out correct default
	//Currently using the same as ip command
	CnciVnic.Link.TxQLen = 500
	CnciVnic.GlobalID = id
	return CnciVnic, nil
}

// GetDevice is used to associate with an existing CnciVnic provided it satisfies
// the needs of a CnciVnic. Returns error if the CnciVnic does not exist
func (v *CnciVnic) getDevice() error {

	if v.GlobalID == "" {
		return netError(v, "getdevice unnamed cnci vnic")
	}

	link, err := netlink.LinkByAlias(v.GlobalID)
	if err != nil {
		return netError(v, "getdevice interface does not exist: %v", v.GlobalID)
	}

	vl, ok := link.(*netlink.Macvtap)
	if !ok {
		return netError(v, "getdevice incorrect interface type %v %v", v.GlobalID, link.Type())
	}

	if link.Type() != "macvtap" {
		return netError(v, "getdevice incorrect interface type %v %v", v.GlobalID, link.Type())
	}

	v.LinkName = vl.Name
	v.Link = vl
	return nil
}

// Create instantiates new vnic
func (v *CnciVnic) create() error {
	var err error

	if v.GlobalID == "" {
		return netError(v, "create cannot create an unnamed cnci vnic")
	}

	if v.LinkName == "" {
		if v.LinkName, err = genIface(v, true); err != nil {
			return netError(v, "create geniface %v %v", v.GlobalID, err)
		}

		if _, err := netlink.LinkByAlias(v.GlobalID); err == nil {
			return netError(v, "create interface exists %v %v", v.GlobalID, err)
		}
	}

	v.Link.Name = v.LinkName
	if v.Link.ParentIndex == 0 {
		return netError(v, "create parent index not set %v %v", v.GlobalID, v.Link)
	}

	if err := netlink.LinkAdd(v.Link); err != nil {
		return netError(v, "create netlink.LinkAdd %v %v", v.GlobalID, err)
	}

	link, err := netlink.LinkByName(v.LinkName)
	if err != nil {
		return netError(v, "create netlink.LinkAdd %v %v", v.GlobalID, err)
	}

	vl, ok := link.(*netlink.Macvtap)
	if !ok {
		return netError(v, "create incorrect interface type %v %v", v.GlobalID, link.Type())
	}

	v.Link = vl

	if err := v.setAlias(v.GlobalID); err != nil {
		err1 := v.destroy()
		return netError(v, "create set alias [%v] [%v] [%v]", v.GlobalID, err, err1)
	}

	if v.MACAddr != nil {
		if err := v.setHardwareAddr(*v.MACAddr); err != nil {
			err1 := v.destroy()
			return netError(v, "create set hardware addr [%v] [%v] [%v] [%v]",
				v.MACAddr.String(), v.GlobalID, err, err1)
		}
	}

	return nil
}

// Destroy a vnic
func (v *CnciVnic) destroy() error {

	if v.Link == nil {
		return netError(v, "destroy invalid Link: %v", v)
	}

	if err := netlink.LinkDel(v.Link); err != nil {
		return netError(v, "destroy link del %v", err)
	}

	return nil

}

// Enable the vnic
func (v *CnciVnic) enable() error {

	if v.Link == nil {
		return netError(v, "enable invalid link: %v", v)
	}

	if err := netlink.LinkSetUp(v.Link); err != nil {
		return netError(v, "enable link up %v", err)
	}

	return nil

}

// Disable the vnic
func (v *CnciVnic) disable() error {

	if v.Link == nil {
		return netError(v, "disable invalid link: %v", v)
	}

	if err := netlink.LinkSetDown(v.Link); err != nil {
		return netError(v, "disable link down %v", err)
	}

	return nil
}

func (v *CnciVnic) setAlias(alias string) error {

	if v.Link == nil {
		return netError(v, "set alias vnic unnitialized")
	}

	if err := netlink.LinkSetAlias(v.Link, alias); err != nil {
		return netError(v, "set alias link set alias %v %v", alias, err)
	}

	return nil
}

func (v *CnciVnic) setHardwareAddr(hwaddr net.HardwareAddr) error {

	if v.Link == nil {
		return netError(v, "set hw addr vnic unnitialized")
	}

	if err := netlink.LinkSetHardwareAddr(v.Link, hwaddr); err != nil {
		return netError(v, "set hwaddr %v %v", hwaddr.String(), err)
	}

	return nil
}
<commit_msg>Networking: Switch CNCI VNIC to macvtap bridge mode<commit_after>//
// Copyright (c) 2016 Intel Corporation
//
// 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 libsnnet

import (
	"net"

	"github.com/vishvananda/netlink"
)

// NewCnciVnic is used to initialize the CnciVnic properties
// This has to be called prior to Create() or GetDevice()
func newCnciVnic(id string) (*CnciVnic, error) {
	CnciVnic := &CnciVnic{}
	CnciVnic.Link = &netlink.Macvtap{}
	CnciVnic.Link.Mode = netlink.MACVLAN_MODE_BRIDGE
	CnciVnic.Link.TxQLen = 500
	CnciVnic.GlobalID = id
	return CnciVnic, nil
}

// GetDevice is used to associate with an existing CnciVnic provided it satisfies
// the needs of a CnciVnic. Returns error if the CnciVnic does not exist
func (v *CnciVnic) getDevice() error {

	if v.GlobalID == "" {
		return netError(v, "getdevice unnamed cnci vnic")
	}

	link, err := netlink.LinkByAlias(v.GlobalID)
	if err != nil {
		return netError(v, "getdevice interface does not exist: %v", v.GlobalID)
	}

	vl, ok := link.(*netlink.Macvtap)
	if !ok {
		return netError(v, "getdevice incorrect interface type %v %v", v.GlobalID, link.Type())
	}

	if link.Type() != "macvtap" {
		return netError(v, "getdevice incorrect interface type %v %v", v.GlobalID, link.Type())
	}

	v.LinkName = vl.Name
	v.Link = vl
	return nil
}

// Create instantiates new vnic
func (v *CnciVnic) create() error {
	var err error

	if v.GlobalID == "" {
		return netError(v, "create cannot create an unnamed cnci vnic")
	}

	if v.LinkName == "" {
		if v.LinkName, err = genIface(v, true); err != nil {
			return netError(v, "create geniface %v %v", v.GlobalID, err)
		}

		if _, err := netlink.LinkByAlias(v.GlobalID); err == nil {
			return netError(v, "create interface exists %v %v", v.GlobalID, err)
		}
	}

	v.Link.Name = v.LinkName
	if v.Link.ParentIndex == 0 {
		return netError(v, "create parent index not set %v %v", v.GlobalID, v.Link)
	}

	if err := netlink.LinkAdd(v.Link); err != nil {
		return netError(v, "create netlink.LinkAdd %v %v", v.GlobalID, err)
	}

	link, err := netlink.LinkByName(v.LinkName)
	if err != nil {
		return netError(v, "create netlink.LinkAdd %v %v", v.GlobalID, err)
	}

	vl, ok := link.(*netlink.Macvtap)
	if !ok {
		return netError(v, "create incorrect interface type %v %v", v.GlobalID, link.Type())
	}

	v.Link = vl

	if err := v.setAlias(v.GlobalID); err != nil {
		err1 := v.destroy()
		return netError(v, "create set alias [%v] [%v] [%v]", v.GlobalID, err, err1)
	}

	if v.MACAddr != nil {
		if err := v.setHardwareAddr(*v.MACAddr); err != nil {
			err1 := v.destroy()
			return netError(v, "create set hardware addr [%v] [%v] [%v] [%v]",
				v.MACAddr.String(), v.GlobalID, err, err1)
		}
	}

	return nil
}

// Destroy a vnic
func (v *CnciVnic) destroy() error {

	if v.Link == nil {
		return netError(v, "destroy invalid Link: %v", v)
	}

	if err := netlink.LinkDel(v.Link); err != nil {
		return netError(v, "destroy link del %v", err)
	}

	return nil

}

// Enable the vnic
func (v *CnciVnic) enable() error {

	if v.Link == nil {
		return netError(v, "enable invalid link: %v", v)
	}

	if err := netlink.LinkSetUp(v.Link); err != nil {
		return netError(v, "enable link up %v", err)
	}

	return nil

}

// Disable the vnic
func (v *CnciVnic) disable() error {

	if v.Link == nil {
		return netError(v, "disable invalid link: %v", v)
	}

	if err := netlink.LinkSetDown(v.Link); err != nil {
		return netError(v, "disable link down %v", err)
	}

	return nil
}

func (v *CnciVnic) setAlias(alias string) error {

	if v.Link == nil {
		return netError(v, "set alias vnic unnitialized")
	}

	if err := netlink.LinkSetAlias(v.Link, alias); err != nil {
		return netError(v, "set alias link set alias %v %v", alias, err)
	}

	return nil
}

func (v *CnciVnic) setHardwareAddr(hwaddr net.HardwareAddr) error {

	if v.Link == nil {
		return netError(v, "set hw addr vnic unnitialized")
	}

	if err := netlink.LinkSetHardwareAddr(v.Link, hwaddr); err != nil {
		return netError(v, "set hwaddr %v %v", hwaddr.String(), err)
	}

	return nil
}
<|endoftext|>


<commit_before>package redis

import (
	"context"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/go-redis/redis/internal"
	"github.com/go-redis/redis/internal/pool"
	"github.com/go-redis/redis/internal/proto"
)

// Nil reply Redis returns when key does not exist.
const Nil = proto.Nil

func init() {
	SetLogger(log.New(os.Stderr, "redis: ", log.LstdFlags|log.Lshortfile))
}

func SetLogger(logger *log.Logger) {
	internal.Logger = logger
}

type baseClient struct {
	opt      *Options
	connPool pool.Pooler

	process           func(Cmder) error
	processPipeline   func([]Cmder) error
	processTxPipeline func([]Cmder) error

	onClose func() error // hook called when client is closed
}

func (c *baseClient) init() {
	c.process = c.defaultProcess
	c.processPipeline = c.defaultProcessPipeline
	c.processTxPipeline = c.defaultProcessTxPipeline
}

func (c *baseClient) String() string {
	return fmt.Sprintf("Redis<%s db:%d>", c.getAddr(), c.opt.DB)
}

func (c *baseClient) newConn() (*pool.Conn, error) {
	cn, err := c.connPool.NewConn()
	if err != nil {
		return nil, err
	}

	if cn.InitedAt.IsZero() {
		if err := c.initConn(cn); err != nil {
			_ = c.connPool.CloseConn(cn)
			return nil, err
		}
	}

	return cn, nil
}

func (c *baseClient) getConn() (*pool.Conn, error) {
	cn, err := c.connPool.Get()
	if err != nil {
		return nil, err
	}

	if cn.InitedAt.IsZero() {
		err := c.initConn(cn)
		if err != nil {
			c.connPool.Remove(cn)
			return nil, err
		}
	}

	return cn, nil
}

func (c *baseClient) releaseConn(cn *pool.Conn, err error) bool {
	if internal.IsBadConn(err, false) {
		c.connPool.Remove(cn)
		return false
	}

	c.connPool.Put(cn)
	return true
}

func (c *baseClient) initConn(cn *pool.Conn) error {
	cn.InitedAt = time.Now()

	if c.opt.Password == "" &&
		c.opt.DB == 0 &&
		!c.opt.readOnly &&
		c.opt.OnConnect == nil {
		return nil
	}

	conn := newConn(c.opt, cn)
	_, err := conn.Pipelined(func(pipe Pipeliner) error {
		if c.opt.Password != "" {
			pipe.Auth(c.opt.Password)
		}

		if c.opt.DB > 0 {
			pipe.Select(c.opt.DB)
		}

		if c.opt.readOnly {
			pipe.ReadOnly()
		}

		return nil
	})
	if err != nil {
		return err
	}

	if c.opt.OnConnect != nil {
		return c.opt.OnConnect(conn)
	}
	return nil
}

// Do creates a Cmd from the args and processes the cmd.
func (c *baseClient) Do(args ...interface{}) *Cmd {
	cmd := NewCmd(args...)
	c.Process(cmd)
	return cmd
}

// WrapProcess wraps function that processes Redis commands.
func (c *baseClient) WrapProcess(
	fn func(oldProcess func(cmd Cmder) error) func(cmd Cmder) error,
) {
	c.process = fn(c.process)
}

func (c *baseClient) Process(cmd Cmder) error {
	return c.process(cmd)
}

func (c *baseClient) defaultProcess(cmd Cmder) error {
	for attempt := 0; attempt <= c.opt.MaxRetries; attempt++ {
		if attempt > 0 {
			time.Sleep(c.retryBackoff(attempt))
		}

		cn, err := c.getConn()
		if err != nil {
			cmd.setErr(err)
			if internal.IsRetryableError(err, true) {
				continue
			}
			return err
		}

		err = cn.WithWriter(c.opt.WriteTimeout, func(wr *proto.Writer) error {
			return writeCmd(wr, cmd)
		})
		if err != nil {
			c.releaseConn(cn, err)
			cmd.setErr(err)
			if internal.IsRetryableError(err, true) {
				continue
			}
			return err
		}

		err = cn.WithReader(c.cmdTimeout(cmd), func(rd *proto.Reader) error {
			return cmd.readReply(rd)
		})
		c.releaseConn(cn, err)
		if err != nil && internal.IsRetryableError(err, cmd.readTimeout() == nil) {
			continue
		}

		return err
	}

	return cmd.Err()
}

func (c *baseClient) retryBackoff(attempt int) time.Duration {
	return internal.RetryBackoff(attempt, c.opt.MinRetryBackoff, c.opt.MaxRetryBackoff)
}

func (c *baseClient) cmdTimeout(cmd Cmder) time.Duration {
	if timeout := cmd.readTimeout(); timeout != nil {
		t := *timeout
		if t == 0 {
			return 0
		}
		return t + 10*time.Second
	}
	return c.opt.ReadTimeout
}

// Close closes the client, releasing any open resources.
//
// It is rare to Close a Client, as the Client is meant to be
// long-lived and shared between many goroutines.
func (c *baseClient) Close() error {
	var firstErr error
	if c.onClose != nil {
		if err := c.onClose(); err != nil && firstErr == nil {
			firstErr = err
		}
	}
	if err := c.connPool.Close(); err != nil && firstErr == nil {
		firstErr = err
	}
	return firstErr
}

func (c *baseClient) getAddr() string {
	return c.opt.Addr
}

func (c *baseClient) WrapProcessPipeline(
	fn func(oldProcess func([]Cmder) error) func([]Cmder) error,
) {
	c.processPipeline = fn(c.processPipeline)
	c.processTxPipeline = fn(c.processTxPipeline)
}

func (c *baseClient) defaultProcessPipeline(cmds []Cmder) error {
	return c.generalProcessPipeline(cmds, c.pipelineProcessCmds)
}

func (c *baseClient) defaultProcessTxPipeline(cmds []Cmder) error {
	return c.generalProcessPipeline(cmds, c.txPipelineProcessCmds)
}

type pipelineProcessor func(*pool.Conn, []Cmder) (bool, error)

func (c *baseClient) generalProcessPipeline(cmds []Cmder, p pipelineProcessor) error {
	for attempt := 0; attempt <= c.opt.MaxRetries; attempt++ {
		if attempt > 0 {
			time.Sleep(c.retryBackoff(attempt))
		}

		cn, err := c.getConn()
		if err != nil {
			setCmdsErr(cmds, err)
			return err
		}

		canRetry, err := p(cn, cmds)

		if err == nil || internal.IsRedisError(err) {
			c.connPool.Put(cn)
			break
		}
		c.connPool.Remove(cn)

		if !canRetry || !internal.IsRetryableError(err, true) {
			break
		}
	}
	return cmdsFirstErr(cmds)
}

func (c *baseClient) pipelineProcessCmds(cn *pool.Conn, cmds []Cmder) (bool, error) {
	err := cn.WithWriter(c.opt.WriteTimeout, func(wr *proto.Writer) error {
		return writeCmd(wr, cmds...)
	})
	if err != nil {
		setCmdsErr(cmds, err)
		return true, err
	}

	err = cn.WithReader(c.opt.ReadTimeout, func(rd *proto.Reader) error {
		return pipelineReadCmds(rd, cmds)
	})
	return true, err
}

func pipelineReadCmds(rd *proto.Reader, cmds []Cmder) error {
	for _, cmd := range cmds {
		err := cmd.readReply(rd)
		if err != nil && !internal.IsRedisError(err) {
			return err
		}
	}
	return nil
}

func (c *baseClient) txPipelineProcessCmds(cn *pool.Conn, cmds []Cmder) (bool, error) {
	err := cn.WithWriter(c.opt.WriteTimeout, func(wr *proto.Writer) error {
		return txPipelineWriteMulti(wr, cmds)
	})
	if err != nil {
		setCmdsErr(cmds, err)
		return true, err
	}

	err = cn.WithReader(c.opt.ReadTimeout, func(rd *proto.Reader) error {
		err := txPipelineReadQueued(rd, cmds)
		if err != nil {
			setCmdsErr(cmds, err)
			return err
		}
		return pipelineReadCmds(rd, cmds)
	})
	return false, err
}

func txPipelineWriteMulti(wr *proto.Writer, cmds []Cmder) error {
	multiExec := make([]Cmder, 0, len(cmds)+2)
	multiExec = append(multiExec, NewStatusCmd("MULTI"))
	multiExec = append(multiExec, cmds...)
	multiExec = append(multiExec, NewSliceCmd("EXEC"))
	return writeCmd(wr, multiExec...)
}

func txPipelineReadQueued(rd *proto.Reader, cmds []Cmder) error {
	// Parse queued replies.
	var statusCmd StatusCmd
	err := statusCmd.readReply(rd)
	if err != nil {
		return err
	}

	for _ = range cmds {
		err = statusCmd.readReply(rd)
		if err != nil && !internal.IsRedisError(err) {
			return err
		}
	}

	// Parse number of replies.
	line, err := rd.ReadLine()
	if err != nil {
		if err == Nil {
			err = TxFailedErr
		}
		return err
	}

	switch line[0] {
	case proto.ErrorReply:
		return proto.ParseErrorReply(line)
	case proto.ArrayReply:
		// ok
	default:
		err := fmt.Errorf("redis: expected '*', but got line %q", line)
		return err
	}

	return nil
}

//------------------------------------------------------------------------------

// Client is a Redis client representing a pool of zero or more
// underlying connections. It's safe for concurrent use by multiple
// goroutines.
type Client struct {
	baseClient
	cmdable

	ctx context.Context
}

// NewClient returns a client to the Redis Server specified by Options.
func NewClient(opt *Options) *Client {
	opt.init()

	c := Client{
		baseClient: baseClient{
			opt:      opt,
			connPool: newConnPool(opt),
		},
	}
	c.baseClient.init()
	c.init()

	return &c
}

func (c *Client) init() {
	c.cmdable.setProcessor(c.Process)
}

func (c *Client) Context() context.Context {
	if c.ctx != nil {
		return c.ctx
	}
	return context.Background()
}

func (c *Client) WithContext(ctx context.Context) *Client {
	if ctx == nil {
		panic("nil context")
	}
	c2 := c.copy()
	c2.ctx = ctx
	return c2
}

func (c *Client) copy() *Client {
	cp := *c
	cp.init()
	return &cp
}

// Options returns read-only Options that were used to create the client.
func (c *Client) Options() *Options {
	return c.opt
}

type PoolStats pool.Stats

// PoolStats returns connection pool stats.
func (c *Client) PoolStats() *PoolStats {
	stats := c.connPool.Stats()
	return (*PoolStats)(stats)
}

func (c *Client) Pipelined(fn func(Pipeliner) error) ([]Cmder, error) {
	return c.Pipeline().Pipelined(fn)
}

func (c *Client) Pipeline() Pipeliner {
	pipe := Pipeline{
		exec: c.processPipeline,
	}
	pipe.statefulCmdable.setProcessor(pipe.Process)
	return &pipe
}

func (c *Client) TxPipelined(fn func(Pipeliner) error) ([]Cmder, error) {
	return c.TxPipeline().Pipelined(fn)
}

// TxPipeline acts like Pipeline, but wraps queued commands with MULTI/EXEC.
func (c *Client) TxPipeline() Pipeliner {
	pipe := Pipeline{
		exec: c.processTxPipeline,
	}
	pipe.statefulCmdable.setProcessor(pipe.Process)
	return &pipe
}

func (c *Client) pubSub() *PubSub {
	pubsub := &PubSub{
		opt: c.opt,

		newConn: func(channels []string) (*pool.Conn, error) {
			return c.newConn()
		},
		closeConn: c.connPool.CloseConn,
	}
	pubsub.init()
	return pubsub
}

// Subscribe subscribes the client to the specified channels.
// Channels can be omitted to create empty subscription.
func (c *Client) Subscribe(channels ...string) *PubSub {
	pubsub := c.pubSub()
	if len(channels) > 0 {
		_ = pubsub.Subscribe(channels...)
	}
	return pubsub
}

// PSubscribe subscribes the client to the given patterns.
// Patterns can be omitted to create empty subscription.
func (c *Client) PSubscribe(channels ...string) *PubSub {
	pubsub := c.pubSub()
	if len(channels) > 0 {
		_ = pubsub.PSubscribe(channels...)
	}
	return pubsub
}

//------------------------------------------------------------------------------

// Conn is like Client, but its pool contains single connection.
type Conn struct {
	baseClient
	statefulCmdable
}

func newConn(opt *Options, cn *pool.Conn) *Conn {
	c := Conn{
		baseClient: baseClient{
			opt:      opt,
			connPool: pool.NewSingleConnPool(cn),
		},
	}
	c.baseClient.init()
	c.statefulCmdable.setProcessor(c.Process)
	return &c
}

func (c *Conn) Pipelined(fn func(Pipeliner) error) ([]Cmder, error) {
	return c.Pipeline().Pipelined(fn)
}

func (c *Conn) Pipeline() Pipeliner {
	pipe := Pipeline{
		exec: c.processPipeline,
	}
	pipe.statefulCmdable.setProcessor(pipe.Process)
	return &pipe
}

func (c *Conn) TxPipelined(fn func(Pipeliner) error) ([]Cmder, error) {
	return c.TxPipeline().Pipelined(fn)
}

// TxPipeline acts like Pipeline, but wraps queued commands with MULTI/EXEC.
func (c *Conn) TxPipeline() Pipeliner {
	pipe := Pipeline{
		exec: c.processTxPipeline,
	}
	pipe.statefulCmdable.setProcessor(pipe.Process)
	return &pipe
}
<commit_msg>remove redundant expression<commit_after>package redis

import (
	"context"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/go-redis/redis/internal"
	"github.com/go-redis/redis/internal/pool"
	"github.com/go-redis/redis/internal/proto"
)

// Nil reply Redis returns when key does not exist.
const Nil = proto.Nil

func init() {
	SetLogger(log.New(os.Stderr, "redis: ", log.LstdFlags|log.Lshortfile))
}

func SetLogger(logger *log.Logger) {
	internal.Logger = logger
}

type baseClient struct {
	opt      *Options
	connPool pool.Pooler

	process           func(Cmder) error
	processPipeline   func([]Cmder) error
	processTxPipeline func([]Cmder) error

	onClose func() error // hook called when client is closed
}

func (c *baseClient) init() {
	c.process = c.defaultProcess
	c.processPipeline = c.defaultProcessPipeline
	c.processTxPipeline = c.defaultProcessTxPipeline
}

func (c *baseClient) String() string {
	return fmt.Sprintf("Redis<%s db:%d>", c.getAddr(), c.opt.DB)
}

func (c *baseClient) newConn() (*pool.Conn, error) {
	cn, err := c.connPool.NewConn()
	if err != nil {
		return nil, err
	}

	if cn.InitedAt.IsZero() {
		if err := c.initConn(cn); err != nil {
			_ = c.connPool.CloseConn(cn)
			return nil, err
		}
	}

	return cn, nil
}

func (c *baseClient) getConn() (*pool.Conn, error) {
	cn, err := c.connPool.Get()
	if err != nil {
		return nil, err
	}

	if cn.InitedAt.IsZero() {
		err := c.initConn(cn)
		if err != nil {
			c.connPool.Remove(cn)
			return nil, err
		}
	}

	return cn, nil
}

func (c *baseClient) releaseConn(cn *pool.Conn, err error) bool {
	if internal.IsBadConn(err, false) {
		c.connPool.Remove(cn)
		return false
	}

	c.connPool.Put(cn)
	return true
}

func (c *baseClient) initConn(cn *pool.Conn) error {
	cn.InitedAt = time.Now()

	if c.opt.Password == "" &&
		c.opt.DB == 0 &&
		!c.opt.readOnly &&
		c.opt.OnConnect == nil {
		return nil
	}

	conn := newConn(c.opt, cn)
	_, err := conn.Pipelined(func(pipe Pipeliner) error {
		if c.opt.Password != "" {
			pipe.Auth(c.opt.Password)
		}

		if c.opt.DB > 0 {
			pipe.Select(c.opt.DB)
		}

		if c.opt.readOnly {
			pipe.ReadOnly()
		}

		return nil
	})
	if err != nil {
		return err
	}

	if c.opt.OnConnect != nil {
		return c.opt.OnConnect(conn)
	}
	return nil
}

// Do creates a Cmd from the args and processes the cmd.
func (c *baseClient) Do(args ...interface{}) *Cmd {
	cmd := NewCmd(args...)
	c.Process(cmd)
	return cmd
}

// WrapProcess wraps function that processes Redis commands.
func (c *baseClient) WrapProcess(
	fn func(oldProcess func(cmd Cmder) error) func(cmd Cmder) error,
) {
	c.process = fn(c.process)
}

func (c *baseClient) Process(cmd Cmder) error {
	return c.process(cmd)
}

func (c *baseClient) defaultProcess(cmd Cmder) error {
	for attempt := 0; attempt <= c.opt.MaxRetries; attempt++ {
		if attempt > 0 {
			time.Sleep(c.retryBackoff(attempt))
		}

		cn, err := c.getConn()
		if err != nil {
			cmd.setErr(err)
			if internal.IsRetryableError(err, true) {
				continue
			}
			return err
		}

		err = cn.WithWriter(c.opt.WriteTimeout, func(wr *proto.Writer) error {
			return writeCmd(wr, cmd)
		})
		if err != nil {
			c.releaseConn(cn, err)
			cmd.setErr(err)
			if internal.IsRetryableError(err, true) {
				continue
			}
			return err
		}

		err = cn.WithReader(c.cmdTimeout(cmd), func(rd *proto.Reader) error {
			return cmd.readReply(rd)
		})
		c.releaseConn(cn, err)
		if err != nil && internal.IsRetryableError(err, cmd.readTimeout() == nil) {
			continue
		}

		return err
	}

	return cmd.Err()
}

func (c *baseClient) retryBackoff(attempt int) time.Duration {
	return internal.RetryBackoff(attempt, c.opt.MinRetryBackoff, c.opt.MaxRetryBackoff)
}

func (c *baseClient) cmdTimeout(cmd Cmder) time.Duration {
	if timeout := cmd.readTimeout(); timeout != nil {
		t := *timeout
		if t == 0 {
			return 0
		}
		return t + 10*time.Second
	}
	return c.opt.ReadTimeout
}

// Close closes the client, releasing any open resources.
//
// It is rare to Close a Client, as the Client is meant to be
// long-lived and shared between many goroutines.
func (c *baseClient) Close() error {
	var firstErr error
	if c.onClose != nil {
		if err := c.onClose(); err != nil && firstErr == nil {
			firstErr = err
		}
	}
	if err := c.connPool.Close(); err != nil && firstErr == nil {
		firstErr = err
	}
	return firstErr
}

func (c *baseClient) getAddr() string {
	return c.opt.Addr
}

func (c *baseClient) WrapProcessPipeline(
	fn func(oldProcess func([]Cmder) error) func([]Cmder) error,
) {
	c.processPipeline = fn(c.processPipeline)
	c.processTxPipeline = fn(c.processTxPipeline)
}

func (c *baseClient) defaultProcessPipeline(cmds []Cmder) error {
	return c.generalProcessPipeline(cmds, c.pipelineProcessCmds)
}

func (c *baseClient) defaultProcessTxPipeline(cmds []Cmder) error {
	return c.generalProcessPipeline(cmds, c.txPipelineProcessCmds)
}

type pipelineProcessor func(*pool.Conn, []Cmder) (bool, error)

func (c *baseClient) generalProcessPipeline(cmds []Cmder, p pipelineProcessor) error {
	for attempt := 0; attempt <= c.opt.MaxRetries; attempt++ {
		if attempt > 0 {
			time.Sleep(c.retryBackoff(attempt))
		}

		cn, err := c.getConn()
		if err != nil {
			setCmdsErr(cmds, err)
			return err
		}

		canRetry, err := p(cn, cmds)

		if err == nil || internal.IsRedisError(err) {
			c.connPool.Put(cn)
			break
		}
		c.connPool.Remove(cn)

		if !canRetry || !internal.IsRetryableError(err, true) {
			break
		}
	}
	return cmdsFirstErr(cmds)
}

func (c *baseClient) pipelineProcessCmds(cn *pool.Conn, cmds []Cmder) (bool, error) {
	err := cn.WithWriter(c.opt.WriteTimeout, func(wr *proto.Writer) error {
		return writeCmd(wr, cmds...)
	})
	if err != nil {
		setCmdsErr(cmds, err)
		return true, err
	}

	err = cn.WithReader(c.opt.ReadTimeout, func(rd *proto.Reader) error {
		return pipelineReadCmds(rd, cmds)
	})
	return true, err
}

func pipelineReadCmds(rd *proto.Reader, cmds []Cmder) error {
	for _, cmd := range cmds {
		err := cmd.readReply(rd)
		if err != nil && !internal.IsRedisError(err) {
			return err
		}
	}
	return nil
}

func (c *baseClient) txPipelineProcessCmds(cn *pool.Conn, cmds []Cmder) (bool, error) {
	err := cn.WithWriter(c.opt.WriteTimeout, func(wr *proto.Writer) error {
		return txPipelineWriteMulti(wr, cmds)
	})
	if err != nil {
		setCmdsErr(cmds, err)
		return true, err
	}

	err = cn.WithReader(c.opt.ReadTimeout, func(rd *proto.Reader) error {
		err := txPipelineReadQueued(rd, cmds)
		if err != nil {
			setCmdsErr(cmds, err)
			return err
		}
		return pipelineReadCmds(rd, cmds)
	})
	return false, err
}

func txPipelineWriteMulti(wr *proto.Writer, cmds []Cmder) error {
	multiExec := make([]Cmder, 0, len(cmds)+2)
	multiExec = append(multiExec, NewStatusCmd("MULTI"))
	multiExec = append(multiExec, cmds...)
	multiExec = append(multiExec, NewSliceCmd("EXEC"))
	return writeCmd(wr, multiExec...)
}

func txPipelineReadQueued(rd *proto.Reader, cmds []Cmder) error {
	// Parse queued replies.
	var statusCmd StatusCmd
	err := statusCmd.readReply(rd)
	if err != nil {
		return err
	}

	for range cmds {
		err = statusCmd.readReply(rd)
		if err != nil && !internal.IsRedisError(err) {
			return err
		}
	}

	// Parse number of replies.
	line, err := rd.ReadLine()
	if err != nil {
		if err == Nil {
			err = TxFailedErr
		}
		return err
	}

	switch line[0] {
	case proto.ErrorReply:
		return proto.ParseErrorReply(line)
	case proto.ArrayReply:
		// ok
	default:
		err := fmt.Errorf("redis: expected '*', but got line %q", line)
		return err
	}

	return nil
}

//------------------------------------------------------------------------------

// Client is a Redis client representing a pool of zero or more
// underlying connections. It's safe for concurrent use by multiple
// goroutines.
type Client struct {
	baseClient
	cmdable

	ctx context.Context
}

// NewClient returns a client to the Redis Server specified by Options.
func NewClient(opt *Options) *Client {
	opt.init()

	c := Client{
		baseClient: baseClient{
			opt:      opt,
			connPool: newConnPool(opt),
		},
	}
	c.baseClient.init()
	c.init()

	return &c
}

func (c *Client) init() {
	c.cmdable.setProcessor(c.Process)
}

func (c *Client) Context() context.Context {
	if c.ctx != nil {
		return c.ctx
	}
	return context.Background()
}

func (c *Client) WithContext(ctx context.Context) *Client {
	if ctx == nil {
		panic("nil context")
	}
	c2 := c.copy()
	c2.ctx = ctx
	return c2
}

func (c *Client) copy() *Client {
	cp := *c
	cp.init()
	return &cp
}

// Options returns read-only Options that were used to create the client.
func (c *Client) Options() *Options {
	return c.opt
}

type PoolStats pool.Stats

// PoolStats returns connection pool stats.
func (c *Client) PoolStats() *PoolStats {
	stats := c.connPool.Stats()
	return (*PoolStats)(stats)
}

func (c *Client) Pipelined(fn func(Pipeliner) error) ([]Cmder, error) {
	return c.Pipeline().Pipelined(fn)
}

func (c *Client) Pipeline() Pipeliner {
	pipe := Pipeline{
		exec: c.processPipeline,
	}
	pipe.statefulCmdable.setProcessor(pipe.Process)
	return &pipe
}

func (c *Client) TxPipelined(fn func(Pipeliner) error) ([]Cmder, error) {
	return c.TxPipeline().Pipelined(fn)
}

// TxPipeline acts like Pipeline, but wraps queued commands with MULTI/EXEC.
func (c *Client) TxPipeline() Pipeliner {
	pipe := Pipeline{
		exec: c.processTxPipeline,
	}
	pipe.statefulCmdable.setProcessor(pipe.Process)
	return &pipe
}

func (c *Client) pubSub() *PubSub {
	pubsub := &PubSub{
		opt: c.opt,

		newConn: func(channels []string) (*pool.Conn, error) {
			return c.newConn()
		},
		closeConn: c.connPool.CloseConn,
	}
	pubsub.init()
	return pubsub
}

// Subscribe subscribes the client to the specified channels.
// Channels can be omitted to create empty subscription.
func (c *Client) Subscribe(channels ...string) *PubSub {
	pubsub := c.pubSub()
	if len(channels) > 0 {
		_ = pubsub.Subscribe(channels...)
	}
	return pubsub
}

// PSubscribe subscribes the client to the given patterns.
// Patterns can be omitted to create empty subscription.
func (c *Client) PSubscribe(channels ...string) *PubSub {
	pubsub := c.pubSub()
	if len(channels) > 0 {
		_ = pubsub.PSubscribe(channels...)
	}
	return pubsub
}

//------------------------------------------------------------------------------

// Conn is like Client, but its pool contains single connection.
type Conn struct {
	baseClient
	statefulCmdable
}

func newConn(opt *Options, cn *pool.Conn) *Conn {
	c := Conn{
		baseClient: baseClient{
			opt:      opt,
			connPool: pool.NewSingleConnPool(cn),
		},
	}
	c.baseClient.init()
	c.statefulCmdable.setProcessor(c.Process)
	return &c
}

func (c *Conn) Pipelined(fn func(Pipeliner) error) ([]Cmder, error) {
	return c.Pipeline().Pipelined(fn)
}

func (c *Conn) Pipeline() Pipeliner {
	pipe := Pipeline{
		exec: c.processPipeline,
	}
	pipe.statefulCmdable.setProcessor(pipe.Process)
	return &pipe
}

func (c *Conn) TxPipelined(fn func(Pipeliner) error) ([]Cmder, error) {
	return c.TxPipeline().Pipelined(fn)
}

// TxPipeline acts like Pipeline, but wraps queued commands with MULTI/EXEC.
func (c *Conn) TxPipeline() Pipeliner {
	pipe := Pipeline{
		exec: c.processTxPipeline,
	}
	pipe.statefulCmdable.setProcessor(pipe.Process)
	return &pipe
}
<|endoftext|>


<commit_before>// Provides an HTTP router that routes requests based on regex patterns.
//
// Different than http.ServeMux, routing is based on Request.RequestURI (unescaped)
// rather than Request.URL.Path (escaped).
// You can't use http.ServeMux if you want to embed a URL in URL. That's why Remux is written.
//
// Named match groups are saved in the Request.
// You can access them by prefixing the group name with ":".
package remux

import (
	"net/http"
	"net/url"
	"regexp"
)

type Remux struct {
	NotFoundHandler http.Handler
	handlers        []*regexHandler
}

type regexHandler struct {
	pattern        *regexp.Regexp
	allowedMethods map[string]struct{}
	http.Handler
}

// Handle adds a new handler for requests matching regex.
// Panics if regex does not compile.
// When a request received, registered handlers are checked in FIFO order.
func (r *Remux) Handle(regex string, handler http.Handler) *regexHandler {
	pattern := regexp.MustCompile(regex)
	h := &regexHandler{pattern, nil, handler}
	r.handlers = append(r.handlers, h)
	return h
}

func (r *Remux) HandleFunc(regex string, f func(http.ResponseWriter, *http.Request)) *regexHandler {
	return r.Handle(regex, http.HandlerFunc(f))
}

func (r Remux) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	for _, h := range r.handlers {
		submatches := h.pattern.FindStringSubmatch(req.RequestURI)
		if len(submatches) == 0 {
			continue
		}
		if h.allowedMethods != nil {
			if _, ok := h.allowedMethods[req.Method]; !ok {
				w.WriteHeader(http.StatusMethodNotAllowed)
				return
			}
		}
		names := h.pattern.SubexpNames()
		params := make(url.Values)
		for i := range names {
			params.Add(":"+names[i], submatches[i])
		}
		req.URL.RawQuery = params.Encode() + "&" + req.URL.RawQuery
		h.ServeHTTP(w, req)
		return
	}
	if r.NotFoundHandler != nil {
		r.NotFoundHandler.ServeHTTP(w, req)
	} else {
		http.NotFound(w, req)
	}
}

// Method restricts the handler to be executed only on specific method.
func (h *regexHandler) Method(m string) *regexHandler {
	if h.allowedMethods == nil {
		h.allowedMethods = make(map[string]struct{})
	}
	h.allowedMethods[m] = struct{}{}
	return h
}

func (h *regexHandler) Head() *regexHandler    { return h.Method("HEAD") }
func (h *regexHandler) Get() *regexHandler     { return h.Method("GET") }
func (h *regexHandler) Post() *regexHandler    { return h.Method("POST") }
func (h *regexHandler) Put() *regexHandler     { return h.Method("PUT") }
func (h *regexHandler) Delete() *regexHandler  { return h.Method("DELETE") }
func (h *regexHandler) Options() *regexHandler { return h.Method("OPTIONS") }
<commit_msg>export Handler<commit_after>// Provides an HTTP router that routes requests based on regex patterns.
//
// Different than http.ServeMux, routing is based on Request.RequestURI (unescaped)
// rather than Request.URL.Path (escaped).
// You can't use http.ServeMux if you want to embed a URL in URL. That's why Remux is written.
//
// Named match groups are saved in the Request.
// You can access them by prefixing the group name with ":".
package remux

import (
	"net/http"
	"net/url"
	"regexp"
)

type Remux struct {
	NotFoundHandler http.Handler
	handlers        []*Handler
}

// Handler is type of returned value from Remux's Handle and HandleFunc methods.
// You can restrict the handler to run only on specified methods by calling Get, Head, Post, Put, Delete and Options methods.
// These methods are chainable.
type Handler struct {
	pattern        *regexp.Regexp
	allowedMethods map[string]struct{}
	http.Handler
}

// Handle adds a new handler for requests matching regex.
// Panics if regex does not compile.
// When a request received, registered handlers are checked in the same order they are registered.
func (r *Remux) Handle(regex string, handler http.Handler) *Handler {
	pattern := regexp.MustCompile(regex)
	h := &Handler{pattern, nil, handler}
	r.handlers = append(r.handlers, h)
	return h
}

func (r *Remux) HandleFunc(regex string, f func(http.ResponseWriter, *http.Request)) *Handler {
	return r.Handle(regex, http.HandlerFunc(f))
}

func (r Remux) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	for _, h := range r.handlers {
		submatches := h.pattern.FindStringSubmatch(req.RequestURI)
		if len(submatches) == 0 {
			continue
		}
		if h.allowedMethods != nil {
			if _, ok := h.allowedMethods[req.Method]; !ok {
				w.WriteHeader(http.StatusMethodNotAllowed)
				return
			}
		}
		names := h.pattern.SubexpNames()
		params := make(url.Values)
		for i := range names {
			params.Add(":"+names[i], submatches[i])
		}
		req.URL.RawQuery = params.Encode() + "&" + req.URL.RawQuery
		h.ServeHTTP(w, req)
		return
	}
	if r.NotFoundHandler != nil {
		r.NotFoundHandler.ServeHTTP(w, req)
	} else {
		http.NotFound(w, req)
	}
}

// method restricts the handler to be executed only on specific method.
func (h *Handler) method(m string) *Handler {
	if h.allowedMethods == nil {
		h.allowedMethods = make(map[string]struct{})
	}
	h.allowedMethods[m] = struct{}{}
	return h
}

func (h *Handler) Head() *Handler    { return h.method("HEAD") }
func (h *Handler) Get() *Handler     { return h.method("GET") }
func (h *Handler) Post() *Handler    { return h.method("POST") }
func (h *Handler) Put() *Handler     { return h.method("PUT") }
func (h *Handler) Delete() *Handler  { return h.method("DELETE") }
func (h *Handler) Options() *Handler { return h.method("OPTIONS") }
<|endoftext|>


<commit_before>package git

/*
#include <git2.h>
*/
import "C"
import "runtime"

type ResetType int

const (
	ResetSoft  ResetType = C.GIT_RESET_SOFT
	ResetMixed ResetType = C.GIT_RESET_MIXED
	ResetHard  ResetType = C.GIT_RESET_HARD
)

func (r *Repository) ResetToCommit(commit *Commit, resetType ResetType, opts *CheckoutOpts) error {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	ret := C.git_reset(r.ptr, commit.ptr, C.git_reset_t(resetType), opts.toC())

	if ret < 0 {
		return MakeGitError(ret)
	}
	return nil
}
<commit_msg>ResetDefaultToCommit maps git_reset_default<commit_after>package git

/*
#include <git2.h>
*/
import "C"
import "runtime"

type ResetType int

const (
	ResetSoft  ResetType = C.GIT_RESET_SOFT
	ResetMixed ResetType = C.GIT_RESET_MIXED
	ResetHard  ResetType = C.GIT_RESET_HARD
)

func (r *Repository) ResetToCommit(commit *Commit, resetType ResetType, opts *CheckoutOpts) error {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	ret := C.git_reset(r.ptr, commit.ptr, C.git_reset_t(resetType), opts.toC())

	if ret < 0 {
		return MakeGitError(ret)
	}
	return nil
}

func (r *Repository) ResetDefaultToCommit(commit *Commit, pathspecs []string) error {
	cpathspecs := C.git_strarray{}
	cpathspecs.count = C.size_t(len(pathspecs))
	cpathspecs.strings = makeCStringsFromStrings(pathspecs)
	defer freeStrarray(&cpathspecs)

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	ret := C.git_reset_default(r.ptr, commit.ptr, &cpathspecs)

	if ret < 0 {
		return MakeGitError(ret)
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/rooby-lang/rooby/bytecode"
	"github.com/rooby-lang/rooby/parser"
	"github.com/rooby-lang/rooby/vm"
	"io/ioutil"
	"os"
	"path"
	"strings"
)

func main() {
	compileOptionPtr := flag.Bool("c", false, "Compile to bytecode")

	flag.Parse()

	filepath := flag.Arg(0)
	args := flag.Args()[1:]

	var fileExt string
	dir, filename := path.Split(filepath)
	splitedFN := strings.Split(filename, ".")

	if len(splitedFN) <= 1 {
		fmt.Printf("Only support eval/compile single file now.")
		return
	}

	filename = splitedFN[0]
	fileExt = splitedFN[1]

	file, err := ioutil.ReadFile(filepath)
	check(err)

	switch fileExt {
	case "ro":
		program := parser.BuildAST(file)

		g := bytecode.NewGenerator(program)
		bytecodes := g.GenerateByteCode(program)

		if !*compileOptionPtr {
			v := vm.New(dir, args)
			v.ExecBytecodes(bytecodes, filepath)
			return
		}

		writeByteCode(bytecodes, dir, filename)
	case "robc":
		bytecodes := string(file)
		v := vm.New(dir, args)
		v.ExecBytecodes(bytecodes, filepath)
	default:
		fmt.Printf("Unknown file extension: %s", fileExt)
	}
}

func writeByteCode(bytecodes, dir, filename string) {
	f, err := os.Create(dir + filename + ".robc")

	if err != nil {
		panic(err)
	}

	f.WriteString(bytecodes)
}

func check(e error) {
	if e != nil {
		panic(e)
	}
}
<commit_msg>Add option for easy profiling.<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/rooby-lang/rooby/bytecode"
	"github.com/rooby-lang/rooby/parser"
	"github.com/rooby-lang/rooby/vm"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"github.com/pkg/profile"
)

func main() {
	compileOptionPtr := flag.Bool("c", false, "Compile to bytecode")
	profileOptionPtr := flag.Bool("p", false, "Profile program execution")

	flag.Parse()

	filepath := flag.Arg(0)
	args := flag.Args()[1:]

	var fileExt string
	dir, filename := path.Split(filepath)
	splitedFN := strings.Split(filename, ".")

	if len(splitedFN) <= 1 {
		fmt.Printf("Only support eval/compile single file now.")
		return
	}

	filename = splitedFN[0]
	fileExt = splitedFN[1]

	file, err := ioutil.ReadFile(filepath)
	check(err)

	if *profileOptionPtr {
		defer profile.Start().Stop()
	}

	switch fileExt {
	case "ro":
		program := parser.BuildAST(file)

		g := bytecode.NewGenerator(program)
		bytecodes := g.GenerateByteCode(program)

		if !*compileOptionPtr {
			v := vm.New(dir, args)
			v.ExecBytecodes(bytecodes, filepath)
			return
		}

		writeByteCode(bytecodes, dir, filename)
	case "robc":
		bytecodes := string(file)
		v := vm.New(dir, args)
		v.ExecBytecodes(bytecodes, filepath)
	default:
		fmt.Printf("Unknown file extension: %s", fileExt)
	}
}

func writeByteCode(bytecodes, dir, filename string) {
	f, err := os.Create(dir + filename + ".robc")

	if err != nil {
		panic(err)
	}

	f.WriteString(bytecodes)
}

func check(e error) {
	if e != nil {
		panic(e)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"encoding/binary"
	"errors"
	"flag"
	"fmt"
	"log"
	"math"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/dchest/siphash"
	"github.com/hashicorp/golang-lru"
	"github.com/miekg/dns"
)

const (
	// BayesianAverageC Constant for the Bayesian average for the RTT
	BayesianAverageC = 10
	// MaxFailures Maximum number of unanswered queries before a server is marked as dead for VacuumPeriod
	MaxFailures = 10
	// MinTTL Minimum TTL
	MinTTL = 60
	// MaxTTL Maximum TTL
	MaxTTL = 604800
	// VacuumPeriod Vacuum period in seconds
	VacuumPeriod = 30
)

// SipHashKey SipHash secret key
type SipHashKey struct {
	k1 uint64
	k2 uint64
}

// UpstreamServer Upstream server
type UpstreamServer struct {
	addr     string
	failures uint
	offline  bool
}

// UpstreamServers List of upstream servers
type UpstreamServers struct {
	lock    sync.RWMutex
	servers []UpstreamServer
	live    []string
}

// UpstreamRTT Keep track of the mean RTT
type UpstreamRTT struct {
	lock  sync.Mutex
	RTT   float64
	count float64
}

var (
	address            = flag.String("listen", ":53", "Address to listen to (TCP and UDP)")
	upstreamServersStr = flag.String("upstream", "8.8.8.8:53,8.8.4.4:53", "Comma-delimited list of upstream servers")
	upstreamServers    *UpstreamServers
	cacheSize          = flag.Int("cachesize", 10000000, "Cache size")
	memSize            = flag.Uint64("memsize", 1*1024, "Memory size in MB")
	minLabelsCount     = flag.Int("minlabels", 2, "Minimum number of labels")
	cache              *lru.ARCCache
	sipHashKey         = SipHashKey{k1: 0, k2: 0}
	pending            = uint32(0)
	maxClients         = flag.Uint("maxclients", 10000, "Maximum number of simultaneous clients")
	maxRTT             = flag.Float64("maxrtt", 0.25, "Maximum mean RTT for upstream queries before marking a server as dead")
	upstreamRtt        UpstreamRTT
)

func parseUpstreamServers(str string) (*UpstreamServers, error) {
	servers := []UpstreamServer{}
	live := []string{}
	for _, addr := range strings.Split(str, ",") {
		server := UpstreamServer{addr: addr}
		servers = append(servers, server)
		live = append(live, addr)
	}
	res := UpstreamServers{servers: servers, live: live}
	return &res, nil
}

func randUint64() uint64 {
	buf := make([]byte, 8)
	length, err := rand.Read(buf)
	if err != nil || length != len(buf) {
		log.Fatal("RNG failure")
	}
	return binary.LittleEndian.Uint64(buf)
}

func main() {
	flag.Parse()
	*memSize *= 1024 * 1024
	if *cacheSize < 2 {
		log.Fatal("Cache size too small")
	}
	cache, _ = lru.NewARC(*cacheSize)
	upstreamServers, _ = parseUpstreamServers(*upstreamServersStr)
	sipHashKey = SipHashKey{k1: randUint64(), k2: randUint64()}
	fmt.Println("RPDNS")
	dns.HandleFunc(".", route)
	udpServer := &dns.Server{Addr: *address, Net: "udp"}
	tcpServer := &dns.Server{Addr: *address, Net: "tcp"}
	go func() {
		log.Fatal(udpServer.ListenAndServe())
	}()
	go func() {
		log.Fatal(tcpServer.ListenAndServe())
	}()
	vacuumThread()
}

// CacheKey Key for a cache entry
type CacheKey struct {
	Name   string `dns:"cdomain-name"`
	Qtype  uint16
	DNSSEC bool
}

// CacheVal Value for a cache entry
type CacheVal struct {
	ValidUntil time.Time
	Response   *dns.Msg
}

func getKey(req *dns.Msg) (*CacheKey, error) {
	questions := req.Question
	if len(questions) != 1 {
		return nil, errors.New("Invalid number of questions")
	}
	question := questions[0]
	if question.Qclass != dns.ClassINET {
		return nil, errors.New("Unsupported question class")
	}
	if dns.CountLabel(question.Name) < *minLabelsCount {
		return nil, errors.New("Not enough labels")
	}
	dnssec := false
	for _, extra := range req.Extra {
		if extra.Header().Rrtype == dns.TypeOPT {
			dnssec = extra.(*dns.OPT).Do()
		}
	}
	CacheKey := CacheKey{Name: strings.ToLower(question.Name),
		Qtype: question.Qtype, DNSSEC: dnssec}
	return &CacheKey, nil
}

func getMaxPayloadSize(req *dns.Msg) uint16 {
	opt := req.IsEdns0()
	if opt == nil {
		return dns.MinMsgSize
	}
	maxPayloadSize := opt.UDPSize()
	if maxPayloadSize < dns.MinMsgSize {
		maxPayloadSize = dns.MinMsgSize
	}
	return maxPayloadSize
}

func pickUpstream(req *dns.Msg) (*string, error) {
	name := strings.ToLower(req.Question[0].Name)
	h := siphash.Hash(sipHashKey.k1, sipHashKey.k2, []byte(name))
	upstreamServers.lock.RLock()
	liveCount := uint64(len(upstreamServers.live))
	if liveCount <= 0 {
		upstreamServers.lock.RUnlock()
		return nil, errors.New("All upstream servers are down")
	}
	i := h / (math.MaxUint64 / liveCount)
	if i >= liveCount {
		i = liveCount - 1
	}
	res := upstreamServers.live[i]
	upstreamServers.lock.RUnlock()
	return &res, nil
}

func markFailed(addr string) {
	upstreamServers.lock.Lock()
	defer upstreamServers.lock.Unlock()
	for i, server := range upstreamServers.servers {
		if server.addr != addr {
			continue
		}
		if server.offline {
			return
		}
		upstreamServers.servers[i].failures++
		if upstreamServers.servers[i].failures < MaxFailures {
			return
		}
		break
	}
	servers := upstreamServers.servers
	live := []string{}
	for i, server := range upstreamServers.servers {
		if server.addr == addr {
			servers[i].offline = true
		} else if server.offline == false {
			live = append(live, server.addr)
		}
	}
	upstreamServers.servers = servers
	upstreamServers.live = live
	log.Printf("[%v] is unresponsive", addr)
}

func resetRTT() {
	upstreamRtt.lock.Lock()
	defer upstreamRtt.lock.Unlock()
	upstreamRtt.count = 0.0
	upstreamRtt.RTT = 0.0
}

func resetUpstreamServers() {
	upstreamServers.lock.Lock()
	defer upstreamServers.lock.Unlock()
	servers := upstreamServers.servers
	if len(servers) == len(upstreamServers.live) {
		return
	}
	live := []string{}
	for i, server := range upstreamServers.servers {
		servers[i].failures = 0
		servers[i].offline = false
		live = append(live, server.addr)
	}
	upstreamServers.servers = servers
	upstreamServers.live = live
	resetRTT()
}

func syncResolve(req *dns.Msg) (*dns.Msg, time.Duration, error) {
	curPending := atomic.AddUint32(&pending, uint32(1))
	defer atomic.AddUint32(&pending, ^uint32(0))
	if uint(curPending) > *maxClients {
		return nil, 0, errors.New("Too many clients")
	}
	addr, err := pickUpstream(req)
	if err != nil {
		return nil, 0, err
	}
	client := &dns.Client{Net: "udp"}
	client.SingleInflight = true
	resolved, rtt, err := client.Exchange(req, *addr)
	if resolved.Truncated {
		client = &dns.Client{Net: "tcp"}
		client.SingleInflight = true
		resolved, rtt, err = client.Exchange(req, *addr)
	}
	if err != nil {
		markFailed(*addr)
		return nil, 0, err
	}
	upstreamRtt.lock.Lock()
	upstreamRtt.count++
	upstreamRtt.RTT += rtt.Seconds()
	meanRTT := upstreamRtt.RTT / (upstreamRtt.count + BayesianAverageC)
	upstreamRtt.lock.Unlock()
	if meanRTT > *maxRTT {
		markFailed(*addr)
	}
	return resolved, rtt, nil
}

func resolve(req *dns.Msg, dnssec bool) (*dns.Msg, error) {
	extra2 := []dns.RR{}
	for _, extra := range req.Extra {
		if extra.Header().Rrtype != dns.TypeOPT {
			extra2 = append(extra2, extra)
		}
	}
	req.Extra = extra2
	req.SetEdns0(dns.MaxMsgSize, dnssec)
	resolved, _, err := syncResolve(req)
	if err != nil {
		return nil, err
	}
	resolved.Compress = true
	return resolved, nil
}

func getMinTTL(resp *dns.Msg) time.Duration {
	ttl := uint32(MaxTTL)
	for _, rr := range resp.Answer {
		if rr.Header().Ttl < ttl {
			ttl = rr.Header().Ttl
		}
	}
	if ttl < MinTTL {
		ttl = MaxTTL
	}
	return time.Duration(ttl) * time.Second
}

func sendTruncated(w dns.ResponseWriter, msgHdr dns.MsgHdr) {
	emptyResp := new(dns.Msg)
	emptyResp.MsgHdr = msgHdr
	emptyResp.Truncated = true
	w.WriteMsg(emptyResp)
}

func vacuumThread() {
	for {
		time.Sleep(VacuumPeriod * time.Second)
		resetUpstreamServers()
		memStats := new(runtime.MemStats)
		runtime.ReadMemStats(memStats)
		if memStats.Alloc > (*memSize)*1024*1024 {
			cache.Purge()
		}
	}
}

func failWithRcode(w dns.ResponseWriter, r *dns.Msg, rCode int) {
	m := new(dns.Msg)
	m.SetRcode(r, rCode)
	w.WriteMsg(m)
}

func handleSpecialNames(w dns.ResponseWriter, req *dns.Msg) bool {
	question := req.Question[0]
	if question.Qtype != dns.TypeANY {
		return false
	}
	m := new(dns.Msg)
	m.Id = req.Id
	hinfo := new(dns.HINFO)
	hinfo.Hdr = dns.RR_Header{Name: question.Name, Rrtype: dns.TypeHINFO,
		Class: dns.ClassINET, Ttl: 86400}
	hinfo.Cpu = "ANY is not supported any more"
	hinfo.Os = "See draft-jabley-dnsop-refuse-any"
	m.Answer = []dns.RR{hinfo}
	w.WriteMsg(m)
	return true
}

func route(w dns.ResponseWriter, req *dns.Msg) {
	keyP, err := getKey(req)
	if err != nil {
		failWithRcode(w, req, dns.RcodeRefused)
		return
	}
	if handleSpecialNames(w, req) {
		return
	}
	maxPayloadSize := getMaxPayloadSize(req)
	var resp *dns.Msg
	cacheValP, _ := cache.Get(*keyP)
	if cacheValP != nil {
		cacheVal := cacheValP.(CacheVal)
		remaining := -time.Since(cacheVal.ValidUntil)
		if remaining > 0 {
			resp = cacheVal.Response.Copy()
			resp.Id = req.Id
			resp.Question = req.Question
		}
	}
	if resp == nil {
		resp, err = resolve(req, keyP.DNSSEC)
		if err != nil {
			dns.HandleFailed(w, req)
			return
		}
		validUntil := time.Now().Add(getMinTTL(resp))
		cache.Add(*keyP, CacheVal{ValidUntil: validUntil, Response: resp})
	}
	packed, _ := resp.Pack()
	packedLen := len(packed)
	if uint16(packedLen) > maxPayloadSize {
		sendTruncated(w, resp.MsgHdr)
	} else {
		w.WriteMsg(resp)
	}
}
<commit_msg>Truncation is not a crime (2)<commit_after>package main

import (
	"crypto/rand"
	"encoding/binary"
	"errors"
	"flag"
	"fmt"
	"log"
	"math"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/dchest/siphash"
	"github.com/hashicorp/golang-lru"
	"github.com/miekg/dns"
)

const (
	// BayesianAverageC Constant for the Bayesian average for the RTT
	BayesianAverageC = 10
	// MaxFailures Maximum number of unanswered queries before a server is marked as dead for VacuumPeriod
	MaxFailures = 10
	// MinTTL Minimum TTL
	MinTTL = 60
	// MaxTTL Maximum TTL
	MaxTTL = 604800
	// VacuumPeriod Vacuum period in seconds
	VacuumPeriod = 30
)

// SipHashKey SipHash secret key
type SipHashKey struct {
	k1 uint64
	k2 uint64
}

// UpstreamServer Upstream server
type UpstreamServer struct {
	addr     string
	failures uint
	offline  bool
}

// UpstreamServers List of upstream servers
type UpstreamServers struct {
	lock    sync.RWMutex
	servers []UpstreamServer
	live    []string
}

// UpstreamRTT Keep track of the mean RTT
type UpstreamRTT struct {
	lock  sync.Mutex
	RTT   float64
	count float64
}

var (
	address            = flag.String("listen", ":53", "Address to listen to (TCP and UDP)")
	upstreamServersStr = flag.String("upstream", "8.8.8.8:53,8.8.4.4:53", "Comma-delimited list of upstream servers")
	upstreamServers    *UpstreamServers
	cacheSize          = flag.Int("cachesize", 10000000, "Cache size")
	memSize            = flag.Uint64("memsize", 1*1024, "Memory size in MB")
	minLabelsCount     = flag.Int("minlabels", 2, "Minimum number of labels")
	cache              *lru.ARCCache
	sipHashKey         = SipHashKey{k1: 0, k2: 0}
	pending            = uint32(0)
	maxClients         = flag.Uint("maxclients", 10000, "Maximum number of simultaneous clients")
	maxRTT             = flag.Float64("maxrtt", 0.25, "Maximum mean RTT for upstream queries before marking a server as dead")
	upstreamRtt        UpstreamRTT
)

func parseUpstreamServers(str string) (*UpstreamServers, error) {
	servers := []UpstreamServer{}
	live := []string{}
	for _, addr := range strings.Split(str, ",") {
		server := UpstreamServer{addr: addr}
		servers = append(servers, server)
		live = append(live, addr)
	}
	res := UpstreamServers{servers: servers, live: live}
	return &res, nil
}

func randUint64() uint64 {
	buf := make([]byte, 8)
	length, err := rand.Read(buf)
	if err != nil || length != len(buf) {
		log.Fatal("RNG failure")
	}
	return binary.LittleEndian.Uint64(buf)
}

func main() {
	flag.Parse()
	*memSize *= 1024 * 1024
	if *cacheSize < 2 {
		log.Fatal("Cache size too small")
	}
	cache, _ = lru.NewARC(*cacheSize)
	upstreamServers, _ = parseUpstreamServers(*upstreamServersStr)
	sipHashKey = SipHashKey{k1: randUint64(), k2: randUint64()}
	fmt.Println("RPDNS")
	dns.HandleFunc(".", route)
	udpServer := &dns.Server{Addr: *address, Net: "udp"}
	tcpServer := &dns.Server{Addr: *address, Net: "tcp"}
	go func() {
		log.Fatal(udpServer.ListenAndServe())
	}()
	go func() {
		log.Fatal(tcpServer.ListenAndServe())
	}()
	vacuumThread()
}

// CacheKey Key for a cache entry
type CacheKey struct {
	Name   string `dns:"cdomain-name"`
	Qtype  uint16
	DNSSEC bool
}

// CacheVal Value for a cache entry
type CacheVal struct {
	ValidUntil time.Time
	Response   *dns.Msg
}

func getKey(req *dns.Msg) (*CacheKey, error) {
	questions := req.Question
	if len(questions) != 1 {
		return nil, errors.New("Invalid number of questions")
	}
	question := questions[0]
	if question.Qclass != dns.ClassINET {
		return nil, errors.New("Unsupported question class")
	}
	if dns.CountLabel(question.Name) < *minLabelsCount {
		return nil, errors.New("Not enough labels")
	}
	dnssec := false
	for _, extra := range req.Extra {
		if extra.Header().Rrtype == dns.TypeOPT {
			dnssec = extra.(*dns.OPT).Do()
		}
	}
	CacheKey := CacheKey{Name: strings.ToLower(question.Name),
		Qtype: question.Qtype, DNSSEC: dnssec}
	return &CacheKey, nil
}

func getMaxPayloadSize(req *dns.Msg) uint16 {
	opt := req.IsEdns0()
	if opt == nil {
		return dns.MinMsgSize
	}
	maxPayloadSize := opt.UDPSize()
	if maxPayloadSize < dns.MinMsgSize {
		maxPayloadSize = dns.MinMsgSize
	}
	return maxPayloadSize
}

func pickUpstream(req *dns.Msg) (*string, error) {
	name := strings.ToLower(req.Question[0].Name)
	h := siphash.Hash(sipHashKey.k1, sipHashKey.k2, []byte(name))
	upstreamServers.lock.RLock()
	liveCount := uint64(len(upstreamServers.live))
	if liveCount <= 0 {
		upstreamServers.lock.RUnlock()
		return nil, errors.New("All upstream servers are down")
	}
	i := h / (math.MaxUint64 / liveCount)
	if i >= liveCount {
		i = liveCount - 1
	}
	res := upstreamServers.live[i]
	upstreamServers.lock.RUnlock()
	return &res, nil
}

func markFailed(addr string) {
	upstreamServers.lock.Lock()
	defer upstreamServers.lock.Unlock()
	for i, server := range upstreamServers.servers {
		if server.addr != addr {
			continue
		}
		if server.offline {
			return
		}
		upstreamServers.servers[i].failures++
		if upstreamServers.servers[i].failures < MaxFailures {
			return
		}
		break
	}
	servers := upstreamServers.servers
	live := []string{}
	for i, server := range upstreamServers.servers {
		if server.addr == addr {
			servers[i].offline = true
		} else if server.offline == false {
			live = append(live, server.addr)
		}
	}
	upstreamServers.servers = servers
	upstreamServers.live = live
	log.Printf("[%v] is unresponsive", addr)
}

func resetRTT() {
	upstreamRtt.lock.Lock()
	defer upstreamRtt.lock.Unlock()
	upstreamRtt.count = 0.0
	upstreamRtt.RTT = 0.0
}

func resetUpstreamServers() {
	upstreamServers.lock.Lock()
	defer upstreamServers.lock.Unlock()
	servers := upstreamServers.servers
	if len(servers) == len(upstreamServers.live) {
		return
	}
	live := []string{}
	for i, server := range upstreamServers.servers {
		servers[i].failures = 0
		servers[i].offline = false
		live = append(live, server.addr)
	}
	upstreamServers.servers = servers
	upstreamServers.live = live
	resetRTT()
}

func syncResolve(req *dns.Msg) (*dns.Msg, time.Duration, error) {
	curPending := atomic.AddUint32(&pending, uint32(1))
	defer atomic.AddUint32(&pending, ^uint32(0))
	if uint(curPending) > *maxClients {
		return nil, 0, errors.New("Too many clients")
	}
	addr, err := pickUpstream(req)
	if err != nil {
		return nil, 0, err
	}
	client := &dns.Client{Net: "udp"}
	client.SingleInflight = true
	resolved, rtt, err := client.Exchange(req, *addr)
	if resolved != nil && resolved.Truncated {
		client = &dns.Client{Net: "tcp"}
		client.SingleInflight = true
		resolved, rtt, err = client.Exchange(req, *addr)
	}
	if err != nil {
		markFailed(*addr)
		return nil, 0, err
	}
	upstreamRtt.lock.Lock()
	upstreamRtt.count++
	upstreamRtt.RTT += rtt.Seconds()
	meanRTT := upstreamRtt.RTT / (upstreamRtt.count + BayesianAverageC)
	upstreamRtt.lock.Unlock()
	if meanRTT > *maxRTT {
		markFailed(*addr)
	}
	return resolved, rtt, nil
}

func resolve(req *dns.Msg, dnssec bool) (*dns.Msg, error) {
	extra2 := []dns.RR{}
	for _, extra := range req.Extra {
		if extra.Header().Rrtype != dns.TypeOPT {
			extra2 = append(extra2, extra)
		}
	}
	req.Extra = extra2
	req.SetEdns0(dns.MaxMsgSize, dnssec)
	resolved, _, err := syncResolve(req)
	if err != nil {
		return nil, err
	}
	resolved.Compress = true
	return resolved, nil
}

func getMinTTL(resp *dns.Msg) time.Duration {
	ttl := uint32(MaxTTL)
	for _, rr := range resp.Answer {
		if rr.Header().Ttl < ttl {
			ttl = rr.Header().Ttl
		}
	}
	if ttl < MinTTL {
		ttl = MaxTTL
	}
	return time.Duration(ttl) * time.Second
}

func sendTruncated(w dns.ResponseWriter, msgHdr dns.MsgHdr) {
	emptyResp := new(dns.Msg)
	emptyResp.MsgHdr = msgHdr
	emptyResp.Truncated = true
	w.WriteMsg(emptyResp)
}

func vacuumThread() {
	for {
		time.Sleep(VacuumPeriod * time.Second)
		resetUpstreamServers()
		memStats := new(runtime.MemStats)
		runtime.ReadMemStats(memStats)
		if memStats.Alloc > (*memSize)*1024*1024 {
			cache.Purge()
		}
	}
}

func failWithRcode(w dns.ResponseWriter, r *dns.Msg, rCode int) {
	m := new(dns.Msg)
	m.SetRcode(r, rCode)
	w.WriteMsg(m)
}

func handleSpecialNames(w dns.ResponseWriter, req *dns.Msg) bool {
	question := req.Question[0]
	if question.Qtype != dns.TypeANY {
		return false
	}
	m := new(dns.Msg)
	m.Id = req.Id
	hinfo := new(dns.HINFO)
	hinfo.Hdr = dns.RR_Header{Name: question.Name, Rrtype: dns.TypeHINFO,
		Class: dns.ClassINET, Ttl: 86400}
	hinfo.Cpu = "ANY is not supported any more"
	hinfo.Os = "See draft-jabley-dnsop-refuse-any"
	m.Answer = []dns.RR{hinfo}
	w.WriteMsg(m)
	return true
}

func route(w dns.ResponseWriter, req *dns.Msg) {
	keyP, err := getKey(req)
	if err != nil {
		failWithRcode(w, req, dns.RcodeRefused)
		return
	}
	if handleSpecialNames(w, req) {
		return
	}
	maxPayloadSize := getMaxPayloadSize(req)
	var resp *dns.Msg
	cacheValP, _ := cache.Get(*keyP)
	if cacheValP != nil {
		cacheVal := cacheValP.(CacheVal)
		remaining := -time.Since(cacheVal.ValidUntil)
		if remaining > 0 {
			resp = cacheVal.Response.Copy()
			resp.Id = req.Id
			resp.Question = req.Question
		}
	}
	if resp == nil {
		resp, err = resolve(req, keyP.DNSSEC)
		if err != nil {
			dns.HandleFailed(w, req)
			return
		}
		validUntil := time.Now().Add(getMinTTL(resp))
		cache.Add(*keyP, CacheVal{ValidUntil: validUntil, Response: resp})
	}
	packed, _ := resp.Pack()
	packedLen := len(packed)
	if uint16(packedLen) > maxPayloadSize {
		sendTruncated(w, resp.MsgHdr)
	} else {
		w.WriteMsg(resp)
	}
}
<|endoftext|>


<commit_before>// Copyright © 2015 Hilko Bengen <bengen@hilluzination.de>. All rights reserved.
// Use of this source code is governed by the license that can be
// found in the LICENSE file.

// Package yara provides bindings to the YARA library.
package yara

/*
#cgo LDFLAGS: -lyara
#include <yara.h>

int callback(int message, void *message_data, void *user_data);
*/
import "C"
import (
	"errors"
	"runtime"
	"time"
	"unsafe"
)

// Rules contains a compiled YARA ruleset.
type Rules struct {
	r *C.YR_RULES
}

// A MatchRule represents a rule successfully matched against a block
// of data.
type MatchRule struct {
	Rule      string
	Namespace string
	Tags      []string
	Meta      map[string]interface{}
	Strings   []MatchString
}

// A MatchString represents a string declared and matched in a rule.
type MatchString struct {
	Name   string
	Offset uint64
	Data   []byte
}

func init() {
	_ = C.yr_initialize()
}

//export newMatch
func newMatch(matches *[]MatchRule, namespace, identifier *C.char) {
	*matches = append(*matches, MatchRule{
		Rule:      C.GoString(identifier),
		Namespace: C.GoString(namespace),
		Tags:      []string{},
		Meta:      map[string]interface{}{},
		Strings:   []MatchString{},
	})
}

//export addMetaInt
func addMetaInt(matches *[]MatchRule, identifier *C.char, value C.int) {
	(*matches)[len(*matches)-1].Meta[C.GoString(identifier)] = int32(value)
}

//export addMetaString
func addMetaString(matches *[]MatchRule, identifier *C.char, value *C.char) {
	(*matches)[len(*matches)-1].Meta[C.GoString(identifier)] = C.GoString(value)
}

//export addMetaBool
func addMetaBool(matches *[]MatchRule, identifier *C.char, value C.int) {
	(*matches)[len(*matches)-1].Meta[C.GoString(identifier)] = bool(value != 0)
}

//export addTag
func addTag(matches *[]MatchRule, tag *C.char) {
	(*matches)[len(*matches)-1].Tags = append((*matches)[len(*matches)-1].Tags, C.GoString(tag))
}

//export addString
func addString(matches *[]MatchRule, identifier *C.char, offset C.uint64_t, data unsafe.Pointer, length C.int) {
	(*matches)[len(*matches)-1].Strings = append(
		(*matches)[len(*matches)-1].Strings,
		MatchString{
			Name:   C.GoString(identifier),
			Offset: uint64(offset),
			Data:   C.GoBytes(data, length),
		})
}

// ScanFlags are used to tweak the behavior of Scan* functions.
type ScanFlags int

const (
	// ScanFlagsFastMode avoids multiple matches of the same string when not necessary.
	ScanFlagsFastMode = C.SCAN_FLAGS_FAST_MODE
	// ScanFlagsProcessMemory causes the scanned data to be
	// interpreted like live, in-prcess memory rather than an on-disk
	// file.
	ScanFlagsProcessMemory = C.SCAN_FLAGS_PROCESS_MEMORY
)

// ScanMem scans an in-memory buffer using the ruleset.
func (r *Rules) ScanMem(buf []byte, flags ScanFlags, timeout time.Duration) (matches []MatchRule, err error) {
	err = newError(C.yr_rules_scan_mem(
		r.r,
		(*C.uint8_t)(unsafe.Pointer(&(buf[0]))),
		C.size_t(len(buf)),
		C.int(flags),
		C.YR_CALLBACK_FUNC(C.callback),
		unsafe.Pointer(&matches),
		C.int(timeout)))
	return
}

// ScanFile scans a file using the ruleset.
func (r *Rules) ScanFile(filename string, flags ScanFlags, timeout time.Duration) (matches []MatchRule, err error) {
	err = newError(C.yr_rules_scan_file(
		r.r,
		C.CString(filename),
		C.int(flags),
		C.YR_CALLBACK_FUNC(C.callback),
		unsafe.Pointer(&matches),
		C.int(timeout)))
	return
}

// ScanProc scans a live process using the ruleset.
func (r *Rules) ScanProc(pid int, flags int, timeout time.Duration) (matches []MatchRule, err error) {
	err = newError(C.yr_rules_scan_proc(
		r.r,
		C.int(pid),
		C.int(flags),
		C.YR_CALLBACK_FUNC(C.callback),
		unsafe.Pointer(&matches),
		C.int(timeout)))
	return
}

// Save writes a compiled ruleset to filename.
func (r *Rules) Save(filename string) (err error) {
	err = newError(C.yr_rules_save(r.r, C.CString(filename)))
	return
}

// LoadRules retrieves compiled ruleset from filename.
func LoadRules(filename string) (rules *Rules, err error) {
	var r *C.YR_RULES
	err = newError(C.yr_rules_load(C.CString(filename), &r))
	if err == nil {
		rules = &Rules{r: r}
		runtime.SetFinalizer(rules, func(r *Rules) {
			C.yr_rules_destroy(r.r)
			r.r = nil
		})
	}
	return
}

// DefineVariable defines a named variable for use by the compiler.
// Boolean, int64, and string types are supported.
func (r *Rules) DefineVariable(name string, value interface{}) (err error) {
	switch value.(type) {
	case bool:
		var v int
		if value.(bool) {
			v = 1
		}
		err = newError(C.yr_rules_define_boolean_variable(
			r.r, C.CString(name), C.int(v)))
	case int64:
		err = newError(C.yr_rules_define_integer_variable(
			r.r, C.CString(name), C.int64_t(value.(int64))))
	case string:
		err = newError(C.yr_rules_define_string_variable(
			r.r, C.CString(name), C.CString(value.(string))))
	default:
		err = errors.New("wrong value type passed to DefineVariable; bool, int64, string are accepted.")
	}
	return
}
<commit_msg>yr_rules_scan_* functions expect timeout to be specified as seconds: Convert time.Duration accordingly.<commit_after>// Copyright © 2015 Hilko Bengen <bengen@hilluzination.de>. All rights reserved.
// Use of this source code is governed by the license that can be
// found in the LICENSE file.

// Package yara provides bindings to the YARA library.
package yara

/*
#cgo LDFLAGS: -lyara
#include <yara.h>

int callback(int message, void *message_data, void *user_data);
*/
import "C"
import (
	"errors"
	"runtime"
	"time"
	"unsafe"
)

// Rules contains a compiled YARA ruleset.
type Rules struct {
	r *C.YR_RULES
}

// A MatchRule represents a rule successfully matched against a block
// of data.
type MatchRule struct {
	Rule      string
	Namespace string
	Tags      []string
	Meta      map[string]interface{}
	Strings   []MatchString
}

// A MatchString represents a string declared and matched in a rule.
type MatchString struct {
	Name   string
	Offset uint64
	Data   []byte
}

func init() {
	_ = C.yr_initialize()
}

//export newMatch
func newMatch(matches *[]MatchRule, namespace, identifier *C.char) {
	*matches = append(*matches, MatchRule{
		Rule:      C.GoString(identifier),
		Namespace: C.GoString(namespace),
		Tags:      []string{},
		Meta:      map[string]interface{}{},
		Strings:   []MatchString{},
	})
}

//export addMetaInt
func addMetaInt(matches *[]MatchRule, identifier *C.char, value C.int) {
	(*matches)[len(*matches)-1].Meta[C.GoString(identifier)] = int32(value)
}

//export addMetaString
func addMetaString(matches *[]MatchRule, identifier *C.char, value *C.char) {
	(*matches)[len(*matches)-1].Meta[C.GoString(identifier)] = C.GoString(value)
}

//export addMetaBool
func addMetaBool(matches *[]MatchRule, identifier *C.char, value C.int) {
	(*matches)[len(*matches)-1].Meta[C.GoString(identifier)] = bool(value != 0)
}

//export addTag
func addTag(matches *[]MatchRule, tag *C.char) {
	(*matches)[len(*matches)-1].Tags = append((*matches)[len(*matches)-1].Tags, C.GoString(tag))
}

//export addString
func addString(matches *[]MatchRule, identifier *C.char, offset C.uint64_t, data unsafe.Pointer, length C.int) {
	(*matches)[len(*matches)-1].Strings = append(
		(*matches)[len(*matches)-1].Strings,
		MatchString{
			Name:   C.GoString(identifier),
			Offset: uint64(offset),
			Data:   C.GoBytes(data, length),
		})
}

// ScanFlags are used to tweak the behavior of Scan* functions.
type ScanFlags int

const (
	// ScanFlagsFastMode avoids multiple matches of the same string when not necessary.
	ScanFlagsFastMode = C.SCAN_FLAGS_FAST_MODE
	// ScanFlagsProcessMemory causes the scanned data to be
	// interpreted like live, in-prcess memory rather than an on-disk
	// file.
	ScanFlagsProcessMemory = C.SCAN_FLAGS_PROCESS_MEMORY
)

// ScanMem scans an in-memory buffer using the ruleset.
func (r *Rules) ScanMem(buf []byte, flags ScanFlags, timeout time.Duration) (matches []MatchRule, err error) {
	err = newError(C.yr_rules_scan_mem(
		r.r,
		(*C.uint8_t)(unsafe.Pointer(&(buf[0]))),
		C.size_t(len(buf)),
		C.int(flags),
		C.YR_CALLBACK_FUNC(C.callback),
		unsafe.Pointer(&matches),
		C.int(timeout/time.Second)))
	return
}

// ScanFile scans a file using the ruleset.
func (r *Rules) ScanFile(filename string, flags ScanFlags, timeout time.Duration) (matches []MatchRule, err error) {
	err = newError(C.yr_rules_scan_file(
		r.r,
		C.CString(filename),
		C.int(flags),
		C.YR_CALLBACK_FUNC(C.callback),
		unsafe.Pointer(&matches),
		C.int(timeout/time.Second)))
	return
}

// ScanProc scans a live process using the ruleset.
func (r *Rules) ScanProc(pid int, flags int, timeout time.Duration) (matches []MatchRule, err error) {
	err = newError(C.yr_rules_scan_proc(
		r.r,
		C.int(pid),
		C.int(flags),
		C.YR_CALLBACK_FUNC(C.callback),
		unsafe.Pointer(&matches),
		C.int(timeout/time.Second)))
	return
}

// Save writes a compiled ruleset to filename.
func (r *Rules) Save(filename string) (err error) {
	err = newError(C.yr_rules_save(r.r, C.CString(filename)))
	return
}

// LoadRules retrieves compiled ruleset from filename.
func LoadRules(filename string) (rules *Rules, err error) {
	var r *C.YR_RULES
	err = newError(C.yr_rules_load(C.CString(filename), &r))
	if err == nil {
		rules = &Rules{r: r}
		runtime.SetFinalizer(rules, func(r *Rules) {
			C.yr_rules_destroy(r.r)
			r.r = nil
		})
	}
	return
}

// DefineVariable defines a named variable for use by the compiler.
// Boolean, int64, and string types are supported.
func (r *Rules) DefineVariable(name string, value interface{}) (err error) {
	switch value.(type) {
	case bool:
		var v int
		if value.(bool) {
			v = 1
		}
		err = newError(C.yr_rules_define_boolean_variable(
			r.r, C.CString(name), C.int(v)))
	case int64:
		err = newError(C.yr_rules_define_integer_variable(
			r.r, C.CString(name), C.int64_t(value.(int64))))
	case string:
		err = newError(C.yr_rules_define_string_variable(
			r.r, C.CString(name), C.CString(value.(string))))
	default:
		err = errors.New("wrong value type passed to DefineVariable; bool, int64, string are accepted.")
	}
	return
}
<|endoftext|>


<commit_before>/*
   Copyright 2010-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.

   This file is licensed under the Apache License, Version 2.0 (the "License").
   You may not use this file except in compliance with the License. A copy of
   the License is located at

    http://aws.amazon.com/apache2.0/

   This file 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 main

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sqs"
)

// Usage:
// go run sqs_receive_message.go
func main() {
	sess := session.Must(session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
	}))

	svc := sqs.New(sess)

	// URL to our queue
	qURL := "QueueURL"

	result, err := svc.ReceiveMessage(&sqs.ReceiveMessageInput{
		AttributeNames: []*string{
			aws.String(sqs.MessageSystemAttributeNameSentTimestamp),
		},
		MessageAttributeNames: []*string{
			aws.String(sqs.QueueAttributeNameAll),
		},
		QueueUrl:            &qURL,
		MaxNumberOfMessages: aws.Int64(10),
		VisibilityTimeout:   aws.Int64(60), // 60 seconds
		WaitTimeSeconds:     aws.Int64(0),
	})
	if err != nil {
		fmt.Println("Error", err)
		return
	}
	if len(result.Messages) == 0 {
		fmt.Println("Received no messages")
		return
	}

	fmt.Printf("Success: %+v\n", result.Messages)
}
<commit_msg>Update sqs_receive_message.go<commit_after>// snippet-comment:[These are tags for the AWS doc team's sample catalog. Do not remove.]
// snippet-sourceauthor:[AWS]
// snippet-sourcedescription:[Receives an SQS message.]
// snippet-keyword:[Amazon Simple Queue Service]
// snippet-keyword:[Amazon SQS]
// snippet-keyword:[ReceiveMessage function]
// snippet-keyword:[Go]
// snippet-service:[sqs]
// snippet-keyword:[Code Sample]
// snippet-sourcetype:[full-example]
// snippet-sourcedate:[2019-01-28]
/*
   Copyright 2010-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.

   This file is licensed under the Apache License, Version 2.0 (the "License").
   You may not use this file except in compliance with the License. A copy of
   the License is located at

    http://aws.amazon.com/apache2.0/

   This file 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.
*/
// snippet-start:[sqs.go.receive_message]
package main

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sqs"
)

// Usage:
// go run sqs_receive_message.go
func main() {
	sess := session.Must(session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
	}))

	svc := sqs.New(sess)

	// URL to our queue
	qURL := "QueueURL"

	result, err := svc.ReceiveMessage(&sqs.ReceiveMessageInput{
		AttributeNames: []*string{
			aws.String(sqs.MessageSystemAttributeNameSentTimestamp),
		},
		MessageAttributeNames: []*string{
			aws.String(sqs.QueueAttributeNameAll),
		},
		QueueUrl:            &qURL,
		MaxNumberOfMessages: aws.Int64(10),
		VisibilityTimeout:   aws.Int64(60), // 60 seconds
		WaitTimeSeconds:     aws.Int64(0),
	})
	if err != nil {
		fmt.Println("Error", err)
		return
	}
	if len(result.Messages) == 0 {
		fmt.Println("Received no messages")
		return
	}

	fmt.Printf("Success: %+v\n", result.Messages)
}
// snippet-end:[sqs.go.receive_message]
<|endoftext|>


<commit_before>package vnc

import (
	"fmt"
	"log"
	"net"
	"reflect"
	"testing"

	"golang.org/x/net/context"
)

func newMockServer(t *testing.T, version string) string {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("error listening: %s", err)
	}

	go func() {
		defer ln.Close()
		c, err := ln.Accept()
		if err != nil {
			t.Fatalf("error accepting conn: %s", err)
		}
		defer c.Close()

		_, err = c.Write([]byte(fmt.Sprintf("RFB %s\n", version)))
		if err != nil {
			t.Fatal("failed writing version")
		}
	}()

	return ln.Addr().String()
}

func TestLowMajorVersion(t *testing.T) {
	nc, err := net.Dial("tcp", newMockServer(t, "002.009"))
	if err != nil {
		t.Fatalf("error connecting to mock server: %s", err)
	}

	_, err = Connect(context.Background(), nc, &ClientConfig{})
	if err == nil {
		t.Fatal("error expected")
	}
	if err != nil {
		if verr, ok := err.(*VNCError); !ok {
			t.Errorf("Client() unexpected %v error: %v", reflect.TypeOf(err), verr)
		}
	}
}

func TestLowMinorVersion(t *testing.T) {
	nc, err := net.Dial("tcp", newMockServer(t, "003.002"))
	if err != nil {
		t.Fatalf("error connecting to mock server: %s", err)
	}

	_, err = Connect(context.Background(), nc, &ClientConfig{})
	if err == nil {
		t.Fatal("error expected")
	}
	if err != nil {
		if verr, ok := err.(*VNCError); !ok {
			t.Errorf("Client() unexpected %v error: %v", reflect.TypeOf(err), verr)
		}
	}
}

func ExampleConnect() {
	// Establish TCP connection to VNC server.
	nc, err := net.Dial("tcp", "127.0.0.1:5900")
	if err != nil {
		log.Fatalf("Error connecting to VNC host. %v", err)
	}

	// Negotiate connection with the server.
	vcc := NewClientConfig("password")
	vc, err := Connect(context.Background(), nc, vcc)
	if err != nil {
		log.Fatalf("Error negotiating connection to VNC host. %v", err)
	}

	// Listen and handle server messages.
	go vc.ListenAndHandle()
	for {
		msg := <-vcc.ServerMessageCh
		switch msg.Type() {
		case FramebufferUpdate:
			log.Println("Received FramebufferUpdate message.")
		default:
			log.Printf("Received message type:%v msg:%v\n", msg.Type(), msg)
		}
	}
}
<commit_msg>added basic tests for ClientConfig{}<commit_after>package vnc

import (
	"fmt"
	"log"
	"net"
	"reflect"
	"testing"

	"golang.org/x/net/context"
)

func newMockServer(t *testing.T, version string) string {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("error listening: %s", err)
	}

	go func() {
		defer ln.Close()
		c, err := ln.Accept()
		if err != nil {
			t.Fatalf("error accepting conn: %s", err)
		}
		defer c.Close()

		_, err = c.Write([]byte(fmt.Sprintf("RFB %s\n", version)))
		if err != nil {
			t.Fatal("failed writing version")
		}
	}()

	return ln.Addr().String()
}

func TestLowMajorVersion(t *testing.T) {
	nc, err := net.Dial("tcp", newMockServer(t, "002.009"))
	if err != nil {
		t.Fatalf("error connecting to mock server: %s", err)
	}

	_, err = Connect(context.Background(), nc, &ClientConfig{})
	if err == nil {
		t.Fatal("error expected")
	}
	if err != nil {
		if verr, ok := err.(*VNCError); !ok {
			t.Errorf("Client() unexpected %v error: %v", reflect.TypeOf(err), verr)
		}
	}
}

func TestLowMinorVersion(t *testing.T) {
	nc, err := net.Dial("tcp", newMockServer(t, "003.002"))
	if err != nil {
		t.Fatalf("error connecting to mock server: %s", err)
	}

	_, err = Connect(context.Background(), nc, &ClientConfig{})
	if err == nil {
		t.Fatal("error expected")
	}
	if err != nil {
		if verr, ok := err.(*VNCError); !ok {
			t.Errorf("Client() unexpected %v error: %v", reflect.TypeOf(err), verr)
		}
	}
}

func TestClientConn(t *testing.T) {
	conn := &ClientConn{}

	if got, want := conn.DesktopName(), ""; got != want {
		t.Errorf("DesktopName() failed; got = %v, want = %v", got, want)
	}
	if got, want := conn.FramebufferHeight(), uint16(0); got != want {
		t.Errorf("FramebufferHeight() failed; got = %v, want = %v", got, want)
	}
	if got, want := conn.FramebufferWidth(), uint16(0); got != want {
		t.Errorf("FramebufferWidth() failed; got = %v, want = %v", got, want)
	}
}

func ExampleConnect() {
	// Establish TCP connection to VNC server.
	nc, err := net.Dial("tcp", "127.0.0.1:5900")
	if err != nil {
		log.Fatalf("Error connecting to VNC host. %v", err)
	}

	// Negotiate connection with the server.
	vcc := NewClientConfig("password")
	vc, err := Connect(context.Background(), nc, vcc)
	if err != nil {
		log.Fatalf("Error negotiating connection to VNC host. %v", err)
	}

	// Listen and handle server messages.
	go vc.ListenAndHandle()
	for {
		msg := <-vcc.ServerMessageCh
		switch msg.Type() {
		case FramebufferUpdate:
			log.Println("Received FramebufferUpdate message.")
		default:
			log.Printf("Received message type:%v msg:%v\n", msg.Type(), msg)
		}
	}
}
<|endoftext|>


<commit_before>package utils

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/golang/glog"
	jsonnet "github.com/strickyak/jsonnet_cgo"
	"k8s.io/apimachinery/pkg/util/yaml"
	"k8s.io/client-go/pkg/runtime"
)

// Read fetches and decodes K8s objects by path.
// TODO: Replace this with something supporting more sophisticated
// content negotiation.
func Read(vm *jsonnet.VM, path string) ([]runtime.Object, error) {
	ext := filepath.Ext(path)
	if ext == ".json" {
		f, err := os.Open(path)
		if err != nil {
			return nil, err
		}
		defer f.Close()
		return jsonReader(f)
	} else if ext == ".yaml" {
		f, err := os.Open(path)
		if err != nil {
			return nil, err
		}
		defer f.Close()
		return yamlReader(f)
	} else if ext == ".jsonnet" {
		return jsonnetReader(vm, path)
	}

	return nil, fmt.Errorf("Unknown file extension: %s", path)
}

func jsonReader(r io.Reader) ([]runtime.Object, error) {
	data, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, err
	}
	obj, _, err := runtime.UnstructuredJSONScheme.Decode(data, nil, nil)
	if err != nil {
		return nil, err
	}
	return []runtime.Object{obj}, nil
}

func yamlReader(r io.ReadCloser) ([]runtime.Object, error) {
	decoder := yaml.NewDocumentDecoder(r)
	ret := []runtime.Object{}
	buf := []byte{}
	for {
		_, err := decoder.Read(buf)
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, err
		}
		jsondata, err := yaml.ToJSON(buf)
		if err != nil {
			return nil, err
		}
		obj, _, err := runtime.UnstructuredJSONScheme.Decode(jsondata, nil, nil)
		if err != nil {
			return nil, err
		}
		ret = append(ret, obj)
	}
	return ret, nil
}

func jsonWalk(obj interface{}) ([]interface{}, error) {
	switch o := obj.(type) {
	case map[string]interface{}:
		if o["kind"] != nil && o["apiVersion"] != nil {
			return []interface{}{o}, nil
		}
		ret := []interface{}{}
		for _, v := range o {
			children, err := jsonWalk(v)
			if err != nil {
				return nil, err
			}
			ret = append(ret, children...)
		}
		return ret, nil
	case []interface{}:
		ret := make([]interface{}, 0, len(o))
		for _, v := range o {
			children, err := jsonWalk(v)
			if err != nil {
				return nil, err
			}
			ret = append(ret, children...)
		}
		return ret, nil
	default:
		return nil, fmt.Errorf("Unexpected object structure: %T", o)
	}
}

func jsonnetReader(vm *jsonnet.VM, path string) ([]runtime.Object, error) {
	jsonstr, err := vm.EvaluateFile(path)
	if err != nil {
		return nil, err
	}

	glog.V(4).Infof("jsonnet result is: %s\n", jsonstr)

	var top interface{}
	if err = json.Unmarshal([]byte(jsonstr), &top); err != nil {
		return nil, err
	}

	objs, err := jsonWalk(top)
	if err != nil {
		return nil, err
	}

	ret := make([]runtime.Object, 0, len(objs))
	for _, v := range objs {
		// TODO: Going to json and back is a bit horrible
		data, err := json.Marshal(v)
		if err != nil {
			return nil, err
		}
		obj, _, err := runtime.UnstructuredJSONScheme.Decode(data, nil, nil)
		if err != nil {
			return nil, err
		}
		ret = append(ret, obj)
	}

	return ret, nil
}

// FlattenToV1 expands any List-type objects into their members, and
// cooerces everything to v1.Unstructured.  Panics if coercion
// encounters an unexpected object type.
func FlattenToV1(objs []runtime.Object) []*runtime.Unstructured {
	ret := make([]*runtime.Unstructured, 0, len(objs))
	for _, obj := range objs {
		switch o := obj.(type) {
		case *runtime.UnstructuredList:
			for _, item := range o.Items {
				ret = append(ret, item)
			}
		case *runtime.Unstructured:
			ret = append(ret, o)
		default:
			panic("Unexpected unstructured object type")
		}
	}
	return ret
}
<commit_msg>Fix reading YAML files<commit_after>package utils

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/golang/glog"
	jsonnet "github.com/strickyak/jsonnet_cgo"
	"k8s.io/client-go/pkg/runtime"
	"k8s.io/client-go/pkg/util/yaml"
)

// Read fetches and decodes K8s objects by path.
// TODO: Replace this with something supporting more sophisticated
// content negotiation.
func Read(vm *jsonnet.VM, path string) ([]runtime.Object, error) {
	ext := filepath.Ext(path)
	if ext == ".json" {
		f, err := os.Open(path)
		if err != nil {
			return nil, err
		}
		defer f.Close()
		return jsonReader(f)
	} else if ext == ".yaml" {
		f, err := os.Open(path)
		if err != nil {
			return nil, err
		}
		defer f.Close()
		return yamlReader(f)
	} else if ext == ".jsonnet" {
		return jsonnetReader(vm, path)
	}

	return nil, fmt.Errorf("Unknown file extension: %s", path)
}

func jsonReader(r io.Reader) ([]runtime.Object, error) {
	data, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, err
	}
	obj, _, err := runtime.UnstructuredJSONScheme.Decode(data, nil, nil)
	if err != nil {
		return nil, err
	}
	return []runtime.Object{obj}, nil
}

func yamlReader(r io.ReadCloser) ([]runtime.Object, error) {
	decoder := yaml.NewYAMLReader(bufio.NewReader(r))
	ret := []runtime.Object{}
	for {
		bytes, err := decoder.Read()
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, err
		}
		glog.V(4).Infof("Read %d bytes of YAML: %s", len(bytes), bytes)
		if len(bytes) == 0 {
			continue
		}
		jsondata, err := yaml.ToJSON(bytes)
		if err != nil {
			return nil, err
		}
		glog.V(4).Infof("Converted to JSON: %s", jsondata)
		obj, _, err := runtime.UnstructuredJSONScheme.Decode(jsondata, nil, nil)
		if err != nil {
			return nil, err
		}
		ret = append(ret, obj)
	}
	return ret, nil
}

func jsonWalk(obj interface{}) ([]interface{}, error) {
	switch o := obj.(type) {
	case map[string]interface{}:
		if o["kind"] != nil && o["apiVersion"] != nil {
			return []interface{}{o}, nil
		}
		ret := []interface{}{}
		for _, v := range o {
			children, err := jsonWalk(v)
			if err != nil {
				return nil, err
			}
			ret = append(ret, children...)
		}
		return ret, nil
	case []interface{}:
		ret := make([]interface{}, 0, len(o))
		for _, v := range o {
			children, err := jsonWalk(v)
			if err != nil {
				return nil, err
			}
			ret = append(ret, children...)
		}
		return ret, nil
	default:
		return nil, fmt.Errorf("Unexpected object structure: %T", o)
	}
}

func jsonnetReader(vm *jsonnet.VM, path string) ([]runtime.Object, error) {
	jsonstr, err := vm.EvaluateFile(path)
	if err != nil {
		return nil, err
	}

	glog.V(4).Infof("jsonnet result is: %s\n", jsonstr)

	var top interface{}
	if err = json.Unmarshal([]byte(jsonstr), &top); err != nil {
		return nil, err
	}

	objs, err := jsonWalk(top)
	if err != nil {
		return nil, err
	}

	ret := make([]runtime.Object, 0, len(objs))
	for _, v := range objs {
		// TODO: Going to json and back is a bit horrible
		data, err := json.Marshal(v)
		if err != nil {
			return nil, err
		}
		obj, _, err := runtime.UnstructuredJSONScheme.Decode(data, nil, nil)
		if err != nil {
			return nil, err
		}
		ret = append(ret, obj)
	}

	return ret, nil
}

// FlattenToV1 expands any List-type objects into their members, and
// cooerces everything to v1.Unstructured.  Panics if coercion
// encounters an unexpected object type.
func FlattenToV1(objs []runtime.Object) []*runtime.Unstructured {
	ret := make([]*runtime.Unstructured, 0, len(objs))
	for _, obj := range objs {
		switch o := obj.(type) {
		case *runtime.UnstructuredList:
			for _, item := range o.Items {
				ret = append(ret, item)
			}
		case *runtime.Unstructured:
			ret = append(ret, o)
		default:
			panic("Unexpected unstructured object type")
		}
	}
	return ret
}
<|endoftext|>


<commit_before>package instructions

import "jvmgo/rtda"

// Load double from local variable
type dload struct {Index8Instruction}
func (self *dload) Execute(thread *rtda.Thread) {
    _dload(thread, uint(self.index))
}

type dload_0 struct {NoOperandsInstruction}
func (self *dload_0) Execute(thread *rtda.Thread) {
    _dload(thread, 0)
}

type dload_1 struct {NoOperandsInstruction}
func (self *dload_1) Execute(thread *rtda.Thread) {
    _dload(thread, 1)
}

type dload_2 struct {NoOperandsInstruction}
func (self *dload_2) Execute(thread *rtda.Thread) {
    _dload(thread, 2)
}

type dload_3 struct {NoOperandsInstruction}
func (self *dload_3) Execute(thread *rtda.Thread) {
    _dload(thread, 3)
}

func _dload(thread *rtda.Thread, index uint) {
    val := thread.CurrentFrame().LocalVars().GetDouble(index)
    thread.CurrentFrame().OperandStack().PushDouble(val)
}
<commit_msg>refactor dload<commit_after>package instructions

import "jvmgo/rtda"

// Load double from local variable
type dload struct {Index8Instruction}
func (self *dload) Execute(frame *rtda.Frame) {
    _dload(frame, uint(self.index))
}

type dload_0 struct {NoOperandsInstruction}
func (self *dload_0) Execute(frame *rtda.Frame) {
    _dload(frame, 0)
}

type dload_1 struct {NoOperandsInstruction}
func (self *dload_1) Execute(frame *rtda.Frame) {
    _dload(frame, 1)
}

type dload_2 struct {NoOperandsInstruction}
func (self *dload_2) Execute(frame *rtda.Frame) {
    _dload(frame, 2)
}

type dload_3 struct {NoOperandsInstruction}
func (self *dload_3) Execute(frame *rtda.Frame) {
    _dload(frame, 3)
}

func _dload(frame *rtda.Frame, index uint) {
    val := frame.LocalVars().GetDouble(index)
    frame.OperandStack().PushDouble(val)
}
<|endoftext|>