<commit_before>package MQTTg

import (
	"net"
	"strconv"
	"time"
)

type Broker struct {
	MyAddr *net.TCPAddr
	// TODO: check whether not good to use addr as key
	Clients   map[string]*Client //map[clientID]*CLient
	TopicRoot *TopicNode
}

func (self *Broker) Start() error {
	addr, err := GetLocalAddr()
	if err != nil {
		return err
	}
	self.MyAddr = addr
	listener, err := net.ListenTCP("tcp4", addr)
	if err != nil {
		// TODO: use channel to return error
		return err
	}
	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			// TODO: use channel to return error
			EmitError(err)
			continue
		}
		client := NewClient("", nil, 0, nil)
		client.Ct = &Transport{conn}
		client.Broker = self
		client.ReadChan = make(chan Message)
		bc := &BrokerSideClient{client}
		go bc.ReadMessage()
		go ReadLoop(bc, bc.ReadChan)
	}
}

func (self *Broker) DisconnectFromBroker(c *Client) {
	w := c.Will
	if w != nil {
		if w.Retain {
			self.TopicRoot.ApplyRetain(w.Topic, w.QoS, w.Message)
		}
		nodes, _ := self.TopicRoot.GetTopicNodes(w.Topic, true)
		for subscriberID, _ := range nodes[0].Subscribers {
			// TODO: check which qos should be used, Will.QoS or requested QoS
			subscriber, _ := self.Clients[subscriberID]
			subscriber.Publish(w.Topic, w.Message, w.QoS, w.Retain)
		}
	}
	self.disconnectProcessing(c)
}

func (self *Broker) disconnectProcessing(c *Client) (err error) {
	c.KeepAliveTimer.Stop()
	err = c.disconnectProcessing()
	if c.CleanSession {
		delete(self.Clients, c.ID)
	}
	return
}

func (self *Broker) ApplyDummyClientID() string {
	return "DummyClientID:" + strconv.Itoa(len(self.Clients)+1)
}

type BrokerSideClient struct {
	*Client
}

func (self *BrokerSideClient) RunClientTimer() {
	<-self.KeepAliveTimer.C
	EmitError(CLIENT_TIMED_OUT)
	self.Broker.DisconnectFromBroker(self.Client)
	// TODO: logging?
}

func (self *BrokerSideClient) recvConnectMessage(m *ConnectMessage) (err error) {
	// NOTICE: when connection error is sent to client, self.Ct.SendMessage()
	//         should be used for avoiding Isconnecting validation
	if m.Protocol.Name != MQTT_3_1_1.Name {
		// server MAY disconnect
		self.disconnectProcessing()
		return INVALID_PROTOCOL_NAME
	}

	if m.Protocol.Level != MQTT_3_1_1.Level {
		// CHECK: Is false correct?
		err = self.Ct.SendMessage(NewConnackMessage(false, UnacceptableProtocolVersion))
		self.disconnectProcessing()
		return INVALID_PROTOCOL_LEVEL
	}

	c, ok := self.Broker.Clients[m.ClientID]
	if ok && c.IsConnecting {
		// TODO: this might cause problem
		// TODO; which should be disconnected, connecting one? or trying to connect one?
		err = self.Ct.SendMessage(NewConnackMessage(false, IdentifierRejected))
		self.disconnectProcessing()
		return CLIENT_ID_IS_USED_ALREADY
	}
	cleanSession := m.Flags&CleanSession_Flag == CleanSession_Flag
	if ok && !cleanSession {
		self.setPreviousSession(c)
	} else if !cleanSession && len(m.ClientID) == 0 {
		err = self.Ct.SendMessage(NewConnackMessage(false, IdentifierRejected))
		self.disconnectProcessing()
		return CLEANSESSION_MUST_BE_TRUE
	}

	sessionPresent := ok
	if cleanSession || !ok {
		// TODO: need to manage QoS base processing
		self.Duration = time.Duration(float32(m.KeepAlive)*1.5) * time.Second
		if len(m.ClientID) == 0 {
			m.ClientID = self.Broker.ApplyDummyClientID()
		}
		self.ID = m.ClientID
		self.User = m.User
		self.KeepAlive = m.KeepAlive
		self.Will = m.Will
		self.CleanSession = cleanSession
		self.KeepAliveTimer = time.NewTimer(self.Duration)
		sessionPresent = false
	}
	self.Broker.Clients[m.ClientID] = self.Client

	if m.Flags&Will_Flag == Will_Flag {
		self.Will = m.Will
		// TODO: consider QoS and Retain as broker need
	} else {

	}

	if m.KeepAlive != 0 {
		go self.RunClientTimer()
	}
	self.IsConnecting = true
	err = self.SendMessage(NewConnackMessage(sessionPresent, Accepted))
	self.Redelivery()
	return err
}

func (self *BrokerSideClient) recvConnackMessage(m *ConnackMessage) (err error) {
	return INVALID_MESSAGE_CAME
}

func (self *BrokerSideClient) recvPublishMessage(m *PublishMessage) (err error) {
	if m.Dup {
		// re-delivered
	} else {
		// first time delivery
	}

	if m.Retain {
		// store tehe application message to designated topic
		data := string(m.Payload)
		if m.QoS == 0 && len(data) > 0 {
			// TODO: warnning, in this case data cannot be stored.
			// discard retained message
			data = ""
		}
		self.Broker.TopicRoot.ApplyRetain(m.TopicName, m.QoS, data)
	}

	nodes, err := self.Broker.TopicRoot.GetTopicNodes(m.TopicName, true)
	if err != nil {
		return err
	}
	for subscriberID, qos := range nodes[0].Subscribers {
		subscriber, _ := self.Broker.Clients[subscriberID]
		subscriber.Publish(m.TopicName, string(m.Payload), qos, false)
	}

	switch m.QoS {
	// in any case, Dub must be 0
	case 0:
		if m.PacketID != 0 {
			return PACKET_ID_SHOULD_BE_ZERO
		}
	case 1:
		err = self.SendMessage(NewPubackMessage(m.PacketID))
	case 2:
		err = self.SendMessage(NewPubrecMessage(m.PacketID))
	}
	return err
}

func (self *BrokerSideClient) recvPubackMessage(m *PubackMessage) (err error) {
	// acknowledge the sent Publish packet
	if m.PacketID > 0 {
		err = self.AckMessage(m.PacketID)
	}
	return err
}

func (self *BrokerSideClient) recvPubrecMessage(m *PubrecMessage) (err error) {
	// acknowledge the sent Publish packet
	err = self.AckMessage(m.PacketID)
	if err != nil {
		return err
	}
	err = self.SendMessage(NewPubrelMessage(m.PacketID))
	return err
}

func (self *BrokerSideClient) recvPubrelMessage(m *PubrelMessage) (err error) {
	// acknowledge the sent Pubrel packet
	err = self.AckMessage(m.PacketID)
	if err != nil {
		return err
	}
	err = self.SendMessage(NewPubcompMessage(m.PacketID))
	return err
}

func (self *BrokerSideClient) recvPubcompMessage(m *PubcompMessage) (err error) {
	// acknowledge the sent Pubrel packet
	err = self.AckMessage(m.PacketID)
	return err
}

func (self *BrokerSideClient) recvSubscribeMessage(m *SubscribeMessage) (err error) {
	// TODO: check The wild card is permitted
	returnCodes := make([]SubscribeReturnCode, 0)
	for _, subTopic := range m.SubscribeTopics {
		// TODO: need to validate wheter there are same topics or not
		edges, err := self.Broker.TopicRoot.GetTopicNodes(subTopic.Topic, true)
		codes := make([]SubscribeReturnCode, len(edges))
		if err != nil {
			for i, _ := range codes {
				codes[i] = SubscribeFailure
			}
		} else {
			for i, edge := range edges {
				edge.Subscribers[self.ID] = subTopic.QoS
				codes[i] = SubscribeReturnCode(subTopic.QoS)
				self.SubTopics = append(self.SubTopics,
					&SubscribeTopic{SubscribeAck,
						edge.FullPath,
						uint8(subTopic.QoS),
					})
				if len(edge.RetainMessage) > 0 {
					// publish retain
					// TODO: check all arguments
					id, err := self.getUsablePacketID()
					EmitError(err)
					err = self.SendMessage(NewPublishMessage(false, edge.RetainQoS, true,
						edge.FullPath, id, []uint8(edge.RetainMessage)))
					EmitError(err)
					// TODO: error validation
				}
			}
		}
		returnCodes = append(returnCodes, codes...)
	}
	// TODO: check whether the number of return codes are correct?
	err = self.SendMessage(NewSubackMessage(m.PacketID, returnCodes))
	return err
}

func (self *BrokerSideClient) recvSubackMessage(m *SubackMessage) (err error) {
	return INVALID_MESSAGE_CAME
}
func (self *BrokerSideClient) recvUnsubscribeMessage(m *UnsubscribeMessage) (err error) {
	if len(m.TopicNames) == 0 {
		// protocol violation
	}

	for _, name := range m.TopicNames {
		self.Broker.TopicRoot.DeleteSubscriber(self.ID, name)
	}
	// TODO: optimize here
	result := []*SubscribeTopic{}
	for _, t := range self.SubTopics {
		del := false
		for _, name := range m.TopicNames {
			if string(t.Topic) == string(name) {
				del = true
			}
		}
		if !del {
			result = append(result, t)
		}
	}
	self.SubTopics = result
	err = self.SendMessage(NewUnsubackMessage(m.PacketID))
	return err
}
func (self *BrokerSideClient) recvUnsubackMessage(m *UnsubackMessage) (err error) {
	return INVALID_MESSAGE_CAME
}

func (self *BrokerSideClient) recvPingreqMessage(m *PingreqMessage) (err error) {
	// Pingresp
	// TODO: calc elapsed time from previous pingreq.
	//       and store the time to duration of Transport
	err = self.SendMessage(NewPingrespMessage())
	if err != nil {
		return err
	}
	if self.KeepAlive != 0 {
		self.ResetTimer()
		go self.RunClientTimer()
	}
	return err
}

func (self *BrokerSideClient) recvPingrespMessage(m *PingrespMessage) (err error) {
	return INVALID_MESSAGE_CAME
}

func (self *BrokerSideClient) recvDisconnectMessage(m *DisconnectMessage) (err error) {
	self.Broker.disconnectProcessing(self.Client)
	// close the client
	return err
}
<commit_msg>use client's Publish for code improvement<commit_after>package MQTTg

import (
	"net"
	"strconv"
	"time"
)

type Broker struct {
	MyAddr *net.TCPAddr
	// TODO: check whether not good to use addr as key
	Clients   map[string]*Client //map[clientID]*CLient
	TopicRoot *TopicNode
}

func (self *Broker) Start() error {
	addr, err := GetLocalAddr()
	if err != nil {
		return err
	}
	self.MyAddr = addr
	listener, err := net.ListenTCP("tcp4", addr)
	if err != nil {
		// TODO: use channel to return error
		return err
	}
	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			// TODO: use channel to return error
			EmitError(err)
			continue
		}
		client := NewClient("", nil, 0, nil)
		client.Ct = &Transport{conn}
		client.Broker = self
		client.ReadChan = make(chan Message)
		bc := &BrokerSideClient{client}
		go bc.ReadMessage()
		go ReadLoop(bc, bc.ReadChan)
	}
}

func (self *Broker) DisconnectFromBroker(c *Client) {
	w := c.Will
	if w != nil {
		if w.Retain {
			self.TopicRoot.ApplyRetain(w.Topic, w.QoS, w.Message)
		}
		nodes, _ := self.TopicRoot.GetTopicNodes(w.Topic, true)
		for subscriberID, _ := range nodes[0].Subscribers {
			// TODO: check which qos should be used, Will.QoS or requested QoS
			subscriber, _ := self.Clients[subscriberID]
			subscriber.Publish(w.Topic, w.Message, w.QoS, w.Retain)
		}
	}
	self.disconnectProcessing(c)
}

