<commit_before>/*
  This file is part of go-ethereum

  go-ethereum is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  go-ethereum is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with go-ethereum.  If not, see <http://www.gnu.org/licenses/>.
*/
package rpc

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"math/big"
	"strings"

	"errors"
	"net"
	"net/http"
	"time"

	"io"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"
)

type hexdata struct {
	data  []byte
	isNil bool
}

func (d *hexdata) String() string {
	return "0x" + common.Bytes2Hex(d.data)
}

func (d *hexdata) MarshalJSON() ([]byte, error) {
	if d.isNil {
		return json.Marshal(nil)
	}
	return json.Marshal(d.String())
}

func newHexData(input interface{}) *hexdata {
	d := new(hexdata)

	if input == nil {
		d.isNil = true
		return d
	}
	switch input := input.(type) {
	case []byte:
		d.data = input
	case common.Hash:
		d.data = input.Bytes()
	case *common.Hash:
		if input == nil {
			d.isNil = true
		} else {
			d.data = input.Bytes()
		}
	case common.Address:
		d.data = input.Bytes()
	case *common.Address:
		if input == nil {
			d.isNil = true
		} else {
			d.data = input.Bytes()
		}
	case types.Bloom:
		d.data = input.Bytes()
	case *types.Bloom:
		if input == nil {
			d.isNil = true
		} else {
			d.data = input.Bytes()
		}
	case *big.Int:
		if input == nil {
			d.isNil = true
		} else {
			d.data = input.Bytes()
		}
	case int64:
		d.data = big.NewInt(input).Bytes()
	case uint64:
		buff := make([]byte, 8)
		binary.BigEndian.PutUint64(buff, input)
		d.data = buff
	case int:
		d.data = big.NewInt(int64(input)).Bytes()
	case uint:
		d.data = big.NewInt(int64(input)).Bytes()
	case int8:
		d.data = big.NewInt(int64(input)).Bytes()
	case uint8:
		d.data = big.NewInt(int64(input)).Bytes()
	case int16:
		d.data = big.NewInt(int64(input)).Bytes()
	case uint16:
		buff := make([]byte, 2)
		binary.BigEndian.PutUint16(buff, input)
		d.data = buff
	case int32:
		d.data = big.NewInt(int64(input)).Bytes()
	case uint32:
		buff := make([]byte, 4)
		binary.BigEndian.PutUint32(buff, input)
		d.data = buff
	case string: // hexstring
		d.data = common.Big(input).Bytes()
	default:
		d.isNil = true
	}

	return d
}

type hexnum struct {
	data  []byte
	isNil bool
}

func (d *hexnum) String() string {
	// Get hex string from bytes
	out := common.Bytes2Hex(d.data)
	// Trim leading 0s
	out = strings.TrimLeft(out, "0")
	// Output "0x0" when value is 0
	if len(out) == 0 {
		out = "0"
	}
	return "0x" + out
}

func (d *hexnum) MarshalJSON() ([]byte, error) {
	if d.isNil {
		return json.Marshal(nil)
	}
	return json.Marshal(d.String())
}

func newHexNum(input interface{}) *hexnum {
	d := new(hexnum)

	d.data = newHexData(input).data

	return d
}

type RpcConfig struct {
	ListenAddress string
	ListenPort    uint
	CorsDomain    string
}

type InvalidTypeError struct {
	method string
	msg    string
}

func (e *InvalidTypeError) Error() string {
	return fmt.Sprintf("invalid type on field %s: %s", e.method, e.msg)
}

func NewInvalidTypeError(method, msg string) *InvalidTypeError {
	return &InvalidTypeError{
		method: method,
		msg:    msg,
	}
}

type InsufficientParamsError struct {
	have int
	want int
}

func (e *InsufficientParamsError) Error() string {
	return fmt.Sprintf("insufficient params, want %d have %d", e.want, e.have)
}

func NewInsufficientParamsError(have int, want int) *InsufficientParamsError {
	return &InsufficientParamsError{
		have: have,
		want: want,
	}
}

type NotImplementedError struct {
	Method string
}

func (e *NotImplementedError) Error() string {
	return fmt.Sprintf("%s method not implemented", e.Method)
}

func NewNotImplementedError(method string) *NotImplementedError {
	return &NotImplementedError{
		Method: method,
	}
}

type NotAvailableError struct {
	Method string
	Reason string
}

func (e *NotAvailableError) Error() string {
	return fmt.Sprintf("%s method not available: %s", e.Method, e.Reason)
}

func NewNotAvailableError(method string, reason string) *NotAvailableError {
	return &NotAvailableError{
		Method: method,
		Reason: reason,
	}
}

type DecodeParamError struct {
	err string
}

func (e *DecodeParamError) Error() string {
	return fmt.Sprintf("could not decode, %s", e.err)

}

func NewDecodeParamError(errstr string) error {
	return &DecodeParamError{
		err: errstr,
	}
}

type ValidationError struct {
	ParamName string
	msg       string
}

func (e *ValidationError) Error() string {
	return fmt.Sprintf("%s not valid, %s", e.ParamName, e.msg)
}

func NewValidationError(param string, msg string) error {
	return &ValidationError{
		ParamName: param,
		msg:       msg,
	}
}

type RpcRequest struct {
	Id      interface{}     `json:"id"`
	Jsonrpc string          `json:"jsonrpc"`
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params"`
}

type RpcSuccessResponse struct {
	Id      interface{} `json:"id"`
	Jsonrpc string      `json:"jsonrpc"`
	Result  interface{} `json:"result"`
}

type RpcErrorResponse struct {
	Id      interface{}     `json:"id"`
	Jsonrpc string          `json:"jsonrpc"`
	Error   *RpcErrorObject `json:"error"`
}

type RpcErrorObject struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
	// Data    interface{} `json:"data"`
}

type listenerHasStoppedError struct {
	msg string
}

func (self listenerHasStoppedError) Error() string {
	return self.msg
}

var listenerStoppedError = listenerHasStoppedError{"Listener stopped"}

// When https://github.com/golang/go/issues/4674 is fixed this could be replaced
type stoppableTCPListener struct {
	*net.TCPListener
	stop chan struct{} // closed when the listener must stop
}

// Wraps the default handler and checks if the RPC service was stopped. In that case it returns an
// error indicating that the service was stopped. This will only happen for connections which are
// kept open (HTTP keep-alive) when the RPC service was shutdown.
func newStoppableHandler(h http.Handler, stop chan struct{}) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		select {
		case <-stop:
			w.Header().Set("Content-Type", "application/json")
			jsonerr := &RpcErrorObject{-32603, "RPC service stopped"}
			send(w, &RpcErrorResponse{Jsonrpc: jsonrpcver, Id: nil, Error: jsonerr})
		default:
			h.ServeHTTP(w, r)
		}
	})
}

// Stop the listener and all accepted and still active connections.
func (self *stoppableTCPListener) Stop() {
	close(self.stop)
}

func newStoppableTCPListener(addr string) (*stoppableTCPListener, error) {
	wl, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}

	if tcpl, ok := wl.(*net.TCPListener); ok {
		stop := make(chan struct{})
		l := &stoppableTCPListener{tcpl, stop}
		return l, nil
	}

	return nil, errors.New("Unable to create TCP listener for RPC service")
}

func (self *stoppableTCPListener) Accept() (net.Conn, error) {
	for {
		self.SetDeadline(time.Now().Add(time.Duration(1 * time.Second)))
		c, err := self.TCPListener.AcceptTCP()

		select {
		case <-self.stop:
			if c != nil { // accept timeout
				c.Close()
			}
			self.TCPListener.Close()
			return nil, listenerStoppedError
		default:
		}

		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() && netErr.Temporary() {
				continue // regular timeout
			}
		}

		return &closableConnection{c, self.stop}, err
	}
}

type closableConnection struct {
	*net.TCPConn
	closed chan struct{}
}

func (self *closableConnection) Read(b []byte) (n int, err error) {
	select {
	case <-self.closed:
		self.TCPConn.Close()
		return 0, io.EOF
	default:
		return self.TCPConn.Read(b)
	}
}
<commit_msg>fix hex conversion bug in RPC for byte slices<commit_after>/*
  This file is part of go-ethereum

  go-ethereum is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  go-ethereum is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with go-ethereum.  If not, see <http://www.gnu.org/licenses/>.
*/
package rpc

import (
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"math/big"
	"strings"

	"errors"
	"net"
	"net/http"
	"time"

	"io"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"
)

type hexdata struct {
	data  []byte
	isNil bool
}

func (d *hexdata) String() string {
	return "0x" + common.Bytes2Hex(d.data)
}

func (d *hexdata) MarshalJSON() ([]byte, error) {
	if d.isNil {
		return json.Marshal(nil)
	}
	return json.Marshal(d.String())
}

func newHexData(input interface{}) *hexdata {
	d := new(hexdata)

	if input == nil {
		d.isNil = true
		return d
	}
	switch input := input.(type) {
	case []byte:
		d.data = input
	case common.Hash:
		d.data = input.Bytes()
	case *common.Hash:
		if input == nil {
			d.isNil = true
		} else {
			d.data = input.Bytes()
		}
	case common.Address:
		d.data = input.Bytes()
	case *common.Address:
		if input == nil {
			d.isNil = true
		} else {
			d.data = input.Bytes()
		}
	case types.Bloom:
		d.data = input.Bytes()
	case *types.Bloom:
		if input == nil {
			d.isNil = true
		} else {
			d.data = input.Bytes()
		}
	case *big.Int:
		if input == nil {
			d.isNil = true
		} else {
			d.data = input.Bytes()
		}
	case int64:
		d.data = big.NewInt(input).Bytes()
	case uint64:
		buff := make([]byte, 8)
		binary.BigEndian.PutUint64(buff, input)
		d.data = buff
	case int:
		d.data = big.NewInt(int64(input)).Bytes()
	case uint:
		d.data = big.NewInt(int64(input)).Bytes()
	case int8:
		d.data = big.NewInt(int64(input)).Bytes()
	case uint8:
		d.data = big.NewInt(int64(input)).Bytes()
	case int16:
		d.data = big.NewInt(int64(input)).Bytes()
	case uint16:
		buff := make([]byte, 2)
		binary.BigEndian.PutUint16(buff, input)
		d.data = buff
	case int32:
		d.data = big.NewInt(int64(input)).Bytes()
	case uint32:
		buff := make([]byte, 4)
		binary.BigEndian.PutUint32(buff, input)
		d.data = buff
	case string: // hexstring
		// aaargh ffs TODO: avoid back-and-forth hex encodings where unneeded
		bytes, err := hex.DecodeString(strings.TrimPrefix(input, "0x"))
		if err != nil {
			d.isNil = true
		} else {
			d.data = bytes
		}
	default:
		d.isNil = true
	}

	return d
}

