<commit_before>package rule

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
	"sync"
)

// Config ...
type Config struct {
	// ChannelOptions embedded on top level.
	ChannelOptions
	// Namespaces – list of namespaces for custom channel options.
	Namespaces []ChannelNamespace
	// RpcOptions embedded on top level.
	RpcOptions
	// RpcNamespaces - list of rpc namespace for custom rpc options.
	RpcNamespaces []RpcNamespace
	// RpcNamespaceBoundary is a string separator which must be put after
	// rpc namespace part in rpc method.
	RpcNamespaceBoundary string
	// ChannelUserBoundary is a string separator which must be set before
	// allowed users part in channel name.
	// ChannelPrivatePrefix is a prefix in channel name which indicates that
	// channel is private.
	ChannelPrivatePrefix string
	// ChannelNamespaceBoundary is a string separator which must be put after
	// namespace part in channel name.
	ChannelNamespaceBoundary string
	// ChannelUserBoundary is a string separator which must be set before
	// allowed users part in channel name.
	ChannelUserBoundary string
	// ChannelUserSeparator separates allowed users in user part of channel name.
	// So you can limit access to channel to limited set of users.
	ChannelUserSeparator string
	// UserSubscribeToPersonal enables automatic subscribing to personal channel
	// by user.  Only users with user ID defined will subscribe to personal
	// channels, anonymous users are ignored.
	UserSubscribeToPersonal bool
	// UserPersonalChannelPrefix defines prefix to be added to user personal channel.
	// This should match one of configured namespace names. By default no namespace
	// used for personal channel.
	UserPersonalChannelNamespace string
	// UserPersonalSingleConnection turns on a mode in which Centrifugo will try to
	// maintain only a single connection for each user in the same moment. As soon as
	// user establishes a connection other connections from the same user will be closed
	// with connection limit reason.
	// This feature works with a help of presence information inside personal channel.
	// So presence should be turned on in personal channel.
	UserPersonalSingleConnection bool
	// ClientInsecure turns on insecure mode for client connections - when it's
	// turned on then no authentication required at all when connecting to Centrifugo,
	// anonymous access and publish allowed for all channels, no connection expire
	// performed. This can be suitable for demonstration or personal usage.
	ClientInsecure bool
	// ClientAnonymous when set to true, allows connect requests without specifying
	// a token or setting Credentials in authentication middleware. The resulting
	// user will have empty string for user ID, meaning user can only subscribe
	// to anonymous channels.
	ClientAnonymous bool
	// ClientConcurrency when set allows processing client commands concurrently
	// with provided concurrency level. By default commands processed sequentially
	// one after another.
	ClientConcurrency int
}

// DefaultConfig has default config options.
var DefaultConfig = Config{
	ChannelPrivatePrefix:     "$", // so private channel will look like "$gossips".
	ChannelNamespaceBoundary: ":", // so namespace "public" can be used as "public:news".
	ChannelUserBoundary:      "#", // so user limited channel is "user#2694" where "2696" is user ID.
	ChannelUserSeparator:     ",", // so several users limited channel is "dialog#2694,3019".
	RpcNamespaceBoundary:     ":", // so rpc namespace "chat" can be used as "chat:get_user_info".
}

func stringInSlice(a string, list []string) bool {
	for _, b := range list {
		if b == a {
			return true
		}
	}
	return false
}

var namePattern = "^[-a-zA-Z0-9_.]{2,}$"
var nameRe = regexp.MustCompile(namePattern)

func ValidateNamespace(ns ChannelNamespace) error {
	name := ns.Name
	match := nameRe.MatchString(name)
	if !match {
		return fmt.Errorf("invalid namespace name – %s (must match %s regular expression)", name, namePattern)
	}
	if err := ValidateChannelOptions(ns.ChannelOptions); err != nil {
		return err
	}
	return nil
}

func ValidateRpcNamespace(ns RpcNamespace) error {
	name := ns.Name
	match := nameRe.MatchString(name)
	if !match {
		return fmt.Errorf("invalid rpc namespace name – %s (must match %s regular expression)", name, namePattern)
	}
	if err := ValidateRpcOptions(ns.RpcOptions); err != nil {
		return err
	}
	return nil
}

func ValidateChannelOptions(c ChannelOptions) error {
	if (c.HistorySize != 0 && c.HistoryTTL == 0) || (c.HistorySize == 0 && c.HistoryTTL != 0) {
		return errors.New("both history size and history ttl required for history")
	}
	if c.Recover && (c.HistorySize == 0 || c.HistoryTTL == 0) {
		return errors.New("both history size and history ttl required for recovery")
	}
	return nil
}

func ValidateRpcOptions(_ RpcOptions) error {
	return nil
}

// Validate validates config and returns error if problems found
func (c *Config) Validate() error {
	if err := ValidateChannelOptions(c.ChannelOptions); err != nil {
		return err
	}
	if err := ValidateRpcOptions(c.RpcOptions); err != nil {
		return err
	}

	usePersonalChannel := c.UserSubscribeToPersonal
	personalChannelNamespace := c.UserPersonalChannelNamespace
	personalSingleConnection := c.UserPersonalSingleConnection
	var validPersonalChannelNamespace bool
	if !usePersonalChannel || personalChannelNamespace == "" {
		validPersonalChannelNamespace = true
		if personalSingleConnection && !c.Presence {
			return fmt.Errorf("presence must be enabled on top level to maintain single connection")
		}
	}

	nss := make([]string, 0, len(c.Namespaces))
	for _, n := range c.Namespaces {
		if stringInSlice(n.Name, nss) {
			return fmt.Errorf("namespace name must be unique: %s", n.Name)
		}
		if err := ValidateNamespace(n); err != nil {
			return fmt.Errorf("namespace %s: %v", n.Name, err)
		}
		if n.Name == personalChannelNamespace {
			validPersonalChannelNamespace = true
			if personalSingleConnection && !n.Presence {
				return fmt.Errorf("presence must be enabled for namespace %s to maintain single connection", n.Name)
			}
		}
		nss = append(nss, n.Name)
	}

	if !validPersonalChannelNamespace {
		return fmt.Errorf("namespace for user personal channel not found: %s", personalChannelNamespace)
	}

	rpcNss := make([]string, 0, len(c.RpcNamespaces))
	for _, n := range c.RpcNamespaces {
		if stringInSlice(n.Name, rpcNss) {
			return fmt.Errorf("rpc namespace name must be unique: %s", n.Name)
		}
		if err := ValidateRpcNamespace(n); err != nil {
			return fmt.Errorf("rpc namespace %s: %v", n.Name, err)
		}
		rpcNss = append(rpcNss, n.Name)
	}

	return nil
}

// Container ...
type Container struct {
	mu     sync.RWMutex
	config Config
}

// NewContainer ...
func NewContainer(config Config) *Container {
	return &Container{
		config: config,
	}
}

// Reload node config.
func (n *Container) Reload(c Config) error {
	if err := c.Validate(); err != nil {
		return err
	}
	n.mu.Lock()
	defer n.mu.Unlock()
	n.config = c
	return nil
}

// namespaceName returns namespace name from channel if exists.
func (n *Container) namespaceName(ch string) string {
	cTrim := strings.TrimPrefix(ch, n.config.ChannelPrivatePrefix)
	if n.config.ChannelNamespaceBoundary != "" && strings.Contains(cTrim, n.config.ChannelNamespaceBoundary) {
		parts := strings.SplitN(cTrim, n.config.ChannelNamespaceBoundary, 2)
		return parts[0]
	}
	return ""
}

// ChannelOptions returns channel options for channel using current channel config.
func (n *Container) ChannelOptions(ch string) (string, ChannelOptions, bool, error) {
	n.mu.RLock()
	defer n.mu.RUnlock()
	return n.config.channelOpts(n.namespaceName(ch))
}

// channelOpts searches for channel options for specified namespace key.
func (c *Config) channelOpts(namespaceName string) (string, ChannelOptions, bool, error) {
	if namespaceName == "" {
		return namespaceName, c.ChannelOptions, true, nil
	}
	for _, n := range c.Namespaces {
		if n.Name == namespaceName {
			return namespaceName, n.ChannelOptions, true, nil
		}
	}
	return namespaceName, ChannelOptions{}, false, nil
}

// PersonalChannel returns personal channel for user based on node configuration.
func (n *Container) PersonalChannel(user string) string {
	config := n.Config()
	if config.UserPersonalChannelNamespace == "" {
		return config.ChannelUserBoundary + user
	}
	return config.UserPersonalChannelNamespace + config.ChannelNamespaceBoundary + config.ChannelUserBoundary + user
}

// Config returns a copy of node Config.
func (n *Container) Config() Config {
	n.mu.RLock()
	c := n.config
	n.mu.RUnlock()
	return c
}

// IsPrivateChannel checks if channel requires token to subscribe. In case of
// token-protected channel subscription request must contain a proper token.
func (n *Container) IsPrivateChannel(ch string) bool {
	n.mu.RLock()
	defer n.mu.RUnlock()
	if n.config.ChannelPrivatePrefix == "" {
		return false
	}
	return strings.HasPrefix(ch, n.config.ChannelPrivatePrefix)
}

// IsUserLimited returns whether channel is user-limited.
func (n *Container) IsUserLimited(ch string) bool {
	n.mu.RLock()
	defer n.mu.RUnlock()
	userBoundary := n.config.ChannelUserBoundary
	if userBoundary == "" {
		return false
	}
	return strings.Contains(ch, userBoundary)
}

// UserAllowed checks if user can subscribe on channel - as channel
// can contain special part in the end to indicate which users allowed
// to subscribe on it.
func (n *Container) UserAllowed(ch string, user string) bool {
	n.mu.RLock()
	defer n.mu.RUnlock()
	userBoundary := n.config.ChannelUserBoundary
	userSeparator := n.config.ChannelUserSeparator
	if userBoundary == "" {
		return true
	}
	if !strings.Contains(ch, userBoundary) {
		return true
	}
	parts := strings.Split(ch, userBoundary)
	if userSeparator == "" {
		return parts[len(parts)-1] == user
	}
	allowedUsers := strings.Split(parts[len(parts)-1], userSeparator)
	for _, allowedUser := range allowedUsers {
		if user == allowedUser {
			return true
		}
	}
	return false
}

// rpcNamespaceName returns rpc namespace name from channel if exists.
func (n *Container) rpcNamespaceName(method string) string {
	if n.config.RpcNamespaceBoundary != "" && strings.Contains(method, n.config.RpcNamespaceBoundary) {
		parts := strings.SplitN(method, n.config.RpcNamespaceBoundary, 2)
		return parts[0]
	}
	return ""
}

// RpcOptions returns rpc options for method using current config.
func (n *Container) RpcOptions(method string) (RpcOptions, bool, error) {
	n.mu.RLock()
	defer n.mu.RUnlock()
	return n.config.rpcOpts(n.rpcNamespaceName(method))
}

// rpcOpts searches for channel options for specified namespace key.
func (c *Config) rpcOpts(namespaceName string) (RpcOptions, bool, error) {
	if namespaceName == "" {
		return c.RpcOptions, true, nil
	}
	for _, n := range c.RpcNamespaces {
		if n.Name == namespaceName {
			return n.RpcOptions, true, nil
		}
	}
	return RpcOptions{}, false, nil
}
<commit_msg>fix tests<commit_after>package rule

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
	"sync"
)

// Config ...
type Config struct {
	// ChannelOptions embedded on top level.
	ChannelOptions
	// Namespaces – list of namespaces for custom channel options.
	Namespaces []ChannelNamespace
	// RpcOptions embedded on top level.
	RpcOptions
	// RpcNamespaces - list of rpc namespace for custom rpc options.
	RpcNamespaces []RpcNamespace
	// RpcNamespaceBoundary is a string separator which must be put after
	// rpc namespace part in rpc method.
	RpcNamespaceBoundary string
	// ChannelUserBoundary is a string separator which must be set before
	// allowed users part in channel name.
	// ChannelPrivatePrefix is a prefix in channel name which indicates that
	// channel is private.
	ChannelPrivatePrefix string
	// ChannelNamespaceBoundary is a string separator which must be put after
	// namespace part in channel name.
	ChannelNamespaceBoundary string
	// ChannelUserBoundary is a string separator which must be set before
	// allowed users part in channel name.
	ChannelUserBoundary string
	// ChannelUserSeparator separates allowed users in user part of channel name.
	// So you can limit access to channel to limited set of users.
	ChannelUserSeparator string
	// UserSubscribeToPersonal enables automatic subscribing to personal channel
	// by user.  Only users with user ID defined will subscribe to personal
	// channels, anonymous users are ignored.
	UserSubscribeToPersonal bool
	// UserPersonalChannelPrefix defines prefix to be added to user personal channel.
	// This should match one of configured namespace names. By default no namespace
	// used for personal channel.
	UserPersonalChannelNamespace string
	// UserPersonalSingleConnection turns on a mode in which Centrifugo will try to
	// maintain only a single connection for each user in the same moment. As soon as
	// user establishes a connection other connections from the same user will be closed
	// with connection limit reason.
	// This feature works with a help of presence information inside personal channel.
	// So presence should be turned on in personal channel.
	UserPersonalSingleConnection bool
	// ClientInsecure turns on insecure mode for client connections - when it's
	// turned on then no authentication required at all when connecting to Centrifugo,
	// anonymous access and publish allowed for all channels, no connection expire
	// performed. This can be suitable for demonstration or personal usage.
	ClientInsecure bool
	// ClientAnonymous when set to true, allows connect requests without specifying
	// a token or setting Credentials in authentication middleware. The resulting
	// user will have empty string for user ID, meaning user can only subscribe
	// to anonymous channels.
	ClientAnonymous bool
	// ClientConcurrency when set allows processing client commands concurrently
	// with provided concurrency level. By default commands processed sequentially
	// one after another.
	ClientConcurrency int
}