func (self *Broker) disconnectProcessing(c *Client) (err error) {
	c.KeepAliveTimer.Stop()
	err = c.disconnectProcessing()
	if c.CleanSession {
		delete(self.Clients, c.ID)
	}
	return
}

func (self *Broker) ApplyDummyClientID() string {
	return "DummyClientID:" + strconv.Itoa(len(self.Clients)+1)
}

type BrokerSideClient struct {
	*Client
}

func (self *BrokerSideClient) RunClientTimer() {
	<-self.KeepAliveTimer.C
	EmitError(CLIENT_TIMED_OUT)
	self.Broker.DisconnectFromBroker(self.Client)
	// TODO: logging?
}

func (self *BrokerSideClient) recvConnectMessage(m *ConnectMessage) (err error) {
	// NOTICE: when connection error is sent to client, self.Ct.SendMessage()
	//         should be used for avoiding Isconnecting validation
	if m.Protocol.Name != MQTT_3_1_1.Name {
		// server MAY disconnect
		self.disconnectProcessing()
		return INVALID_PROTOCOL_NAME
	}

	if m.Protocol.Level != MQTT_3_1_1.Level {
		// CHECK: Is false correct?
		err = self.Ct.SendMessage(NewConnackMessage(false, UnacceptableProtocolVersion))
		self.disconnectProcessing()
		return INVALID_PROTOCOL_LEVEL
	}

	c, ok := self.Broker.Clients[m.ClientID]
	if ok && c.IsConnecting {
		// TODO: this might cause problem
		// TODO; which should be disconnected, connecting one? or trying to connect one?
		err = self.Ct.SendMessage(NewConnackMessage(false, IdentifierRejected))
		self.disconnectProcessing()
		return CLIENT_ID_IS_USED_ALREADY
	}
	cleanSession := m.Flags&CleanSession_Flag == CleanSession_Flag
	if ok && !cleanSession {
		self.setPreviousSession(c)
	} else if !cleanSession && len(m.ClientID) == 0 {
		err = self.Ct.SendMessage(NewConnackMessage(false, IdentifierRejected))
		self.disconnectProcessing()
		return CLEANSESSION_MUST_BE_TRUE
	}

	sessionPresent := ok
	if cleanSession || !ok {
		// TODO: need to manage QoS base processing
		self.Duration = time.Duration(float32(m.KeepAlive)*1.5) * time.Second
		if len(m.ClientID) == 0 {
			m.ClientID = self.Broker.ApplyDummyClientID()
		}
		self.ID = m.ClientID
		self.User = m.User
		self.KeepAlive = m.KeepAlive
		self.Will = m.Will
		self.CleanSession = cleanSession
		self.KeepAliveTimer = time.NewTimer(self.Duration)
		sessionPresent = false
	}
	self.Broker.Clients[m.ClientID] = self.Client

	if m.Flags&Will_Flag == Will_Flag {
		self.Will = m.Will
		// TODO: consider QoS and Retain as broker need
	} else {

	}

	if m.KeepAlive != 0 {
		go self.RunClientTimer()
	}
	self.IsConnecting = true
	err = self.SendMessage(NewConnackMessage(sessionPresent, Accepted))
	self.Redelivery()
	return err
}

func (self *BrokerSideClient) recvConnackMessage(m *ConnackMessage) (err error) {
	return INVALID_MESSAGE_CAME
}

func (self *BrokerSideClient) recvPublishMessage(m *PublishMessage) (err error) {
	if m.Dup {
		// re-delivered
	} else {
		// first time delivery
	}

	if m.Retain {
		// store tehe application message to designated topic
		data := string(m.Payload)
		if m.QoS == 0 && len(data) > 0 {
			// TODO: warnning, in this case data cannot be stored.
			// discard retained message
			data = ""
		}
		self.Broker.TopicRoot.ApplyRetain(m.TopicName, m.QoS, data)
	}

	nodes, err := self.Broker.TopicRoot.GetTopicNodes(m.TopicName, true)
	if err != nil {
		return err
	}
	for subscriberID, qos := range nodes[0].Subscribers {
		subscriber, _ := self.Broker.Clients[subscriberID]
		subscriber.Publish(m.TopicName, string(m.Payload), qos, false)
	}

	switch m.QoS {
	// in any case, Dub must be 0
	case 0:
		if m.PacketID != 0 {
			return PACKET_ID_SHOULD_BE_ZERO
		}
	case 1:
		err = self.SendMessage(NewPubackMessage(m.PacketID))
	case 2:
		err = self.SendMessage(NewPubrecMessage(m.PacketID))
	}
	return err
}

func (self *BrokerSideClient) recvPubackMessage(m *PubackMessage) (err error) {
	// acknowledge the sent Publish packet
	if m.PacketID > 0 {
		err = self.AckMessage(m.PacketID)
	}
	return err
}

func (self *BrokerSideClient) recvPubrecMessage(m *PubrecMessage) (err error) {
	// acknowledge the sent Publish packet
	err = self.AckMessage(m.PacketID)
	if err != nil {
		return err
	}
	err = self.SendMessage(NewPubrelMessage(m.PacketID))
	return err
}

func (self *BrokerSideClient) recvPubrelMessage(m *PubrelMessage) (err error) {
	// acknowledge the sent Pubrel packet
	err = self.AckMessage(m.PacketID)
	if err != nil {
		return err
	}
	err = self.SendMessage(NewPubcompMessage(m.PacketID))
	return err
}

func (self *BrokerSideClient) recvPubcompMessage(m *PubcompMessage) (err error) {
	// acknowledge the sent Pubrel packet
	err = self.AckMessage(m.PacketID)
	return err
}

func (self *BrokerSideClient) recvSubscribeMessage(m *SubscribeMessage) (err error) {
	// TODO: check The wild card is permitted
	returnCodes := make([]SubscribeReturnCode, 0)
	for _, subTopic := range m.SubscribeTopics {
		// TODO: need to validate wheter there are same topics or not
		edges, err := self.Broker.TopicRoot.GetTopicNodes(subTopic.Topic, true)
		codes := make([]SubscribeReturnCode, len(edges))
		if err != nil {
			for i, _ := range codes {
				codes[i] = SubscribeFailure
			}
		} else {
			for i, edge := range edges {
				edge.Subscribers[self.ID] = subTopic.QoS
				codes[i] = SubscribeReturnCode(subTopic.QoS)
				self.SubTopics = append(self.SubTopics,
					&SubscribeTopic{SubscribeAck,
						edge.FullPath,
						uint8(subTopic.QoS),
					})
				if len(edge.RetainMessage) > 0 {
					// publish retain
					// TODO: check all arguments
					err = self.Publish(edge.FullPath, edge.RetainMessage, edge.RetainQoS, true)
					EmitError(err)
				}
			}
		}
		returnCodes = append(returnCodes, codes...)
	}
	// TODO: check whether the number of return codes are correct?
	err = self.SendMessage(NewSubackMessage(m.PacketID, returnCodes))
	return err
}

func (self *BrokerSideClient) recvSubackMessage(m *SubackMessage) (err error) {
	return INVALID_MESSAGE_CAME
}
func (self *BrokerSideClient) recvUnsubscribeMessage(m *UnsubscribeMessage) (err error) {
	if len(m.TopicNames) == 0 {
		// protocol violation
	}

	for _, name := range m.TopicNames {
		self.Broker.TopicRoot.DeleteSubscriber(self.ID, name)
	}
	// TODO: optimize here
	result := []*SubscribeTopic{}
	for _, t := range self.SubTopics {
		del := false
		for _, name := range m.TopicNames {
			if string(t.Topic) == string(name) {
				del = true
			}
		}
		if !del {
			result = append(result, t)
		}
	}
	self.SubTopics = result
	err = self.SendMessage(NewUnsubackMessage(m.PacketID))
	return err
}
func (self *BrokerSideClient) recvUnsubackMessage(m *UnsubackMessage) (err error) {
	return INVALID_MESSAGE_CAME
}

func (self *BrokerSideClient) recvPingreqMessage(m *PingreqMessage) (err error) {
	// Pingresp
	// TODO: calc elapsed time from previous pingreq.
	//       and store the time to duration of Transport
	err = self.SendMessage(NewPingrespMessage())
	if err != nil {
		return err
	}
	if self.KeepAlive != 0 {
		self.ResetTimer()
		go self.RunClientTimer()
	}
	return err
}

func (self *BrokerSideClient) recvPingrespMessage(m *PingrespMessage) (err error) {
	return INVALID_MESSAGE_CAME
}

func (self *BrokerSideClient) recvDisconnectMessage(m *DisconnectMessage) (err error) {
	self.Broker.disconnectProcessing(self.Client)
	// close the client
	return err
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"log"
	"strings"

	"github.com/dt/thile/gen"
	"github.com/foursquare/gohfile"
)

type ThriftRpcImpl struct {
	*hfile.CollectionSet
}

func (cs *ThriftRpcImpl) GetValuesSingle(req *gen.SingleHFileKeyRequest) (r *gen.SingleHFileKeyResponse, err error) {
	if Settings.debug {
		log.Printf("[GetValuesSingle] %s (%d keys)\n", *req.HfileName, len(req.SortedKeys))
	}
	hfile, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	reader := hfile.GetScanner()
	defer reader.Release()

	if req.PerKeyValueLimit != nil {
		// TODO(davidt) impl
		log.Println("[GetValuesSingle] PerKeyValueLimit. oh well.")
	}

	if req.CountOnly != nil {
		// TODO(davidt) impl
		log.Println("[GetValuesSingle] CountOnly. oh well.")
	}

	res := new(gen.SingleHFileKeyResponse)
	res.Values = make(map[int32][]byte)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		if Settings.debug {
			log.Printf("[GetValuesSingle] key: %s\n", hex.EncodeToString(key))
		}
		if idx > 0 && bytes.Equal(req.SortedKeys[idx-1], key) {
			if prev, ok := res.Values[int32(idx-1)]; ok {
				res.Values[int32(idx)] = prev
				found++
			}
			continue
		}
		value, err, ok := reader.GetFirst(key)
		if err != nil {
			return nil, err
		}
		if ok {
			found++
			res.Values[int32(idx)] = value
		}
	}

	if Settings.debug {
		log.Printf("[GetValuesSingle] %s found %d of %d.\n", *req.HfileName, found, len(req.SortedKeys))
	}
	res.KeyCount = &found
	return res, nil
}

func (cs *ThriftRpcImpl) GetValuesMulti(req *gen.SingleHFileKeyRequest) (r *gen.MultiHFileKeyResponse, err error) {
	if Settings.debug {
		log.Println("[GetValuesMulti]", len(req.SortedKeys))
	}

	hfile, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	reader := hfile.GetScanner()
	defer reader.Release()

	res := new(gen.MultiHFileKeyResponse)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		values, err := reader.GetAll(key)
		if err != nil {
			return nil, err
		}
		if len(values) > 0 {
			found += int32(len(values))
			res.Values[int32(idx)] = values
		}
	}

	res.KeyCount = &found
	return res, nil

}

func (cs *ThriftRpcImpl) GetValuesForPrefixes(req *gen.PrefixRequest) (r *gen.PrefixResponse, err error) {
	res := new(gen.PrefixResponse)
	if reader, err := cs.ReaderFor(*req.HfileName); err != nil {
		return nil, err
	} else {
		i := reader.GetIterator()
		defer i.Release()
		if res.Values, err = i.AllForPrfixes(req.SortedKeys); err != nil {
			return nil, err
		} else {
			return res, nil
		}
	}
}

func (cs *ThriftRpcImpl) GetValuesMultiSplitKeys(req *gen.MultiHFileSplitKeyRequest) (r *gen.KeyToValuesResponse, err error) {
	res := make(map[string][][]byte)
	reader, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	scanner := reader.GetScanner()
	defer scanner.Release()

	for _, parts := range RevProduct(req.SplitKey) {
		// TODO(davidt): avoid allocing concated key by adding split-key search lower down.
		key := bytes.Join(parts, nil)

		if values, err := scanner.GetAll(key); err != nil {
			return nil, err
		} else if len(values) > 0 {
			res[string(key)] = values
		}
	}
	return &gen.KeyToValuesResponse{res}, nil
}

