<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package router

import (
	"io"
	"sync"

	"github.com/TheThingsNetwork/ttn/api"
	"github.com/TheThingsNetwork/ttn/api/discovery"
	"github.com/TheThingsNetwork/ttn/api/gateway"
	"github.com/TheThingsNetwork/ttn/utils/errors"

	"golang.org/x/net/context"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
)

// GatewayClient is used as the main client for Gateways to communicate with the Router
type GatewayClient interface {
	SendGatewayStatus(*gateway.Status) error
	SendUplink(*UplinkMessage) error
	Subscribe() (<-chan *DownlinkMessage, <-chan error, error)
	Unsbscribe() error
	Activate(*DeviceActivationRequest) (*DeviceActivationResponse, error)
	Close() error
}

// NewClient returns a new Client
func NewClient(routerAnnouncement *discovery.Announcement) (*Client, error) {
	conn, err := routerAnnouncement.Dial()
	if err != nil {
		return nil, err
	}
	client := &Client{
		conn:   conn,
		client: NewRouterClient(conn),
	}
	return client, nil
}

// Client is a wrapper around RouterClient
type Client struct {
	mutex    sync.Mutex
	conn     *grpc.ClientConn
	client   RouterClient
	gateways []GatewayClient
}

// ForGateway returns a GatewayClient that is configured for the provided gateway
func (c *Client) ForGateway(gatewayID string, tokenFunc func() string) GatewayClient {
	defer c.mutex.Unlock()
	c.mutex.Lock()
	gatewayClient := NewGatewayClient(c.client, gatewayID, tokenFunc)
	c.gateways = append(c.gateways, gatewayClient)
	return gatewayClient
}

// Close purges the cache and closes the connection with the Router
func (c *Client) Close() error {
	defer c.mutex.Unlock()
	c.mutex.Lock()
	for _, gateway := range c.gateways {
		gateway.Close()
	}
	return c.conn.Close()
}

// NewGatewayClient returns a new GatewayClient
func NewGatewayClient(client RouterClient, gatewayID string, tokenFunc func() string) GatewayClient {
	gatewayClient := &gatewayClient{
		client:    client,
		id:        gatewayID,
		tokenFunc: tokenFunc,
	}
	return gatewayClient
}

type gatewayClient struct {
	mutex             sync.Mutex
	id                string
	tokenFunc         func() string
	client            RouterClient
	gatewayStatus     Router_GatewayStatusClient
	stopGatewayStatus chan bool
	uplink            Router_UplinkClient
	stopUplink        chan bool
	downlink          Router_SubscribeClient
	stopDownlink      chan bool
}

func (c *gatewayClient) getContext() context.Context {
	md := metadata.Pairs(
		"id", c.id,
		"token", c.tokenFunc(),
	)
	gatewayContext := metadata.NewContext(context.Background(), md)
	return gatewayContext
}

func (c *gatewayClient) setupGatewayStatus() error {
	api.GetLogger().Debugf("Setting up gateway status stream for %s...", c.id)
	gatewayStatusClient, err := c.client.GatewayStatus(c.getContext())
	if err != nil {
		return err
	}
	c.gatewayStatus = gatewayStatusClient
	c.stopGatewayStatus = make(chan bool)
	go func() {
		var msg interface{}
		for {
			select {
			case <-c.stopGatewayStatus:
				return
			default:
				if err := gatewayStatusClient.RecvMsg(msg); err != nil {
					api.GetLogger().Warnf("Error in gateway status stream for %s: %s", c.id, err.Error())
					c.teardownGatewayStatus()
					return
				}
				api.GetLogger().Debugf("Received: %v", msg)
			}
		}
	}()
	return nil
}

func (c *gatewayClient) teardownGatewayStatus() {
	defer c.mutex.Unlock()
	c.mutex.Lock()
	if c.gatewayStatus != nil {
		api.GetLogger().Debugf("Closing gateway status stream for %s...", c.id)
		close(c.stopGatewayStatus)
		c.gatewayStatus.CloseSend()
		c.gatewayStatus = nil
	}
}

func (c *gatewayClient) SendGatewayStatus(status *gateway.Status) error {
	defer c.mutex.Unlock()
	c.mutex.Lock()
	if c.gatewayStatus == nil {
		if err := c.setupGatewayStatus(); err != nil {
			return errors.FromGRPCError(err)
		}
	}
	if err := c.gatewayStatus.Send(status); err != nil {
		if err == io.EOF {
			api.GetLogger().Warnf("Could not send gateway status for %s on closed stream", c.id)
			go c.teardownGatewayStatus()
			return errors.FromGRPCError(err)
		}
		api.GetLogger().Warnf("Error sending gateway status for %s: %s", c.id, err.Error())
		return errors.FromGRPCError(err)
	}
	return nil
}

func (c *gatewayClient) setupUplink() error {
	api.GetLogger().Debugf("Setting up uplink stream for %s...", c.id)
	uplinkClient, err := c.client.Uplink(c.getContext())
	if err != nil {
		return err
	}
	c.uplink = uplinkClient
	c.stopUplink = make(chan bool)
	go func() {
		var msg interface{}
		for {
			select {
			case <-c.stopUplink:
				return
			default:
				if err := uplinkClient.RecvMsg(msg); err != nil {
					api.GetLogger().Warnf("Error in uplink stream for %s: %s", c.id, err.Error())
					c.teardownUplink()
					return
				}
				api.GetLogger().Debugf("Received: %v", msg)

			}
		}
	}()
	return nil
}

func (c *gatewayClient) teardownUplink() {
	defer c.mutex.Unlock()
	c.mutex.Lock()
	if c.uplink != nil {
		api.GetLogger().Debugf("Closing uplink stream for %s...", c.id)
		close(c.stopUplink)
		c.uplink.CloseSend()
		c.uplink = nil
	}
}

func (c *gatewayClient) SendUplink(uplink *UplinkMessage) error {
	defer c.mutex.Unlock()
	c.mutex.Lock()
	if c.uplink == nil {
		if err := c.setupUplink(); err != nil {
			return errors.FromGRPCError(err)
		}
	}
	if err := c.uplink.Send(uplink); err != nil {
		if err == io.EOF {
			api.GetLogger().Warnf("Could not send uplink for %s on closed stream", c.id)
			go c.teardownUplink()
			return errors.FromGRPCError(err)
		}
		api.GetLogger().Warnf("Error sending uplink for %s: %s", c.id, err.Error())
		return errors.FromGRPCError(err)
	}
	return nil
}

func (c *gatewayClient) setupDownlink() error {
	api.GetLogger().Debugf("Setting up downlink stream for %s...", c.id)
	ctx, cancel := context.WithCancel(c.getContext())
	downlinkClient, err := c.client.Subscribe(ctx, &SubscribeRequest{})
	if err != nil {
		return err
	}
	c.stopDownlink = make(chan bool)
	go func() {
		<-c.stopDownlink
		cancel()
	}()
	c.downlink = downlinkClient
	return nil
}

func (c *gatewayClient) teardownDownlink() {
	defer c.mutex.Unlock()
	c.mutex.Lock()
	if c.downlink != nil {
		api.GetLogger().Debugf("Closing downlink stream for %s...", c.id)
		close(c.stopDownlink)
		c.downlink.CloseSend()
		c.downlink = nil
	}
}

func (c *gatewayClient) Subscribe() (<-chan *DownlinkMessage, <-chan error, error) {
	defer c.mutex.Unlock()
	c.mutex.Lock()
	if c.downlink == nil {
		if err := c.setupDownlink(); err != nil {
			return nil, nil, errors.FromGRPCError(err)
		}
	}
	downChan := make(chan *DownlinkMessage)
	errChan := make(chan error)
	go func() {
		defer func() {
			close(downChan)
			close(errChan)
		}()
		for {
			select {
			case <-c.stopDownlink:
				return
			default:
				downlink, err := c.downlink.Recv()
				if err != nil {
					if grpc.Code(err) == codes.Canceled {
						api.GetLogger().Debugf("Downlink stream for %s was canceled", c.id)
						errChan <- nil
					} else {
						api.GetLogger().Warnf("Error receiving gateway downlink for %s: %s", c.id, err.Error())
						errChan <- errors.FromGRPCError(err)
					}
					c.teardownDownlink()
					return
				}
				downChan <- downlink
			}
		}
	}()
	return downChan, errChan, nil
}

func (c *gatewayClient) Unsbscribe() error {
	c.teardownDownlink()
	return nil
}

func (c *gatewayClient) Activate(req *DeviceActivationRequest) (*DeviceActivationResponse, error) {
	return c.client.Activate(c.getContext(), req)
}

func (c *gatewayClient) Close() error {
	c.teardownGatewayStatus()
	c.teardownUplink()
	c.teardownDownlink()
	return nil
}
<commit_msg>Fix typo in Router client<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package router

import (
	"io"
	"sync"

	"github.com/TheThingsNetwork/ttn/api"
	"github.com/TheThingsNetwork/ttn/api/discovery"
	"github.com/TheThingsNetwork/ttn/api/gateway"
	"github.com/TheThingsNetwork/ttn/utils/errors"

	"golang.org/x/net/context"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
)

// GatewayClient is used as the main client for Gateways to communicate with the Router
type GatewayClient interface {
	SendGatewayStatus(*gateway.Status) error
	SendUplink(*UplinkMessage) error
	Subscribe() (<-chan *DownlinkMessage, <-chan error, error)
	Unsubscribe() error
	Activate(*DeviceActivationRequest) (*DeviceActivationResponse, error)
	Close() error
}

// NewClient returns a new Client
func NewClient(routerAnnouncement *discovery.Announcement) (*Client, error) {
	conn, err := routerAnnouncement.Dial()
	if err != nil {
		return nil, err
	}
	client := &Client{
		conn:   conn,
		client: NewRouterClient(conn),
	}
	return client, nil
}

// Client is a wrapper around RouterClient
type Client struct {
	mutex    sync.Mutex
	conn     *grpc.ClientConn
	client   RouterClient
	gateways []GatewayClient
}

// ForGateway returns a GatewayClient that is configured for the provided gateway
func (c *Client) ForGateway(gatewayID string, tokenFunc func() string) GatewayClient {
	defer c.mutex.Unlock()
	c.mutex.Lock()
	gatewayClient := NewGatewayClient(c.client, gatewayID, tokenFunc)
	c.gateways = append(c.gateways, gatewayClient)
	return gatewayClient
}

// Close purges the cache and closes the connection with the Router
func (c *Client) Close() error {
	defer c.mutex.Unlock()
	c.mutex.Lock()
	for _, gateway := range c.gateways {
		gateway.Close()
	}
	return c.conn.Close()
}

// NewGatewayClient returns a new GatewayClient
func NewGatewayClient(client RouterClient, gatewayID string, tokenFunc func() string) GatewayClient {
	gatewayClient := &gatewayClient{
		client:    client,
		id:        gatewayID,
		tokenFunc: tokenFunc,
	}
	return gatewayClient
}

type gatewayClient struct {
	mutex             sync.Mutex
	id                string
	tokenFunc         func() string
	client            RouterClient
	gatewayStatus     Router_GatewayStatusClient
	stopGatewayStatus chan bool
	uplink            Router_UplinkClient
	stopUplink        chan bool
	downlink          Router_SubscribeClient
	stopDownlink      chan bool
}

func (c *gatewayClient) getContext() context.Context {
	md := metadata.Pairs(
		"id", c.id,
		"token", c.tokenFunc(),
	)
	gatewayContext := metadata.NewContext(context.Background(), md)
	return gatewayContext
}

func (c *gatewayClient) setupGatewayStatus() error {
	api.GetLogger().Debugf("Setting up gateway status stream for %s...", c.id)
	gatewayStatusClient, err := c.client.GatewayStatus(c.getContext())
	if err != nil {
		return err
	}
	c.gatewayStatus = gatewayStatusClient
	c.stopGatewayStatus = make(chan bool)
	go func() {
		var msg interface{}
		for {
			select {
			case <-c.stopGatewayStatus:
				return
			default:
				if err := gatewayStatusClient.RecvMsg(msg); err != nil {
					api.GetLogger().Warnf("Error in gateway status stream for %s: %s", c.id, err.Error())
					c.teardownGatewayStatus()
					return
				}
				api.GetLogger().Debugf("Received: %v", msg)
			}
		}
	}()
	return nil
}

func (c *gatewayClient) teardownGatewayStatus() {
	defer c.mutex.Unlock()
	c.mutex.Lock()
	if c.gatewayStatus != nil {
		api.GetLogger().Debugf("Closing gateway status stream for %s...", c.id)
		close(c.stopGatewayStatus)
		c.gatewayStatus.CloseSend()
		c.gatewayStatus = nil
	}
}

func (c *gatewayClient) SendGatewayStatus(status *gateway.Status) error {
	defer c.mutex.Unlock()
	c.mutex.Lock()
	if c.gatewayStatus == nil {
		if err := c.setupGatewayStatus(); err != nil {
			return errors.FromGRPCError(err)
		}
	}
	if err := c.gatewayStatus.Send(status); err != nil {
		if err == io.EOF {
			api.GetLogger().Warnf("Could not send gateway status for %s on closed stream", c.id)
			go c.teardownGatewayStatus()
			return errors.FromGRPCError(err)
		}
		api.GetLogger().Warnf("Error sending gateway status for %s: %s", c.id, err.Error())
		return errors.FromGRPCError(err)
	}
	return nil
}