// DefaultConfig has default config options.
var DefaultConfig = Config{
	ChannelPrivatePrefix:     "$", // so private channel will look like "$gossips".
	ChannelNamespaceBoundary: ":", // so namespace "public" can be used as "public:news".
	ChannelUserBoundary:      "#", // so user limited channel is "user#2694" where "2696" is user ID.
	ChannelUserSeparator:     ",", // so several users limited channel is "dialog#2694,3019".
	RpcNamespaceBoundary:     ":", // so rpc namespace "chat" can be used as "chat:get_user_info".
}

func stringInSlice(a string, list []string) bool {
	for _, b := range list {
		if b == a {
			return true
		}
	}
	return false
}

var namePattern = "^[-a-zA-Z0-9_.]{2,}$"
var nameRe = regexp.MustCompile(namePattern)

func ValidateNamespace(ns ChannelNamespace) error {
	name := ns.Name
	match := nameRe.MatchString(name)
	if !match {
		return fmt.Errorf("invalid namespace name – %s (must match %s regular expression)", name, namePattern)
	}
	if err := ValidateChannelOptions(ns.ChannelOptions); err != nil {
		return err
	}
	return nil
}

func ValidateRpcNamespace(ns RpcNamespace) error {
	name := ns.Name
	match := nameRe.MatchString(name)
	if !match {
		return fmt.Errorf("invalid rpc namespace name – %s (must match %s regular expression)", name, namePattern)
	}
	if err := ValidateRpcOptions(ns.RpcOptions); err != nil {
		return err
	}
	return nil
}

func ValidateChannelOptions(c ChannelOptions) error {
	if (c.HistorySize != 0 && c.HistoryTTL == 0) || (c.HistorySize == 0 && c.HistoryTTL != 0) {
		return errors.New("both history size and history ttl required for history")
	}
	if c.Recover && (c.HistorySize == 0 || c.HistoryTTL == 0) {
		return errors.New("both history size and history ttl required for recovery")
	}
	return nil
}

func ValidateRpcOptions(_ RpcOptions) error {
	return nil
}

// Validate validates config and returns error if problems found
func (c *Config) Validate() error {
	if err := ValidateChannelOptions(c.ChannelOptions); err != nil {
		return err
	}
	if err := ValidateRpcOptions(c.RpcOptions); err != nil {
		return err
	}

	usePersonalChannel := c.UserSubscribeToPersonal
	personalChannelNamespace := c.UserPersonalChannelNamespace
	personalSingleConnection := c.UserPersonalSingleConnection
	var validPersonalChannelNamespace bool
	if !usePersonalChannel || personalChannelNamespace == "" {
		validPersonalChannelNamespace = true
		if personalSingleConnection && !c.Presence {
			return fmt.Errorf("presence must be enabled on top level to maintain single connection")
		}
	}

	nss := make([]string, 0, len(c.Namespaces))
	for _, n := range c.Namespaces {
		if stringInSlice(n.Name, nss) {
			return fmt.Errorf("namespace name must be unique: %s", n.Name)
		}
		if err := ValidateNamespace(n); err != nil {
			return fmt.Errorf("namespace %s: %v", n.Name, err)
		}
		if n.Name == personalChannelNamespace {
			validPersonalChannelNamespace = true
			if personalSingleConnection && !n.Presence {
				return fmt.Errorf("presence must be enabled for namespace %s to maintain single connection", n.Name)
			}
		}
		nss = append(nss, n.Name)
	}

	if !validPersonalChannelNamespace {
		return fmt.Errorf("namespace for user personal channel not found: %s", personalChannelNamespace)
	}

	rpcNss := make([]string, 0, len(c.RpcNamespaces))
	for _, n := range c.RpcNamespaces {
		if stringInSlice(n.Name, rpcNss) {
			return fmt.Errorf("rpc namespace name must be unique: %s", n.Name)
		}
		if err := ValidateRpcNamespace(n); err != nil {
			return fmt.Errorf("rpc namespace %s: %v", n.Name, err)
		}
		rpcNss = append(rpcNss, n.Name)
	}

	return nil
}

// Container ...
type Container struct {
	mu     sync.RWMutex
	config Config
}

// NewContainer ...
func NewContainer(config Config) *Container {
	return &Container{
		config: config,
	}
}

// Reload node config.
func (n *Container) Reload(c Config) error {
	if err := c.Validate(); err != nil {
		return err
	}
	n.mu.Lock()
	defer n.mu.Unlock()
	n.config = c
	return nil
}

// namespaceName returns namespace name from channel if exists.
func (n *Container) namespaceName(ch string) string {
	cTrim := strings.TrimPrefix(ch, n.config.ChannelPrivatePrefix)
	if n.config.ChannelNamespaceBoundary != "" && strings.Contains(cTrim, n.config.ChannelNamespaceBoundary) {
		parts := strings.SplitN(cTrim, n.config.ChannelNamespaceBoundary, 2)
		return parts[0]
	}
	return ""
}

// ChannelOptions returns channel options for channel using current channel config.
func (n *Container) ChannelOptions(ch string) (string, ChannelOptions, bool, error) {
	n.mu.RLock()
	defer n.mu.RUnlock()
	nsName := n.namespaceName(ch)
	chOpts, ok, err := n.config.channelOpts(nsName)
	return nsName, chOpts, ok, err
}

// channelOpts searches for channel options for specified namespace key.
func (c *Config) channelOpts(namespaceName string) (ChannelOptions, bool, error) {
	if namespaceName == "" {
		return c.ChannelOptions, true, nil
	}
	for _, n := range c.Namespaces {
		if n.Name == namespaceName {
			return n.ChannelOptions, true, nil
		}
	}
	return ChannelOptions{}, false, nil
}

// PersonalChannel returns personal channel for user based on node configuration.
func (n *Container) PersonalChannel(user string) string {
	config := n.Config()
	if config.UserPersonalChannelNamespace == "" {
		return config.ChannelUserBoundary + user
	}
	return config.UserPersonalChannelNamespace + config.ChannelNamespaceBoundary + config.ChannelUserBoundary + user
}

// Config returns a copy of node Config.
func (n *Container) Config() Config {
	n.mu.RLock()
	c := n.config
	n.mu.RUnlock()
	return c
}

// IsPrivateChannel checks if channel requires token to subscribe. In case of
// token-protected channel subscription request must contain a proper token.
func (n *Container) IsPrivateChannel(ch string) bool {
	n.mu.RLock()
	defer n.mu.RUnlock()
	if n.config.ChannelPrivatePrefix == "" {
		return false
	}
	return strings.HasPrefix(ch, n.config.ChannelPrivatePrefix)
}

// IsUserLimited returns whether channel is user-limited.
func (n *Container) IsUserLimited(ch string) bool {
	n.mu.RLock()
	defer n.mu.RUnlock()
	userBoundary := n.config.ChannelUserBoundary
	if userBoundary == "" {
		return false
	}
	return strings.Contains(ch, userBoundary)
}

// UserAllowed checks if user can subscribe on channel - as channel
// can contain special part in the end to indicate which users allowed
// to subscribe on it.
func (n *Container) UserAllowed(ch string, user string) bool {
	n.mu.RLock()
	defer n.mu.RUnlock()
	userBoundary := n.config.ChannelUserBoundary
	userSeparator := n.config.ChannelUserSeparator
	if userBoundary == "" {
		return true
	}
	if !strings.Contains(ch, userBoundary) {
		return true
	}
	parts := strings.Split(ch, userBoundary)
	if userSeparator == "" {
		return parts[len(parts)-1] == user
	}
	allowedUsers := strings.Split(parts[len(parts)-1], userSeparator)
	for _, allowedUser := range allowedUsers {
		if user == allowedUser {
			return true
		}
	}
	return false
}

// rpcNamespaceName returns rpc namespace name from channel if exists.
func (n *Container) rpcNamespaceName(method string) string {
	if n.config.RpcNamespaceBoundary != "" && strings.Contains(method, n.config.RpcNamespaceBoundary) {
		parts := strings.SplitN(method, n.config.RpcNamespaceBoundary, 2)
		return parts[0]
	}
	return ""
}

// RpcOptions returns rpc options for method using current config.
func (n *Container) RpcOptions(method string) (RpcOptions, bool, error) {
	n.mu.RLock()
	defer n.mu.RUnlock()
	return n.config.rpcOpts(n.rpcNamespaceName(method))
}

// rpcOpts searches for channel options for specified namespace key.
func (c *Config) rpcOpts(namespaceName string) (RpcOptions, bool, error) {
	if namespaceName == "" {
		return c.RpcOptions, true, nil
	}
	for _, n := range c.RpcNamespaces {
		if n.Name == namespaceName {
			return n.RpcOptions, true, nil
		}
	}
	return RpcOptions{}, false, nil
}
<|endoftext|>


<commit_before>// Package adhier provides an algorithm for adaptive hierarchical interpolation
// with local refinements.
package adhier

import (
	"fmt"
	"math"
)

const (
	bufferInitCount  = 200
	bufferGrowFactor = 2
)

// Grid is the interface that an sparse grid should satisfy in order to be used
// in the algorithm.
type Grid interface {
	Dimensionality() uint16
	ComputeNodes(levels []uint8, orders []uint32) []float64
	ComputeChildren(levels []uint8, orders []uint32) ([]uint8, []uint32)
}

// Basis is the interface that a functional basis should satisfy in order to be
// used in the algorithm.
type Basis interface {
	Evaluate(levels []uint8, orders []uint32, point []float64) float64
}

// Self represents a particular instantiation of the algorithm.
type Self struct {
	grid     Grid
	basis    Basis
	outCount uint16
	minLevel uint8
	maxLevel uint8
	absError float64
	relError float64
}

// New creates an instance of the algorithm for the given sparse grid and
// functional basis.
func New(grid Grid, basis Basis, outCount uint16) *Self {
	return &Self{
		grid:     grid,
		basis:    basis,
		outCount: outCount,
		minLevel: 1,
		maxLevel: 9,
		absError: 1e-4,
		relError: 1e-2,
	}
}

// Surrogate is the result of Construct, which represents an interpolant for a
// function.
type Surrogate struct {
	level     uint8
	inCount   uint16
	outCount  uint16
	nodeCount uint32

	levels    []uint8
	orders    []uint32
	surpluses []float64
}

func (s *Surrogate) initialize(inCount, outCount uint16) {
	ins := bufferInitCount * inCount
	outs := bufferInitCount * outCount

	s.inCount = inCount
	s.outCount = outCount
	s.nodeCount = bufferInitCount

	s.levels = make([]uint8, ins)
	s.orders = make([]uint32, ins)
	s.surpluses = make([]float64, outs)
}

func (s *Surrogate) finalize(level uint8, nodeCount uint32) {
	ins := nodeCount * uint32(s.inCount)
	outs := nodeCount * uint32(s.outCount)

	s.level = level
	s.nodeCount = nodeCount

	s.levels = s.levels[0:ins]
	s.orders = s.orders[0:ins]
	s.surpluses = s.surpluses[0:outs]
}

func (s *Surrogate) resize(nodeCount uint32) {
	if nodeCount <= s.nodeCount {
		return
	}

	if count := bufferGrowFactor * s.nodeCount; count > nodeCount {
		nodeCount = count
	}

	// New sizes
	ins := nodeCount * uint32(s.inCount)
	outs := nodeCount * uint32(s.outCount)

	levels := make([]uint8, ins)
	orders := make([]uint32, ins)
	surpluses := make([]float64, outs)

	// Old sizes
	ins = s.nodeCount * uint32(s.inCount)
	outs = s.nodeCount * uint32(s.outCount)

	copy(levels, s.levels[0:ins])
	copy(orders, s.orders[0:ins])
	copy(surpluses, s.surpluses[0:outs])

	s.nodeCount = nodeCount

	s.levels = levels
	s.orders = orders
	s.surpluses = surpluses
}

// String returns a string containing human-friendly information about the
// surrogate.
func (s *Surrogate) String() string {
	return fmt.Sprintf("Surrogate{ inputs: %d, outputs: %d, levels: %d, nodes: %d }",
		s.inCount, s.outCount, s.level, s.nodeCount)
}