func (cs *ThriftRpcImpl) GetIterator(req *gen.IteratorRequest) (*gen.IteratorResponse, error) {
	// 	HfileName     *string `thrift:"hfileName,1" json:"hfileName"`
	// 	IncludeValues *bool   `thrift:"includeValues,2" json:"includeValues"`
	// 	LastKey       []byte  `thrift:"lastKey,3" json:"lastKey"`
	// 	SkipKeys      *int32  `thrift:"skipKeys,4" json:"skipKeys"`
	// 	ResponseLimit *int32  `thrift:"responseLimit,5" json:"responseLimit"`
	// 	EndKey        []byte  `thrift:"endKey,6" json:"endKey"`
	var err error

	if req.ResponseLimit == nil {
		return nil, fmt.Errorf("Missing limit.")
	}
	limit := int(*req.ResponseLimit)

	reader, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	it := reader.GetIterator()
	defer it.Release()

	remaining := false

	if req.LastKey != nil {
		remaining, err = it.Seek(req.LastKey)
	} else {
		remaining, err = it.Next()
	}

	if err != nil {
		return nil, err
	}

	res := new(gen.IteratorResponse)

	if !remaining {
		return res, nil
	}

	skipKeys := int32(0)
	lastKey := it.Key()

	if toSkip := req.GetSkipKeys(); toSkip > 0 {
		for i := int32(0); i < toSkip && remaining; i++ {
			if bytes.Equal(lastKey, it.Key()) {
				skipKeys = skipKeys + 1
			} else {
				skipKeys = 0
			}

			lastKey = it.Key()

			remaining, err = it.Next()
			if err != nil {
				return nil, err
			}
		}
		if !remaining {
			return res, nil
		}
	}

	if req.EndKey != nil {
		remaining = remaining && !hfile.After(it.Key(), req.EndKey)
	}

	r := make([]*gen.KeyValueItem, 0)
	for i := 0; i < limit && remaining; i++ {
		v := []byte{}
		if req.IncludeValues == nil || *req.IncludeValues {
			v = it.Value()
		}
		r = append(r, &gen.KeyValueItem{it.Key(), v})

		if bytes.Equal(lastKey, it.Key()) {
			skipKeys = skipKeys + 1
		} else {
			skipKeys = 0
		}
		lastKey = it.Key()

		remaining, err = it.Next()
		if err != nil {
			return nil, err
		}
		if req.EndKey != nil {
			remaining = remaining && !hfile.After(it.Key(), req.EndKey)
		}
	}
	return &gen.IteratorResponse{r, lastKey, &skipKeys}, nil
}

func GetCollectionInfo(r *hfile.Reader, keySampleSize int) (*gen.HFileInfo, error) {
	i := new(gen.HFileInfo)
	i.Name = &r.Name
	i.Path = &r.SourcePath
	c := int64(r.EntryCount)
	i.NumElements = &c
	i.FirstKey, _ = r.FirstKey()
	return i, nil
}

func (cs *ThriftRpcImpl) GetInfo(req *gen.InfoRequest) (r []*gen.HFileInfo, err error) {
	require := ""
	if req.IsSetHfileName() {
		require := req.GetHfileName()
		if require != "" && !strings.ContainsRune(require, '/') {
			require = require + "/"
		}
	}

	sample := 0
	if req.IsSetNumRandomKeys() {
		sample = int(*req.NumRandomKeys)
	}

	for name, reader := range cs.Collections {
		if require == "" || strings.HasPrefix(name, require) {
			if i, err := GetCollectionInfo(reader, sample); err != nil {
				return nil, err
			} else {
				r = append(r, i)
			}
		}
	}

	return r, nil
}

func (cs *ThriftRpcImpl) TestTimeout(waitInMillis int32) (r int32, err error) {
	return 0, fmt.Errorf("Not implemented")
}
<commit_msg>disable strict key-order checking for perf<commit_after>package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"log"
	"strings"

	"github.com/dt/thile/gen"
	"github.com/foursquare/gohfile"
)

type ThriftRpcImpl struct {
	*hfile.CollectionSet
}

func (cs *ThriftRpcImpl) GetValuesSingle(req *gen.SingleHFileKeyRequest) (r *gen.SingleHFileKeyResponse, err error) {
	if Settings.debug {
		log.Printf("[GetValuesSingle] %s (%d keys)\n", *req.HfileName, len(req.SortedKeys))
	}
	hfile, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	reader := hfile.GetScanner()
	// TODO: clients should request strict during dev/testing?
	reader.EnforceKeyOrder = false
	defer reader.Release()

	if req.PerKeyValueLimit != nil {
		// TODO(davidt) impl
		log.Println("[GetValuesSingle] PerKeyValueLimit. oh well.")
	}

	if req.CountOnly != nil {
		// TODO(davidt) impl
		log.Println("[GetValuesSingle] CountOnly. oh well.")
	}

	res := new(gen.SingleHFileKeyResponse)
	res.Values = make(map[int32][]byte)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		if Settings.debug {
			log.Printf("[GetValuesSingle] key: %s\n", hex.EncodeToString(key))
		}
		if idx > 0 && bytes.Equal(req.SortedKeys[idx-1], key) {
			if prev, ok := res.Values[int32(idx-1)]; ok {
				res.Values[int32(idx)] = prev
				found++
			}
			continue
		}
		value, err, ok := reader.GetFirst(key)
		if err != nil {
			return nil, err
		}
		if ok {
			found++
			res.Values[int32(idx)] = value
		}
	}

	if Settings.debug {
		log.Printf("[GetValuesSingle] %s found %d of %d.\n", *req.HfileName, found, len(req.SortedKeys))
	}
	res.KeyCount = &found
	return res, nil
}

func (cs *ThriftRpcImpl) GetValuesMulti(req *gen.SingleHFileKeyRequest) (r *gen.MultiHFileKeyResponse, err error) {
	if Settings.debug {
		log.Println("[GetValuesMulti]", len(req.SortedKeys))
	}

	hfile, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	reader := hfile.GetScanner()
	defer reader.Release()

	res := new(gen.MultiHFileKeyResponse)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		values, err := reader.GetAll(key)
		if err != nil {
			return nil, err
		}
		if len(values) > 0 {
			found += int32(len(values))
			res.Values[int32(idx)] = values
		}
	}

	res.KeyCount = &found
	return res, nil

}

func (cs *ThriftRpcImpl) GetValuesForPrefixes(req *gen.PrefixRequest) (r *gen.PrefixResponse, err error) {
	res := new(gen.PrefixResponse)
	if reader, err := cs.ReaderFor(*req.HfileName); err != nil {
		return nil, err
	} else {
		i := reader.GetIterator()
		defer i.Release()
		if res.Values, err = i.AllForPrfixes(req.SortedKeys); err != nil {
			return nil, err
		} else {
			return res, nil
		}
	}
}

func (cs *ThriftRpcImpl) GetValuesMultiSplitKeys(req *gen.MultiHFileSplitKeyRequest) (r *gen.KeyToValuesResponse, err error) {
	res := make(map[string][][]byte)
	reader, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	scanner := reader.GetScanner()
	defer scanner.Release()

	for _, parts := range RevProduct(req.SplitKey) {
		// TODO(davidt): avoid allocing concated key by adding split-key search lower down.
		key := bytes.Join(parts, nil)

		if values, err := scanner.GetAll(key); err != nil {
			return nil, err
		} else if len(values) > 0 {
			res[string(key)] = values
		}
	}
	return &gen.KeyToValuesResponse{res}, nil
}

func (cs *ThriftRpcImpl) GetIterator(req *gen.IteratorRequest) (*gen.IteratorResponse, error) {
	// 	HfileName     *string `thrift:"hfileName,1" json:"hfileName"`
	// 	IncludeValues *bool   `thrift:"includeValues,2" json:"includeValues"`
	// 	LastKey       []byte  `thrift:"lastKey,3" json:"lastKey"`
	// 	SkipKeys      *int32  `thrift:"skipKeys,4" json:"skipKeys"`
	// 	ResponseLimit *int32  `thrift:"responseLimit,5" json:"responseLimit"`
	// 	EndKey        []byte  `thrift:"endKey,6" json:"endKey"`
	var err error

	if req.ResponseLimit == nil {
		return nil, fmt.Errorf("Missing limit.")
	}
	limit := int(*req.ResponseLimit)

	reader, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	it := reader.GetIterator()
	defer it.Release()

	remaining := false

	if req.LastKey != nil {
		remaining, err = it.Seek(req.LastKey)
	} else {
		remaining, err = it.Next()
	}

	if err != nil {
		return nil, err
	}

	res := new(gen.IteratorResponse)

	if !remaining {
		return res, nil
	}

	skipKeys := int32(0)
	lastKey := it.Key()

	if toSkip := req.GetSkipKeys(); toSkip > 0 {
		for i := int32(0); i < toSkip && remaining; i++ {
			if bytes.Equal(lastKey, it.Key()) {
				skipKeys = skipKeys + 1
			} else {
				skipKeys = 0
			}

			lastKey = it.Key()

			remaining, err = it.Next()
			if err != nil {
				return nil, err
			}
		}
		if !remaining {
			return res, nil
		}
	}

	if req.EndKey != nil {
		remaining = remaining && !hfile.After(it.Key(), req.EndKey)
	}

	r := make([]*gen.KeyValueItem, 0)
	for i := 0; i < limit && remaining; i++ {
		v := []byte{}
		if req.IncludeValues == nil || *req.IncludeValues {
			v = it.Value()
		}
		r = append(r, &gen.KeyValueItem{it.Key(), v})

		if bytes.Equal(lastKey, it.Key()) {
			skipKeys = skipKeys + 1
		} else {
			skipKeys = 0
		}
		lastKey = it.Key()

		remaining, err = it.Next()
		if err != nil {
			return nil, err
		}
		if req.EndKey != nil {
			remaining = remaining && !hfile.After(it.Key(), req.EndKey)
		}
	}
	return &gen.IteratorResponse{r, lastKey, &skipKeys}, nil
}

func GetCollectionInfo(r *hfile.Reader, keySampleSize int) (*gen.HFileInfo, error) {
	i := new(gen.HFileInfo)
	i.Name = &r.Name
	i.Path = &r.SourcePath
	c := int64(r.EntryCount)
	i.NumElements = &c
	i.FirstKey, _ = r.FirstKey()
	return i, nil
}

func (cs *ThriftRpcImpl) GetInfo(req *gen.InfoRequest) (r []*gen.HFileInfo, err error) {
	require := ""
	if req.IsSetHfileName() {
		require := req.GetHfileName()
		if require != "" && !strings.ContainsRune(require, '/') {
			require = require + "/"
		}
	}

	sample := 0
	if req.IsSetNumRandomKeys() {
		sample = int(*req.NumRandomKeys)
	}

	for name, reader := range cs.Collections {
		if require == "" || strings.HasPrefix(name, require) {
			if i, err := GetCollectionInfo(reader, sample); err != nil {
				return nil, err
			} else {
				r = append(r, i)
			}
		}
	}

	return r, nil
}

func (cs *ThriftRpcImpl) TestTimeout(waitInMillis int32) (r int32, err error) {
	return 0, fmt.Errorf("Not implemented")
}
<|endoftext|>


<commit_before>/*
  Go Language Raspberry Pi Interface
  (c) Copyright David Thorpe 2016-2020
  All Rights Reserved
  For Licensing and Usage information, please see LICENSE.md
*/

package gopi

import (
	"context"
	"fmt"
	"net"
	"strconv"
	"strings"
	"time"
)

////////////////////////////////////////////////////////////////////////////////
// TYPES

// RPCServiceRecord defines a service which can be registered or discovered
// on the network
type RPCServiceRecord struct {
	Name    string
	Service string
	Host    string
	Port    uint16
	Addrs   []net.IP
	Txt     []string
}