func (c *gatewayClient) setupUplink() error {
	api.GetLogger().Debugf("Setting up uplink stream for %s...", c.id)
	uplinkClient, err := c.client.Uplink(c.getContext())
	if err != nil {
		return err
	}
	c.uplink = uplinkClient
	c.stopUplink = make(chan bool)
	go func() {
		var msg interface{}
		for {
			select {
			case <-c.stopUplink:
				return
			default:
				if err := uplinkClient.RecvMsg(msg); err != nil {
					api.GetLogger().Warnf("Error in uplink stream for %s: %s", c.id, err.Error())
					c.teardownUplink()
					return
				}
				api.GetLogger().Debugf("Received: %v", msg)

			}
		}
	}()
	return nil
}

func (c *gatewayClient) teardownUplink() {
	defer c.mutex.Unlock()
	c.mutex.Lock()
	if c.uplink != nil {
		api.GetLogger().Debugf("Closing uplink stream for %s...", c.id)
		close(c.stopUplink)
		c.uplink.CloseSend()
		c.uplink = nil
	}
}

func (c *gatewayClient) SendUplink(uplink *UplinkMessage) error {
	defer c.mutex.Unlock()
	c.mutex.Lock()
	if c.uplink == nil {
		if err := c.setupUplink(); err != nil {
			return errors.FromGRPCError(err)
		}
	}
	if err := c.uplink.Send(uplink); err != nil {
		if err == io.EOF {
			api.GetLogger().Warnf("Could not send uplink for %s on closed stream", c.id)
			go c.teardownUplink()
			return errors.FromGRPCError(err)
		}
		api.GetLogger().Warnf("Error sending uplink for %s: %s", c.id, err.Error())
		return errors.FromGRPCError(err)
	}
	return nil
}

func (c *gatewayClient) setupDownlink() error {
	api.GetLogger().Debugf("Setting up downlink stream for %s...", c.id)
	ctx, cancel := context.WithCancel(c.getContext())
	downlinkClient, err := c.client.Subscribe(ctx, &SubscribeRequest{})
	if err != nil {
		return err
	}
	c.stopDownlink = make(chan bool)
	go func() {
		<-c.stopDownlink
		cancel()
	}()
	c.downlink = downlinkClient
	return nil
}

func (c *gatewayClient) teardownDownlink() {
	defer c.mutex.Unlock()
	c.mutex.Lock()
	if c.downlink != nil {
		api.GetLogger().Debugf("Closing downlink stream for %s...", c.id)
		close(c.stopDownlink)
		c.downlink.CloseSend()
		c.downlink = nil
	}
}

func (c *gatewayClient) Subscribe() (<-chan *DownlinkMessage, <-chan error, error) {
	defer c.mutex.Unlock()
	c.mutex.Lock()
	if c.downlink == nil {
		if err := c.setupDownlink(); err != nil {
			return nil, nil, errors.FromGRPCError(err)
		}
	}
	downChan := make(chan *DownlinkMessage)
	errChan := make(chan error)
	go func() {
		defer func() {
			close(downChan)
			close(errChan)
		}()
		for {
			select {
			case <-c.stopDownlink:
				return
			default:
				downlink, err := c.downlink.Recv()
				if err != nil {
					if grpc.Code(err) == codes.Canceled {
						api.GetLogger().Debugf("Downlink stream for %s was canceled", c.id)
						errChan <- nil
					} else {
						api.GetLogger().Warnf("Error receiving gateway downlink for %s: %s", c.id, err.Error())
						errChan <- errors.FromGRPCError(err)
					}
					c.teardownDownlink()
					return
				}
				downChan <- downlink
			}
		}
	}()
	return downChan, errChan, nil
}

func (c *gatewayClient) Unsubscribe() error {
	c.teardownDownlink()
	return nil
}

func (c *gatewayClient) Activate(req *DeviceActivationRequest) (*DeviceActivationResponse, error) {
	return c.client.Activate(c.getContext(), req)
}

func (c *gatewayClient) Close() error {
	c.teardownGatewayStatus()
	c.teardownUplink()
	c.teardownDownlink()
	return nil
}
<|endoftext|>


<commit_before>/*
Package hamt64 implements a functional Hash Array Mapped Trie (HAMT).
It is called hamt64 because this package is using 64 nodes for each level of
the Trie. The term functional is used to imply immutable and persistent.

The 60bits of hash are separated into ten 6bit values that constitue the hash
path of any Key in this Trie. However, not all ten levels of the Trie are used.
As many levels (ten or less) are used to find a unique location
for the leaf to be placed within the Trie.

If all ten levels of the Trie are used for two or more key/val pairs then a
special collision leaf will be used to store those key/val pairs,  at the tenth
level of the Trie.
*/
package hamt64

import (
	"fmt"
	"strings"

	"github.com/lleo/go-hamt/key"
)

// Nbits constant is the number of bits(6) a 60bit hash value is split into,
// to provied the indexes of a HAMT.
const Nbits uint = 6

// MaxDepth constant is the maximum depth(9) of Nbits values that constitute
// the path in a HAMT, from [0..MaxDepth]for a total of MaxDepth+1(10) levels.
// Nbits*(MaxDepth+1) == HASHBITS (ie 6*(9+1) == 60).
const MaxDepth uint = 9

// TableCapacity constant is the number of table entries in a each node of
// a HAMT datastructure; its value is 1<<Nbits (ie 2^6 == 64).
const TableCapacity uint = 1 << Nbits

func hashPathMask(depth uint) uint64 {
	return uint64(1<<((depth)*Nbits)) - 1
}

// Create a string of the form "/%02d/%02d..." to describe a hashPath of
// a given depth.
//
// If you want hashPathStrig() to include the current idx, you Must
// add one to depth. You may need to do this because you are creating
// a table to be put at the idx'th slot of the current table.
func hashPathString(hashPath uint64, depth uint) string {
	if depth == 0 {
		return "/"
	}
	var strs = make([]string, depth)

	for d := uint(0); d < depth; d++ {
		var idx = index(hashPath, d)
		strs[d] = fmt.Sprintf("%02d", idx)
	}

	return "/" + strings.Join(strs, "/")
}

func hash60String(h60 uint64) string {
	return hashPathString(h60, MaxDepth)
}

//indexMask() generates a Nbits(6-bit) mask for a given depth
func indexMask(depth uint) uint64 {
	return uint64(uint8(1<<Nbits)-1) << (depth * Nbits)
}

//index() calculates a Nbits(6-bit) integer based on the hash and depth
func index(h60 uint64, depth uint) uint {
	var idxMask = indexMask(depth)
	var idx = uint((h60 & idxMask) >> (depth * Nbits))
	return idx
}

//buildHashPath(hashPath, idx, depth)
func buildHashPath(hashPath uint64, idx, depth uint) uint64 {
	return hashPath | uint64(idx<<(depth*Nbits))
}

type keyVal struct {
	key key.Key
	val interface{}
}

// GradeTables variable controls whether Hamt structures will upgrade/
// downgrade compressed/full tables. This variable and FullTableInit
// should not be changed during the lifetime of any Hamt structure.
// Default: true
var GradeTables = true

// FullTableInit variable controls whether the initial new table type is
// fullTable, else the initial new table type is compressedTable.
// Default: false
var FullTableInit = false

// UpgradeThreshold is a variable that defines when a compressedTable meats
// or exceeds that number of entries, then that table will be upgraded to
// a fullTable. This only applies when HybridTables option is chosen.
// The current value is TableCapacity/2.
var UpgradeThreshold = TableCapacity / 2

// DowngradeThreshold is a variable that defines when a fullTable becomes
// lower than that number of entries, then that table will be downgraded to
// a compressedTable. This only applies when HybridTables option is chosen.
// The current value is TableCapacity/4.
var DowngradeThreshold = TableCapacity / 4

type Hamt struct {
	root     tableI
	nentries uint
}

func (h Hamt) IsEmpty() bool {
	//return h.root == nil
	//return h.nentries == 0
	//return h.root == nil && h.nentries == 0
	return h == Hamt{}
}

func createRootTable(leaf leafI) tableI {
	if FullTableInit {
		return createRootFullTable(leaf)
	}
	return createRootCompressedTable(leaf)
}

func createTable(depth uint, leaf1 leafI, k key.Key, v interface{}) tableI {
	//var hashPath = k.Hash60() & hashPathMask(depth)
	var leaf2 = *newFlatLeaf(k, v)

	if FullTableInit {
		return createFullTable(depth, leaf1, leaf2)
	}
	return createCompressedTable(depth, leaf1, leaf2)
}

// copyUp is ONLY called on a fresh copy of the current Hamt. Hence, modifying
// it is allowed.
func (nh *Hamt) copyUp(oldTable, newTable tableI, path pathT) {
	if path.isEmpty() {
		nh.root = newTable
		return
	}

	var depth = uint(len(path))
	var parentDepth = depth - 1

	var oldParent = path.pop()

	var parentIdx = index(oldTable.Hash60(), parentDepth)
	var newParent tableI
	if newTable == nil {
		newParent = oldParent.remove(parentIdx)
	} else {
		newParent = oldParent.replace(parentIdx, newTable)
	}

	nh.copyUp(oldParent, newParent, path) //recurses at most MaxDepth-1 times

	return
}

// Get(k) retrieves the value for a given key from the Hamt. The bool
// represents whether the key was found.
func (h Hamt) Get(k key.Key) (interface{}, bool) {
	if h.IsEmpty() {
		return nil, false
	}

	var h60 = k.Hash60()

	var curTable = h.root

	for depth := uint(0); depth <= MaxDepth; depth++ {
		var idx = index(h60, depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			break
		}

		if leaf, ok := curNode.(leafI); ok {
			if leaf.Hash60() == h60 {
				return leaf.get(k)
			}
			return nil, false
		}

		//else curNode MUST BE A tableI
		curTable = curNode.(tableI)
	}
	// curNode == nil || depth > MaxDepth

	return nil, false
}

// Put new key/val pair into Hamt, returning a new persistant Hamt and a bool
// indicating if the key/val pair was added(true) or mearly updated(false).
func (h Hamt) Put(k key.Key, v interface{}) (Hamt, bool) {
	var nh = h

	if h.IsEmpty() {
		nh.root = createRootTable(newFlatLeaf(k, v))
		nh.nentries++
		return nh, true
	}

	var newTable tableI
	var added bool

	// for-loop state is path, curTable and depth.
	var path = newPathT()
	var curTable = h.root
	var depth uint

	for depth = 0; depth < MaxDepth; depth++ {
		var idx = index(k.Hash60(), depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			newTable = curTable.insert(idx, newFlatLeaf(k, v))
			added = true
			break
		}

		if curLeaf, isLeaf := curNode.(leafI); isLeaf {
			if curLeaf.Hash60() == k.Hash60() {
				var newLeaf leafI
				newLeaf, added = curLeaf.put(k, v)
				newTable = curTable.replace(idx, newLeaf)
				break
			}

			var tmpTable = createTable(depth+1, curLeaf, k, v)
			newTable = curTable.replace(idx, tmpTable)
			added = true
			break
		}

		path.push(curTable)
		curTable = curNode.(tableI)
	}
	if depth == MaxDepth {
		var idx = index(k.Hash60(), depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			newTable = curTable.insert(idx, newFlatLeaf(k, v))
			added = true
		} else {
			var curLeaf = curNode.(leafI)
			var newLeaf leafI
			newLeaf, added = curLeaf.put(k, v)
			newTable = curTable.replace(idx, newLeaf)
		}
	}

	if added {
		nh.nentries++
	}
	nh.copyUp(curTable, newTable, path)

	return nh, added
}

// Del returns a new Hamt, the value deleted, and a boolean that
// specifies whether or not the key was deleted (eg it didn't exist to start
// with). Therefor you must always test deleted before using the new *Hamt
// value.
func (h Hamt) Del(k key.Key) (Hamt, interface{}, bool) {
	var nh = h

	var val interface{}
	var deleted bool
	var newTable tableI

	// for-loop state is path, curTable, and depth.
	var path = newPathT()
	var curTable = h.root
	var depth uint
	for depth = 0; depth < MaxDepth; depth++ {
		var idx = index(k.Hash60(), depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			//return h, nil, false
			return nh, val, deleted
		}

		if curLeaf, ok := curNode.(leafI); ok {
			var newLeaf leafI
			newLeaf, val, deleted = curLeaf.del(k)

			if !deleted {
				//return h, nil, false
				return nh, val, deleted
			}

			if newLeaf == nil {
				newTable = curTable.remove(idx)
			} else {
				newTable = curTable.replace(idx, newLeaf)
			}

			break
		}

		// curNode is NOT nil & NOT a leafI, so curNode MUST BE a tableI.
		// We are going to loop, so update loop state like path and curTable.
		// for-loop will handle updating depth.

		path.push(curTable)
		curTable = curNode.(tableI)
	}
	if depth == MaxDepth {
		var idx = index(k.Hash60(), depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			return nh, val, deleted
		}

		var curLeaf = curNode.(leafI)

		var newLeaf leafI
		newLeaf, val, deleted = curLeaf.del(k)

		if !deleted {
			return nh, val, deleted
		}

		if newLeaf == nil {
			newTable = curTable.remove(idx)
		} else {
			newTable = curTable.replace(idx, newLeaf)
		}
	}

	nh.nentries--
	nh.copyUp(curTable, newTable, path)

	return nh, val, deleted
}

