<commit_before>package socket

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"github.com/davyxu/cellnet"
	"io"
	"net"
	"sync"
)

const (
	PackageHeaderSize = 8 // MsgID(uint32) + Ser(uint16) + Size(uint16)
)

type TLVStream struct {
	recvser      uint16
	sendser      uint16
	conn         net.Conn
	sendtagGuard sync.RWMutex

	outputWriter     *bufio.Writer
	outputHeadBuffer *bytes.Buffer

	inputHeadBuffer []byte
	headReader      *bytes.Reader

	maxPacketSize int
}

var (
	ErrPackageTagNotMatch     = errors.New("ReadPacket: package tag not match")
	ErrPackageDataSizeInvalid = errors.New("ReadPacket: package crack, invalid size")
	ErrPackageTooBig          = errors.New("ReadPacket: package too big")
)

func (self *TLVStream) SetMaxPacketSize(size int) {
	self.maxPacketSize = size
}

// 从socket读取1个封包,并返回
func (self *TLVStream) Read() (msgid uint32, data []byte, err error) {

	if _, err = self.headReader.Seek(0, 0); err != nil {
		return
	}

	if _, err = io.ReadFull(self.conn, self.inputHeadBuffer); err != nil {
		return
	}

	// 读取ID
	if err = binary.Read(self.headReader, binary.LittleEndian, &msgid); err != nil {
		return
	}

	// 读取序号
	var ser uint16
	if err = binary.Read(self.headReader, binary.LittleEndian, &ser); err != nil {
		return
	}

	// 读取整包大小
	var fullsize uint16
	if err = binary.Read(self.headReader, binary.LittleEndian, &fullsize); err != nil {
		return
	}

	// 封包太大
	if self.maxPacketSize > 0 && int(fullsize) > self.maxPacketSize {
		err = ErrPackageTooBig
		return
	}

	// 序列号不匹配
	if self.recvser != ser {
		err = ErrPackageTagNotMatch
		return
	}

	dataSize := fullsize - PackageHeaderSize
	if dataSize < 0 {
		err = ErrPackageDataSizeInvalid
		return
	}

	// 读取数据
	msgBytes := make([]byte, dataSize)
	if _, err = io.ReadFull(self.conn, msgBytes); err != nil {
		return
	}

	data = msgBytes

	// 增加序列号值
	self.recvser++

	return
}

// 将一个封包发送到socket
func (self *TLVStream) Write(msgid uint32, data []byte) (err error) {

	// 防止将Send放在go内造成的多线程冲突问题
	self.sendtagGuard.Lock()
	defer self.sendtagGuard.Unlock()

	self.outputHeadBuffer.Reset()

	// 写ID
	if err = binary.Write(self.outputHeadBuffer, binary.LittleEndian, msgid); err != nil {
		return err
	}

	// 写序号
	if err = binary.Write(self.outputHeadBuffer, binary.LittleEndian, self.sendser); err != nil {
		return err
	}

	// 写包大小
	if err = binary.Write(self.outputHeadBuffer, binary.LittleEndian, uint16(len(data)+PackageHeaderSize)); err != nil {
		return err
	}

	// 发包头
	if err = self.writeFull(self.outputHeadBuffer.Bytes()); err != nil {
		return err
	}

	// 发包内容
	if err = self.writeFull(data); err != nil {
		return err
	}

	// 增加序号值
	self.sendser++

	return
}

// 完整发送所有封包
func (self *TLVStream) writeFull(p []byte) error {

	sizeToWrite := len(p)

	for {

		n, err := self.outputWriter.Write(p)

		if err != nil {
			return err
		}

		if n >= sizeToWrite {
			break
		}

		copy(p[0:sizeToWrite-n], p[n:sizeToWrite])
		sizeToWrite -= n
	}

	return nil

}

const sendTotalTryCount = 100

func (self *TLVStream) Flush() error {

	var err error
	for tryTimes := 0; tryTimes < sendTotalTryCount; tryTimes++ {

		err = self.outputWriter.Flush()

		// 如果没写完, flush底层会将没发完的buff准备好, 我们只需要重新调一次flush
		if err != io.ErrShortWrite {
			break
		}
	}

	return err
}

// 关闭
func (self *TLVStream) Close() error {
	return self.conn.Close()
}

// 裸socket操作
func (self *TLVStream) Raw() net.Conn {
	return self.conn
}

// 封包流 relay模式: 在封包头有clientid信息
func NewTLVStream(conn net.Conn) *TLVStream {

	s := &TLVStream{
		conn:             conn,
		recvser:          1,
		sendser:          1,
		outputWriter:     bufio.NewWriter(conn),
		outputHeadBuffer: bytes.NewBuffer([]byte{}),
		inputHeadBuffer:  make([]byte, PackageHeaderSize),
	}
	s.headReader = bytes.NewReader(s.inputHeadBuffer)

	return s
}

func errToResult(err error) cellnet.Result {

	if err == nil {
		return cellnet.Result_OK
	}

	switch err {
	case ErrPackageTagNotMatch, ErrPackageDataSizeInvalid, ErrPackageTooBig:
		return cellnet.Result_PackageCrack
	}

	switch n := err.(type) {
	case net.Error:
		if n.Timeout() {
			return cellnet.Result_SocketTimeout
		}
	}

	return cellnet.Result_SocketError
}
<commit_msg>修复: 包大小超过64k时封包大小发送错误问题, 包size现在扩展到32位无符号整形, 整个包头扩展到10字节<commit_after>package socket

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"github.com/davyxu/cellnet"
	"io"
	"net"
	"sync"
)

const (
	PackageHeaderSize = 10 // MsgID(uint32) + Ser(uint16) + Size(uint32)
)

type TLVStream struct {
	recvser      uint16
	sendser      uint16
	conn         net.Conn
	sendtagGuard sync.RWMutex

	outputWriter     *bufio.Writer
	outputHeadBuffer *bytes.Buffer

	inputHeadBuffer []byte
	headReader      *bytes.Reader

	maxPacketSize int
}

var (
	ErrPackageTagNotMatch     = errors.New("ReadPacket: package tag not match")
	ErrPackageDataSizeInvalid = errors.New("ReadPacket: package crack, invalid size")
	ErrPackageTooBig          = errors.New("ReadPacket: package too big")
)

func (self *TLVStream) SetMaxPacketSize(size int) {
	self.maxPacketSize = size
}

// 从socket读取1个封包,并返回
func (self *TLVStream) Read() (msgid uint32, data []byte, err error) {

	if _, err = self.headReader.Seek(0, 0); err != nil {
		return
	}

	if _, err = io.ReadFull(self.conn, self.inputHeadBuffer); err != nil {
		return
	}

	// 读取ID
	if err = binary.Read(self.headReader, binary.LittleEndian, &msgid); err != nil {
		return
	}

	// 读取序号
	var ser uint16
	if err = binary.Read(self.headReader, binary.LittleEndian, &ser); err != nil {
		return
	}

	// 读取整包大小
	var fullsize uint32
	if err = binary.Read(self.headReader, binary.LittleEndian, &fullsize); err != nil {
		return
	}

	// 封包太大
	if self.maxPacketSize > 0 && int(fullsize) > self.maxPacketSize {
		err = ErrPackageTooBig
		return
	}

	// 序列号不匹配
	if self.recvser != ser {
		err = ErrPackageTagNotMatch
		return
	}

	dataSize := fullsize - PackageHeaderSize
	if dataSize < 0 {
		err = ErrPackageDataSizeInvalid
		return
	}

	// 读取数据
	msgBytes := make([]byte, dataSize)
	if _, err = io.ReadFull(self.conn, msgBytes); err != nil {
		return
	}

	data = msgBytes

	// 增加序列号值
	self.recvser++

	return
}

// 将一个封包发送到socket
func (self *TLVStream) Write(msgid uint32, data []byte) (err error) {

	// 防止将Send放在go内造成的多线程冲突问题
	self.sendtagGuard.Lock()
	defer self.sendtagGuard.Unlock()

	self.outputHeadBuffer.Reset()

	// 写ID
	if err = binary.Write(self.outputHeadBuffer, binary.LittleEndian, msgid); err != nil {
		return err
	}

	// 写序号
	if err = binary.Write(self.outputHeadBuffer, binary.LittleEndian, self.sendser); err != nil {
		return err
	}

	// 写包大小
	if err = binary.Write(self.outputHeadBuffer, binary.LittleEndian, uint32(len(data)+PackageHeaderSize)); err != nil {
		return err
	}

	// 发包头
	if err = self.writeFull(self.outputHeadBuffer.Bytes()); err != nil {
		return err
	}

	// 发包内容
	if err = self.writeFull(data); err != nil {
		return err
	}

	// 增加序号值
	self.sendser++

	return
}

// 完整发送所有封包
func (self *TLVStream) writeFull(p []byte) error {

	sizeToWrite := len(p)

	for {

		n, err := self.outputWriter.Write(p)

		if err != nil {
			return err
		}

		if n >= sizeToWrite {
			break
		}

		copy(p[0:sizeToWrite-n], p[n:sizeToWrite])
		sizeToWrite -= n
	}

	return nil

}

const sendTotalTryCount = 100

func (self *TLVStream) Flush() error {

	var err error
	for tryTimes := 0; tryTimes < sendTotalTryCount; tryTimes++ {

		err = self.outputWriter.Flush()

		// 如果没写完, flush底层会将没发完的buff准备好, 我们只需要重新调一次flush
		if err != io.ErrShortWrite {
			break
		}
	}

	return err
}

// 关闭
func (self *TLVStream) Close() error {
	return self.conn.Close()
}

// 裸socket操作
func (self *TLVStream) Raw() net.Conn {
	return self.conn
}

// 封包流 relay模式: 在封包头有clientid信息
func NewTLVStream(conn net.Conn) *TLVStream {

	s := &TLVStream{
		conn:             conn,
		recvser:          1,
		sendser:          1,
		outputWriter:     bufio.NewWriter(conn),
		outputHeadBuffer: bytes.NewBuffer([]byte{}),
		inputHeadBuffer:  make([]byte, PackageHeaderSize),
	}
	s.headReader = bytes.NewReader(s.inputHeadBuffer)

	return s
}

func errToResult(err error) cellnet.Result {

	if err == nil {
		return cellnet.Result_OK
	}

	switch err {
	case ErrPackageTagNotMatch, ErrPackageDataSizeInvalid, ErrPackageTooBig:
		return cellnet.Result_PackageCrack
	}

	switch n := err.(type) {
	case net.Error:
		if n.Timeout() {
			return cellnet.Result_SocketTimeout
		}
	}

	return cellnet.Result_SocketError
}
<|endoftext|>


<commit_before>package main

import (
	"math/rand"
	"strings"

	"github.com/animenotifier/arn"

	"github.com/animenotifier/arn/search"
	"github.com/bwmarrin/discordgo"
)

// OnMessageCreate is called every time a new message is created on any channel.
func OnMessageCreate(s *discordgo.Session, msg *discordgo.MessageCreate) {
	// Ignore all messages created by the bot itself
	if msg.Author.ID == s.State.User.ID {
		return
	}

	if msg.Content == "!help" || msg.Content == "!commands" {
		s.ChannelMessageSend(msg.ChannelID, `
**!a** [anime search term]
**!animelist** [username]
**!play** [status text]
**!randomquote**
**!source**
**!region** [region]`)
	}

	// Has the bot been mentioned?
	for _, user := range msg.Mentions {
		if user.ID == discord.State.User.ID {
			s.ChannelMessageSend(msg.ChannelID, msg.Author.Mention()+" :heart:")
			return
		}
	}

	// Anime search
	if strings.HasPrefix(msg.Content, "!a ") {
		term := msg.Content[len("!a "):]
		animes := search.Anime(term, 3)
		message := ""

		for _, anime := range animes {
			message += "https://notify.moe" + anime.Link() + "\n"
		}

		if len(animes) == 0 {
			message = "Sorry, I couldn't find anything using that term."
		}

		s.ChannelMessageSend(msg.ChannelID, message)
		return
	}

	// Anime list of user
	if strings.HasPrefix(msg.Content, "!animelist ") {
		s.ChannelMessageSend(msg.ChannelID, "https://notify.moe/+"+strings.Split(msg.Content, " ")[1]+"/animelist/watching")
		return
	}

	// Play status
	if strings.HasPrefix(msg.Content, "!play ") {
		s.UpdateStatus(0, msg.Content[len("!play "):])
		return
	}

	// Random quote
	if msg.Content == "!randomquote" {
		allQuotes := arn.FilterQuotes(func(quote *arn.Quote) bool {
			return !quote.IsDraft && quote.IsValid()
		})

		quote := allQuotes[rand.Intn(len(allQuotes))]
		s.ChannelMessageSend(msg.ChannelID, "https://notify.moe"+quote.Link())
		return
	}

	// GitHub source of the bot
	if msg.Content == "!source" {
		s.ChannelMessageSend(msg.ChannelID, msg.Author.Mention()+" B-baaaaaaaka! Y..you...you want to...TOUCH MY CODE?!\n\nhttps://github.com/animenotifier/notify.moe/tree/go/bots/discord")
		return
	}
}
<commit_msg>Some more work on the discord bot<commit_after>package main