type (
	RPCFlag      uint // RPCFlag is a set of flags modifying behavior
	RPCEventType uint // RPCEventType is an enumeration of event types
)

////////////////////////////////////////////////////////////////////////////////
// INTERFACES

// RPCServiceDiscovery will lookup services and classes of service
type RPCServiceDiscovery interface {
	// Lookup service instances by name
	Lookup(ctx context.Context, service string) ([]RPCServiceRecord, error)

	// Return list of service names
	EnumerateServices(ctx context.Context) ([]string, error)

	// Implements gopi.Unit
	Unit
}

// RPCServiceRegister will register services
type RPCServiceRegister interface {
	// Register service record, and de-register when deadline is exceeded
	Register(ctx context.Context, record RPCServiceRecord) error

	// Implements gopi.Unit
	Unit
}

// RPCServer serves requests for one or more services
type RPCServer interface {
	// Start an RPC server in currently running thread.
	// The method will not return until Stop is called
	// which needs to be done in a different thread
	Start() error

	// Stop RPC server. If halt is true then it immediately
	// ends the server without waiting for current requests to
	// be served
	Stop(halt bool) error

	// Return address the server is bound to, or nil if
	// the server is not running
	Addr() net.Addr

	// Implements gopi.Unit
	Unit
}

// RPCService is a driver which implements all the necessary
// methods to handle remote calls
type RPCService interface {
	// CancelRequests is called by the server to gracefully end any
	// on-going streaming requests, but before the service is shutdown
	CancelRequests() error

	// Implements gopi.Unit
	Unit
}

// RPCClientPool represents a pool of clients and connections to
// remove services
type RPCClientPool interface {
	// Lookup service records by parameter - returns records
	// which match either a service name up to max number of
	// records (or zero for unlimited). Will wait for new records
	// and block until cancelled
	Lookup(ctx context.Context, addr string, max uint) ([]RPCServiceRecord, error)

	// Connect to service by service record
	Connect(RPCServiceRecord, RPCFlag) (RPCClientConn, error)

	// Connect to service by address and port
	ConnectAddr(net.IP, uint16) (RPCClientConn, error)

	// Connect to service by unix socket
	ConnectFifo(string) (RPCClientConn, error)

	// Disconnect from service
	Disconnect(RPCClientConn) error

	// Implements gopi.Unit
	Unit
}

// RPCClientConn implements a single client connection for
// communicating with an RPC server
type RPCClientConn interface {
	// Return address for connection, or nil if not connected
	Addr() net.IP

	// Return port for connection
	Port() uint16

	// Return service names supported by connection
	Services() ([]string, error)

	// Mutex locking to ensure one request at a time
	Lock()
	Unlock()

	// Implements gopi.Unit
	Unit
}

// RPCEvent is emitted on service discovery and server events
type RPCEvent interface {
	// Type of event
	Type() RPCEventType

	// Service record associated with event
	Service() RPCServiceRecord

	// Time-to-live value for event
	TTL() time.Duration

	// Implements gopi.Event
	Event
}

////////////////////////////////////////////////////////////////////////////////
// CONSTANTS

const (
	RPC_FLAG_NONE          RPCFlag = 0
	RPC_FLAG_INET_UDP      RPCFlag = (1 << iota) >> 1 // Use UDP protocol (TCP assumed otherwise)
	RPC_FLAG_INET_V4                                  // Use V4 addressing
	RPC_FLAG_INET_V6                                  // Use V6 addressing
	RPC_FLAG_SERVICE_FIRST                            // Use first service
	RPC_FLAG_SERVICE_ANY                              // Use any service
	RPC_FLAG_MIN           = RPC_FLAG_INET_UDP
	RPC_FLAG_MAX           = RPC_FLAG_SERVICE_ANY
)

const (
	RPC_EVENT_NONE            RPCEventType = iota
	RPC_EVENT_SERVER_STARTED               // RPC Server started
	RPC_EVENT_SERVER_STOPPED               // RPC Server stopped
	RPC_EVENT_SERVICE_ADDED                // Service instance lookup (new)
	RPC_EVENT_SERVICE_UPDATED              // Service instance lookup (updated)
	RPC_EVENT_SERVICE_REMOVED              // Service instance lookup (removed)
	RPC_EVENT_SERVICE_EXPIRED              // Service instance lookup (expired)
	RPC_EVENT_SERVICE_NAME                 // Service name discovered
	RPC_EVENT_SERVICE_RECORD               // Service record lookup
	RPC_EVENT_CLIENT_CONNECTED
	RPC_EVENT_CLIENT_DISCONNECTED
	RPC_EVENT_MAX = RPC_EVENT_CLIENT_DISCONNECTED
)

////////////////////////////////////////////////////////////////////////////////
// STRINGIFY

func (f RPCFlag) String() string {
	str := ""
	if f == 0 {
		return f.FlagString()
	}
	for v := RPC_FLAG_MIN; v <= RPC_FLAG_MAX; v <<= 1 {
		if f&v == v {
			str += "|" + v.FlagString()
		}
	}
	return strings.TrimPrefix(str, "|")
}

func (v RPCFlag) FlagString() string {
	switch v {
	case RPC_FLAG_NONE:
		return "RPC_FLAG_NONE"
	case RPC_FLAG_INET_UDP:
		return "RPC_FLAG_INET_UDP"
	case RPC_FLAG_INET_V4:
		return "RPC_FLAG_INET_V4"
	case RPC_FLAG_INET_V6:
		return "RPC_FLAG_INET_V6"
	case RPC_FLAG_SERVICE_FIRST:
		return "RPC_FLAG_SERVICE_FIRST"
	case RPC_FLAG_SERVICE_ANY:
		return "RPC_FLAG_SERVICE_ANY"
	default:
		return "[?? Invalid PlatformType value]"
	}
}

func (t RPCEventType) String() string {
	switch t {
	case RPC_EVENT_NONE:
		return "RPC_EVENT_NONE"
	case RPC_EVENT_SERVER_STARTED:
		return "RPC_EVENT_SERVER_STARTED"
	case RPC_EVENT_SERVER_STOPPED:
		return "RPC_EVENT_SERVER_STOPPED"
	case RPC_EVENT_SERVICE_ADDED:
		return "RPC_EVENT_SERVICE_ADDED"
	case RPC_EVENT_SERVICE_UPDATED:
		return "RPC_EVENT_SERVICE_UPDATED"
	case RPC_EVENT_SERVICE_REMOVED:
		return "RPC_EVENT_SERVICE_REMOVED"
	case RPC_EVENT_SERVICE_EXPIRED:
		return "RPC_EVENT_SERVICE_EXPIRED"
	case RPC_EVENT_SERVICE_NAME:
		return "RPC_EVENT_SERVICE_NAME"
	case RPC_EVENT_SERVICE_RECORD:
		return "RPC_EVENT_SERVICE_RECORD"
	case RPC_EVENT_CLIENT_CONNECTED:
		return "RPC_EVENT_CLIENT_CONNECTED"
	case RPC_EVENT_CLIENT_DISCONNECTED:
		return "RPC_EVENT_CLIENT_DISCONNECTED"
	default:
		return "[?? Invalid RPCEventType value]"
	}
}

func (this RPCServiceRecord) String() string {
	str := "<RPCServiceRecord name=" + strconv.Quote(this.Name)
	if this.Service != "" {
		str += " service=" + strconv.Quote(this.Service)
	}
	if this.Host != "" {
		str += " host=" + strconv.Quote(this.Host)
	}
	if this.Port != 0 {
		str += " port=" + fmt.Sprint(this.Port)
	}
	if len(this.Addrs) > 0 {
		str += " addrs="
		for _, addr := range this.Addrs {
			str += addr.String() + ","
		}
		str = strings.TrimSuffix(str, ",")
	}
	if len(this.Txt) > 0 {
		str += " txt=" + fmt.Sprint(this.Txt)
	}
	str += ">"
	return str
}
<commit_msg>Updated RPC to change Addr to a string and remove port method<commit_after>/*
  Go Language Raspberry Pi Interface
  (c) Copyright David Thorpe 2016-2020
  All Rights Reserved
  For Licensing and Usage information, please see LICENSE.md
*/

package gopi

import (
	"context"
	"fmt"
	"net"
	"strconv"
	"strings"
	"time"
)

////////////////////////////////////////////////////////////////////////////////
// TYPES

// RPCServiceRecord defines a service which can be registered or discovered
// on the network
type RPCServiceRecord struct {
	Name    string
	Service string
	Host    string
	Port    uint16
	Addrs   []net.IP
	Txt     []string
}

type (
	RPCFlag      uint // RPCFlag is a set of flags modifying behavior
	RPCEventType uint // RPCEventType is an enumeration of event types
)

////////////////////////////////////////////////////////////////////////////////
// INTERFACES

// RPCServiceDiscovery will lookup services and classes of service
type RPCServiceDiscovery interface {
	// Lookup service instances by name
	Lookup(ctx context.Context, service string) ([]RPCServiceRecord, error)

	// Return list of service names
	EnumerateServices(ctx context.Context) ([]string, error)

	// Implements gopi.Unit
	Unit
}

// RPCServiceRegister will register services
type RPCServiceRegister interface {
	// Register service record, and de-register when deadline is exceeded
	Register(ctx context.Context, record RPCServiceRecord) error

	// Implements gopi.Unit
	Unit
}

// RPCServer serves requests for one or more services
type RPCServer interface {
	// Start an RPC server in currently running thread.
	// The method will not return until Stop is called
	// which needs to be done in a different thread
	Start() error

	// Stop RPC server. If halt is true then it immediately
	// ends the server without waiting for current requests to
	// be served
	Stop(halt bool) error

	// Return address the server is bound to, or nil if
	// the server is not running
	Addr() net.Addr

	// Implements gopi.Unit
	Unit
}

// RPCService is a driver which implements all the necessary
// methods to handle remote calls
type RPCService interface {
	// CancelRequests is called by the server to gracefully end any
	// on-going streaming requests, but before the service is shutdown
	CancelRequests() error

	// Implements gopi.Unit
	Unit
}

// RPCClientPool represents a pool of clients and connections to
// remove services
type RPCClientPool interface {
	// Lookup service records by parameter - returns records
	// which match either a service name up to max number of
	// records (or zero for unlimited). Will wait for new records
	// and block until cancelled
	Lookup(ctx context.Context, addr string, max uint) ([]RPCServiceRecord, error)

	// Connect to service by service record
	Connect(RPCServiceRecord, RPCFlag) (RPCClientConn, error)

	// Connect to service by address and port
	ConnectAddr(net.IP, uint16) (RPCClientConn, error)

	// Connect to service by unix socket
	ConnectFifo(string) (RPCClientConn, error)

	// Disconnect from service
	Disconnect(RPCClientConn) error

	// Implements gopi.Unit
	Unit
}

// RPCClientConn implements a single client connection for
// communicating with an RPC server
type RPCClientConn interface {
	// Return connection information, or nil if not connected
	Addr() string

	// Return service names supported by connection
	Services() ([]string, error)

	// Mutex locking to ensure one request at a time
	Lock()
	Unlock()

	// Implements gopi.Unit
	Unit
}

// RPCEvent is emitted on service discovery and server events
type RPCEvent interface {
	// Type of event
	Type() RPCEventType

	// Service record associated with event
	Service() RPCServiceRecord

	// Time-to-live value for event
	TTL() time.Duration

	// Implements gopi.Event
	Event
}

////////////////////////////////////////////////////////////////////////////////
// CONSTANTS

const (
	RPC_FLAG_NONE          RPCFlag = 0
	RPC_FLAG_INET_UDP      RPCFlag = (1 << iota) >> 1 // Use UDP protocol (TCP assumed otherwise)
	RPC_FLAG_INET_V4                                  // Use V4 addressing
	RPC_FLAG_INET_V6                                  // Use V6 addressing
	RPC_FLAG_SERVICE_FIRST                            // Use first service
	RPC_FLAG_SERVICE_ANY                              // Use any service
	RPC_FLAG_MIN           = RPC_FLAG_INET_UDP
	RPC_FLAG_MAX           = RPC_FLAG_SERVICE_ANY
)