func (h Hamt) String() string {
	return fmt.Sprintf("Hamt{ nentries: %d, root: %s }", h.nentries, h.root)
}

func (h Hamt) LongString(indent string) string {
	var str string
	if h.root != nil {
		str = indent + fmt.Sprintf("Hamt{ nentries: %d, root:\n", h.nentries)
		str += indent + h.root.LongString(indent, 0)
		str += indent + "}"
		return str
	} else {
		str = indent + fmt.Sprintf("Hamt{ nentries: %d, root: nil }", h.nentries)
	}
	return str
}
<commit_msg>UpgradeThreshold const changed from 1/2 to 2/3 TableCapacity<commit_after>/*
Package hamt64 implements a functional Hash Array Mapped Trie (HAMT).
It is called hamt64 because this package is using 64 nodes for each level of
the Trie. The term functional is used to imply immutable and persistent.

The 60bits of hash are separated into ten 6bit values that constitue the hash
path of any Key in this Trie. However, not all ten levels of the Trie are used.
As many levels (ten or less) are used to find a unique location
for the leaf to be placed within the Trie.

If all ten levels of the Trie are used for two or more key/val pairs then a
special collision leaf will be used to store those key/val pairs,  at the tenth
level of the Trie.
*/
package hamt64

import (
	"fmt"
	"strings"

	"github.com/lleo/go-hamt/key"
)

// Nbits constant is the number of bits(6) a 60bit hash value is split into,
// to provied the indexes of a HAMT.
const Nbits uint = 6

// MaxDepth constant is the maximum depth(9) of Nbits values that constitute
// the path in a HAMT, from [0..MaxDepth]for a total of MaxDepth+1(10) levels.
// Nbits*(MaxDepth+1) == HASHBITS (ie 6*(9+1) == 60).
const MaxDepth uint = 9

// TableCapacity constant is the number of table entries in a each node of
// a HAMT datastructure; its value is 1<<Nbits (ie 2^6 == 64).
const TableCapacity uint = 1 << Nbits

func hashPathMask(depth uint) uint64 {
	return uint64(1<<((depth)*Nbits)) - 1
}

// Create a string of the form "/%02d/%02d..." to describe a hashPath of
// a given depth.
//
// If you want hashPathStrig() to include the current idx, you Must
// add one to depth. You may need to do this because you are creating
// a table to be put at the idx'th slot of the current table.
func hashPathString(hashPath uint64, depth uint) string {
	if depth == 0 {
		return "/"
	}
	var strs = make([]string, depth)

	for d := uint(0); d < depth; d++ {
		var idx = index(hashPath, d)
		strs[d] = fmt.Sprintf("%02d", idx)
	}

	return "/" + strings.Join(strs, "/")
}

func hash60String(h60 uint64) string {
	return hashPathString(h60, MaxDepth)
}

//indexMask() generates a Nbits(6-bit) mask for a given depth
func indexMask(depth uint) uint64 {
	return uint64(uint8(1<<Nbits)-1) << (depth * Nbits)
}

//index() calculates a Nbits(6-bit) integer based on the hash and depth
func index(h60 uint64, depth uint) uint {
	var idxMask = indexMask(depth)
	var idx = uint((h60 & idxMask) >> (depth * Nbits))
	return idx
}

//buildHashPath(hashPath, idx, depth)
func buildHashPath(hashPath uint64, idx, depth uint) uint64 {
	return hashPath | uint64(idx<<(depth*Nbits))
}

type keyVal struct {
	key key.Key
	val interface{}
}

// GradeTables variable controls whether Hamt structures will upgrade/
// downgrade compressed/full tables. This variable and FullTableInit
// should not be changed during the lifetime of any Hamt structure.
// Default: true
var GradeTables = true

// FullTableInit variable controls whether the initial new table type is
// fullTable, else the initial new table type is compressedTable.
// Default: false
var FullTableInit = false

// UpgradeThreshold is a variable that defines when a compressedTable meats
// or exceeds that number of entries, then that table will be upgraded to
// a fullTable. This only applies when HybridTables option is chosen.
// The current value is TableCapacity/2.
var UpgradeThreshold = TableCapacity * 2 / 3

// DowngradeThreshold is a variable that defines when a fullTable becomes
// lower than that number of entries, then that table will be downgraded to
// a compressedTable. This only applies when HybridTables option is chosen.
// The current value is TableCapacity/4.
var DowngradeThreshold = TableCapacity / 4

type Hamt struct {
	root     tableI
	nentries uint
}

func (h Hamt) IsEmpty() bool {
	//return h.root == nil
	//return h.nentries == 0
	//return h.root == nil && h.nentries == 0
	return h == Hamt{}
}

func createRootTable(leaf leafI) tableI {
	if FullTableInit {
		return createRootFullTable(leaf)
	}
	return createRootCompressedTable(leaf)
}

func createTable(depth uint, leaf1 leafI, k key.Key, v interface{}) tableI {
	//var hashPath = k.Hash60() & hashPathMask(depth)
	var leaf2 = *newFlatLeaf(k, v)

	if FullTableInit {
		return createFullTable(depth, leaf1, leaf2)
	}
	return createCompressedTable(depth, leaf1, leaf2)
}

// copyUp is ONLY called on a fresh copy of the current Hamt. Hence, modifying
// it is allowed.
func (nh *Hamt) copyUp(oldTable, newTable tableI, path pathT) {
	if path.isEmpty() {
		nh.root = newTable
		return
	}

	var depth = uint(len(path))
	var parentDepth = depth - 1

	var oldParent = path.pop()

	var parentIdx = index(oldTable.Hash60(), parentDepth)
	var newParent tableI
	if newTable == nil {
		newParent = oldParent.remove(parentIdx)
	} else {
		newParent = oldParent.replace(parentIdx, newTable)
	}

	nh.copyUp(oldParent, newParent, path) //recurses at most MaxDepth-1 times

	return
}

// Get(k) retrieves the value for a given key from the Hamt. The bool
// represents whether the key was found.
func (h Hamt) Get(k key.Key) (interface{}, bool) {
	if h.IsEmpty() {
		return nil, false
	}

	var h60 = k.Hash60()

	var curTable = h.root

	for depth := uint(0); depth <= MaxDepth; depth++ {
		var idx = index(h60, depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			break
		}

		if leaf, ok := curNode.(leafI); ok {
			if leaf.Hash60() == h60 {
				return leaf.get(k)
			}
			return nil, false
		}

		//else curNode MUST BE A tableI
		curTable = curNode.(tableI)
	}
	// curNode == nil || depth > MaxDepth

	return nil, false
}

// Put new key/val pair into Hamt, returning a new persistant Hamt and a bool
// indicating if the key/val pair was added(true) or mearly updated(false).
func (h Hamt) Put(k key.Key, v interface{}) (Hamt, bool) {
	var nh = h

	if h.IsEmpty() {
		nh.root = createRootTable(newFlatLeaf(k, v))
		nh.nentries++
		return nh, true
	}

	var newTable tableI
	var added bool

	// for-loop state is path, curTable and depth.
	var path = newPathT()
	var curTable = h.root
	var depth uint

	for depth = 0; depth < MaxDepth; depth++ {
		var idx = index(k.Hash60(), depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			newTable = curTable.insert(idx, newFlatLeaf(k, v))
			added = true
			break
		}

		if curLeaf, isLeaf := curNode.(leafI); isLeaf {
			if curLeaf.Hash60() == k.Hash60() {
				var newLeaf leafI
				newLeaf, added = curLeaf.put(k, v)
				newTable = curTable.replace(idx, newLeaf)
				break
			}

			var tmpTable = createTable(depth+1, curLeaf, k, v)
			newTable = curTable.replace(idx, tmpTable)
			added = true
			break
		}

		path.push(curTable)
		curTable = curNode.(tableI)
	}
	if depth == MaxDepth {
		var idx = index(k.Hash60(), depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			newTable = curTable.insert(idx, newFlatLeaf(k, v))
			added = true
		} else {
			var curLeaf = curNode.(leafI)
			var newLeaf leafI
			newLeaf, added = curLeaf.put(k, v)
			newTable = curTable.replace(idx, newLeaf)
		}
	}

	if added {
		nh.nentries++
	}
	nh.copyUp(curTable, newTable, path)

	return nh, added
}

// Del returns a new Hamt, the value deleted, and a boolean that
// specifies whether or not the key was deleted (eg it didn't exist to start
// with). Therefor you must always test deleted before using the new *Hamt
// value.
func (h Hamt) Del(k key.Key) (Hamt, interface{}, bool) {
	var nh = h

	var val interface{}
	var deleted bool
	var newTable tableI

	// for-loop state is path, curTable, and depth.
	var path = newPathT()
	var curTable = h.root
	var depth uint
	for depth = 0; depth < MaxDepth; depth++ {
		var idx = index(k.Hash60(), depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			//return h, nil, false
			return nh, val, deleted
		}

		if curLeaf, ok := curNode.(leafI); ok {
			var newLeaf leafI
			newLeaf, val, deleted = curLeaf.del(k)

			if !deleted {
				//return h, nil, false
				return nh, val, deleted
			}

			if newLeaf == nil {
				newTable = curTable.remove(idx)
			} else {
				newTable = curTable.replace(idx, newLeaf)
			}

			break
		}

		// curNode is NOT nil & NOT a leafI, so curNode MUST BE a tableI.
		// We are going to loop, so update loop state like path and curTable.
		// for-loop will handle updating depth.

		path.push(curTable)
		curTable = curNode.(tableI)
	}
	if depth == MaxDepth {
		var idx = index(k.Hash60(), depth)
		var curNode = curTable.get(idx)

		if curNode == nil {
			return nh, val, deleted
		}

		var curLeaf = curNode.(leafI)

		var newLeaf leafI
		newLeaf, val, deleted = curLeaf.del(k)

		if !deleted {
			return nh, val, deleted
		}

		if newLeaf == nil {
			newTable = curTable.remove(idx)
		} else {
			newTable = curTable.replace(idx, newLeaf)
		}
	}

	nh.nentries--
	nh.copyUp(curTable, newTable, path)

	return nh, val, deleted
}

func (h Hamt) String() string {
	return fmt.Sprintf("Hamt{ nentries: %d, root: %s }", h.nentries, h.root)
}

func (h Hamt) LongString(indent string) string {
	var str string
	if h.root != nil {
		str = indent + fmt.Sprintf("Hamt{ nentries: %d, root:\n", h.nentries)
		str += indent + h.root.LongString(indent, 0)
		str += indent + "}"
		return str
	} else {
		str = indent + fmt.Sprintf("Hamt{ nentries: %d, root: nil }", h.nentries)
	}
	return str
}
<|endoftext|>


<commit_before>package storage

import (
	"strconv"
	"sync"
)

type NeedleValue struct {
	Key    Key
	Offset uint32 `comment:"Volume offset"` //since aligned to 8 bytes, range is 4G*8=32G
	Size   uint32 `comment:"Size of the data portion"`
}

const (
	batch = 100000
)

type Key uint64

func (k Key) String() string {
	return strconv.FormatUint(uint64(k), 10)
}

type CompactSection struct {
	sync.RWMutex
	values   []NeedleValue
	overflow map[Key]NeedleValue
	start    Key
	end      Key
	counter  int
}

func NewCompactSection(start Key) *CompactSection {
	return &CompactSection{
		values:   make([]NeedleValue, batch),
		overflow: make(map[Key]NeedleValue),
		start:    start,
	}
}

//return old entry size
func (cs *CompactSection) Set(key Key, offset uint32, size uint32) uint32 {
	ret := uint32(0)
	if key > cs.end {
		cs.end = key
	}
	cs.Lock()
	if i := cs.binarySearchValues(key); i >= 0 {
		ret = cs.values[i].Size
		//println("key", key, "old size", ret)
		cs.values[i].Offset, cs.values[i].Size = offset, size
	} else {
		needOverflow := cs.counter >= batch
		needOverflow = needOverflow || cs.counter > 0 && cs.values[cs.counter-1].Key > key
		if needOverflow {
			//println("start", cs.start, "counter", cs.counter, "key", key)
			if oldValue, found := cs.overflow[key]; found {
				ret = oldValue.Size
			}
			cs.overflow[key] = NeedleValue{Key: key, Offset: offset, Size: size}
		} else {
			p := &cs.values[cs.counter]
			p.Key, p.Offset, p.Size = key, offset, size
			//println("added index", cs.counter, "key", key, cs.values[cs.counter].Key)
			cs.counter++
		}
	}
	cs.Unlock()
	return ret
}