// Construct takes a function and yields a surrogate/interpolant for it, which
// can be further fed to Evaluate for the actual interpolation.
func (self *Self) Construct(target func([]float64) []float64) *Surrogate {
	inc := uint32(self.grid.Dimensionality())
	outc := uint32(self.outCount)

	surrogate := new(Surrogate)
	surrogate.initialize(uint16(inc), uint16(outc))

	// Assume level 0 has only one node, and its order is 0.
	level := uint8(0)
	newc := uint32(1)
	levels := make([]uint8, newc*inc)
	orders := make([]uint32, newc*inc)

	oldc := uint32(0)
	nodeCount := uint32(0)

	var i, j, k, l uint32

	value := make([]float64, outc)

	minValue := make([]float64, outc)
	maxValue := make([]float64, outc)

	minValue[0] = math.Inf(1)
	maxValue[0] = math.Inf(-1)
	for i = 1; i < outc; i++ {
		minValue[i] = minValue[i-1]
		maxValue[i] = maxValue[i-1]
	}

	for {
		surrogate.resize(oldc + newc)

		copy(surrogate.levels[oldc*inc:], levels)
		copy(surrogate.orders[oldc*inc:], orders)

		nodes := self.grid.ComputeNodes(levels, orders)
		values := target(nodes)

		// Compute the surpluses corresponding to the active nodes.
		for i, k = 0, oldc*outc; i < newc; i++ {
			evaluate(self.basis, surrogate, inc, outc, oldc,
				nodes[i*inc:(i+1)*inc], value)

			for j = 0; j < outc; j++ {
				surrogate.surpluses[k] = values[i*outc+j] - value[j]
				k++
			}
		}

		nodeCount += newc

		if level >= self.maxLevel {
			break
		}

		// Keep track of the maximal and minimal values of the function.
		for i, k = 0, 0; i < newc; i++ {
			for j = 0; j < outc; j++ {
				if values[k] < minValue[j] {
					minValue[j] = values[k]
				}
				if values[k] > maxValue[j] {
					maxValue[j] = values[k]
				}
				k++
			}
		}

		if level >= self.minLevel {
			k, l = 0, 0

			for i = 0; i < newc; i++ {
				refine := false

				for j = 0; j < outc; j++ {
					absError := math.Abs(surrogate.surpluses[(oldc+i)*outc+j])

					if absError > self.absError {
						refine = true
						break
					}

					relError := absError / (maxValue[j] - minValue[j])

					if relError > self.relError {
						refine = true
						break
					}
				}

				if !refine {
					l += inc
					continue
				}

				if k != l {
					// Shift everything, assuming a lot of refinements.
					copy(levels[k:], levels[l:])
					copy(orders[k:], orders[l:])
					l = k
				}

				k += inc
				l += inc
			}

			levels = levels[0:k]
			orders = orders[0:k]
		}

		levels, orders = self.grid.ComputeChildren(levels, orders)

		oldc += newc
		newc = uint32(len(levels)) / inc

		if newc == 0 {
			break
		}

		level++
	}

	surrogate.finalize(level, nodeCount)
	return surrogate
}

// Evaluate takes a surrogate produced by Construct and evaluates it at the
// given points.
func (self *Self) Evaluate(surrogate *Surrogate, points []float64) []float64 {
	inc := uint32(surrogate.inCount)
	outc := uint32(surrogate.outCount)

	pointCount := uint32(len(points)) / inc

	values := make([]float64, pointCount*outc)

	for i := uint32(0); i < pointCount; i++ {
		evaluate(self.basis, surrogate, inc, outc, surrogate.nodeCount,
			points[i*inc:], values[i*outc:])
	}

	return values
}

func evaluate(basis Basis, surrogate *Surrogate, inc, outc, nodeCount uint32,
	point []float64, value []float64) {

	var i, j uint32 = 1, 0

	// Rewrite value in case it is dirty (not zeroed).
	weight := basis.Evaluate(surrogate.levels, surrogate.orders, point)
	for ; j < outc; j++ {
		value[j] = surrogate.surpluses[j] * weight
	}

	for ; i < nodeCount; i++ {
		weight = basis.Evaluate(surrogate.levels[i*inc:], surrogate.orders[i*inc:], point)
		for j = 0; j < outc; j++ {
			value[j] += surrogate.surpluses[i*outc+j] * weight
		}
	}
}
<commit_msg>Fixed an overflow bug when the subject function has many outputs<commit_after>// Package adhier provides an algorithm for adaptive hierarchical interpolation
// with local refinements.
package adhier

import (
	"fmt"
	"math"
)

const (
	bufferInitCount  = 200
	bufferGrowFactor = 2
)

// Grid is the interface that an sparse grid should satisfy in order to be used
// in the algorithm.
type Grid interface {
	Dimensionality() uint16
	ComputeNodes(levels []uint8, orders []uint32) []float64
	ComputeChildren(levels []uint8, orders []uint32) ([]uint8, []uint32)
}

// Basis is the interface that a functional basis should satisfy in order to be
// used in the algorithm.
type Basis interface {
	Evaluate(levels []uint8, orders []uint32, point []float64) float64
}

// Self represents a particular instantiation of the algorithm.
type Self struct {
	grid     Grid
	basis    Basis
	outCount uint16
	minLevel uint8
	maxLevel uint8
	absError float64
	relError float64
}

// New creates an instance of the algorithm for the given sparse grid and
// functional basis.
func New(grid Grid, basis Basis, outCount uint16) *Self {
	return &Self{
		grid:     grid,
		basis:    basis,
		outCount: outCount,
		minLevel: 1,
		maxLevel: 9,
		absError: 1e-4,
		relError: 1e-2,
	}
}

// Surrogate is the result of Construct, which represents an interpolant for a
// function.
type Surrogate struct {
	level     uint8
	inCount   uint16
	outCount  uint16
	nodeCount uint32

	levels    []uint8
	orders    []uint32
	surpluses []float64
}

func (s *Surrogate) initialize(inCount, outCount uint16) {
	ins := bufferInitCount * uint32(inCount)
	outs := bufferInitCount * uint32(outCount)

	s.inCount = inCount
	s.outCount = outCount
	s.nodeCount = bufferInitCount

	s.levels = make([]uint8, ins)
	s.orders = make([]uint32, ins)
	s.surpluses = make([]float64, outs)
}

func (s *Surrogate) finalize(level uint8, nodeCount uint32) {
	ins := nodeCount * uint32(s.inCount)
	outs := nodeCount * uint32(s.outCount)

	s.level = level
	s.nodeCount = nodeCount

	s.levels = s.levels[0:ins]
	s.orders = s.orders[0:ins]
	s.surpluses = s.surpluses[0:outs]
}

func (s *Surrogate) resize(nodeCount uint32) {
	if nodeCount <= s.nodeCount {
		return
	}

	if count := bufferGrowFactor * s.nodeCount; count > nodeCount {
		nodeCount = count
	}

	// New sizes
	ins := nodeCount * uint32(s.inCount)
	outs := nodeCount * uint32(s.outCount)

	levels := make([]uint8, ins)
	orders := make([]uint32, ins)
	surpluses := make([]float64, outs)

	// Old sizes
	ins = s.nodeCount * uint32(s.inCount)
	outs = s.nodeCount * uint32(s.outCount)

	copy(levels, s.levels[0:ins])
	copy(orders, s.orders[0:ins])
	copy(surpluses, s.surpluses[0:outs])

	s.nodeCount = nodeCount

	s.levels = levels
	s.orders = orders
	s.surpluses = surpluses
}

// String returns a string containing human-friendly information about the
// surrogate.
func (s *Surrogate) String() string {
	return fmt.Sprintf("Surrogate{ inputs: %d, outputs: %d, levels: %d, nodes: %d }",
		s.inCount, s.outCount, s.level, s.nodeCount)
}

// Construct takes a function and yields a surrogate/interpolant for it, which
// can be further fed to Evaluate for the actual interpolation.
func (self *Self) Construct(target func([]float64) []float64) *Surrogate {
	inc := uint32(self.grid.Dimensionality())
	outc := uint32(self.outCount)

	surrogate := new(Surrogate)
	surrogate.initialize(uint16(inc), uint16(outc))

	// Assume level 0 has only one node, and its order is 0.
	level := uint8(0)
	newc := uint32(1)
	levels := make([]uint8, newc*inc)
	orders := make([]uint32, newc*inc)

	oldc := uint32(0)
	nodeCount := uint32(0)

	var i, j, k, l uint32

	value := make([]float64, outc)

	minValue := make([]float64, outc)
	maxValue := make([]float64, outc)

	minValue[0] = math.Inf(1)
	maxValue[0] = math.Inf(-1)
	for i = 1; i < outc; i++ {
		minValue[i] = minValue[i-1]
		maxValue[i] = maxValue[i-1]
	}

	for {
		surrogate.resize(oldc + newc)

		copy(surrogate.levels[oldc*inc:], levels)
		copy(surrogate.orders[oldc*inc:], orders)

		nodes := self.grid.ComputeNodes(levels, orders)
		values := target(nodes)

		// Compute the surpluses corresponding to the active nodes.
		for i, k = 0, oldc*outc; i < newc; i++ {
			evaluate(self.basis, surrogate, inc, outc, oldc,
				nodes[i*inc:(i+1)*inc], value)

			for j = 0; j < outc; j++ {
				surrogate.surpluses[k] = values[i*outc+j] - value[j]
				k++
			}
		}

		nodeCount += newc

		if level >= self.maxLevel {
			break
		}

		// Keep track of the maximal and minimal values of the function.
		for i, k = 0, 0; i < newc; i++ {
			for j = 0; j < outc; j++ {
				if values[k] < minValue[j] {
					minValue[j] = values[k]
				}
				if values[k] > maxValue[j] {
					maxValue[j] = values[k]
				}
				k++
			}
		}

		if level >= self.minLevel {
			k, l = 0, 0

			for i = 0; i < newc; i++ {
				refine := false

				for j = 0; j < outc; j++ {
					absError := math.Abs(surrogate.surpluses[(oldc+i)*outc+j])

					if absError > self.absError {
						refine = true
						break
					}

					relError := absError / (maxValue[j] - minValue[j])

					if relError > self.relError {
						refine = true
						break
					}
				}

				if !refine {
					l += inc
					continue
				}

				if k != l {
					// Shift everything, assuming a lot of refinements.
					copy(levels[k:], levels[l:])
					copy(orders[k:], orders[l:])
					l = k
				}

				k += inc
				l += inc
			}

			levels = levels[0:k]
			orders = orders[0:k]
		}

		levels, orders = self.grid.ComputeChildren(levels, orders)

		oldc += newc
		newc = uint32(len(levels)) / inc

		if newc == 0 {
			break
		}

		level++
	}

	surrogate.finalize(level, nodeCount)
	return surrogate
}

// Evaluate takes a surrogate produced by Construct and evaluates it at the
// given points.
func (self *Self) Evaluate(surrogate *Surrogate, points []float64) []float64 {
	inc := uint32(surrogate.inCount)
	outc := uint32(surrogate.outCount)

	pointCount := uint32(len(points)) / inc

	values := make([]float64, pointCount*outc)

	for i := uint32(0); i < pointCount; i++ {
		evaluate(self.basis, surrogate, inc, outc, surrogate.nodeCount,
			points[i*inc:], values[i*outc:])
	}

	return values
}

