<commit_before>package allocation

import (
	"encoding/binary"
	"net"
	"sync"
	"time"

	"github.com/gortc/turn"
	"github.com/pion/logging"
	"github.com/pion/stun"
	"github.com/pion/turn/internal/ipnet"
	"github.com/pkg/errors"
)

// Allocation is tied to a FiveTuple and relays traffic
// use CreateAllocation and GetAllocation to operate
type Allocation struct {
	RelayAddr           net.Addr
	Protocol            Protocol
	TurnSocket          net.PacketConn
	RelaySocket         net.PacketConn
	fiveTuple           *FiveTuple
	permissionsLock     sync.RWMutex
	permissions         map[string]*Permission
	channelBindingsLock sync.RWMutex
	channelBindings     []*ChannelBind
	lifetimeTimer       *time.Timer
	closed              chan interface{}
	log                 logging.LeveledLogger
}

func addr2Fingerprint(addr net.Addr) string {
	switch a := addr.(type) {
	case *net.UDPAddr:
		return a.IP.String()
	case *net.TCPAddr: // Do we really need this case?
		return a.IP.String()
	}
	return "" // shoud never happen
}

// NewAllocation creates a new instance of NewAllocation.
func NewAllocation(turnSocket net.PacketConn, fiveTuple *FiveTuple, log logging.LeveledLogger) *Allocation {
	return &Allocation{
		TurnSocket:  turnSocket,
		fiveTuple:   fiveTuple,
		permissions: make(map[string]*Permission, 64),
		closed:      make(chan interface{}),
		log:         log,
	}
}

// GetPermission gets the Permission from the allocation
func (a *Allocation) GetPermission(addr net.Addr) *Permission {
	a.permissionsLock.RLock()
	defer a.permissionsLock.RUnlock()

	fingerprint := addr2Fingerprint(addr)
	return a.permissions[fingerprint]
}

// AddPermission adds a new permission to the allocation
func (a *Allocation) AddPermission(p *Permission) {
	fingerprint := addr2Fingerprint(p.Addr)

	a.permissionsLock.RLock()
	existedPermission, ok := a.permissions[fingerprint]
	a.permissionsLock.RUnlock()
	if ok {
		existedPermission.refresh(permissionTimeout)
		return
	}

	p.allocation = a
	a.permissionsLock.Lock()
	a.permissions[fingerprint] = p
	a.permissionsLock.Unlock()

	p.start(permissionTimeout)
}

// RemovePermission removes the net.Addr's fingerprint from the allocation's permissions
func (a *Allocation) RemovePermission(addr net.Addr) {
	a.permissionsLock.Lock()
	defer a.permissionsLock.Unlock()
	fingerprint := addr2Fingerprint(addr)
	delete(a.permissions, fingerprint)
}

// AddChannelBind adds a new ChannelBind to the allocation, it also updates the
// permissions needed for this ChannelBind
func (a *Allocation) AddChannelBind(c *ChannelBind, lifetime time.Duration) error {
	// Check that this channel id isn't bound to another transport address, and
	// that this transport address isn't bound to another channel id.
	channelByID := a.GetChannelByID(c.ID)
	channelByPeer := a.GetChannelByAddr(c.Peer)
	if channelByID != channelByPeer {
		return errors.Errorf("You cannot use the same channel number with different peer")
	}

	// Add or refresh this channel.
	if channelByID == nil {
		a.channelBindingsLock.Lock()
		defer a.channelBindingsLock.Unlock()

		c.allocation = a
		a.channelBindings = append(a.channelBindings, c)
		c.start(lifetime)

		// Channel binds also refresh permissions.
		a.AddPermission(NewPermission(c.Peer, a.log))
	} else {
		channelByID.refresh(lifetime)

		// Channel binds also refresh permissions.
		a.AddPermission(NewPermission(channelByID.Peer, a.log))
	}

	return nil
}

// RemoveChannelBind removes the ChannelBind from this allocation by id
func (a *Allocation) RemoveChannelBind(id turn.ChannelNumber) bool {
	a.channelBindingsLock.Lock()
	defer a.channelBindingsLock.Unlock()

	for i := len(a.channelBindings) - 1; i >= 0; i-- {
		if a.channelBindings[i].ID == id {
			a.channelBindings = append(a.channelBindings[:i], a.channelBindings[i+1:]...)
			return true
		}
	}

	return false
}

// GetChannelByID gets the ChannelBind from this allocation by id
func (a *Allocation) GetChannelByID(id turn.ChannelNumber) *ChannelBind {
	a.channelBindingsLock.RLock()
	defer a.channelBindingsLock.RUnlock()
	for _, cb := range a.channelBindings {
		if cb.ID == id {
			return cb
		}
	}
	return nil
}

// GetChannelByAddr gets the ChannelBind from this allocation by net.Addr
func (a *Allocation) GetChannelByAddr(addr net.Addr) *ChannelBind {
	a.channelBindingsLock.RLock()
	defer a.channelBindingsLock.RUnlock()
	for _, cb := range a.channelBindings {
		if ipnet.AddrEqual(cb.Peer, addr) {
			return cb
		}
	}
	return nil
}

// Refresh updates the allocations lifetime
func (a *Allocation) Refresh(lifetime time.Duration) {
	if lifetime == 0 {
		if !a.lifetimeTimer.Stop() {
			a.log.Errorf("Failed to stop allocation timer for %v", a.fiveTuple)
		}
		return
	}
	if !a.lifetimeTimer.Reset(lifetime) {
		a.log.Errorf("Failed to reset allocation timer for %v", a.fiveTuple)
	}
}

// Close closes the allocation
func (a *Allocation) Close() error {
	select {
	case <-a.closed:
		return nil
	default:
	}
	close(a.closed)

	a.lifetimeTimer.Stop()

	a.permissionsLock.RLock()
	for _, p := range a.permissions {
		p.lifetimeTimer.Stop()
	}
	a.permissionsLock.RUnlock()

	a.channelBindingsLock.RLock()
	for _, c := range a.channelBindings {
		c.lifetimeTimer.Stop()
	}
	a.channelBindingsLock.RUnlock()

	return a.RelaySocket.Close()
}

//  https://tools.ietf.org/html/rfc5766#section-10.3
//  When the server receives a UDP datagram at a currently allocated
//  relayed transport address, the server looks up the allocation
//  associated with the relayed transport address.  The server then
//  checks to see whether the set of permissions for the allocation allow
//  the relaying of the UDP datagram as described in Section 8.
//
//  If relaying is permitted, then the server checks if there is a
//  channel bound to the peer that sent the UDP datagram (see
//  Section 11).  If a channel is bound, then processing proceeds as
//  described in Section 11.7.
//
//  If relaying is permitted but no channel is bound to the peer, then
//  the server forms and sends a Data indication.  The Data indication
//  MUST contain both an XOR-PEER-ADDRESS and a DATA attribute.  The DATA
//  attribute is set to the value of the 'data octets' field from the
//  datagram, and the XOR-PEER-ADDRESS attribute is set to the source
//  transport address of the received UDP datagram.  The Data indication
//  is then sent on the 5-tuple associated with the allocation.

const rtpMTU = 1500

func (a *Allocation) packetHandler(m *Manager) {
	buffer := make([]byte, rtpMTU)

	for {
		n, srcAddr, err := a.RelaySocket.ReadFrom(buffer)
		if err != nil {
			m.DeleteAllocation(a.fiveTuple)
			return
		}

		a.log.Debugf("relay socket %s received %d bytes from %s",
			a.RelaySocket.LocalAddr().String(),
			n,
			srcAddr.String())

		if channel := a.GetChannelByAddr(a.RelaySocket.LocalAddr()); channel != nil {
			channelData := make([]byte, 4)
			binary.BigEndian.PutUint16(channelData[0:], uint16(channel.ID))
			binary.BigEndian.PutUint16(channelData[2:], uint16(n))
			channelData = append(channelData, buffer[:n]...)

			if _, err = a.TurnSocket.WriteTo(channelData, a.fiveTuple.SrcAddr); err != nil {
				a.log.Errorf("Failed to send ChannelData from allocation %v %v", srcAddr, err)
			}
		} else if p := a.GetPermission(srcAddr); p != nil {
			udpAddr := srcAddr.(*net.UDPAddr)
			peerAddressAttr := turn.PeerAddress{IP: udpAddr.IP, Port: udpAddr.Port}
			dataAttr := turn.Data(buffer[:n])

			msg, err := stun.Build(stun.TransactionID, stun.NewType(stun.MethodData, stun.ClassIndication), peerAddressAttr, dataAttr)
			if err != nil {
				a.log.Errorf("Failed to send DataIndication from allocation %v %v", srcAddr, err)
			}
			a.log.Debugf("relaying message to client at %s", a.fiveTuple.SrcAddr.String())
			if _, err = a.TurnSocket.WriteTo(msg.Raw, a.fiveTuple.SrcAddr); err != nil {
				a.log.Errorf("Failed to send DataIndication from allocation %v %v", srcAddr, err)
			}

		} else {
			a.log.Errorf("Packet unhandled in relay src %v", srcAddr)
		}
	}

}
<commit_msg>Adopt the suggestion from review feedback<commit_after>package allocation

import (
	"encoding/binary"
	"net"
	"sync"
	"time"

	"github.com/gortc/turn"
	"github.com/pion/logging"
	"github.com/pion/stun"
	"github.com/pion/turn/internal/ipnet"
	"github.com/pkg/errors"
)

// Allocation is tied to a FiveTuple and relays traffic
// use CreateAllocation and GetAllocation to operate
type Allocation struct {
	RelayAddr           net.Addr
	Protocol            Protocol
	TurnSocket          net.PacketConn
	RelaySocket         net.PacketConn
	fiveTuple           *FiveTuple
	permissionsLock     sync.RWMutex
	permissions         map[string]*Permission
	channelBindingsLock sync.RWMutex
	channelBindings     []*ChannelBind
	lifetimeTimer       *time.Timer
	closed              chan interface{}
	log                 logging.LeveledLogger
}

func addr2IPFingerprint(addr net.Addr) string {
	switch a := addr.(type) {
	case *net.UDPAddr:
		return a.IP.String()
	case *net.TCPAddr: // Do we really need this case?
		return a.IP.String()
	}
	return "" // shoud never happen
}

// NewAllocation creates a new instance of NewAllocation.
func NewAllocation(turnSocket net.PacketConn, fiveTuple *FiveTuple, log logging.LeveledLogger) *Allocation {
	return &Allocation{
		TurnSocket:  turnSocket,
		fiveTuple:   fiveTuple,
		permissions: make(map[string]*Permission, 64),
		closed:      make(chan interface{}),
		log:         log,
	}
}

// GetPermission gets the Permission from the allocation
func (a *Allocation) GetPermission(addr net.Addr) *Permission {
	a.permissionsLock.RLock()
	defer a.permissionsLock.RUnlock()

	fingerprint := addr2IPFingerprint(addr)
	return a.permissions[fingerprint]
}

// AddPermission adds a new permission to the allocation
func (a *Allocation) AddPermission(p *Permission) {
	fingerprint := addr2IPFingerprint(p.Addr)

	a.permissionsLock.RLock()
	existedPermission, ok := a.permissions[fingerprint]
	a.permissionsLock.RUnlock()
	if ok {
		existedPermission.refresh(permissionTimeout)
		return
	}

	p.allocation = a
	a.permissionsLock.Lock()
	a.permissions[fingerprint] = p
	a.permissionsLock.Unlock()

	p.start(permissionTimeout)
}

// RemovePermission removes the net.Addr's fingerprint from the allocation's permissions
func (a *Allocation) RemovePermission(addr net.Addr) {
	a.permissionsLock.Lock()
	defer a.permissionsLock.Unlock()
	fingerprint := addr2IPFingerprint(addr)
	delete(a.permissions, fingerprint)
}

// AddChannelBind adds a new ChannelBind to the allocation, it also updates the
// permissions needed for this ChannelBind
func (a *Allocation) AddChannelBind(c *ChannelBind, lifetime time.Duration) error {
	// Check that this channel id isn't bound to another transport address, and
	// that this transport address isn't bound to another channel id.
	channelByID := a.GetChannelByID(c.ID)
	channelByPeer := a.GetChannelByAddr(c.Peer)
	if channelByID != channelByPeer {
		return errors.Errorf("You cannot use the same channel number with different peer")
	}

	// Add or refresh this channel.
	if channelByID == nil {
		a.channelBindingsLock.Lock()
		defer a.channelBindingsLock.Unlock()

		c.allocation = a
		a.channelBindings = append(a.channelBindings, c)
		c.start(lifetime)

		// Channel binds also refresh permissions.
		a.AddPermission(NewPermission(c.Peer, a.log))
	} else {
		channelByID.refresh(lifetime)

		// Channel binds also refresh permissions.
		a.AddPermission(NewPermission(channelByID.Peer, a.log))
	}

	return nil
}