//return old entry size
func (cs *CompactSection) Delete(key Key) uint32 {
	cs.Lock()
	ret := uint32(0)
	if i := cs.binarySearchValues(key); i >= 0 {
		if cs.values[i].Size > 0 {
			ret = cs.values[i].Size
			cs.values[i].Size = 0
		}
	}
	if v, found := cs.overflow[key]; found {
		delete(cs.overflow, key)
		ret = v.Size
	}
	cs.Unlock()
	return ret
}
func (cs *CompactSection) Get(key Key) (*NeedleValue, bool) {
	cs.RLock()
	if v, ok := cs.overflow[key]; ok {
		cs.RUnlock()
		return &v, true
	}
	if i := cs.binarySearchValues(key); i >= 0 {
		cs.RUnlock()
		return &cs.values[i], true
	}
	cs.RUnlock()
	return nil, false
}
func (cs *CompactSection) binarySearchValues(key Key) int {
	l, h := 0, cs.counter-1
	if h >= 0 && cs.values[h].Key < key {
		return -2
	}
	//println("looking for key", key)
	for l <= h {
		m := (l + h) / 2
		//println("mid", m, "key", cs.values[m].Key, cs.values[m].Offset, cs.values[m].Size)
		if cs.values[m].Key < key {
			l = m + 1
		} else if key < cs.values[m].Key {
			h = m - 1
		} else {
			//println("found", m)
			return m
		}
	}
	return -1
}

//This map assumes mostly inserting increasing keys
type CompactMap struct {
	list []*CompactSection
}

func NewCompactMap() CompactMap {
	return CompactMap{}
}

func (cm *CompactMap) Set(key Key, offset uint32, size uint32) uint32 {
	x := cm.binarySearchCompactSection(key)
	if x < 0 {
		//println(x, "creating", len(cm.list), "section, starting", key)
		cm.list = append(cm.list, NewCompactSection(key))
		x = len(cm.list) - 1
		//keep compact section sorted by start
		for x > 0 {
			if cm.list[x-1].start > cm.list[x].start {
				cm.list[x-1], cm.list[x] = cm.list[x], cm.list[x-1]
				x = x - 1
			} else {
				break
			}
		}
	}
	return cm.list[x].Set(key, offset, size)
}
func (cm *CompactMap) Delete(key Key) uint32 {
	x := cm.binarySearchCompactSection(key)
	if x < 0 {
		return uint32(0)
	}
	return cm.list[x].Delete(key)
}
func (cm *CompactMap) Get(key Key) (*NeedleValue, bool) {
	x := cm.binarySearchCompactSection(key)
	if x < 0 {
		return nil, false
	}
	return cm.list[x].Get(key)
}
func (cm *CompactMap) binarySearchCompactSection(key Key) int {
	l, h := 0, len(cm.list)-1
	if h < 0 {
		return -5
	}
	if cm.list[h].start <= key {
		if cm.list[h].counter < batch || key <= cm.list[h].end {
			return h
		}
		return -4
	}
	for l <= h {
		m := (l + h) / 2
		if key < cm.list[m].start {
			h = m - 1
		} else { // cm.list[m].start <= key
			if cm.list[m+1].start <= key {
				l = m + 1
			} else {
				return m
			}
		}
	}
	return -3
}

// Visit visits all entries or stop if any error when visiting
func (cm *CompactMap) Visit(visit func(NeedleValue) error) error {
	for _, cs := range cm.list {
		cs.RLock()
		for _, v := range cs.overflow {
			if err := visit(v); err != nil {
				cs.RUnlock()
				return err
			}
		}
		for _, v := range cs.values {
			if _, found := cs.overflow[v.Key]; !found {
				if err := visit(v); err != nil {
					cs.RUnlock()
					return err
				}
			}
		}
		cs.RUnlock()
	}
	return nil
}
<commit_msg>lock<commit_after>package storage

import (
	"strconv"
	"sync"
)

type NeedleValue struct {
	Key    Key
	Offset uint32 `comment:"Volume offset"` //since aligned to 8 bytes, range is 4G*8=32G
	Size   uint32 `comment:"Size of the data portion"`
}

const (
	batch = 100000
)

type Key uint64

func (k Key) String() string {
	return strconv.FormatUint(uint64(k), 10)
}

type CompactSection struct {
	sync.RWMutex
	values   []NeedleValue
	overflow map[Key]NeedleValue
	start    Key
	end      Key
	counter  int
}

func NewCompactSection(start Key) *CompactSection {
	return &CompactSection{
		values:   make([]NeedleValue, batch),
		overflow: make(map[Key]NeedleValue),
		start:    start,
	}
}

//return old entry size
func (cs *CompactSection) Set(key Key, offset uint32, size uint32) uint32 {
	ret := uint32(0)
	cs.Lock()
	if key > cs.end {
		cs.end = key
	}
	if i := cs.binarySearchValues(key); i >= 0 {
		ret = cs.values[i].Size
		//println("key", key, "old size", ret)
		cs.values[i].Offset, cs.values[i].Size = offset, size
	} else {
		needOverflow := cs.counter >= batch
		needOverflow = needOverflow || cs.counter > 0 && cs.values[cs.counter-1].Key > key
		if needOverflow {
			//println("start", cs.start, "counter", cs.counter, "key", key)
			if oldValue, found := cs.overflow[key]; found {
				ret = oldValue.Size
			}
			cs.overflow[key] = NeedleValue{Key: key, Offset: offset, Size: size}
		} else {
			p := &cs.values[cs.counter]
			p.Key, p.Offset, p.Size = key, offset, size
			//println("added index", cs.counter, "key", key, cs.values[cs.counter].Key)
			cs.counter++
		}
	}
	cs.Unlock()
	return ret
}

//return old entry size
func (cs *CompactSection) Delete(key Key) uint32 {
	cs.Lock()
	ret := uint32(0)
	if i := cs.binarySearchValues(key); i >= 0 {
		if cs.values[i].Size > 0 {
			ret = cs.values[i].Size
			cs.values[i].Size = 0
		}
	}
	if v, found := cs.overflow[key]; found {
		delete(cs.overflow, key)
		ret = v.Size
	}
	cs.Unlock()
	return ret
}
func (cs *CompactSection) Get(key Key) (*NeedleValue, bool) {
	cs.RLock()
	if v, ok := cs.overflow[key]; ok {
		cs.RUnlock()
		return &v, true
	}
	if i := cs.binarySearchValues(key); i >= 0 {
		cs.RUnlock()
		return &cs.values[i], true
	}
	cs.RUnlock()
	return nil, false
}
func (cs *CompactSection) binarySearchValues(key Key) int {
	l, h := 0, cs.counter-1
	if h >= 0 && cs.values[h].Key < key {
		return -2
	}
	//println("looking for key", key)
	for l <= h {
		m := (l + h) / 2
		//println("mid", m, "key", cs.values[m].Key, cs.values[m].Offset, cs.values[m].Size)
		if cs.values[m].Key < key {
			l = m + 1
		} else if key < cs.values[m].Key {
			h = m - 1
		} else {
			//println("found", m)
			return m
		}
	}
	return -1
}

//This map assumes mostly inserting increasing keys
type CompactMap struct {
	list []*CompactSection
}

func NewCompactMap() CompactMap {
	return CompactMap{}
}

func (cm *CompactMap) Set(key Key, offset uint32, size uint32) uint32 {
	x := cm.binarySearchCompactSection(key)
	if x < 0 {
		//println(x, "creating", len(cm.list), "section, starting", key)
		cm.list = append(cm.list, NewCompactSection(key))
		x = len(cm.list) - 1
		//keep compact section sorted by start
		for x > 0 {
			if cm.list[x-1].start > cm.list[x].start {
				cm.list[x-1], cm.list[x] = cm.list[x], cm.list[x-1]
				x = x - 1
			} else {
				break
			}
		}
	}
	return cm.list[x].Set(key, offset, size)
}
func (cm *CompactMap) Delete(key Key) uint32 {
	x := cm.binarySearchCompactSection(key)
	if x < 0 {
		return uint32(0)
	}
	return cm.list[x].Delete(key)
}
func (cm *CompactMap) Get(key Key) (*NeedleValue, bool) {
	x := cm.binarySearchCompactSection(key)
	if x < 0 {
		return nil, false
	}
	return cm.list[x].Get(key)
}
func (cm *CompactMap) binarySearchCompactSection(key Key) int {
	l, h := 0, len(cm.list)-1
	if h < 0 {
		return -5
	}
	if cm.list[h].start <= key {
		if cm.list[h].counter < batch || key <= cm.list[h].end {
			return h
		}
		return -4
	}
	for l <= h {
		m := (l + h) / 2
		if key < cm.list[m].start {
			h = m - 1
		} else { // cm.list[m].start <= key
			if cm.list[m+1].start <= key {
				l = m + 1
			} else {
				return m
			}
		}
	}
	return -3
}

// Visit visits all entries or stop if any error when visiting
func (cm *CompactMap) Visit(visit func(NeedleValue) error) error {
	for _, cs := range cm.list {
		cs.RLock()
		for _, v := range cs.overflow {
			if err := visit(v); err != nil {
				cs.RUnlock()
				return err
			}
		}
		for _, v := range cs.values {
			if _, found := cs.overflow[v.Key]; !found {
				if err := visit(v); err != nil {
					cs.RUnlock()
					return err
				}
			}
		}
		cs.RUnlock()
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"log"
	"net"
	"os"
)

import "bufio"

func main() {
	reader := bufio.NewReader(os.Stdin)

	addr, err := net.ResolveUDPAddr("udp", "localhost:2003")
	if err != nil {
		log.Fatal(err)
	}

	conn, err := net.Dial("udp", addr.String())
	if err != nil {
		log.Fatal(err)
	}

	for {
		line, err := reader.ReadString('\n')

		if err != nil {
			if err == io.EOF {
				break
			}
			log.Fatal(err)
			break
		}

		if _, err := conn.Write([]byte(line)); err != nil {
			log.Fatal(err)
		}

	}
}
<commit_msg>Remove old unused directory<commit_after><|endoftext|>


<commit_before>package globalnamespacerbac

import (
	"fmt"
	"k8s.io/apimachinery/pkg/types"
	"reflect"
	"sort"
	"strings"

	"github.com/rancher/rancher/pkg/namespace"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	rbacv1 "github.com/rancher/types/apis/rbac.authorization.k8s.io/v1"
	"github.com/rancher/types/config"

	k8srbacv1 "k8s.io/api/rbac/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
)

const (
	allAccess                 = "all"
	updateAccess              = "update"
	readOnlyAccess            = "readonly"
	MultiClusterAppResource   = "multiclusterapps"
	GlobalDNSResource         = "globaldnses"
	GlobalDNSProviderResource = "globaldnsproviders"
	CloudCredentialResource   = "secrets"
	CreatorIDAnn              = "field.cattle.io/creatorId"
)

func CreateRoleAndRoleBinding(resource string, name string, UID types.UID, members []v3.Member, creatorID string,
	managementContext *config.ManagementContext, apiGroup ...string) error {
	/* Create 3 Roles containing the resource (multiclusterapp or globalDNS), and the current multiclusterapp/globalDNS in resourceNames list
	1. Role with all verbs (everything access, includes creator); name multiclusterapp.Name + "-ma" / (globalDNS.Name + "-ga")
	2. Role with "get", "list" "watch" verbs (ReadOnly access); name multiclusterapp.Name + "-mr" / (globalDNS.Name + "-gr")
	3. Role with "update" verb (Upgrade access); name multiclusterapp.Name + "-mu" / (globalDNS.Name + "-gu")
	*/
	api := []string{"management.cattle.io"}
	if len(apiGroup) > 0 {
		api = apiGroup
	}
	if _, err := createRole(resource, allAccess, name, managementContext, api); err != nil {
		return err
	}

	// Create a roleBinding referring the role with everything access, and containing creator of the multiclusterapp, along with
	// any members that have everything access
	var allAccessSubjects, readOnlyAccessSubjects, updateAccessSubjects []k8srbacv1.Subject
	allAccessSubjects = append(allAccessSubjects, k8srbacv1.Subject{Kind: "User", Name: creatorID, APIGroup: rbacv1.GroupName})
	for _, m := range members {
		s, err := buildSubjectForMember(m, managementContext)
		if err != nil {
			return err
		}
		switch m.AccessType {
		case allAccess:
			allAccessSubjects = append(allAccessSubjects, s)
		case updateAccess:
			updateAccessSubjects = append(updateAccessSubjects, s)
		case readOnlyAccess:
			readOnlyAccessSubjects = append(readOnlyAccessSubjects, s)
		default:
			readOnlyAccessSubjects = append(readOnlyAccessSubjects, s)
		}
	}

	if _, err := createRoleBinding(allAccess, name, UID, allAccessSubjects, managementContext, resource, api); err != nil {
		return err
	}

	// Check if there are members with readonly or update access; if found then create rolebindings for those
	if len(readOnlyAccessSubjects) > 0 {
		if _, err := createRole(resource, readOnlyAccess, name, managementContext, api); err != nil {
			return err
		}
		if _, err := createRoleBinding(readOnlyAccess, name, UID, readOnlyAccessSubjects, managementContext, resource, api); err != nil {
			return err
		}
	}
	if len(updateAccessSubjects) > 0 {
		if _, err := createRole(resource, updateAccess, name, managementContext, api); err != nil {
			return err
		}
		if _, err := createRoleBinding(updateAccess, name, UID, updateAccessSubjects, managementContext, resource, api); err != nil {
			return err
		}
	}
	return nil
}

func createRole(resource string, roleAccess string, resourceName string,
	managementContext *config.ManagementContext, apiGroups []string) (*k8srbacv1.Role, error) {
	roleName, verbs := getRoleNameAndVerbs(roleAccess, resourceName, resource)
	newRole := &k8srbacv1.Role{
		ObjectMeta: metav1.ObjectMeta{
			Name:      roleName,
			Namespace: namespace.GlobalNamespace,
		},
		Rules: []k8srbacv1.PolicyRule{
			{
				APIGroups:     apiGroups,
				Resources:     []string{resource},
				ResourceNames: []string{resourceName},
				Verbs:         verbs,
			},
		},
	}
	role, err := managementContext.RBAC.Roles("").GetNamespaced(namespace.GlobalNamespace, roleName, metav1.GetOptions{})
	if err != nil {
		if apierrors.IsNotFound(err) {
			role, err = managementContext.RBAC.Roles("").Create(newRole)
			if err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	} else if role != nil {
		if !reflect.DeepEqual(newRole, role) {
			toUpdate := newRole.DeepCopy()
			updated, err := managementContext.RBAC.Roles("").Update(toUpdate)
			if err != nil {
				return updated, err
			}
		}
	}
	return role, nil
}

func createRoleBinding(roleAccess string, name string, UID types.UID,
	subjects []k8srbacv1.Subject, managementContext *config.ManagementContext, resource string, apiGroups []string) (*k8srbacv1.RoleBinding, error) {
	roleName, _ := getRoleNameAndVerbs(roleAccess, name, resource)
	// we can define the rolebinding first, since if it's not already present we can call create. And if it's present then we'll
	// still need to compare the current members' list
	sort.Slice(subjects, func(i, j int) bool { return subjects[i].Name < subjects[j].Name })

	apiVersion := "management.cattle.io/v3"
	if apiGroups[0] != "management.cattle.io" {
		apiVersion = "v1"
	}
	ownerReference := metav1.OwnerReference{
		APIVersion: apiVersion,
		Kind:       resource,
		Name:       name,
		UID:        UID,
	}

	newRoleBinding := &k8srbacv1.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Name:            roleName,
			Namespace:       namespace.GlobalNamespace,
			OwnerReferences: []metav1.OwnerReference{ownerReference},
		},
		RoleRef: k8srbacv1.RoleRef{
			Name: roleName,
			Kind: "Role",
		},
		Subjects: subjects,
	}
	roleBinding, err := managementContext.RBAC.RoleBindings("").GetNamespaced(namespace.GlobalNamespace, roleName, metav1.GetOptions{})
	if err != nil {
		if apierrors.IsNotFound(err) {
			_, err = managementContext.RBAC.RoleBindings("").Create(newRoleBinding)
			if err != nil && !apierrors.IsAlreadyExists(err) {
				return nil, err
			}
		} else {
			return nil, err
		}
	} else if roleBinding != nil {
		if !reflect.DeepEqual(roleBinding, newRoleBinding) {
			toUpdate := newRoleBinding.DeepCopy()
			updated, err := managementContext.RBAC.RoleBindings("").Update(toUpdate)
			if err != nil {
				return updated, err
			}
		}
	}
	return newRoleBinding, nil
}