import (
	"math/rand"
	"strings"

	"github.com/animenotifier/arn"

	"github.com/animenotifier/arn/search"
	"github.com/bwmarrin/discordgo"
)

// OnMessageCreate is called every time a new message is created on any channel.
func OnMessageCreate(s *discordgo.Session, msg *discordgo.MessageCreate) {
	// Ignore all messages created by the bot itself
	if msg.Author.ID == s.State.User.ID {
		return
	}

	if msg.Content == "!help" || msg.Content == "!commands" {
		s.ChannelMessageSend(msg.ChannelID, `
**!a** [anime search term]
**!animelist** [username]
**!play** [status text]
**!randomquote**
**!source**
**!region** [region]`)
	}

	// Has the bot been mentioned?
	for _, user := range msg.Mentions {
		if user.ID == discord.State.User.ID {
			s.ChannelMessageSend(msg.ChannelID, msg.Author.Mention()+" :heart:")
			return
		}
	}

	// Anime search
	if strings.HasPrefix(msg.Content, "!a ") {
		term := msg.Content[len("!a "):]
		animes := search.Anime(term, 3)
		message := ""

		for _, anime := range animes {
			message += "https://notify.moe" + anime.Link() + "\n"
		}

		if len(animes) == 0 {
			message = "Sorry, I couldn't find anything using that term."
		}

		s.ChannelMessageSend(msg.ChannelID, message)
		return
	}

	// Anime list of user
	if strings.HasPrefix(msg.Content, "!animelist ") {
		s.ChannelMessageSend(msg.ChannelID, "https://notify.moe/+"+strings.Split(msg.Content, " ")[1]+"/animelist/watching")
		return
	}

	// Play status
	if strings.HasPrefix(msg.Content, "!play ") {
		s.UpdateStatus(0, msg.Content[len("!play "):])
		return
	}

	// Random quote
	if msg.Content == "!randomquote" {
		allQuotes := arn.FilterQuotes(func(quote *arn.Quote) bool {
			return !quote.IsDraft && quote.IsValid()
		})

		quote := allQuotes[rand.Intn(len(allQuotes))]
		s.ChannelMessageSend(msg.ChannelID, "https://notify.moe"+quote.Link())
		return
	}

	// GitHub source of the bot
	if msg.Content == "!source" {
		s.ChannelMessageSend(msg.ChannelID, msg.Author.Mention()+" B-baaaaaaaka! Y..you...you want to...TOUCH MY CODE?!\n\nhttps://github.com/animenotifier/notify.moe/tree/go/bots/discord")
		return
	}

	// Set the specific region role for the user
	if strings.HasPrefix(msg.Content, "!region ") {

		regions := map[string]string{
			"africa":    "465876853236826112",
			"america":   "465876808311635979",
			"asia":      "465876834031108096",
			"australia": "465876893036707840",
			"europe":    "465876773029019659",
		}

		region := msg.Content[len("!region "):]

		// check to make sure the region is in the region map
		if _, ok := regions[region]; ok {
			// try to set the role
			c, _ := s.Channel(msg.ChannelID)
			err := s.GuildMemberRoleAdd(c.GuildID, msg.Author.ID, regions[region])

			if err != nil {
				s.ChannelMessageSend(msg.ChannelID, "The region role could not be set!")
				return
			}

			s.ChannelMessageSend(msg.ChannelID, "The "+region+" role has been set on your user!")

		} else {
			s.ChannelMessageSend(msg.ChannelID, "This is not a region!")
		}

		return
	}
}
<|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.

package strings

import (
	"errors"
	"io"
	"unicode/utf8"
)

// A Reader implements the io.Reader, io.ReaderAt, io.Seeker, io.WriterTo,
// io.ByteScanner, and io.RuneScanner interfaces by reading
// from a string.
// The zero value for Reader operates like a Reader of an empty string.
type Reader struct {
	s        string
	i        int64 // current reading index
	prevRune int   // index of previous rune; or < 0
}

// Len returns the number of bytes of the unread portion of the
// string.
func (r *Reader) Len() int {
	if r.i >= int64(len(r.s)) {
		return 0
	}
	return int(int64(len(r.s)) - r.i)
}

// Size returns the original length of the underlying string.
// Size is the number of bytes available for reading via ReadAt.
// The returned value is always the same and is not affected by calls
// to any other method.
func (r *Reader) Size() int64 { return int64(len(r.s)) }

// Read implements the io.Reader interface.
func (r *Reader) Read(b []byte) (n int, err error) {
	if r.i >= int64(len(r.s)) {
		return 0, io.EOF
	}
	r.prevRune = -1
	n = copy(b, r.s[r.i:])
	r.i += int64(n)
	return
}

// ReadAt implements the io.ReaderAt interface.
func (r *Reader) ReadAt(b []byte, off int64) (n int, err error) {
	// cannot modify state - see io.ReaderAt
	if off < 0 {
		return 0, errors.New("strings.Reader.ReadAt: negative offset")
	}
	if off >= int64(len(r.s)) {
		return 0, io.EOF
	}
	n = copy(b, r.s[off:])
	if n < len(b) {
		err = io.EOF
	}
	return
}

// ReadByte implements the io.ByteReader interface.
func (r *Reader) ReadByte() (byte, error) {
	r.prevRune = -1
	if r.i >= int64(len(r.s)) {
		return 0, io.EOF
	}
	b := r.s[r.i]
	r.i++
	return b, nil
}

// UnreadByte implements the io.ByteScanner interface.
func (r *Reader) UnreadByte() error {
	if r.i <= 0 {
		return errors.New("strings.Reader.UnreadByte: at beginning of string")
	}
	r.prevRune = -1
	r.i--
	return nil
}

// ReadRune implements the io.RuneReader interface.
func (r *Reader) ReadRune() (ch rune, size int, err error) {
	if r.i >= int64(len(r.s)) {
		r.prevRune = -1
		return 0, 0, io.EOF
	}
	r.prevRune = int(r.i)
	if c := r.s[r.i]; c < utf8.RuneSelf {
		r.i++
		return rune(c), 1, nil
	}
	ch, size = utf8.DecodeRuneInString(r.s[r.i:])
	r.i += int64(size)
	return
}

// UnreadRune implements the io.RuneScanner interface.
func (r *Reader) UnreadRune() error {
	if r.i <= 0 {
		return errors.New("strings.Reader.UnreadRune: at beginning of string")
	}
	if r.prevRune < 0 {
		return errors.New("strings.Reader.UnreadRune: previous operation was not ReadRune")
	}
	r.i = int64(r.prevRune)
	r.prevRune = -1
	return nil
}

// Seek implements the io.Seeker interface.
func (r *Reader) Seek(offset int64, whence int) (int64, error) {
	r.prevRune = -1
	var abs int64
	switch whence {
	case io.SeekStart:
		abs = offset
	case io.SeekCurrent:
		abs = r.i + offset
	case io.SeekEnd:
		abs = int64(len(r.s)) + offset
	default:
		return 0, errors.New("strings.Reader.Seek: invalid whence")
	}
	if abs < 0 {
		return 0, errors.New("strings.Reader.Seek: negative position")
	}
	r.i = abs
	return abs, nil
}

// WriteTo implements the io.WriterTo interface.
func (r *Reader) WriteTo(w io.Writer) (n int64, err error) {
	r.prevRune = -1
	if r.i >= int64(len(r.s)) {
		return 0, nil
	}
	s := r.s[r.i:]
	m, err := io.WriteString(w, s)
	if m > len(s) {
		panic("strings.Reader.WriteTo: invalid WriteString count")
	}
	r.i += int64(m)
	n = int64(m)
	if m != len(s) && err == nil {
		err = io.ErrShortWrite
	}
	return
}

// Reset resets the Reader to be reading from s.
func (r *Reader) Reset(s string) { *r = Reader{s, 0, -1} }

// NewReader returns a new Reader reading from s.
// It is similar to bytes.NewBufferString but more efficient and read-only.
func NewReader(s string) *Reader { return &Reader{s, 0, -1} }
<commit_msg>strings: complete Reader doc string<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.

package strings

import (
	"errors"
	"io"
	"unicode/utf8"
)

// A Reader implements the io.Reader, io.ReaderAt, io.ByteReader, io.ByteScanner,
// io.RuneReader, io.RuneScanner, io.Seeker, and io.WriterTo interfaces by reading
// from a string.
// The zero value for Reader operates like a Reader of an empty string.
type Reader struct {
	s        string
	i        int64 // current reading index
	prevRune int   // index of previous rune; or < 0
}

// Len returns the number of bytes of the unread portion of the
// string.
func (r *Reader) Len() int {
	if r.i >= int64(len(r.s)) {
		return 0
	}
	return int(int64(len(r.s)) - r.i)
}

// Size returns the original length of the underlying string.
// Size is the number of bytes available for reading via ReadAt.
// The returned value is always the same and is not affected by calls
// to any other method.
func (r *Reader) Size() int64 { return int64(len(r.s)) }

// Read implements the io.Reader interface.
func (r *Reader) Read(b []byte) (n int, err error) {
	if r.i >= int64(len(r.s)) {
		return 0, io.EOF
	}
	r.prevRune = -1
	n = copy(b, r.s[r.i:])
	r.i += int64(n)
	return
}

// ReadAt implements the io.ReaderAt interface.
func (r *Reader) ReadAt(b []byte, off int64) (n int, err error) {
	// cannot modify state - see io.ReaderAt
	if off < 0 {
		return 0, errors.New("strings.Reader.ReadAt: negative offset")
	}
	if off >= int64(len(r.s)) {
		return 0, io.EOF
	}
	n = copy(b, r.s[off:])
	if n < len(b) {
		err = io.EOF
	}
	return
}

// ReadByte implements the io.ByteReader interface.
func (r *Reader) ReadByte() (byte, error) {
	r.prevRune = -1
	if r.i >= int64(len(r.s)) {
		return 0, io.EOF
	}
	b := r.s[r.i]
	r.i++
	return b, nil
}

// UnreadByte implements the io.ByteScanner interface.
func (r *Reader) UnreadByte() error {
	if r.i <= 0 {
		return errors.New("strings.Reader.UnreadByte: at beginning of string")
	}
	r.prevRune = -1
	r.i--
	return nil
}

// ReadRune implements the io.RuneReader interface.
func (r *Reader) ReadRune() (ch rune, size int, err error) {
	if r.i >= int64(len(r.s)) {
		r.prevRune = -1
		return 0, 0, io.EOF
	}
	r.prevRune = int(r.i)
	if c := r.s[r.i]; c < utf8.RuneSelf {
		r.i++
		return rune(c), 1, nil
	}
	ch, size = utf8.DecodeRuneInString(r.s[r.i:])
	r.i += int64(size)
	return
}

// UnreadRune implements the io.RuneScanner interface.
func (r *Reader) UnreadRune() error {
	if r.i <= 0 {
		return errors.New("strings.Reader.UnreadRune: at beginning of string")
	}
	if r.prevRune < 0 {
		return errors.New("strings.Reader.UnreadRune: previous operation was not ReadRune")
	}
	r.i = int64(r.prevRune)
	r.prevRune = -1
	return nil
}

// Seek implements the io.Seeker interface.
func (r *Reader) Seek(offset int64, whence int) (int64, error) {
	r.prevRune = -1
	var abs int64
	switch whence {
	case io.SeekStart:
		abs = offset
	case io.SeekCurrent:
		abs = r.i + offset
	case io.SeekEnd:
		abs = int64(len(r.s)) + offset
	default:
		return 0, errors.New("strings.Reader.Seek: invalid whence")
	}
	if abs < 0 {
		return 0, errors.New("strings.Reader.Seek: negative position")
	}
	r.i = abs
	return abs, nil
}

// WriteTo implements the io.WriterTo interface.
func (r *Reader) WriteTo(w io.Writer) (n int64, err error) {
	r.prevRune = -1
	if r.i >= int64(len(r.s)) {
		return 0, nil
	}
	s := r.s[r.i:]
	m, err := io.WriteString(w, s)
	if m > len(s) {
		panic("strings.Reader.WriteTo: invalid WriteString count")
	}
	r.i += int64(m)
	n = int64(m)
	if m != len(s) && err == nil {
		err = io.ErrShortWrite
	}
	return
}

// Reset resets the Reader to be reading from s.
func (r *Reader) Reset(s string) { *r = Reader{s, 0, -1} }

// NewReader returns a new Reader reading from s.
// It is similar to bytes.NewBufferString but more efficient and read-only.
func NewReader(s string) *Reader { return &Reader{s, 0, -1} }
<|endoftext|>


<commit_before>package main

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/esiqveland/notify"
	"github.com/godbus/dbus"

	"cloud.google.com/go/firestore"
)

const (
	TwitchID            = "TWITCH_ID"
	TwitchAuth          = "TWITCH_AUTH"
	FirestoreApp        = "TWITCH_FIRESTORE_APP"
	FirestoreCollection = "FIRESTORE_COLLECTION"
	FirestoreDocID      = "FIRESTORE_DOC_ID"
)

type TwitchChannel struct {
	DisplayName string `json:"display_name"`
	Game        string `json:"game"`
	Status      string `json:"status"`
}