// RemoveChannelBind removes the ChannelBind from this allocation by id
func (a *Allocation) RemoveChannelBind(id turn.ChannelNumber) bool {
	a.channelBindingsLock.Lock()
	defer a.channelBindingsLock.Unlock()

	for i := len(a.channelBindings) - 1; i >= 0; i-- {
		if a.channelBindings[i].ID == id {
			a.channelBindings = append(a.channelBindings[:i], a.channelBindings[i+1:]...)
			return true
		}
	}

	return false
}

// GetChannelByID gets the ChannelBind from this allocation by id
func (a *Allocation) GetChannelByID(id turn.ChannelNumber) *ChannelBind {
	a.channelBindingsLock.RLock()
	defer a.channelBindingsLock.RUnlock()
	for _, cb := range a.channelBindings {
		if cb.ID == id {
			return cb
		}
	}
	return nil
}

// GetChannelByAddr gets the ChannelBind from this allocation by net.Addr
func (a *Allocation) GetChannelByAddr(addr net.Addr) *ChannelBind {
	a.channelBindingsLock.RLock()
	defer a.channelBindingsLock.RUnlock()
	for _, cb := range a.channelBindings {
		if ipnet.AddrEqual(cb.Peer, addr) {
			return cb
		}
	}
	return nil
}

// Refresh updates the allocations lifetime
func (a *Allocation) Refresh(lifetime time.Duration) {
	if lifetime == 0 {
		if !a.lifetimeTimer.Stop() {
			a.log.Errorf("Failed to stop allocation timer for %v", a.fiveTuple)
		}
		return
	}
	if !a.lifetimeTimer.Reset(lifetime) {
		a.log.Errorf("Failed to reset allocation timer for %v", a.fiveTuple)
	}
}

// Close closes the allocation
func (a *Allocation) Close() error {
	select {
	case <-a.closed:
		return nil
	default:
	}
	close(a.closed)

	a.lifetimeTimer.Stop()

	a.permissionsLock.RLock()
	for _, p := range a.permissions {
		p.lifetimeTimer.Stop()
	}
	a.permissionsLock.RUnlock()

	a.channelBindingsLock.RLock()
	for _, c := range a.channelBindings {
		c.lifetimeTimer.Stop()
	}
	a.channelBindingsLock.RUnlock()

	return a.RelaySocket.Close()
}

//  https://tools.ietf.org/html/rfc5766#section-10.3
//  When the server receives a UDP datagram at a currently allocated
//  relayed transport address, the server looks up the allocation
//  associated with the relayed transport address.  The server then
//  checks to see whether the set of permissions for the allocation allow
//  the relaying of the UDP datagram as described in Section 8.
//
//  If relaying is permitted, then the server checks if there is a
//  channel bound to the peer that sent the UDP datagram (see
//  Section 11).  If a channel is bound, then processing proceeds as
//  described in Section 11.7.
//
//  If relaying is permitted but no channel is bound to the peer, then
//  the server forms and sends a Data indication.  The Data indication
//  MUST contain both an XOR-PEER-ADDRESS and a DATA attribute.  The DATA
//  attribute is set to the value of the 'data octets' field from the
//  datagram, and the XOR-PEER-ADDRESS attribute is set to the source
//  transport address of the received UDP datagram.  The Data indication
//  is then sent on the 5-tuple associated with the allocation.

const rtpMTU = 1500

func (a *Allocation) packetHandler(m *Manager) {
	buffer := make([]byte, rtpMTU)

	for {
		n, srcAddr, err := a.RelaySocket.ReadFrom(buffer)
		if err != nil {
			m.DeleteAllocation(a.fiveTuple)
			return
		}

		a.log.Debugf("relay socket %s received %d bytes from %s",
			a.RelaySocket.LocalAddr().String(),
			n,
			srcAddr.String())

		if channel := a.GetChannelByAddr(a.RelaySocket.LocalAddr()); channel != nil {
			channelData := make([]byte, 4)
			binary.BigEndian.PutUint16(channelData[0:], uint16(channel.ID))
			binary.BigEndian.PutUint16(channelData[2:], uint16(n))
			channelData = append(channelData, buffer[:n]...)

			if _, err = a.TurnSocket.WriteTo(channelData, a.fiveTuple.SrcAddr); err != nil {
				a.log.Errorf("Failed to send ChannelData from allocation %v %v", srcAddr, err)
			}
		} else if p := a.GetPermission(srcAddr); p != nil {
			udpAddr := srcAddr.(*net.UDPAddr)
			peerAddressAttr := turn.PeerAddress{IP: udpAddr.IP, Port: udpAddr.Port}
			dataAttr := turn.Data(buffer[:n])

			msg, err := stun.Build(stun.TransactionID, stun.NewType(stun.MethodData, stun.ClassIndication), peerAddressAttr, dataAttr)
			if err != nil {
				a.log.Errorf("Failed to send DataIndication from allocation %v %v", srcAddr, err)
			}
			a.log.Debugf("relaying message to client at %s", a.fiveTuple.SrcAddr.String())
			if _, err = a.TurnSocket.WriteTo(msg.Raw, a.fiveTuple.SrcAddr); err != nil {
				a.log.Errorf("Failed to send DataIndication from allocation %v %v", srcAddr, err)
			}

		} else {
			a.log.Errorf("Packet unhandled in relay src %v", srcAddr)
		}
	}

}
<|endoftext|>


<commit_before>package ultimate

import (
	"context"
	"fmt"
	"log"
	"reflect"
	"sync"
	"time"

	"github.com/fatih/color"
	"github.com/hellodudu/Ultimate/global"
	world_message "github.com/hellodudu/Ultimate/proto"
)

var ArenaMatchSectionNum int32 = 8 // arena section num

type Arena struct {
	mapRecord     map[int64]*world_message.ArenaRecord // all player's arena record
	listMatchPool []map[int64]struct{}                 // 8 level match pool
	listRecReq    map[int64]struct{}                   // list to request player arena record
	chMatchWait   chan int64                           // match wait player channel
	endTime       int32                                `sql:"arena_end_time"`

	ctx      context.Context
	cancel   context.CancelFunc
	mu       sync.Mutex
	chDBInit chan struct{}
}

func NewArena(ctx context.Context) (*Arena, error) {
	arena := &Arena{
		mapRecord:     make(map[int64]*world_message.ArenaRecord),
		listMatchPool: make([]map[int64]struct{}, ArenaMatchSectionNum),
		listRecReq:    make(map[int64]struct{}, 1000),
		chMatchWait:   make(chan int64, 1000),
		chDBInit:      make(chan struct{}, 1),
	}

	for n := 0; int32(n) < ArenaMatchSectionNum; n++ {
		arena.listMatchPool[n] = make(map[int64]struct{})
	}

	arena.ctx, arena.cancel = context.WithCancel(ctx)

	return arena, nil
}

func GetSectionIndexByScore(score int32) int32 {
	if score <= 1000 {
		return 0
	} else if score <= 1100 {
		return 1
	} else if score <= 1200 {
		return 2
	} else if score <= 1300 {
		return 3
	} else if score <= 1400 {
		return 4
	} else if score <= 1500 {
		return 5
	} else if score <= 1600 {
		return 6
	} else {
		return 7
	}
}

func (arena *Arena) Stop() {
	arena.cancel()
}

func (arena *Arena) Run() {
	<-arena.chDBInit

	for {
		select {
		// context canceled
		case <-arena.ctx.Done():
			log.Println(color.RedString("arena context done!"))
			return

		// matching request
		case id := <-arena.chMatchWait:
			log.Println(color.CyanString("player:", id, " start arena matching!"))
			arena.UpdateMatching(id)

		default:
			t := time.Now()
			arena.UpdateRecordRequest()
			d := time.Since(t)
			time.Sleep(time.Millisecond - d)

		}
	}
}

func (arena *Arena) LoadFromDB() {
	f, ok := reflect.TypeOf(*arena).FieldByName("endTime")
	if !ok {
		log.Println(color.YellowString("cannot find arena's endTime field!"))
		return
	}

	query := fmt.Sprintf("select %s from global where id = %d", f.Tag.Get("sql"), global.UltimateID)

	rows, err := Instance().dbMgr.Query(query)
	if err != nil {
		log.Println(color.YellowString("cannot load arena's endTime from dbMgr!"))
		return
	}

	for rows.Next() {
		if err := rows.Scan(&arena.endTime); err != nil {
			log.Println(color.YellowString("arena load from db failed:", err))
		}
		log.Println(color.CyanString("arena load from db success:", arena.endTime))
	}

	arena.chDBInit <- struct{}{}

}

func (arena *Arena) UpdateMatching(id int64) {
	// not enough targets
	if len(arena.mapRecord) < 2 {
		return
	}

	// get player arena section
	srcRec, ok := arena.mapRecord[id]
	if !ok {
		log.Println(color.YellowString("cannot find player:", id, " 's arena record yet!"))
		return
	}

	secIdx := GetSectionIndexByScore(srcRec.ArenaScore)
	for k := range arena.listMatchPool[secIdx] {
		// peek self then continue
		if k == id {
			continue
		}

		dstRec, ok := arena.mapRecord[k]
		if !ok {
			continue
		}

		info := Instance().GetGameMgr().GetPlayerInfoByID(k)
		if info == nil {
			continue
		}

		if world := Instance().GetWorldSession().GetWorldByID(info.ServerId); world != nil {
			msg := &world_message.MUW_ArenaStartBattle{
				AttackId:     k,
				TargetRecord: dstRec,
			}
			world.SendProtoMessage(msg)
		}
	}
}

func (arena *Arena) UpdateRecordRequest() {
	for k := range arena.listRecReq {
		if _, ok := arena.mapRecord[k]; ok {
			arena.chMatchWait <- k

			arena.mu.Lock()
			delete(arena.listRecReq, k)
			arena.mu.Unlock()
		}
	}
}

func (arena *Arena) Matching(w *World, playerID int64) {
	_, ok := arena.mapRecord[playerID]
	if ok {
		// add to match request
		arena.chMatchWait <- playerID

	} else {
		// request player record
		msg := &world_message.MUW_ArenaAddRecord{
			PlayerId: playerID,
		}

		w.SendProtoMessage(msg)

		// add to record request list
		arena.mu.Lock()
		arena.listRecReq[playerID] = struct{}{}
		arena.mu.Unlock()
	}
}

func (arena *Arena) AddRecord(rec *world_message.ArenaRecord) {
	if _, ok := arena.mapRecord[rec.PlayerId]; ok {
		return
	}

	arena.mu.Lock()

	// add to record
	arena.mapRecord[rec.PlayerId] = rec

	// add to matching cache
	index := GetSectionIndexByScore(rec.ArenaScore)
	arena.listMatchPool[index][rec.PlayerId] = struct{}{}

	arena.mu.Unlock()

}

func (arena *Arena) ReorderRecord(rec *world_message.ArenaRecord, preSection, newSection int32) {
	arena.mu.Lock()

	delete(arena.listMatchPool[preSection], rec.PlayerId)
	arena.listMatchPool[newSection][rec.PlayerId] = struct{}{}

	arena.mu.Unlock()
}

func (arena *Arena) BattleResult(atkID int64, tarID int64, win bool) {
	log.Println(color.CyanString("arena battle result:", atkID, tarID, win))

	atkRec, ok := arena.mapRecord[atkID]
	if !ok {
		log.Println(color.YellowString("arena battle result return without record:", atkID, tarID, win))
		return
	}

	preSection := GetSectionIndexByScore(atkRec.ArenaScore)
	atkRec.ArenaScore += 10
	newSection := GetSectionIndexByScore(atkRec.ArenaScore)
	if preSection != newSection {
		arena.ReorderRecord(atkRec, preSection, newSection)
	}
}
<commit_msg>arena end time refresh when init<commit_after>package ultimate

import (
	"context"
	"fmt"
	"log"
	"reflect"
	"sync"
	"time"

	"github.com/fatih/color"
	"github.com/hellodudu/Ultimate/global"
	world_message "github.com/hellodudu/Ultimate/proto"
)

var ArenaMatchSectionNum int32 = 8 // arena section num

type Arena struct {
	mapRecord     map[int64]*world_message.ArenaRecord // all player's arena record
	listMatchPool []map[int64]struct{}                 // 8 level match pool
	listRecReq    map[int64]struct{}                   // list to request player arena record
	chMatchWait   chan int64                           // match wait player channel
	endTime       int32                                `sql:"arena_end_time"`

	ctx      context.Context
	cancel   context.CancelFunc
	mu       sync.Mutex
	chDBInit chan struct{}
}

func NewArena(ctx context.Context) (*Arena, error) {
	arena := &Arena{
		mapRecord:     make(map[int64]*world_message.ArenaRecord),
		listMatchPool: make([]map[int64]struct{}, ArenaMatchSectionNum),
		listRecReq:    make(map[int64]struct{}, 1000),
		chMatchWait:   make(chan int64, 1000),
		chDBInit:      make(chan struct{}, 1),
	}

	for n := 0; int32(n) < ArenaMatchSectionNum; n++ {
		arena.listMatchPool[n] = make(map[int64]struct{})
	}

	arena.ctx, arena.cancel = context.WithCancel(ctx)

	return arena, nil
}