func evaluate(basis Basis, surrogate *Surrogate, inc, outc, nodeCount uint32,
	point []float64, value []float64) {

	var i, j uint32 = 1, 0

	// Rewrite value in case it is dirty (not zeroed).
	weight := basis.Evaluate(surrogate.levels, surrogate.orders, point)
	for ; j < outc; j++ {
		value[j] = surrogate.surpluses[j] * weight
	}

	for ; i < nodeCount; i++ {
		weight = basis.Evaluate(surrogate.levels[i*inc:], surrogate.orders[i*inc:], point)
		for j = 0; j < outc; j++ {
			value[j] += surrogate.surpluses[i*outc+j] * weight
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"testing"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// A LXD node which is already configured for networking can be converted to a
// single-node LXD cluster.
func TestCluster_Bootstrap(t *testing.T) {
	daemon, cleanup := newTestDaemon(t)
	defer cleanup()

	// Simulate what happens when running "lxd init", where a PUT /1.0
	// request is issued to set both core.https_address and
	// cluster.https_address to the same value.
	f := clusterFixture{t: t}
	f.EnableNetworkingWithClusterAddress(daemon, "")

	client := f.ClientUnix(daemon)

	cluster := api.ClusterPut{}
	cluster.ServerName = "buzz"
	cluster.Enabled = true
	op, err := client.UpdateCluster(cluster, "")
	require.NoError(t, err)
	require.NoError(t, op.Wait())

	server, _, err := client.GetServer()
	require.NoError(t, err)
	assert.True(t, client.IsClustered())
	assert.Equal(t, "buzz", server.Environment.ServerName)
}

// Check the cluster API on a non-clustered server.
func TestCluster_Get(t *testing.T) {
	daemon, cleanup := newTestDaemon(t)
	defer cleanup()

	client, err := lxd.ConnectLXDUnix(daemon.UnixSocket(), nil)
	require.NoError(t, err)

	cluster, _, err := client.GetCluster()
	require.NoError(t, err)
	assert.Equal(t, "", cluster.ServerName)
	assert.False(t, cluster.Enabled)
}

// A LXD node can be renamed.
func TestCluster_RenameNode(t *testing.T) {
	daemon, cleanup := newTestDaemon(t)
	defer cleanup()

	f := clusterFixture{t: t}
	f.EnableNetworking(daemon, "")

	client := f.ClientUnix(daemon)

	cluster := api.ClusterPut{}
	cluster.ServerName = "buzz"
	cluster.Enabled = true
	op, err := client.UpdateCluster(cluster, "")
	require.NoError(t, err)
	require.NoError(t, op.Wait())

	node := api.ClusterMemberPost{ServerName: "rusp"}
	err = client.RenameClusterMember("buzz", node)
	require.NoError(t, err)

	_, _, err = client.GetClusterMember("rusp")
	require.NoError(t, err)
}

// Test helper for cluster-related APIs.
type clusterFixture struct {
	t       *testing.T
	clients map[*Daemon]lxd.InstanceServer
	daemons []*Daemon
}

// Enable networking in the given daemon. The password is optional and can be
// an empty string.
func (f *clusterFixture) EnableNetworking(daemon *Daemon, password string) {
	port, err := shared.AllocatePort()
	require.NoError(f.t, err)

	address := fmt.Sprintf("127.0.0.1:%d", port)

	client := f.ClientUnix(daemon)
	server, _, err := client.GetServer()
	require.NoError(f.t, err)
	serverPut := server.Writable()
	serverPut.Config["core.https_address"] = address
	serverPut.Config["core.trust_password"] = password

	require.NoError(f.t, client.UpdateServer(serverPut, ""))
}

// Enable networking in the given daemon, and set cluster.https_address to the
// same value as core.https address. The password is optional and can be an
// empty string.
func (f *clusterFixture) EnableNetworkingWithClusterAddress(daemon *Daemon, password string) {
	port, err := shared.AllocatePort()
	require.NoError(f.t, err)

	address := fmt.Sprintf("127.0.0.1:%d", port)

	client := f.ClientUnix(daemon)
	server, _, err := client.GetServer()
	require.NoError(f.t, err)
	serverPut := server.Writable()
	serverPut.Config["core.https_address"] = address
	serverPut.Config["core.trust_password"] = password
	serverPut.Config["cluster.https_address"] = address

	require.NoError(f.t, client.UpdateServer(serverPut, ""))
}

// Get a client for the given daemon connected via UNIX socket, creating one if
// needed.
func (f *clusterFixture) ClientUnix(daemon *Daemon) lxd.InstanceServer {
	if f.clients == nil {
		f.clients = make(map[*Daemon]lxd.InstanceServer)
	}
	client, ok := f.clients[daemon]
	if !ok {
		var err error
		client, err = lxd.ConnectLXDUnix(daemon.UnixSocket(), nil)
		require.NoError(f.t, err)
	}
	return client
}
<commit_msg>lxd: Removes unused struct field.<commit_after>package main

import (
	"fmt"
	"testing"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// A LXD node which is already configured for networking can be converted to a
// single-node LXD cluster.
func TestCluster_Bootstrap(t *testing.T) {
	daemon, cleanup := newTestDaemon(t)
	defer cleanup()

	// Simulate what happens when running "lxd init", where a PUT /1.0
	// request is issued to set both core.https_address and
	// cluster.https_address to the same value.
	f := clusterFixture{t: t}
	f.EnableNetworkingWithClusterAddress(daemon, "")

	client := f.ClientUnix(daemon)

	cluster := api.ClusterPut{}
	cluster.ServerName = "buzz"
	cluster.Enabled = true
	op, err := client.UpdateCluster(cluster, "")
	require.NoError(t, err)
	require.NoError(t, op.Wait())

	server, _, err := client.GetServer()
	require.NoError(t, err)
	assert.True(t, client.IsClustered())
	assert.Equal(t, "buzz", server.Environment.ServerName)
}

// Check the cluster API on a non-clustered server.
func TestCluster_Get(t *testing.T) {
	daemon, cleanup := newTestDaemon(t)
	defer cleanup()

	client, err := lxd.ConnectLXDUnix(daemon.UnixSocket(), nil)
	require.NoError(t, err)

	cluster, _, err := client.GetCluster()
	require.NoError(t, err)
	assert.Equal(t, "", cluster.ServerName)
	assert.False(t, cluster.Enabled)
}

// A LXD node can be renamed.
func TestCluster_RenameNode(t *testing.T) {
	daemon, cleanup := newTestDaemon(t)
	defer cleanup()

	f := clusterFixture{t: t}
	f.EnableNetworking(daemon, "")

	client := f.ClientUnix(daemon)

	cluster := api.ClusterPut{}
	cluster.ServerName = "buzz"
	cluster.Enabled = true
	op, err := client.UpdateCluster(cluster, "")
	require.NoError(t, err)
	require.NoError(t, op.Wait())

	node := api.ClusterMemberPost{ServerName: "rusp"}
	err = client.RenameClusterMember("buzz", node)
	require.NoError(t, err)

	_, _, err = client.GetClusterMember("rusp")
	require.NoError(t, err)
}

// Test helper for cluster-related APIs.
type clusterFixture struct {
	t       *testing.T
	clients map[*Daemon]lxd.InstanceServer
}

// Enable networking in the given daemon. The password is optional and can be
// an empty string.
func (f *clusterFixture) EnableNetworking(daemon *Daemon, password string) {
	port, err := shared.AllocatePort()
	require.NoError(f.t, err)

	address := fmt.Sprintf("127.0.0.1:%d", port)

	client := f.ClientUnix(daemon)
	server, _, err := client.GetServer()
	require.NoError(f.t, err)
	serverPut := server.Writable()
	serverPut.Config["core.https_address"] = address
	serverPut.Config["core.trust_password"] = password

	require.NoError(f.t, client.UpdateServer(serverPut, ""))
}

// Enable networking in the given daemon, and set cluster.https_address to the
// same value as core.https address. The password is optional and can be an
// empty string.
func (f *clusterFixture) EnableNetworkingWithClusterAddress(daemon *Daemon, password string) {
	port, err := shared.AllocatePort()
	require.NoError(f.t, err)

	address := fmt.Sprintf("127.0.0.1:%d", port)

	client := f.ClientUnix(daemon)
	server, _, err := client.GetServer()
	require.NoError(f.t, err)
	serverPut := server.Writable()
	serverPut.Config["core.https_address"] = address
	serverPut.Config["core.trust_password"] = password
	serverPut.Config["cluster.https_address"] = address

	require.NoError(f.t, client.UpdateServer(serverPut, ""))
}

// Get a client for the given daemon connected via UNIX socket, creating one if
// needed.
func (f *clusterFixture) ClientUnix(daemon *Daemon) lxd.InstanceServer {
	if f.clients == nil {
		f.clients = make(map[*Daemon]lxd.InstanceServer)
	}
	client, ok := f.clients[daemon]
	if !ok {
		var err error
		client, err = lxd.ConnectLXDUnix(daemon.UnixSocket(), nil)
		require.NoError(f.t, err)
	}
	return client
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"github.com/spf13/cobra"
)

/*
#ifndef _GNU_SOURCE
#define _GNU_SOURCE 1
#endif
#include <fcntl.h>
#include <libgen.h>
#include <sched.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/capability.h>
#include <sys/fsuid.h>
#include <sys/prctl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/vfs.h>
#include <sys/xattr.h>
#include <unistd.h>

#include "include/memory_utils.h"

extern char* advance_arg(bool required);
extern int dosetns(int pid, char *nstype);

static inline bool same_fsinfo(struct stat *s1, struct stat *s2,
			       struct statfs *sfs1, struct statfs *sfs2)
{
	return ((sfs1->f_type == sfs2->f_type) && (s1->st_dev == s2->st_dev));
}

static int fstat_fstatfs(int fd, struct stat *s, struct statfs *sfs)
{
	if (fstat(fd, s))
		return -1;

	if (fstatfs(fd, sfs))
		return -1;

	return 0;
}

static bool chdirchroot_in_mntns(int cwd_fd, int root_fd)
{
	ssize_t len;
	char buf[PATH_MAX];

	if (fchdir(cwd_fd))
		return false;

	len = readlinkat(root_fd, "", buf, sizeof(buf));
	if (len < 0 || len >= sizeof(buf))
		return false;
	buf[len] = '\0';

	if (chroot(buf))
		return false;

	return true;
}

// Expects command line to be in the form:
// <PID> <root-uid> <root-gid> <path> <mode> <dev>
static void forkmknod()
{
	__do_close_prot_errno int cwd_fd = -EBADF, mnt_fd = -EBADF, root_fd = -EBADF, target_dir_fd = -EBADF;
	char *cur = NULL, *target = NULL, *target_dir = NULL, *target_host = NULL;
	int ret;
	char path[PATH_MAX];
	mode_t mode;
	dev_t dev;
	pid_t pid;
	uid_t fsuid, uid;
	gid_t fsgid, gid;
	struct statfs sfs;
	cap_t caps;

	pid = atoi(advance_arg(true));
	target = advance_arg(true);
	mode = atoi(advance_arg(true));
	dev = atoi(advance_arg(true));
	target_host = advance_arg(true);
	uid = atoi(advance_arg(true));
	gid = atoi(advance_arg(true));
	fsuid = atoi(advance_arg(true));
	fsgid = atoi(advance_arg(true));

	snprintf(path, sizeof(path), "/proc/%d/ns/mnt", pid);
	mnt_fd = open(path, O_RDONLY | O_CLOEXEC);
	if (mnt_fd < 0) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	snprintf(path, sizeof(path), "/proc/%d/root", pid);
	root_fd = open(path, O_PATH | O_RDONLY | O_CLOEXEC | O_NOFOLLOW);
	if (root_fd < 0) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	snprintf(path, sizeof(path), "/proc/%d/cwd", pid);
	cwd_fd = open(path, O_PATH | O_RDONLY | O_CLOEXEC);
	if (cwd_fd < 0) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	if (setns(mnt_fd, CLONE_NEWNS)) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	if (!chdirchroot_in_mntns(cwd_fd, root_fd)) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	snprintf(path, sizeof(path), "%s", target);
	target_dir = dirname(path);
	target_dir_fd = open(target_dir, O_PATH | O_RDONLY | O_CLOEXEC | O_DIRECTORY);
	if (target_dir_fd < 0) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	caps = cap_get_pid(pid);
	if (!caps) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	ret = prctl(PR_SET_KEEPCAPS, 1);
	if (ret) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	ret = setegid(gid);
	if (ret) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	setfsgid(fsgid);

	ret = seteuid(uid);
	if (ret) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	setfsuid(fsuid);

	ret = cap_set_proc(caps);
	if (ret) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	ret = fstatfs(target_dir_fd, &sfs);
	if (ret) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	if (sfs.f_flags & MS_NODEV) {
		fprintf(stderr, "%d", EPERM);
		_exit(EXIT_FAILURE);
	}

	// basename() can modify its argument so accessing target_host is
	// invalid from now on.
	ret = mknodat(target_dir_fd, target, mode, dev);
	if (ret) {
		if (errno == EPERM)
			fprintf(stderr, "%d", ENOMEDIUM);
		else
			fprintf(stderr, "%d", errno);

		_exit(EXIT_FAILURE);
	}
}

#ifndef CLONE_NEWCGROUP
#define CLONE_NEWCGROUP	0x02000000
#endif

const char *ns_names[] = { "user", "pid", "uts", "ipc", "net", "cgroup", NULL };

static bool setnsat(int ns_fd, const char *ns)
{
	__do_close_prot_errno int fd = -EBADF;

	fd = openat(ns_fd, ns, O_RDONLY | O_CLOEXEC);
	if (fd < 0)
		return false;

	return setns(fd, 0) == 0;
}

static bool change_creds(int ns_fd, cap_t caps, uid_t nsuid, gid_t nsgid, uid_t nsfsuid, gid_t nsfsgid)
{
	__do_close_prot_errno int fd = -EBADF;

	if (prctl(PR_SET_KEEPCAPS, 1))
		return false;

	for (const char **ns = ns_names; ns && *ns; ns++) {
		if (!setnsat(ns_fd, *ns))
			return false;
	}

	if (setegid(nsgid))
		return false;

	setfsgid(nsfsgid);

	if (seteuid(nsuid))
		return false;

	setfsuid(nsfsuid);

	if (cap_set_proc(caps))
		return false;

	return true;
}

static void forksetxattr()
{
	__do_close_prot_errno int cwd_fd = -EBADF, mnt_fd = -EBADF, ns_fd = -EBADF, root_fd = -EBADF, target_fd = -EBADF;
	int flags = 0;
	char *name, *target;
	char path[PATH_MAX];
	uid_t nsfsuid, nsuid;
	gid_t nsfsgid, nsgid;
	pid_t pid = 0;
	cap_t caps;
	cap_flag_value_t flag;
	int whiteout;
	void *data;
	size_t size;

	pid = atoi(advance_arg(true));
	nsuid = atoi(advance_arg(true));
	nsgid = atoi(advance_arg(true));
	nsfsuid = atoi(advance_arg(true));
	nsfsgid = atoi(advance_arg(true));
	name = advance_arg(true);
	target = advance_arg(true);
	flags = atoi(advance_arg(true));
	whiteout = atoi(advance_arg(true));
	size = atoi(advance_arg(true));
	data = advance_arg(true);

	snprintf(path, sizeof(path), "/proc/%d/ns", pid);
	ns_fd = open(path, O_PATH | O_RDONLY | O_CLOEXEC | O_DIRECTORY);
	if (ns_fd < 0) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	snprintf(path, sizeof(path), "/proc/%d/root", pid);
	root_fd = open(path, O_PATH | O_RDONLY | O_CLOEXEC | O_NOFOLLOW);
	if (root_fd < 0) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	snprintf(path, sizeof(path), "/proc/%d/ns/mnt", pid);
	mnt_fd = open(path, O_RDONLY | O_CLOEXEC);
	if (mnt_fd < 0) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	snprintf(path, sizeof(path), "/proc/%d/cwd", pid);
	cwd_fd = open(path, O_PATH | O_RDONLY | O_CLOEXEC);
	if (cwd_fd < 0) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	if (setns(mnt_fd, CLONE_NEWNS)) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	if (!chdirchroot_in_mntns(cwd_fd, root_fd)) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	target_fd = open(target, O_RDONLY | O_CLOEXEC);
	if (target_fd < 0) {
		fprintf(stderr, "%d", errno);
		_exit(EXIT_FAILURE);
	}

	caps = cap_get_pid(pid);
	if (!caps) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	if (whiteout == 1) {
		if (cap_get_flag(caps,  CAP_SYS_ADMIN, CAP_EFFECTIVE, &flag) != 0) {
			fprintf(stderr, "%d", EPERM);
			_exit(EXIT_FAILURE);
		}

		if (flag == CAP_CLEAR) {
			fprintf(stderr, "%d", EPERM);
			_exit(EXIT_FAILURE);
		}
	}

	if (whiteout == 1) {
		if (fsetxattr(target_fd, "trusted.overlay.opaque", "y", 1, flags)) {
			fprintf(stderr, "%d", errno);
			_exit(EXIT_FAILURE);
		}
	} else {
		if (!change_creds(ns_fd, caps, nsuid, nsgid, nsfsuid, nsfsgid)) {
			fprintf(stderr, "%d", EFAULT);
			_exit(EXIT_FAILURE);
		}

		if (fsetxattr(target_fd, name, data, size, flags)) {
			fprintf(stderr, "%d", errno);
			_exit(EXIT_FAILURE);
		}
	}
}

void forksyscall()
{
	char *syscall = NULL;

	// Check that we're root
	if (geteuid() != 0)
		_exit(EXIT_FAILURE);

	// Get the subcommand
	syscall = advance_arg(false);
	if (syscall == NULL ||
	    (strcmp(syscall, "--help") == 0 ||
	     strcmp(syscall, "--version") == 0 || strcmp(syscall, "-h") == 0))
		_exit(EXIT_SUCCESS);

	if (strcmp(syscall, "mknod") == 0)
		forkmknod();
	else if (strcmp(syscall, "setxattr") == 0)
		forksetxattr();
	else
		_exit(EXIT_FAILURE);

	_exit(EXIT_SUCCESS);
}
*/
// #cgo CFLAGS: -std=gnu11 -Wvla
// #cgo LDFLAGS: -lcap
import "C"

type cmdForksyscall struct {
	global *cmdGlobal
}

func (c *cmdForksyscall) Command() *cobra.Command {
	// Main subcommand
	cmd := &cobra.Command{}
	cmd.Use = "forksyscall <syscall> <PID> <path> <mode> <dev>"
	cmd.Short = "Perform syscall operations"
	cmd.Long = `Description:
  Perform syscall operations

  This set of internal commands are used for all seccom-based container syscall
  operations.
`
	cmd.RunE = c.Run
	cmd.Hidden = true

	return cmd
}

func (c *cmdForksyscall) Run(cmd *cobra.Command, args []string) error {
	return fmt.Errorf("This command should have been intercepted in cgo")
}
<commit_msg>forksyscall: introduce acquire_basic_creds()<commit_after>package main

import (
	"fmt"

	"github.com/spf13/cobra"
)

/*
#ifndef _GNU_SOURCE
#define _GNU_SOURCE 1
#endif
#include <fcntl.h>
#include <libgen.h>
#include <sched.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/capability.h>
#include <sys/fsuid.h>
#include <sys/prctl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/vfs.h>
#include <sys/xattr.h>
#include <unistd.h>

#include "include/memory_utils.h"

extern char* advance_arg(bool required);
extern int dosetns(int pid, char *nstype);

static inline bool same_fsinfo(struct stat *s1, struct stat *s2,
			       struct statfs *sfs1, struct statfs *sfs2)
{
	return ((sfs1->f_type == sfs2->f_type) && (s1->st_dev == s2->st_dev));
}

static int fstat_fstatfs(int fd, struct stat *s, struct statfs *sfs)
{
	if (fstat(fd, s))
		return -1;

	if (fstatfs(fd, sfs))
		return -1;

	return 0;
}

static bool chdirchroot_in_mntns(int cwd_fd, int root_fd)
{
	ssize_t len;
	char buf[PATH_MAX];

	if (fchdir(cwd_fd))
		return false;

	len = readlinkat(root_fd, "", buf, sizeof(buf));
	if (len < 0 || len >= sizeof(buf))
		return false;
	buf[len] = '\0';

	if (chroot(buf))
		return false;

	return true;
}

static bool acquire_basic_creds(pid_t pid)
{
	__do_close_prot_errno int cwd_fd = -EBADF, mnt_fd = -EBADF, root_fd = -EBADF;
	char buf[256];

	snprintf(buf, sizeof(buf), "/proc/%d/ns/mnt", pid);
	mnt_fd = open(buf, O_RDONLY | O_CLOEXEC);
	if (mnt_fd < 0)
		return false;

	snprintf(buf, sizeof(buf), "/proc/%d/root", pid);
	root_fd = open(buf, O_PATH | O_RDONLY | O_CLOEXEC | O_NOFOLLOW);
	if (root_fd < 0)
		return false;

	snprintf(buf, sizeof(buf), "/proc/%d/cwd", pid);
	cwd_fd = open(buf, O_PATH | O_RDONLY | O_CLOEXEC);
	if (cwd_fd < 0)
		return false;

	if (setns(mnt_fd, CLONE_NEWNS))
		return false;

	return chdirchroot_in_mntns(cwd_fd, root_fd);
}

// Expects command line to be in the form:
// <PID> <root-uid> <root-gid> <path> <mode> <dev>
static void forkmknod()
{
	__do_close_prot_errno int target_dir_fd = -EBADF;
	char *cur = NULL, *target = NULL, *target_dir = NULL, *target_host = NULL;
	int ret;
	char path[PATH_MAX];
	mode_t mode;
	dev_t dev;
	pid_t pid;
	uid_t fsuid, uid;
	gid_t fsgid, gid;
	struct statfs sfs;
	cap_t caps;

	pid = atoi(advance_arg(true));
	target = advance_arg(true);
	mode = atoi(advance_arg(true));
	dev = atoi(advance_arg(true));
	target_host = advance_arg(true);
	uid = atoi(advance_arg(true));
	gid = atoi(advance_arg(true));
	fsuid = atoi(advance_arg(true));
	fsgid = atoi(advance_arg(true));

	if (!acquire_basic_creds(pid)) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	snprintf(path, sizeof(path), "%s", target);
	target_dir = dirname(path);
	target_dir_fd = open(target_dir, O_PATH | O_RDONLY | O_CLOEXEC | O_DIRECTORY);
	if (target_dir_fd < 0) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	caps = cap_get_pid(pid);
	if (!caps) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	ret = prctl(PR_SET_KEEPCAPS, 1);
	if (ret) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	ret = setegid(gid);
	if (ret) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	setfsgid(fsgid);

	ret = seteuid(uid);
	if (ret) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	setfsuid(fsuid);

	ret = cap_set_proc(caps);
	if (ret) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	ret = fstatfs(target_dir_fd, &sfs);
	if (ret) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	if (sfs.f_flags & MS_NODEV) {
		fprintf(stderr, "%d", EPERM);
		_exit(EXIT_FAILURE);
	}

	// basename() can modify its argument so accessing target_host is
	// invalid from now on.
	ret = mknodat(target_dir_fd, target, mode, dev);
	if (ret) {
		if (errno == EPERM)
			fprintf(stderr, "%d", ENOMEDIUM);
		else
			fprintf(stderr, "%d", errno);

		_exit(EXIT_FAILURE);
	}
}

#ifndef CLONE_NEWCGROUP
#define CLONE_NEWCGROUP	0x02000000
#endif

const char *ns_names[] = { "user", "pid", "uts", "ipc", "net", "cgroup", NULL };

static bool setnsat(int ns_fd, const char *ns)
{
	__do_close_prot_errno int fd = -EBADF;

	fd = openat(ns_fd, ns, O_RDONLY | O_CLOEXEC);
	if (fd < 0)
		return false;

	return setns(fd, 0) == 0;
}

static bool change_creds(int ns_fd, cap_t caps, uid_t nsuid, gid_t nsgid, uid_t nsfsuid, gid_t nsfsgid)
{
	__do_close_prot_errno int fd = -EBADF;

	if (prctl(PR_SET_KEEPCAPS, 1))
		return false;

	for (const char **ns = ns_names; ns && *ns; ns++) {
		if (!setnsat(ns_fd, *ns))
			return false;
	}

	if (setegid(nsgid))
		return false;

	setfsgid(nsfsgid);

	if (seteuid(nsuid))
		return false;

	setfsuid(nsfsuid);

	if (cap_set_proc(caps))
		return false;

	return true;
}

static void forksetxattr()
{
	__do_close_prot_errno int ns_fd = -EBADF, target_fd = -EBADF;
	int flags = 0;
	char *name, *target;
	char path[PATH_MAX];
	uid_t nsfsuid, nsuid;
	gid_t nsfsgid, nsgid;
	pid_t pid = 0;
	cap_t caps;
	cap_flag_value_t flag;
	int whiteout;
	void *data;
	size_t size;

	pid = atoi(advance_arg(true));
	nsuid = atoi(advance_arg(true));
	nsgid = atoi(advance_arg(true));
	nsfsuid = atoi(advance_arg(true));
	nsfsgid = atoi(advance_arg(true));
	name = advance_arg(true);
	target = advance_arg(true);
	flags = atoi(advance_arg(true));
	whiteout = atoi(advance_arg(true));
	size = atoi(advance_arg(true));
	data = advance_arg(true);

	snprintf(path, sizeof(path), "/proc/%d/ns", pid);
	ns_fd = open(path, O_PATH | O_RDONLY | O_CLOEXEC | O_DIRECTORY);
	if (ns_fd < 0) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	if (!acquire_basic_creds(pid)) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	target_fd = open(target, O_RDONLY | O_CLOEXEC);
	if (target_fd < 0) {
		fprintf(stderr, "%d", errno);
		_exit(EXIT_FAILURE);
	}

	caps = cap_get_pid(pid);
	if (!caps) {
		fprintf(stderr, "%d", ENOANO);
		_exit(EXIT_FAILURE);
	}

	if (whiteout == 1) {
		if (cap_get_flag(caps,  CAP_SYS_ADMIN, CAP_EFFECTIVE, &flag) != 0) {
			fprintf(stderr, "%d", EPERM);
			_exit(EXIT_FAILURE);
		}

		if (flag == CAP_CLEAR) {
			fprintf(stderr, "%d", EPERM);
			_exit(EXIT_FAILURE);
		}
	}

	if (whiteout == 1) {
		if (fsetxattr(target_fd, "trusted.overlay.opaque", "y", 1, flags)) {
			fprintf(stderr, "%d", errno);
			_exit(EXIT_FAILURE);
		}
	} else {
		if (!change_creds(ns_fd, caps, nsuid, nsgid, nsfsuid, nsfsgid)) {
			fprintf(stderr, "%d", EFAULT);
			_exit(EXIT_FAILURE);
		}

		if (fsetxattr(target_fd, name, data, size, flags)) {
			fprintf(stderr, "%d", errno);
			_exit(EXIT_FAILURE);
		}
	}
}

void forksyscall()
{
	char *syscall = NULL;

	// Check that we're root
	if (geteuid() != 0)
		_exit(EXIT_FAILURE);

	// Get the subcommand
	syscall = advance_arg(false);
	if (syscall == NULL ||
	    (strcmp(syscall, "--help") == 0 ||
	     strcmp(syscall, "--version") == 0 || strcmp(syscall, "-h") == 0))
		_exit(EXIT_SUCCESS);

	if (strcmp(syscall, "mknod") == 0)
		forkmknod();
	else if (strcmp(syscall, "setxattr") == 0)
		forksetxattr();
	else
		_exit(EXIT_FAILURE);

	_exit(EXIT_SUCCESS);
}
*/
// #cgo CFLAGS: -std=gnu11 -Wvla
// #cgo LDFLAGS: -lcap
import "C"

type cmdForksyscall struct {
	global *cmdGlobal
}

func (c *cmdForksyscall) Command() *cobra.Command {
	// Main subcommand
	cmd := &cobra.Command{}
	cmd.Use = "forksyscall <syscall> <PID> <path> <mode> <dev>"
	cmd.Short = "Perform syscall operations"
	cmd.Long = `Description:
  Perform syscall operations

  This set of internal commands are used for all seccom-based container syscall
  operations.
`
	cmd.RunE = c.Run
	cmd.Hidden = true

	return cmd
}

func (c *cmdForksyscall) Run(cmd *cobra.Command, args []string) error {
	return fmt.Errorf("This command should have been intercepted in cgo")
}
<|endoftext|>


<commit_before>package kernel

import (
	"testing"
)

func TestBuildInitrd(t *testing.T) {
	if name, err := BuildInitrd("/lib/modules/3.14.33-1-lts"); err != nil {
		t.Errorf("%s", err)
	} else {
		t.Errorf("%s", name)

	}

}
<commit_msg>Remove a miss added test file<commit_after><|endoftext|>


<commit_before>package handler

import (
	"OnlineJudge/base"
	"OnlineJudge/db"
	"OnlineJudge/irpc"
	"OnlineJudge/models"
	"OnlineJudge/pbgen/api"

	"log"
	"time"
)

func (this *AdminHandler) Submit(response *api.SubmitResponse, req *api.SubmitRequest) {
	defer PanicHandler(response, this.debug)
	Submit_BuildResponse(this.dbu, response, req,
		this.session.GetUserId(), this.session.GetIPAddr(), true, this.debug)
}

func (this *BasicHandler) Submit(response *api.SubmitResponse, req *api.SubmitRequest) {
	MakeResponseError(response, this.debug, PBLoginRequired, nil)
}

// Need to be tested
// Depend on MetaProblems, OJInfo,
func (this *UserHandler) Submit(response *api.SubmitResponse, req *api.SubmitRequest) {
	defer PanicHandler(response, this.debug)
	Submit_BuildResponse(this.dbu, response, req,
		this.session.GetUserId(), this.session.GetIPAddr(), false, this.debug)
}

func Submit_BuildResponse(
	dbu *db.DBUtil,
	response *api.SubmitResponse,
	req *api.SubmitRequest,
	user_id int64,
	ip_addr string,
	use_hide bool,
	debug bool) {

	// Parse ProblemSid
	pid, err := base.ParseSid(req.GetProblemSid())
	if err != nil {
		MakeResponseError(response, debug, PBBadRequest, err)
		return
	}
	tx := dbu.MustBegin()
	defer dbu.Rollback()

	mp, err := models.Query_MetaProblem_By_OJName_OJPid(tx, pid.OJName,
		pid.OJPid, []string{"meta_pid", "hide", "is_spj"}, nil)
	PanicOnError(err)

	if mp.MetaPid == 0 {
		MakeResponseError(response, debug, PBProblemNotFound, nil)
		return
	}

	// if visible
	if mp.Hide == true && !use_hide {
		MakeResponseError(response, debug, PBProblemNotFound, nil)
		return
	}

	// Add Submission
	subm := models.NewSubmissionModel()
	sub := &models.Submission{
		Status:       "Pending",
		StatusCode:   "wt",
		SubmitTime:   time.Now(),
		Code:         req.GetCode(),
		SubmitIPAddr: ip_addr,
		IsShared:     req.GetIsShared(),

		IsContest: false,
		IsSpj:     mp.IsSpj,
		MetaPidFK: mp.MetaPid,
		UserIdFK:  user_id,
		LangIdFK:  req.GetLanguageId(),
	}
	run_id, err := subm.Insert(tx, sub)
	PanicOnError(err)

	response.RunId = run_id
	dbu.MustCommit()

	// Use RPC to call Daemon to judge the submission
	dbu.MustBegin()
	defer dbu.Rollback()

	helper := irpc.NewHelper()

	if err := helper.Connect(); err != nil {
		// Log the error
		log.Println(err)
		if err := subm.SetSystemError(tx, run_id); err != nil {
			PanicOnError(err)
		}
		dbu.MustCommit()
		return
	}
	defer helper.Disconnect()

	helper.NewClient()
	res, err := helper.StartJudging(run_id)

	if err != nil || res.Received != true {
		// Log the error
		log.Println(err)
		if err := subm.SetSystemError(tx, run_id); err != nil {
			PanicOnError(err)
		}
		dbu.MustCommit()
		return
	}
}
<commit_msg>still fix bugs in submit.go<commit_after>package handler

import (
	"OnlineJudge/base"
	"OnlineJudge/db"
	"OnlineJudge/irpc"
	"OnlineJudge/models"
	"OnlineJudge/pbgen/api"

	"log"
	"time"
)

func (this *AdminHandler) Submit(response *api.SubmitResponse, req *api.SubmitRequest) {
	defer PanicHandler(response, this.debug)
	Submit_BuildResponse(this.dbu, response, req,
		this.session.GetUserId(), this.session.GetIPAddr(), true, this.debug)
}

func (this *BasicHandler) Submit(response *api.SubmitResponse, req *api.SubmitRequest) {
	MakeResponseError(response, this.debug, PBLoginRequired, nil)
}

// Need to be tested
// Depend on MetaProblems, OJInfo,
func (this *UserHandler) Submit(response *api.SubmitResponse, req *api.SubmitRequest) {
	defer PanicHandler(response, this.debug)
	Submit_BuildResponse(this.dbu, response, req,
		this.session.GetUserId(), this.session.GetIPAddr(), false, this.debug)
}

func Submit_BuildResponse(
	dbu *db.DBUtil,
	response *api.SubmitResponse,
	req *api.SubmitRequest,
	user_id int64,
	ip_addr string,
	use_hide bool,
	debug bool) {

	// Parse ProblemSid
	pid, err := base.ParseSid(req.GetProblemSid())
	if err != nil {
		MakeResponseError(response, debug, PBBadRequest, err)
		return
	}
	tx := dbu.MustBegin()
	defer dbu.Rollback()

	mp, err := models.Query_MetaProblem_By_OJName_OJPid(tx, pid.OJName,
		pid.OJPid, []string{"meta_pid", "hide", "is_spj"}, nil)
	PanicOnError(err)

	if mp.MetaPid == 0 {
		MakeResponseError(response, debug, PBProblemNotFound, nil)
		return
	}

	// if visible
	if mp.Hide == true && !use_hide {
		MakeResponseError(response, debug, PBProblemNotFound, nil)
		return
	}

	// Add Submission
	subm := models.NewSubmissionModel()
	sub := &models.Submission{
		Status:       "Pending",
		StatusCode:   "wt",
		SubmitTime:   time.Now(),
		Code:         req.GetCode(),
		SubmitIPAddr: ip_addr,
		IsShared:     req.GetIsShared(),

		IsContest: false,
		IsSpj:     mp.IsSpj,
		MetaPidFK: mp.MetaPid,
		UserIdFK:  user_id,
		LangIdFK:  req.GetLanguageId(),
	}
	run_id, err := subm.Insert(tx, sub)
	PanicOnError(err)

	response.RunId = run_id
	dbu.MustCommit()

	// Use RPC to call Daemon to judge the submission
	go func(dbu *db.DBUtil) {
		tx := dbu.MustBegin()
		defer dbu.Rollback()

		helper := irpc.NewHelper()

		if err := helper.Connect(); err != nil {
			// Log the error
			log.Println("connect error:" + err.Error())
			if err := subm.SetSystemError(tx, run_id); err != nil {
				log.Println("set error:" + err.Error())

			}
			dbu.Commit()
			return
		}
		defer helper.Disconnect()

		helper.NewClient()
		res, err := helper.StartJudging(run_id)

		if err != nil || res.Received != true {
			// Log the error
			log.Println("start judging error:" + err.Error())
			if err := subm.SetSystemError(tx, run_id); err != nil {
				log.Println("set error:" + err.Error())
			}
			dbu.Commit()
			return
		}
	}(dbu)
}
<|endoftext|>


<commit_before>package gorma

const intTmpl = `
type {{$typeName}}Storage interface {
	List(ctx *app.List{{demodel $typeName}}Context) []{{$typeName}}
	Get(ctx *app.Show{{demodel $typeName }}Context) ({{$typeName}}, error)
	Add(ctx *app.Create{{demodel $typeName}}Context) ({{$typeName}}, error)
	Update(ctx *app.Update{{demodel $typeName}}Context) (error)
	Delete(ctx *app.Delete{{demodel $typeName}}Context) (error)
}
`
<commit_msg>trying it<commit_after>package gorma

const intTmpl = `
 {{ $typeName :=  gotypename . 0}}{{$typeName := demodel $typeName}}
type {{$typeName}}Storage interface {
	List(ctx *app.List{{demodel $typeName}}Context) []{{$typeName}}
	Get(ctx *app.Show{{demodel $typeName }}Context) ({{$typeName}}, error)
	Add(ctx *app.Create{{demodel $typeName}}Context) ({{$typeName}}, error)
	Update(ctx *app.Update{{demodel $typeName}}Context) (error)
	Delete(ctx *app.Delete{{demodel $typeName}}Context) (error)
}
`
<|endoftext|>


<commit_before>package lamprey

import (
	"fmt"
	"time"
)

import (
	"butler/null"
	"golog"
	tp "tritium/proto"
	"tritium/whale"
)

type Lamprey struct {
	*whale.Whale
}

type Ctx struct {
	*whale.WhaleContext
	*Lamprey
	level int
}

func NewEngine(logger *golog.Logger) *Lamprey {
	e := &Lamprey{}
	e.Whale = whale.NewEngine(logger)
	return e
}

func (eng *Lamprey) Run(transform *tp.Transform, rrules []*tp.RewriteRule, input interface{}, vars map[string]string, deadline time.Time) (output string, exports [][]string, logs []string) {
	ctx := &Ctx{Lamprey: eng}
	ctx.WhaleContext = whale.NewEngineCtx(eng.Whale, vars, transform, rrules, deadline)

	ctx.Yields = append(ctx.Yields, &whale.YieldBlock{Vars: make(map[string]interface{})})
	ctx.UsePackage(transform.Pkg)
	scope := &whale.Scope{Value: input.(string)}
	obj := transform.Objects[0]
	ctx.Filename = null.GetString(obj.Name)
	ctx.level = 0
	ctx.RunInstruction(scope, obj.Root)
	output = scope.Value.(string)
	exports = ctx.Exports
	logs = ctx.Logs
	return
}

func (ctx *Ctx) RunChildren(scope *whale.Scope, ins *tp.Instruction) (returnValue interface{}) {
	for _, child := range ins.Children {
		returnValue = ctx.RunInstruction(scope, child)
	}
	return
}

func (ctx *Ctx) RunInstruction(scope *whale.Scope, ins *tp.Instruction) (returnValue interface{}) {
	defer func() {
		if x := recover(); x != nil {
			err, ok := x.(error)
			errString := ""
			if ok {
				errString = err.Error()
			} else {
				errString = x.(string)
			}
			if ctx.HadError == false {
				ctx.HadError = true
				errString = errString + "\n" + ins.Type.String() + " " + null.GetString(ins.Value) + "\n\n\nTritium Stack\n=========\n\n"
			}
			errString = errString + ctx.FileAndLine(ins) + "\n"
			panic(errString)
		}
	}()

	// If our object is invalid, then skip it
	if null.GetBool(ins.IsValid) == false {
		panic("Invalid instruction. Should have stopped before linking!")
	}
	indent := ""
	for i := 0; i < ctx.level; i++ {
		indent += ".."
	}

	ctx.level++

	returnValue = ""
	switch *ins.Type {
	case tp.Instruction_BLOCK:
		fmt.Printf("%sInstruction block; scope value:\n%v\n", indent, scope.Value)
		returnValue = ctx.RunChildren(scope, ins)
		fmt.Printf("%sInstruction block returns: %v\n", indent, returnValue)
	case tp.Instruction_TEXT:
		fmt.Printf("%sInstruction text\n", indent)
		returnValue = null.GetString(ins.Value)
		fmt.Printf("%sInstruction text returns: %v\n", indent, returnValue)
	case tp.Instruction_LOCAL_VAR:
		fmt.Printf("%sEval local var\n", indent)
		name := null.GetString(ins.Value)
		vars := ctx.Vars()
		if len(ins.Arguments) > 0 {
			vars[name] = ctx.RunInstruction(scope, ins.Arguments[0])
		}
		fmt.Printf("%s%%%s: %v\n", indent, name, vars[name])
		if len(ins.Children) > 0 {
			ts := &whale.Scope{Value: ctx.Vars()[name]}
			ctx.RunChildren(ts, ins)
			vars[name] = ts.Value
		}
		fmt.Printf("%sreturn: %v\n", indent, vars[name])
		returnValue = vars[name]
	case tp.Instruction_IMPORT:
		obj := ctx.Objects[int(null.GetInt32(ins.ObjectId))]
		curFile := ctx.Filename
		ctx.Filename = null.GetString(obj.Name)
		ctx.RunChildren(scope, obj.Root)
		ctx.Filename = curFile
		fmt.Printf("%sImport: %q\n", indent, ctx.Filename)
	case tp.Instruction_FUNCTION_CALL:
		fun := ctx.Functions[int(null.GetInt32(ins.FunctionId))]
		fmt.Printf("%s%s\n", indent, fun.Name)
		fmt.Printf("%sscope.Value berfore loading vars:\n%v\n", indent, scope.Value)
		args := make([]interface{}, len(ins.Arguments))
		for i, argIns := range ins.Arguments {
			args[i] = ctx.RunInstruction(scope, argIns)
		}
		debugInfo := fmt.Sprintf("%s(", fun.Name)

		for index, v := range args {
			debugInfo += fmt.Sprintf("%q", v)
			if index < len(args)-1 {
				debugInfo += ", "
			}
		}
		debugInfo += ")\n"
		fmt.Printf("%s%s", indent, debugInfo)
		fmt.Printf("%sscope.Value before calling:\n%v\n", indent, scope.Value)

		if null.GetBool(fun.BuiltIn) {
			if f := whale.LookupBuiltIn(fun.Name); f != nil {
				returnValue = f(ctx, scope, ins, args)
			} else {
				panic("missing function: " + fun.Name)
			}
		} else {
			// We are using a user-defined function
			//println("Resetting localvar")
			// Setup the new local var
			vars := make(map[string]interface{}, len(args))
			for i, arg := range fun.Args {
				vars[null.GetString(arg.Name)] = args[i]
			}
			yieldBlock := &whale.YieldBlock{
				Ins:  ins,
				Vars: vars,
			}
			// PUSH!
			ctx.PushYieldBlock(yieldBlock)
			fmt.Printf("%srun children\n", indent)
			returnValue = ctx.RunChildren(scope, fun.Instruction)
			// POP!
			ctx.PopYieldBlock()
		}

		fmt.Printf("%sscope.Value:\n%v\n", indent, scope.Value)
		fmt.Printf("%sreturnValue:\n%v\n", indent, returnValue)
		fmt.Printf("%s%s done\n\n", indent, fun.Name)
	}
	ctx.level--

	return
}
<commit_msg>print info to stderr so it can display when debugging ts<commit_after>package lamprey

import (
	"fmt"
	"time"
	"os"
)

import (
	"butler/null"
	"golog"
	tp "tritium/proto"
	"tritium/whale"
)

type Lamprey struct {
	*whale.Whale
}

type Ctx struct {
	*whale.WhaleContext
	*Lamprey
	level int
}

func NewEngine(logger *golog.Logger) *Lamprey {
	e := &Lamprey{}
	e.Whale = whale.NewEngine(logger)
	return e
}

func (eng *Lamprey) Run(transform *tp.Transform, rrules []*tp.RewriteRule, input interface{}, vars map[string]string, deadline time.Time) (output string, exports [][]string, logs []string) {
	ctx := &Ctx{Lamprey: eng}
	ctx.WhaleContext = whale.NewEngineCtx(eng.Whale, vars, transform, rrules, deadline)

	ctx.Yields = append(ctx.Yields, &whale.YieldBlock{Vars: make(map[string]interface{})})
	ctx.UsePackage(transform.Pkg)
	scope := &whale.Scope{Value: input.(string)}
	obj := transform.Objects[0]
	ctx.Filename = null.GetString(obj.Name)
	ctx.level = 0
	ctx.RunInstruction(scope, obj.Root)
	output = scope.Value.(string)
	exports = ctx.Exports
	logs = ctx.Logs
	return
}

func (ctx *Ctx) RunChildren(scope *whale.Scope, ins *tp.Instruction) (returnValue interface{}) {
	for _, child := range ins.Children {
		returnValue = ctx.RunInstruction(scope, child)
	}
	return
}

func (ctx *Ctx) RunInstruction(scope *whale.Scope, ins *tp.Instruction) (returnValue interface{}) {
	defer func() {
		if x := recover(); x != nil {
			err, ok := x.(error)
			errString := ""
			if ok {
				errString = err.Error()
			} else {
				errString = x.(string)
			}
			if ctx.HadError == false {
				ctx.HadError = true
				errString = errString + "\n" + ins.Type.String() + " " + null.GetString(ins.Value) + "\n\n\nTritium Stack\n=========\n\n"
			}
			errString = errString + ctx.FileAndLine(ins) + "\n"
			panic(errString)
		}
	}()

	// If our object is invalid, then skip it
	if null.GetBool(ins.IsValid) == false {
		panic("Invalid instruction. Should have stopped before linking!")
	}
	indent := ""
	for i := 0; i < ctx.level; i++ {
		indent += ".."
	}

	ctx.level++
	returnValue = ""
	switch *ins.Type {
	case tp.Instruction_BLOCK:
		fmt.Fprintf(os.Stderr, "%sInstruction block; scope value:\n%v\n", indent, scope.Value)
		returnValue = ctx.RunChildren(scope, ins)
		fmt.Fprintf(os.Stderr, "%sInstruction block returns: %v\n", indent, returnValue)
	case tp.Instruction_TEXT:
		fmt.Fprintf(os.Stderr, "%sInstruction text\n", indent)
		returnValue = null.GetString(ins.Value)
		fmt.Fprintf(os.Stderr, "%sInstruction text returns: %v\n", indent, returnValue)
	case tp.Instruction_LOCAL_VAR:
		fmt.Fprintf(os.Stderr, "%sEval local var\n", indent)
		name := null.GetString(ins.Value)
		vars := ctx.Vars()
		if len(ins.Arguments) > 0 {
			vars[name] = ctx.RunInstruction(scope, ins.Arguments[0])
		}
		fmt.Fprintf(os.Stderr, "%s%%%s: %v\n", indent, name, vars[name])
		if len(ins.Children) > 0 {
			ts := &whale.Scope{Value: ctx.Vars()[name]}
			ctx.RunChildren(ts, ins)
			vars[name] = ts.Value
		}
		fmt.Fprintf(os.Stderr, "%sreturn: %v\n", indent, vars[name])
		returnValue = vars[name]
	case tp.Instruction_IMPORT:
		obj := ctx.Objects[int(null.GetInt32(ins.ObjectId))]
		curFile := ctx.Filename
		ctx.Filename = null.GetString(obj.Name)
		ctx.RunChildren(scope, obj.Root)
		ctx.Filename = curFile
		fmt.Fprintf(os.Stderr, "%sImport: %q\n", indent, ctx.Filename)
	case tp.Instruction_FUNCTION_CALL:
		fun := ctx.Functions[int(null.GetInt32(ins.FunctionId))]
		fmt.Fprintf(os.Stderr, "%s%s\n", indent, fun.Name)
		fmt.Fprintf(os.Stderr, "%sscope.Value berfore loading vars:\n%v\n", indent, scope.Value)
		args := make([]interface{}, len(ins.Arguments))
		for i, argIns := range ins.Arguments {
			args[i] = ctx.RunInstruction(scope, argIns)
		}
		debugInfo := fmt.Sprintf("%s(", fun.Name)

		for index, v := range args {
			debugInfo += fmt.Sprintf("%q", v)
			if index < len(args)-1 {
				debugInfo += ", "
			}
		}
		debugInfo += ")\n"
		fmt.Fprintf(os.Stderr, "%s%s", indent, debugInfo)
		fmt.Fprintf(os.Stderr, "%sscope.Value before calling:\n%v\n", indent, scope.Value)

		if null.GetBool(fun.BuiltIn) {
			if f := whale.LookupBuiltIn(fun.Name); f != nil {
				returnValue = f(ctx, scope, ins, args)
			} else {
				panic("missing function: " + fun.Name)
			}
		} else {
			// We are using a user-defined function
			//println("Resetting localvar")
			// Setup the new local var
			vars := make(map[string]interface{}, len(args))
			for i, arg := range fun.Args {
				vars[null.GetString(arg.Name)] = args[i]
			}
			yieldBlock := &whale.YieldBlock{
				Ins:  ins,
				Vars: vars,
			}
			// PUSH!
			ctx.PushYieldBlock(yieldBlock)
			fmt.Fprintf(os.Stderr, "%srun children\n", indent)
			returnValue = ctx.RunChildren(scope, fun.Instruction)
			// POP!
			ctx.PopYieldBlock()
		}

		fmt.Fprintf(os.Stderr, "%sscope.Value:\n%v\n", indent, scope.Value)
		fmt.Fprintf(os.Stderr, "%sreturnValue:\n%v\n", indent, returnValue)
		fmt.Fprintf(os.Stderr, "%s%s done\n\n", indent, fun.Name)
	}
	ctx.level--

	return
}
<|endoftext|>


<commit_before>package db

import "time"

// MockDB is an in-memory database repository implementing the DB interface
// used for testing
type MockDB struct {
	installations map[int]GHInstallation // installationID -> exists
	err           error
	Tools         []Tool
}

// Ensure MockDB implements DB
var _ DB = (*MockDB)(nil)

// NewMockDB returns an MockDB
func NewMockDB() *MockDB {
	return &MockDB{
		installations: make(map[int]GHInstallation),
	}
}

func (db *MockDB) ForceError(err error) {
	db.err = err
}

// AddGHInstallation implements DB interface
func (db *MockDB) AddGHInstallation(installationID, accountID, senderID int) error {
	db.installations[installationID] = GHInstallation{
		InstallationID: installationID,
		AccountID:      accountID,
		SenderID:       senderID,
	}
	return db.err
}

// RemoveGHInstallation implements DB interface
func (db *MockDB) RemoveGHInstallation(installationID int) error {
	delete(db.installations, installationID)
	return db.err
}

// EnableGHInstallation enables a gh installation
func (db *MockDB) EnableGHInstallation(installationID int) error {
	install := db.installations[installationID]
	install.enabledAt = time.Unix(1, 0)
	db.installations[installationID] = install
	return db.err
}

// GetGHInstallation implements DB interface
func (db *MockDB) GetGHInstallation(installationID int) (*GHInstallation, error) {
	if installation, ok := db.installations[installationID]; ok {
		return &installation, db.err
	}
	return nil, db.err
}

// ListTools implements DB interface
func (db *MockDB) ListTools() ([]Tool, error) {
	return db.Tools, nil
}

// StartAnalysis implements the DB interface.
func (db *MockDB) StartAnalysis(ghInstallationID, repositoryID int) (*Analysis, error) {
	analysis := NewAnalysis()
	analysis.ID = 99
	return analysis, nil
}

// FinishAnalysis implements the DB interface.
func (db *MockDB) FinishAnalysis(analysisID int, status AnalysisStatus, analysis *Analysis) error {
	return nil
}

// GetAnalysis implements the DB interface.
func (db *MockDB) GetAnalysis(analysisID int) (*Analysis, error) {
	return nil, nil
}
<commit_msg>Golint issue<commit_after>package db

import "time"

// MockDB is an in-memory database repository implementing the DB interface
// used for testing
type MockDB struct {
	installations map[int]GHInstallation // installationID -> exists
	err           error
	Tools         []Tool
}

// Ensure MockDB implements DB
var _ DB = (*MockDB)(nil)

// NewMockDB returns an MockDB
func NewMockDB() *MockDB {
	return &MockDB{
		installations: make(map[int]GHInstallation),
	}
}

// ForceError forces MockDB to return err on all methods that return an error.
func (db *MockDB) ForceError(err error) {
	db.err = err
}

// AddGHInstallation implements DB interface
func (db *MockDB) AddGHInstallation(installationID, accountID, senderID int) error {
	db.installations[installationID] = GHInstallation{
		InstallationID: installationID,
		AccountID:      accountID,
		SenderID:       senderID,
	}
	return db.err
}

// RemoveGHInstallation implements DB interface
func (db *MockDB) RemoveGHInstallation(installationID int) error {
	delete(db.installations, installationID)
	return db.err
}

// EnableGHInstallation enables a gh installation
func (db *MockDB) EnableGHInstallation(installationID int) error {
	install := db.installations[installationID]
	install.enabledAt = time.Unix(1, 0)
	db.installations[installationID] = install
	return db.err
}

// GetGHInstallation implements DB interface
func (db *MockDB) GetGHInstallation(installationID int) (*GHInstallation, error) {
	if installation, ok := db.installations[installationID]; ok {
		return &installation, db.err
	}
	return nil, db.err
}

// ListTools implements DB interface
func (db *MockDB) ListTools() ([]Tool, error) {
	return db.Tools, nil
}

// StartAnalysis implements the DB interface.
func (db *MockDB) StartAnalysis(ghInstallationID, repositoryID int) (*Analysis, error) {
	analysis := NewAnalysis()
	analysis.ID = 99
	return analysis, nil
}

// FinishAnalysis implements the DB interface.
func (db *MockDB) FinishAnalysis(analysisID int, status AnalysisStatus, analysis *Analysis) error {
	return nil
}

// GetAnalysis implements the DB interface.
func (db *MockDB) GetAnalysis(analysisID int) (*Analysis, error) {
	return nil, nil
}
<|endoftext|>


<commit_before>package marshalling

type Marshaler interface {
	Marshal(interface{}) error
}

type Unmarshaler interface {
	Unmarshal(interface{}) error
}

type MarshalerUnmarshaler interface {
	Marshaler
	Unmarshaler
}
<commit_msg>Put in proper package.<commit_after>package marshaling

type Marshaler interface {
	Marshal(interface{}) error
}

type Unmarshaler interface {
	Unmarshal(interface{}) error
}

type MarshalUnmarshaler interface {
	Marshaler
	Unmarshaler
}
<|endoftext|>


<commit_before>// Copyright 2018 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package sqlbuilder_test

import (
	"database/sql"
	"testing"

	"github.com/issue9/assert"
	"github.com/issue9/orm/v2"
	"github.com/issue9/orm/v2/dialect"
	"github.com/issue9/orm/v2/internal/sqltest"
	"github.com/issue9/orm/v2/sqlbuilder"

	_ "github.com/go-sql-driver/mysql"
	_ "github.com/lib/pq"
	_ "github.com/mattn/go-sqlite3"
)

var (
	_ sqlbuilder.SQLer       = &sqlbuilder.SelectStmt{}
	_ sqlbuilder.WhereStmter = &sqlbuilder.SelectStmt{}
)

func TestSelect_sqlite3(t *testing.T) {
	a := assert.New(t)
	e, err := orm.NewDB("sqlite3", "./test.db", "test_", dialect.Sqlite3())
	a.NotError(err)
	s := sqlbuilder.Select(e, e.Dialect()).Select("c1", "column2 as c2", "c3").
		From("table").
		And("c1=?", 1).
		Or("c2=@c2", sql.Named("c2", 2)).
		Limit(10, 0).
		Desc("c1")
	a.NotNil(s)
	query, args, err := s.SQL()
	a.NotError(err)
	a.Equal(args, []interface{}{1, sql.Named("c2", 2), 10, 0})
	sqltest.Equal(a, query, "select c1,column2 as c2,c3 from table where c1=? or c2=@c2 order by c1 desc limit ? offset ?")

	// count
	s.Count("count(*) as cnt")
	query, args, err = s.SQL()
	a.NotError(err)
	a.Equal(args, []interface{}{1, sql.Named("c2", 2)})
	sqltest.Equal(a, query, "select count(*) as cnt from table where c1=? or c2=@c2")

	// reset
	s.Reset()
	query, args, err = s.SQL()
	a.Error(err)

	s.Distinct().
		Select("c1,c2,c3").
		Join("left", "users as u", "a.id=u.id").
		Where("id=?", 5).
		Asc("id").
		From("table1 as t")
	query, args, err = s.SQL()
	a.NotError(err)
	a.Equal(args, []interface{}{5})
	sqltest.Equal(a, query, "select distinct c1,c2,c3 from table1 as t left join users as u on a.id=u.id where id=? order by id asc")

	// s.reset，没有  where
	s.Reset()
	s.Select("c1,c2").From("#tb1")
	query, args, err = s.SQL()
	a.NotError(err).Empty(args)
	sqltest.Equal(a, query, "select c1,c2 from #tb1")
}
<commit_msg>[sqlbuilder] 添加部分测试功能<commit_after>// Copyright 2018 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package sqlbuilder_test

import (
	"database/sql"
	"testing"

	"github.com/issue9/assert"
	"github.com/issue9/orm/v2"
	"github.com/issue9/orm/v2/dialect"
	"github.com/issue9/orm/v2/internal/sqltest"
	"github.com/issue9/orm/v2/internal/testconfig"
	"github.com/issue9/orm/v2/sqlbuilder"

	_ "github.com/go-sql-driver/mysql"
	_ "github.com/lib/pq"
	_ "github.com/mattn/go-sqlite3"
)

var (
	_ sqlbuilder.SQLer       = &sqlbuilder.SelectStmt{}
	_ sqlbuilder.WhereStmter = &sqlbuilder.SelectStmt{}
)

func TestSelect_Query(t *testing.T) {
	a := assert.New(t)
	db := createTable(a)
	defer testconfig.CloseDB(db, a)
	insertData(a, db)

	sql := sqlbuilder.Select(db, db.Dialect()).Select("*").
		From("#user").
		Where("id<?", 5).
		Desc("id")

	id, err := sql.QueryInt("id")
	a.NotError(err).
		Equal(id, 4)

	obj := &user{}
	size, err := sql.QueryObject(true, obj)
	a.NotError(err).Equal(1, size)
	a.Equal(obj.ID, 4)

	cnt, err := sql.Count("count(*) AS cnt").QueryInt("cnt")
	a.NotError(err).
		Equal(cnt, 4)
}

func TestSelect_sqlite3(t *testing.T) {
	a := assert.New(t)
	e, err := orm.NewDB("sqlite3", "./test.db", "test_", dialect.Sqlite3())
	a.NotError(err)
	s := sqlbuilder.Select(e, e.Dialect()).Select("c1", "column2 as c2", "c3").
		From("table").
		And("c1=?", 1).
		Or("c2=@c2", sql.Named("c2", 2)).
		Limit(10, 0).
		Desc("c1")
	a.NotNil(s)
	query, args, err := s.SQL()
	a.NotError(err)
	a.Equal(args, []interface{}{1, sql.Named("c2", 2), 10, 0})
	sqltest.Equal(a, query, "select c1,column2 as c2,c3 from table where c1=? or c2=@c2 order by c1 desc limit ? offset ?")

	// count
	s.Count("count(*) as cnt")
	query, args, err = s.SQL()
	a.NotError(err)
	a.Equal(args, []interface{}{1, sql.Named("c2", 2)})
	sqltest.Equal(a, query, "select count(*) as cnt from table where c1=? or c2=@c2")

	// reset
	s.Reset()
	query, args, err = s.SQL()
	a.Error(err)

	s.Distinct().
		Select("c1,c2,c3").
		Join("left", "users as u", "a.id=u.id").
		Where("id=?", 5).
		Asc("id").
		From("table1 as t")
	query, args, err = s.SQL()
	a.NotError(err)
	a.Equal(args, []interface{}{5})
	sqltest.Equal(a, query, "select distinct c1,c2,c3 from table1 as t left join users as u on a.id=u.id where id=? order by id asc")

	// s.reset，没有  where
	s.Reset()
	s.Select("c1,c2").From("#tb1")
	query, args, err = s.SQL()
	a.NotError(err).Empty(args)
	sqltest.Equal(a, query, "select c1,c2 from #tb1")
}
<|endoftext|>


<commit_before>/*
Copyright 2022 Gravitational, 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 ec2

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/gravitational/teleport/api/types"
	"github.com/gravitational/teleport/lib/cloud/aws"
	"github.com/gravitational/teleport/lib/utils"
	"github.com/gravitational/trace"
	"github.com/jonboulle/clockwork"
	"github.com/sirupsen/logrus"
)

const (
	// AWSNamespace is used as the namespace prefix for any labels
	// imported from AWS.
	AWSNamespace = "aws"
	// ec2LabelUpdatePeriod is the period for updating EC2 labels.
	ec2LabelUpdatePeriod = time.Hour
)

// Config is the configuration for the EC2 label service.
type Config struct {
	Client aws.InstanceMetadata
	Clock  clockwork.Clock
	Log    logrus.FieldLogger
}

func (conf *Config) checkAndSetDefaults(ctx context.Context) error {
	if conf.Client == nil {
		client, err := utils.NewInstanceMetadataClient(ctx)
		if err != nil {
			return trace.Wrap(err)
		}
		conf.Client = client
	}
	if conf.Clock == nil {
		conf.Clock = clockwork.NewRealClock()
	}
	if conf.Log == nil {
		conf.Log = logrus.WithField(trace.Component, "ec2labels")
	}
	return nil
}

// EC2 is a service that periodically imports tags from EC2 via instance
// metadata.
type EC2 struct {
	c      *Config
	mu     sync.RWMutex
	labels map[string]string

	closeCh chan struct{}
}

func New(ctx context.Context, c *Config) (*EC2, error) {
	if err := c.checkAndSetDefaults(ctx); err != nil {
		return nil, trace.Wrap(err)
	}
	return &EC2{
		c:       c,
		labels:  make(map[string]string),
		closeCh: make(chan struct{}),
	}, nil
}

// Get returns the list of updated EC2 labels.
func (l *EC2) Get() map[string]string {
	l.mu.RLock()
	defer l.mu.RUnlock()
	labels := make(map[string]string)
	for k, v := range l.labels {
		labels[k] = v
	}
	return labels
}

// Apply adds EC2 labels to the provided resource.
func (l *EC2) Apply(r types.ResourceWithLabels) {
	labels := l.Get()
	for k, v := range r.GetStaticLabels() {
		labels[k] = v
	}
	r.SetStaticLabels(labels)
}

// Sync will block and synchronously update EC2 labels.
func (l *EC2) Sync(ctx context.Context) error {
	m := make(map[string]string)

	tags, err := l.c.Client.GetTagKeys(ctx)
	if err != nil {
		if trace.IsNotFound(err) {
			l.c.Log.Warningf("could not fetch tags, please ensure 'allow instance tags in metadata' is enabled on the instance")
			return nil
		}
		return trace.Wrap(err)
	}

	currentLabels := l.Get()
	var errors []error
	for _, t := range tags {
		if !types.IsValidLabelKey(t) {
			l.c.Log.Debugf("Skipping EC2 tag %q, not a valid label key.", t)
			continue
		}

		key := toAWSLabel(t)
		value, err := l.c.Client.GetTagValue(ctx, t)
		if err != nil {
			errors = append(errors, err)
			// If we know the key exists but GetTagValue failed, use the current value
			// if it exists.
			if currentValue, ok := currentLabels[key]; ok {
				m[key] = currentValue
			} else {
				m[key] = ""
			}
		} else {
			m[key] = value
		}

	}

	l.mu.Lock()
	defer l.mu.Unlock()
	l.labels = m
	return trace.NewAggregate(errors...)
}

// Start will start a loop that continually keeps EC2 labels updated.
func (l *EC2) Start(ctx context.Context) {
	go l.periodicUpdateLabels(ctx)
}

func (l *EC2) periodicUpdateLabels(ctx context.Context) {
	ticker := l.c.Clock.NewTicker(ec2LabelUpdatePeriod)
	defer ticker.Stop()

	for {
		if err := l.Sync(ctx); err != nil {
			l.c.Log.Warningf("Error fetching EC2 tags: %v", err)
		}
		select {
		case <-ticker.Chan():
		case <-ctx.Done():
			return
		}
	}
}

// toAWSLabel formats labels coming from EC2.
func toAWSLabel(key string) string {
	return fmt.Sprintf("%s/%s", AWSNamespace, key)
}
<commit_msg>Lower EC2 label log frequency (#15179)<commit_after>/*
Copyright 2022 Gravitational, 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 ec2

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/gravitational/teleport/api/types"
	"github.com/gravitational/teleport/lib/cloud/aws"
	"github.com/gravitational/teleport/lib/utils"
	"github.com/gravitational/trace"
	"github.com/jonboulle/clockwork"
	"github.com/sirupsen/logrus"
)

const (
	// AWSNamespace is used as the namespace prefix for any labels
	// imported from AWS.
	AWSNamespace = "aws"
	// ec2LabelUpdatePeriod is the period for updating EC2 labels.
	ec2LabelUpdatePeriod = time.Hour
)

// Config is the configuration for the EC2 label service.
type Config struct {
	Client aws.InstanceMetadata
	Clock  clockwork.Clock
	Log    logrus.FieldLogger
}

func (conf *Config) checkAndSetDefaults(ctx context.Context) error {
	if conf.Client == nil {
		client, err := utils.NewInstanceMetadataClient(ctx)
		if err != nil {
			return trace.Wrap(err)
		}
		conf.Client = client
	}
	if conf.Clock == nil {
		conf.Clock = clockwork.NewRealClock()
	}
	if conf.Log == nil {
		conf.Log = logrus.WithField(trace.Component, "ec2labels")
	}
	return nil
}

// EC2 is a service that periodically imports tags from EC2 via instance
// metadata.
type EC2 struct {
	c      *Config
	mu     sync.RWMutex
	labels map[string]string

	closeCh chan struct{}
	// loggedInstanceTagsNotFound is used to ensure that the error message for
	// incorrectly configured instance metadata is only logged once.
	loggedInstanceTagsNotFound bool
}

func New(ctx context.Context, c *Config) (*EC2, error) {
	if err := c.checkAndSetDefaults(ctx); err != nil {
		return nil, trace.Wrap(err)
	}
	return &EC2{
		c:       c,
		labels:  make(map[string]string),
		closeCh: make(chan struct{}),
	}, nil
}

// Get returns the list of updated EC2 labels.
func (l *EC2) Get() map[string]string {
	l.mu.RLock()
	defer l.mu.RUnlock()
	labels := make(map[string]string)
	for k, v := range l.labels {
		labels[k] = v
	}
	return labels
}

// Apply adds EC2 labels to the provided resource.
func (l *EC2) Apply(r types.ResourceWithLabels) {
	labels := l.Get()
	for k, v := range r.GetStaticLabels() {
		labels[k] = v
	}
	r.SetStaticLabels(labels)
}

// Sync will block and synchronously update EC2 labels.
func (l *EC2) Sync(ctx context.Context) error {
	m := make(map[string]string)

	tags, err := l.c.Client.GetTagKeys(ctx)
	if err != nil {
		if trace.IsNotFound(err) {
			// Only show the error the first time around.
			if !l.loggedInstanceTagsNotFound {
				l.loggedInstanceTagsNotFound = true
				l.c.Log.Warning("Could not fetch EC2 instance's tags, please ensure 'allow instance tags in metadata' is enabled on the instance.")
			}
			return nil
		}
		return trace.Wrap(err)
	}
	l.loggedInstanceTagsNotFound = false

	currentLabels := l.Get()
	var errors []error
	for _, t := range tags {
		if !types.IsValidLabelKey(t) {
			l.c.Log.Debugf("Skipping EC2 tag %q, not a valid label key.", t)
			continue
		}

		key := toAWSLabel(t)
		value, err := l.c.Client.GetTagValue(ctx, t)
		if err != nil {
			errors = append(errors, err)
			// If we know the key exists but GetTagValue failed, use the current value
			// if it exists.
			if currentValue, ok := currentLabels[key]; ok {
				m[key] = currentValue
			} else {
				m[key] = ""
			}
		} else {
			m[key] = value
		}

	}

	l.mu.Lock()
	defer l.mu.Unlock()
	l.labels = m
	return trace.NewAggregate(errors...)
}

// Start will start a loop that continually keeps EC2 labels updated.
func (l *EC2) Start(ctx context.Context) {
	go l.periodicUpdateLabels(ctx)
}

func (l *EC2) periodicUpdateLabels(ctx context.Context) {
	ticker := l.c.Clock.NewTicker(ec2LabelUpdatePeriod)
	defer ticker.Stop()

	for {
		if err := l.Sync(ctx); err != nil {
			l.c.Log.Warningf("Error fetching EC2 tags: %v", err)
		}
		select {
		case <-ticker.Chan():
		case <-ctx.Done():
			return
		}
	}
}

// toAWSLabel formats labels coming from EC2.
func toAWSLabel(key string) string {
	return fmt.Sprintf("%s/%s", AWSNamespace, key)
}
<|endoftext|>