type TwitchData struct {
	LastModified time.Time
	Channels     map[string]*TwitchChannel
}

type TwitchDataWithErr struct {
	*TwitchData
	error
}

var TwitchAuthHeaders = http.Header{
	"Accept":        []string{"application/vnd.twitchtv.v5+json"},
	"Client-ID":     []string{},
	"Authorization": []string{"OAuth "},
}

const TwitchAPIFollowed = "https://api.twitch.tv/kraken/streams/followed"

func init() {
	TwitchAuthHeaders.Set("Client-ID", os.Getenv(TwitchID))
	TwitchAuthHeaders.Set("Authorization",
		TwitchAuthHeaders.Get("Authorization")+os.Getenv(TwitchAuth))
}

func OldChannels(client *firestore.Client, ctx context.Context) TwitchDataWithErr {

	snapshot, err := client.Collection(os.Getenv(FirestoreCollection)).
		Doc(os.Getenv(FirestoreDocID)).
		Get(ctx)
	if err != nil {
		return TwitchDataWithErr{nil, err}
	}
	var doc TwitchData
	if err := snapshot.DataTo(&doc); err != nil {
		return TwitchDataWithErr{nil, err}
	}
	return TwitchDataWithErr{&doc, nil}
}

func NewChannels() TwitchDataWithErr {
	req, err := http.NewRequest("GET", TwitchAPIFollowed, nil)
	if err != nil {
		return TwitchDataWithErr{nil, err}
	}
	req.Header = TwitchAuthHeaders
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return TwitchDataWithErr{nil, err}
	}
	defer resp.Body.Close()
	var data struct {
		Streams []struct {
			Channel struct {
				Name string `json:"name"`
				*TwitchChannel
			} `json:"channel"`
		} `json:"streams"`
	}
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return TwitchDataWithErr{nil, err}
	}
	ret := TwitchData{
		LastModified: time.Now(),
		Channels:     make(map[string]*TwitchChannel),
	}
	for _, stream := range data.Streams {
		ret.Channels[stream.Channel.Name] = stream.Channel.TwitchChannel
	}
	return TwitchDataWithErr{&ret, nil}
}

func equal(data1, data2 *TwitchChannel) bool {
	return data1.DisplayName == data2.DisplayName &&
		data1.Game == data2.Game &&
		data1.Status == data2.Status
}

func main() {

	// initialize dbus connection
	conn, err := dbus.SessionBus()
	if err != nil {
		fmt.Println(err)
		return
	}
	notifier, err := notify.New(conn)
	if err != nil {
		fmt.Println(err)
		return
	}
	defer notifier.Close()

	// initialize firestore connection
	ctx := context.Background()
	client, err := firestore.NewClient(ctx, os.Getenv(FirestoreApp))
	if err != nil {
		fmt.Println(err)
		return
	}
	defer client.Close()

	// execute requests in parallel
	var newChannelsRes TwitchDataWithErr
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		// get fresh channels from twitch
		newChannelsRes = NewChannels()
		wg.Done()
	}()
	// get old channels from firestore
	oldChannelsRes := OldChannels(client, ctx)
	if oldChannelsRes.error != nil {
		fmt.Println(oldChannelsRes.error)
		return
	}
	wg.Wait()
	if newChannelsRes.error != nil {
		fmt.Println(newChannelsRes.error)
		return
	}
	// handle offline channels
	// (in old but not in new -> offline)
	for ch, data := range oldChannelsRes.Channels {
		if _, ok := newChannelsRes.Channels[ch]; !ok {
			if _, err := notifier.SendNotification(notify.Notification{
				Summary: data.DisplayName,
				Body:    "is no longer online",
			}); err != nil {
				fmt.Println(err)
				return
			}
		}
	}

	// handle new channels and updated channel headers
	for ch, dataNew := range newChannelsRes.Channels {
		dataOld, ok := oldChannelsRes.Channels[ch]
		if !ok || !equal(dataNew, dataOld) {
			if _, err := notifier.SendNotification(notify.Notification{
				Summary: fmt.Sprintf("%v is playing %v", dataNew.DisplayName, dataNew.Game),
				Body:    dataNew.Status,
			}); err != nil {
				fmt.Println(err)
				return
			}
		}
	}

	_, err = client.Collection(os.Getenv(FirestoreCollection)).
		Doc(os.Getenv(FirestoreDocID)).
		Set(ctx, newChannelsRes.TwitchData)
	if err != nil {
		fmt.Println(err)
		return
	}

}
<commit_msg>get rid of ...Err structure<commit_after>package main

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/esiqveland/notify"
	"github.com/godbus/dbus"

	"cloud.google.com/go/firestore"
)

const (
	TwitchID            = "TWITCH_ID"
	TwitchAuth          = "TWITCH_AUTH"
	FirestoreApp        = "TWITCH_FIRESTORE_APP"
	FirestoreCollection = "FIRESTORE_COLLECTION"
	FirestoreDocID      = "FIRESTORE_DOC_ID"
)

type TwitchChannel struct {
	DisplayName string `json:"display_name"`
	Game        string `json:"game"`
	Status      string `json:"status"`
}

// ToDo get rid of LastModofied field because
// firestore already stores timepoint metadata
// snapshot.UpdateTime etc
type TwitchData struct {
	LastModified time.Time
	Channels     map[string]*TwitchChannel
}

var TwitchAuthHeaders = http.Header{
	"Accept":        []string{"application/vnd.twitchtv.v5+json"},
	"Client-ID":     []string{},
	"Authorization": []string{"OAuth "},
}

const TwitchAPIFollowed = "https://api.twitch.tv/kraken/streams/followed"

func init() {
	TwitchAuthHeaders.Set("Client-ID", os.Getenv(TwitchID))
	TwitchAuthHeaders.Set("Authorization",
		TwitchAuthHeaders.Get("Authorization")+os.Getenv(TwitchAuth))
}

func OldChannels(client *firestore.Client, ctx context.Context) (*TwitchData, error) {

	snapshot, err := client.Collection(os.Getenv(FirestoreCollection)).
		Doc(os.Getenv(FirestoreDocID)).
		Get(ctx)
	if err != nil {
		return nil, err
	}
	var doc TwitchData
	if err := snapshot.DataTo(&doc); err != nil {
		return nil, err
	}
	return &doc, nil
}

func NewChannels() (*TwitchData, error) {
	req, err := http.NewRequest("GET", TwitchAPIFollowed, nil)
	if err != nil {
		return nil, err
	}
	req.Header = TwitchAuthHeaders
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	var data struct {
		Streams []struct {
			Channel struct {
				Name string `json:"name"`
				*TwitchChannel
			} `json:"channel"`
		} `json:"streams"`
	}
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return nil, err
	}
	ret := TwitchData{
		LastModified: time.Now(),
		Channels:     make(map[string]*TwitchChannel),
	}
	for _, stream := range data.Streams {
		ret.Channels[stream.Channel.Name] = stream.Channel.TwitchChannel
	}
	return &ret, nil
}

func equal(data1, data2 *TwitchChannel) bool {
	return data1.DisplayName == data2.DisplayName &&
		data1.Game == data2.Game &&
		data1.Status == data2.Status
}

func main() {

	// initialize dbus connection
	conn, err := dbus.SessionBus()
	if err != nil {
		fmt.Println(err)
		return
	}
	notifier, err := notify.New(conn)
	if err != nil {
		fmt.Println(err)
		return
	}
	defer notifier.Close()

	// initialize firestore connection
	ctx := context.Background()
	client, err := firestore.NewClient(ctx, os.Getenv(FirestoreApp))
	if err != nil {
		fmt.Println(err)
		return
	}
	defer client.Close()

	// execute requests in parallel
	var newChannelsData *TwitchData
	var newChannelsErr error
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		// get fresh channels from twitch
		newChannelsData, newChannelsErr = NewChannels()
		wg.Done()
	}()
	// get old channels from firestore
	oldChannelsData, oldChannelsErr := OldChannels(client, ctx)
	if oldChannelsErr != nil {
		fmt.Println(oldChannelsErr)
		return
	}
	wg.Wait()
	if newChannelsErr != nil {
		fmt.Println(newChannelsErr)
		return
	}
	// handle offline channels
	// (in old but not in new -> offline)
	for ch, data := range oldChannelsData.Channels {
		if _, ok := newChannelsData.Channels[ch]; !ok {
			if _, err := notifier.SendNotification(notify.Notification{
				Summary: data.DisplayName,
				Body:    "is no longer online",
			}); err != nil {
				fmt.Println(err)
				return
			}
		}
	}

	// handle new channels and updated channel headers
	for ch, dataNew := range newChannelsData.Channels {
		dataOld, ok := oldChannelsData.Channels[ch]
		if !ok || !equal(dataNew, dataOld) {
			if _, err := notifier.SendNotification(notify.Notification{
				Summary: fmt.Sprintf("%v is playing %v", dataNew.DisplayName, dataNew.Game),
				Body:    dataNew.Status,
			}); err != nil {
				fmt.Println(err)
				return
			}
		}
	}

	// send new channels to firestore
	_, err = client.Collection(os.Getenv(FirestoreCollection)).
		Doc(os.Getenv(FirestoreDocID)).
		Set(ctx, newChannelsData)
	if err != nil {
		fmt.Println(err)
		return
	}

}
<|endoftext|>


<commit_before>package git

import (
	"bytes"
	"expvar"
	"fmt"
	"io"
	"net/http"
	"sort"
	"strings"
)

func Fetch(gitURL string, haves map[string]struct{}, msgW io.Writer,
	bwCounter *expvar.Int) (refs map[string]string, r io.Reader, err error) {

	req, err := http.NewRequest("GET", gitURL+"/info/refs?service=git-upload-pack", nil)
	if err != nil {
		return
	}
	req.Header.Set("User-Agent", "github.com/thecodearchive/gitarchive/git")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, nil, fmt.Errorf("GET /info/refs: %d", resp.StatusCode)
	}
	refs, err = ParseSmartResponse(resp.Body)
	if err != nil {
		return
	}

	for name := range refs {
		if strings.HasPrefix(name, "refs/pull/") {
			delete(refs, name)
		}
	}

	var wants []string
	for name, ref := range refs {
		if _, ok := haves[ref]; ok {
			continue
		}
		wants = append(wants, refs[name])
	}
	sort.Strings(wants)

	if len(wants) == 0 {
		return
	}

	body := &bytes.Buffer{}
	last := ""
	for _, want := range wants {
		if last == want {
			continue
		}
		command := "want " + want
		if last == "" {
			command += " ofs-delta side-band-64k thin-pack"
		}
		command += "\n"
		body.WriteString(fmt.Sprintf("%04x%s", len(command)+4, command))
		last = want
	}
	body.WriteString("0000")
	for have := range haves { // TODO: sort the haves
		command := "have " + have + "\n"
		body.WriteString(fmt.Sprintf("%04x%s", len(command)+4, command))
	}
	body.WriteString("0009done\n")

	req, err = http.NewRequest("POST", gitURL+"/git-upload-pack", body)
	if err != nil {
		return
	}
	req.Header.Set("Content-Type", "application/x-git-upload-pack-request")
	req.Header.Set("Accept", "application/x-git-upload-pack-result")
	req.Header.Set("User-Agent", "github.com/thecodearchive/gitarchive/git")
	resp, err = http.DefaultClient.Do(req)
	if err != nil {
		return
	}
	if resp.StatusCode != 200 {
		return nil, nil, fmt.Errorf("POST /git-upload-pack: %d", resp.StatusCode)
	}

	sbr := &sideBandReader{Upstream: resp.Body, MsgW: msgW}
	cr := &countingReader{Upstream: sbr, Counter: bwCounter}

	// TODO peek into reader to make sure it's not empty.
	// Needs to be > 32 bytes (hdr is 12, trailer is 20).

	return refs, cr, nil
}
<commit_msg>git: peek to check a fetched packfile is not empty<commit_after>package git

import (
	"bytes"
	"expvar"
	"fmt"
	"io"
	"net/http"
	"sort"
	"strings"
)