func GetSectionIndexByScore(score int32) int32 {
	if score <= 1000 {
		return 0
	} else if score <= 1100 {
		return 1
	} else if score <= 1200 {
		return 2
	} else if score <= 1300 {
		return 3
	} else if score <= 1400 {
		return 4
	} else if score <= 1500 {
		return 5
	} else if score <= 1600 {
		return 6
	} else {
		return 7
	}
}

func (arena *Arena) Stop() {
	arena.cancel()
}

func (arena *Arena) Run() {
	<-arena.chDBInit

	for {
		select {
		// context canceled
		case <-arena.ctx.Done():
			log.Println(color.RedString("arena context done!"))
			return

		// matching request
		case id := <-arena.chMatchWait:
			log.Println(color.CyanString("player:", id, " start arena matching!"))
			arena.UpdateMatching(id)

		default:
			t := time.Now()
			arena.UpdateRecordRequest()
			d := time.Since(t)
			time.Sleep(time.Millisecond - d)

		}
	}
}

func (arena *Arena) LoadFromDB() {
	f, ok := reflect.TypeOf(*arena).FieldByName("endTime")
	if !ok {
		log.Println(color.YellowString("cannot find arena's endTime field!"))
		return
	}

	query := fmt.Sprintf("select %s from global where id = %d", f.Tag.Get("sql"), global.UltimateID)

	rows, err := Instance().dbMgr.Query(query)
	if err != nil {
		log.Println(color.YellowString("cannot load arena's endTime from dbMgr!"))
		return
	}

	for rows.Next() {
		if err := rows.Scan(&arena.endTime); err != nil {
			log.Println(color.YellowString("arena load from db failed:", err))
		}
		log.Println(color.CyanString("arena load from db success:", arena.endTime))
	}

	// if arena endtime was expired, set a new endtime one month later
	if int32(time.Now().Unix()) > arena.endTime {
		arena.endTime = int32(time.Now().Add(time.Hour * 24 * 30).Unix())
		query := fmt.Sprintf("update global set arena_end_time = %d", arena.endTime)
		Instance().dbMgr.Exec(query)
	}

	arena.chDBInit <- struct{}{}
}

func (arena *Arena) UpdateMatching(id int64) {
	// not enough targets
	if len(arena.mapRecord) < 2 {
		return
	}

	// get player arena section
	srcRec, ok := arena.mapRecord[id]
	if !ok {
		log.Println(color.YellowString("cannot find player:", id, " 's arena record yet!"))
		return
	}

	secIdx := GetSectionIndexByScore(srcRec.ArenaScore)
	for k := range arena.listMatchPool[secIdx] {
		// peek self then continue
		if k == id {
			continue
		}

		dstRec, ok := arena.mapRecord[k]
		if !ok {
			continue
		}

		info := Instance().GetGameMgr().GetPlayerInfoByID(k)
		if info == nil {
			continue
		}

		if world := Instance().GetWorldSession().GetWorldByID(info.ServerId); world != nil {
			msg := &world_message.MUW_ArenaStartBattle{
				AttackId:     k,
				TargetRecord: dstRec,
			}
			world.SendProtoMessage(msg)
		}
	}
}

func (arena *Arena) UpdateRecordRequest() {
	for k := range arena.listRecReq {
		if _, ok := arena.mapRecord[k]; ok {
			arena.chMatchWait <- k

			arena.mu.Lock()
			delete(arena.listRecReq, k)
			arena.mu.Unlock()
		}
	}
}

func (arena *Arena) Matching(w *World, playerID int64) {
	_, ok := arena.mapRecord[playerID]
	if ok {
		// add to match request
		arena.chMatchWait <- playerID

	} else {
		// request player record
		msg := &world_message.MUW_ArenaAddRecord{
			PlayerId: playerID,
		}

		w.SendProtoMessage(msg)

		// add to record request list
		arena.mu.Lock()
		arena.listRecReq[playerID] = struct{}{}
		arena.mu.Unlock()
	}
}

func (arena *Arena) AddRecord(rec *world_message.ArenaRecord) {
	if _, ok := arena.mapRecord[rec.PlayerId]; ok {
		return
	}

	arena.mu.Lock()

	// add to record
	arena.mapRecord[rec.PlayerId] = rec

	// add to matching cache
	index := GetSectionIndexByScore(rec.ArenaScore)
	arena.listMatchPool[index][rec.PlayerId] = struct{}{}

	arena.mu.Unlock()

}

func (arena *Arena) ReorderRecord(rec *world_message.ArenaRecord, preSection, newSection int32) {
	arena.mu.Lock()

	delete(arena.listMatchPool[preSection], rec.PlayerId)
	arena.listMatchPool[newSection][rec.PlayerId] = struct{}{}

	arena.mu.Unlock()
}

func (arena *Arena) BattleResult(atkID int64, tarID int64, win bool) {
	log.Println(color.CyanString("arena battle result:", atkID, tarID, win))

	atkRec, ok := arena.mapRecord[atkID]
	if !ok {
		log.Println(color.YellowString("arena battle result return without record:", atkID, tarID, win))
		return
	}

	preSection := GetSectionIndexByScore(atkRec.ArenaScore)
	atkRec.ArenaScore += 10
	newSection := GetSectionIndexByScore(atkRec.ArenaScore)
	if preSection != newSection {
		arena.ReorderRecord(atkRec, preSection, newSection)
	}
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/bouk/httprouter"
	"github.com/influxdata/chronograf"
	idgen "github.com/influxdata/chronograf/id"
)

const (
	// DefaultWidth is used if not specified
	DefaultWidth = 4
	// DefaultHeight is used if not specified
	DefaultHeight = 4
)

type dashboardCellLinks struct {
	Self string `json:"self"` // Self link mapping to this resource
}

type dashboardCellResponse struct {
	chronograf.DashboardCell
	Links dashboardCellLinks `json:"links"`
}

func newCellResponse(dID chronograf.DashboardID, cell chronograf.DashboardCell) dashboardCellResponse {
	base := "/chronograf/v1/dashboards"
	if cell.Queries == nil {
		cell.Queries = []chronograf.DashboardQuery{}
	}
	if cell.CellColors == nil {
		cell.CellColors = []chronograf.CellColor{}
	}

	// Copy to handle race condition
	newAxes := make(map[string]chronograf.Axis, len(cell.Axes))
	for k, v := range cell.Axes {
		newAxes[k] = v
	}

	// ensure x, y, and y2 axes always returned
	for _, lbl := range []string{"x", "y", "y2"} {
		if _, found := newAxes[lbl]; !found {
			newAxes[lbl] = chronograf.Axis{
				Bounds: []string{"", ""},
			}
		}
	}
	cell.Axes = newAxes

	return dashboardCellResponse{
		DashboardCell: cell,
		Links: dashboardCellLinks{
			Self: fmt.Sprintf("%s/%d/cells/%s", base, dID, cell.ID),
		},
	}
}

func newCellResponses(dID chronograf.DashboardID, dcells []chronograf.DashboardCell) []dashboardCellResponse {
	cells := make([]dashboardCellResponse, len(dcells))
	for i, cell := range dcells {
		cells[i] = newCellResponse(dID, cell)
	}
	return cells
}

// ValidDashboardCellRequest verifies that the dashboard cells have a query and
// have the correct axes specified
func ValidDashboardCellRequest(c *chronograf.DashboardCell) error {
	if c == nil {
		return fmt.Errorf("Chronograf dashboard cell was nil")
	}

	CorrectWidthHeight(c)
	for _, q := range c.Queries {
		if err := ValidateQueryConfig(&q.QueryConfig); err != nil {
			return err
		}
	}
	MoveTimeShift(c)
	err := HasCorrectAxes(c)
	if err != nil {
		return err
	}
	if err = HasCorrectColors(c); err != nil {
		return err
	}
	return HasCorrectLegend(c)
}

// HasCorrectAxes verifies that only permitted axes exist within a DashboardCell
func HasCorrectAxes(c *chronograf.DashboardCell) error {
	for label, axis := range c.Axes {
		if !oneOf(label, "x", "y", "y2") {
			return chronograf.ErrInvalidAxis
		}

		if !oneOf(axis.Scale, "linear", "log", "") {
			return chronograf.ErrInvalidAxis
		}

		if !oneOf(axis.Base, "10", "2", "") {
			return chronograf.ErrInvalidAxis
		}
	}

	return nil
}

// HasCorrectColors verifies that the format of each color is correct
func HasCorrectColors(c *chronograf.DashboardCell) error {
	for _, color := range c.CellColors {
		if !oneOf(color.Type, "max", "min", "threshold", "text", "background") {
			return chronograf.ErrInvalidColorType
		}
		if len(color.Hex) != 7 {
			return chronograf.ErrInvalidColor
		}
	}
	return nil
}

// HasCorrectLegend verifies that the format of the legend is correct
func HasCorrectLegend(c *chronograf.DashboardCell) error {
	// No legend set
	if c.Legend.Type == "" && c.Legend.Orientation == "" {
		return nil
	}

	if c.Legend.Type == "" || c.Legend.Orientation == "" {
		return chronograf.ErrInvalidLegend
	}
	if !oneOf(c.Legend.Orientation, "top", "bottom", "right", "left") {
		return chronograf.ErrInvalidLegendOrient
	}

	// Remember! if we add other types, update ErrInvalidLegendType
	if !oneOf(c.Legend.Type, "static") {
		return chronograf.ErrInvalidLegendType
	}
	return nil
}

// oneOf reports whether a provided string is a member of a variadic list of
// valid options
func oneOf(prop string, validOpts ...string) bool {
	for _, valid := range validOpts {
		if prop == valid {
			return true
		}
	}
	return false
}

// CorrectWidthHeight changes the cell to have at least the
// minimum width and height
func CorrectWidthHeight(c *chronograf.DashboardCell) {
	if c.W < 1 {
		c.W = DefaultWidth
	}
	if c.H < 1 {
		c.H = DefaultHeight
	}
}

// MoveTimeShift moves TimeShift from the QueryConfig to the DashboardQuery
func MoveTimeShift(c *chronograf.DashboardCell) {
	for i, query := range c.Queries {
		query.Shifts = query.QueryConfig.Shifts
		c.Queries[i] = query
	}
}

// AddQueryConfig updates a cell by converting InfluxQL into queryconfigs
// If influxql cannot be represented by a full query config, then, the
// query config's raw text is set to the command.
func AddQueryConfig(c *chronograf.DashboardCell) {
	for i, q := range c.Queries {
		qc := ToQueryConfig(q.Command)
		qc.Shifts = append([]chronograf.TimeShift(nil), q.Shifts...)
		q.Shifts = nil
		q.QueryConfig = qc
		c.Queries[i] = q
	}
}

// DashboardCells returns all cells from a dashboard within the store
func (s *Service) DashboardCells(w http.ResponseWriter, r *http.Request) {
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	ctx := r.Context()
	e, err := s.Store.Dashboards(ctx).Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	boards := newDashboardResponse(e)
	cells := boards.Cells
	encodeJSON(w, http.StatusOK, cells, s.Logger)
}

// NewDashboardCell adds a cell to an existing dashboard
func (s *Service) NewDashboardCell(w http.ResponseWriter, r *http.Request) {
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	ctx := r.Context()
	dash, err := s.Store.Dashboards(ctx).Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}
	var cell chronograf.DashboardCell
	if err := json.NewDecoder(r.Body).Decode(&cell); err != nil {
		invalidJSON(w, s.Logger)
		return
	}

	if err := ValidDashboardCellRequest(&cell); err != nil {
		invalidData(w, err, s.Logger)
		return
	}

	ids := &idgen.UUID{}
	cid, err := ids.Generate()
	if err != nil {
		msg := fmt.Sprintf("Error creating cell ID of dashboard %d: %v", id, err)
		Error(w, http.StatusInternalServerError, msg, s.Logger)
		return
	}
	cell.ID = cid

	dash.Cells = append(dash.Cells, cell)
	if err := s.Store.Dashboards(ctx).Update(ctx, dash); err != nil {
		msg := fmt.Sprintf("Error adding cell %s to dashboard %d: %v", cid, id, err)
		Error(w, http.StatusInternalServerError, msg, s.Logger)
		return
	}

	boards := newDashboardResponse(dash)
	for _, cell := range boards.Cells {
		if cell.ID == cid {
			encodeJSON(w, http.StatusOK, cell, s.Logger)
			return
		}
	}
}

// DashboardCellID gets a specific cell from an existing dashboard
func (s *Service) DashboardCellID(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	dash, err := s.Store.Dashboards(ctx).Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	boards := newDashboardResponse(dash)
	cid := httprouter.GetParamFromContext(ctx, "cid")
	for _, cell := range boards.Cells {
		if cell.ID == cid {
			encodeJSON(w, http.StatusOK, cell, s.Logger)
			return
		}
	}
	notFound(w, id, s.Logger)
}