func getRoleNameAndVerbs(roleAccess string, resourceName string, resource string) (string, []string) {
	var roleName string
	var verbs []string
	switch resource {
	case MultiClusterAppResource:
		resourceName += "-m"
	case GlobalDNSResource:
		resourceName += "-g"
	}
	switch roleAccess {
	case allAccess:
		roleName = resourceName + "a"
		verbs = []string{"*"}
	case updateAccess:
		roleName = resourceName + "u"
		verbs = []string{"update", "get", "list", "watch"}
	case readOnlyAccess:
		roleName = resourceName + "r"
		verbs = []string{"get", "list", "watch"}
	}
	return roleName, verbs
}

func buildSubjectForMember(member v3.Member, managementContext *config.ManagementContext) (k8srbacv1.Subject, error) {
	var name, kind string
	member, err := checkAndSetUserFields(member, managementContext)
	if err != nil {
		return k8srbacv1.Subject{}, err
	}

	if member.UserName != "" {
		name = member.UserName
		kind = "User"
	}

	if member.GroupPrincipalName != "" {
		if name != "" {
			return k8srbacv1.Subject{}, fmt.Errorf("member %v has both username and groupPrincipalId set", member)
		}
		name = member.GroupPrincipalName
		kind = "Group"
	}

	if name == "" {
		return k8srbacv1.Subject{}, fmt.Errorf("member %v doesn't have any name fields set", member)
	}

	return k8srbacv1.Subject{
		Kind:     kind,
		Name:     name,
		APIGroup: rbacv1.GroupName,
	}, nil
}

func checkAndSetUserFields(m v3.Member, managementContext *config.ManagementContext) (v3.Member, error) {
	if m.GroupPrincipalName != "" || (m.UserPrincipalName != "" && m.UserName != "") {
		return m, nil
	}
	if m.UserPrincipalName != "" && m.UserName == "" {
		displayName := m.DisplayName
		user, err := managementContext.UserManager.EnsureUser(m.UserPrincipalName, displayName)
		if err != nil {
			return m, err
		}

		m.UserName = user.Name
		return m, nil
	}
	return m, nil
}

func GetMemberGroups(members []v3.Member) map[string]string {
	groups := make(map[string]string)
	for _, m := range members {
		if m.GroupPrincipalName != "" {
			groups[m.GroupPrincipalName] = m.AccessType
		}
	}
	return groups
}

func GetUpdatedMembers(projects []string, members []v3.Member, prtbLister v3.ProjectRoleTemplateBindingLister) ([]v3.Member, error) {
	var updatedUsers, updatedGroups, updatedMembers []v3.Member
	users := make(map[string]bool)
	groups := make(map[string]bool)
	for _, p := range projects {
		split := strings.SplitN(p, ":", 2)
		if len(split) != 2 {
			return updatedMembers, fmt.Errorf("Target project name %s is invalid", p)
		}
		projectNS := split[1]
		prtbs, err := prtbLister.List(projectNS, labels.NewSelector())
		if err != nil {
			return updatedMembers, err
		}
		prtbMap := make(map[string]bool)
		// get all PRTBs for this project
		for _, prtb := range prtbs {
			if prtb.UserPrincipalName != "" && prtb.DeletionTimestamp == nil {
				prtbMap[prtb.UserPrincipalName] = true
			} else if prtb.GroupPrincipalName != "" && prtb.DeletionTimestamp == nil {
				prtbMap[prtb.GroupPrincipalName] = true
			}
		}

		// all members should have PRTBs in this project, if not, remove them
		for _, mem := range members {
			if mem.UserPrincipalName != "" && prtbMap[mem.UserPrincipalName] && !users[mem.UserPrincipalName] {
				users[mem.UserPrincipalName] = true
				updatedUsers = append(updatedUsers, mem)
			} else if mem.GroupPrincipalName != "" && prtbMap[mem.GroupPrincipalName] && !groups[mem.GroupPrincipalName] {
				groups[mem.GroupPrincipalName] = true
				updatedGroups = append(updatedGroups, mem)
			}
		}
	}
	sort.Slice(updatedUsers, func(i, j int) bool { return updatedUsers[i].UserPrincipalName < updatedUsers[j].UserPrincipalName })
	sort.Slice(updatedGroups, func(i, j int) bool { return updatedGroups[i].GroupPrincipalName < updatedGroups[j].GroupPrincipalName })
	updatedMembers = append(updatedUsers, updatedGroups...)
	return updatedMembers, nil
}

func GetCurrentMembers(members []v3.Member) []v3.Member {
	var memberUsers, memberGroups []v3.Member
	for _, m := range members {
		if m.UserPrincipalName != "" {
			memberUsers = append(memberUsers, m)
		} else if m.GroupPrincipalName != "" {
			memberGroups = append(memberGroups, m)
		}
	}
	sort.Slice(memberUsers, func(i, j int) bool { return memberUsers[i].UserPrincipalName < memberUsers[j].UserPrincipalName })
	sort.Slice(memberGroups, func(i, j int) bool { return memberGroups[i].GroupPrincipalName < memberGroups[j].GroupPrincipalName })
	memberUsers = append(memberUsers, memberGroups...)
	return memberUsers
}
<commit_msg>clean roles for global namespace resources on deletion<commit_after>package globalnamespacerbac

import (
	"fmt"
	"k8s.io/apimachinery/pkg/types"
	"reflect"
	"sort"
	"strings"

	"github.com/rancher/rancher/pkg/namespace"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	rbacv1 "github.com/rancher/types/apis/rbac.authorization.k8s.io/v1"
	"github.com/rancher/types/config"

	k8srbacv1 "k8s.io/api/rbac/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
)

const (
	allAccess                 = "all"
	updateAccess              = "update"
	readOnlyAccess            = "readonly"
	MultiClusterAppResource   = "multiclusterapps"
	GlobalDNSResource         = "globaldnses"
	GlobalDNSProviderResource = "globaldnsproviders"
	CloudCredentialResource   = "secrets"
	CreatorIDAnn              = "field.cattle.io/creatorId"
)

func CreateRoleAndRoleBinding(resource string, name string, UID types.UID, members []v3.Member, creatorID string,
	managementContext *config.ManagementContext, apiGroup ...string) error {
	/* Create 3 Roles containing the resource (multiclusterapp or globalDNS), and the current multiclusterapp/globalDNS in resourceNames list
	1. Role with all verbs (everything access, includes creator); name multiclusterapp.Name + "-ma" / (globalDNS.Name + "-ga")
	2. Role with "get", "list" "watch" verbs (ReadOnly access); name multiclusterapp.Name + "-mr" / (globalDNS.Name + "-gr")
	3. Role with "update" verb (Upgrade access); name multiclusterapp.Name + "-mu" / (globalDNS.Name + "-gu")
	*/
	api := []string{"management.cattle.io"}
	if len(apiGroup) > 0 {
		api = apiGroup
	}
	if _, err := createRole(resource, allAccess, name, UID, managementContext, api); err != nil {
		return err
	}

	// Create a roleBinding referring the role with everything access, and containing creator of the multiclusterapp, along with
	// any members that have everything access
	var allAccessSubjects, readOnlyAccessSubjects, updateAccessSubjects []k8srbacv1.Subject
	allAccessSubjects = append(allAccessSubjects, k8srbacv1.Subject{Kind: "User", Name: creatorID, APIGroup: rbacv1.GroupName})
	for _, m := range members {
		s, err := buildSubjectForMember(m, managementContext)
		if err != nil {
			return err
		}
		switch m.AccessType {
		case allAccess:
			allAccessSubjects = append(allAccessSubjects, s)
		case updateAccess:
			updateAccessSubjects = append(updateAccessSubjects, s)
		case readOnlyAccess:
			readOnlyAccessSubjects = append(readOnlyAccessSubjects, s)
		default:
			readOnlyAccessSubjects = append(readOnlyAccessSubjects, s)
		}
	}

	if _, err := createRoleBinding(allAccess, name, UID, allAccessSubjects, managementContext, resource, api); err != nil {
		return err
	}

	// Check if there are members with readonly or update access; if found then create rolebindings for those
	if len(readOnlyAccessSubjects) > 0 {
		if _, err := createRole(resource, readOnlyAccess, name, UID, managementContext, api); err != nil {
			return err
		}
		if _, err := createRoleBinding(readOnlyAccess, name, UID, readOnlyAccessSubjects, managementContext, resource, api); err != nil {
			return err
		}
	}
	if len(updateAccessSubjects) > 0 {
		if _, err := createRole(resource, updateAccess, name, UID, managementContext, api); err != nil {
			return err
		}
		if _, err := createRoleBinding(updateAccess, name, UID, updateAccessSubjects, managementContext, resource, api); err != nil {
			return err
		}
	}
	return nil
}