// Fetch fetches the git repo at gitURL proving the refs in haves as
// haves.
//
// Sideband messages from the git server are writetn to msgW, and the
// number of bytes fetched is incremented in bwCounter. bwCounter is
// incremented here to get fine-grained metrics.
func Fetch(gitURL string, haves map[string]struct{}, msgW io.Writer,
	bwCounter *expvar.Int) (refs map[string]string, r io.ReadCloser, err error) {

	req, err := http.NewRequest("GET", gitURL+"/info/refs?service=git-upload-pack", nil)
	if err != nil {
		return
	}
	req.Header.Set("User-Agent", "github.com/thecodearchive/gitarchive/git")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, nil, fmt.Errorf("GET /info/refs: %d", resp.StatusCode)
	}
	refs, err = ParseSmartResponse(resp.Body)
	if err != nil {
		return
	}

	for name := range refs {
		if strings.HasPrefix(name, "refs/pull/") {
			delete(refs, name)
		}
	}

	var wants []string
	for name, ref := range refs {
		if _, ok := haves[ref]; ok {
			continue
		}
		wants = append(wants, refs[name])
	}
	sort.Strings(wants)

	if len(wants) == 0 {
		return
	}

	body := &bytes.Buffer{}
	last := ""
	for _, want := range wants {
		if last == want {
			continue
		}
		command := "want " + want
		if last == "" {
			command += " ofs-delta side-band-64k thin-pack"
		}
		command += "\n"
		body.WriteString(fmt.Sprintf("%04x%s", len(command)+4, command))
		last = want
	}
	body.WriteString("0000")
	for have := range haves { // TODO: sort the haves
		command := "have " + have + "\n"
		body.WriteString(fmt.Sprintf("%04x%s", len(command)+4, command))
	}
	body.WriteString("0009done\n")

	req, err = http.NewRequest("POST", gitURL+"/git-upload-pack", body)
	if err != nil {
		return
	}
	req.Header.Set("Content-Type", "application/x-git-upload-pack-request")
	req.Header.Set("Accept", "application/x-git-upload-pack-result")
	req.Header.Set("User-Agent", "github.com/thecodearchive/gitarchive/git")
	resp, err = http.DefaultClient.Do(req)
	if err != nil {
		return
	}
	if resp.StatusCode != 200 {
		return nil, nil, fmt.Errorf("POST /git-upload-pack: %d", resp.StatusCode)
	}

	sbr := &sideBandReader{Upstream: resp.Body, MsgW: msgW}
	cr := &countingReader{Upstream: sbr, Counter: bwCounter}

	// Peek into the first 32 bytes to make sure it's not an empty
	// packfile.
	// Needs to be > 32 bytes (hdr is 12, trailer is 20).
	var buf bytes.Buffer
	n, err := io.CopyN(&buf, cr, 64)
	if err == io.EOF && n == 32 {
		resp.Body.Close()
		return refs, nil, nil
	}
	if err != nil {
		resp.Body.Close()
		return
	}
	rc := struct {
		io.Reader
		io.Closer
	}{
		Reader: io.MultiReader(&buf, cr),
		Closer: resp.Body,
	}
	return refs, rc, nil
}
<|endoftext|>


<commit_before>// Copyright (C) 2017 Jennal(jennalcn@gmail.com). All rights reserved.
//
// Licensed under the MIT License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// http://opensource.org/licenses/MIT
//
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.

package service

import (
	"reflect"
	"sync"
	"time"

	"fmt"

	"github.com/jennal/goplay/aop"
	"github.com/jennal/goplay/encode"
	"github.com/jennal/goplay/filter"
	"github.com/jennal/goplay/filter/heartbeat"
	"github.com/jennal/goplay/helpers"
	"github.com/jennal/goplay/log"
	"github.com/jennal/goplay/pkg"
	"github.com/jennal/goplay/router"
	"github.com/jennal/goplay/session"
	"github.com/jennal/goplay/transfer"
)

const (
	REQUEST_TIMEOUT = 3 * time.Second
)

type requestCallbacks struct {
	startTime      time.Time
	successCallbak *Method
	failCallback   *Method
}

type ServiceClient struct {
	*session.Session

	router  *router.Router
	filters []filter.IFilter

	requestCbsMutex sync.Mutex
	requestCbs      map[pkg.PackageIDType]*requestCallbacks

	pushCbsMutex sync.Mutex
	pushCbs      map[string][]*Method
}

func NewServiceClient(cli transfer.IClient) *ServiceClient {
	result := &ServiceClient{
		Session: session.NewSession(cli),

		router: nil,
		filters: []filter.IFilter{
			heartbeat.NewHeartBeatManager(),
		},

		requestCbs: make(map[pkg.PackageIDType]*requestCallbacks),
		pushCbs:    make(map[string][]*Method),
	}
	result.setupEventLoop()

	return result
}

func (self *ServiceClient) SetRouter(router *router.Router) {
	self.router = router
}

func (self *ServiceClient) RegistFilter(filter filter.IFilter) {
	self.filters = append(self.filters, filter)
}

func (self *ServiceClient) SetFilters(filters []filter.IFilter) {
	self.filters = filters
}

func (s *ServiceClient) checkTimeoutLoop() {
	for {
		if !s.IsConnected() {
			break
		}

		ids := []pkg.PackageIDType{}

		s.requestCbsMutex.Lock()
		for id, item := range s.requestCbs {
			if time.Since(item.startTime) > REQUEST_TIMEOUT {
				ids = append(ids, id)
				item.failCallback.Call(pkg.NewErrorMessage(pkg.STAT_ERR_TIMEOUT, "Request Timeout"))
			}
		}

		for _, id := range ids {
			delete(s.requestCbs, id)
		}
		s.requestCbsMutex.Unlock()

		time.Sleep(REQUEST_TIMEOUT)
	}
}

func (s *ServiceClient) setupEventLoop() {
	s.AddListener(ON_SERVICE_DOWN, func(ok bool) {
		// log.Log(ON_SERVICE_DOWN)
		s.Disconnect()
	})

	exitChan := make(chan int, 1)
	s.On(transfer.EVENT_CLIENT_CONNECTED, s, func(client transfer.IClient) {
		if s.filters != nil && len(s.filters) > 0 {
			for _, filter := range s.filters {
				if !filter.OnNewClient(s.Session) {
					return
				}
			}
		}

		go s.checkTimeoutLoop()
		go func() {
			aop.Recover(func() {
			Loop:
				for {
					select {
					case <-exitChan:
						break Loop
					default:
						sess := session.NewSession(s)
						sess.Bind(s.ID)
						header, bodyBuf, err := s.Recv()
						if err != nil {
							log.Errorf("Recv:\n\terr => %v\n\theader => %#v\n\tbody => %#v | %v", err, header, bodyBuf, string(bodyBuf))
							s.Disconnect()
							break Loop
						}

						// if header.Type != pkg.PKG_HEARTBEAT && header.Type != pkg.PKG_HEARTBEAT_RESPONSE {
						// 	log.Logf("Recv:\n\theader => %#v\n\tbody => %#v | %v\n\terr => %v\n", header, bodyBuf, string(bodyBuf), err)
						// }

						if (header.Type & pkg.PKG_RPC) == pkg.PKG_RPC {
							sess.BindClientID(header.ClientID)
							s.BindClientID(header.ClientID) /* FIXME: this may not be good enough */
						}

						//filters
						if s.filters != nil && len(s.filters) > 0 {
							for _, filter := range s.filters {
								if !filter.OnRecv(s.Session, header, bodyBuf) {
									goto Loop
								}
							}
						}

						switch header.Type {
						case pkg.PKG_REQUEST, pkg.PKG_RPC_REQUEST:
							if s.router != nil {
								results, err := s.callRouteFunc(sess, header, bodyBuf)
								if err != nil {
									log.Errorf("CallRouteFunc:\n\terr => %v\n\theader => %#v\n\tbody => %#v | %v", err, header, bodyBuf, string(bodyBuf))
									sess.Disconnect()
									break Loop
								}
								// fmt.Printf(" => Loop result: %#v\n", results)
								err = s.response(sess, header, results)
								if err != nil {
									log.Errorf("Response:\n\terr => %v\n\theader => %#v\n\tresults => %#v", err, header, results)
									sess.Disconnect()
									break Loop
								}
							}
						case pkg.PKG_NOTIFY, pkg.PKG_RPC_NOTIFY:
							if s.router != nil {
								_, err := s.callRouteFunc(sess, header, bodyBuf)
								if err != nil {
									log.Errorf("CallRouteFunc:\n\terr => %v\n\theader => %#v\n\tbody => %#v | %v", err, header, bodyBuf, string(bodyBuf))
									sess.Disconnect()
									break Loop
								}
							}
						case pkg.PKG_PUSH, pkg.PKG_RPC_PUSH:
							s.recvPush(header, bodyBuf)
						case pkg.PKG_RESPONSE, pkg.PKG_RPC_RESPONSE:
							s.recvResponse(header, bodyBuf)
						case pkg.PKG_HEARTBEAT, pkg.PKG_HEARTBEAT_RESPONSE:
							fallthrough
						default:
							log.Errorf("Can't reach here!!\n\terr => %v\n\theader => %#v\n\tbody => %#v", err, header, bodyBuf)
							break
						}
					}
				}
			}, func(err interface{}) {
				if e, ok := err.(error); ok {
					log.Error(e)
				} else if err != nil {
					log.Errorf("%#v", err)
				} else {
					log.Errorf("%#v", err)
				}

				s.Disconnect()
			})
		}()
	})
	s.Once(transfer.EVENT_CLIENT_DISCONNECTED, s, func(cli transfer.IClient) {
		exitChan <- 1
	})
}

func (s *ServiceClient) callRouteFunc(sess *session.Session, header *pkg.Header, bodyBuf []byte) ([]interface{}, error) {
	/*
	 * 1. find route func
	 * 2. unmarshal data
	 * 3. call route func
	 */
	method := s.router.Get(header.Route)
	if method == nil {
		return nil, log.NewErrorf("Can't find method with route: %s", header.Route)
	}
	val := method.NewArg(2)
	// fmt.Printf("Service.callRouteFunc: %#v => %v\n", val, reflect.TypeOf(val))
	decoder := encode.GetEncodeDecoder(header.Encoding)
	err := decoder.Unmarshal(bodyBuf, val)
	if err != nil {
		return nil, log.NewErrorf("Service.callRouteFunc decoder.Unmarshal failed: %v", err)
	}
	// fmt.Printf("Service.callRouteFunc: %#v => %v\n", val, reflect.TypeOf(val))

	var result []interface{}
	aop.Recover(func() {
		result = method.Call(sess, helpers.GetValueFromPtr(val))
	}, func(e interface{}) {
		err = e.(error)
	})

	return result, err
}

func (s *ServiceClient) response(sess *session.Session, header *pkg.Header, results []interface{}) error {
	respHeader := *header
	if header.Type == pkg.PKG_RPC_REQUEST {
		respHeader.Type = pkg.PKG_RPC_RESPONSE
	} else {
		respHeader.Type = pkg.PKG_RESPONSE
	}

	if results == nil || len(results) <= 0 {
		return sess.Send(&respHeader, []byte{})
	}

	result := results[0]
	/* check error != nil */
	if len(results) == 2 && !reflect.ValueOf(results[1]).IsNil() {
		respHeader.Status = pkg.STAT_ERR
		result = results[1]
	}

	// fmt.Println("result:", result)

	encoder := encode.GetEncodeDecoder(header.Encoding)
	body, err := encoder.Marshal(result)
	if err != nil {
		return err
	}

	return sess.Send(&respHeader, body)
}

func (s *ServiceClient) recvPush(header *pkg.Header, body []byte) {
	s.pushCbsMutex.Lock()
	list, ok := s.pushCbs[header.Route]
	s.pushCbsMutex.Unlock()

	if !ok {
		return
	}

	for _, item := range list {
		val := item.NewArg(0)
		s.Encoder.Unmarshal(body, val)
		// log.Log("==========>\t", i, "\t", val)
		item.Call(helpers.GetValueFromPtr(val))
	}
}

func (s *ServiceClient) recvResponse(header *pkg.Header, body []byte) {
	s.requestCbsMutex.Lock()
	cbs, ok := s.requestCbs[header.ID]
	if ok {
		delete(s.requestCbs, header.ID)
	}
	s.requestCbsMutex.Unlock()

	if !ok {
		return
	}

	// log.Logf("%v %v %v", header.Status, body, string(body))
	if header.Status == pkg.STAT_OK {
		val := cbs.successCallbak.NewArg(0)
		err := s.Encoder.Unmarshal(body, val)
		if err == nil {
			cbs.successCallbak.Call(helpers.GetValueFromPtr(val))
			return
		}
	} else {
		val := cbs.failCallback.NewArg(0)
		err := s.Encoder.Unmarshal(body, val)
		if err == nil {
			cbs.failCallback.Call(helpers.GetValueFromPtr(val))
			return
		}
	}

	cbs.failCallback.Call(pkg.NewErrorMessage(
		pkg.STAT_ERR_DECODE_FAILED,
		fmt.Sprintf("decode body failed: %#v | %v", body, string(body))))
}

func (s *ServiceClient) Request(route string, data interface{}, succCb interface{}, failCb func(*pkg.ErrorMessage)) error {
	header := s.NewHeader(pkg.PKG_RPC_REQUEST, s.Encoding, route)
	header = pkg.NewRpcHeader(header, s.ClientID)
	cbs := requestCallbacks{
		successCallbak: NewMethod(succCb),
		failCallback:   NewMethod(failCb),
		startTime:      time.Now(),
	}

	s.requestCbsMutex.Lock()
	s.requestCbs[header.ID] = &cbs
	s.requestCbsMutex.Unlock()

	buf, err := s.Encoder.Marshal(data)
	if err != nil {
		return err
	}
	return s.Send(header, buf)
}

func (s *ServiceClient) Notify(route string, data interface{}) error {
	header := s.NewHeader(pkg.PKG_RPC_NOTIFY, s.Encoding, route)
	header = pkg.NewRpcHeader(header, s.ClientID)
	buf, err := s.Encoder.Marshal(data)
	if err != nil {
		return err
	}
	return s.Send(header, buf)
}