type hexnum struct {
	data  []byte
	isNil bool
}

func (d *hexnum) String() string {
	// Get hex string from bytes
	out := common.Bytes2Hex(d.data)
	// Trim leading 0s
	out = strings.TrimLeft(out, "0")
	// Output "0x0" when value is 0
	if len(out) == 0 {
		out = "0"
	}
	return "0x" + out
}

func (d *hexnum) MarshalJSON() ([]byte, error) {
	if d.isNil {
		return json.Marshal(nil)
	}
	return json.Marshal(d.String())
}

func newHexNum(input interface{}) *hexnum {
	d := new(hexnum)

	d.data = newHexData(input).data

	return d
}

type RpcConfig struct {
	ListenAddress string
	ListenPort    uint
	CorsDomain    string
}

type InvalidTypeError struct {
	method string
	msg    string
}

func (e *InvalidTypeError) Error() string {
	return fmt.Sprintf("invalid type on field %s: %s", e.method, e.msg)
}

func NewInvalidTypeError(method, msg string) *InvalidTypeError {
	return &InvalidTypeError{
		method: method,
		msg:    msg,
	}
}

type InsufficientParamsError struct {
	have int
	want int
}

func (e *InsufficientParamsError) Error() string {
	return fmt.Sprintf("insufficient params, want %d have %d", e.want, e.have)
}

func NewInsufficientParamsError(have int, want int) *InsufficientParamsError {
	return &InsufficientParamsError{
		have: have,
		want: want,
	}
}

type NotImplementedError struct {
	Method string
}

func (e *NotImplementedError) Error() string {
	return fmt.Sprintf("%s method not implemented", e.Method)
}

func NewNotImplementedError(method string) *NotImplementedError {
	return &NotImplementedError{
		Method: method,
	}
}

type NotAvailableError struct {
	Method string
	Reason string
}

func (e *NotAvailableError) Error() string {
	return fmt.Sprintf("%s method not available: %s", e.Method, e.Reason)
}

func NewNotAvailableError(method string, reason string) *NotAvailableError {
	return &NotAvailableError{
		Method: method,
		Reason: reason,
	}
}

type DecodeParamError struct {
	err string
}

func (e *DecodeParamError) Error() string {
	return fmt.Sprintf("could not decode, %s", e.err)

}

func NewDecodeParamError(errstr string) error {
	return &DecodeParamError{
		err: errstr,
	}
}

type ValidationError struct {
	ParamName string
	msg       string
}

func (e *ValidationError) Error() string {
	return fmt.Sprintf("%s not valid, %s", e.ParamName, e.msg)
}

func NewValidationError(param string, msg string) error {
	return &ValidationError{
		ParamName: param,
		msg:       msg,
	}
}

type RpcRequest struct {
	Id      interface{}     `json:"id"`
	Jsonrpc string          `json:"jsonrpc"`
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params"`
}

type RpcSuccessResponse struct {
	Id      interface{} `json:"id"`
	Jsonrpc string      `json:"jsonrpc"`
	Result  interface{} `json:"result"`
}

type RpcErrorResponse struct {
	Id      interface{}     `json:"id"`
	Jsonrpc string          `json:"jsonrpc"`
	Error   *RpcErrorObject `json:"error"`
}

type RpcErrorObject struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
	// Data    interface{} `json:"data"`
}

type listenerHasStoppedError struct {
	msg string
}

func (self listenerHasStoppedError) Error() string {
	return self.msg
}

var listenerStoppedError = listenerHasStoppedError{"Listener stopped"}

// When https://github.com/golang/go/issues/4674 is fixed this could be replaced
type stoppableTCPListener struct {
	*net.TCPListener
	stop chan struct{} // closed when the listener must stop
}

// Wraps the default handler and checks if the RPC service was stopped. In that case it returns an
// error indicating that the service was stopped. This will only happen for connections which are
// kept open (HTTP keep-alive) when the RPC service was shutdown.
func newStoppableHandler(h http.Handler, stop chan struct{}) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		select {
		case <-stop:
			w.Header().Set("Content-Type", "application/json")
			jsonerr := &RpcErrorObject{-32603, "RPC service stopped"}
			send(w, &RpcErrorResponse{Jsonrpc: jsonrpcver, Id: nil, Error: jsonerr})
		default:
			h.ServeHTTP(w, r)
		}
	})
}

// Stop the listener and all accepted and still active connections.
func (self *stoppableTCPListener) Stop() {
	close(self.stop)
}

func newStoppableTCPListener(addr string) (*stoppableTCPListener, error) {
	wl, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}

	if tcpl, ok := wl.(*net.TCPListener); ok {
		stop := make(chan struct{})
		l := &stoppableTCPListener{tcpl, stop}
		return l, nil
	}

	return nil, errors.New("Unable to create TCP listener for RPC service")
}

func (self *stoppableTCPListener) Accept() (net.Conn, error) {
	for {
		self.SetDeadline(time.Now().Add(time.Duration(1 * time.Second)))
		c, err := self.TCPListener.AcceptTCP()

		select {
		case <-self.stop:
			if c != nil { // accept timeout
				c.Close()
			}
			self.TCPListener.Close()
			return nil, listenerStoppedError
		default:
		}

		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() && netErr.Temporary() {
				continue // regular timeout
			}
		}

		return &closableConnection{c, self.stop}, err
	}
}

type closableConnection struct {
	*net.TCPConn
	closed chan struct{}
}

func (self *closableConnection) Read(b []byte) (n int, err error) {
	select {
	case <-self.closed:
		self.TCPConn.Close()
		return 0, io.EOF
	default:
		return self.TCPConn.Read(b)
	}
}
<|endoftext|>


<commit_before>package MQTTg

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

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

func (self *Broker) Start() error {
	addr, err := GetLocalAddr()
	fmt.Println(addr)
	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
		}
		bc := &BrokerSideClient{
			ClientInfo: &ClientInfo{
				Ct:             &Transport{conn},
				IsConnecting:   false,
				ID:             "",
				User:           nil,
				KeepAlive:      0,
				Will:           nil,
				PacketIDMap:    make(map[uint16]Message, 0),
				CleanSession:   false,
				KeepAliveTimer: time.NewTimer(0),
				Duration:       0,
				WriteChan:      make(chan Message),
			},
			SubTopics: make([]*SubscribeTopic, 0),
			Broker:    self,
		}
		go bc.ReadLoop(bc) // TODO: use single Loop function
		go bc.WriteLoop()
	}
}

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

			var id uint16 = 0
			var err error
			if w.QoS > 0 {
				id, err = subscriber.getUsablePacketID()
				if err != nil {
					panic(err)
				}
			}
			pub := NewPublishMessage(false, w.QoS, w.Retain, w.Topic, id, []uint8(w.Message))
			subscriber.WriteChan <- pub
		}
	}
	if self.IsConnecting {
		self.KeepAliveTimer.Stop()
		if self.CleanSession {
			delete(self.Broker.Clients, self.ID)
		}
	}
	err = self.disconnectBase()
	return err
}

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

type BrokerSideClient struct {
	*ClientInfo
	SubTopics []*SubscribeTopic
	Broker    *Broker
}

func (self *BrokerSideClient) RunClientTimer() {
	<-self.KeepAliveTimer.C
	EmitError(CLIENT_TIMED_OUT)
	self.disconnectProcessing()
	// 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.ClientInfo

	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
	connack := NewConnackMessage(sessionPresent, Accepted)
	self.WriteChan <- connack
	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]
		var id uint16 = 0
		var err error
		if qos > 0 {
			id, err = subscriber.getUsablePacketID()
			if err != nil {
				panic(err)
			}
		}
		pub := NewPublishMessage(false, qos, false, m.TopicName, id, m.Payload)
		subscriber.WriteChan <- pub
	}

	switch m.QoS {
	// in any case, Dub must be 0
	case 0:
		if m.PacketID != 0 {
			return PACKET_ID_SHOULD_BE_ZERO
		}
	case 1:
		puback := NewPubackMessage(m.PacketID)
		self.WriteChan <- puback
	case 2:
		pubrec := NewPubrecMessage(m.PacketID)
		self.WriteChan <- pubrec
	}
	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
	}
	pubrel := NewPubrelMessage(m.PacketID)
	self.WriteChan <- pubrel
	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
	}
	pubcomp := NewPubcompMessage(m.PacketID)
	self.WriteChan <- pubcomp
	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
					var id uint16
					if edge.RetainQoS > 0 {
						id, err = self.getUsablePacketID()
						if err != nil {
							return err
						}
					}
					pub := NewPublishMessage(false, edge.RetainQoS, true, edge.FullPath, id, []uint8(edge.RetainMessage))
					self.WriteChan <- pub
					EmitError(err)
				}
			}
		}
		returnCodes = append(returnCodes, codes...)
	}
	// TODO: check whether the number of return codes are correct?
	suback := NewSubackMessage(m.PacketID, returnCodes)
	self.WriteChan <- suback
	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
	unsuback := NewUnsubackMessage(m.PacketID)

	self.WriteChan <- unsuback
	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
	pingresp := NewPingrespMessage()
	self.WriteChan <- pingresp
	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.Will = nil
	self.disconnectProcessing()
	// close the client
	return err
}
<commit_msg>fix trivial<commit_after>package MQTTg

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

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