// RemoveDashboardCell removes a specific cell from an existing dashboard
func (s *Service) RemoveDashboardCell(w http.ResponseWriter, r *http.Request) {
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	ctx := r.Context()
	dash, err := s.Store.Dashboards(ctx).Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	cid := httprouter.GetParamFromContext(ctx, "cid")
	cellid := -1
	for i, cell := range dash.Cells {
		if cell.ID == cid {
			cellid = i
			break
		}
	}
	if cellid == -1 {
		notFound(w, id, s.Logger)
		return
	}

	dash.Cells = append(dash.Cells[:cellid], dash.Cells[cellid+1:]...)
	if err := s.Store.Dashboards(ctx).Update(ctx, dash); err != nil {
		msg := fmt.Sprintf("Error removing cell %s from dashboard %d: %v", cid, id, err)
		Error(w, http.StatusInternalServerError, msg, s.Logger)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

// ReplaceDashboardCell replaces a cell entirely within an existing dashboard
func (s *Service) ReplaceDashboardCell(w http.ResponseWriter, r *http.Request) {
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	ctx := r.Context()
	dash, err := s.Store.Dashboards(ctx).Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	cid := httprouter.GetParamFromContext(ctx, "cid")
	cellid := -1
	for i, cell := range dash.Cells {
		if cell.ID == cid {
			cellid = i
			break
		}
	}
	if cellid == -1 {
		notFound(w, cid, s.Logger)
		return
	}

	var cell chronograf.DashboardCell
	if err := json.NewDecoder(r.Body).Decode(&cell); err != nil {
		invalidJSON(w, s.Logger)
		return
	}

	for i, a := range cell.Axes {
		if len(a.Bounds) == 0 {
			a.Bounds = []string{"", ""}
			cell.Axes[i] = a
		}
	}

	if err := ValidDashboardCellRequest(&cell); err != nil {
		invalidData(w, err, s.Logger)
		return
	}
	cell.ID = cid

	dash.Cells[cellid] = cell
	if err := s.Store.Dashboards(ctx).Update(ctx, dash); err != nil {
		msg := fmt.Sprintf("Error updating cell %s in dashboard %d: %v", cid, id, err)
		Error(w, http.StatusInternalServerError, msg, s.Logger)
		return
	}

	res := newCellResponse(dash.ID, cell)
	encodeJSON(w, http.StatusOK, res, s.Logger)
}
<commit_msg>Add "Scale" type to accepted cell color types<commit_after>package server

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/bouk/httprouter"
	"github.com/influxdata/chronograf"
	idgen "github.com/influxdata/chronograf/id"
)

const (
	// DefaultWidth is used if not specified
	DefaultWidth = 4
	// DefaultHeight is used if not specified
	DefaultHeight = 4
)

type dashboardCellLinks struct {
	Self string `json:"self"` // Self link mapping to this resource
}

type dashboardCellResponse struct {
	chronograf.DashboardCell
	Links dashboardCellLinks `json:"links"`
}

func newCellResponse(dID chronograf.DashboardID, cell chronograf.DashboardCell) dashboardCellResponse {
	base := "/chronograf/v1/dashboards"
	if cell.Queries == nil {
		cell.Queries = []chronograf.DashboardQuery{}
	}
	if cell.CellColors == nil {
		cell.CellColors = []chronograf.CellColor{}
	}

	// Copy to handle race condition
	newAxes := make(map[string]chronograf.Axis, len(cell.Axes))
	for k, v := range cell.Axes {
		newAxes[k] = v
	}

	// ensure x, y, and y2 axes always returned
	for _, lbl := range []string{"x", "y", "y2"} {
		if _, found := newAxes[lbl]; !found {
			newAxes[lbl] = chronograf.Axis{
				Bounds: []string{"", ""},
			}
		}
	}
	cell.Axes = newAxes

	return dashboardCellResponse{
		DashboardCell: cell,
		Links: dashboardCellLinks{
			Self: fmt.Sprintf("%s/%d/cells/%s", base, dID, cell.ID),
		},
	}
}

func newCellResponses(dID chronograf.DashboardID, dcells []chronograf.DashboardCell) []dashboardCellResponse {
	cells := make([]dashboardCellResponse, len(dcells))
	for i, cell := range dcells {
		cells[i] = newCellResponse(dID, cell)
	}
	return cells
}

// ValidDashboardCellRequest verifies that the dashboard cells have a query and
// have the correct axes specified
func ValidDashboardCellRequest(c *chronograf.DashboardCell) error {
	if c == nil {
		return fmt.Errorf("Chronograf dashboard cell was nil")
	}

	CorrectWidthHeight(c)
	for _, q := range c.Queries {
		if err := ValidateQueryConfig(&q.QueryConfig); err != nil {
			return err
		}
	}
	MoveTimeShift(c)
	err := HasCorrectAxes(c)
	if err != nil {
		return err
	}
	if err = HasCorrectColors(c); err != nil {
		return err
	}
	return HasCorrectLegend(c)
}

// HasCorrectAxes verifies that only permitted axes exist within a DashboardCell
func HasCorrectAxes(c *chronograf.DashboardCell) error {
	for label, axis := range c.Axes {
		if !oneOf(label, "x", "y", "y2") {
			return chronograf.ErrInvalidAxis
		}

		if !oneOf(axis.Scale, "linear", "log", "") {
			return chronograf.ErrInvalidAxis
		}

		if !oneOf(axis.Base, "10", "2", "") {
			return chronograf.ErrInvalidAxis
		}
	}

	return nil
}

// HasCorrectColors verifies that the format of each color is correct
func HasCorrectColors(c *chronograf.DashboardCell) error {
	for _, color := range c.CellColors {
		if !oneOf(color.Type, "max", "min", "threshold", "text", "background", "scale") {
			return chronograf.ErrInvalidColorType
		}
		if len(color.Hex) != 7 {
			return chronograf.ErrInvalidColor
		}
	}
	return nil
}

// HasCorrectLegend verifies that the format of the legend is correct
func HasCorrectLegend(c *chronograf.DashboardCell) error {
	// No legend set
	if c.Legend.Type == "" && c.Legend.Orientation == "" {
		return nil
	}

	if c.Legend.Type == "" || c.Legend.Orientation == "" {
		return chronograf.ErrInvalidLegend
	}
	if !oneOf(c.Legend.Orientation, "top", "bottom", "right", "left") {
		return chronograf.ErrInvalidLegendOrient
	}

	// Remember! if we add other types, update ErrInvalidLegendType
	if !oneOf(c.Legend.Type, "static") {
		return chronograf.ErrInvalidLegendType
	}
	return nil
}

// oneOf reports whether a provided string is a member of a variadic list of
// valid options
func oneOf(prop string, validOpts ...string) bool {
	for _, valid := range validOpts {
		if prop == valid {
			return true
		}
	}
	return false
}

// CorrectWidthHeight changes the cell to have at least the
// minimum width and height
func CorrectWidthHeight(c *chronograf.DashboardCell) {
	if c.W < 1 {
		c.W = DefaultWidth
	}
	if c.H < 1 {
		c.H = DefaultHeight
	}
}

// MoveTimeShift moves TimeShift from the QueryConfig to the DashboardQuery
func MoveTimeShift(c *chronograf.DashboardCell) {
	for i, query := range c.Queries {
		query.Shifts = query.QueryConfig.Shifts
		c.Queries[i] = query
	}
}

// AddQueryConfig updates a cell by converting InfluxQL into queryconfigs
// If influxql cannot be represented by a full query config, then, the
// query config's raw text is set to the command.
func AddQueryConfig(c *chronograf.DashboardCell) {
	for i, q := range c.Queries {
		qc := ToQueryConfig(q.Command)
		qc.Shifts = append([]chronograf.TimeShift(nil), q.Shifts...)
		q.Shifts = nil
		q.QueryConfig = qc
		c.Queries[i] = q
	}
}

// DashboardCells returns all cells from a dashboard within the store
func (s *Service) DashboardCells(w http.ResponseWriter, r *http.Request) {
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	ctx := r.Context()
	e, err := s.Store.Dashboards(ctx).Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	boards := newDashboardResponse(e)
	cells := boards.Cells
	encodeJSON(w, http.StatusOK, cells, s.Logger)
}

// NewDashboardCell adds a cell to an existing dashboard
func (s *Service) NewDashboardCell(w http.ResponseWriter, r *http.Request) {
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	ctx := r.Context()
	dash, err := s.Store.Dashboards(ctx).Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}
	var cell chronograf.DashboardCell
	if err := json.NewDecoder(r.Body).Decode(&cell); err != nil {
		invalidJSON(w, s.Logger)
		return
	}

	if err := ValidDashboardCellRequest(&cell); err != nil {
		invalidData(w, err, s.Logger)
		return
	}

	ids := &idgen.UUID{}
	cid, err := ids.Generate()
	if err != nil {
		msg := fmt.Sprintf("Error creating cell ID of dashboard %d: %v", id, err)
		Error(w, http.StatusInternalServerError, msg, s.Logger)
		return
	}
	cell.ID = cid

	dash.Cells = append(dash.Cells, cell)
	if err := s.Store.Dashboards(ctx).Update(ctx, dash); err != nil {
		msg := fmt.Sprintf("Error adding cell %s to dashboard %d: %v", cid, id, err)
		Error(w, http.StatusInternalServerError, msg, s.Logger)
		return
	}

	boards := newDashboardResponse(dash)
	for _, cell := range boards.Cells {
		if cell.ID == cid {
			encodeJSON(w, http.StatusOK, cell, s.Logger)
			return
		}
	}
}

// DashboardCellID gets a specific cell from an existing dashboard
func (s *Service) DashboardCellID(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	dash, err := s.Store.Dashboards(ctx).Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	boards := newDashboardResponse(dash)
	cid := httprouter.GetParamFromContext(ctx, "cid")
	for _, cell := range boards.Cells {
		if cell.ID == cid {
			encodeJSON(w, http.StatusOK, cell, s.Logger)
			return
		}
	}
	notFound(w, id, s.Logger)
}

// RemoveDashboardCell removes a specific cell from an existing dashboard
func (s *Service) RemoveDashboardCell(w http.ResponseWriter, r *http.Request) {
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	ctx := r.Context()
	dash, err := s.Store.Dashboards(ctx).Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	cid := httprouter.GetParamFromContext(ctx, "cid")
	cellid := -1
	for i, cell := range dash.Cells {
		if cell.ID == cid {
			cellid = i
			break
		}
	}
	if cellid == -1 {
		notFound(w, id, s.Logger)
		return
	}

	dash.Cells = append(dash.Cells[:cellid], dash.Cells[cellid+1:]...)
	if err := s.Store.Dashboards(ctx).Update(ctx, dash); err != nil {
		msg := fmt.Sprintf("Error removing cell %s from dashboard %d: %v", cid, id, err)
		Error(w, http.StatusInternalServerError, msg, s.Logger)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

// ReplaceDashboardCell replaces a cell entirely within an existing dashboard
func (s *Service) ReplaceDashboardCell(w http.ResponseWriter, r *http.Request) {
	id, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), s.Logger)
		return
	}

	ctx := r.Context()
	dash, err := s.Store.Dashboards(ctx).Get(ctx, chronograf.DashboardID(id))
	if err != nil {
		notFound(w, id, s.Logger)
		return
	}

	cid := httprouter.GetParamFromContext(ctx, "cid")
	cellid := -1
	for i, cell := range dash.Cells {
		if cell.ID == cid {
			cellid = i
			break
		}
	}
	if cellid == -1 {
		notFound(w, cid, s.Logger)
		return
	}

	var cell chronograf.DashboardCell
	if err := json.NewDecoder(r.Body).Decode(&cell); err != nil {
		invalidJSON(w, s.Logger)
		return
	}

	for i, a := range cell.Axes {
		if len(a.Bounds) == 0 {
			a.Bounds = []string{"", ""}
			cell.Axes[i] = a
		}
	}

	if err := ValidDashboardCellRequest(&cell); err != nil {
		invalidData(w, err, s.Logger)
		return
	}
	cell.ID = cid

	dash.Cells[cellid] = cell
	if err := s.Store.Dashboards(ctx).Update(ctx, dash); err != nil {
		msg := fmt.Sprintf("Error updating cell %s in dashboard %d: %v", cid, id, err)
		Error(w, http.StatusInternalServerError, msg, s.Logger)
		return
	}

	res := newCellResponse(dash.ID, cell)
	encodeJSON(w, http.StatusOK, res, s.Logger)
}
<|endoftext|>


<commit_before>package loraserver

import (
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/brocaar/loraserver/models"
	"github.com/brocaar/lorawan"
	"github.com/brocaar/lorawan/band"
)

// Band is the ISM band configuration to use
var Band band.Band

// Server represents a LoRaWAN network-server.
type Server struct {
	ctx Context
	wg  sync.WaitGroup
}

// NewServer creates a new server.
func NewServer(ctx Context) *Server {
	return &Server{
		ctx: ctx,
	}
}

// Start starts the server.
func (s *Server) Start() error {
	go func() {
		s.wg.Add(1)
		handleRXPackets(s.ctx)
		s.wg.Done()
	}()
	go func() {
		s.wg.Add(1)
		handleTXPayloads(s.ctx)
		s.wg.Done()
	}()
	return nil
}