func createRole(resource string, roleAccess string, resourceName string, UID types.UID,
	managementContext *config.ManagementContext, apiGroups []string) (*k8srbacv1.Role, error) {
	roleName, verbs := getRoleNameAndVerbs(roleAccess, resourceName, resource)
	apiVersion := "management.cattle.io/v3"
	if apiGroups[0] != "management.cattle.io" {
		apiVersion = "v1"
	}
	ownerReference := metav1.OwnerReference{
		APIVersion: apiVersion,
		Kind:       resource,
		Name:       resourceName,
		UID:        UID,
	}
	newRole := &k8srbacv1.Role{
		ObjectMeta: metav1.ObjectMeta{
			Name:            roleName,
			Namespace:       namespace.GlobalNamespace,
			OwnerReferences: []metav1.OwnerReference{ownerReference},
		},
		Rules: []k8srbacv1.PolicyRule{
			{
				APIGroups:     apiGroups,
				Resources:     []string{resource},
				ResourceNames: []string{resourceName},
				Verbs:         verbs,
			},
		},
	}
	role, err := managementContext.RBAC.Roles("").GetNamespaced(namespace.GlobalNamespace, roleName, metav1.GetOptions{})
	if err != nil {
		if apierrors.IsNotFound(err) {
			role, err = managementContext.RBAC.Roles("").Create(newRole)
			if err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	} else if role != nil {
		if !reflect.DeepEqual(newRole, role) {
			toUpdate := newRole.DeepCopy()
			updated, err := managementContext.RBAC.Roles("").Update(toUpdate)
			if err != nil {
				return updated, err
			}
		}
	}
	return role, nil
}

func createRoleBinding(roleAccess string, name string, UID types.UID,
	subjects []k8srbacv1.Subject, managementContext *config.ManagementContext, resource string, apiGroups []string) (*k8srbacv1.RoleBinding, error) {
	roleName, _ := getRoleNameAndVerbs(roleAccess, name, resource)
	// we can define the rolebinding first, since if it's not already present we can call create. And if it's present then we'll
	// still need to compare the current members' list
	sort.Slice(subjects, func(i, j int) bool { return subjects[i].Name < subjects[j].Name })

	apiVersion := "management.cattle.io/v3"
	if apiGroups[0] != "management.cattle.io" {
		apiVersion = "v1"
	}
	ownerReference := metav1.OwnerReference{
		APIVersion: apiVersion,
		Kind:       resource,
		Name:       name,
		UID:        UID,
	}

	newRoleBinding := &k8srbacv1.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Name:            roleName,
			Namespace:       namespace.GlobalNamespace,
			OwnerReferences: []metav1.OwnerReference{ownerReference},
		},
		RoleRef: k8srbacv1.RoleRef{
			Name: roleName,
			Kind: "Role",
		},
		Subjects: subjects,
	}
	roleBinding, err := managementContext.RBAC.RoleBindings("").GetNamespaced(namespace.GlobalNamespace, roleName, metav1.GetOptions{})
	if err != nil {
		if apierrors.IsNotFound(err) {
			_, err = managementContext.RBAC.RoleBindings("").Create(newRoleBinding)
			if err != nil && !apierrors.IsAlreadyExists(err) {
				return nil, err
			}
		} else {
			return nil, err
		}
	} else if roleBinding != nil {
		if !reflect.DeepEqual(roleBinding, newRoleBinding) {
			toUpdate := newRoleBinding.DeepCopy()
			updated, err := managementContext.RBAC.RoleBindings("").Update(toUpdate)
			if err != nil {
				return updated, err
			}
		}
	}
	return newRoleBinding, nil
}

func getRoleNameAndVerbs(roleAccess string, resourceName string, resource string) (string, []string) {
	var roleName string
	var verbs []string
	switch resource {
	case MultiClusterAppResource:
		resourceName += "-m"
	case GlobalDNSResource:
		resourceName += "-g"
	}
	switch roleAccess {
	case allAccess:
		roleName = resourceName + "a"
		verbs = []string{"*"}
	case updateAccess:
		roleName = resourceName + "u"
		verbs = []string{"update", "get", "list", "watch"}
	case readOnlyAccess:
		roleName = resourceName + "r"
		verbs = []string{"get", "list", "watch"}
	}
	return roleName, verbs
}

func buildSubjectForMember(member v3.Member, managementContext *config.ManagementContext) (k8srbacv1.Subject, error) {
	var name, kind string
	member, err := checkAndSetUserFields(member, managementContext)
	if err != nil {
		return k8srbacv1.Subject{}, err
	}

	if member.UserName != "" {
		name = member.UserName
		kind = "User"
	}

	if member.GroupPrincipalName != "" {
		if name != "" {
			return k8srbacv1.Subject{}, fmt.Errorf("member %v has both username and groupPrincipalId set", member)
		}
		name = member.GroupPrincipalName
		kind = "Group"
	}

	if name == "" {
		return k8srbacv1.Subject{}, fmt.Errorf("member %v doesn't have any name fields set", member)
	}

	return k8srbacv1.Subject{
		Kind:     kind,
		Name:     name,
		APIGroup: rbacv1.GroupName,
	}, nil
}

func checkAndSetUserFields(m v3.Member, managementContext *config.ManagementContext) (v3.Member, error) {
	if m.GroupPrincipalName != "" || (m.UserPrincipalName != "" && m.UserName != "") {
		return m, nil
	}
	if m.UserPrincipalName != "" && m.UserName == "" {
		displayName := m.DisplayName
		user, err := managementContext.UserManager.EnsureUser(m.UserPrincipalName, displayName)
		if err != nil {
			return m, err
		}

		m.UserName = user.Name
		return m, nil
	}
	return m, nil
}

func GetMemberGroups(members []v3.Member) map[string]string {
	groups := make(map[string]string)
	for _, m := range members {
		if m.GroupPrincipalName != "" {
			groups[m.GroupPrincipalName] = m.AccessType
		}
	}
	return groups
}

func GetUpdatedMembers(projects []string, members []v3.Member, prtbLister v3.ProjectRoleTemplateBindingLister) ([]v3.Member, error) {
	var updatedUsers, updatedGroups, updatedMembers []v3.Member
	users := make(map[string]bool)
	groups := make(map[string]bool)
	for _, p := range projects {
		split := strings.SplitN(p, ":", 2)
		if len(split) != 2 {
			return updatedMembers, fmt.Errorf("Target project name %s is invalid", p)
		}
		projectNS := split[1]
		prtbs, err := prtbLister.List(projectNS, labels.NewSelector())
		if err != nil {
			return updatedMembers, err
		}
		prtbMap := make(map[string]bool)
		// get all PRTBs for this project
		for _, prtb := range prtbs {
			if prtb.UserPrincipalName != "" && prtb.DeletionTimestamp == nil {
				prtbMap[prtb.UserPrincipalName] = true
			} else if prtb.GroupPrincipalName != "" && prtb.DeletionTimestamp == nil {
				prtbMap[prtb.GroupPrincipalName] = true
			}
		}

		// all members should have PRTBs in this project, if not, remove them
		for _, mem := range members {
			if mem.UserPrincipalName != "" && prtbMap[mem.UserPrincipalName] && !users[mem.UserPrincipalName] {
				users[mem.UserPrincipalName] = true
				updatedUsers = append(updatedUsers, mem)
			} else if mem.GroupPrincipalName != "" && prtbMap[mem.GroupPrincipalName] && !groups[mem.GroupPrincipalName] {
				groups[mem.GroupPrincipalName] = true
				updatedGroups = append(updatedGroups, mem)
			}
		}
	}
	sort.Slice(updatedUsers, func(i, j int) bool { return updatedUsers[i].UserPrincipalName < updatedUsers[j].UserPrincipalName })
	sort.Slice(updatedGroups, func(i, j int) bool { return updatedGroups[i].GroupPrincipalName < updatedGroups[j].GroupPrincipalName })
	updatedMembers = append(updatedUsers, updatedGroups...)
	return updatedMembers, nil
}

func GetCurrentMembers(members []v3.Member) []v3.Member {
	var memberUsers, memberGroups []v3.Member
	for _, m := range members {
		if m.UserPrincipalName != "" {
			memberUsers = append(memberUsers, m)
		} else if m.GroupPrincipalName != "" {
			memberGroups = append(memberGroups, m)
		}
	}
	sort.Slice(memberUsers, func(i, j int) bool { return memberUsers[i].UserPrincipalName < memberUsers[j].UserPrincipalName })
	sort.Slice(memberGroups, func(i, j int) bool { return memberGroups[i].GroupPrincipalName < memberGroups[j].GroupPrincipalName })
	memberUsers = append(memberUsers, memberGroups...)
	return memberUsers
}
<|endoftext|>


<commit_before>package apps

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

	. "github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	. "github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry-incubator/cf-test-helpers/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
	archive_helpers "github.com/pivotal-golang/archiver/extractor/test_helper"
)

var _ = Describe("Buildpack cache", func() {
	var (
		appName       string
		BuildpackName string

		appPath string

		buildpackPath        string
		buildpackArchivePath string
	)

	matchingFilename := func(appName string) string {
		return fmt.Sprintf("buildpack-for-buildpack-cache-test-%s", appName)
	}

	BeforeEach(func() {
		AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
			BuildpackName = RandomName()
			appName = RandomName()

			tmpdir, err := ioutil.TempDir(os.TempDir(), "matching-app")
			Expect(err).ToNot(HaveOccurred())

			appPath = tmpdir

			tmpdir, err = ioutil.TempDir(os.TempDir(), "matching-buildpack")
			Expect(err).ToNot(HaveOccurred())

			buildpackPath = tmpdir
			buildpackArchivePath = path.Join(buildpackPath, "buildpack.zip")

			archive_helpers.CreateZipArchive(buildpackArchivePath, []archive_helpers.ArchiveFile{
				{
					Name: "bin/compile",
					Body: `#!/usr/bin/env bash

mkdir -p $1 $2
if [ -f "$2/cached-file" ]; then
	cp $2/cached-file $1/content
else
	echo "cache not found" > $1/content
fi

echo "here's a cache" > $2/cached-file
`,
				},
				{
					Name: "bin/detect",
					Body: fmt.Sprintf(`#!/bin/bash

if [ -f "${1}/%s" ]; then
  echo Buildpack that needs cache
else
  echo no
  exit 1
fi
`, matchingFilename(appName)),
				},
				{
					Name: "bin/release",
					Body: `#!/usr/bin/env bash

content=$(cat $1/content)

cat <<EOF
---
config_vars:
  PATH: bin:/usr/local/bin:/usr/bin:/bin
  FROM_BUILD_PACK: "yes"
default_process_types:
  web: while true; do { echo -e 'HTTP/1.1 200 OK\r\n'; echo "custom buildpack contents - $content"; } | nc -l \$PORT; done
EOF
`,
				},
			})

			_, err = os.Create(path.Join(appPath, matchingFilename(appName)))
			Expect(err).ToNot(HaveOccurred())

			_, err = os.Create(path.Join(appPath, "some-file"))
			Expect(err).ToNot(HaveOccurred())

			createBuildpack := Cf("create-buildpack", BuildpackName, buildpackArchivePath, "0").Wait(DEFAULT_TIMEOUT)
			Expect(createBuildpack).Should(Exit(0))
			Expect(createBuildpack).Should(Say("Creating"))
			Expect(createBuildpack).Should(Say("OK"))
			Expect(createBuildpack).Should(Say("Uploading"))
			Expect(createBuildpack).Should(Say("OK"))
		})
	})

	AfterEach(func() {
		AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
			Expect(Cf("delete-buildpack", BuildpackName, "-f").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		})
	})

	It("uses the buildpack cache after first staging", func() {
		push := Cf("push", appName, "-p", appPath).Wait(CF_PUSH_TIMEOUT)
		Expect(push).To(Exit(0))

		Eventually(func() string {
			return helpers.CurlAppRoot(appName)
		}, DEFAULT_TIMEOUT).Should(ContainSubstring("custom buildpack contents - cache not found"))

		restage := Cf("restage", appName).Wait(CF_PUSH_TIMEOUT)
		Expect(restage).To(Exit(0))

		Eventually(func() string {
			return helpers.CurlAppRoot(appName)
		}, DEFAULT_TIMEOUT).Should(ContainSubstring("custom buildpack contents - here's a cache"))
	})
})
<commit_msg>Allow buildpack cache time to be uploaded.<commit_after>package apps

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"time"

	. "github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	. "github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry-incubator/cf-test-helpers/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
	archive_helpers "github.com/pivotal-golang/archiver/extractor/test_helper"
)

var _ = Describe("Buildpack cache", func() {
	var (
		appName       string
		BuildpackName string

		appPath string

		buildpackPath        string
		buildpackArchivePath string
	)

	matchingFilename := func(appName string) string {
		return fmt.Sprintf("buildpack-for-buildpack-cache-test-%s", appName)
	}

	BeforeEach(func() {
		AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
			BuildpackName = RandomName()
			appName = RandomName()

			tmpdir, err := ioutil.TempDir(os.TempDir(), "matching-app")
			Expect(err).ToNot(HaveOccurred())

			appPath = tmpdir

			tmpdir, err = ioutil.TempDir(os.TempDir(), "matching-buildpack")
			Expect(err).ToNot(HaveOccurred())

			buildpackPath = tmpdir
			buildpackArchivePath = path.Join(buildpackPath, "buildpack.zip")

			archive_helpers.CreateZipArchive(buildpackArchivePath, []archive_helpers.ArchiveFile{
				{
					Name: "bin/compile",
					Body: `#!/usr/bin/env bash

mkdir -p $1 $2
if [ -f "$2/cached-file" ]; then
	cp $2/cached-file $1/content
else
	echo "cache not found" > $1/content
fi

echo "here's a cache" > $2/cached-file
`,
				},
				{
					Name: "bin/detect",
					Body: fmt.Sprintf(`#!/bin/bash

if [ -f "${1}/%s" ]; then
  echo Buildpack that needs cache
else
  echo no
  exit 1
fi
`, matchingFilename(appName)),
				},
				{
					Name: "bin/release",
					Body: `#!/usr/bin/env bash

content=$(cat $1/content)

cat <<EOF
---
config_vars:
  PATH: bin:/usr/local/bin:/usr/bin:/bin
  FROM_BUILD_PACK: "yes"
default_process_types:
  web: while true; do { echo -e 'HTTP/1.1 200 OK\r\n'; echo "custom buildpack contents - $content"; } | nc -l \$PORT; done
EOF
`,
				},
			})

			_, err = os.Create(path.Join(appPath, matchingFilename(appName)))
			Expect(err).ToNot(HaveOccurred())

			_, err = os.Create(path.Join(appPath, "some-file"))
			Expect(err).ToNot(HaveOccurred())

			createBuildpack := Cf("create-buildpack", BuildpackName, buildpackArchivePath, "0").Wait(DEFAULT_TIMEOUT)
			Expect(createBuildpack).Should(Exit(0))
			Expect(createBuildpack).Should(Say("Creating"))
			Expect(createBuildpack).Should(Say("OK"))
			Expect(createBuildpack).Should(Say("Uploading"))
			Expect(createBuildpack).Should(Say("OK"))
		})
	})

	AfterEach(func() {
		AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
			Expect(Cf("delete-buildpack", BuildpackName, "-f").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		})
	})

	It("uses the buildpack cache after first staging", func() {
		push := Cf("push", appName, "-p", appPath).Wait(CF_PUSH_TIMEOUT)
		Expect(push).To(Exit(0))

		Eventually(func() string {
			return helpers.CurlAppRoot(appName)
		}, DEFAULT_TIMEOUT).Should(ContainSubstring("custom buildpack contents - cache not found"))

		time.Sleep(DEFAULT_TIMEOUT)

		restage := Cf("restage", appName).Wait(CF_PUSH_TIMEOUT)
		Expect(restage).To(Exit(0))

		Eventually(func() string {
			return helpers.CurlAppRoot(appName)
		}, DEFAULT_TIMEOUT).Should(ContainSubstring("custom buildpack contents - here's a cache"))
	})
})
<|endoftext|>