func (s *ServiceClient) Push(route string, data interface{}) error {
	return s.Session.Push(route, data)
}

func (s *ServiceClient) AddListener(route string, callback interface{}) {
	s.pushCbsMutex.Lock()
	defer s.pushCbsMutex.Unlock()

	list, ok := s.pushCbs[route]
	if !ok {
		list = make([]*Method, 0)
		s.pushCbs[route] = list
	}

	s.pushCbs[route] = append(s.pushCbs[route], NewMethod(callback))
}
<commit_msg>fix different sesion on handler issue<commit_after>// Copyright (C) 2017 Jennal(jennalcn@gmail.com). All rights reserved.
//
// Licensed under the MIT License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// http://opensource.org/licenses/MIT
//
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.

package service

import (
	"reflect"
	"sync"
	"time"

	"fmt"

	"github.com/jennal/goplay/aop"
	"github.com/jennal/goplay/encode"
	"github.com/jennal/goplay/filter"
	"github.com/jennal/goplay/filter/heartbeat"
	"github.com/jennal/goplay/helpers"
	"github.com/jennal/goplay/log"
	"github.com/jennal/goplay/pkg"
	"github.com/jennal/goplay/router"
	"github.com/jennal/goplay/session"
	"github.com/jennal/goplay/transfer"
)

const (
	REQUEST_TIMEOUT = 3 * time.Second
)

type requestCallbacks struct {
	startTime      time.Time
	successCallbak *Method
	failCallback   *Method
}

type ServiceClient struct {
	*session.Session
	sessionManager *session.SessionManager

	router  *router.Router
	filters []filter.IFilter

	requestCbsMutex sync.Mutex
	requestCbs      map[pkg.PackageIDType]*requestCallbacks

	pushCbsMutex sync.Mutex
	pushCbs      map[string][]*Method
}

func NewServiceClient(cli transfer.IClient) *ServiceClient {
	result := &ServiceClient{
		Session:        session.NewSession(cli),
		sessionManager: session.NewSessionManager(),

		router: nil,
		filters: []filter.IFilter{
			heartbeat.NewHeartBeatManager(),
		},

		requestCbs: make(map[pkg.PackageIDType]*requestCallbacks),
		pushCbs:    make(map[string][]*Method),
	}
	result.setupEventLoop()

	return result
}

func (self *ServiceClient) SetRouter(router *router.Router) {
	self.router = router
}

func (self *ServiceClient) RegistFilter(filter filter.IFilter) {
	self.filters = append(self.filters, filter)
}

func (self *ServiceClient) SetFilters(filters []filter.IFilter) {
	self.filters = filters
}

func (s *ServiceClient) checkTimeoutLoop() {
	for {
		if !s.IsConnected() {
			break
		}

		ids := []pkg.PackageIDType{}

		s.requestCbsMutex.Lock()
		for id, item := range s.requestCbs {
			if time.Since(item.startTime) > REQUEST_TIMEOUT {
				ids = append(ids, id)
				item.failCallback.Call(pkg.NewErrorMessage(pkg.STAT_ERR_TIMEOUT, "Request Timeout"))
			}
		}

		for _, id := range ids {
			delete(s.requestCbs, id)
		}
		s.requestCbsMutex.Unlock()

		time.Sleep(REQUEST_TIMEOUT)
	}
}

func (s *ServiceClient) setupEventLoop() {
	s.AddListener(ON_SERVICE_DOWN, func(ok bool) {
		// log.Log(ON_SERVICE_DOWN)
		s.Disconnect()
	})

	exitChan := make(chan int, 1)
	s.On(transfer.EVENT_CLIENT_CONNECTED, s, func(client transfer.IClient) {
		if s.filters != nil && len(s.filters) > 0 {
			for _, filter := range s.filters {
				if !filter.OnNewClient(s.Session) {
					return
				}
			}
		}

		go s.checkTimeoutLoop()
		go func() {
			aop.Recover(func() {
			Loop:
				for {
					select {
					case <-exitChan:
						break Loop
					default:
						header, bodyBuf, err := s.Recv()
						if err != nil {
							log.Errorf("Recv:\n\terr => %v\n\theader => %#v\n\tbody => %#v | %v", err, header, bodyBuf, string(bodyBuf))
							s.Disconnect()
							break Loop
						}

						// if header.Type != pkg.PKG_HEARTBEAT && header.Type != pkg.PKG_HEARTBEAT_RESPONSE {
						// 	log.Logf("Recv:\n\theader => %#v\n\tbody => %#v | %v\n\terr => %v\n", header, bodyBuf, string(bodyBuf), err)
						// }

						sess := s.sessionManager.GetSessionByID(s.ID, header.ClientID)
						if sess == nil {
							sess = session.NewSession(s)
							sess.Bind(s.ID)
							sess.BindClientID(header.ClientID)
							s.sessionManager.Add(sess)
						}

						if (header.Type & pkg.PKG_RPC) == pkg.PKG_RPC {
							s.BindClientID(header.ClientID) /* FIXME: this may not be good enough */
						}

						//filters
						if s.filters != nil && len(s.filters) > 0 {
							for _, filter := range s.filters {
								if !filter.OnRecv(s.Session, header, bodyBuf) {
									goto Loop
								}
							}
						}

						switch header.Type {
						case pkg.PKG_REQUEST, pkg.PKG_RPC_REQUEST:
							if s.router != nil {
								results, err := s.callRouteFunc(sess, header, bodyBuf)
								if err != nil {
									log.Errorf("CallRouteFunc:\n\terr => %v\n\theader => %#v\n\tbody => %#v | %v", err, header, bodyBuf, string(bodyBuf))
									sess.Disconnect()
									break Loop
								}
								// fmt.Printf(" => Loop result: %#v\n", results)
								err = s.response(sess, header, results)
								if err != nil {
									log.Errorf("Response:\n\terr => %v\n\theader => %#v\n\tresults => %#v", err, header, results)
									sess.Disconnect()
									break Loop
								}
							}
						case pkg.PKG_NOTIFY, pkg.PKG_RPC_NOTIFY:
							if s.router != nil {
								_, err := s.callRouteFunc(sess, header, bodyBuf)
								if err != nil {
									log.Errorf("CallRouteFunc:\n\terr => %v\n\theader => %#v\n\tbody => %#v | %v", err, header, bodyBuf, string(bodyBuf))
									sess.Disconnect()
									break Loop
								}
							}
						case pkg.PKG_PUSH, pkg.PKG_RPC_PUSH:
							s.recvPush(header, bodyBuf)
						case pkg.PKG_RESPONSE, pkg.PKG_RPC_RESPONSE:
							s.recvResponse(header, bodyBuf)
						case pkg.PKG_HEARTBEAT, pkg.PKG_HEARTBEAT_RESPONSE:
							fallthrough
						default:
							log.Errorf("Can't reach here!!\n\terr => %v\n\theader => %#v\n\tbody => %#v", err, header, bodyBuf)
							break
						}
					}
				}
			}, func(err interface{}) {
				if e, ok := err.(error); ok {
					log.Error(e)
				} else if err != nil {
					log.Errorf("%#v", err)
				} else {
					log.Errorf("%#v", err)
				}

				s.Disconnect()
			})
		}()
	})
	s.Once(transfer.EVENT_CLIENT_DISCONNECTED, s, func(cli transfer.IClient) {
		exitChan <- 1
	})
}

func (s *ServiceClient) callRouteFunc(sess *session.Session, header *pkg.Header, bodyBuf []byte) ([]interface{}, error) {
	/*
	 * 1. find route func
	 * 2. unmarshal data
	 * 3. call route func
	 */
	method := s.router.Get(header.Route)
	if method == nil {
		return nil, log.NewErrorf("Can't find method with route: %s", header.Route)
	}
	val := method.NewArg(2)
	// fmt.Printf("Service.callRouteFunc: %#v => %v\n", val, reflect.TypeOf(val))
	decoder := encode.GetEncodeDecoder(header.Encoding)
	err := decoder.Unmarshal(bodyBuf, val)
	if err != nil {
		return nil, log.NewErrorf("Service.callRouteFunc decoder.Unmarshal failed: %v", err)
	}
	// fmt.Printf("Service.callRouteFunc: %#v => %v\n", val, reflect.TypeOf(val))

	var result []interface{}
	aop.Recover(func() {
		result = method.Call(sess, helpers.GetValueFromPtr(val))
	}, func(e interface{}) {
		err = e.(error)
	})

	return result, err
}

func (s *ServiceClient) response(sess *session.Session, header *pkg.Header, results []interface{}) error {
	respHeader := *header
	if header.Type == pkg.PKG_RPC_REQUEST {
		respHeader.Type = pkg.PKG_RPC_RESPONSE
	} else {
		respHeader.Type = pkg.PKG_RESPONSE
	}

	if results == nil || len(results) <= 0 {
		return sess.Send(&respHeader, []byte{})
	}

	result := results[0]
	/* check error != nil */
	if len(results) == 2 && !reflect.ValueOf(results[1]).IsNil() {
		respHeader.Status = pkg.STAT_ERR
		result = results[1]
	}

	// fmt.Println("result:", result)

	encoder := encode.GetEncodeDecoder(header.Encoding)
	body, err := encoder.Marshal(result)
	if err != nil {
		return err
	}

	return sess.Send(&respHeader, body)
}

func (s *ServiceClient) recvPush(header *pkg.Header, body []byte) {
	s.pushCbsMutex.Lock()
	list, ok := s.pushCbs[header.Route]
	s.pushCbsMutex.Unlock()

	if !ok {
		return
	}

	for _, item := range list {
		val := item.NewArg(0)
		s.Encoder.Unmarshal(body, val)
		// log.Log("==========>\t", i, "\t", val)
		item.Call(helpers.GetValueFromPtr(val))
	}
}

func (s *ServiceClient) recvResponse(header *pkg.Header, body []byte) {
	s.requestCbsMutex.Lock()
	cbs, ok := s.requestCbs[header.ID]
	if ok {
		delete(s.requestCbs, header.ID)
	}
	s.requestCbsMutex.Unlock()

	if !ok {
		return
	}

	// log.Logf("%v %v %v", header.Status, body, string(body))
	if header.Status == pkg.STAT_OK {
		val := cbs.successCallbak.NewArg(0)
		err := s.Encoder.Unmarshal(body, val)
		if err == nil {
			cbs.successCallbak.Call(helpers.GetValueFromPtr(val))
			return
		}
	} else {
		val := cbs.failCallback.NewArg(0)
		err := s.Encoder.Unmarshal(body, val)
		if err == nil {
			cbs.failCallback.Call(helpers.GetValueFromPtr(val))
			return
		}
	}

	cbs.failCallback.Call(pkg.NewErrorMessage(
		pkg.STAT_ERR_DECODE_FAILED,
		fmt.Sprintf("decode body failed: %#v | %v", body, string(body))))
}

func (s *ServiceClient) Request(route string, data interface{}, succCb interface{}, failCb func(*pkg.ErrorMessage)) error {
	header := s.NewHeader(pkg.PKG_RPC_REQUEST, s.Encoding, route)
	header = pkg.NewRpcHeader(header, s.ClientID)
	cbs := requestCallbacks{
		successCallbak: NewMethod(succCb),
		failCallback:   NewMethod(failCb),
		startTime:      time.Now(),
	}

	s.requestCbsMutex.Lock()
	s.requestCbs[header.ID] = &cbs
	s.requestCbsMutex.Unlock()

	buf, err := s.Encoder.Marshal(data)
	if err != nil {
		return err
	}
	return s.Send(header, buf)
}

func (s *ServiceClient) Notify(route string, data interface{}) error {
	header := s.NewHeader(pkg.PKG_RPC_NOTIFY, s.Encoding, route)
	header = pkg.NewRpcHeader(header, s.ClientID)
	buf, err := s.Encoder.Marshal(data)
	if err != nil {
		return err
	}
	return s.Send(header, buf)
}

func (s *ServiceClient) Push(route string, data interface{}) error {
	return s.Session.Push(route, data)
}

func (s *ServiceClient) AddListener(route string, callback interface{}) {
	s.pushCbsMutex.Lock()
	defer s.pushCbsMutex.Unlock()

	list, ok := s.pushCbs[route]
	if !ok {
		list = make([]*Method, 0)
		s.pushCbs[route] = list
	}

	s.pushCbs[route] = append(s.pushCbs[route], NewMethod(callback))
}
<|endoftext|>


<commit_before>// +build js

package atomic

import (
	"unsafe"
)

func SwapInt32(addr *int32, new int32) int32 {
	old := *addr
	*addr = new
	return old
}

func SwapInt64(addr *int64, new int64) int64 {
	old := *addr
	*addr = new
	return old
}

func SwapUint32(addr *uint32, new uint32) uint32 {
	old := *addr
	*addr = new
	return old
}

func SwapUint64(addr *uint64, new uint64) uint64 {
	old := *addr
	*addr = new
	return old
}

func SwapUintptr(addr *uintptr, new uintptr) uintptr {
	old := *addr
	*addr = new
	return old
}

func SwapPointer(addr *unsafe.Pointer, new unsafe.Pointer) unsafe.Pointer {
	old := *addr
	*addr = new
	return old
}