const (
	RPC_EVENT_NONE            RPCEventType = iota
	RPC_EVENT_SERVER_STARTED               // RPC Server started
	RPC_EVENT_SERVER_STOPPED               // RPC Server stopped
	RPC_EVENT_SERVICE_ADDED                // Service instance lookup (new)
	RPC_EVENT_SERVICE_UPDATED              // Service instance lookup (updated)
	RPC_EVENT_SERVICE_REMOVED              // Service instance lookup (removed)
	RPC_EVENT_SERVICE_EXPIRED              // Service instance lookup (expired)
	RPC_EVENT_SERVICE_NAME                 // Service name discovered
	RPC_EVENT_SERVICE_RECORD               // Service record lookup
	RPC_EVENT_CLIENT_CONNECTED
	RPC_EVENT_CLIENT_DISCONNECTED
	RPC_EVENT_MAX = RPC_EVENT_CLIENT_DISCONNECTED
)

////////////////////////////////////////////////////////////////////////////////
// STRINGIFY

func (f RPCFlag) String() string {
	str := ""
	if f == 0 {
		return f.FlagString()
	}
	for v := RPC_FLAG_MIN; v <= RPC_FLAG_MAX; v <<= 1 {
		if f&v == v {
			str += "|" + v.FlagString()
		}
	}
	return strings.TrimPrefix(str, "|")
}

func (v RPCFlag) FlagString() string {
	switch v {
	case RPC_FLAG_NONE:
		return "RPC_FLAG_NONE"
	case RPC_FLAG_INET_UDP:
		return "RPC_FLAG_INET_UDP"
	case RPC_FLAG_INET_V4:
		return "RPC_FLAG_INET_V4"
	case RPC_FLAG_INET_V6:
		return "RPC_FLAG_INET_V6"
	case RPC_FLAG_SERVICE_FIRST:
		return "RPC_FLAG_SERVICE_FIRST"
	case RPC_FLAG_SERVICE_ANY:
		return "RPC_FLAG_SERVICE_ANY"
	default:
		return "[?? Invalid PlatformType value]"
	}
}

func (t RPCEventType) String() string {
	switch t {
	case RPC_EVENT_NONE:
		return "RPC_EVENT_NONE"
	case RPC_EVENT_SERVER_STARTED:
		return "RPC_EVENT_SERVER_STARTED"
	case RPC_EVENT_SERVER_STOPPED:
		return "RPC_EVENT_SERVER_STOPPED"
	case RPC_EVENT_SERVICE_ADDED:
		return "RPC_EVENT_SERVICE_ADDED"
	case RPC_EVENT_SERVICE_UPDATED:
		return "RPC_EVENT_SERVICE_UPDATED"
	case RPC_EVENT_SERVICE_REMOVED:
		return "RPC_EVENT_SERVICE_REMOVED"
	case RPC_EVENT_SERVICE_EXPIRED:
		return "RPC_EVENT_SERVICE_EXPIRED"
	case RPC_EVENT_SERVICE_NAME:
		return "RPC_EVENT_SERVICE_NAME"
	case RPC_EVENT_SERVICE_RECORD:
		return "RPC_EVENT_SERVICE_RECORD"
	case RPC_EVENT_CLIENT_CONNECTED:
		return "RPC_EVENT_CLIENT_CONNECTED"
	case RPC_EVENT_CLIENT_DISCONNECTED:
		return "RPC_EVENT_CLIENT_DISCONNECTED"
	default:
		return "[?? Invalid RPCEventType value]"
	}
}