<commit_before>package core

import (
	"errors"
	"log"
	"os"
	"time"

	"github.com/mdlayher/goset"
	"github.com/romanoff/fsmonitor"
)

// validSet is a set of valid file extensions which we should scan as media, as they are the ones
// which TagLib is capable of reading
var validSet = set.New(".ape", ".flac", ".m4a", ".mp3", ".mpc", ".ogg", ".wma", ".wv")

// fsQueue is a queue of tasks to be performed by the filesystem, such as media and orphan scans
var fsQueue = make(chan fsTask, 10)

// fsSource is the data source used to scan for media files (could be filesystem, memory, etc)
var fsSource fileSource

// fsTask is the interface which defines a filesystem task, such as a media scan, or an orphan scan
type fsTask interface {
	Folders() (string, string)
	SetFolders(string, string)
	Scan(string, string, chan struct{}) error
	Verbose(bool)
}

// fileSource represents a source from which files can be scanned and indexed
type fileSource interface {
	MediaScan(string, bool, chan struct{}) error
	OrphanScan(string, string, bool, chan struct{}) error
}

// fsManager handles fsWalker processes, and communicates back and forth with the manager goroutine
func fsManager(mediaFolder string, fsKillChan chan struct{}) {
	log.Println("fs: starting...")

	// Initialize a queue to cancel filesystem tasks
	cancelQueue := make(chan chan struct{}, 10)

	// Set up the data source (typically filesystem, unless in test mode)
	fsSource = fsFileSource{}
	if os.Getenv("WAVEPIPE_TEST") == "1" {
		// Mock file source
		fsSource = memFileSource{}
	}

	// Queue an initial, verbose orphan scan
	o := new(fsOrphanScan)
	o.SetFolders(mediaFolder, "")
	o.Verbose(true)
	fsQueue <- o

	// Queue a media scan
	m := new(fsMediaScan)
	m.SetFolders(mediaFolder, "")
	m.Verbose(true)
	fsQueue <- m

	// Invoke task queue via goroutine, so it can be halted via the manager
	go func() {
		for {
			select {
			// Trigger a fsTask from queue
			case task := <-fsQueue:
				// Create a channel to halt the scan
				cancelChan := make(chan struct{})
				cancelQueue <- cancelChan

				// Retrieve the folders to use with scan
				baseFolder, subFolder := task.Folders()

				// Start the scan
				if err := task.Scan(baseFolder, subFolder, cancelChan); err != nil {
					log.Println(err)
				}

				// On completion, close the cancel channel
				cancelChan = <-cancelQueue
				close(cancelChan)
			}
		}
	}()

	// Create a filesystem watcher
	watcher, err := fsmonitor.NewWatcher()
	if err != nil {
		log.Println(err)
		return
	}

	// Wait for events on goroutine
	go func() {
		// Recently modified files set, used as a rate-limiter to prevent modify events from flooding
		// the select statement.  The filesystem watcher appears to fire off as many events as
		// there are files in the filesystem, so this will block most of those.
		recentModifySet := set.New()

		for {
			select {
			// Event occurred
			case ev := <-watcher.Event:
				switch {
				// On create or modify, trigger a media scan
				case ev.IsModify():
					// Add file to set, stopping it from propogating if the event was recently triggered
					if !recentModifySet.Add(ev.Name) {
						break
					}

					// Remove file from rate-limiting set after a couple seconds
					go func() {
						<-time.After(2 * time.Second)
						recentModifySet.Remove(ev.Name)
					}()

					fallthrough
				// On create, trigger a media scan
				case ev.IsCreate():
					// Invoke a slight delay to enable file creation
					<-time.After(250 * time.Millisecond)

					// Scan item as the "base folder", so it just adds this item
					m := new(fsMediaScan)
					m.SetFolders(ev.Name, "")
					m.Verbose(false)
					fsQueue <- m
				// On delete, trigger an orphan scan
				case ev.IsDelete():
					// Invoke a slight delay to enable file deletion
					<-time.After(250 * time.Millisecond)

					// Scan item as the "subfolder", so it just removes this item
					o := new(fsOrphanScan)
					o.SetFolders("", ev.Name)
					o.Verbose(false)
					fsQueue <- o
				}
			// Watcher errors
			case err := <-watcher.Error:
				log.Println(err)
				return
			}
		}
	}()

	// Watch media folder
	if err := watcher.Watch(mediaFolder); err != nil {
		log.Println(err)
		return
	}
	log.Println("fs: watching folder:", mediaFolder)

	// Trigger manager events via channel
	for {
		select {
		// Stop filesystem manager
		case <-fsKillChan:
			// Halt any in-progress tasks
			log.Println("fs: halting tasks")
			for i := 0; i < len(cancelQueue); i++ {
				// Receive a channel
				f := <-cancelQueue
				if f == nil {
					continue
				}

				// Send termination
				f <- struct{}{}
				log.Println("fs: task halted")
			}

			// Inform manager that shutdown is complete
			log.Println("fs: stopped!")
			fsKillChan <- struct{}{}
			return
		}
	}
}

// fsMediaScan represents a filesystem task which scans the given path for new media
type fsMediaScan struct {
	baseFolder string
	subFolder  string
	verbose    bool
}

// Folders returns the base folder and subfolder for use with a scanning task
func (fs *fsMediaScan) Folders() (string, string) {
	return fs.baseFolder, fs.subFolder
}

// SetFolders sets the base folder and subfolder for use with a scanning task
func (fs *fsMediaScan) SetFolders(baseFolder string, subFolder string) {
	fs.baseFolder = baseFolder
	fs.subFolder = subFolder
}

// Verbose sets the verbosity level of the scanning task
func (fs *fsMediaScan) Verbose(verbose bool) {
	fs.verbose = verbose
}

// Scan scans for media files in a specified path, and queues them up for inclusion
// in the wavepipe database
func (fs *fsMediaScan) Scan(mediaFolder string, subFolder string, walkCancelChan chan struct{}) error {
	// Media scans are comprehensive, so subfolder has no purpose
	if subFolder != "" {
		return errors.New("media scan: subfolder not valid for media scan operation")
	}

	// Scan for media using the specified file source
	return fsSource.MediaScan(mediaFolder, fs.verbose, walkCancelChan)
}

// fsOrphanScan represents a filesystem task which scans the given path for orphaned media
type fsOrphanScan struct {
	baseFolder string
	subFolder  string
	verbose    bool
}

// Folders returns the base folder and subfolder for use with a scanning task
func (fs *fsOrphanScan) Folders() (string, string) {
	return fs.baseFolder, fs.subFolder
}

// SetFolders sets the base folder and subfolder for use with a scanning task
func (fs *fsOrphanScan) SetFolders(baseFolder string, subFolder string) {
	fs.baseFolder = baseFolder
	fs.subFolder = subFolder
}

// Verbose sets the verbosity level of the scanning task
func (fs *fsOrphanScan) Verbose(verbose bool) {
	fs.verbose = verbose
}

// Scan scans for media files which have been removed from the media directory, and removes
// them as appropriate.  An orphan is defined as follows:
//   - Artist: no more songs contain this artist's ID
//   - Album: no more songs contain this album's ID
//   - Song: song is no longer present in the filesystem
// The baseFolder is the root location of the media folder.  As wavepipe currently supports only
// one media folder, any media which does not reside in this folder is orphaned.  If left blank,
// only the subFolder will be checked.
// The subFolder is the current file location, under the baseFolder.  This is used to allow for
// quick scans of a small subsection of the directory, such as on a filesystem change.  Any files
// which are in the database, but do not exist on disk, will be orphaned and removed.
func (fs *fsOrphanScan) Scan(baseFolder string, subFolder string, orphanCancelChan chan struct{}) error {
	// If both folders are empty, there is nothing to do
	if baseFolder == "" && subFolder == "" {
		return errors.New("orphan scan: no base folder or subfolder")
	}

	// Scan for orphans using the specified file source
	return fsSource.OrphanScan(baseFolder, subFolder, fs.verbose, orphanCancelChan)
}
<commit_msg>Add experimental support for IsRename() event<commit_after>package core

import (
	"errors"
	"log"
	"os"
	"time"

	"github.com/mdlayher/goset"
	"github.com/romanoff/fsmonitor"
)

// validSet is a set of valid file extensions which we should scan as media, as they are the ones
// which TagLib is capable of reading
var validSet = set.New(".ape", ".flac", ".m4a", ".mp3", ".mpc", ".ogg", ".wma", ".wv")

// fsQueue is a queue of tasks to be performed by the filesystem, such as media and orphan scans
var fsQueue = make(chan fsTask, 10)

// fsSource is the data source used to scan for media files (could be filesystem, memory, etc)
var fsSource fileSource

// fsTask is the interface which defines a filesystem task, such as a media scan, or an orphan scan
type fsTask interface {
	Folders() (string, string)
	SetFolders(string, string)
	Scan(string, string, chan struct{}) error
	Verbose(bool)
}

// fileSource represents a source from which files can be scanned and indexed
type fileSource interface {
	MediaScan(string, bool, chan struct{}) error
	OrphanScan(string, string, bool, chan struct{}) error
}

// fsManager handles fsWalker processes, and communicates back and forth with the manager goroutine
func fsManager(mediaFolder string, fsKillChan chan struct{}) {
	log.Println("fs: starting...")

	// Initialize a queue to cancel filesystem tasks
	cancelQueue := make(chan chan struct{}, 10)

	// Set up the data source (typically filesystem, unless in test mode)
	fsSource = fsFileSource{}
	if os.Getenv("WAVEPIPE_TEST") == "1" {
		// Mock file source
		fsSource = memFileSource{}
	}

	// Queue an initial, verbose orphan scan
	o := new(fsOrphanScan)
	o.SetFolders(mediaFolder, "")
	o.Verbose(true)
	fsQueue <- o

	// Queue a media scan
	m := new(fsMediaScan)
	m.SetFolders(mediaFolder, "")
	m.Verbose(true)
	fsQueue <- m

	// Invoke task queue via goroutine, so it can be halted via the manager
	go func() {
		for {
			select {
			// Trigger a fsTask from queue
			case task := <-fsQueue:
				// Create a channel to halt the scan
				cancelChan := make(chan struct{})
				cancelQueue <- cancelChan

				// Retrieve the folders to use with scan
				baseFolder, subFolder := task.Folders()

				// Start the scan
				if err := task.Scan(baseFolder, subFolder, cancelChan); err != nil {
					log.Println(err)
				}

				// On completion, close the cancel channel
				cancelChan = <-cancelQueue
				close(cancelChan)
			}
		}
	}()

	// Create a filesystem watcher
	watcher, err := fsmonitor.NewWatcher()
	if err != nil {
		log.Println(err)
		return
	}

	// Wait for events on goroutine
	go func() {
		// Recently modified/renamed files sets, used as rate-limiters to prevent modify
		// events from flooding the select statement.  The filesystem watcher may fire an
		// excessive number of events, so these will block the extras for a couple seconds.
		recentModifySet := set.New()
		recentRenameSet := set.New()

		for {
			select {
			// Event occurred
			case ev := <-watcher.Event:
				switch {
				// On modify, trigger a media scan
				case ev.IsModify():
					// Add file to set, stopping it from propogating if the event was recently triggered
					if !recentModifySet.Add(ev.Name) {
						break
					}

					// Remove file from rate-limiting set after a couple seconds
					go func() {
						<-time.After(2 * time.Second)
						recentModifySet.Remove(ev.Name)
					}()

					fallthrough
				// On create, trigger a media scan
				case ev.IsCreate():
					// Invoke a slight delay to enable file creation
					<-time.After(250 * time.Millisecond)

					// Scan item as the "base folder", so it just adds this item
					m := new(fsMediaScan)
					m.SetFolders(ev.Name, "")
					m.Verbose(false)
					fsQueue <- m
				// On rename, trigger an orphan scan
				case ev.IsRename():
					// Add file to set, stopping it from propogating if the event was recently triggered
					if !recentRenameSet.Add(ev.Name) {
						break
					}

					// Remove file from rate-limiting set after a couple seconds
					go func() {
						<-time.After(2 * time.Second)
						recentRenameSet.Remove(ev.Name)
					}()

					fallthrough
				// On delete, trigger an orphan scan
				case ev.IsDelete():
					// Invoke a slight delay to enable file deletion
					<-time.After(250 * time.Millisecond)

					// Scan item as the "subfolder", so it just removes this item
					o := new(fsOrphanScan)
					o.SetFolders("", ev.Name)
					o.Verbose(false)
					fsQueue <- o
				}
			// Watcher errors
			case err := <-watcher.Error:
				log.Println(err)
				return
			}
		}
	}()

	// Watch media folder
	if err := watcher.Watch(mediaFolder); err != nil {
		log.Println(err)
		return
	}
	log.Println("fs: watching folder:", mediaFolder)

	// Trigger manager events via channel
	for {
		select {
		// Stop filesystem manager
		case <-fsKillChan:
			// Halt any in-progress tasks
			log.Println("fs: halting tasks")
			for i := 0; i < len(cancelQueue); i++ {
				// Receive a channel
				f := <-cancelQueue
				if f == nil {
					continue
				}

				// Send termination
				f <- struct{}{}
				log.Println("fs: task halted")
			}

			// Inform manager that shutdown is complete
			log.Println("fs: stopped!")
			fsKillChan <- struct{}{}
			return
		}
	}
}