func CompareAndSwapInt32(addr *int32, old, new int32) bool {
	if *addr == old {
		*addr = new
		return true
	}
	return false
}

func CompareAndSwapInt64(addr *int64, old, new int64) bool {
	if *addr == old {
		*addr = new
		return true
	}
	return false
}

func CompareAndSwapUint32(addr *uint32, old, new uint32) bool {
	if *addr == old {
		*addr = new
		return true
	}
	return false
}

func CompareAndSwapUint64(addr *uint64, old, new uint64) bool {
	if *addr == old {
		*addr = new
		return true
	}
	return false
}

func CompareAndSwapUintptr(addr *uintptr, old, new uintptr) bool {
	if *addr == old {
		*addr = new
		return true
	}
	return false
}

func CompareAndSwapPointer(addr *unsafe.Pointer, old, new unsafe.Pointer) bool {
	if *addr == old {
		*addr = new
		return true
	}
	return false
}

func AddInt32(addr *int32, delta int32) int32 {
	new := *addr + delta
	*addr = new
	return new
}

func AddUint32(addr *uint32, delta uint32) uint32 {
	new := *addr + delta
	*addr = new
	return new
}

func AddInt64(addr *int64, delta int64) int64 {
	new := *addr + delta
	*addr = new
	return new
}

func AddUint64(addr *uint64, delta uint64) uint64 {
	new := *addr + delta
	*addr = new
	return new
}

func AddUintptr(addr *uintptr, delta uintptr) uintptr {
	new := *addr + delta
	*addr = new
	return new
}

func LoadInt32(addr *int32) int32 {
	return *addr
}

func LoadInt64(addr *int64) int64 {
	return *addr
}

func LoadUint32(addr *uint32) uint32 {
	return *addr
}

func LoadUint64(addr *uint64) uint64 {
	return *addr
}

func LoadUintptr(addr *uintptr) uintptr {
	return *addr
}

func LoadPointer(addr *unsafe.Pointer) unsafe.Pointer {
	return *addr
}

func StoreInt32(addr *int32, val int32) {
	*addr = val
}

func StoreInt64(addr *int64, val int64) {
	*addr = val
}

func StoreUint32(addr *uint32, val uint32) {
	*addr = val
}

func StoreUint64(addr *uint64, val uint64) {
	*addr = val
}

func StoreUintptr(addr *uintptr, val uintptr) {
	*addr = val
}

func StorePointer(addr *unsafe.Pointer, val unsafe.Pointer) {
	*addr = val
}
<commit_msg>go1.4: added native implementations for sync/atomic.Value<commit_after>// +build js

package atomic

import (
	"unsafe"

	"github.com/gopherjs/gopherjs/js"
)

func SwapInt32(addr *int32, new int32) int32 {
	old := *addr
	*addr = new
	return old
}

func SwapInt64(addr *int64, new int64) int64 {
	old := *addr
	*addr = new
	return old
}

func SwapUint32(addr *uint32, new uint32) uint32 {
	old := *addr
	*addr = new
	return old
}

func SwapUint64(addr *uint64, new uint64) uint64 {
	old := *addr
	*addr = new
	return old
}

func SwapUintptr(addr *uintptr, new uintptr) uintptr {
	old := *addr
	*addr = new
	return old
}

func SwapPointer(addr *unsafe.Pointer, new unsafe.Pointer) unsafe.Pointer {
	old := *addr
	*addr = new
	return old
}

func CompareAndSwapInt32(addr *int32, old, new int32) bool {
	if *addr == old {
		*addr = new
		return true
	}
	return false
}

func CompareAndSwapInt64(addr *int64, old, new int64) bool {
	if *addr == old {
		*addr = new
		return true
	}
	return false
}

func CompareAndSwapUint32(addr *uint32, old, new uint32) bool {
	if *addr == old {
		*addr = new
		return true
	}
	return false
}

func CompareAndSwapUint64(addr *uint64, old, new uint64) bool {
	if *addr == old {
		*addr = new
		return true
	}
	return false
}

func CompareAndSwapUintptr(addr *uintptr, old, new uintptr) bool {
	if *addr == old {
		*addr = new
		return true
	}
	return false
}

func CompareAndSwapPointer(addr *unsafe.Pointer, old, new unsafe.Pointer) bool {
	if *addr == old {
		*addr = new
		return true
	}
	return false
}

func AddInt32(addr *int32, delta int32) int32 {
	new := *addr + delta
	*addr = new
	return new
}

func AddUint32(addr *uint32, delta uint32) uint32 {
	new := *addr + delta
	*addr = new
	return new
}

func AddInt64(addr *int64, delta int64) int64 {
	new := *addr + delta
	*addr = new
	return new
}

func AddUint64(addr *uint64, delta uint64) uint64 {
	new := *addr + delta
	*addr = new
	return new
}

func AddUintptr(addr *uintptr, delta uintptr) uintptr {
	new := *addr + delta
	*addr = new
	return new
}

func LoadInt32(addr *int32) int32 {
	return *addr
}

func LoadInt64(addr *int64) int64 {
	return *addr
}

func LoadUint32(addr *uint32) uint32 {
	return *addr
}

func LoadUint64(addr *uint64) uint64 {
	return *addr
}

func LoadUintptr(addr *uintptr) uintptr {
	return *addr
}

func LoadPointer(addr *unsafe.Pointer) unsafe.Pointer {
	return *addr
}

func StoreInt32(addr *int32, val int32) {
	*addr = val
}

func StoreInt64(addr *int64, val int64) {
	*addr = val
}

func StoreUint32(addr *uint32, val uint32) {
	*addr = val
}

func StoreUint64(addr *uint64, val uint64) {
	*addr = val
}

func StoreUintptr(addr *uintptr, val uintptr) {
	*addr = val
}

func StorePointer(addr *unsafe.Pointer, val unsafe.Pointer) {
	*addr = val
}

func (v *Value) Load() (x interface{}) {
	return v.v
}

func (v *Value) Store(x interface{}) {
	if x == nil {
		panic("sync/atomic: store of nil value into Value")
	}
	if v.v != nil && js.InternalObject(x).Get("constructor") != js.InternalObject(v.v).Get("constructor") {
		panic("sync/atomic: store of inconsistently typed value into Value")
	}
	v.v = x
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"reflect"
	"strconv"
)

// Arguments are the strings given at the end of the command line.
type argument struct {
	field    string
	desc     string
	position int
	variadic bool
	required bool
	value    string
	values   []string
}

// Reflect the gathered information into the concrete action instance.
func (arg *argument) reflectTo(value reflect.Value) (e error) {
	if arg.value == "" && len(arg.values) == 0 {
		if arg.required {
			return fmt.Errorf("required argument not set")
		}
		return nil
	}

	switch {
	case arg.variadic:
		field := value.FieldByName(arg.field)
		return setSlice(field, arg.values)
	default:
		field := value.FieldByName(arg.field)
		return arg.setField(field, arg.value)
	}
}

func (arg *argument) setField(target reflect.Value, source string) (e error) {
	switch target.Kind() {
	case reflect.String:
		target.SetString(source)
	case reflect.Int, reflect.Int64:
		i, e := strconv.ParseInt(source, 10, 64)
		if e != nil {
			return fmt.Errorf(`argument %q at index "%d" has wrong type`, arg.field, arg.position)
		}
		target.SetInt(i)
	default:
		return fmt.Errorf("invalid type %q", target.Type().String())
	}
	return nil
}

func setSlice(target reflect.Value, source []string) (e error) {
	switch target.Type().Elem().Kind() {
	case reflect.String:
		target.Set(reflect.ValueOf(source))
	case reflect.Int:
		result := make([]int, 0, len(source))
		for _, i := range source {
			j, e := strconv.Atoi(i)
			if e != nil {
				return e
			}
			result = append(result, j)
		}
		target.Set(reflect.ValueOf(result))
	default:
		return fmt.Errorf("..invalid type %q", target.Type().String())
	}
	return nil
}

func (arg *argument) setValue(value string) {
	if arg.variadic {
		arg.values = append(arg.values, value)
	} else {
		arg.value = value
	}
}

func (a *argument) description() string {
	desc := fmt.Sprintf("    %s", a.shortDescription())
	desc += fmt.Sprintf("%-*s", 30-len(desc), " ") + a.desc
	return desc
}

func (a *argument) shortDescription() (desc string) {
	desc = "<" + a.field + ">"
	if !a.required {
		desc += "?"
	}
	if a.variadic {
		desc += "..."
	}
	return desc
}

func (a *action) argumentForPosition(argIdx int) *argument {
	for idx := range a.args {
		arg := a.args[idx]
		if arg.position == argIdx || (arg.variadic && arg.position <= argIdx) {
			return arg
		}
	}
	return nil
}

func (a *action) createArgument(field reflect.StructField, value reflect.Value, tagMap map[string]string) (e error) {
	if e := validateTagMap(tagMap, "type", "desc", "required"); e != nil {
		return fmt.Errorf("[argument:%s] %s", field.Name, e.Error())
	}

	arg := &argument{field: field.Name, position: 0}

	arg.required, e = handleRequired(tagMap)
	if e != nil {
		return e
	}

	arg.variadic = field.Type.Kind() == reflect.Slice

	arg.desc = handleDescription(tagMap)

	if len(a.args) == 0 {
		arg.position = 0
	} else {
		prevArg := a.args[len(a.args)-1]
		if prevArg.variadic {
			return fmt.Errorf("only last argument can be variadic")
		}
		arg.position = prevArg.position + 1
	}

	a.args = append(a.args, arg)
	return nil
}
<commit_msg>cli: try to respect preset default values for required arguments<commit_after>package cli

import (
	"fmt"
	"reflect"
	"strconv"
)

// Arguments are the strings given at the end of the command line.
type argument struct {
	field    string
	desc     string
	position int
	variadic bool
	required bool
	value    string
	values   []string
}

// Reflect the gathered information into the concrete action instance.
func (arg *argument) reflectTo(value reflect.Value) (e error) {
	if arg.value == "" && len(arg.values) == 0 {
		if arg.required && isFieldBlank(value, arg.field) {
			return fmt.Errorf("required argument not set")
		}
		return nil
	}

	switch {
	case arg.variadic:
		field := value.FieldByName(arg.field)
		return setSlice(field, arg.values)
	default:
		field := value.FieldByName(arg.field)
		return arg.setField(field, arg.value)
	}
}

func (arg *argument) setField(target reflect.Value, source string) (e error) {
	switch target.Kind() {
	case reflect.String:
		target.SetString(source)
	case reflect.Int, reflect.Int64:
		i, e := strconv.ParseInt(source, 10, 64)
		if e != nil {
			return fmt.Errorf(`argument %q at index "%d" has wrong type`, arg.field, arg.position)
		}
		target.SetInt(i)
	default:
		return fmt.Errorf("invalid type %q", target.Type().String())
	}
	return nil
}

func isFieldBlank(value reflect.Value, name string) bool {
	if value.Type().Kind() == reflect.Ptr {
		return true
	}
	switch c := value.FieldByName(name).Interface().(type) {
	case string:
		return c == ""
	case int:
		return c == 0
	default:
		return true
	}
}

func setSlice(target reflect.Value, source []string) (e error) {
	switch target.Type().Elem().Kind() {
	case reflect.String:
		target.Set(reflect.ValueOf(source))
	case reflect.Int:
		result := make([]int, 0, len(source))
		for _, i := range source {
			j, e := strconv.Atoi(i)
			if e != nil {
				return e
			}
			result = append(result, j)
		}
		target.Set(reflect.ValueOf(result))
	default:
		return fmt.Errorf("..invalid type %q", target.Type().String())
	}
	return nil
}

func (arg *argument) setValue(value string) {
	if arg.variadic {
		arg.values = append(arg.values, value)
	} else {
		arg.value = value
	}
}

func (a *argument) description() string {
	desc := fmt.Sprintf("    %s", a.shortDescription())
	desc += fmt.Sprintf("%-*s", 30-len(desc), " ") + a.desc
	return desc
}

func (a *argument) shortDescription() (desc string) {
	desc = "<" + a.field + ">"
	if !a.required {
		desc += "?"
	}
	if a.variadic {
		desc += "..."
	}
	return desc
}

func (a *action) argumentForPosition(argIdx int) *argument {
	for idx := range a.args {
		arg := a.args[idx]
		if arg.position == argIdx || (arg.variadic && arg.position <= argIdx) {
			return arg
		}
	}
	return nil
}

func (a *action) createArgument(field reflect.StructField, value reflect.Value, tagMap map[string]string) (e error) {
	if e := validateTagMap(tagMap, "type", "desc", "required"); e != nil {
		return fmt.Errorf("[argument:%s] %s", field.Name, e.Error())
	}

	arg := &argument{field: field.Name, position: 0}

	arg.required, e = handleRequired(tagMap)
	if e != nil {
		return e
	}

	arg.variadic = field.Type.Kind() == reflect.Slice

	arg.desc = handleDescription(tagMap)

	if len(a.args) == 0 {
		arg.position = 0
	} else {
		prevArg := a.args[len(a.args)-1]
		if prevArg.variadic {
			return fmt.Errorf("only last argument can be variadic")
		}
		arg.position = prevArg.position + 1
	}

	a.args = append(a.args, arg)
	return nil
}
<|endoftext|>


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