// Stop closes the gateway and application backends and waits for the
// server to complete the pending packets and actions.
func (s *Server) Stop() error {
	if err := s.ctx.Gateway.Close(); err != nil {
		return fmt.Errorf("close gateway backend error: %s", err)
	}
	if err := s.ctx.Application.Close(); err != nil {
		return fmt.Errorf("close application backend error: %s", err)
	}

	log.Info("waiting for pending packets to complete")
	s.wg.Wait()
	return nil
}

func handleTXPayloads(ctx Context) {
	var wg sync.WaitGroup
	for txPayload := range ctx.Application.TXPayloadChan() {
		go func(txPayload models.TXPayload) {
			wg.Add(1)
			if err := addTXPayloadToQueue(ctx.RedisPool, txPayload); err != nil {
				log.WithField("dev_eui", txPayload.DevEUI).Errorf("add tx payload to queue error: %s", err)
			}
			wg.Done()
		}(txPayload)
	}
	wg.Wait()
}

func handleRXPackets(ctx Context) {
	var wg sync.WaitGroup
	for rxPacket := range ctx.Gateway.RXPacketChan() {
		go func(rxPacket models.RXPacket) {
			wg.Add(1)
			if err := handleRXPacket(ctx, rxPacket); err != nil {
				data, _ := rxPacket.PHYPayload.MarshalText()
				log.WithField("data_base64", string(data)).Errorf("processing rx packet error: %s", err)
			}
			wg.Done()
		}(rxPacket)
	}
	wg.Wait()
}

func handleRXPacket(ctx Context, rxPacket models.RXPacket) error {
	switch rxPacket.PHYPayload.MHDR.MType {
	case lorawan.JoinRequest:
		return validateAndCollectJoinRequestPacket(ctx, rxPacket)
	case lorawan.UnconfirmedDataUp, lorawan.ConfirmedDataUp:
		return validateAndCollectDataUpRXPacket(ctx, rxPacket)
	default:
		return fmt.Errorf("unknown MType: %v", rxPacket.PHYPayload.MHDR.MType)
	}
}

func validateAndCollectDataUpRXPacket(ctx Context, rxPacket models.RXPacket) error {
	// MACPayload must be of type *lorawan.MACPayload
	macPL, ok := rxPacket.PHYPayload.MACPayload.(*lorawan.MACPayload)
	if !ok {
		return fmt.Errorf("expected *lorawan.MACPayload, got: %T", rxPacket.PHYPayload.MACPayload)
	}

	// get the session data
	ns, err := getNodeSession(ctx.RedisPool, macPL.FHDR.DevAddr)
	if err != nil {
		return err
	}

	// validate and get the full int32 FCnt
	fullFCnt, ok := validateAndGetFullFCntUp(ns, macPL.FHDR.FCnt)
	if !ok {
		log.WithFields(log.Fields{
			"dev_addr":    macPL.FHDR.DevAddr,
			"dev_eui":     ns.DevEUI,
			"packet_fcnt": macPL.FHDR.FCnt,
			"server_fcnt": ns.FCntUp,
		}).Warning("invalid FCnt")
		return errors.New("invalid FCnt or too many dropped frames")
	}
	macPL.FHDR.FCnt = fullFCnt

	// validate MIC
	micOK, err := rxPacket.PHYPayload.ValidateMIC(ns.NwkSKey)
	if err != nil {
		return fmt.Errorf("validate MIC error: %s", err)
	}
	if !micOK {
		return errors.New("invalid MIC")
	}

	if macPL.FPort != nil {
		if *macPL.FPort == 0 {
			// decrypt FRMPayload with NwkSKey when FPort == 0
			if err := rxPacket.PHYPayload.DecryptFRMPayload(ns.NwkSKey); err != nil {
				return fmt.Errorf("decrypt FRMPayload error: %s", err)
			}
		} else {
			if err := rxPacket.PHYPayload.DecryptFRMPayload(ns.AppSKey); err != nil {
				return fmt.Errorf("decrypt FRMPayload error: %s", err)
			}
		}
	}

	return collectAndCallOnce(ctx.RedisPool, rxPacket, func(rxPackets RXPackets) error {
		return handleCollectedDataUpPackets(ctx, rxPackets)
	})
}

func handleCollectedDataUpPackets(ctx Context, rxPackets RXPackets) error {
	if len(rxPackets) == 0 {
		return errors.New("packet collector returned 0 packets")
	}
	rxPacket := rxPackets[0]

	var macs []string
	for _, p := range rxPackets {
		macs = append(macs, p.RXInfo.MAC.String())
	}

	macPL, ok := rxPacket.PHYPayload.MACPayload.(*lorawan.MACPayload)
	if !ok {
		return fmt.Errorf("expected *lorawan.MACPayload, got: %T", rxPacket.PHYPayload.MACPayload)
	}

	ns, err := getNodeSession(ctx.RedisPool, macPL.FHDR.DevAddr)
	if err != nil {
		return err
	}

	log.WithFields(log.Fields{
		"dev_eui":  ns.DevEUI,
		"gw_count": len(rxPackets),
		"gw_macs":  strings.Join(macs, ", "),
		"mtype":    rxPackets[0].PHYPayload.MHDR.MType,
	}).Info("packet(s) collected")

	// send rx info notification to be used by the network-controller
	if err = sendRXInfoNotification(ctx, ns.AppEUI, ns.DevEUI, rxPackets); err != nil {
		return fmt.Errorf("send rx info notification error: %s", err)
	}

	if macPL.FPort != nil {
		if *macPL.FPort == 0 {
			log.Warn("todo: implement FPort == 0 packets")
		} else {
			var data []byte

			// it is possible that the FRMPayload is empty, in this case only
			// the FPort will be send
			if len(macPL.FRMPayload) == 1 {
				dataPL, ok := macPL.FRMPayload[0].(*lorawan.DataPayload)
				if !ok {
					return errors.New("FRMPayload must be of type *lorawan.DataPayload")
				}
				data = dataPL.Bytes
			}

			err = ctx.Application.Send(ns.DevEUI, ns.AppEUI, models.RXPayload{
				DevEUI:       ns.DevEUI,
				GatewayCount: len(rxPackets),
				FPort:        *macPL.FPort,
				RSSI:         rxPacket.RXInfo.RSSI,
				Data:         data,
			})
			if err != nil {
				return fmt.Errorf("send rx payload to application error: %s", err)
			}
		}
	}

	// sync counter with that of the device
	ns.FCntUp = macPL.FHDR.FCnt
	if err := saveNodeSession(ctx.RedisPool, ns); err != nil {
		return err
	}

	// handle downlink (ACK)
	time.Sleep(CollectDataDownWait)
	if err := handleDataDownReply(ctx, rxPacket, ns); err != nil {
		return fmt.Errorf("handling downlink data for node %s failed: %s", ns.DevEUI, err)
	}

	return nil
}