// fsMediaScan represents a filesystem task which scans the given path for new media
type fsMediaScan struct {
	baseFolder string
	subFolder  string
	verbose    bool
}

// Folders returns the base folder and subfolder for use with a scanning task
func (fs *fsMediaScan) Folders() (string, string) {
	return fs.baseFolder, fs.subFolder
}

// SetFolders sets the base folder and subfolder for use with a scanning task
func (fs *fsMediaScan) SetFolders(baseFolder string, subFolder string) {
	fs.baseFolder = baseFolder
	fs.subFolder = subFolder
}

// Verbose sets the verbosity level of the scanning task
func (fs *fsMediaScan) Verbose(verbose bool) {
	fs.verbose = verbose
}

// Scan scans for media files in a specified path, and queues them up for inclusion
// in the wavepipe database
func (fs *fsMediaScan) Scan(mediaFolder string, subFolder string, walkCancelChan chan struct{}) error {
	// Media scans are comprehensive, so subfolder has no purpose
	if subFolder != "" {
		return errors.New("media scan: subfolder not valid for media scan operation")
	}

	// Scan for media using the specified file source
	return fsSource.MediaScan(mediaFolder, fs.verbose, walkCancelChan)
}

// fsOrphanScan represents a filesystem task which scans the given path for orphaned media
type fsOrphanScan struct {
	baseFolder string
	subFolder  string
	verbose    bool
}

// Folders returns the base folder and subfolder for use with a scanning task
func (fs *fsOrphanScan) Folders() (string, string) {
	return fs.baseFolder, fs.subFolder
}

// SetFolders sets the base folder and subfolder for use with a scanning task
func (fs *fsOrphanScan) SetFolders(baseFolder string, subFolder string) {
	fs.baseFolder = baseFolder
	fs.subFolder = subFolder
}

// Verbose sets the verbosity level of the scanning task
func (fs *fsOrphanScan) Verbose(verbose bool) {
	fs.verbose = verbose
}

// Scan scans for media files which have been removed from the media directory, and removes
// them as appropriate.  An orphan is defined as follows:
//   - Artist: no more songs contain this artist's ID
//   - Album: no more songs contain this album's ID
//   - Song: song is no longer present in the filesystem
// The baseFolder is the root location of the media folder.  As wavepipe currently supports only
// one media folder, any media which does not reside in this folder is orphaned.  If left blank,
// only the subFolder will be checked.
// The subFolder is the current file location, under the baseFolder.  This is used to allow for
// quick scans of a small subsection of the directory, such as on a filesystem change.  Any files
// which are in the database, but do not exist on disk, will be orphaned and removed.
func (fs *fsOrphanScan) Scan(baseFolder string, subFolder string, orphanCancelChan chan struct{}) error {
	// If both folders are empty, there is nothing to do
	if baseFolder == "" && subFolder == "" {
		return errors.New("orphan scan: no base folder or subfolder")
	}

	// Scan for orphans using the specified file source
	return fsSource.OrphanScan(baseFolder, subFolder, fs.verbose, orphanCancelChan)
}
<|endoftext|>


<commit_before>package aphid

import (
	// External requirements.
	_ "github.com/BurntSushi/toml"
	_ "github.com/gdamore/mangos"
	_ "github.com/huandu/skiplist"
	_ "github.com/microcosm-cc/bluemonday"
	_ "github.com/russross/blackfriday"
	_ "golang.org/x/crypto/scrypt"
	_ "github.com/syndtr/goleveldb/leveldb"
)

import (
	// Yak requirements.
	_ "github.com/strickyak/redhed"
	_ "github.com/strickyak/rye"
	_ "github.com/yak-labs/chirp-lang"
	_ "github.com/yak-labs/chirp-lang/goapi/default"
	_ "github.com/yak-labs/chirp-lang/http"
	_ "github.com/yak-labs/chirp-lang/img"
	_ "github.com/yak-labs/chirp-lang/posix"
	_ "github.com/yak-labs/chirp-lang/rpc"
	_ "github.com/yak-labs/chirp-lang/ryba"
)

import (
	_ "io"
)
<commit_msg>add github.com/nfnt/resize to ahpid.<commit_after>package aphid

import (
	// External requirements.
	_ "github.com/BurntSushi/toml"
	_ "github.com/gdamore/mangos"
	_ "github.com/huandu/skiplist"
	_ "github.com/microcosm-cc/bluemonday"
	_ "github.com/nfnt/resize"
	_ "github.com/russross/blackfriday"
	_ "github.com/syndtr/goleveldb/leveldb"
	_ "golang.org/x/crypto/scrypt"
)

import (
	// Yak requirements.
	_ "github.com/strickyak/redhed"
	_ "github.com/strickyak/rye"
	_ "github.com/yak-labs/chirp-lang"
	_ "github.com/yak-labs/chirp-lang/goapi/default"
	_ "github.com/yak-labs/chirp-lang/http"
	_ "github.com/yak-labs/chirp-lang/img"
	_ "github.com/yak-labs/chirp-lang/posix"
	_ "github.com/yak-labs/chirp-lang/rpc"
	_ "github.com/yak-labs/chirp-lang/ryba"
)

import (
	_ "io"
)
<|endoftext|>


<commit_before>package main

import (
	"log"

	"bitbucket.org/firstrow/logvoyage/models"
	"bitbucket.org/firstrow/logvoyage/shared/config"

	jwt "github.com/dgrijalva/jwt-go"
	"gopkg.in/kataras/iris.v6"
	"gopkg.in/kataras/iris.v6/adaptors/httprouter"
)

var (
	app      *iris.Framework
	response Response
)

type Response struct {
}

// Success responses with 200 code.
// Note: only fist body argument will be passed to the response.
func (r Response) Success(ctx *iris.Context, body ...interface{}) {
	if len(body) > 0 {
		ctx.JSON(200, map[string]interface{}{"success": true, "body": body[0]})
	} else {
		ctx.JSON(200, map[string]interface{}{"success": true})
	}
}

// Error returns 200 OK response with json field "errors" with error descrioption.
// This function should be used to display validation or other expected errors.
// errors may be string or array of hashes.
func (r Response) Error(ctx *iris.Context, err interface{}) {
	ctx.JSON(200, map[string]interface{}{"errors": err})

}

// Panic responses with 503 error.
func (r Response) Panic(ctx *iris.Context, err error) {
	// TODO: Send orignal error to issue tracker.
	ctx.JSON(503, map[string]interface{}{"errors": err.Error()})
}

// Forbidden responses with 401 code, means user does not valid credentials to access handler.
func (r Response) Forbidden(ctx *iris.Context) {
	ctx.StopExecution()
	ctx.JSON(401, map[string]string{"errors": "Authentication failed"})
}

// authMiddleware performs authentication
func authMiddleware(ctx *iris.Context) {
	tokenString := ctx.RequestHeader("Authentication")

	token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			log.Printf("Unexpected signing method: %v", token.Header["alg"])
		}
		return []byte(config.Get("secret")), nil
	})

	if err != nil {
		response.Forbidden(ctx)
		return
	}

	if claims, ok := token.Claims.(jwt.MapClaims); ok && token.Valid {
		user, res := models.FindUserByID(claims["user_id"])

		if res.Error != nil {
			if res.RecordNotFound() {
				response.Forbidden(ctx)
			} else {
				response.Error(ctx, "User not found")
			}
			return
		}

		ctx.Set("user", user)
	} else {
		response.Forbidden(ctx)
		return
	}

	ctx.Next()
}

func init() {
	response = Response{}

	app = iris.New()
	app.Adapt(httprouter.New())
	app.Adapt(iris.DevLogger())

	userAPI := app.Party("/users")
	{
		userAPI.Post("/", UsersCreate)
		userAPI.Post("/login", UsersLogin)
	}

	projectAPI := app.Party("/projects", authMiddleware)
	{
		projectAPI.Post("/", ProjectsCreate)
		projectAPI.Get("/", ProjectsList)
	}
}

func main() {
	app.Listen("127.0.0.1:3000")
}
<commit_msg>Added simple cors<commit_after>package main

import (
	"log"

	"bitbucket.org/firstrow/logvoyage/models"
	"bitbucket.org/firstrow/logvoyage/shared/config"

	jwt "github.com/dgrijalva/jwt-go"
	"gopkg.in/kataras/iris.v6"
	"gopkg.in/kataras/iris.v6/adaptors/httprouter"
)

var (
	app      *iris.Framework
	response Response
)

type Response struct {
}

// Success responses with 200 code.
// Note: only fist body argument will be passed to the response.
func (r Response) Success(ctx *iris.Context, body ...interface{}) {
	if len(body) > 0 {
		ctx.JSON(200, map[string]interface{}{"success": true, "body": body[0]})
	} else {
		ctx.JSON(200, map[string]interface{}{"success": true})
	}
}

// Error returns 200 OK response with json field "errors" with error descrioption.
// This function should be used to display validation or other expected errors.
// errors may be string or array of hashes.
func (r Response) Error(ctx *iris.Context, err interface{}) {
	ctx.JSON(200, map[string]interface{}{"errors": err})

}

// Panic responses with 503 error.
func (r Response) Panic(ctx *iris.Context, err error) {
	// TODO: Send orignal error to issue tracker.
	ctx.JSON(503, map[string]interface{}{"errors": err.Error()})
}

// Forbidden responses with 401 code, means user does not valid credentials to access handler.
func (r Response) Forbidden(ctx *iris.Context) {
	ctx.StopExecution()
	ctx.JSON(401, map[string]string{"errors": "Authentication failed"})
}

// authMiddleware performs authentication
func authMiddleware(ctx *iris.Context) {
	tokenString := ctx.RequestHeader("Authentication")

	token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			log.Printf("Unexpected signing method: %v", token.Header["alg"])
		}
		return []byte(config.Get("secret")), nil
	})

	if err != nil {
		response.Forbidden(ctx)
		return
	}

	if claims, ok := token.Claims.(jwt.MapClaims); ok && token.Valid {
		user, res := models.FindUserByID(claims["user_id"])

		if res.Error != nil {
			if res.RecordNotFound() {
				response.Forbidden(ctx)
			} else {
				response.Error(ctx, "User not found")
			}
			return
		}

		ctx.Set("user", user)
	} else {
		response.Forbidden(ctx)
		return
	}

	ctx.Next()
}

func corsMiddleware(ctx *iris.Context) {
	if ctx.Request.Method != "OPTIONS" {
		ctx.Next()
	} else {
		ctx.SetHeader("Access-Control-Allow-Origin", "*")
		ctx.SetHeader("Access-Control-Allow-Methods", "GET,POST,PUT,PATCH,DELETE,OPTIONS")
		ctx.SetHeader("Access-Control-Allow-Headers", "authorization, origin, content-type, accept")
		ctx.SetHeader("Allow", "HEAD,GET,POST,PUT,PATCH,DELETE,OPTIONS")
		ctx.SetHeader("Content-Type", "application/json")
		ctx.SetStatusCode(200)
		ctx.StopExecution()
	}
}

func init() {
	response = Response{}

	app = iris.New()

	app.Adapt(
		httprouter.New(),
		iris.DevLogger(),
	)

	userAPI := app.Party("/users")
	{
		userAPI.Post("/", UsersCreate)
		userAPI.Post("/login", UsersLogin)
	}

	projectAPI := app.Party("/projects", authMiddleware)
	{
		projectAPI.Post("/", ProjectsCreate)
		projectAPI.Get("/", ProjectsList)
	}
}

func main() {
	app.Listen("127.0.0.1:3000")
}
<|endoftext|>