package cmd

import (
	"encoding/json"
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/control-center/serviced/cli/api"
	"github.com/control-center/serviced/domain/pool"
)

// Initializer for serviced pool subcommands
func (c *ServicedCli) initPool() {
	c.app.Commands = append(c.app.Commands, cli.Command{
		Name:        "pool",
		Usage:       "Administers pool data",
		Description: "",
		Subcommands: []cli.Command{
			{
				Name:         "list",
				Usage:        "Lists all pools",
				Description:  "serviced pool list [POOLID]",
				BashComplete: c.printPoolsFirst,
				Action:       c.cmdPoolList,
				Flags: []cli.Flag{
					cli.BoolFlag{"verbose, v", "Show JSON format"},
				},
			}, {
				Name:  "add",
				Usage: "Adds a new resource pool",
				//Description:  "serviced pool add POOLID CORE_LIMIT MEMORY_LIMIT PRIORITY REALM",
				Description:  "serviced pool add POOLID PRIORITY",
				BashComplete: nil,
				Action:       c.cmdPoolAdd,
			}, {
				Name:         "remove",
				ShortName:    "rm",
				Usage:        "Removes an existing resource pool",
				Description:  "serviced pool remove POOLID ...",
				BashComplete: c.printPoolsAll,
				Action:       c.cmdPoolRemove,
			}, {
				Name:         "list-ips",
				Usage:        "Lists the IP addresses for a resource pool",
				Description:  "serviced pool list-ips POOLID",
				BashComplete: c.printPoolsFirst,
				Action:       c.cmdPoolListIPs,
				Flags: []cli.Flag{
					cli.BoolFlag{"verbose, v", "Show JSON format"},
				},
			}, {
				Name:         "add-virtual-ip",
				Usage:        "Add a virtual IP address to a pool",
				Description:  "serviced pool add-virtual-ip POOLID IPADDRESS NETMASK BINDINTERFACE",
				BashComplete: c.printPoolsFirst,
				Action:       c.cmdAddVirtualIP,
			}, {
				Name:         "remove-virtual-ip",
				Usage:        "Remove a virtual IP address from a pool",
				Description:  "serviced pool remove-virtual-ip POOLID IPADDRESS",
				BashComplete: c.printPoolsFirst,
				Action:       c.cmdRemoveVirtualIP,
			},
		},
	})
}

// Returns a list of available pools
func (c *ServicedCli) pools() (data []string) {
	pools, err := c.driver.GetResourcePools()
	if err != nil || pools == nil || len(pools) == 0 {
		return
	}

	data = make([]string, len(pools))
	for i, p := range pools {
		data[i] = p.ID
	}

	return
}

// Bash-completion command that prints the list of available pools as the
// first argument
func (c *ServicedCli) printPoolsFirst(ctx *cli.Context) {
	if len(ctx.Args()) > 0 {
		return
	}
	fmt.Println(strings.Join(c.pools(), "\n"))
}

// Bash-completion command that prints the list of available pools as all
// arguments
func (c *ServicedCli) printPoolsAll(ctx *cli.Context) {
	args := ctx.Args()
	pools := c.pools()

	for _, p := range pools {
		for _, a := range args {
			if p == a {
				goto next
			}
		}
		fmt.Println(p)
	next:
	}
}

// serviced pool list [POOLID]
func (c *ServicedCli) cmdPoolList(ctx *cli.Context) {
	if len(ctx.Args()) > 0 {
		poolID := ctx.Args()[0]
		if pool, err := c.driver.GetResourcePool(poolID); err != nil {
			fmt.Fprintln(os.Stderr, err)
		} else if pool == nil {
			fmt.Fprintln(os.Stderr, "pool not found")
		} else if jsonPool, err := json.MarshalIndent(pool, " ", "  "); err != nil {
			fmt.Fprintf(os.Stderr, "failed to marshal resource pool: %s", err)
		} else {
			fmt.Println(string(jsonPool))
		}
		return
	}

	pools, err := c.driver.GetResourcePools()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return
	} else if pools == nil || len(pools) == 0 {
		fmt.Fprintln(os.Stderr, "no resource pools found")
		return
	}

	if ctx.Bool("verbose") {
		if jsonPool, err := json.MarshalIndent(pools, " ", "  "); err != nil {
			fmt.Fprintf(os.Stderr, "failed to marshal resource pool list: %s", err)
		} else {
			fmt.Println(string(jsonPool))
		}
	} else {
		tablePool := newtable(0, 8, 2)
		tablePool.printrow("ID", "PARENT" /*"CORE", "MEM",*/, "PRI")
		for _, p := range pools {
			tablePool.printrow(p.ID, p.ParentID /*p.CoreLimit, p.MemoryLimit,*/, p.Priority)
		}
		tablePool.flush()
	}
}

// serviced pool add POOLID PRIORITY
func (c *ServicedCli) cmdPoolAdd(ctx *cli.Context) {
	args := ctx.Args()
	if len(args) < 2 {
		fmt.Printf("Incorrect Usage.\n\n")
		cli.ShowCommandHelp(ctx, "add")
		return
	}

	var err error

	cfg := api.PoolConfig{}
	cfg.PoolID = args[0]

	/* Disabled until enforced. See ZEN-11450
	cfg.CoreLimit, err = strconv.Atoi(args[1])
	if err != nil {
		fmt.Println("CORE_LIMIT must be a number")
		return
	}

	cfg.MemoryLimit, err = strconv.ParseUint(args[2], 10, 64)
	if err != nil {
		fmt.Println("MEMORY_LIMIT must be a number")
		return
	}
	*/

	cfg.Priority, err = strconv.Atoi(args[1])
	if err != nil {
		fmt.Println("PRIORITY must be a number")
		return
	}

	/* TODO: 1.1
	if len(args) > 2 {
		cfg.Realm = args[2]
	}
	*/

	if pool, err := c.driver.AddResourcePool(cfg); err != nil {
		fmt.Fprintln(os.Stderr, err)
	} else if pool == nil {
		fmt.Fprintln(os.Stderr, "received nil resource pool")
	} else {
		fmt.Println(pool.ID)
	}
}

// serviced pool remove POOLID ...
func (c *ServicedCli) cmdPoolRemove(ctx *cli.Context) {
	args := ctx.Args()
	if len(args) < 1 {
		fmt.Printf("Incorrect Usage.\n\n")
		cli.ShowCommandHelp(ctx, "remove")
	}

	for _, id := range args {
		if p, err := c.driver.GetResourcePool(id); err != nil {
			fmt.Fprintf(os.Stderr, "%s: %s\n", id, err)
		} else if p == nil {
			fmt.Fprintf(os.Stderr, "%s: pool not found", id)
		} else if err := c.driver.RemoveResourcePool(id); err != nil {
			fmt.Fprintf(os.Stderr, "%s: %s\n", id, err)
		} else {
			fmt.Println(id)
		}
	}
}

// serviced pool list-ips POOLID
func (c *ServicedCli) cmdPoolListIPs(ctx *cli.Context) {
	args := ctx.Args()
	if len(args) < 1 {
		fmt.Printf("Incorrect Usage.\n\n")
		cli.ShowCommandHelp(ctx, "list-ips")
		return
	}

	if poolIps, err := c.driver.GetPoolIPs(args[0]); err != nil {
		fmt.Fprintln(os.Stderr, err)
		return
	} else if poolIps.HostIPs == nil || (len(poolIps.HostIPs) == 0 && len(poolIps.VirtualIPs) == 0) {
		fmt.Fprintln(os.Stderr, "no resource pool IPs found")
		return
	} else if ctx.Bool("verbose") {
		if jsonPoolIP, err := json.MarshalIndent(poolIps.HostIPs, " ", "  "); err != nil {
			fmt.Fprintf(os.Stderr, "failed to marshal resource pool IPs: %s", err)
		} else {
			fmt.Println(string(jsonPoolIP))
		}
	} else {
		tableIPs := newtable(0, 10, 2)
		tableIPs.printrow("Interface Name", "IP Address", "Type")
		for _, ip := range poolIps.HostIPs {
			tableIPs.printrow(ip.InterfaceName, ip.IPAddress, "static")
		}
		for _, ip := range poolIps.VirtualIPs {
			tableIPs.printrow("", ip.IP, "virtual")
		}
		tableIPs.flush()
	}
}

// serviced pool add-virtual-ip POOLID IPADDRESS NETMASK BINDINTERFACE
func (c *ServicedCli) cmdAddVirtualIP(ctx *cli.Context) {
	args := ctx.Args()
	if len(args) < 1 || len(args) > 4 {
		fmt.Printf("Incorrect Usage.\n\n")
		cli.ShowCommandHelp(ctx, "add-virtual-ip")
		return
	}

	requestVirtualIP := pool.VirtualIP{PoolID: args[0], IP: args[1], Netmask: args[2], BindInterface: args[3]}
	if err := c.driver.AddVirtualIP(requestVirtualIP); err != nil {
		fmt.Fprintln(os.Stderr, err)
		return
	}

	fmt.Println("Added virtual IP:", args[1])
}

// serviced pool remove-virtual-ip POOLID IPADDRESS
func (c *ServicedCli) cmdRemoveVirtualIP(ctx *cli.Context) {
	args := ctx.Args()
	if len(args) < 2 {
		fmt.Printf("Incorrect Usage.\n\n")
		cli.ShowCommandHelp(ctx, "remove-virtual-ip")
		return
	}

	requestVirtualIP := pool.VirtualIP{PoolID: args[0], IP: args[1], Netmask: "", BindInterface: ""}
	if err := c.driver.RemoveVirtualIP(requestVirtualIP); err != nil {
		fmt.Fprintln(os.Stderr, err)
		return
	} else {
		fmt.Printf("Removed virtual IP: %v from pool %v\n", args[1], args[0])
	}
}
<commit_msg>CC-250: adding VIP in CLI tracebacks when incorrectly used<commit_after>// Copyright 2014 The Serviced Authors.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"encoding/json"
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/control-center/serviced/cli/api"
	"github.com/control-center/serviced/domain/pool"
)

// Initializer for serviced pool subcommands
func (c *ServicedCli) initPool() {
	c.app.Commands = append(c.app.Commands, cli.Command{
		Name:        "pool",
		Usage:       "Administers pool data",
		Description: "",
		Subcommands: []cli.Command{
			{
				Name:         "list",
				Usage:        "Lists all pools",
				Description:  "serviced pool list [POOLID]",
				BashComplete: c.printPoolsFirst,
				Action:       c.cmdPoolList,
				Flags: []cli.Flag{
					cli.BoolFlag{"verbose, v", "Show JSON format"},
				},
			}, {
				Name:  "add",
				Usage: "Adds a new resource pool",
				//Description:  "serviced pool add POOLID CORE_LIMIT MEMORY_LIMIT PRIORITY REALM",
				Description:  "serviced pool add POOLID PRIORITY",
				BashComplete: nil,
				Action:       c.cmdPoolAdd,
			}, {
				Name:         "remove",
				ShortName:    "rm",
				Usage:        "Removes an existing resource pool",
				Description:  "serviced pool remove POOLID ...",
				BashComplete: c.printPoolsAll,
				Action:       c.cmdPoolRemove,
			}, {
				Name:         "list-ips",
				Usage:        "Lists the IP addresses for a resource pool",
				Description:  "serviced pool list-ips POOLID",
				BashComplete: c.printPoolsFirst,
				Action:       c.cmdPoolListIPs,
				Flags: []cli.Flag{
					cli.BoolFlag{"verbose, v", "Show JSON format"},
				},
			}, {
				Name:         "add-virtual-ip",
				Usage:        "Add a virtual IP address to a pool",
				Description:  "serviced pool add-virtual-ip POOLID IPADDRESS NETMASK BINDINTERFACE",
				BashComplete: c.printPoolsFirst,
				Action:       c.cmdAddVirtualIP,
			}, {
				Name:         "remove-virtual-ip",
				Usage:        "Remove a virtual IP address from a pool",
				Description:  "serviced pool remove-virtual-ip POOLID IPADDRESS",
				BashComplete: c.printPoolsFirst,
				Action:       c.cmdRemoveVirtualIP,
			},
		},
	})
}

// Returns a list of available pools
func (c *ServicedCli) pools() (data []string) {
	pools, err := c.driver.GetResourcePools()
	if err != nil || pools == nil || len(pools) == 0 {
		return
	}

	data = make([]string, len(pools))
	for i, p := range pools {
		data[i] = p.ID
	}

	return
}

// Bash-completion command that prints the list of available pools as the
// first argument
func (c *ServicedCli) printPoolsFirst(ctx *cli.Context) {
	if len(ctx.Args()) > 0 {
		return
	}
	fmt.Println(strings.Join(c.pools(), "\n"))
}

// Bash-completion command that prints the list of available pools as all
// arguments
func (c *ServicedCli) printPoolsAll(ctx *cli.Context) {
	args := ctx.Args()
	pools := c.pools()

	for _, p := range pools {
		for _, a := range args {
			if p == a {
				goto next
			}
		}
		fmt.Println(p)
	next:
	}
}