func handleDataDownReply(ctx Context, rxPacket models.RXPacket, ns models.NodeSession) error {
	macPL, ok := rxPacket.PHYPayload.MACPayload.(*lorawan.MACPayload)
	if !ok {
		return fmt.Errorf("expected *lorawan.MACPayload, got: %T", rxPacket.PHYPayload.MACPayload)
	}

	// the last payload was received by the node
	if macPL.FHDR.FCtrl.ACK {
		txPayload, err := clearInProcessTXPayload(ctx.RedisPool, ns.DevEUI)
		if err != nil {
			return err
		}
		ns.FCntDown++
		if err = saveNodeSession(ctx.RedisPool, ns); err != nil {
			return err
		}
		if txPayload != nil {
			err = ctx.Application.SendNotification(ns.DevEUI, ns.AppEUI, models.ACKNotificationType, models.ACKNotification{
				Reference: txPayload.Reference,
				DevEUI:    ns.DevEUI,
			})
			if err != nil {
				return err
			}
		}
	}

	// check if there are payloads pending in the queue
	txPayload, remaining, err := getTXPayloadAndRemainingFromQueue(ctx.RedisPool, ns.DevEUI)
	if err != nil {
		return err
	}

	// nothing pending in the queue and no need to ACK RXPacket
	if rxPacket.PHYPayload.MHDR.MType != lorawan.ConfirmedDataUp && txPayload == nil {
		return nil
	}

	// get TX DR
	uplinkDR, err := Band.GetDataRate(rxPacket.RXInfo.DataRate)
	if err != nil {
		return err
	}
	// get TX channel
	uplinkChannel, err := Band.GetChannel(rxPacket.RXInfo.Frequency, uplinkDR)
	if err != nil {
		return err
	}
	// get RX1 channel
	rx1Channel := Band.GetRX1Channel(uplinkChannel)
	// get RX1 DR
	rx1DR := Band.RX1DataRate[uplinkDR][ns.RX1DROffset]
	// get rx delay
	rxDelay := Band.ReceiveDelay1
	if ns.RXDelay > 0 {
		rxDelay = time.Duration(ns.RXDelay) * time.Second
	}

	phy := lorawan.PHYPayload{
		MHDR: lorawan.MHDR{
			MType: lorawan.UnconfirmedDataDown,
			Major: lorawan.LoRaWANR1,
		},
	}
	macPL = &lorawan.MACPayload{
		FHDR: lorawan.FHDR{
			DevAddr: ns.DevAddr,
			FCtrl: lorawan.FCtrl{
				ACK: rxPacket.PHYPayload.MHDR.MType == lorawan.ConfirmedDataUp, // set ACK to true when received packet needs an ACK
			},
			FCnt: ns.FCntDown,
		},
	}
	phy.MACPayload = macPL

	// add the payload from the queue
	if txPayload != nil {
		// validate the max payload size
		if len(txPayload.Data) > Band.MaxPayloadSize[rx1DR].N {
			// remove the payload from the queue regarding confirmed or not
			if _, err := clearInProcessTXPayload(ctx.RedisPool, ns.DevEUI); err != nil {
				return err
			}

			log.WithFields(log.Fields{
				"dev_eui":             ns.DevEUI,
				"data_rate":           rx1DR,
				"frmpayload_size":     len(txPayload.Data),
				"max_frmpayload_size": Band.MaxPayloadSize[rx1DR].N,
			}).Warning("downlink payload max size exceeded")
			err = ctx.Application.SendNotification(ns.DevEUI, ns.AppEUI, models.ErrorNotificationType, models.ErrorNotification{
				Reference: txPayload.Reference,
				DevEUI:    ns.DevEUI,
				Message:   fmt.Sprintf("downlink payload max size exceeded (dr: %d, allowed: %d, got: %d)", rx1DR, Band.MaxPayloadSize[rx1DR].N, len(txPayload.Data)),
			})
			if err != nil {
				return err
			}
		} else {
			// remove the payload from the queue when not confirmed
			if !txPayload.Confirmed {
				if _, err := clearInProcessTXPayload(ctx.RedisPool, ns.DevEUI); err != nil {
					return err
				}
			}

			macPL.FHDR.FCtrl.FPending = remaining
			if txPayload.Confirmed {
				phy.MHDR.MType = lorawan.ConfirmedDataDown
			}
			macPL.FPort = &txPayload.FPort
			macPL.FRMPayload = []lorawan.Payload{
				&lorawan.DataPayload{Bytes: txPayload.Data},
			}
		}
	}

	// when the payload did not pass the validation and there is no ACK set,
	// there is nothing to send
	if !macPL.FHDR.FCtrl.ACK && len(macPL.FRMPayload) == 0 {
		return nil
	}

	// if there is no payload set, encrypt will just do nothing
	if err := phy.EncryptFRMPayload(ns.AppSKey); err != nil {
		return fmt.Errorf("encrypt FRMPayload error: %s", err)
	}

	if err := phy.SetMIC(ns.NwkSKey); err != nil {
		return fmt.Errorf("set MIC error: %s", err)
	}

	txPacket := models.TXPacket{
		TXInfo: models.TXInfo{
			MAC:       rxPacket.RXInfo.MAC,
			Timestamp: rxPacket.RXInfo.Timestamp + uint32(rxDelay/time.Microsecond),
			Frequency: Band.DownlinkChannels[rx1Channel].Frequency,
			Power:     Band.DefaultTXPower,
			DataRate:  Band.DataRates[rx1DR],
			CodeRate:  rxPacket.RXInfo.CodeRate,
		},
		PHYPayload: phy,
	}

	// window 1
	if err := ctx.Gateway.Send(txPacket); err != nil {
		return fmt.Errorf("send tx packet (rx window 1) to gateway error: %s", err)
	}

	// increment the FCntDown when MType != ConfirmedDataDown. In case of
	// ConfirmedDataDown we increment on ACK.
	if phy.MHDR.MType != lorawan.ConfirmedDataDown {
		ns.FCntDown++
		if err := saveNodeSession(ctx.RedisPool, ns); err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>Use GetRX1DataRateForOffset which performs input validation.<commit_after>package loraserver

import (
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/brocaar/loraserver/models"
	"github.com/brocaar/lorawan"
	"github.com/brocaar/lorawan/band"
)

// Band is the ISM band configuration to use
var Band band.Band

// Server represents a LoRaWAN network-server.
type Server struct {
	ctx Context
	wg  sync.WaitGroup
}

// NewServer creates a new server.
func NewServer(ctx Context) *Server {
	return &Server{
		ctx: ctx,
	}
}

// Start starts the server.
func (s *Server) Start() error {
	go func() {
		s.wg.Add(1)
		handleRXPackets(s.ctx)
		s.wg.Done()
	}()
	go func() {
		s.wg.Add(1)
		handleTXPayloads(s.ctx)
		s.wg.Done()
	}()
	return nil
}

// Stop closes the gateway and application backends and waits for the
// server to complete the pending packets and actions.
func (s *Server) Stop() error {
	if err := s.ctx.Gateway.Close(); err != nil {
		return fmt.Errorf("close gateway backend error: %s", err)
	}
	if err := s.ctx.Application.Close(); err != nil {
		return fmt.Errorf("close application backend error: %s", err)
	}

	log.Info("waiting for pending packets to complete")
	s.wg.Wait()
	return nil
}

func handleTXPayloads(ctx Context) {
	var wg sync.WaitGroup
	for txPayload := range ctx.Application.TXPayloadChan() {
		go func(txPayload models.TXPayload) {
			wg.Add(1)
			if err := addTXPayloadToQueue(ctx.RedisPool, txPayload); err != nil {
				log.WithField("dev_eui", txPayload.DevEUI).Errorf("add tx payload to queue error: %s", err)
			}
			wg.Done()
		}(txPayload)
	}
	wg.Wait()
}

func handleRXPackets(ctx Context) {
	var wg sync.WaitGroup
	for rxPacket := range ctx.Gateway.RXPacketChan() {
		go func(rxPacket models.RXPacket) {
			wg.Add(1)
			if err := handleRXPacket(ctx, rxPacket); err != nil {
				data, _ := rxPacket.PHYPayload.MarshalText()
				log.WithField("data_base64", string(data)).Errorf("processing rx packet error: %s", err)
			}
			wg.Done()
		}(rxPacket)
	}
	wg.Wait()
}

func handleRXPacket(ctx Context, rxPacket models.RXPacket) error {
	switch rxPacket.PHYPayload.MHDR.MType {
	case lorawan.JoinRequest:
		return validateAndCollectJoinRequestPacket(ctx, rxPacket)
	case lorawan.UnconfirmedDataUp, lorawan.ConfirmedDataUp:
		return validateAndCollectDataUpRXPacket(ctx, rxPacket)
	default:
		return fmt.Errorf("unknown MType: %v", rxPacket.PHYPayload.MHDR.MType)
	}
}

func validateAndCollectDataUpRXPacket(ctx Context, rxPacket models.RXPacket) error {
	// MACPayload must be of type *lorawan.MACPayload
	macPL, ok := rxPacket.PHYPayload.MACPayload.(*lorawan.MACPayload)
	if !ok {
		return fmt.Errorf("expected *lorawan.MACPayload, got: %T", rxPacket.PHYPayload.MACPayload)
	}

	// get the session data
	ns, err := getNodeSession(ctx.RedisPool, macPL.FHDR.DevAddr)
	if err != nil {
		return err
	}

	// validate and get the full int32 FCnt
	fullFCnt, ok := validateAndGetFullFCntUp(ns, macPL.FHDR.FCnt)
	if !ok {
		log.WithFields(log.Fields{
			"dev_addr":    macPL.FHDR.DevAddr,
			"dev_eui":     ns.DevEUI,
			"packet_fcnt": macPL.FHDR.FCnt,
			"server_fcnt": ns.FCntUp,
		}).Warning("invalid FCnt")
		return errors.New("invalid FCnt or too many dropped frames")
	}
	macPL.FHDR.FCnt = fullFCnt

	// validate MIC
	micOK, err := rxPacket.PHYPayload.ValidateMIC(ns.NwkSKey)
	if err != nil {
		return fmt.Errorf("validate MIC error: %s", err)
	}
	if !micOK {
		return errors.New("invalid MIC")
	}

	if macPL.FPort != nil {
		if *macPL.FPort == 0 {
			// decrypt FRMPayload with NwkSKey when FPort == 0
			if err := rxPacket.PHYPayload.DecryptFRMPayload(ns.NwkSKey); err != nil {
				return fmt.Errorf("decrypt FRMPayload error: %s", err)
			}
		} else {
			if err := rxPacket.PHYPayload.DecryptFRMPayload(ns.AppSKey); err != nil {
				return fmt.Errorf("decrypt FRMPayload error: %s", err)
			}
		}
	}

	return collectAndCallOnce(ctx.RedisPool, rxPacket, func(rxPackets RXPackets) error {
		return handleCollectedDataUpPackets(ctx, rxPackets)
	})
}

func handleCollectedDataUpPackets(ctx Context, rxPackets RXPackets) error {
	if len(rxPackets) == 0 {
		return errors.New("packet collector returned 0 packets")
	}
	rxPacket := rxPackets[0]

	var macs []string
	for _, p := range rxPackets {
		macs = append(macs, p.RXInfo.MAC.String())
	}

	macPL, ok := rxPacket.PHYPayload.MACPayload.(*lorawan.MACPayload)
	if !ok {
		return fmt.Errorf("expected *lorawan.MACPayload, got: %T", rxPacket.PHYPayload.MACPayload)
	}

	ns, err := getNodeSession(ctx.RedisPool, macPL.FHDR.DevAddr)
	if err != nil {
		return err
	}

	log.WithFields(log.Fields{
		"dev_eui":  ns.DevEUI,
		"gw_count": len(rxPackets),
		"gw_macs":  strings.Join(macs, ", "),
		"mtype":    rxPackets[0].PHYPayload.MHDR.MType,
	}).Info("packet(s) collected")

	// send rx info notification to be used by the network-controller
	if err = sendRXInfoNotification(ctx, ns.AppEUI, ns.DevEUI, rxPackets); err != nil {
		return fmt.Errorf("send rx info notification error: %s", err)
	}

	if macPL.FPort != nil {
		if *macPL.FPort == 0 {
			log.Warn("todo: implement FPort == 0 packets")
		} else {
			var data []byte

			// it is possible that the FRMPayload is empty, in this case only
			// the FPort will be send
			if len(macPL.FRMPayload) == 1 {
				dataPL, ok := macPL.FRMPayload[0].(*lorawan.DataPayload)
				if !ok {
					return errors.New("FRMPayload must be of type *lorawan.DataPayload")
				}
				data = dataPL.Bytes
			}

			err = ctx.Application.Send(ns.DevEUI, ns.AppEUI, models.RXPayload{
				DevEUI:       ns.DevEUI,
				GatewayCount: len(rxPackets),
				FPort:        *macPL.FPort,
				RSSI:         rxPacket.RXInfo.RSSI,
				Data:         data,
			})
			if err != nil {
				return fmt.Errorf("send rx payload to application error: %s", err)
			}
		}
	}

	// sync counter with that of the device
	ns.FCntUp = macPL.FHDR.FCnt
	if err := saveNodeSession(ctx.RedisPool, ns); err != nil {
		return err
	}

	// handle downlink (ACK)
	time.Sleep(CollectDataDownWait)
	if err := handleDataDownReply(ctx, rxPacket, ns); err != nil {
		return fmt.Errorf("handling downlink data for node %s failed: %s", ns.DevEUI, err)
	}

	return nil
}

func handleDataDownReply(ctx Context, rxPacket models.RXPacket, ns models.NodeSession) error {
	macPL, ok := rxPacket.PHYPayload.MACPayload.(*lorawan.MACPayload)
	if !ok {
		return fmt.Errorf("expected *lorawan.MACPayload, got: %T", rxPacket.PHYPayload.MACPayload)
	}

	// the last payload was received by the node
	if macPL.FHDR.FCtrl.ACK {
		txPayload, err := clearInProcessTXPayload(ctx.RedisPool, ns.DevEUI)
		if err != nil {
			return err
		}
		ns.FCntDown++
		if err = saveNodeSession(ctx.RedisPool, ns); err != nil {
			return err
		}
		if txPayload != nil {
			err = ctx.Application.SendNotification(ns.DevEUI, ns.AppEUI, models.ACKNotificationType, models.ACKNotification{
				Reference: txPayload.Reference,
				DevEUI:    ns.DevEUI,
			})
			if err != nil {
				return err
			}
		}
	}

	// check if there are payloads pending in the queue
	txPayload, remaining, err := getTXPayloadAndRemainingFromQueue(ctx.RedisPool, ns.DevEUI)
	if err != nil {
		return err
	}

	// nothing pending in the queue and no need to ACK RXPacket
	if rxPacket.PHYPayload.MHDR.MType != lorawan.ConfirmedDataUp && txPayload == nil {
		return nil
	}

	// get TX DR
	uplinkDR, err := Band.GetDataRate(rxPacket.RXInfo.DataRate)
	if err != nil {
		return err
	}
	// get TX channel
	uplinkChannel, err := Band.GetChannel(rxPacket.RXInfo.Frequency, uplinkDR)
	if err != nil {
		return err
	}
	// get RX1 channel
	rx1Channel := Band.GetRX1Channel(uplinkChannel)
	// get RX1 DR
	rx1DR, err := Band.GetRX1DataRateForOffset(uplinkDR, int(ns.RX1DROffset))
	if err != nil {
		return err
	}
	// get rx delay
	rxDelay := Band.ReceiveDelay1
	if ns.RXDelay > 0 {
		rxDelay = time.Duration(ns.RXDelay) * time.Second
	}

	phy := lorawan.PHYPayload{
		MHDR: lorawan.MHDR{
			MType: lorawan.UnconfirmedDataDown,
			Major: lorawan.LoRaWANR1,
		},
	}
	macPL = &lorawan.MACPayload{
		FHDR: lorawan.FHDR{
			DevAddr: ns.DevAddr,
			FCtrl: lorawan.FCtrl{
				ACK: rxPacket.PHYPayload.MHDR.MType == lorawan.ConfirmedDataUp, // set ACK to true when received packet needs an ACK
			},
			FCnt: ns.FCntDown,
		},
	}
	phy.MACPayload = macPL

	// add the payload from the queue
	if txPayload != nil {
		// validate the max payload size
		if len(txPayload.Data) > Band.MaxPayloadSize[rx1DR].N {
			// remove the payload from the queue regarding confirmed or not
			if _, err := clearInProcessTXPayload(ctx.RedisPool, ns.DevEUI); err != nil {
				return err
			}

			log.WithFields(log.Fields{
				"dev_eui":             ns.DevEUI,
				"data_rate":           rx1DR,
				"frmpayload_size":     len(txPayload.Data),
				"max_frmpayload_size": Band.MaxPayloadSize[rx1DR].N,
			}).Warning("downlink payload max size exceeded")
			err = ctx.Application.SendNotification(ns.DevEUI, ns.AppEUI, models.ErrorNotificationType, models.ErrorNotification{
				Reference: txPayload.Reference,
				DevEUI:    ns.DevEUI,
				Message:   fmt.Sprintf("downlink payload max size exceeded (dr: %d, allowed: %d, got: %d)", rx1DR, Band.MaxPayloadSize[rx1DR].N, len(txPayload.Data)),
			})
			if err != nil {
				return err
			}
		} else {
			// remove the payload from the queue when not confirmed
			if !txPayload.Confirmed {
				if _, err := clearInProcessTXPayload(ctx.RedisPool, ns.DevEUI); err != nil {
					return err
				}
			}

			macPL.FHDR.FCtrl.FPending = remaining
			if txPayload.Confirmed {
				phy.MHDR.MType = lorawan.ConfirmedDataDown
			}
			macPL.FPort = &txPayload.FPort
			macPL.FRMPayload = []lorawan.Payload{
				&lorawan.DataPayload{Bytes: txPayload.Data},
			}
		}
	}

	// when the payload did not pass the validation and there is no ACK set,
	// there is nothing to send
	if !macPL.FHDR.FCtrl.ACK && len(macPL.FRMPayload) == 0 {
		return nil
	}

	// if there is no payload set, encrypt will just do nothing
	if err := phy.EncryptFRMPayload(ns.AppSKey); err != nil {
		return fmt.Errorf("encrypt FRMPayload error: %s", err)
	}

	if err := phy.SetMIC(ns.NwkSKey); err != nil {
		return fmt.Errorf("set MIC error: %s", err)
	}

	txPacket := models.TXPacket{
		TXInfo: models.TXInfo{
			MAC:       rxPacket.RXInfo.MAC,
			Timestamp: rxPacket.RXInfo.Timestamp + uint32(rxDelay/time.Microsecond),
			Frequency: Band.DownlinkChannels[rx1Channel].Frequency,
			Power:     Band.DefaultTXPower,
			DataRate:  Band.DataRates[rx1DR],
			CodeRate:  rxPacket.RXInfo.CodeRate,
		},
		PHYPayload: phy,
	}

	// window 1
	if err := ctx.Gateway.Send(txPacket); err != nil {
		return fmt.Errorf("send tx packet (rx window 1) to gateway error: %s", err)
	}

	// increment the FCntDown when MType != ConfirmedDataDown. In case of
	// ConfirmedDataDown we increment on ACK.
	if phy.MHDR.MType != lorawan.ConfirmedDataDown {
		ns.FCntDown++
		if err := saveNodeSession(ctx.RedisPool, ns); err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

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

	"github.com/coreos/go-systemd/unit"
	"github.com/urfave/cli"
)

const (
	networkBase     = "/etc/systemd/network"
	ethernetService = "wired.service"
	wirelessService = "wireless.service"
)

//Ethernet is the ehternet configuration.
type Ethernet struct {
	Network
}

//Wifi is the wifi configuration.
type Wifi struct {
	Network
	Username string `json:"username"`
	Password string `json:"password"`
}

//UnitFile is an interface for systemd uni file
type UnitFile interface {
	ToSystemdUnit() ([]*unit.UnitOption, error)
}

func (w Wifi) ToSystemdUnit() ([]*unit.UnitOption, error) {
	if w.Interface == "" {
		w.Network.Interface = "wlan0"
	}
	return w.Network.ToSystemdUnit()
}

//CreateSystemdFile creates a file that has systemd unit file content.
func CreateSystemdFile(u UnitFile, filename string, mode os.FileMode, out ...io.Writer) error {
	x, err := u.ToSystemdUnit()
	if err != nil {
		return err
	}
	r := unit.Serialize(x)
	if len(out) > 0 {
		_, err := io.Copy(out[0], r)
		return err
	}
	f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, mode)
	if err != nil {
		return err
	}
	defer func() {
		_ = f.Close()
	}()
	_, err = io.Copy(f, r)
	return err
}

func main() {
	app := cli.NewApp()
	app.Version = "0.1.0"
	app.Name = "fconf"
	app.Usage = "fessbox configuration manager"
	app.Commands = []cli.Command{
		{
			Name:    "ethernet",
			Aliases: []string{"e"},
			Usage:   "configures ethernet with systemd",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "name",
					Usage: "The name of the unit file",
					Value: ethernetService,
				},
				cli.StringFlag{
					Name:  "dir",
					Usage: "The directory in which to write the file",
					Value: networkBase,
				},
			},
			Action: ethernet,
		},
		{
			Name:    "wifi-client",
			Aliases: []string{"w"},
			Usage:   "configures wifi client with systemd",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "name",
					Usage: "The name of the unit file",
					Value: wirelessService,
				},
				cli.StringFlag{
					Name:  "dir",
					Usage: "The directory in which to write the file",
					Value: networkBase,
				},
			},
			Action: wifiClient,
		},
	}
	err := app.Run(os.Args)
	if err != nil {
		log.Fatalf("fconf: %v", err)
	}
}