func (self *Broker) Start() error {
	addr, err := GetLocalAddr()
	fmt.Println(addr)
	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
		}
		bc := &BrokerSideClient{
			ClientInfo: &ClientInfo{
				Ct:             &Transport{conn},
				IsConnecting:   false,
				ID:             "",
				User:           nil,
				KeepAlive:      0,
				Will:           nil,
				PacketIDMap:    make(map[uint16]Message, 0),
				CleanSession:   false,
				KeepAliveTimer: time.NewTimer(0),
				Duration:       0,
				WriteChan:      make(chan Message),
			},
			SubTopics: make([]*SubscribeTopic, 0),
			Broker:    self,
		}
		go bc.ReadLoop(bc) // TODO: use single Loop function
		go bc.WriteLoop()
	}
}

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

			var id uint16 = 0
			var err error
			if w.QoS > 0 {
				id, err = subscriber.getUsablePacketID()
				if err != nil {
					panic(err)
				}
			}
			pub := NewPublishMessage(false, w.QoS, w.Retain, w.Topic, id, []uint8(w.Message))
			subscriber.WriteChan <- pub
		}
	}
	if self.IsConnecting {
		self.KeepAliveTimer.Stop()
		if self.CleanSession {
			delete(broker.Clients, self.ID)
		}
	}
	err = self.disconnectBase()
	return err
}

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

type BrokerSideClient struct {
	*ClientInfo
	SubTopics []*SubscribeTopic
	Broker    *Broker
}

func (self *BrokerSideClient) RunClientTimer() {
	<-self.KeepAliveTimer.C
	EmitError(CLIENT_TIMED_OUT)
	self.disconnectProcessing()
	// 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.ClientInfo

	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
	connack := NewConnackMessage(sessionPresent, Accepted)
	self.WriteChan <- connack
	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]
		var id uint16 = 0
		var err error
		if qos > 0 {
			id, err = subscriber.getUsablePacketID()
			if err != nil {
				panic(err)
			}
		}
		pub := NewPublishMessage(false, qos, false, m.TopicName, id, m.Payload)
		subscriber.WriteChan <- pub
	}

	switch m.QoS {
	// in any case, Dub must be 0
	case 0:
		if m.PacketID != 0 {
			return PACKET_ID_SHOULD_BE_ZERO
		}
	case 1:
		puback := NewPubackMessage(m.PacketID)
		self.WriteChan <- puback
	case 2:
		pubrec := NewPubrecMessage(m.PacketID)
		self.WriteChan <- pubrec
	}
	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
	}
	pubrel := NewPubrelMessage(m.PacketID)
	self.WriteChan <- pubrel
	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
	}
	pubcomp := NewPubcompMessage(m.PacketID)
	self.WriteChan <- pubcomp
	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
					var id uint16
					if edge.RetainQoS > 0 {
						id, err = self.getUsablePacketID()
						if err != nil {
							return err
						}
					}
					pub := NewPublishMessage(false, edge.RetainQoS, true, edge.FullPath, id, []uint8(edge.RetainMessage))
					self.WriteChan <- pub
					EmitError(err)
				}
			}
		}
		returnCodes = append(returnCodes, codes...)
	}
	// TODO: check whether the number of return codes are correct?
	suback := NewSubackMessage(m.PacketID, returnCodes)
	self.WriteChan <- suback
	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
	unsuback := NewUnsubackMessage(m.PacketID)

	self.WriteChan <- unsuback
	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
	pingresp := NewPingrespMessage()
	self.WriteChan <- pingresp
	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.Will = nil
	self.disconnectProcessing()
	// close the client
	return err
}
<|endoftext|>


<commit_before>package MQTTg

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

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

func (self *Broker) Start() error {
	addr, err := GetLocalAddr()
	fmt.Println(addr)
	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
		}
		bc := NewBrokerSideClient(&Transport{conn}, self)
		go bc.ReadLoop(bc) // TODO: use single Loop function
		go bc.WriteLoop()
	}
}

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

			var id uint16 = 0
			var err error
			if w.QoS > 0 {
				id, err = subscriber.getUsablePacketID()
				if err != nil {
					panic(err)
				}
			}
			pub := NewPublishMessage(false, w.QoS, w.Retain, w.Topic, id, []uint8(w.Message))
			subscriber.WriteChan <- pub
		}
	}
	if self.IsConnecting {
		self.KeepAliveTimer.Stop()
		if self.CleanSession {
			delete(broker.Clients, self.ID)
		}
	}
	err = self.disconnectBase()
	return err
}

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

type BrokerSideClient struct {
	*ClientInfo
	SubTopics []*SubscribeTopic
	Broker    *Broker
}

func NewBrokerSideClient(ct *Transport, broker *Broker) *BrokerSideClient {
	return &BrokerSideClient{
		ClientInfo: &ClientInfo{
			Ct:             ct,
			IsConnecting:   false,
			ID:             "",
			User:           nil,
			KeepAlive:      0,
			Will:           nil,
			PacketIDMap:    make(map[uint16]Message, 0),
			CleanSession:   false,
			KeepAliveTimer: time.NewTimer(0),
			Duration:       0,
			WriteChan:      make(chan Message),
		},
		SubTopics: make([]*SubscribeTopic, 0),
		Broker:    broker,
	}
}

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

func (self *BrokerSideClient) setPreviousSession(prevSession *BrokerSideClient) {
	self.SubTopics = prevSession.SubTopics

	self.PacketIDMap = prevSession.PacketIDMap
	self.CleanSession = prevSession.CleanSession
	self.Will = prevSession.Will
	self.Duration = prevSession.Duration
	self.KeepAliveTimer = time.NewTimer(self.Duration)
	self.KeepAlive = prevSession.KeepAlive
	// TODO: authorize here
	self.User = prevSession.User
}

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

	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
	connack := NewConnackMessage(sessionPresent, Accepted)
	self.WriteChan <- connack
	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]
		var id uint16 = 0
		var err error
		if qos > 0 {
			id, err = subscriber.getUsablePacketID()
			if err != nil {
				panic(err)
			}
		}
		pub := NewPublishMessage(false, qos, false, m.TopicName, id, m.Payload)
		subscriber.WriteChan <- pub
	}

	switch m.QoS {
	// in any case, Dub must be 0
	case 0:
		if m.PacketID != 0 {
			return PACKET_ID_SHOULD_BE_ZERO
		}
	case 1:
		puback := NewPubackMessage(m.PacketID)
		self.WriteChan <- puback
	case 2:
		pubrec := NewPubrecMessage(m.PacketID)
		self.WriteChan <- pubrec
	}
	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
	}
	pubrel := NewPubrelMessage(m.PacketID)
	self.WriteChan <- pubrel
	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
	}
	pubcomp := NewPubcompMessage(m.PacketID)
	self.WriteChan <- pubcomp
	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
					var id uint16
					if edge.RetainQoS > 0 {
						id, err = self.getUsablePacketID()
						if err != nil {
							return err
						}
					}
					pub := NewPublishMessage(false, edge.RetainQoS, true, edge.FullPath, id, []uint8(edge.RetainMessage))
					self.WriteChan <- pub
					EmitError(err)
				}
			}
		}
		returnCodes = append(returnCodes, codes...)
	}
	// TODO: check whether the number of return codes are correct?
	suback := NewSubackMessage(m.PacketID, returnCodes)
	self.WriteChan <- suback
	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
	unsuback := NewUnsubackMessage(m.PacketID)

	self.WriteChan <- unsuback
	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
	pingresp := NewPingrespMessage()
	self.WriteChan <- pingresp
	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.Will = nil
	self.disconnectProcessing()
	// close the client
	return err
}
<commit_msg>add QoS downgrade<commit_after>package MQTTg

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

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

func (self *Broker) Start() error {
	addr, err := GetLocalAddr()
	fmt.Println(addr)
	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
		}
		bc := NewBrokerSideClient(&Transport{conn}, self)
		go bc.ReadLoop(bc) // TODO: use single Loop function
		go bc.WriteLoop()
	}
}

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

			var id uint16 = 0
			var err error
			if w.QoS > 0 {
				id, err = subscriber.getUsablePacketID()
				if err != nil {
					panic(err)
				}
			}
			pub := NewPublishMessage(false, w.QoS, w.Retain, w.Topic, id, []uint8(w.Message))
			subscriber.WriteChan <- pub
		}
	}
	if self.IsConnecting {
		self.KeepAliveTimer.Stop()
		if self.CleanSession {
			delete(broker.Clients, self.ID)
		}
	}
	err = self.disconnectBase()
	return err
}

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

type BrokerSideClient struct {
	*ClientInfo
	SubTopics []*SubscribeTopic
	Broker    *Broker
}

func NewBrokerSideClient(ct *Transport, broker *Broker) *BrokerSideClient {
	return &BrokerSideClient{
		ClientInfo: &ClientInfo{
			Ct:             ct,
			IsConnecting:   false,
			ID:             "",
			User:           nil,
			KeepAlive:      0,
			Will:           nil,
			PacketIDMap:    make(map[uint16]Message, 0),
			CleanSession:   false,
			KeepAliveTimer: time.NewTimer(0),
			Duration:       0,
			WriteChan:      make(chan Message),
		},
		SubTopics: make([]*SubscribeTopic, 0),
		Broker:    broker,
	}
}

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