// serviced pool list [POOLID]
func (c *ServicedCli) cmdPoolList(ctx *cli.Context) {
	if len(ctx.Args()) > 0 {
		poolID := ctx.Args()[0]
		if pool, err := c.driver.GetResourcePool(poolID); err != nil {
			fmt.Fprintln(os.Stderr, err)
		} else if pool == nil {
			fmt.Fprintln(os.Stderr, "pool not found")
		} else if jsonPool, err := json.MarshalIndent(pool, " ", "  "); err != nil {
			fmt.Fprintf(os.Stderr, "failed to marshal resource pool: %s", err)
		} else {
			fmt.Println(string(jsonPool))
		}
		return
	}

	pools, err := c.driver.GetResourcePools()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return
	} else if pools == nil || len(pools) == 0 {
		fmt.Fprintln(os.Stderr, "no resource pools found")
		return
	}

	if ctx.Bool("verbose") {
		if jsonPool, err := json.MarshalIndent(pools, " ", "  "); err != nil {
			fmt.Fprintf(os.Stderr, "failed to marshal resource pool list: %s", err)
		} else {
			fmt.Println(string(jsonPool))
		}
	} else {
		tablePool := newtable(0, 8, 2)
		tablePool.printrow("ID", "PARENT" /*"CORE", "MEM",*/, "PRI")
		for _, p := range pools {
			tablePool.printrow(p.ID, p.ParentID /*p.CoreLimit, p.MemoryLimit,*/, p.Priority)
		}
		tablePool.flush()
	}
}

// serviced pool add POOLID PRIORITY
func (c *ServicedCli) cmdPoolAdd(ctx *cli.Context) {
	args := ctx.Args()
	if len(args) < 2 {
		fmt.Printf("Incorrect Usage.\n\n")
		cli.ShowCommandHelp(ctx, "add")
		return
	}

	var err error

	cfg := api.PoolConfig{}
	cfg.PoolID = args[0]

	/* Disabled until enforced. See ZEN-11450
	cfg.CoreLimit, err = strconv.Atoi(args[1])
	if err != nil {
		fmt.Println("CORE_LIMIT must be a number")
		return
	}

	cfg.MemoryLimit, err = strconv.ParseUint(args[2], 10, 64)
	if err != nil {
		fmt.Println("MEMORY_LIMIT must be a number")
		return
	}
	*/

	cfg.Priority, err = strconv.Atoi(args[1])
	if err != nil {
		fmt.Println("PRIORITY must be a number")
		return
	}

	/* TODO: 1.1
	if len(args) > 2 {
		cfg.Realm = args[2]
	}
	*/

	if pool, err := c.driver.AddResourcePool(cfg); err != nil {
		fmt.Fprintln(os.Stderr, err)
	} else if pool == nil {
		fmt.Fprintln(os.Stderr, "received nil resource pool")
	} else {
		fmt.Println(pool.ID)
	}
}

// serviced pool remove POOLID ...
func (c *ServicedCli) cmdPoolRemove(ctx *cli.Context) {
	args := ctx.Args()
	if len(args) < 1 {
		fmt.Printf("Incorrect Usage.\n\n")
		cli.ShowCommandHelp(ctx, "remove")
	}

	for _, id := range args {
		if p, err := c.driver.GetResourcePool(id); err != nil {
			fmt.Fprintf(os.Stderr, "%s: %s\n", id, err)
		} else if p == nil {
			fmt.Fprintf(os.Stderr, "%s: pool not found", id)
		} else if err := c.driver.RemoveResourcePool(id); err != nil {
			fmt.Fprintf(os.Stderr, "%s: %s\n", id, err)
		} else {
			fmt.Println(id)
		}
	}
}

// serviced pool list-ips POOLID
func (c *ServicedCli) cmdPoolListIPs(ctx *cli.Context) {
	args := ctx.Args()
	if len(args) < 1 {
		fmt.Printf("Incorrect Usage.\n\n")
		cli.ShowCommandHelp(ctx, "list-ips")
		return
	}

	if poolIps, err := c.driver.GetPoolIPs(args[0]); err != nil {
		fmt.Fprintln(os.Stderr, err)
		return
	} else if poolIps.HostIPs == nil || (len(poolIps.HostIPs) == 0 && len(poolIps.VirtualIPs) == 0) {
		fmt.Fprintln(os.Stderr, "no resource pool IPs found")
		return
	} else if ctx.Bool("verbose") {
		if jsonPoolIP, err := json.MarshalIndent(poolIps.HostIPs, " ", "  "); err != nil {
			fmt.Fprintf(os.Stderr, "failed to marshal resource pool IPs: %s", err)
		} else {
			fmt.Println(string(jsonPoolIP))
		}
	} else {
		tableIPs := newtable(0, 10, 2)
		tableIPs.printrow("Interface Name", "IP Address", "Type")
		for _, ip := range poolIps.HostIPs {
			tableIPs.printrow(ip.InterfaceName, ip.IPAddress, "static")
		}
		for _, ip := range poolIps.VirtualIPs {
			tableIPs.printrow("", ip.IP, "virtual")
		}
		tableIPs.flush()
	}
}

// serviced pool add-virtual-ip POOLID IPADDRESS NETMASK BINDINTERFACE
func (c *ServicedCli) cmdAddVirtualIP(ctx *cli.Context) {
	args := ctx.Args()
	if len(args) != 4 {
		fmt.Printf("Incorrect Usage.\n\n")
		cli.ShowCommandHelp(ctx, "add-virtual-ip")
		return
	}

	requestVirtualIP := pool.VirtualIP{PoolID: args[0], IP: args[1], Netmask: args[2], BindInterface: args[3]}
	if err := c.driver.AddVirtualIP(requestVirtualIP); err != nil {
		fmt.Fprintln(os.Stderr, err)
		return
	}

	fmt.Println("Added virtual IP:", args[1])
}

// serviced pool remove-virtual-ip POOLID IPADDRESS
func (c *ServicedCli) cmdRemoveVirtualIP(ctx *cli.Context) {
	args := ctx.Args()
	if len(args) < 2 {
		fmt.Printf("Incorrect Usage.\n\n")
		cli.ShowCommandHelp(ctx, "remove-virtual-ip")
		return
	}

	requestVirtualIP := pool.VirtualIP{PoolID: args[0], IP: args[1], Netmask: "", BindInterface: ""}
	if err := c.driver.RemoveVirtualIP(requestVirtualIP); err != nil {
		fmt.Fprintln(os.Stderr, err)
		return
	} else {
		fmt.Printf("Removed virtual IP: %v from pool %v\n", args[1], args[0])
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"time"

	"github.com/johanneswuerbach/nfsexports"
)

type result struct {
	value string
	err   error
}

func GetIP(uuid string) (string, error) {
	var ip string
	timeout := make(chan bool, 1)
	value := make(chan result, 1)

	go func() {
		time.Sleep(15 * time.Second)
		timeout <- true
	}()

	go func() {
		for ip == "" {
			time.Sleep(time.Second)
			matchRe := regexp.MustCompile(`.*name=([a-fA-F0-9\-]+)\n.*ip_address=([0-9\.]+)`)

			file, err := os.Open("/var/db/dhcpd_leases")
			if err != nil {
				if os.IsNotExist(err) {
					continue
				}
				value <- result{"", err}
				return
			}

			defer file.Close()
			leases, err := ioutil.ReadAll(file)
			if err != nil {
				value <- result{"", err}
				return
			}

			lines := string(leases)
			matches := matchRe.FindAllStringSubmatch(lines, -1)
			for _, match := range matches {
				if match[1] == uuid {
					value <- result{match[2], nil}
					break
				}
			}
		}
	}()

	select {
	case res := <-value:
		return res.value, res.err
	case <-timeout:
		return "", fmt.Errorf("Failed to find an IP for the virtual machine in 15 seconds")
	}
}

func AddHost(hostname, ip string) error {
	if hostname == "" {
		hostname = "local.docker"
	}

	ipRe := regexp.MustCompile(`.*# added by dlite$`)
	ipLine := ip + " " + hostname + " # added by dlite"

	file, err := os.OpenFile("/etc/hosts", os.O_RDWR, 0644)
	if err != nil {
		return err
	}

	defer file.Close()

	hosts, err := ioutil.ReadAll(file)
	if err != nil {
		return err
	}

	lines := strings.Split(string(hosts), "\n")
	added := false
	for i, line := range lines {
		if ipRe.MatchString(line) {
			lines[i] = ipLine
			added = true
			break
		}
	}

	if !added {
		lines = append(lines, ipLine)
	}

	_, err = file.Seek(0, 0)
	if err != nil {
		return err
	}

	_, err = file.Write([]byte(strings.Join(lines, "\n") + "\n"))
	return err
}

func RemoveHost() error {
	ipRe := regexp.MustCompile(`.*# added by dlite$`)

	file, err := os.OpenFile("/etc/hosts", os.O_RDWR, 0644)
	if err != nil {
		return err
	}

	defer file.Close()

	hosts, err := ioutil.ReadAll(file)
	if err != nil {
		return err
	}

	lines := strings.Split(string(hosts), "\n")
	for i, line := range lines {
		if ipRe.MatchString(line) {
			lines = append(lines[:i], lines[i+1:]...)
			break
		}
	}

	_, err = file.Seek(0, 0)
	if err != nil {
		return err
	}

	n, err := file.Write([]byte(strings.Join(lines, "\n")))
	if err != nil {
		return err
	}

	return file.Truncate(int64(n))
}

func AddExport(uuid, share string) error {
	if share == "" {
		share = "/Users"
	}

	export := fmt.Sprintf("%s -network 192.168.64.0 -mask 255.255.255.0 -alldirs -maproot=root:wheel", share)
	_, err := nfsexports.Add("", "dlite", export)
	if err != nil {
		return err
	}

	err = nfsexports.ReloadDaemon()
	if err != nil {
		return exec.Command("sudo", "nfsd", "start").Run()
	}

	return err
}

func RemoveExport() error {
	_, err := nfsexports.Remove("", "dlite")
	if err != nil {
		return err
	}

	err = nfsexports.ReloadDaemon()
	if err != nil {
		return exec.Command("sudo", "nfsd", "start").Run()
	}

	return err
}
<commit_msg>cleanup addhost and removehost, remove addexport and removeexport<commit_after>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"regexp"
	"time"
)

type result struct {
	value string
	err   error
}

func GetIP(uuid string) (string, error) {
	var ip string
	timeout := make(chan bool, 1)
	value := make(chan result, 1)

	go func() {
		time.Sleep(15 * time.Second)
		timeout <- true
	}()

	go func() {
		for ip == "" {
			time.Sleep(time.Second)
			matchRe := regexp.MustCompile(`.*name=([a-fA-F0-9\-]+)\n.*ip_address=([0-9\.]+)`)

			file, err := os.Open("/var/db/dhcpd_leases")
			if err != nil {
				if os.IsNotExist(err) {
					continue
				}
				value <- result{"", err}
				return
			}

			defer file.Close()
			leases, err := ioutil.ReadAll(file)
			if err != nil {
				value <- result{"", err}
				return
			}

			lines := string(leases)
			matches := matchRe.FindAllStringSubmatch(lines, -1)
			for _, match := range matches {
				if match[1] == uuid {
					value <- result{match[2], nil}
					break
				}
			}
		}
	}()

	select {
	case res := <-value:
		return res.value, res.err
	case <-timeout:
		return "", fmt.Errorf("Failed to find an IP for the virtual machine in 15 seconds")
	}
}

func AddHost(hostname, ip string) error {
	path := "/etc/hosts"
	hosts, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}

	startMarker := []byte("# begin dlite")
	endMarker := []byte("# end dlite\n")

	begin := bytes.Index(hosts, startMarker)
	end := bytes.Index(hosts, endMarker)

	var temp []byte
	if begin > -1 && end > -1 {
		temp = append(hosts[:begin], hosts[end + len(endMarker):]...)
		temp = append(bytes.TrimSpace(temp), '\n')
	} else {
		temp = hosts
	}

	if len(temp) > 0 && !bytes.HasSuffix(temp, []byte("\n")) {
		temp = append(temp, []byte("\n")...)
	}

	entry := fmt.Sprintf("# begin dlite\n%s %s\n# end dlite\n", ip, hostname)
	temp = append(temp, []byte(entry)...)
	return ioutil.WriteFile(path, temp, 0644)
}

func RemoveHost() error {
	path := "/etc/hosts"
	hosts, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}

	startMarker := []byte("# begin dlite")
	endMarker := []byte("# end dlite\n")

	begin := bytes.Index(hosts, startMarker)
	end := bytes.Index(hosts, endMarker)

	if begin == -1 && end == -1 {
		return nil
	}

	temp := append(hosts[:begin], hosts[end + len(endMarker):]...)
	temp = append(bytes.TrimSpace(temp), '\n')
	if len(temp) > 0 && !bytes.HasSuffix(temp, []byte("\n")) {
		temp = append(temp, []byte("\n")...)
	}
	return ioutil.WriteFile(path, temp, 0644)
}
<|endoftext|>