func ethernet(ctx *cli.Context) error {
	base := ctx.String("dir")
	name := ctx.String("name")
	src := ctx.Args().First()
	if src == "" {
		return errors.New("fconf: missing argument")
	}
	b, err := ioutil.ReadFile(src)
	if err != nil {
		return err
	}
	e := Ethernet{}
	err = json.Unmarshal(b, &e)
	if err != nil {
		return err
	}
	err = checkDir(base)
	if err != nil {
		return err
	}
	filename := filepath.Join(base, name)
	err = CreateSystemdFile(e, filename, 0644)
	if err != nil {
		return err
	}
	fmt.Printf("successful written ethernet configuration to %s \n", filename)
	return nil
}

// Checks if the directory exists. If the directory doesnt exist, this function
// will create the directory with permission 0755.
//
// The directory created will recursively create subdirectory. It will behave
// something like mkdir -p /dir/subdir.
func checkDir(dir string) error {
	_, err := os.Stat(dir)
	if os.IsNotExist(err) {
		err = os.MkdirAll(dir, 07755)
		if err != nil {
			return err
		}
	}
	return nil
}

func wifiClient(ctx *cli.Context) error {
	base := ctx.String("dir")
	name := ctx.String("name")
	src := ctx.Args().First()
	if src == "" {
		return errors.New("fconf: missing argument")
	}
	b, err := ioutil.ReadFile(src)
	if err != nil {
		return err
	}
	e := Wifi{}
	err = json.Unmarshal(b, &e)
	if err != nil {
		return err
	}
	err = checkDir(base)
	if err != nil {
		return err
	}
	filename := filepath.Join(base, name)
	err = CreateSystemdFile(e, filename, 0644)
	if err != nil {
		return err
	}
	fmt.Printf("successful written wifi configuration to %s \n", filename)
	return nil
}
<commit_msg>Set default interfaces for Wifi and ethernet<commit_after>package main

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

	"github.com/coreos/go-systemd/unit"
	"github.com/urfave/cli"
)

const (
	networkBase     = "/etc/systemd/network"
	ethernetService = "wired.service"
	wirelessService = "wireless.service"
)

//Ethernet is the ehternet configuration.
type Ethernet struct {
	Network
}

func (e Ethernet) ToSystemdUnit() ([]*unit.UnitOption, error) {
	if e.Interface == "" {
		e.Interface = "eth0"
	}
	return e.Network.ToSystemdUnit()
}

//Wifi is the wifi configuration.
type Wifi struct {
	Network
	Username string `json:"username"`
	Password string `json:"password"`
}

//UnitFile is an interface for systemd uni file
type UnitFile interface {
	ToSystemdUnit() ([]*unit.UnitOption, error)
}

func (w Wifi) ToSystemdUnit() ([]*unit.UnitOption, error) {
	if w.Interface == "" {
		w.Interface = "wlan0"
	}
	return w.Network.ToSystemdUnit()
}

//CreateSystemdFile creates a file that has systemd unit file content.
func CreateSystemdFile(u UnitFile, filename string, mode os.FileMode, out ...io.Writer) error {
	x, err := u.ToSystemdUnit()
	if err != nil {
		return err
	}
	r := unit.Serialize(x)
	if len(out) > 0 {
		_, err := io.Copy(out[0], r)
		return err
	}
	f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, mode)
	if err != nil {
		return err
	}
	defer func() {
		_ = f.Close()
	}()
	_, err = io.Copy(f, r)
	return err
}

func main() {
	app := cli.NewApp()
	app.Version = "0.1.0"
	app.Name = "fconf"
	app.Usage = "fessbox configuration manager"
	app.Commands = []cli.Command{
		{
			Name:    "ethernet",
			Aliases: []string{"e"},
			Usage:   "configures ethernet with systemd",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "name",
					Usage: "The name of the unit file",
					Value: ethernetService,
				},
				cli.StringFlag{
					Name:  "dir",
					Usage: "The directory in which to write the file",
					Value: networkBase,
				},
			},
			Action: ethernet,
		},
		{
			Name:    "wifi-client",
			Aliases: []string{"w"},
			Usage:   "configures wifi client with systemd",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "name",
					Usage: "The name of the unit file",
					Value: wirelessService,
				},
				cli.StringFlag{
					Name:  "dir",
					Usage: "The directory in which to write the file",
					Value: networkBase,
				},
			},
			Action: wifiClient,
		},
	}
	err := app.Run(os.Args)
	if err != nil {
		log.Fatalf("fconf: %v", err)
	}
}

func ethernet(ctx *cli.Context) error {
	base := ctx.String("dir")
	name := ctx.String("name")
	src := ctx.Args().First()
	if src == "" {
		return errors.New("fconf: missing argument")
	}
	b, err := ioutil.ReadFile(src)
	if err != nil {
		return err
	}
	e := Ethernet{}
	err = json.Unmarshal(b, &e)
	if err != nil {
		return err
	}
	err = checkDir(base)
	if err != nil {
		return err
	}
	filename := filepath.Join(base, name)
	err = CreateSystemdFile(e, filename, 0644)
	if err != nil {
		return err
	}
	fmt.Printf("successful written ethernet configuration to %s \n", filename)
	return nil
}

// Checks if the directory exists. If the directory doesnt exist, this function
// will create the directory with permission 0755.
//
// The directory created will recursively create subdirectory. It will behave
// something like mkdir -p /dir/subdir.
func checkDir(dir string) error {
	_, err := os.Stat(dir)
	if os.IsNotExist(err) {
		err = os.MkdirAll(dir, 07755)
		if err != nil {
			return err
		}
	}
	return nil
}

func wifiClient(ctx *cli.Context) error {
	base := ctx.String("dir")
	name := ctx.String("name")
	src := ctx.Args().First()
	if src == "" {
		return errors.New("fconf: missing argument")
	}
	b, err := ioutil.ReadFile(src)
	if err != nil {
		return err
	}
	e := Wifi{}
	err = json.Unmarshal(b, &e)
	if err != nil {
		return err
	}
	err = checkDir(base)
	if err != nil {
		return err
	}
	filename := filepath.Join(base, name)
	err = CreateSystemdFile(e, filename, 0644)
	if err != nil {
		return err
	}
	fmt.Printf("successful written wifi configuration to %s \n", filename)
	return nil
}
<|endoftext|>


<commit_before>package sqlite3schema

import (
	"database/sql"
	"fmt"
	"math"
	"strings"

	"github.com/gimpldo/ba-prototype-go/sqlschema"
)

// ElementFound = SQLite database schema element found (after a search).
// Schema elements are also known as "schema objects".
type ElementFound struct {
	ElemType sqlschema.ElemTypeCode

	// page number -1 means that the rootpage column was NULL in 'sqlite_master'
	RootPage int32

	Name      string
	TableName string
	CreateSQL string
}

func ReadFromDB(db *sql.DB, namePrefix, nameSuffix string) ([]ElementFound, error) {
	var err error
	err = checkSafeNameChars(namePrefix)
	if err != nil {
		panic(err)
	}
	err = checkSafeNameChars(nameSuffix)
	if err != nil {
		panic(err)
	}
	rows, err := db.Query(
		"SELECT type, rootpage, name, tbl_name, sql FROM sqlite_master WHERE name LIKE ?",
		namePrefix+"%"+nameSuffix)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	return readSQLiteMasterRows(rows)
}

func readSQLiteMasterRows(rows *sql.Rows) ([]ElementFound, error) {
	var result []ElementFound

	for rows.Next() {
		var (
			// From 'https://www.sqlite.org/fileformat.html':
			// "For rows that define views, triggers, and
			// virtual tables, the rootpage column is 0 or NULL."
			//
			// The maximum page number is 2147483646 (2**31 - 2)
			// but there is no standard type 'sql.NullInt32'.
			//
			elemRootPage sql.NullInt64

			// From 'https://www.sqlite.org/fileformat.html':
			// "sqlite_master.sql is NULL for the internal indexes
			// that are automatically created by
			// UNIQUE or PRIMARY KEY constraints."
			//
			elemSQL sql.NullString

			elemTypeText string
			elemFound    ElementFound
		)
		err := rows.Scan(&elemTypeText, &elemRootPage,
			&elemFound.Name, &elemFound.TableName, &elemSQL)
		if err != nil {
			return result, err
		}

		if elemRootPage.Valid {
			if elemRootPage.Int64 < 0 {
				return result, fmt.Errorf(
					"Negative root page number %d (type %q, name %q, tbl_name %q)",
					elemRootPage.Int64, elemTypeText, elemFound.Name, elemFound.TableName)
			}
			if elemRootPage.Int64 >= math.MaxInt32 {
				return result, fmt.Errorf(
					"Root page number too big: %d (type %q, name %q, tbl_name %q)",
					elemRootPage.Int64, elemTypeText, elemFound.Name, elemFound.TableName)
			}
			elemFound.RootPage = int32(elemRootPage.Int64)
		} else {
			// 'sqlite_master.rootpage' is NULL:
			elemFound.RootPage = -1
		}

		if elemSQL.Valid {
			elemFound.CreateSQL = elemSQL.String
		} else {
			// 'sqlite_master.sql' is NULL:
			if !(elemTypeText == "index" &&
				strings.HasPrefix(elemFound.Name, "sqlite_autoindex_")) {
				return result, fmt.Errorf(
					"Unexpected NULL in 'sqlite_master.sql' (type %q, name %q, tbl_name %q, rootpage %d)",
					elemTypeText, elemFound.Name, elemFound.TableName, elemFound.RootPage)
			}
		}

		switch elemTypeText {
		case "table":
			elemFound.ElemType = sqlschema.TableElem
			if elemFound.Name != elemFound.TableName {
				return result, fmt.Errorf(
					"Bad table: name %q != tbl_name %q",
					elemFound.Name, elemFound.TableName)
			}
		case "index":
			elemFound.ElemType = sqlschema.IndexElem
		case "view":
			elemFound.ElemType = sqlschema.ViewElem
			if elemFound.Name != elemFound.TableName {
				return result, fmt.Errorf(
					"Bad view: name %q != tbl_name %q",
					elemFound.Name, elemFound.TableName)
			}
		case "trigger":
			elemFound.ElemType = sqlschema.TriggerElem
		default:
			return result, fmt.Errorf(
				"Unexpected element type %q (name %q, tbl_name %q)",
				elemTypeText, elemFound.Name, elemFound.TableName)
		}

		result = append(result, elemFound)
	}
	if err := rows.Err(); err != nil {
		return result, err
	}

	return result, nil
}