func (self *BrokerSideClient) setPreviousSession(prevSession *BrokerSideClient) {
	self.SubTopics = prevSession.SubTopics

	self.PacketIDMap = prevSession.PacketIDMap
	self.CleanSession = prevSession.CleanSession
	self.Will = prevSession.Will
	self.Duration = prevSession.Duration
	self.KeepAliveTimer = time.NewTimer(self.Duration)
	self.KeepAlive = prevSession.KeepAlive
	// TODO: authorize here
	self.User = prevSession.User
}

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

	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
	connack := NewConnackMessage(sessionPresent, Accepted)
	self.WriteChan <- connack
	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, reqQoS := range nodes[0].Subscribers {
		subscriber, _ := self.Broker.Clients[subscriberID]
		var id uint16 = 0
		var err error
		if m.QoS > 0 {
			id, err = subscriber.getUsablePacketID()
			if err != nil {
				panic(err)
			}
		}
		qos := m.QoS
		if reqQoS < m.QoS {
			// downgrade the QoS
			qos = reqQoS
		}

		pub := NewPublishMessage(false, qos, false, m.TopicName, id, m.Payload)
		subscriber.WriteChan <- pub
	}

	switch m.QoS {
	// in any case, Dub must be 0
	case 0:
		if m.PacketID != 0 {
			return PACKET_ID_SHOULD_BE_ZERO
		}
	case 1:
		puback := NewPubackMessage(m.PacketID)
		self.WriteChan <- puback
	case 2:
		pubrec := NewPubrecMessage(m.PacketID)
		self.WriteChan <- pubrec
	}
	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
	}
	pubrel := NewPubrelMessage(m.PacketID)
	self.WriteChan <- pubrel
	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
	}
	pubcomp := NewPubcompMessage(m.PacketID)
	self.WriteChan <- pubcomp
	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
					var id uint16
					if edge.RetainQoS > 0 {
						id, err = self.getUsablePacketID
						if err != nil {
							return err
						}
					}
					qos := edge.RetainQoS
					if subTopic.QoS < edge.RetainQoS {
						qos = subTopic.QoS //downgrade the QoS
					}
					pub := NewPublishMessage(false, qos, true, edge.FullPath, id, []uint8(edge.RetainMessage))
					self.WriteChan <- pub
					EmitError(err)
				}
			}
		}
		returnCodes = append(returnCodes, codes...)
	}
	// TODO: check whether the number of return codes are correct?
	suback := NewSubackMessage(m.PacketID, returnCodes)
	self.WriteChan <- suback
	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
	unsuback := NewUnsubackMessage(m.PacketID)

	self.WriteChan <- unsuback
	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
	pingresp := NewPingrespMessage()
	self.WriteChan <- pingresp
	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.Will = nil
	self.disconnectProcessing()
	// close the client
	return err
}
<|endoftext|>


<commit_before>package editor

import (
	"fmt"
	"strings"

	"github.com/elpinal/coco3/screen"
)

type exCommand struct {
	name string
	fn   func(*commandline, []string) continuity
}

// exComands represents a table of Ex commands and corresponding functions.
// The order is important. Precede commands have higher precedence.
var exCommands = []exCommand{
	{"help", (*commandline).help},
	{"delete", (*commandline).delete},
	{"quit", (*commandline).quit},
	{"substitute", (*commandline).substitute},
}

type commandline struct {
	streamSet
	*editor

	basic *basic
}

func newCommandline(s streamSet, e *editor) *commandline {
	return &commandline{
		streamSet: s,
		editor:    e,
		basic:     &basic{},
	}
}

func (e *commandline) Mode() mode {
	return modeCommandline
}

func (e *commandline) Position() int {
	return e.basic.pos + 1
}

func (e *commandline) Runes() []rune {
	return e.buf
}

func (e *commandline) Message() []rune {
	return append([]rune{':'}, e.basic.buf...)
}

func (e *commandline) Highlight() *screen.Hi {
	return nil
}

func (e *commandline) Run() (end continuity, next modeChanger, err error) {
	r, _, err := e.streamSet.in.ReadRune()
	if err != nil {
		return end, next, err
	}
	switch r {
	case CharCtrlM, CharCtrlJ:
	case CharEscape, CharCtrlC:
		next = norm()
		return end, next, err
	case CharBackspace, CharCtrlH:
		if len(e.basic.buf) == 0 {
			next = norm()
			return
		}
		e.basic.delete(e.basic.pos-1, e.basic.pos)
	case CharCtrlB:
		e.basic.move(0)
	case CharCtrlE:
		e.basic.move(len(e.basic.buf))
	case CharCtrlU:
		e.basic.delete(0, e.basic.pos)
	case CharCtrlW:
		// FIXME: It's redundant.
		ed := newEditor()
		ed.pos = e.basic.pos
		ed.buf = e.basic.buf
		pos := ed.pos
		ed.wordBackward()
		e.basic.delete(pos, ed.pos)
		return
	default:
		e.basic.insert([]rune{r}, e.basic.pos)
	}
	if r != CharCtrlM && r != CharCtrlJ {
		return
	}
	next = norm()
	var candidate exCommand
	s := string(e.basic.buf)
	if s == "" {
		return
	}
	args := strings.Split(s, " ")
	s = args[0]
	for _, cmd := range exCommands {
		if !strings.HasPrefix(cmd.name, s) {
			continue
		}
		if cmd.name == s {
			end = cmd.fn(e, nil)
			return
		}
		if candidate.name == "" {
			candidate = cmd
		}
	}
	if candidate.name != "" {
		end = candidate.fn(e, nil)
		return
	}
	err = fmt.Errorf("not a command: %q", s)
	return
}

func (e *commandline) quit(args []string) continuity {
	return exit
}

func (e *commandline) delete(args []string) (_ continuity) {
	e.editor.delete(0, len(e.editor.buf))
	return
}

func (e *commandline) help(args []string) continuity {
	e.buf = []rune("help")
	e.pos = 4
	return execute
}

func (e *commandline) substitute(args []string) (_ continuity) {
	if len(args) != 2 {
		return
	}
	return
}
<commit_msg>Pass arguments to ex-command<commit_after>package editor

import (
	"fmt"
	"strings"

	"github.com/elpinal/coco3/screen"
)

type exCommand struct {
	name string
	fn   func(*commandline, []string) continuity
}

// exComands represents a table of Ex commands and corresponding functions.
// The order is important. Precede commands have higher precedence.
var exCommands = []exCommand{
	{"help", (*commandline).help},
	{"delete", (*commandline).delete},
	{"quit", (*commandline).quit},
	{"substitute", (*commandline).substitute},
}

type commandline struct {
	streamSet
	*editor

	basic *basic
}

func newCommandline(s streamSet, e *editor) *commandline {
	return &commandline{
		streamSet: s,
		editor:    e,
		basic:     &basic{},
	}
}

func (e *commandline) Mode() mode {
	return modeCommandline
}

func (e *commandline) Position() int {
	return e.basic.pos + 1
}

func (e *commandline) Runes() []rune {
	return e.buf
}

func (e *commandline) Message() []rune {
	return append([]rune{':'}, e.basic.buf...)
}

func (e *commandline) Highlight() *screen.Hi {
	return nil
}

func (e *commandline) Run() (end continuity, next modeChanger, err error) {
	r, _, err := e.streamSet.in.ReadRune()
	if err != nil {
		return end, next, err
	}
	switch r {
	case CharCtrlM, CharCtrlJ:
	case CharEscape, CharCtrlC:
		next = norm()
		return end, next, err
	case CharBackspace, CharCtrlH:
		if len(e.basic.buf) == 0 {
			next = norm()
			return
		}
		e.basic.delete(e.basic.pos-1, e.basic.pos)
	case CharCtrlB:
		e.basic.move(0)
	case CharCtrlE:
		e.basic.move(len(e.basic.buf))
	case CharCtrlU:
		e.basic.delete(0, e.basic.pos)
	case CharCtrlW:
		// FIXME: It's redundant.
		ed := newEditor()
		ed.pos = e.basic.pos
		ed.buf = e.basic.buf
		pos := ed.pos
		ed.wordBackward()
		e.basic.delete(pos, ed.pos)
		return
	default:
		e.basic.insert([]rune{r}, e.basic.pos)
	}
	if r != CharCtrlM && r != CharCtrlJ {
		return
	}
	next = norm()
	var candidate exCommand
	s := string(e.basic.buf)
	if s == "" {
		return
	}
	args := strings.Split(s, " ")
	s = args[0]
	args = args[1:]
	for _, cmd := range exCommands {
		if !strings.HasPrefix(cmd.name, s) {
			continue
		}
		if cmd.name == s {
			end = cmd.fn(e, args)
			return
		}
		if candidate.name == "" {
			candidate = cmd
		}
	}
	if candidate.name != "" {
		end = candidate.fn(e, args)
		return
	}
	err = fmt.Errorf("not a command: %q", s)
	return
}

func (e *commandline) quit(args []string) continuity {
	return exit
}

func (e *commandline) delete(args []string) (_ continuity) {
	e.editor.delete(0, len(e.editor.buf))
	return
}

func (e *commandline) help(args []string) continuity {
	e.buf = []rune("help")
	e.pos = 4
	return execute
}

func (e *commandline) substitute(args []string) (_ continuity) {
	if len(args) != 2 {
		return
	}
	return
}
<|endoftext|>


<commit_before>package common

import (
	"fmt"
	"strings"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
)

//DescribeVpcs mocks an ec2.DescribeVpcsOutput for a given input
func (m *mockEC2Conn) DescribeVpcs(input *ec2.DescribeVpcsInput) (*ec2.DescribeVpcsOutput, error) {
	m.lock.Lock()
	m.copyImageCount++
	m.lock.Unlock()

	if input == nil || len(input.VpcIds) == 0 {
		return nil, fmt.Errorf("oops looks like we need more input")
	}

	var isDefault bool
	vpcID := aws.StringValue(input.VpcIds[0])

	//only one default VPC per region
	if strings.Contains("vpc-default-id", vpcID) {
		isDefault = true
	}

	output := &ec2.DescribeVpcsOutput{
		Vpcs: []*ec2.Vpc{
			&ec2.Vpc{IsDefault: aws.Bool(isDefault),
				VpcId: aws.String(vpcID),
			},
		},
	}
	return output, nil
}