func (this RPCServiceRecord) String() string {
	str := "<RPCServiceRecord name=" + strconv.Quote(this.Name)
	if this.Service != "" {
		str += " service=" + strconv.Quote(this.Service)
	}
	if this.Host != "" {
		str += " host=" + strconv.Quote(this.Host)
	}
	if this.Port != 0 {
		str += " port=" + fmt.Sprint(this.Port)
	}
	if len(this.Addrs) > 0 {
		str += " addrs="
		for _, addr := range this.Addrs {
			str += addr.String() + ","
		}
		str = strings.TrimSuffix(str, ",")
	}
	if len(this.Txt) > 0 {
		str += " txt=" + fmt.Sprint(this.Txt)
	}
	str += ">"
	return str
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package main

// This file implements go name generation for thrift identifiers.
// It has to match the Apache Thrift generated names.

import "strings"

// goKeywords taken from https://golang.org/ref/spec#Keywords (and added error).
var goKeywords = map[string]bool{
	"error":       true,
	"break":       true,
	"default":     true,
	"func":        true,
	"interface":   true,
	"select":      true,
	"case":        true,
	"defer":       true,
	"go":          true,
	"map":         true,
	"struct":      true,
	"chan":        true,
	"else":        true,
	"goto":        true,
	"package":     true,
	"switch":      true,
	"const":       true,
	"fallthrough": true,
	"if":          true,
	"range":       true,
	"type":        true,
	"continue":    true,
	"for":         true,
	"import":      true,
	"return":      true,
	"var":         true,
}

func goName(name string) string {
	// Thrift Identifier from IDL: ( Letter | '_' ) ( Letter | Digit | '.' | '_' )*
	// Go identifier from spec: letter { letter | unicode_digit } .
	// Go letter allows underscore, so the only difference is period. However, periods cannot
	// actaully be used - this seems to be a bug in the IDL.
	if _, ok := goKeywords[name]; ok {
		// The thrift compiler appends _a1 for any clashes with go keywords.
		name += "_a1"
	}

	name = camelCase(name)
	return name
}

// camelCase takes a name with underscores such as my_arg and returns camelCase (e.g. myArg).
func camelCase(name string) string {
	parts := strings.Split(name, "_")
	for i := 1; i < len(parts); i++ {
		name := parts[i]
		if name == "" {
			continue
		}

		// If the first letter is uppercase, Thrift keeps the underscore.
		if strings.ToUpper(name[0:1]) == name[0:1] {
			name = "_" + name
		} else {
			name = strings.ToUpper(name[0:1]) + name[1:]
		}

		parts[i] = name
	}
	return strings.Join(parts, "")
}

func avoidThriftClash(name string) string {
	if strings.HasSuffix(name, "Result") ||
		strings.HasSuffix(name, "Args") ||
		strings.HasPrefix(name, "New") {
		return name + "_"
	}
	return name
}

// goPublicName returns a go identifier that is exported.
func goPublicName(name string) string {
	// Public names cannot clash with goKeywords as they are all lowercase.
	name = camelCase(name)
	name = strings.ToUpper(name[0:1]) + name[1:]
	return name
}

// goPublicFieldName returns the name of the field as used in a struct.
func goPublicFieldName(name string) string {
	name = camcelCase(name)
	name = strings.ToUpper(name[0:1]) + name[1:]
	return avoidThriftClash(name)
}

var thriftToGo = map[string]string{
	"bool":   "bool",
	"byte":   "int8",
	"i16":    "int16",
	"i32":    "int32",
	"i64":    "int64",
	"double": "float64",
	"string": "string",
}
<commit_msg>Refactor calls to camelCase<commit_after>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package main

// This file implements go name generation for thrift identifiers.
// It has to match the Apache Thrift generated names.

import "strings"

// goKeywords taken from https://golang.org/ref/spec#Keywords (and added error).
var goKeywords = map[string]bool{
	"error":       true,
	"break":       true,
	"default":     true,
	"func":        true,
	"interface":   true,
	"select":      true,
	"case":        true,
	"defer":       true,
	"go":          true,
	"map":         true,
	"struct":      true,
	"chan":        true,
	"else":        true,
	"goto":        true,
	"package":     true,
	"switch":      true,
	"const":       true,
	"fallthrough": true,
	"if":          true,
	"range":       true,
	"type":        true,
	"continue":    true,
	"for":         true,
	"import":      true,
	"return":      true,
	"var":         true,
}

func goName(name string) string {
	// Thrift Identifier from IDL: ( Letter | '_' ) ( Letter | Digit | '.' | '_' )*
	// Go identifier from spec: letter { letter | unicode_digit } .
	// Go letter allows underscore, so the only difference is period. However, periods cannot
	// actaully be used - this seems to be a bug in the IDL.
	if _, ok := goKeywords[name]; ok {
		// The thrift compiler appends _a1 for any clashes with go keywords.
		name += "_a1"
	}

	name = camelCase(name, false /* publicName */)
	return name
}

// camelCase takes a name with underscores such as my_arg and returns camelCase (e.g. myArg).
func camelCase(name string, publicName bool) string {
	parts := strings.Split(name, "_")
	for i := 1; i < len(parts); i++ {
		name := parts[i]
		if name == "" {
			continue
		}

		// If the first letter is uppercase, Thrift keeps the underscore.
		if strings.ToUpper(name[0:1]) == name[0:1] {
			name = "_" + name
		} else {
			name = strings.ToUpper(name[0:1]) + name[1:]
		}

		parts[i] = name
	}
	return strings.Join(parts, "")
}

func avoidThriftClash(name string) string {
	if strings.HasSuffix(name, "Result") ||
		strings.HasSuffix(name, "Args") ||
		strings.HasPrefix(name, "New") {
		return name + "_"
	}
	return name
}

// goPublicName returns a go identifier that is exported.
func goPublicName(name string) string {
	// Public names cannot clash with goKeywords as they are all lowercase.
	name = camelCase(name, false /* publicName */)
	name = strings.ToUpper(name[0:1]) + name[1:]
	return name
}

// goPublicFieldName returns the name of the field as used in a struct.
func goPublicFieldName(name string) string {
	return avoidThriftClash(goPublicName(name))
}

var thriftToGo = map[string]string{
	"bool":   "bool",
	"byte":   "int8",
	"i16":    "int16",
	"i32":    "int32",
	"i64":    "int64",
	"double": "float64",
	"string": "string",
}
<|endoftext|>


<commit_before>// Copyright 2017 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package buildsource

import (
	"encoding/hex"
	"fmt"

	"golang.org/x/net/context"

	"go.chromium.org/gae/service/datastore"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/sync/parallel"

	"go.chromium.org/luci/milo/api/config"
	"go.chromium.org/luci/milo/common"
	"go.chromium.org/luci/milo/common/model"
	"go.chromium.org/luci/milo/frontend/ui"
)

// ConsoleRow is one row of a particular console.
//
// It has the git commit for the row, as well as a mapping of BuilderID to any
// BuildSummaries which reported using this commit.
//
// For Builder definitions which had more than one BuilderID assoaciated with
// them, they're indexed by the first BuilderID in the console's definition. So
// if the Console has a Builder with name: "A", and name: "B", only "A" will
// show up the Builds map (but if you look at the []BuildSummary, you'll see
// builds from both "A" and "B").
type ConsoleRow struct {
	Commit string
	Builds map[BuilderID][]*model.BuildSummary
}

// GetConsoleRows returns a row-oriented collection of BuildSummary
// objects. Each row corresponds to the similarly-indexed commit in the
// `commits` slice.
func GetConsoleRows(c context.Context, project string, console *config.Console, commits []string) ([]*ConsoleRow, error) {
	rawCommits := make([][]byte, len(commits))
	for i, c := range commits {
		var err error
		if rawCommits[i], err = hex.DecodeString(c); err != nil {
			return nil, errors.Annotate(err, "bad commit[%d]: %q", i, c).Err()
		}
	}

	// Maps BuilderID -> Column BuilderID (see ConsoleRow).
	columnMap := map[string]BuilderID{}
	for _, b := range console.Builders {
		columnID := b.Name[0]
		columnMap[columnID] = BuilderID(columnID)
		for _, n := range b.Name[1:] {
			columnMap[n] = BuilderID(columnID)
		}
	}

	ret := make([]*ConsoleRow, len(commits))
	url := console.RepoUrl
	// HACK(iannucci): This little hack should be removed when console definitions
	// no longer use a manifest name of "REVISION". REVISION was used to index the
	// 'got_revision' value before manifests were implemented.
	if console.ManifestName == "REVISION" {
		url = ""
	}
	partialKey := model.NewPartialManifestKey(project, console.Id, console.ManifestName, url)
	q := datastore.NewQuery("BuildSummary")
	err := parallel.WorkPool(16, func(ch chan<- func() error) {
		for i := range rawCommits {
			i := i
			r := &ConsoleRow{Commit: commits[i]}
			ret[i] = r
			ch <- func() error {
				fullQ := q.Eq("ManifestKeys", partialKey.AddRevision(rawCommits[i]))
				return datastore.Run(c, fullQ, func(bs *model.BuildSummary) {
					if bs.Experimental && !console.IncludeExperimentalBuilds {
						return
					}
					if i, ok := columnMap[bs.BuilderID]; ok {
						if r.Builds == nil {
							r.Builds = map[BuilderID][]*model.BuildSummary{}
						}
						r.Builds[i] = append(r.Builds[i], bs)
					}
				})
			}
		}
	})

	return ret, err
}

// ConsolePreview is mapping of builder IDs to each builder's latest build.
//
// This reflects a console preview, which is a console table except rendered
// with only the builder's latest build.
type ConsolePreview map[BuilderID]*model.BuildSummary

// GetConsoleSummariesFromIDs returns a list of console summaries from the datastore
// using a slice of console IDs as input.
//
// This list of console summaries directly corresponds to the input list of
// console IDs.
//
// projectID is the project being served in the current request.
func GetConsoleSummariesFromIDs(c context.Context, consoleIDs []common.ConsoleID, projectID string) (
	map[common.ConsoleID]*ui.ConsoleSummary, error) {

	// Get the console definitions and builders, then rearrange them into console summaries.
	defs, err := common.GetConsoles(c, consoleIDs)
	if err != nil {
		return nil, err
	}
	return GetConsoleSummariesFromDefs(c, defs, projectID)
}

// GetConsoleSummariesFromDefs returns a map of consoleID -> summary from the
// datastore using a slice of console definitions as input.
//
// This expects all builders in all consoles coming from the same projectID.
func GetConsoleSummariesFromDefs(c context.Context, consoleEnts []*common.Console, projectID string) (
	map[common.ConsoleID]*ui.ConsoleSummary, error) {

	// Maps consoleID -> console config definition.
	consoles := make(map[common.ConsoleID]*config.Console, len(consoleEnts))

	// Maps the BuilderID to the per-console pointer-to-summary in the summaries
	// map. Note that builders with multiple builderIDs in the same console will
	// all map to the same BuilderSummary.
	columns := map[BuilderID]map[common.ConsoleID]*model.BuilderSummary{}

	// The return result.
	summaries := map[common.ConsoleID]*ui.ConsoleSummary{}

	for _, ent := range consoleEnts {
		cid := ent.ConsoleID()
		consoles[cid] = &ent.Def

		summaries[cid] = &ui.ConsoleSummary{
			Builders: make([]*model.BuilderSummary, len(ent.Def.Builders)),
			Name: ui.NewLink(
				ent.ID,
				fmt.Sprintf("/p/%s/g/%s/console", ent.ProjectID(), ent.ID),
				fmt.Sprintf("Console %s in project %s", ent.ID, ent.ProjectID()),
			),
		}
		// Some builders might not have any builds yet, but we still want an (empty)
		// bubble to show up here.
		for i := range summaries[cid].Builders {
			summaries[cid].Builders[i] = &model.BuilderSummary{}
		}

		for i, column := range ent.Def.Builders {
			for _, rawName := range column.Name {
				name := BuilderID(rawName)
				// Find/populate the BuilderID -> {console: summary}
				colMap, ok := columns[name]
				if !ok {
					colMap = map[common.ConsoleID]*model.BuilderSummary{}
					columns[name] = colMap
				}

				colMap[cid] = summaries[cid].Builders[i]
			}
		}
	}

	// Now grab ALL THE DATA.
	bs := make([]*model.BuilderSummary, 0, len(columns))
	for builderID := range columns {
		bs = append(bs, &model.BuilderSummary{
			BuilderID: string(builderID),
			// TODO: change builder ID format to include project id.
			ProjectID: projectID,
		})
	}
	if err := datastore.Get(c, bs); err != nil {
		me := err.(errors.MultiError)
		lme := errors.NewLazyMultiError(len(me))
		for i, ierr := range me {
			if ierr == datastore.ErrNoSuchEntity {
				ierr = nil  // ignore ErrNoSuchEntity
				bs[i] = nil // nil out the BuilderSummary, want to skip this below
			}
			lme.Assign(i, ierr)
		}

		// Return an error only if we encouter an error other than datastore.ErrNoSuchEntity.
		if err := lme.Get(); err != nil {
			return nil, err
		}
	}

	// Now we have the mapping from BuilderID -> summary, and ALL THE DATA, map
	// the data back into the summaries.
	for _, summary := range bs {
		if summary == nil { // We got ErrNoSuchEntity above
			continue
		}

		for cid, curSummary := range columns[BuilderID(summary.BuilderID)] {
			cons := consoles[cid]

			// If this console doesn't show experimental builds, skip all summaries of
			// experimental builds.
			if !cons.IncludeExperimentalBuilds && summary.LastFinishedExperimental {
				continue
			}

			// If the new summary's build was created before the current summary's
			// build, skip it.
			if summary.LastFinishedCreated.Before((*curSummary).LastFinishedCreated) {
				continue
			}

			// Looks like this is the best summary for this slot so far, so save it.
			*curSummary = *summary
		}
	}

	return summaries, nil
}
<commit_msg>[milo] Fix crash due to missing ProjectID<commit_after>// Copyright 2017 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package buildsource

import (
	"encoding/hex"
	"fmt"

	"golang.org/x/net/context"

	"go.chromium.org/gae/service/datastore"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/sync/parallel"

	"go.chromium.org/luci/milo/api/config"
	"go.chromium.org/luci/milo/common"
	"go.chromium.org/luci/milo/common/model"
	"go.chromium.org/luci/milo/frontend/ui"
)

// ConsoleRow is one row of a particular console.
//
// It has the git commit for the row, as well as a mapping of BuilderID to any
// BuildSummaries which reported using this commit.
//
// For Builder definitions which had more than one BuilderID assoaciated with
// them, they're indexed by the first BuilderID in the console's definition. So
// if the Console has a Builder with name: "A", and name: "B", only "A" will
// show up the Builds map (but if you look at the []BuildSummary, you'll see
// builds from both "A" and "B").
type ConsoleRow struct {
	Commit string
	Builds map[BuilderID][]*model.BuildSummary
}

// GetConsoleRows returns a row-oriented collection of BuildSummary
// objects. Each row corresponds to the similarly-indexed commit in the
// `commits` slice.
func GetConsoleRows(c context.Context, project string, console *config.Console, commits []string) ([]*ConsoleRow, error) {
	rawCommits := make([][]byte, len(commits))
	for i, c := range commits {
		var err error
		if rawCommits[i], err = hex.DecodeString(c); err != nil {
			return nil, errors.Annotate(err, "bad commit[%d]: %q", i, c).Err()
		}
	}

	// Maps BuilderID -> Column BuilderID (see ConsoleRow).
	columnMap := map[string]BuilderID{}
	for _, b := range console.Builders {
		columnID := b.Name[0]
		columnMap[columnID] = BuilderID(columnID)
		for _, n := range b.Name[1:] {
			columnMap[n] = BuilderID(columnID)
		}
	}

	ret := make([]*ConsoleRow, len(commits))
	url := console.RepoUrl
	// HACK(iannucci): This little hack should be removed when console definitions
	// no longer use a manifest name of "REVISION". REVISION was used to index the
	// 'got_revision' value before manifests were implemented.
	if console.ManifestName == "REVISION" {
		url = ""
	}
	partialKey := model.NewPartialManifestKey(project, console.Id, console.ManifestName, url)
	q := datastore.NewQuery("BuildSummary")
	err := parallel.WorkPool(16, func(ch chan<- func() error) {
		for i := range rawCommits {
			i := i
			r := &ConsoleRow{Commit: commits[i]}
			ret[i] = r
			ch <- func() error {
				fullQ := q.Eq("ManifestKeys", partialKey.AddRevision(rawCommits[i]))
				return datastore.Run(c, fullQ, func(bs *model.BuildSummary) {
					if bs.Experimental && !console.IncludeExperimentalBuilds {
						return
					}
					if i, ok := columnMap[bs.BuilderID]; ok {
						if r.Builds == nil {
							r.Builds = map[BuilderID][]*model.BuildSummary{}
						}
						r.Builds[i] = append(r.Builds[i], bs)
					}
				})
			}
		}
	})

	return ret, err
}

// ConsolePreview is mapping of builder IDs to each builder's latest build.
//
// This reflects a console preview, which is a console table except rendered
// with only the builder's latest build.
type ConsolePreview map[BuilderID]*model.BuildSummary

// GetConsoleSummariesFromIDs returns a list of console summaries from the datastore
// using a slice of console IDs as input.
//
// This list of console summaries directly corresponds to the input list of
// console IDs.
//
// projectID is the project being served in the current request.
func GetConsoleSummariesFromIDs(c context.Context, consoleIDs []common.ConsoleID, projectID string) (
	map[common.ConsoleID]*ui.ConsoleSummary, error) {

	// Get the console definitions and builders, then rearrange them into console summaries.
	defs, err := common.GetConsoles(c, consoleIDs)
	if err != nil {
		return nil, err
	}
	return GetConsoleSummariesFromDefs(c, defs, projectID)
}

// GetConsoleSummariesFromDefs returns a map of consoleID -> summary from the
// datastore using a slice of console definitions as input.
//
// This expects all builders in all consoles coming from the same projectID.
func GetConsoleSummariesFromDefs(c context.Context, consoleEnts []*common.Console, projectID string) (
	map[common.ConsoleID]*ui.ConsoleSummary, error) {

	// Maps consoleID -> console config definition.
	consoles := make(map[common.ConsoleID]*config.Console, len(consoleEnts))

	// Maps the BuilderID to the per-console pointer-to-summary in the summaries
	// map. Note that builders with multiple builderIDs in the same console will
	// all map to the same BuilderSummary.
	columns := map[BuilderID]map[common.ConsoleID]*model.BuilderSummary{}

	// The return result.
	summaries := map[common.ConsoleID]*ui.ConsoleSummary{}

	for _, ent := range consoleEnts {
		cid := ent.ConsoleID()
		consoles[cid] = &ent.Def

		summaries[cid] = &ui.ConsoleSummary{
			Builders: make([]*model.BuilderSummary, len(ent.Def.Builders)),
			Name: ui.NewLink(
				ent.ID,
				fmt.Sprintf("/p/%s/g/%s/console", ent.ProjectID(), ent.ID),
				fmt.Sprintf("Console %s in project %s", ent.ID, ent.ProjectID()),
			),
		}

		for i, column := range ent.Def.Builders {
			s := &model.BuilderSummary{
				BuilderID: column.Name[0],
				ProjectID: projectID,
			}
			summaries[cid].Builders[i] = s
			for _, rawName := range column.Name {
				name := BuilderID(rawName)
				// Find/populate the BuilderID -> {console: summary}
				colMap, ok := columns[name]
				if !ok {
					colMap = map[common.ConsoleID]*model.BuilderSummary{}
					columns[name] = colMap
				}

				colMap[cid] = s
			}
		}
	}

	// Now grab ALL THE DATA.
	bs := make([]*model.BuilderSummary, 0, len(columns))
	for builderID := range columns {
		bs = append(bs, &model.BuilderSummary{
			BuilderID: string(builderID),
			// TODO: change builder ID format to include project id.
			ProjectID: projectID,
		})
	}
	if err := datastore.Get(c, bs); err != nil {
		me := err.(errors.MultiError)
		lme := errors.NewLazyMultiError(len(me))
		for i, ierr := range me {
			if ierr == datastore.ErrNoSuchEntity {
				logging.Infof(c, "Missing builder: %s", bs[i].BuilderID)
				ierr = nil  // ignore ErrNoSuchEntity
				bs[i] = nil // nil out the BuilderSummary, want to skip this below
			}
			lme.Assign(i, ierr)
		}

		// Return an error only if we encouter an error other than datastore.ErrNoSuchEntity.
		if err := lme.Get(); err != nil {
			return nil, err
		}
	}

	// Now we have the mapping from BuilderID -> summary, and ALL THE DATA, map
	// the data back into the summaries.
	for _, summary := range bs {
		if summary == nil { // We got ErrNoSuchEntity above
			continue
		}

		for cid, curSummary := range columns[BuilderID(summary.BuilderID)] {
			cons := consoles[cid]

			// If this console doesn't show experimental builds, skip all summaries of
			// experimental builds.
			if !cons.IncludeExperimentalBuilds && summary.LastFinishedExperimental {
				continue
			}

			// If the new summary's build was created before the current summary's
			// build, skip it.
			if summary.LastFinishedCreated.Before((*curSummary).LastFinishedCreated) {
				continue
			}

			// Looks like this is the best summary for this slot so far, so save it.
			*curSummary = *summary
		}
	}

	return summaries, nil
}
<|endoftext|>


<commit_before>package zoneFileParser

import (
	"bufio"
	"bytes"
	"encoding/hex"
	"errors"
	"fmt"
	"rains/rainslib"
	"strconv"

	log "github.com/inconshreveable/log15"
)

//Parser can be used to parse RAINS zone files
type Parser struct {
}

//ParseZoneFile returns all assertions contained in the given zonefile
func (p Parser) ParseZoneFile(zoneFile []byte) ([]*rainslib.AssertionSection, error) {
	assertions := []*rainslib.AssertionSection{}
	scanner := bufio.NewScanner(bytes.NewReader(zoneFile))
	scanner.Split(bufio.ScanWords)
	scanner.Scan()
	if scanner.Text() != ":Z:" {
		log.Warn("zoneFile malformed. It does not start with :Z:")
	}
	scanner.Scan()
	context := scanner.Text()
	scanner.Scan()
	zone := scanner.Text()
	scanner.Scan()
	if scanner.Text() != "[" {
		log.Warn("zonFile malformed. Expected '['")
	}
	scanner.Scan()
	for scanner.Text() != "]" {
		switch scanner.Text() {
		case ":A:":
			a, err := parseAssertion(context, zone, scanner)
			if err != nil {
				return nil, err
			}
			assertions = append(assertions, a)
		case ":S:":
			asserts, err := parseShard(context, zone, scanner)
			if err != nil {
				return nil, err
			}
			assertions = append(assertions, asserts...)
		default:
			return nil, fmt.Errorf("Expected a shard or assertion inside the zone but got=%s", scanner.Text())
		}
		scanner.Scan() //reads in the next section's type or exit the loop in case of ']'
	}
	return assertions, nil
}

func parseShard(context, zone string, scanner *bufio.Scanner) ([]*rainslib.AssertionSection, error) {
	assertions := []*rainslib.AssertionSection{}
	scanner.Scan()
	if scanner.Text() != "[" {
		log.Warn("zonFile malformed. Expected '['")
	}
	scanner.Scan()
	for scanner.Text() != "]" {
		if scanner.Text() != ":A:" {
			return nil, fmt.Errorf("zone file malformed. Expected Assertion inside shard but got=%s", scanner.Text())
		}
		a, err := parseAssertion(context, zone, scanner)
		if err != nil {
			return nil, err
		}
		assertions = append(assertions, a)
		scanner.Scan()
	}
	return assertions, nil
}

//parseAssertion parses the assertions content and returns an assertion section
func parseAssertion(context, zone string, scanner *bufio.Scanner) (*rainslib.AssertionSection, error) {
	scanner.Scan()
	name := scanner.Text()
	scanner.Scan()
	if scanner.Text() != "[" {
		log.Warn("zonFile malformed. Expected '['")
	}
	scanner.Scan()
	objects := []rainslib.Object{}
	for scanner.Text() != "]" {
		switch scanner.Text() {
		case ":name:":
			scanner.Scan()
			objects = append(objects, rainslib.Object{Type: rainslib.OTName, Value: scanner.Text()})
		case ":ip6:":
			scanner.Scan()
			objects = append(objects, rainslib.Object{Type: rainslib.OTIP6Addr, Value: scanner.Text()})
		case ":ip4:":
			scanner.Scan()
			objects = append(objects, rainslib.Object{Type: rainslib.OTIP4Addr, Value: scanner.Text()})
		case ":redir:":
			scanner.Scan()
			objects = append(objects, rainslib.Object{Type: rainslib.OTRedirection, Value: scanner.Text()})
		case ":deleg:":
			delegation, err := getPublicKey(scanner)
			if err != nil {
				return nil, err
			}
			objects = append(objects, rainslib.Object{Type: rainslib.OTDelegation, Value: delegation})
		case ":nameset:":
			scanner.Scan()
			objects = append(objects, rainslib.Object{Type: rainslib.OTNameset, Value: scanner.Text()})
		case ":cert:":
			cert, err := getCertObject(scanner)
			if err != nil {
				return nil, err
			}
			objects = append(objects, rainslib.Object{Type: rainslib.OTCertInfo, Value: cert})
		case ":srv:":
			srvInfo := rainslib.ServiceInfo{}
			scanner.Scan()
			srvInfo.Name = scanner.Text()
			scanner.Scan()
			portNr, err := strconv.Atoi(scanner.Text())
			if err != nil {
				return nil, err
			}
			srvInfo.Port = uint16(portNr)
			scanner.Scan()
			prio, err := strconv.Atoi(scanner.Text())
			if err != nil {
				return nil, err
			}
			srvInfo.Priority = uint(prio)
			objects = append(objects, rainslib.Object{Type: rainslib.OTServiceInfo, Value: srvInfo})
		case ":regr:":
			scanner.Scan()
			objects = append(objects, rainslib.Object{Type: rainslib.OTRegistrar, Value: scanner.Text()})
		case ":regt:":
			scanner.Scan()
			objects = append(objects, rainslib.Object{Type: rainslib.OTRegistrant, Value: scanner.Text()})
		case ":infra:":
			infrastructureKey, err := getPublicKey(scanner)
			if err != nil {
				return nil, err
			}
			objects = append(objects, rainslib.Object{Type: rainslib.OTDelegation, Value: infrastructureKey})
		case ":extra:":
			//TODO CFE not yet implemented
			return nil, errors.New("TODO CFE not yet implemented")
		default:
			return nil, fmt.Errorf("Encountered non existing object type: %s", scanner.Text())
		}
		scanner.Scan() //scan next object type
	}
	a := &rainslib.AssertionSection{Context: context, SubjectZone: zone, SubjectName: name, Content: objects}
	log.Debug("parsed Assertion", "assertion", *a)
	return a, nil
}

func getCertObject(scanner *bufio.Scanner) (rainslib.CertificateObject, error) {
	scanner.Scan()
	certType, err := getCertPT(scanner.Text())
	if err != nil {
		return rainslib.CertificateObject{}, err
	}
	scanner.Scan()
	usage, err := getCertUsage(scanner.Text())
	if err != nil {
		return rainslib.CertificateObject{}, err
	}
	scanner.Scan()
	hashAlgo, err := getCertHashType(scanner.Text())
	if err != nil {
		return rainslib.CertificateObject{}, err
	}
	scanner.Scan()
	cert := rainslib.CertificateObject{
		Type:     certType,
		Usage:    usage,
		HashAlgo: hashAlgo,
		Data:     scanner.Bytes(),
	}
	return cert, nil
}

func getCertPT(certType string) (rainslib.ProtocolType, error) {
	switch certType {
	case "0":
		return rainslib.PTUnspecified, nil
	case "1":
		return rainslib.PTTLS, nil
	default:
		return rainslib.ProtocolType(-1), errors.New("Encountered non existing certificate protocol type")
	}
}

func getCertUsage(certType string) (rainslib.CertificateUsage, error) {
	switch certType {
	case "2":
		return rainslib.CUTrustAnchor, nil
	case "3":
		return rainslib.CUEndEntity, nil
	default:
		return rainslib.CertificateUsage(-1), errors.New("Encountered non existing certificate usage")
	}
}

func getCertHashType(certType string) (rainslib.HashAlgorithmType, error) {
	switch certType {
	case "0":
		return rainslib.NoHashAlgo, nil
	case "1":
		return rainslib.Sha256, nil
	case "2":
		return rainslib.Sha384, nil
	case "3":
		return rainslib.Sha512, nil
	default:
		return rainslib.HashAlgorithmType(-1), errors.New("Encountered non existing certificate hash algorithm")
	}
}

func getPublicKey(scanner *bufio.Scanner) (rainslib.PublicKey, error) {
	scanner.Scan()
	keyAlgoType, err := getKeyAlgoType(scanner.Text())
	if err != nil {
		return rainslib.PublicKey{}, err
	}
	scanner.Scan()
	publicKey := rainslib.PublicKey{Type: keyAlgoType}
	switch keyAlgoType {
	case rainslib.Ed25519:
		return decodePublicKey(scanner, publicKey)
	case rainslib.Ed448:
		return decodePublicKey(scanner, publicKey)
	case rainslib.Ecdsa256:
		log.Warn("Not yet implemented")
		publicKey.Key = rainslib.Ecdsa256PublicKey{}
	case rainslib.Ecdsa384:
		log.Warn("Not yet implemented")
		publicKey.Key = rainslib.Ecdsa384PublicKey{}
	default:
		return rainslib.PublicKey{}, fmt.Errorf("Encountered non existing signature algorithm type. Got:%T", keyAlgoType)
	}
	return publicKey, nil
}

func decodePublicKey(scanner *bufio.Scanner, publicKey rainslib.PublicKey) (rainslib.PublicKey, error) {
	pKey, err := hex.DecodeString(scanner.Text())
	if err != nil {
		return publicKey, err
	}
	if len(pKey) == 32 {
		key := rainslib.Ed25519PublicKey{}
		copy(key[:], pKey)
		publicKey.Key = key
		return publicKey, nil
	}
	if len(pKey) == 57 {
		key := rainslib.Ed448PublicKey{}
		copy(key[:], pKey)
		publicKey.Key = key
		return publicKey, nil
	}
	return publicKey, fmt.Errorf("public key length is not 32 or 57. Got:%d", len(pKey))
}

func getKeyAlgoType(keyAlgoType string) (rainslib.SignatureAlgorithmType, error) {
	switch keyAlgoType {
	case "ed25519":
		return rainslib.Ed25519, nil
	case "ed448":
		return rainslib.Ed448, nil
	case "ecdsa256":
		return rainslib.Ecdsa256, nil
	case "ecdsa384":
		return rainslib.Ecdsa384, nil
	default:
		return rainslib.SignatureAlgorithmType(-1), fmt.Errorf("Encountered non existing signature algorithm type. Got:%s", keyAlgoType)
	}
}
<commit_msg>improved zone file parser logging. Added line numbers and improved log messages<commit_after>package zoneFileParser

import (
	"bufio"
	"bytes"
	"encoding/hex"
	"errors"
	"fmt"
	"rains/rainslib"
	"strconv"
	"strings"

	log "github.com/inconshreveable/log15"
)

//Parser can be used to parse RAINS zone files
type Parser struct {
}

var lineNrLogger log.Logger

//ParseZoneFile returns all assertions contained in the given zonefile
func (p Parser) ParseZoneFile(zoneFile []byte) ([]*rainslib.AssertionSection, error) {
	assertions := []*rainslib.AssertionSection{}
	scanner := NewWordScanner(zoneFile)
	lineNrLogger = log.New("lineNr", log.Lazy{scanner.LineNumber})
	scanner.Scan()
	if scanner.Text() != ":Z:" {
		lineNrLogger.Error("zoneFile malformed.", "expected", ":Z:", "got", scanner.Text())
		return []*rainslib.AssertionSection{}, errors.New("ZoneFile malformed")
	}
	scanner.Scan()
	context := scanner.Text()
	scanner.Scan()
	zone := scanner.Text()
	scanner.Scan()
	if scanner.Text() != "[" {
		lineNrLogger.Error("zonFile malformed.", "expected", "[", "got", scanner.Text())
		return []*rainslib.AssertionSection{}, errors.New("ZoneFile malformed")
	}
	scanner.Scan()
	for scanner.Text() != "]" {
		switch scanner.Text() {
		case ":A:":
			a, err := parseAssertion(context, zone, scanner)
			if err != nil {
				return nil, err
			}
			assertions = append(assertions, a)
		case ":S:":
			asserts, err := parseShard(context, zone, scanner)
			if err != nil {
				return nil, err
			}
			assertions = append(assertions, asserts...)
		default:
			lineNrLogger.Error("zonFile malformed.", "expected", ":A: or :S:", "got", scanner.Text())
			return nil, errors.New("ZoneFile malformed")
		}
		scanner.Scan() //reads in the next section's type or exit the loop in case of ']'
	}
	return assertions, nil
}

func parseShard(context, zone string, scanner *WordScanner) ([]*rainslib.AssertionSection, error) {
	assertions := []*rainslib.AssertionSection{}
	scanner.Scan()
	if scanner.Text() != "[" {
		lineNrLogger.Error("zonFile malformed.", "expected", "[", "got", scanner.Text())
		return []*rainslib.AssertionSection{}, errors.New("ZoneFile malformed")
	}
	scanner.Scan()
	for scanner.Text() != "]" {
		if scanner.Text() != ":A:" {
			lineNrLogger.Error("zonFile malformed.", "expected", ":A:", "got", scanner.Text())
			return nil, errors.New("ZoneFile malformed")
		}
		a, err := parseAssertion(context, zone, scanner)
		if err != nil {
			return nil, err
		}
		assertions = append(assertions, a)
		scanner.Scan()
	}
	return assertions, nil
}

//parseAssertion parses the assertions content and returns an assertion section
func parseAssertion(context, zone string, scanner *WordScanner) (*rainslib.AssertionSection, error) {
	scanner.Scan()
	name := scanner.Text()
	scanner.Scan()
	if scanner.Text() != "[" {
		lineNrLogger.Error("zonFile malformed.", "expected", "[", "got", scanner.Text())
		return &rainslib.AssertionSection{}, errors.New("ZoneFile malformed")
	}
	scanner.Scan()
	objects := []rainslib.Object{}
	for scanner.Text() != "]" {
		switch scanner.Text() {
		case ":name:":
			scanner.Scan()
			objects = append(objects, rainslib.Object{Type: rainslib.OTName, Value: scanner.Text()})
		case ":ip6:":
			scanner.Scan()
			objects = append(objects, rainslib.Object{Type: rainslib.OTIP6Addr, Value: scanner.Text()})
		case ":ip4:":
			scanner.Scan()
			objects = append(objects, rainslib.Object{Type: rainslib.OTIP4Addr, Value: scanner.Text()})
		case ":redir:":
			scanner.Scan()
			objects = append(objects, rainslib.Object{Type: rainslib.OTRedirection, Value: scanner.Text()})
		case ":deleg:":
			delegation, err := getPublicKey(scanner)
			if err != nil {
				return nil, err
			}
			objects = append(objects, rainslib.Object{Type: rainslib.OTDelegation, Value: delegation})
		case ":nameset:":
			scanner.Scan()
			objects = append(objects, rainslib.Object{Type: rainslib.OTNameset, Value: scanner.Text()})
		case ":cert:":
			cert, err := getCertObject(scanner)
			if err != nil {
				return nil, err
			}
			objects = append(objects, rainslib.Object{Type: rainslib.OTCertInfo, Value: cert})
		case ":srv:":
			srvInfo := rainslib.ServiceInfo{}
			scanner.Scan()
			srvInfo.Name = scanner.Text()
			scanner.Scan()
			portNr, err := strconv.Atoi(scanner.Text())
			if err != nil {
				lineNrLogger.Error("zonFile malformed.", "expected", "a number", "got", scanner.Text())
				return nil, err
			}
			srvInfo.Port = uint16(portNr)
			scanner.Scan()
			prio, err := strconv.Atoi(scanner.Text())
			if err != nil {
				lineNrLogger.Error("zonFile malformed.", "expected", "a number", "got", scanner.Text())
				return nil, err
			}
			srvInfo.Priority = uint(prio)
			objects = append(objects, rainslib.Object{Type: rainslib.OTServiceInfo, Value: srvInfo})
		case ":regr:":
			scanner.Scan()
			objects = append(objects, rainslib.Object{Type: rainslib.OTRegistrar, Value: scanner.Text()})
		case ":regt:":
			scanner.Scan()
			objects = append(objects, rainslib.Object{Type: rainslib.OTRegistrant, Value: scanner.Text()})
		case ":infra:":
			infrastructureKey, err := getPublicKey(scanner)
			if err != nil {
				return nil, err
			}
			objects = append(objects, rainslib.Object{Type: rainslib.OTDelegation, Value: infrastructureKey})
		case ":extra:":
			//TODO CFE not yet implemented
			return nil, errors.New("TODO CFE not yet implemented")
		default:
			lineNrLogger.Error("zonFile malformed.", "expected", ":<objectType>: (e.g. :ip4:)", "got", scanner.Text())
			return nil, errors.New("ZoneFile malformed")
		}
		scanner.Scan() //scan next object type
	}
	a := &rainslib.AssertionSection{Context: context, SubjectZone: zone, SubjectName: name, Content: objects}
	log.Debug("parsed Assertion", "assertion", *a)
	return a, nil
}

func getCertObject(scanner *WordScanner) (rainslib.CertificateObject, error) {
	scanner.Scan()
	certType, err := getCertPT(scanner.Text())
	if err != nil {
		return rainslib.CertificateObject{}, err
	}
	scanner.Scan()
	usage, err := getCertUsage(scanner.Text())
	if err != nil {
		return rainslib.CertificateObject{}, err
	}
	scanner.Scan()
	hashAlgo, err := getCertHashType(scanner.Text())
	if err != nil {
		return rainslib.CertificateObject{}, err
	}
	scanner.Scan()
	cert := rainslib.CertificateObject{
		Type:     certType,
		Usage:    usage,
		HashAlgo: hashAlgo,
		Data:     []byte(scanner.Text()),
	}
	return cert, nil
}

func getCertPT(certType string) (rainslib.ProtocolType, error) {
	switch certType {
	case "0":
		return rainslib.PTUnspecified, nil
	case "1":
		return rainslib.PTTLS, nil
	default:
		lineNrLogger.Error("zonFile malformed.", "expected", "certificate protocol type identifier", "got", certType)
		return rainslib.ProtocolType(-1), errors.New("encountered non existing certificate protocol type id")
	}
}

func getCertUsage(usageType string) (rainslib.CertificateUsage, error) {
	switch usageType {
	case "2":
		return rainslib.CUTrustAnchor, nil
	case "3":
		return rainslib.CUEndEntity, nil
	default:
		lineNrLogger.Error("zonFile malformed.", "expected", "certificate usage identifier", "got", usageType)
		return rainslib.CertificateUsage(-1), errors.New("encountered non existing certificate usage")
	}
}

func getCertHashType(hashType string) (rainslib.HashAlgorithmType, error) {
	switch hashType {
	case "0":
		return rainslib.NoHashAlgo, nil
	case "1":
		return rainslib.Sha256, nil
	case "2":
		return rainslib.Sha384, nil
	case "3":
		return rainslib.Sha512, nil
	default:
		lineNrLogger.Error("zonFile malformed.", "expected", "certificate hash algo identifier", "got", hashType)
		return rainslib.HashAlgorithmType(-1), errors.New("encountered non existing certificate hash algorithm")
	}
}

func getPublicKey(scanner *WordScanner) (rainslib.PublicKey, error) {
	scanner.Scan()
	keyAlgoType, err := getKeyAlgoType(scanner.Text())
	if err != nil {
		return rainslib.PublicKey{}, err
	}
	scanner.Scan()
	publicKey := rainslib.PublicKey{Type: keyAlgoType}
	switch keyAlgoType {
	case rainslib.Ed25519:
		return decodePublicKey(scanner, publicKey)
	case rainslib.Ed448:
		return decodePublicKey(scanner, publicKey)
	case rainslib.Ecdsa256:
		log.Warn("Not yet implemented")
		publicKey.Key = rainslib.Ecdsa256PublicKey{}
	case rainslib.Ecdsa384:
		log.Warn("Not yet implemented")
		publicKey.Key = rainslib.Ecdsa384PublicKey{}
	default:
		lineNrLogger.Error("zonFile malformed.", "expected", "key algorithm type identifier", "got", keyAlgoType)
		return rainslib.PublicKey{}, errors.New("encountered non existing signature algorithm type")
	}
	return publicKey, nil
}

func decodePublicKey(scanner *WordScanner, publicKey rainslib.PublicKey) (rainslib.PublicKey, error) {
	pKey, err := hex.DecodeString(scanner.Text())
	if err != nil {
		return publicKey, err
	}
	if len(pKey) == 32 {
		key := rainslib.Ed25519PublicKey{}
		copy(key[:], pKey)
		publicKey.Key = key
		return publicKey, nil
	}
	if len(pKey) == 57 {
		key := rainslib.Ed448PublicKey{}
		copy(key[:], pKey)
		publicKey.Key = key
		return publicKey, nil
	}
	return publicKey, fmt.Errorf("public key length is not 32 or 57. got:%d", len(pKey))
}

func getKeyAlgoType(keyAlgoType string) (rainslib.SignatureAlgorithmType, error) {
	switch keyAlgoType {
	case "ed25519":
		return rainslib.Ed25519, nil
	case "ed448":
		return rainslib.Ed448, nil
	case "ecdsa256":
		return rainslib.Ecdsa256, nil
	case "ecdsa384":
		return rainslib.Ecdsa384, nil
	default:
		lineNrLogger.Error("zonFile malformed.", "expected", "signature algorithm type identifier", "got", keyAlgoType)
		return rainslib.SignatureAlgorithmType(-1), errors.New("encountered non existing signature algorithm type")
	}
}

//NewWordScanner returns a WordScanner
func NewWordScanner(data []byte) *WordScanner {
	scanner := bufio.NewScanner(bytes.NewReader(data))
	scanner.Split(bufio.ScanWords)
	return &WordScanner{data: data, scanner: scanner, wordsRead: 0}
}

//WordScanner uses bufio.Scanner to scan words of the input. Additionally it keeps track of the line (of the input) on which the scanner currently is
type WordScanner struct {
	data      []byte
	scanner   *bufio.Scanner
	wordsRead int
}

//Scan moves the pointer to the next token of the scan
func (ws *WordScanner) Scan() bool {
	ws.wordsRead++
	return ws.scanner.Scan()
}

//Text returns the value of the current Token as a string
func (ws *WordScanner) Text() string {
	return ws.scanner.Text()
}

//LineNumber returns the linenumber of the input data where the token pointer of the scanner currently is.
func (ws *WordScanner) LineNumber() int {
	lineScanner := bufio.NewScanner(bytes.NewReader(ws.data))
	i := 0
	lineNr := 1
	for lineScanner.Scan() && i < ws.wordsRead {
		scanner := bufio.NewScanner(strings.NewReader(lineScanner.Text()))
		scanner.Split(bufio.ScanWords)
		for scanner.Scan() {
			i++
			if i == ws.wordsRead {
				return lineNr
			}
		}
		lineNr++
	}
	return lineNr
}
<|endoftext|>