// Note: Right now, uppercase letters are not accepted because this software
// does not plan to use them in database schema element/object names,
// and the official SQLite style (D. Richard Hipp's) does not use them.
// It would be an easy change to allow uppercase, should the need appear.
func checkSafeNameChars(nameFragment string) error {
	for i, ch := range nameFragment {
		switch {
		case 'a' <= ch && ch <= 'z':
		case '0' <= ch && ch <= '9':
		case ch == '_':
		default:
			return fmt.Errorf("Name fragment is not safe: %x at %d.", ch, i)
		}
		// Problem: case ch == '_' will be misinterpreted as wildcard.
		// We want to have underscores in name fragments but
		// underscore has special meaning in 'LIKE'.
		// Not (yet) worth the trouble to quote the underscores with
		// an 'ESCAPE' character specified in the query. TODO!
	}
	return nil
}
<commit_msg>Transform 'checkSafeNameChars' function into an escaping func, 'checkEscapeForLike', with checking as a preliminary step. This is what is needed in this package ('.../util/sqlite3schema').<commit_after>package sqlite3schema

import (
	"bytes"
	"database/sql"
	"fmt"
	"math"
	"strings"

	"github.com/gimpldo/ba-prototype-go/sqlschema"
	"github.com/pkg/errors"
)

// ElementFound = SQLite database schema element found (after a search).
// Schema elements are also known as "schema objects".
type ElementFound struct {
	ElemType sqlschema.ElemTypeCode

	// page number -1 means that the rootpage column was NULL in 'sqlite_master'
	RootPage int32

	Name      string
	TableName string
	CreateSQL string
}

func ReadFromDB(db *sql.DB, namePrefix, nameSuffix string) ([]ElementFound, error) {
	// Must be the same character as in the SQL 'ESCAPE' clause below:
	const escapeCharForLike = '!'

	var err error

	escPrefix, err := checkEscapeForLike(namePrefix, escapeCharForLike)
	if err != nil {
		return nil, errors.Wrapf(err, "Prefix not acceptable")
	}
	escSuffix, err := checkEscapeForLike(nameSuffix, escapeCharForLike)
	if err != nil {
		return nil, errors.Wrapf(err, "Suffix not acceptable")
	}

	rows, err := db.Query(
		"SELECT type, rootpage, name, tbl_name, sql FROM sqlite_master WHERE name LIKE ? ESCAPE '!'",
		escPrefix+"%"+escSuffix)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	return readSQLiteMasterRows(rows)
}

func readSQLiteMasterRows(rows *sql.Rows) ([]ElementFound, error) {
	var result []ElementFound

	for rows.Next() {
		var (
			// From 'https://www.sqlite.org/fileformat.html':
			// "For rows that define views, triggers, and
			// virtual tables, the rootpage column is 0 or NULL."
			//
			// The maximum page number is 2147483646 (2**31 - 2)
			// but there is no standard type 'sql.NullInt32'.
			//
			elemRootPage sql.NullInt64

			// From 'https://www.sqlite.org/fileformat.html':
			// "sqlite_master.sql is NULL for the internal indexes
			// that are automatically created by
			// UNIQUE or PRIMARY KEY constraints."
			//
			elemSQL sql.NullString

			elemTypeText string
			elemFound    ElementFound
		)
		err := rows.Scan(&elemTypeText, &elemRootPage,
			&elemFound.Name, &elemFound.TableName, &elemSQL)
		if err != nil {
			return result, err
		}

		if elemRootPage.Valid {
			if elemRootPage.Int64 < 0 {
				return result, fmt.Errorf(
					"Negative root page number %d (type %q, name %q, tbl_name %q)",
					elemRootPage.Int64, elemTypeText, elemFound.Name, elemFound.TableName)
			}
			if elemRootPage.Int64 >= math.MaxInt32 {
				return result, fmt.Errorf(
					"Root page number too big: %d (type %q, name %q, tbl_name %q)",
					elemRootPage.Int64, elemTypeText, elemFound.Name, elemFound.TableName)
			}
			elemFound.RootPage = int32(elemRootPage.Int64)
		} else {
			// 'sqlite_master.rootpage' is NULL:
			elemFound.RootPage = -1
		}

		if elemSQL.Valid {
			elemFound.CreateSQL = elemSQL.String
		} else {
			// 'sqlite_master.sql' is NULL:
			if !(elemTypeText == "index" &&
				strings.HasPrefix(elemFound.Name, "sqlite_autoindex_")) {
				return result, fmt.Errorf(
					"Unexpected NULL in 'sqlite_master.sql' (type %q, name %q, tbl_name %q, rootpage %d)",
					elemTypeText, elemFound.Name, elemFound.TableName, elemFound.RootPage)
			}
		}

		switch elemTypeText {
		case "table":
			elemFound.ElemType = sqlschema.TableElem
			if elemFound.Name != elemFound.TableName {
				return result, fmt.Errorf(
					"Bad table: name %q != tbl_name %q",
					elemFound.Name, elemFound.TableName)
			}
		case "index":
			elemFound.ElemType = sqlschema.IndexElem
		case "view":
			elemFound.ElemType = sqlschema.ViewElem
			if elemFound.Name != elemFound.TableName {
				return result, fmt.Errorf(
					"Bad view: name %q != tbl_name %q",
					elemFound.Name, elemFound.TableName)
			}
		case "trigger":
			elemFound.ElemType = sqlschema.TriggerElem
		default:
			return result, fmt.Errorf(
				"Unexpected element type %q (name %q, tbl_name %q)",
				elemTypeText, elemFound.Name, elemFound.TableName)
		}

		result = append(result, elemFound)
	}
	if err := rows.Err(); err != nil {
		return result, err
	}

	return result, nil
}

// checkEscapeForLike uses a very strict (whitelist) approach to validation.
//
// Right now, uppercase letters are not accepted because
// there is no plan or need for this software to use uppercase
// in database schema element/object names, and
// the official SQLite style (D. Richard Hipp's) does not use them either.
// It would be an easy change to allow uppercase, should the need appear.
//
func checkEscapeForLike(nameFragment string, escapeChar rune) (escaped string, err error) {
	var buf bytes.Buffer

	for i, ch := range nameFragment {
		switch {
		case ch == escapeChar: // must be first case
			return "", fmt.Errorf("Name fragment contains the escape char (%x) at %d.", ch, i)
			/*
				If we wanted to escape it, it could have been done as follows:
				buf.WriteRune(escapeChar) // first use: to escape (itself, here)
				buf.WriteRune(escapeChar) // second use: to include it in the LIKE

				But we don't need to support the escape char in names, and
				it's not easy to handle reliably an arbitrary escape char that
				could be one of the characters that we accept or reject below!
			*/
		case 'a' <= ch && ch <= 'z':
			buf.WriteRune(ch)
		case '0' <= ch && ch <= '9':
			buf.WriteRune(ch)
		case ch == '_': // Underscore would be misinterpreted as wildcard, escape it:
			buf.WriteRune(escapeChar)
			buf.WriteRune(ch)
		default:
			return "", fmt.Errorf("Name fragment is not safe: %x at %d.", ch, i)
		}
	}

	return buf.String(), nil
}
<|endoftext|>


<commit_before>package version

import (
	"reflect"
	"testing"

	"github.com/rusenask/keel/types"
)

func TestGetVersionFromImageName(t *testing.T) {
	type args struct {
		name string
	}
	tests := []struct {
		name    string
		args    args
		want    *types.Version
		wantErr bool
	}{
		{
			name:    "image",
			args:    args{name: "karolis/webhook-demo:1.4.5"},
			want:    &types.Version{Major: 1, Minor: 4, Patch: 5},
			wantErr: false,
		},
		{
			name:    "image latest",
			args:    args{name: "karolis/webhook-demo:latest"},
			wantErr: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := GetVersionFromImageName(tt.args.name)
			if (err != nil) != tt.wantErr {
				t.Errorf("GetVersionFromImageName() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("GetVersionFromImageName() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestShouldUpdate(t *testing.T) {
	type args struct {
		current *types.Version
		new     *types.Version
		policy  types.PolicyType
	}
	tests := []struct {
		name    string
		args    args
		want    bool
		wantErr bool
	}{
		{
			name: "new minor increase, policy all",
			args: args{
				current: &types.Version{Major: 1, Minor: 4, Patch: 5},
				new:     &types.Version{Major: 1, Minor: 4, Patch: 6},
				policy:  types.PolicyTypeAll,
			},
			want:    true,
			wantErr: false,
		},
		{
			name: "no increase, policy all",
			args: args{
				current: &types.Version{Major: 1, Minor: 4, Patch: 5},
				new:     &types.Version{Major: 1, Minor: 4, Patch: 5},
				policy:  types.PolicyTypeAll,
			},
			want:    false,
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := ShouldUpdate(tt.args.current, tt.args.new, tt.args.policy)
			if (err != nil) != tt.wantErr {
				t.Errorf("ShouldUpdate() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if got != tt.want {
				t.Errorf("ShouldUpdate() = %v, want %v", got, tt.want)
			}
		})
	}
}
<commit_msg>version tests<commit_after>package version

import (
	"reflect"
	"testing"

	"github.com/rusenask/keel/types"
)

func TestGetVersionFromImageName(t *testing.T) {
	type args struct {
		name string
	}
	tests := []struct {
		name    string
		args    args
		want    *types.Version
		wantErr bool
	}{
		{
			name:    "image",
			args:    args{name: "karolis/webhook-demo:1.4.5"},
			want:    &types.Version{Major: 1, Minor: 4, Patch: 5},
			wantErr: false,
		},
		{
			name:    "image latest",
			args:    args{name: "karolis/webhook-demo:latest"},
			wantErr: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := GetVersionFromImageName(tt.args.name)
			if (err != nil) != tt.wantErr {
				t.Errorf("GetVersionFromImageName() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("GetVersionFromImageName() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestShouldUpdate(t *testing.T) {
	type args struct {
		current *types.Version
		new     *types.Version
		policy  types.PolicyType
	}
	tests := []struct {
		name    string
		args    args
		want    bool
		wantErr bool
	}{
		{
			name: "new lower, policy all",
			args: args{
				current: &types.Version{Major: 1, Minor: 4, Patch: 5},
				new:     &types.Version{Major: 1, Minor: 4, Patch: 3},
				policy:  types.PolicyTypeAll,
			},
			want:    false,
			wantErr: false,
		},
		{
			name: "new minor increase, policy all",
			args: args{
				current: &types.Version{Major: 1, Minor: 4, Patch: 5},
				new:     &types.Version{Major: 1, Minor: 4, Patch: 6},
				policy:  types.PolicyTypeAll,
			},
			want:    true,
			wantErr: false,
		},
		{
			name: "no increase, policy all",
			args: args{
				current: &types.Version{Major: 1, Minor: 4, Patch: 5},
				new:     &types.Version{Major: 1, Minor: 4, Patch: 5},
				policy:  types.PolicyTypeAll,
			},
			want:    false,
			wantErr: false,
		},
		{
			name: "minor increase, policy major",
			args: args{
				current: &types.Version{Major: 1, Minor: 4, Patch: 5},
				new:     &types.Version{Major: 1, Minor: 5, Patch: 5},
				policy:  types.PolicyTypeMajor,
			},
			want:    false,
			wantErr: false,
		},
		{
			name: "major increase, policy major",
			args: args{
				current: &types.Version{Major: 1, Minor: 4, Patch: 5},
				new:     &types.Version{Major: 2, Minor: 4, Patch: 5},
				policy:  types.PolicyTypeMajor,
			},
			want:    true,
			wantErr: false,
		},
		{
			name: "patch increase, policy patch",
			args: args{
				current: &types.Version{Major: 1, Minor: 4, Patch: 5},
				new:     &types.Version{Major: 1, Minor: 4, Patch: 6},
				policy:  types.PolicyTypePatch,
			},
			want:    true,
			wantErr: false,
		},
		{
			name: "patch AND major increase, policy patch",
			args: args{
				current: &types.Version{Major: 1, Minor: 4, Patch: 5},
				new:     &types.Version{Major: 2, Minor: 4, Patch: 6},
				policy:  types.PolicyTypePatch,
			},
			want:    false,
			wantErr: false,
		},
		{
			name: "patch same, policy patch",
			args: args{
				current: &types.Version{Major: 1, Minor: 4, Patch: 5},
				new:     &types.Version{Major: 1, Minor: 4, Patch: 5},
				policy:  types.PolicyTypePatch,
			},
			want:    false,
			wantErr: false,
		},
		{
			name: "minor increase, policy minor",
			args: args{
				current: &types.Version{Major: 1, Minor: 4, Patch: 5},
				new:     &types.Version{Major: 1, Minor: 5, Patch: 5},
				policy:  types.PolicyTypeMinor,
			},
			want:    true,
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := ShouldUpdate(tt.args.current, tt.args.new, tt.args.policy)
			if (err != nil) != tt.wantErr {
				t.Errorf("ShouldUpdate() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if got != tt.want {
				t.Errorf("ShouldUpdate() = %v, want %v", got, tt.want)
			}
		})
	}
}
<|endoftext|>