func TestStepPreValidate_checkVpc(t *testing.T) {
	tt := []struct {
		name          string
		step          StepPreValidate
		errorExpected bool
	}{
		{"DefaultVpc", StepPreValidate{VpcId: "vpc-default-id"}, false},
		{"NonDefaultVpcNoSubnet", StepPreValidate{VpcId: "vpc-1234567890"}, true},
		{"NonDefaultVpcWithSubnet", StepPreValidate{VpcId: "vpc-1234567890", SubnetId: "subnet-1234567890"}, false},
		{"SubnetWithNoVpc", StepPreValidate{SubnetId: "subnet-1234567890"}, false},
		{"NoVpcInformation", StepPreValidate{}, false},
	}

	mockConn, err := getMockConn(nil, "")
	if err != nil {
		t.Fatal("unable to get a mock connection")
	}

	for _, tc := range tt {
		t.Run(tc.name, func(t *testing.T) {
			err := tc.step.checkVpc(mockConn)

			if tc.errorExpected && err == nil {
				t.Errorf("expected a validation error for %q but got %q", tc.name, err)
			}

			if !tc.errorExpected && err != nil {
				t.Errorf("expected a validation to pass for %q but got %q", tc.name, err)
			}
		})
	}

}
<commit_msg>test/builder/amazon/common/step_pre_validate: Update DescribeVpcs mock<commit_after>package common

import (
	"fmt"
	"strings"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
)

//DescribeVpcs mocks an ec2.DescribeVpcsOutput for a given input
func (m *mockEC2Conn) DescribeVpcs(input *ec2.DescribeVpcsInput) (*ec2.DescribeVpcsOutput, error) {

	if input == nil || aws.StringValue(input.VpcIds[0]) == "" {
		return nil, fmt.Errorf("oops looks like we need more input")
	}

	var isDefault bool
	vpcID := aws.StringValue(input.VpcIds[0])

	//only one default VPC per region
	if strings.Contains("vpc-default-id", vpcID) {
		isDefault = true
	}

	output := &ec2.DescribeVpcsOutput{
		Vpcs: []*ec2.Vpc{
			&ec2.Vpc{IsDefault: aws.Bool(isDefault),
				VpcId: aws.String(vpcID),
			},
		},
	}
	return output, nil
}

func TestStepPreValidate_checkVpc(t *testing.T) {
	tt := []struct {
		name          string
		step          StepPreValidate
		errorExpected bool
	}{
		{"DefaultVpc", StepPreValidate{VpcId: "vpc-default-id"}, false},
		{"NonDefaultVpcNoSubnet", StepPreValidate{VpcId: "vpc-1234567890"}, true},
		{"NonDefaultVpcWithSubnet", StepPreValidate{VpcId: "vpc-1234567890", SubnetId: "subnet-1234567890"}, false},
		{"SubnetWithNoVpc", StepPreValidate{SubnetId: "subnet-1234567890"}, false},
		{"NoVpcInformation", StepPreValidate{}, false},
	}

	mockConn, err := getMockConn(nil, "")
	if err != nil {
		t.Fatal("unable to get a mock connection")
	}

	for _, tc := range tt {
		t.Run(tc.name, func(t *testing.T) {
			err := tc.step.checkVpc(mockConn)

			if tc.errorExpected && err == nil {
				t.Errorf("expected a validation error for %q but got %q", tc.name, err)
			}

			if !tc.errorExpected && err != nil {
				t.Errorf("expected a validation to pass for %q but got %q", tc.name, err)
			}
		})
	}

}
<|endoftext|>


<commit_before>package boomer

import (
	"os"
	"testing"
	"time"
)

func TestRunTasksForTest(t *testing.T) {
	count := 0
	taskA := &Task{
		Name: "increaseCount",
		Fn: func() {
			count++
		},
	}
	runTasks = "increaseCount,foobar"
	runTasksForTest(taskA)

	if count != 1 {
		t.Error("count is", count, "expected: 1")
	}
}

func TestStartMemoryProfile(t *testing.T) {
	if _, err := os.Stat("mem.pprof"); os.IsExist(err) {
		os.Remove("mem.pprof")
	}
	startMemoryProfile("mem.pprof", 3*time.Second)
	time.Sleep(4 * time.Second)
	if _, err := os.Stat("mem.pprof"); os.IsNotExist(err) {
		t.Error("File mem.pprof is not generated")
	} else {
		os.Remove("mem.pprof")
	}
}

func TestStartCPUProfile(t *testing.T) {
	if _, err := os.Stat("cpu.pprof"); os.IsExist(err) {
		os.Remove("cpu.pprof")
	}
	startCPUProfile("cpu.pprof", 3*time.Second)
	time.Sleep(4 * time.Second)
	if _, err := os.Stat("cpu.pprof"); os.IsNotExist(err) {
		t.Error("File cpu.pprof is not generated")
	} else {
		os.Remove("cpu.pprof")
	}
}
<commit_msg>add test for initBoomer<commit_after>package boomer

import (
	"os"
	"testing"
	"time"
)

func TestInitBoomer(t *testing.T) {
	initBoomer()
	defer Events.Unsubscribe("request_success", requestSuccessHandler)
	defer Events.Unsubscribe("request_failure", requestFailureHandler)

	defer func() {
		err := recover()
		if err == nil {
			t.Error("It should panic if initBoomer is called more than once.")
		}
	}()
	initBoomer()
}

func TestRunTasksForTest(t *testing.T) {
	count := 0
	taskA := &Task{
		Name: "increaseCount",
		Fn: func() {
			count++
		},
	}
	runTasks = "increaseCount,foobar"
	runTasksForTest(taskA)

	if count != 1 {
		t.Error("count is", count, "expected: 1")
	}
}

func TestStartMemoryProfile(t *testing.T) {
	if _, err := os.Stat("mem.pprof"); os.IsExist(err) {
		os.Remove("mem.pprof")
	}
	startMemoryProfile("mem.pprof", 3*time.Second)
	time.Sleep(4 * time.Second)
	if _, err := os.Stat("mem.pprof"); os.IsNotExist(err) {
		t.Error("File mem.pprof is not generated")
	} else {
		os.Remove("mem.pprof")
	}
}

func TestStartCPUProfile(t *testing.T) {
	if _, err := os.Stat("cpu.pprof"); os.IsExist(err) {
		os.Remove("cpu.pprof")
	}
	startCPUProfile("cpu.pprof", 3*time.Second)
	time.Sleep(4 * time.Second)
	if _, err := os.Stat("cpu.pprof"); os.IsNotExist(err) {
		t.Error("File cpu.pprof is not generated")
	} else {
		os.Remove("cpu.pprof")
	}
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Primitive HTTP client. See RFC 2616.

package http

import (
	"bufio"
	"encoding/base64"
	"fmt"
	"io"
	"net"
	"os"
	"strings"
)

// Given a string of the form "host", "host:port", or "[ipv6::address]:port",
// return true if the string includes a port.
func hasPort(s string) bool { return strings.LastIndex(s, ":") > strings.LastIndex(s, "]") }

// Used in Send to implement io.ReadCloser by bundling together the
// io.BufReader through which we read the response, and the underlying
// network connection.
type readClose struct {
	io.Reader
	io.Closer
}

// Send issues an HTTP request.  Caller should close resp.Body when done reading it.
//
// TODO: support persistent connections (multiple requests on a single connection).
// send() method is nonpublic because, when we refactor the code for persistent
// connections, it may no longer make sense to have a method with this signature.
func send(req *Request) (resp *Response, err os.Error) {
	if req.URL.Scheme != "http" {
		return nil, &badStringError{"unsupported protocol scheme", req.URL.Scheme}
	}

	addr := req.URL.Host
	if !hasPort(addr) {
		addr += ":http"
	}
	info := req.URL.Userinfo
	if len(info) > 0 {
		enc := base64.URLEncoding
		encoded := make([]byte, enc.EncodedLen(len(info)))
		enc.Encode(encoded, strings.Bytes(info))
		if req.Header == nil {
			req.Header = make(map[string]string)
		}
		req.Header["Authorization"] = "Basic " + string(encoded)
	}
	conn, err := net.Dial("tcp", "", addr)
	if err != nil {
		return nil, err
	}

	err = req.Write(conn)
	if err != nil {
		conn.Close()
		return nil, err
	}

	reader := bufio.NewReader(conn)
	resp, err = ReadResponse(reader, req.Method)
	if err != nil {
		conn.Close()
		return nil, err
	}

	resp.Body = readClose{resp.Body, conn}

	return
}

// True if the specified HTTP status code is one for which the Get utility should
// automatically redirect.
func shouldRedirect(statusCode int) bool {
	switch statusCode {
	case StatusMovedPermanently, StatusFound, StatusSeeOther, StatusTemporaryRedirect:
		return true
	}
	return false
}

// Get issues a GET to the specified URL.  If the response is one of the following
// redirect codes, it follows the redirect, up to a maximum of 10 redirects:
//
//    301 (Moved Permanently)
//    302 (Found)
//    303 (See Other)
//    307 (Temporary Redirect)
//
// finalURL is the URL from which the response was fetched -- identical to the
// input URL unless redirects were followed.
//
// Caller should close r.Body when done reading it.
func Get(url string) (r *Response, finalURL string, err os.Error) {
	// TODO: if/when we add cookie support, the redirected request shouldn't
	// necessarily supply the same cookies as the original.
	// TODO: set referrer header on redirects.
	for redirect := 0; ; redirect++ {
		if redirect >= 10 {
			err = os.ErrorString("stopped after 10 redirects")
			break
		}

		var req Request
		if req.URL, err = ParseURL(url); err != nil {
			break
		}
		if r, err = send(&req); err != nil {
			break
		}
		if shouldRedirect(r.StatusCode) {
			r.Body.Close()
			if url = r.GetHeader("Location"); url == "" {
				err = os.ErrorString(fmt.Sprintf("%d response missing Location header", r.StatusCode))
				break
			}
			continue
		}
		finalURL = url
		return
	}

	err = &URLError{"Get", url, err}
	return
}

// Post issues a POST to the specified URL.
//
// Caller should close r.Body when done reading it.
func Post(url string, bodyType string, body io.Reader) (r *Response, err os.Error) {
	var req Request
	req.Method = "POST"
	req.Body = nopCloser{body}
	req.Header = map[string]string{
		"Content-Type": bodyType,
		"Transfer-Encoding": "chunked",
	}

	req.URL, err = ParseURL(url)
	if err != nil {
		return nil, err
	}

	return send(&req)
}

type nopCloser struct {
	io.Reader
}

func (nopCloser) Close() os.Error { return nil }
<commit_msg>http: fix bug in Post<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Primitive HTTP client. See RFC 2616.

package http

import (
	"bufio"
	"encoding/base64"
	"fmt"
	"io"
	"net"
	"os"
	"strings"
)

// Given a string of the form "host", "host:port", or "[ipv6::address]:port",
// return true if the string includes a port.
func hasPort(s string) bool { return strings.LastIndex(s, ":") > strings.LastIndex(s, "]") }

// Used in Send to implement io.ReadCloser by bundling together the
// io.BufReader through which we read the response, and the underlying
// network connection.
type readClose struct {
	io.Reader
	io.Closer
}

// Send issues an HTTP request.  Caller should close resp.Body when done reading it.
//
// TODO: support persistent connections (multiple requests on a single connection).
// send() method is nonpublic because, when we refactor the code for persistent
// connections, it may no longer make sense to have a method with this signature.
func send(req *Request) (resp *Response, err os.Error) {
	if req.URL.Scheme != "http" {
		return nil, &badStringError{"unsupported protocol scheme", req.URL.Scheme}
	}

	addr := req.URL.Host
	if !hasPort(addr) {
		addr += ":http"
	}
	info := req.URL.Userinfo
	if len(info) > 0 {
		enc := base64.URLEncoding
		encoded := make([]byte, enc.EncodedLen(len(info)))
		enc.Encode(encoded, strings.Bytes(info))
		if req.Header == nil {
			req.Header = make(map[string]string)
		}
		req.Header["Authorization"] = "Basic " + string(encoded)
	}
	conn, err := net.Dial("tcp", "", addr)
	if err != nil {
		return nil, err
	}

	err = req.Write(conn)
	if err != nil {
		conn.Close()
		return nil, err
	}

	reader := bufio.NewReader(conn)
	resp, err = ReadResponse(reader, req.Method)
	if err != nil {
		conn.Close()
		return nil, err
	}

	resp.Body = readClose{resp.Body, conn}

	return
}

// True if the specified HTTP status code is one for which the Get utility should
// automatically redirect.
func shouldRedirect(statusCode int) bool {
	switch statusCode {
	case StatusMovedPermanently, StatusFound, StatusSeeOther, StatusTemporaryRedirect:
		return true
	}
	return false
}

// Get issues a GET to the specified URL.  If the response is one of the following
// redirect codes, it follows the redirect, up to a maximum of 10 redirects:
//
//    301 (Moved Permanently)
//    302 (Found)
//    303 (See Other)
//    307 (Temporary Redirect)
//
// finalURL is the URL from which the response was fetched -- identical to the
// input URL unless redirects were followed.
//
// Caller should close r.Body when done reading it.
func Get(url string) (r *Response, finalURL string, err os.Error) {
	// TODO: if/when we add cookie support, the redirected request shouldn't
	// necessarily supply the same cookies as the original.
	// TODO: set referrer header on redirects.
	for redirect := 0; ; redirect++ {
		if redirect >= 10 {
			err = os.ErrorString("stopped after 10 redirects")
			break
		}

		var req Request
		if req.URL, err = ParseURL(url); err != nil {
			break
		}
		if r, err = send(&req); err != nil {
			break
		}
		if shouldRedirect(r.StatusCode) {
			r.Body.Close()
			if url = r.GetHeader("Location"); url == "" {
				err = os.ErrorString(fmt.Sprintf("%d response missing Location header", r.StatusCode))
				break
			}
			continue
		}
		finalURL = url
		return
	}

	err = &URLError{"Get", url, err}
	return
}

// Post issues a POST to the specified URL.
//
// Caller should close r.Body when done reading it.
func Post(url string, bodyType string, body io.Reader) (r *Response, err os.Error) {
	var req Request
	req.Method = "POST"
	req.ProtoMajor = 1
	req.ProtoMinor = 1
	req.Body = nopCloser{body}
	req.Header = map[string]string{
		"Content-Type": bodyType,
	}
	req.TransferEncoding = []string{"chunked"}

	req.URL, err = ParseURL(url)
	if err != nil {
		return nil, err
	}

	return send(&req)
}

type nopCloser struct {
	io.Reader
}

func (nopCloser) Close() os.Error { return nil }
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"flag"
	"fmt"
	_ "github.com/mattn/go-sqlite3"
	"github.com/paulkramme/toml"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"
)

func main() {
	fmt.Println("BTSOOT - Copyright (c) 2017 Paul Kramme All Rights Reserved.")

	ConfigLocation := flag.String("config", "./btsoot.conf", "Specifies configfile location")
	flag.Parse()

	var Config Configuration

	_, err := toml.DecodeFile(*ConfigLocation, &Config)
	if err != nil {
		fmt.Println("Couldn't find or open config file.")
		panic(err)
	}

	f, err := os.OpenFile(Config.LogFileLocation, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		log.Fatalln(err)
	}
	defer f.Close()
	log.SetOutput(f)

	// NOTE: Create a spacer in the log
	log.Println("\n\n\n\n\n")

	db, err := sql.Open("sqlite3", Config.DBFileLocation)
	if err != nil {
		log.Fatal(err)
	}
	err = db.Ping()
	if err != nil {
		panic(err)
	}
	log.Println("Database connection established.")

	ProcessList := CreateMasterProcessList()

	// NOTE: Init standard threads...
	go UpdateProcess(ProcessList[UpdateThreadID], Config)
	go WebServer(ProcessList[WebserverThreadID], Config)
	go ScanningProcess(ProcessList[ScanThreadID], Config)
	signals := make(chan os.Signal)

	log.Println("Startup complete...")

	// NOTE: Wait for SIGINT
	signal.Notify(signals, syscall.SIGINT)
	<-signals
	log.Println("Exiting...")
	KillAll(ProcessList)
	time.Sleep(1 * time.Second) // wait for scanners
}
<commit_msg>Add sql database inits<commit_after>package main

import (
	"database/sql"
	"flag"
	"fmt"
	_ "github.com/mattn/go-sqlite3"
	"github.com/paulkramme/toml"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"
)

func main() {
	fmt.Println("BTSOOT - Copyright (c) 2017 Paul Kramme All Rights Reserved.")

	ConfigLocation := flag.String("config", "./btsoot.conf", "Specifies configfile location")
	flag.Parse()

	var Config Configuration

	_, err := toml.DecodeFile(*ConfigLocation, &Config)
	if err != nil {
		fmt.Println("Couldn't find or open config file.")
		panic(err)
	}

	f, err := os.OpenFile(Config.LogFileLocation, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		log.Fatalln(err)
	}
	defer f.Close()
	log.SetOutput(f)

	// NOTE: Create a spacer in the log
	log.Println("\n\n\n\n\n")

	db, err := sql.Open("sqlite3", Config.DBFileLocation)
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()
	err = db.Ping()
	if err != nil {
		panic(err)
	}

	err = DatabaseSetup(db)
	if err != nil {
		panic(err)
	}

	ProcessList := CreateMasterProcessList()

	// NOTE: Init standard threads...
	go UpdateProcess(ProcessList[UpdateThreadID], Config)
	go WebServer(ProcessList[WebserverThreadID], Config)
	go ScanningProcess(ProcessList[ScanThreadID], Config)
	signals := make(chan os.Signal)

	log.Println("Startup complete")
	fmt.Println("Startup complete")

	// NOTE: Wait for SIGINT
	signal.Notify(signals, syscall.SIGINT)
	<-signals
	log.Println("SIGINT received")
	KillAll(ProcessList)
	time.Sleep(1 * time.Second) // wait for scanners
}

func DatabaseSetup(db *sql.DB) (err error) {
	_, err = db.Exec(`
		CREATE TABLE IF NOT EXISTS blocks(
			name text PRIMARY KEY NOT NULL,
			interval VARCHAR,
			path text NOT NULL,
			destination text NOT NULL,
			absolutepath BOOL)`)
	if err != nil {
		return err
	}
	_, err = db.Exec(`
		CREATE TABLE IF NOT EXISTS files(
			block text NOT NULL,
			checksum TEXT NOT NULL,
			filename text,
			absolutepath text,
			FOREIGN KEY (block) REFERENCES blocks(name))`)
	if err != nil {
		return err
	}
	_, err = db.Exec(
		`CREATE TABLE IF NOT EXISTS users(
			username text,
			password hash)`)
	if err != nil {
		return err
	}
	/*
	_, err = db.Exec("CREATE TABLE IF NOT EXISTS version")
	if err != nil {
		return err
	}
	*/
	return
}
<|endoftext|>


<commit_before>package editor

import (
	"fmt"
	"strings"

	"github.com/elpinal/coco3/screen"
)

type exCommand struct {
	name string
	fn   func(*commandline, []string) continuity
}

// exComands represents a table of Ex commands and corresponding functions.
// The order is important. Precede commands have higher precedence.
var exCommands = []exCommand{
	{"quit", (*commandline).quit},
}

type commandline struct {
	streamSet
	*editor

	basic *basic
}

func (e *commandline) Mode() mode {
	return modeCommandline
}

func (e *commandline) Position() int {
	return e.basic.pos + 1
}

func (e *commandline) Runes() []rune {
	return e.buf
}

func (e *commandline) Message() []rune {
	return append([]rune{':'}, e.basic.buf...)
}

func (e *commandline) Highlight() *screen.Hi {
	return nil
}

func (e *commandline) Run() (end continuity, next mode, err error) {
	next = modeCommandline
	r, _, err := e.streamSet.in.ReadRune()
	if err != nil {
		return end, next, err
	}
	switch r {
	case CharCtrlM:
	case CharEscape:
		return end, modeNormal, err
	case CharBackspace, CharCtrlH:
		if len(e.basic.buf) == 0 {
			next = modeNormal
			return
		}
		e.basic.delete(e.basic.pos-1, e.basic.pos)
	default:
		e.basic.insert([]rune{r}, e.basic.pos)
	}
	if r != CharCtrlM {
		return
	}
	next = modeNormal
	var candidate exCommand
	s := string(e.basic.buf)
	for _, cmd := range exCommands {
		if !strings.HasPrefix(cmd.name, s) {
			continue
		}
		if cmd.name == s {
			end = cmd.fn(e, nil)
			return
		}
		if candidate.name == "" {
			candidate = cmd
		}
	}
	if candidate.name != "" {
		end = candidate.fn(e, nil)
		return
	}
	err = fmt.Errorf("not a command: %q", s)
	return
}

func (e *commandline) quit(args []string) continuity {
	return exit
}
<commit_msg>Implement c_CTRL-B<commit_after>package editor

import (
	"fmt"
	"strings"

	"github.com/elpinal/coco3/screen"
)

type exCommand struct {
	name string
	fn   func(*commandline, []string) continuity
}

// exComands represents a table of Ex commands and corresponding functions.
// The order is important. Precede commands have higher precedence.
var exCommands = []exCommand{
	{"quit", (*commandline).quit},
}

type commandline struct {
	streamSet
	*editor

	basic *basic
}

func (e *commandline) Mode() mode {
	return modeCommandline
}

func (e *commandline) Position() int {
	return e.basic.pos + 1
}

func (e *commandline) Runes() []rune {
	return e.buf
}

func (e *commandline) Message() []rune {
	return append([]rune{':'}, e.basic.buf...)
}

func (e *commandline) Highlight() *screen.Hi {
	return nil
}

func (e *commandline) Run() (end continuity, next mode, err error) {
	next = modeCommandline
	r, _, err := e.streamSet.in.ReadRune()
	if err != nil {
		return end, next, err
	}
	switch r {
	case CharCtrlM:
	case CharEscape:
		return end, modeNormal, err
	case CharBackspace, CharCtrlH:
		if len(e.basic.buf) == 0 {
			next = modeNormal
			return
		}
		e.basic.delete(e.basic.pos-1, e.basic.pos)
	case CharCtrlB:
		e.basic.move(0)
	default:
		e.basic.insert([]rune{r}, e.basic.pos)
	}
	if r != CharCtrlM {
		return
	}
	next = modeNormal
	var candidate exCommand
	s := string(e.basic.buf)
	for _, cmd := range exCommands {
		if !strings.HasPrefix(cmd.name, s) {
			continue
		}
		if cmd.name == s {
			end = cmd.fn(e, nil)
			return
		}
		if candidate.name == "" {
			candidate = cmd
		}
	}
	if candidate.name != "" {
		end = candidate.fn(e, nil)
		return
	}
	err = fmt.Errorf("not a command: %q", s)
	return
}

func (e *commandline) quit(args []string) continuity {
	return exit
}
<|endoftext|>


<commit_before>package bot

import (
	"fmt"
	"github.com/darkliquid/go-ircevent"
	"github.com/darkliquid/leader1/utils"
	"github.com/darkliquid/leader1/debug"
	"sort"
	"strings"
	"time"
)

func (bot *Bot) RunBuiltinCommands(event *irc.Event) {
	args := strings.Split(strings.TrimSpace(event.Message()), " ")
	command := args[0]

	// Bin the command from the arg list
	if len(args) > 1 {
		args = args[1:]
	} else {
		args = []string{}
	}

	// Get the current channel privileges for the nick sending this command
	privs, ok := bot.state.GetPrivs(event.Arguments[0], event.Nick)

	// Commands must be run by known users
	switch {
	case ok && command == "!reload":
		if privs.Owner || privs.Admin || privs.Op {
			bot.pm.InitJS()
			utils.IRCAction(bot.conn, event.Arguments[0], "has reloaded its plugins")
		} else {
			utils.IRCAction(bot.conn, event.Arguments[0], fmt.Sprintf("slaps %s's hands away from the op only controls", event.Nick))
		}
	case command == "!ping":
		bot.conn.Privmsg(event.Arguments[0], fmt.Sprintf("%s: PONG!", event.Nick))
	case ok && command == "!quit":
		if privs.Owner || privs.Admin || privs.Op {
			bot.Quit()
		} else {
			utils.IRCAction(bot.conn, event.Arguments[0], fmt.Sprintf("slaps %s's hands away from the op only controls", event.Nick))
		}
	case command == "!help":
		if len(args) == 0 {
			bot.ShowCommandList(event.Arguments[0], event.Nick)
		} else {
			bot.ShowCommandHelp(event.Arguments[0], event.Nick, args[0])
		}
	case ok && command == "!import" && len(args) >= 2 && bot.cfg.Irc.StaffChannel == event.Arguments[0]:
		if privs.Owner {
			overwrite := false
			if len(args) == 3 {
				overwrite = args[2] == "overwrite"
			}
			if err := bot.pm.ImportPlugin(event.Nick, args[0], args[1], overwrite); err != nil {
				bot.conn.Privmsg(bot.cfg.Irc.StaffChannel, fmt.Sprintf("ALERT: %s tried to use !import with %s and got this error: %s", event.Nick, args[0], err.Error()))
			} else {
				bot.conn.Privmsg(bot.cfg.Irc.StaffChannel, fmt.Sprintf("ALERT: %s successfully used !import with %s", event.Nick, args[0]))
			}
		} else {
			utils.IRCAction(bot.conn, event.Arguments[0], fmt.Sprintf("slaps %s's hands away from the op only controls", event.Nick))
		}
	case ok && command == "!debug" && bot.cfg.Irc.StaffChannel == event.Arguments[0]:
		if privs.Owner || privs.Admin || privs.Op {
			switch {
			case len(args) > 0 && args[0] == "on":
				port, alreadyRunning := debug.StartDebugServer()
				if alreadyRunning {
					bot.conn.Privmsg(bot.cfg.Irc.StaffChannel, fmt.Sprintf("%s: Debug server running on port %s", event.Nick, port))
				} else {
					bot.conn.Privmsg(bot.cfg.Irc.StaffChannel, fmt.Sprintf("%s: Debug server started on port %s", event.Nick, port))
				}
			case len(args) > 0 && args[0] == "off":
				debug.StopDebugServer()
				bot.conn.Privmsg(bot.cfg.Irc.StaffChannel, fmt.Sprintf("%s: Debug server stopped", event.Nick))
			case len(args) > 0 && args[0] == "status":
				status := debug.DebugServerStatus()
				bot.conn.Privmsg(bot.cfg.Irc.StaffChannel, fmt.Sprintf("%s: Debug server is %s", event.Nick, status))
			default:
				bot.conn.Privmsg(bot.cfg.Irc.StaffChannel, fmt.Sprintf("%s: usage - !debug [on|off]", event.Nick))
			}
		} else {
			utils.IRCAction(bot.conn, event.Arguments[0], fmt.Sprintf("slaps %s's hands away from the op only controls", event.Nick))
		}
	}
}

// Handler for reclaiming a stolen nick
func (bot *Bot) ReclaimNick(event *irc.Event) {
	if thief := bot.state.GetNick(bot.cfg.Irc.Nick); thief != nil {
		// Recover nick from thieves
		bot.conn.Privmsg("NickServ", fmt.Sprintf("RECOVER %s %s", bot.cfg.Irc.Nick, bot.cfg.Irc.NickPass))
		time.Sleep(time.Second)
		bot.conn.Privmsg("NickServ", fmt.Sprintf("RELEASE %s %s", bot.cfg.Irc.Nick, bot.cfg.Irc.NickPass))
	}
	bot.SetBotState(event)
}

// Automatically give voice to users in channels in which the bot has Op
func (bot *Bot) AutoVoice(event *irc.Event) {
	if bot.cfg.Irc.AutoVoice {
		time.Sleep(time.Second) // Wait a second before we bother

		privs, ok := bot.state.GetPrivs(event.Arguments[0], event.Nick)

		// No need to autovoice
		if ok && (privs.Owner || privs.Admin || privs.Op || privs.HalfOp || privs.Voice) {
			return
		}

		// Set mode
		bot.conn.Mode(event.Arguments[0], fmt.Sprintf("+v %s", event.Nick))
	}
}

// Function for setting up the botstate
func (bot *Bot) SetBotState(event *irc.Event) {
	if ghost := bot.state.GetNick(bot.cfg.Irc.Nick); ghost != nil && ghost != bot.state.Me() {
		// GHOST the old nick
		bot.conn.Privmsg("NickServ", fmt.Sprintf("GHOST %s %s", bot.cfg.Irc.Nick, bot.cfg.Irc.NickPass))
	}

	// Set up the nick
	bot.conn.Nick(bot.cfg.Irc.Nick)

	// Identify as the nick owner
	bot.conn.Privmsg("NickServ", fmt.Sprintf("IDENTIFY %s", bot.cfg.Irc.NickPass))

	// Tell IRC I'm a bot
	bot.conn.Mode(bot.cfg.Irc.Nick, "+B")
}

// Function for re-joining channels
func (bot *Bot) JoinChannels(event *irc.Event) {
	time.Sleep(time.Second) // Wait a second before we bother

	if _, ok := bot.state.GetPrivs(bot.cfg.Irc.NormalChannel, bot.state.Me().Nick); !ok {
		bot.conn.Join(bot.cfg.Irc.NormalChannel)
	}

	if _, ok := bot.state.GetPrivs(bot.cfg.Irc.StaffChannel, bot.state.Me().Nick); !ok {
		bot.conn.Join(bot.cfg.Irc.StaffChannel)
	}
}

// Print out the commands available
func (bot *Bot) ShowCommandList(source, nick string) {
	var commands []string = make([]string, 0)
	commands = append(commands, "!reload", "!ping", "!quit", "!help", "!import", "!debug")

	for cmd, _ := range bot.pm.CommandHelp() {
		commands = append(commands, cmd)
	}
	sort.Strings(commands)
	bot.conn.Privmsg(source, fmt.Sprintf("%s: available commands are - %s", nick, strings.Join(commands, ", ")))
}

// Print out the commands available
func (bot *Bot) ShowCommandHelp(source, nick, cmd string) {
	message := "unknown command, run `!help` to see what commands are available."
	switch cmd {
	case "!ping":
		message = fmt.Sprintf("makes `%s` reply with PONG!", bot.state.Me().Nick)
	case "!reload":
		message = "reloads the plugins"
	case "!quit":
		message = fmt.Sprintf("makes `%s` quit IRC", bot.state.Me().Nick)
	case "!help":
		message = "shows this message, smart ass"
	case "!import":
		message = "call like !import [url] [name] [overwrite], imports the plugin at [url] into [name].js and loads it into the bot. Will not overwrite unless [overwrite] is set to 'overwrite'"
	case "!debug":
		message = "starts/stops debugging server for inspecting the bot"
	default:
		if help, ok := bot.pm.CommandHelp()[cmd]; ok {
			message = help
		}
	}

	bot.conn.Privmsg(source, fmt.Sprintf("%s: %s - %s", nick, cmd, message))
}
<commit_msg>Minor tweaks<commit_after>package bot

import (
	"fmt"
	"github.com/darkliquid/go-ircevent"
	"github.com/darkliquid/leader1/utils"
	"github.com/darkliquid/leader1/debug"
	"sort"
	"strings"
	"time"
)

func (bot *Bot) RunBuiltinCommands(event *irc.Event) {
	args := strings.Split(strings.TrimSpace(event.Message()), " ")
	command := args[0]

	// Bin the command from the arg list
	if len(args) > 1 {
		args = args[1:]
	} else {
		args = []string{}
	}

	// Get the current channel privileges for the nick sending this command
	privs, ok := bot.state.GetPrivs(event.Arguments[0], event.Nick)

	// Commands must be run by known users
	switch {
	case ok && command == "!reload":
		if privs.Owner || privs.Admin || privs.Op {
			bot.pm.InitJS()
			utils.IRCAction(bot.conn, event.Arguments[0], "has reloaded its plugins")
		} else {
			utils.IRCAction(bot.conn, event.Arguments[0], fmt.Sprintf("slaps %s's hands away from the op only controls", event.Nick))
		}
	case command == "!ping":
		bot.conn.Privmsg(event.Arguments[0], fmt.Sprintf("%s: PONG!", event.Nick))
	case ok && command == "!quit":
		if privs.Owner || privs.Admin || privs.Op {
			bot.Quit()
		} else {
			utils.IRCAction(bot.conn, event.Arguments[0], fmt.Sprintf("slaps %s's hands away from the op only controls", event.Nick))
		}
	case command == "!help":
		if len(args) == 0 {
			bot.ShowCommandList(event.Arguments[0], event.Nick)
		} else {
			bot.ShowCommandHelp(event.Arguments[0], event.Nick, args[0])
		}
	case ok && command == "!import" && len(args) >= 2 && bot.cfg.Irc.StaffChannel == event.Arguments[0]:
		if privs.Owner {
			overwrite := false
			if len(args) == 3 {
				overwrite = args[2] == "overwrite"
			}
			if err := bot.pm.ImportPlugin(event.Nick, args[0], args[1], overwrite); err != nil {
				bot.conn.Privmsg(bot.cfg.Irc.StaffChannel, fmt.Sprintf("ALERT: %s tried to use !import with %s and got this error: %s", event.Nick, args[0], err.Error()))
			} else {
				bot.conn.Privmsg(bot.cfg.Irc.StaffChannel, fmt.Sprintf("ALERT: %s successfully used !import with %s", event.Nick, args[0]))
			}
		} else {
			utils.IRCAction(bot.conn, event.Arguments[0], fmt.Sprintf("slaps %s's hands away from the op only controls", event.Nick))
		}
	case ok && command == "!debug" && bot.cfg.Irc.StaffChannel == event.Arguments[0]:
		if privs.Owner || privs.Admin || privs.Op {
			switch {
			case len(args) > 0 && args[0] == "on":
				port, alreadyRunning := debug.StartDebugServer()
				if alreadyRunning {
					bot.conn.Privmsg(bot.cfg.Irc.StaffChannel, fmt.Sprintf("%s: Debug server running on port %s", event.Nick, port))
				} else {
					bot.conn.Privmsg(bot.cfg.Irc.StaffChannel, fmt.Sprintf("%s: Debug server started on port %s", event.Nick, port))
				}
				bot.cfg.Debug = true
			case len(args) > 0 && args[0] == "off":
				debug.StopDebugServer()
				bot.conn.Privmsg(bot.cfg.Irc.StaffChannel, fmt.Sprintf("%s: Debug server stopped", event.Nick))
				bot.cfg.Debug = false
			case len(args) > 0 && args[0] == "status":
				status := debug.DebugServerStatus()
				bot.conn.Privmsg(bot.cfg.Irc.StaffChannel, fmt.Sprintf("%s: Debug server is %s", event.Nick, status))
			default:
				bot.conn.Privmsg(bot.cfg.Irc.StaffChannel, fmt.Sprintf("%s: usage - !debug [on|off]", event.Nick))
			}
		} else {
			utils.IRCAction(bot.conn, event.Arguments[0], fmt.Sprintf("slaps %s's hands away from the op only controls", event.Nick))
		}
	}
}

// Handler for reclaiming a stolen nick
func (bot *Bot) ReclaimNick(event *irc.Event) {
	if thief := bot.state.GetNick(bot.cfg.Irc.Nick); thief != nil {
		// Recover nick from thieves
		bot.conn.Privmsg("NickServ", fmt.Sprintf("RECOVER %s %s", bot.cfg.Irc.Nick, bot.cfg.Irc.NickPass))
		time.Sleep(time.Second)
		bot.conn.Privmsg("NickServ", fmt.Sprintf("RELEASE %s %s", bot.cfg.Irc.Nick, bot.cfg.Irc.NickPass))
	}
	bot.SetBotState(event)
}

// Automatically give voice to users in channels in which the bot has Op
func (bot *Bot) AutoVoice(event *irc.Event) {
	if bot.cfg.Irc.AutoVoice {
		time.Sleep(time.Second) // Wait a second before we bother

		privs, ok := bot.state.GetPrivs(event.Arguments[0], event.Nick)

		// No need to autovoice
		if ok && (privs.Owner || privs.Admin || privs.Op || privs.HalfOp || privs.Voice) {
			return
		}

		// Set mode
		bot.conn.Mode(event.Arguments[0], fmt.Sprintf("+v %s", event.Nick))
	}
}

// Function for setting up the botstate
func (bot *Bot) SetBotState(event *irc.Event) {
	if ghost := bot.state.GetNick(bot.cfg.Irc.Nick); ghost != nil && ghost != bot.state.Me() {
		// GHOST the old nick
		bot.conn.Privmsg("NickServ", fmt.Sprintf("GHOST %s %s", bot.cfg.Irc.Nick, bot.cfg.Irc.NickPass))
	}

	// Set up the nick
	bot.conn.Nick(bot.cfg.Irc.Nick)

	// Identify as the nick owner
	bot.conn.Privmsg("NickServ", fmt.Sprintf("IDENTIFY %s", bot.cfg.Irc.NickPass))

	// Tell IRC I'm a bot
	bot.conn.Mode(bot.cfg.Irc.Nick, "+B")
}

// Function for re-joining channels
func (bot *Bot) JoinChannels(event *irc.Event) {
	time.Sleep(time.Second) // Wait a second before we bother

	if _, ok := bot.state.GetPrivs(bot.cfg.Irc.NormalChannel, bot.state.Me().Nick); !ok {
		bot.conn.Join(bot.cfg.Irc.NormalChannel)
	}

	if _, ok := bot.state.GetPrivs(bot.cfg.Irc.StaffChannel, bot.state.Me().Nick); !ok {
		bot.conn.Join(bot.cfg.Irc.StaffChannel)
	}
}

// Print out the commands available
func (bot *Bot) ShowCommandList(source, nick string) {
	var commands []string = make([]string, 0)
	commands = append(commands, "!reload", "!ping", "!quit", "!help", "!import", "!debug")

	for cmd, _ := range bot.pm.CommandHelp() {
		commands = append(commands, cmd)
	}
	sort.Strings(commands)
	bot.conn.Privmsg(source, fmt.Sprintf("%s: available commands are: %s", nick, strings.Join(commands, ", ")))
}

// Print out the commands available
func (bot *Bot) ShowCommandHelp(source, nick, cmd string) {
	message := "unknown command, run `!help` to see what commands are available."
	switch cmd {
	case "!ping":
		message = fmt.Sprintf("makes `%s` reply with PONG!", bot.state.Me().Nick)
	case "!reload":
		message = "reloads the plugins"
	case "!quit":
		message = fmt.Sprintf("makes `%s` quit IRC", bot.state.Me().Nick)
	case "!help":
		message = "shows this message, smart ass"
	case "!import":
		message = "call like !import [url] [name] [overwrite], imports the plugin at [url] into [name].js and loads it into the bot. Will not overwrite unless [overwrite] is set to 'overwrite'"
	case "!debug":
		message = "starts/stops debugging server for inspecting the bot"
	default:
		if help, ok := bot.pm.CommandHelp()[cmd]; ok {
			message = help
		}
	}

	bot.conn.Privmsg(source, fmt.Sprintf("%s: %s - %s", nick, cmd, message))
}
<|endoftext|>