<commit_before>package models

import (
	"errors"
	"socialapi/request"
	"time"

	"github.com/koding/bongo"
)

type MessageReply struct {
	// unique identifier of the MessageReply
	Id int64 `json:"id,string"`

	// Id of the interacted message
	MessageId int64 `json:"messageId,string"     sql:"NOT NULL"`

	// Id of the reply
	ReplyId int64 `json:"replyId,string"         sql:"NOT NULL"`

	// holds troll, unsafe, etc
	MetaBits MetaBits `json:"-"`

	// Creation of the MessageReply
	CreatedAt time.Time `json:"createdAt"         sql:"NOT NULL"`
}

func (m MessageReply) GetId() int64 {
	return m.Id
}

func (m MessageReply) TableName() string {
	return "api.message_reply"
}

func NewMessageReply() *MessageReply {
	return &MessageReply{}
}

func (m *MessageReply) BeforeCreate() {
	if m.ReplyId == 0 {
		return
	}

	cm := NewChannelMessage()
	cm.Id = m.ReplyId
	if res, err := cm.isExemptContent(); err == nil && res {
		m.MetaBits = updateTrollModeBit(m.MetaBits)
	}
}

func (m *MessageReply) BeforeUpdate() {
	if m.ReplyId == 0 {
		return
	}

	cm := NewChannelMessage()
	cm.Id = m.ReplyId
	if res, err := cm.isExemptContent(); err == nil && res {
		m.MetaBits = updateTrollModeBit(m.MetaBits)
	}
}

func (m *MessageReply) AfterCreate() {
	bongo.B.AfterCreate(m)
}

func (m *MessageReply) AfterUpdate() {
	bongo.B.AfterUpdate(m)
}

func (m MessageReply) AfterDelete() {
	bongo.B.AfterDelete(m)
}

func (m *MessageReply) ById(id int64) error {
	return bongo.B.ById(m, id)
}

func (m *MessageReply) Create() error {
	return bongo.B.Create(m)
}

func (m *MessageReply) CreateRaw() error {
	insertSql := "INSERT INTO " +
		m.TableName() +
		` ("message_id","reply_id","created_at") VALUES ($1,$2,$3) ` +
		"RETURNING ID"
	return bongo.B.DB.CommonDB().
		QueryRow(insertSql, m.MessageId, m.ReplyId, m.CreatedAt).
		Scan(&m.Id)
}

func (m *MessageReply) Some(data interface{}, q *bongo.Query) error {
	return bongo.B.Some(m, data, q)
}

func (m *MessageReply) One(q *bongo.Query) error {
	return bongo.B.One(m, m, q)
}

func (m *MessageReply) Delete() error {
	selector := map[string]interface{}{
		"message_id": m.MessageId,
		"reply_id":   m.ReplyId,
	}

	if err := m.One(bongo.NewQS(selector)); err != nil {
		return err
	}

	err := bongo.B.Delete(m)
	if err != nil {
		return err
	}

	return nil
}

func (m *MessageReply) DeleteByOrQuery(messageId int64) error {
	var messageReplies []MessageReply
	query := bongo.B.DB.Table(m.TableName())
	query = query.Where("message_id = ? or reply_id = ?", messageId, messageId)

	if err := query.Find(&messageReplies).Error; err != bongo.RecordNotFound && err != nil {
		return err
	}

	if messageReplies == nil {
		return nil
	}

	if len(messageReplies) == 0 {
		return nil
	}

	for _, messageReply := range messageReplies {
		err := bongo.B.Delete(messageReply)
		if err != nil {
			return err
		}
	}

	return nil
}

func (m *MessageReply) List(query *request.Query) ([]ChannelMessage, error) {
	return m.fetchMessages(query)
}

func (m *MessageReply) ListAll() ([]ChannelMessage, error) {
	query := request.NewQuery()
	query.Limit = 0
	query.Skip = 0
	return m.fetchMessages(query)
}

func (m *MessageReply) fetchMessages(query *request.Query) ([]ChannelMessage, error) {
	var replies []int64

	if m.MessageId == 0 {
		return nil, errors.New("messageId is not set")
	}

	q := &bongo.Query{
		Selector: map[string]interface{}{
			"message_id": m.MessageId,
		},
		Pluck:      "reply_id",
		Pagination: *bongo.NewPagination(query.Limit, query.Skip),
		Sort:       map[string]string{"created_at": "DESC"},
	}

	q.AddScope(RemoveTrollContent(m, query.ShowExempt))

	bongoQuery := bongo.B.BuildQuery(m, q)
	if !query.From.IsZero() {
		bongoQuery = bongoQuery.Where("created_at < ?", query.From)
	}

	if err := bongo.CheckErr(
		bongoQuery.Pluck(q.Pluck, &replies),
	); err != nil {
		return nil, err
	}

	parent := NewChannelMessage()
	channelMessageReplies, err := parent.FetchByIds(replies)
	if err != nil {
		return nil, err
	}

	return channelMessageReplies, nil
}

func (m *MessageReply) UnreadCount(messageId int64, addedAt time.Time) (int, error) {
	if messageId == 0 {
		return 0, errors.New("messageId is not set")
	}

	if addedAt.IsZero() {
		return 0, errors.New("last seen at date is not valid - it is zero")
	}

	return bongo.B.Count(
		m,
		"message_id = ? and created_at > ?",
		messageId,
		addedAt.UTC().Format(time.RFC3339),
	)
}

func (m *MessageReply) Count(q *request.Query) (int, error) {
	if m.MessageId == 0 {
		return 0, errors.New("messageId is not set")
	}

	query := &bongo.Query{
		Selector: map[string]interface{}{
			"message_id": m.MessageId,
		},
	}

	query.AddScope(RemoveTrollContent(
		m, q.ShowExempt,
	))

	return m.CountWithQuery(query)
}

func (m *MessageReply) CountWithQuery(q *bongo.Query) (int, error) {
	return bongo.B.CountWithQuery(m, q)
}

func (m *MessageReply) FetchParent() (*ChannelMessage, error) {
	parent := NewChannelMessage()

	if m.ReplyId != 0 {
		if err := parent.ById(m.MessageId); err != nil {
			return nil, err
		}

		return parent, nil
	}

	if m.ReplyId == 0 {
		return nil, errors.New("replyId is not set")
	}

	q := &bongo.Query{
		Selector: map[string]interface{}{
			"reply_id": m.ReplyId,
		},
	}

	if err := m.One(q); err != nil {
		return nil, err
	}

	if err := parent.ById(m.MessageId); err != nil {
		return nil, err
	}

	return parent, nil
}

func (m *MessageReply) FetchReply() (*ChannelMessage, error) {
	reply := NewChannelMessage()

	replyId, err := m.getReplyId()
	if err != nil {
		return nil, err
	}

	if err := reply.ById(replyId); err != nil {
		return nil, err
	}

	return reply, nil
}

func (m *MessageReply) getReplyId() (int64, error) {
	if m.Id == 0 && m.ReplyId == 0 {
		return 0, errors.New("required ids are not set")
	}

	if m.ReplyId != 0 {
		return m.ReplyId, nil
	}

	if m.Id == 0 {
		// shouldnt come here
		return 0, errors.New("couldnt fetch replyId")
	}

	if err := m.ById(m.Id); err != nil {
		return 0, err
	}

	return m.ReplyId, nil
}
<commit_msg>social: add update function into message_reply<commit_after>package models

import (
	"errors"
	"socialapi/request"
	"time"

	"github.com/koding/bongo"
)

type MessageReply struct {
	// unique identifier of the MessageReply
	Id int64 `json:"id,string"`

	// Id of the interacted message
	MessageId int64 `json:"messageId,string"     sql:"NOT NULL"`

	// Id of the reply
	ReplyId int64 `json:"replyId,string"         sql:"NOT NULL"`

	// holds troll, unsafe, etc
	MetaBits MetaBits `json:"-"`

	// Creation of the MessageReply
	CreatedAt time.Time `json:"createdAt"         sql:"NOT NULL"`
}

func (m MessageReply) GetId() int64 {
	return m.Id
}

func (m MessageReply) TableName() string {
	return "api.message_reply"
}

func NewMessageReply() *MessageReply {
	return &MessageReply{}
}

func (m *MessageReply) BeforeCreate() {
	if m.ReplyId == 0 {
		return
	}

	cm := NewChannelMessage()
	cm.Id = m.ReplyId
	if res, err := cm.isExemptContent(); err == nil && res {
		m.MetaBits = updateTrollModeBit(m.MetaBits)
	}
}

func (m *MessageReply) BeforeUpdate() {
	if m.ReplyId == 0 {
		return
	}

	cm := NewChannelMessage()
	cm.Id = m.ReplyId
	if res, err := cm.isExemptContent(); err == nil && res {
		m.MetaBits = updateTrollModeBit(m.MetaBits)
	}
}

func (m *MessageReply) AfterCreate() {
	bongo.B.AfterCreate(m)
}

func (m *MessageReply) AfterUpdate() {
	bongo.B.AfterUpdate(m)
}

func (m MessageReply) AfterDelete() {
	bongo.B.AfterDelete(m)
}

func (m *MessageReply) ById(id int64) error {
	return bongo.B.ById(m, id)
}

func (m *MessageReply) Create() error {
	return bongo.B.Create(m)
}

func (m *MessageReply) Update() error {
	return bongo.B.Update(m)
}

func (m *MessageReply) CreateRaw() error {
	insertSql := "INSERT INTO " +
		m.TableName() +
		` ("message_id","reply_id","created_at") VALUES ($1,$2,$3) ` +
		"RETURNING ID"
	return bongo.B.DB.CommonDB().
		QueryRow(insertSql, m.MessageId, m.ReplyId, m.CreatedAt).
		Scan(&m.Id)
}

func (m *MessageReply) Some(data interface{}, q *bongo.Query) error {
	return bongo.B.Some(m, data, q)
}

func (m *MessageReply) One(q *bongo.Query) error {
	return bongo.B.One(m, m, q)
}

func (m *MessageReply) Delete() error {
	selector := map[string]interface{}{
		"message_id": m.MessageId,
		"reply_id":   m.ReplyId,
	}

	if err := m.One(bongo.NewQS(selector)); err != nil {
		return err
	}

	err := bongo.B.Delete(m)
	if err != nil {
		return err
	}

	return nil
}

func (m *MessageReply) DeleteByOrQuery(messageId int64) error {
	var messageReplies []MessageReply
	query := bongo.B.DB.Table(m.TableName())
	query = query.Where("message_id = ? or reply_id = ?", messageId, messageId)

	if err := query.Find(&messageReplies).Error; err != bongo.RecordNotFound && err != nil {
		return err
	}

	if messageReplies == nil {
		return nil
	}

	if len(messageReplies) == 0 {
		return nil
	}

	for _, messageReply := range messageReplies {
		err := bongo.B.Delete(messageReply)
		if err != nil {
			return err
		}
	}

	return nil
}

func (m *MessageReply) List(query *request.Query) ([]ChannelMessage, error) {
	return m.fetchMessages(query)
}

func (m *MessageReply) ListAll() ([]ChannelMessage, error) {
	query := request.NewQuery()
	query.Limit = 0
	query.Skip = 0
	return m.fetchMessages(query)
}

func (m *MessageReply) fetchMessages(query *request.Query) ([]ChannelMessage, error) {
	var replies []int64

	if m.MessageId == 0 {
		return nil, errors.New("messageId is not set")
	}

	q := &bongo.Query{
		Selector: map[string]interface{}{
			"message_id": m.MessageId,
		},
		Pluck:      "reply_id",
		Pagination: *bongo.NewPagination(query.Limit, query.Skip),
		Sort:       map[string]string{"created_at": "DESC"},
	}

	q.AddScope(RemoveTrollContent(m, query.ShowExempt))

	bongoQuery := bongo.B.BuildQuery(m, q)
	if !query.From.IsZero() {
		bongoQuery = bongoQuery.Where("created_at < ?", query.From)
	}

	if err := bongo.CheckErr(
		bongoQuery.Pluck(q.Pluck, &replies),
	); err != nil {
		return nil, err
	}

	parent := NewChannelMessage()
	channelMessageReplies, err := parent.FetchByIds(replies)
	if err != nil {
		return nil, err
	}

	return channelMessageReplies, nil
}

func (m *MessageReply) UnreadCount(messageId int64, addedAt time.Time) (int, error) {
	if messageId == 0 {
		return 0, errors.New("messageId is not set")
	}

	if addedAt.IsZero() {
		return 0, errors.New("last seen at date is not valid - it is zero")
	}

	return bongo.B.Count(
		m,
		"message_id = ? and created_at > ?",
		messageId,
		addedAt.UTC().Format(time.RFC3339),
	)
}

func (m *MessageReply) Count(q *request.Query) (int, error) {
	if m.MessageId == 0 {
		return 0, errors.New("messageId is not set")
	}

	query := &bongo.Query{
		Selector: map[string]interface{}{
			"message_id": m.MessageId,
		},
	}

	query.AddScope(RemoveTrollContent(
		m, q.ShowExempt,
	))

	return m.CountWithQuery(query)
}

func (m *MessageReply) CountWithQuery(q *bongo.Query) (int, error) {
	return bongo.B.CountWithQuery(m, q)
}

func (m *MessageReply) FetchParent() (*ChannelMessage, error) {
	parent := NewChannelMessage()

	if m.ReplyId != 0 {
		if err := parent.ById(m.MessageId); err != nil {
			return nil, err
		}

		return parent, nil
	}

	if m.ReplyId == 0 {
		return nil, errors.New("replyId is not set")
	}

	q := &bongo.Query{
		Selector: map[string]interface{}{
			"reply_id": m.ReplyId,
		},
	}

	if err := m.One(q); err != nil {
		return nil, err
	}

	if err := parent.ById(m.MessageId); err != nil {
		return nil, err
	}

	return parent, nil
}

func (m *MessageReply) FetchReply() (*ChannelMessage, error) {
	reply := NewChannelMessage()

	replyId, err := m.getReplyId()
	if err != nil {
		return nil, err
	}

	if err := reply.ById(replyId); err != nil {
		return nil, err
	}

	return reply, nil
}

func (m *MessageReply) getReplyId() (int64, error) {
	if m.Id == 0 && m.ReplyId == 0 {
		return 0, errors.New("required ids are not set")
	}

	if m.ReplyId != 0 {
		return m.ReplyId, nil
	}

	if m.Id == 0 {
		// shouldnt come here
		return 0, errors.New("couldnt fetch replyId")
	}

	if err := m.ById(m.Id); err != nil {
		return 0, err
	}

	return m.ReplyId, nil
}
<|endoftext|>


<commit_before>package models

import (
	// "errors"
	"fmt"
	"github.com/jinzhu/gorm"
	"github.com/koding/bongo"
	"time"
)

var (
	FETCH_LIMIT = 50
)

type Notification struct {
	Id                    int64     `json:"id"`
	AccountId             int64     `json:"accountId"             sql:"NOT NULL"`
	NotificationContentId int64     `json:"notificationContentId" sql:"NOT NULL"`
	Glanced               bool      `json:"glanced"               sql:"NOT NULL"`
	UpdatedAt             time.Time `json:"updatedAt"             sql:"NOT NULL"`
}

func (n *Notification) GetId() int64 {
	return n.Id
}

func (n *Notification) TableName() string {
	return "notification"
}

func NewNotification() *Notification {
	return &Notification{}
}

func (n *Notification) One(q *bongo.Query) error {
	return bongo.B.One(n, n, q)
}

func (n *Notification) Create() error {
	s := map[string]interface{}{
		"account_id":              n.AccountId,
		"notification_content_id": n.NotificationContentId,
	}
	q := bongo.NewQS(s)
	if err := n.One(q); err != nil {
		if err != gorm.RecordNotFound {
			return err
		}

		return bongo.B.Create(n)
	}

	n.Glanced = false

	return bongo.B.Update(n)
}

func (n *Notification) List(q *Query) ([]NotificationContainer, error) {
	q.Limit = 8

	result, err := n.getDecoratedList(q)
	if err != nil {
		return nil, err
	}

	result, err = populateActors(result)
	if err != nil {
		return nil, err
	}

	return result, nil
}

func (n *Notification) Some(data interface{}, q *bongo.Query) error {

	return bongo.B.Some(n, data, q)
}

func (n *Notification) fetchByAccountId(q *Query) ([]Notification, error) {
	var notifications []Notification
	query := &bongo.Query{
		Selector: map[string]interface{}{
			"account_id": q.AccountId,
		},
		Sort: map[string]string{
			"updated_at": "desc",
		},
		Pagination: bongo.Pagination{
			Limit: FETCH_LIMIT,
			Skip:  q.Skip,
		},
	}
	if err := bongo.B.Some(n, &notifications, query); err != nil {
		return nil, err
	}

	return notifications, nil
}

// prepareNotifications
func (n *Notification) getDecoratedList(q *Query) ([]NotificationContainer, error) {
	result := make([]NotificationContainer, 0)
	resultMap := make(map[string]struct{}, 0)

	var err error
	result, err = n.decorateContents(result, &resultMap, q)
	if err != nil {
		return nil, err
	}

	return result, nil
}

func (n *Notification) decorateContents(result []NotificationContainer, resultMap *map[string]struct{}, q *Query) ([]NotificationContainer, error) {

	nList, err := n.fetchByAccountId(q)
	if nList == nil {
		return result, nil
	}

	// fetch all notification content relationships
	ncMap, err := fetchRelatedContent(nList)
	if err != nil {
		return nil, err
	}

	for _, n := range nList {
		nc := ncMap[n.NotificationContentId]
		key := prepareResultKey(&nc)
		if _, ok := (*resultMap)[key]; !ok {
			container := buildNotificationContainer(&nc)
			(*resultMap)[key] = struct{}{}
			result = append(result, container)
			if len(result) == q.Limit {
				return result, nil
			}
		}
	}

	if len(nList) == FETCH_LIMIT {
		q.Skip += FETCH_LIMIT
		return n.decorateContents(result, resultMap, q)
	}

	return result, nil
}

func buildNotificationContainer(nc *NotificationContent) NotificationContainer {
	return NotificationContainer{
		TargetId: nc.TargetId,
		Type:     nc.Type,
	}
}

func prepareResultKey(nc *NotificationContent) string {
	return fmt.Sprintf("%s_%d", nc.Type, nc.TargetId)
}

func fetchRelatedContent(nl []Notification) (map[int64]NotificationContent, error) {
	notificationContentIds := make([]int64, 0)
	for _, n := range nl {
		notificationContentIds = append(notificationContentIds, n.NotificationContentId)
	}
	nc := NewNotificationContent()
	return nc.FetchMapByIds(notificationContentIds)
}

// fetch 3 actors and total count of actors. actors are from different tables
func populateActors(ncList []NotificationContainer) ([]NotificationContainer, error) {
	result := make([]NotificationContainer, 0)

	for _, n := range ncList {
		notificationType, err := CreateNotificationType(n.Type)
		if err != nil {
			return nil, err
		}

		notificationType.SetTargetId(n.TargetId)

		actors, err := notificationType.FetchActors()
		// instead of interrupting process we can just proceed here
		if err != nil {
			return nil, err
		}

		n.Actors = *actors
		result = append(result, n)
	}
	return result, nil
}
<commit_msg>SocialApi: item limit is added for notification list<commit_after>package models

import (
	// "errors"
	"fmt"
	"github.com/jinzhu/gorm"
	"github.com/koding/bongo"
	"math"
	"time"
)

var (
	FETCH_LIMIT = 50
)

type Notification struct {
	Id                    int64     `json:"id"`
	AccountId             int64     `json:"accountId"             sql:"NOT NULL"`
	NotificationContentId int64     `json:"notificationContentId" sql:"NOT NULL"`
	Glanced               bool      `json:"glanced"               sql:"NOT NULL"`
	UpdatedAt             time.Time `json:"updatedAt"             sql:"NOT NULL"`
}

func (n *Notification) GetId() int64 {
	return n.Id
}

func (n *Notification) TableName() string {
	return "notification"
}

func NewNotification() *Notification {
	return &Notification{}
}

func (n *Notification) One(q *bongo.Query) error {
	return bongo.B.One(n, n, q)
}

func (n *Notification) Create() error {
	s := map[string]interface{}{
		"account_id":              n.AccountId,
		"notification_content_id": n.NotificationContentId,
	}
	q := bongo.NewQS(s)
	if err := n.One(q); err != nil {
		if err != gorm.RecordNotFound {
			return err
		}

		return bongo.B.Create(n)
	}

	n.Glanced = false

	return bongo.B.Update(n)
}

func (n *Notification) List(q *Query) ([]NotificationContainer, error) {
	q.Limit = 8
	limit := math.Min(float64(q.Limit), 8.0)
	q.Limit = int(limit)

	result, err := n.getDecoratedList(q)
	if err != nil {
		return nil, err
	}

	result, err = populateActors(result)
	if err != nil {
		return nil, err
	}

	return result, nil
}

func (n *Notification) Some(data interface{}, q *bongo.Query) error {

	return bongo.B.Some(n, data, q)
}

func (n *Notification) fetchByAccountId(q *Query) ([]Notification, error) {
	var notifications []Notification
	query := &bongo.Query{
		Selector: map[string]interface{}{
			"account_id": q.AccountId,
		},
		Sort: map[string]string{
			"updated_at": "desc",
		},
		Pagination: bongo.Pagination{
			Limit: FETCH_LIMIT,
			Skip:  q.Skip,
		},
	}
	if err := bongo.B.Some(n, &notifications, query); err != nil {
		return nil, err
	}

	return notifications, nil
}

// prepareNotifications
func (n *Notification) getDecoratedList(q *Query) ([]NotificationContainer, error) {
	result := make([]NotificationContainer, 0)
	resultMap := make(map[string]struct{}, 0)

	var err error
	result, err = n.decorateContents(result, &resultMap, q)
	if err != nil {
		return nil, err
	}

	return result, nil
}

func (n *Notification) decorateContents(result []NotificationContainer, resultMap *map[string]struct{}, q *Query) ([]NotificationContainer, error) {

	nList, err := n.fetchByAccountId(q)
	if nList == nil {
		return result, nil
	}

	// fetch all notification content relationships
	ncMap, err := fetchRelatedContent(nList)
	if err != nil {
		return nil, err
	}

	for _, n := range nList {
		nc := ncMap[n.NotificationContentId]
		key := prepareResultKey(&nc)
		if _, ok := (*resultMap)[key]; !ok {
			container := buildNotificationContainer(&nc)
			(*resultMap)[key] = struct{}{}
			result = append(result, container)
			if len(result) == q.Limit {
				return result, nil
			}
		}
	}

	if len(nList) == FETCH_LIMIT {
		q.Skip += FETCH_LIMIT
		return n.decorateContents(result, resultMap, q)
	}

	return result, nil
}

func buildNotificationContainer(nc *NotificationContent) NotificationContainer {
	return NotificationContainer{
		TargetId: nc.TargetId,
		Type:     nc.Type,
	}
}

func prepareResultKey(nc *NotificationContent) string {
	return fmt.Sprintf("%s_%d", nc.Type, nc.TargetId)
}

func fetchRelatedContent(nl []Notification) (map[int64]NotificationContent, error) {
	notificationContentIds := make([]int64, 0)
	for _, n := range nl {
		notificationContentIds = append(notificationContentIds, n.NotificationContentId)
	}
	nc := NewNotificationContent()
	return nc.FetchMapByIds(notificationContentIds)
}

// fetch 3 actors and total count of actors. actors are from different tables
func populateActors(ncList []NotificationContainer) ([]NotificationContainer, error) {
	result := make([]NotificationContainer, 0)

	for _, n := range ncList {
		notificationType, err := CreateNotificationType(n.Type)
		if err != nil {
			return nil, err
		}

		notificationType.SetTargetId(n.TargetId)

		actors, err := notificationType.FetchActors()
		// instead of interrupting process we can just proceed here
		if err != nil {
			return nil, err
		}

		n.Actors = *actors
		result = append(result, n)
	}
	return result, nil
}
<|endoftext|>


<commit_before>// aoe provides the implementation of an ATA over Ethernet server, backed by a Torus block volume.
package aoe

import (
	"fmt"
	"net"
	"os"
	"sync"
	"syscall"
	"time"

	"github.com/coreos/torus/block"

	"github.com/coreos/pkg/capnslog"
	"github.com/mdlayher/aoe"
	"github.com/mdlayher/raw"
	"golang.org/x/net/bpf"
	"golang.org/x/net/context"
)

var (
	clog          = capnslog.NewPackageLogger("github.com/coreos/torus", "aoe")
	broadcastAddr = net.HardwareAddr([]byte{0xff, 0xff, 0xff, 0xff, 0xff, 0xff})
)

type Server struct {
	dfs *block.BlockVolume

	dev Device

	ctx    context.Context
	cancel context.CancelFunc
	wg     *sync.WaitGroup

	major uint16
	minor uint8

	usingBPF bool
}

// ServerOptions specifies options for a Server.
type ServerOptions struct {
	// Major and Minor specify the major and minor address of an AoE server.
	// Typically, all AoE devices on a single server will share the same
	// major address, but have different minor addresses.
	//
	// It is important to note that all AoE servers on the same layer 2
	// network must have different major and minor addresses.
	Major uint16
	Minor uint8
}

// DefaultServerOptions is the default ServerOptions configuration used
// by NewServer when none is specified.
var DefaultServerOptions = &ServerOptions{
	Major: 1,
	Minor: 1,
}

// NewServer creates a new Server which utilizes the specified block volume.
// If options is nil, DefaultServerOptions will be used.
func NewServer(b *block.BlockVolume, options *ServerOptions) (*Server, error) {
	if options == nil {
		options = DefaultServerOptions
	}

	f, err := b.OpenBlockFile()
	if err != nil {
		return nil, err
	}

	if err := f.Sync(); err != nil {
		return nil, err
	}

	ctx, cancel := context.WithCancel(context.Background())
	wg := &sync.WaitGroup{}

	fd := &FileDevice{f}

	as := &Server{
		dfs:    b,
		dev:    fd,
		ctx:    ctx,
		cancel: cancel,
		wg:     wg,
		major:  options.Major,
		minor:  options.Minor,
	}

	return as, nil
}

func (s *Server) advertise(iface *Interface) error {
	// little hack to trigger a broadcast
	from := &raw.Addr{
		HardwareAddr: broadcastAddr,
	}

	fr := &Frame{
		Header: aoe.Header{
			// Must specify our server's major/minor address or else
			// the request will be filtered
			Major:   s.major,
			Minor:   s.minor,
			Command: aoe.CommandQueryConfigInformation,
			Arg: &aoe.ConfigArg{
				Command: aoe.ConfigCommandRead,
			},
		},
	}

	_, err := s.handleFrame(from, iface, fr)
	return err
}

// A bpfPacketConn is a net.PacketConn which can have a BPF program attached
// to it for kernel-level filtering.
type bpfPacketConn interface {
	net.PacketConn
	SetBPF(filter []bpf.RawInstruction) error
}

// raw.Conn is a bpfPacketConn.
var _ bpfPacketConn = &raw.Conn{}

func (s *Server) Serve(iface *Interface) error {
	// If available, attach a BPF filter to net.PacketConn.
	if bp, ok := iface.PacketConn.(bpfPacketConn); ok {
		clog.Debugf("attaching BPF program to %T", bp)
		if err := bp.SetBPF(s.mustAssembleBPF(iface.MTU)); err != nil {
			// If user does not have permission to attach a BPF filter to an
			// interface, continue without one
			if !os.IsPermission(err) {
				return err
			}

			clog.Debugf("permission denied, continuing with BPF filter disabled")
		} else {
			clog.Debugf("BPF filter enabled")
			s.usingBPF = true
		}
	}

	clog.Tracef("beginning server loop on %+v", iface)

	// Start goroutine to sync device at regular intervals, and halt when
	// the Server's Close method is called.
	s.wg.Add(1)
	go func() {
		for {
			if err := s.dev.Sync(); err != nil {
				clog.Warningf("failed to sync %s: %v", s.dev, err)
			}

			select {
			case <-time.After(5 * time.Second):
			case <-s.ctx.Done():
				clog.Debugf("stopping device sync goroutine")
				s.wg.Done()
				return
			}
		}
	}()

	// broadcast ourselves
	if err := s.advertise(iface); err != nil {
		clog.Errorf("advertisement failed: %v", err)
		return err
	}

	for {
		payload := make([]byte, iface.MTU)
		n, addr, err := iface.ReadFrom(payload)
		if err != nil {
			clog.Errorf("ReadFrom failed: %v", err)
			// will be syscall.EBADF if the conn from raw closed
			if err == syscall.EBADF {
				break
			}
			return err
		}

		// resize payload
		payload = payload[:n]

		var f Frame
		if err := f.UnmarshalBinary(payload); err != nil {
			clog.Errorf("Failed to unmarshal frame: %v", err)
			continue
		}

		clog.Debugf("recv %d %s %+v", n, addr, f.Header)
		//clog.Debugf("recv arg %+v", f.Header.Arg)

		s.handleFrame(addr, iface, &f)
	}

	return nil
}

func (s *Server) handleFrame(from net.Addr, iface *Interface, f *Frame) (int, error) {
	hdr := &f.Header

	// If a BPF filter cannot be used, filter requests not bound for
	// this server in userspace.
	if !s.usingBPF {
		// Ignore client requests that are not being broadcast or sent to
		// our major/minor address combination.
		if hdr.Major != aoe.BroadcastMajor && hdr.Major != s.major {
			clog.Debugf("ignoring header with major address %d", hdr.Major)
			return 0, nil
		}
		if hdr.Minor != aoe.BroadcastMinor && hdr.Minor != s.minor {
			clog.Debugf("ignoring header with minor address %d", hdr.Minor)
			return 0, nil
		}
	}

	sender := &FrameSender{
		orig:  f,
		dst:   from.(*raw.Addr).HardwareAddr,
		src:   iface.HardwareAddr,
		conn:  iface.PacketConn,
		major: s.major,
		minor: s.minor,
	}

	switch hdr.Command {
	case aoe.CommandIssueATACommand:
		n, err := aoe.ServeATA(sender, hdr, s.dev)
		if err != nil {
			clog.Errorf("ServeATA failed: %v", err)
			switch err {
			case aoe.ErrInvalidATARequest:
				return sender.SendError(aoe.ErrorBadArgumentParameter)
			default:
				return sender.SendError(aoe.ErrorDeviceUnavailable)
			}
		}

		return n, nil
	case aoe.CommandQueryConfigInformation:
		cfgarg := f.Header.Arg.(*aoe.ConfigArg)
		clog.Debugf("cfgarg: %+v", cfgarg)

		switch cfgarg.Command {
		case aoe.ConfigCommandRead:
			hdr.Arg = defaultConfig(iface.MTU)

			return sender.Send(hdr)
		}

		return sender.SendError(aoe.ErrorUnrecognizedCommandCode)
	case aoe.CommandMACMaskList:
		fallthrough
	case aoe.CommandReserveRelease:
		fallthrough
	default:
		return sender.SendError(aoe.ErrorUnrecognizedCommandCode)
	}
}

func (s *Server) Close() error {
	clog.Debugf("canceling background goroutines")
	s.cancel()
	clog.Debugf("waiting for background goroutines to stop")
	s.wg.Wait()
	return s.dev.Close()
}

// defaultConfig is the default AoE server configuration.
func defaultConfig(mtu int) *aoe.ConfigArg {
	return &aoe.ConfigArg{
		// if < 2, linux aoe handles it poorly
		BufferCount:     2,
		FirmwareVersion: 0,
		// naive, but works.
		SectorCount:  uint8(mtu / 512),
		Version:      aoe.Version,
		Command:      aoe.ConfigCommandRead,
		StringLength: 0,
		String:       []byte{},
	}
}

// mustAssembleBPF assembles a BPF program to filter out packets not bound
// for this server.
func (s *Server) mustAssembleBPF(mtu int) []bpf.RawInstruction {
	// This BPF program filters out packets that are not bound for this server
	// by checking against both the AoE broadcast addresses and this server's
	// major/minor address combination.  The structure of the incoming ethernet
	// frame and AoE header is as follows:
	//
	// Offset | Length | Comment
	// -------------------------
	//   00   |   06   | Ethernet destination MAC address
	//   06   |   06   | Ethernet source MAC address
	//   12   |   02   | Ethernet EtherType
	// -------------------------
	//   14   |   01   | AoE version + flags
	//   15   |   01   | AoE error
	//   16   |   02   | AoE major address
	//   18   |   01   | AoE minor address
	//
	// Thus, our BPF program needs to check for:
	//  - major address: offset 16, length 2
	//  - minor address: offset 18, length 1
	const (
		majorOffset = 16
		majorLen    = 2
		minorOffset = 18
		minorLen    = 1
	)

	// TODO(mdlayher): this routine likely belongs in package AoE, once the server
	// component is more complete.

	prog, err := bpf.Assemble([]bpf.Instruction{
		// Load major address value from AoE header
		bpf.LoadAbsolute{
			Off:  majorOffset,
			Size: majorLen,
		},
		// If major address is equal to broadcast address, jump to minor address
		// filtering
		bpf.JumpIf{
			Cond:     bpf.JumpEqual,
			Val:      uint32(aoe.BroadcastMajor),
			SkipTrue: 2,
		},
		// If major address is equal to our server's, jump to minor address
		// filtering
		bpf.JumpIf{
			Cond:     bpf.JumpEqual,
			Val:      uint32(s.major),
			SkipTrue: 1,
		},
		// Major address is not our server's or broadcast address
		bpf.RetConstant{
			Val: 0,
		},
		// Load minor address value from AoE header
		bpf.LoadAbsolute{
			Off:  minorOffset,
			Size: minorLen,
		},
		// If minor address is equal to broadcast address, jump to accept packet
		bpf.JumpIf{
			Cond:     bpf.JumpEqual,
			Val:      uint32(aoe.BroadcastMinor),
			SkipTrue: 2,
		},
		// If minor address is equal to our server's, jump to accept packet
		bpf.JumpIf{
			Cond:     bpf.JumpEqual,
			Val:      uint32(s.minor),
			SkipTrue: 1,
		},
		// Minor address is not our server's or broadcast address
		bpf.RetConstant{
			Val: 0,
		},
		// Accept the packet bytes up to the interface's MTU
		bpf.RetConstant{
			Val: uint32(mtu),
		},
	})
	if err != nil {
		panic(fmt.Sprintf("failed to assemble BPF program: %v", err))
	}

	return prog
}
<commit_msg>aoe: tiny: output error when failed to handle frame<commit_after>// aoe provides the implementation of an ATA over Ethernet server, backed by a Torus block volume.
package aoe

import (
	"fmt"
	"net"
	"os"
	"sync"
	"syscall"
	"time"

	"github.com/coreos/torus/block"

	"github.com/coreos/pkg/capnslog"
	"github.com/mdlayher/aoe"
	"github.com/mdlayher/raw"
	"golang.org/x/net/bpf"
	"golang.org/x/net/context"
)

var (
	clog          = capnslog.NewPackageLogger("github.com/coreos/torus", "aoe")
	broadcastAddr = net.HardwareAddr([]byte{0xff, 0xff, 0xff, 0xff, 0xff, 0xff})
)

type Server struct {
	dfs *block.BlockVolume

	dev Device

	ctx    context.Context
	cancel context.CancelFunc
	wg     *sync.WaitGroup

	major uint16
	minor uint8

	usingBPF bool
}

// ServerOptions specifies options for a Server.
type ServerOptions struct {
	// Major and Minor specify the major and minor address of an AoE server.
	// Typically, all AoE devices on a single server will share the same
	// major address, but have different minor addresses.
	//
	// It is important to note that all AoE servers on the same layer 2
	// network must have different major and minor addresses.
	Major uint16
	Minor uint8
}

// DefaultServerOptions is the default ServerOptions configuration used
// by NewServer when none is specified.
var DefaultServerOptions = &ServerOptions{
	Major: 1,
	Minor: 1,
}

// NewServer creates a new Server which utilizes the specified block volume.
// If options is nil, DefaultServerOptions will be used.
func NewServer(b *block.BlockVolume, options *ServerOptions) (*Server, error) {
	if options == nil {
		options = DefaultServerOptions
	}

	f, err := b.OpenBlockFile()
	if err != nil {
		return nil, err
	}

	if err := f.Sync(); err != nil {
		return nil, err
	}

	ctx, cancel := context.WithCancel(context.Background())
	wg := &sync.WaitGroup{}

	fd := &FileDevice{f}

	as := &Server{
		dfs:    b,
		dev:    fd,
		ctx:    ctx,
		cancel: cancel,
		wg:     wg,
		major:  options.Major,
		minor:  options.Minor,
	}

	return as, nil
}

func (s *Server) advertise(iface *Interface) error {
	// little hack to trigger a broadcast
	from := &raw.Addr{
		HardwareAddr: broadcastAddr,
	}

	fr := &Frame{
		Header: aoe.Header{
			// Must specify our server's major/minor address or else
			// the request will be filtered
			Major:   s.major,
			Minor:   s.minor,
			Command: aoe.CommandQueryConfigInformation,
			Arg: &aoe.ConfigArg{
				Command: aoe.ConfigCommandRead,
			},
		},
	}

	_, err := s.handleFrame(from, iface, fr)
	return err
}

// A bpfPacketConn is a net.PacketConn which can have a BPF program attached
// to it for kernel-level filtering.
type bpfPacketConn interface {
	net.PacketConn
	SetBPF(filter []bpf.RawInstruction) error
}

// raw.Conn is a bpfPacketConn.
var _ bpfPacketConn = &raw.Conn{}

func (s *Server) Serve(iface *Interface) error {
	// If available, attach a BPF filter to net.PacketConn.
	if bp, ok := iface.PacketConn.(bpfPacketConn); ok {
		clog.Debugf("attaching BPF program to %T", bp)
		if err := bp.SetBPF(s.mustAssembleBPF(iface.MTU)); err != nil {
			// If user does not have permission to attach a BPF filter to an
			// interface, continue without one
			if !os.IsPermission(err) {
				return err
			}

			clog.Debugf("permission denied, continuing with BPF filter disabled")
		} else {
			clog.Debugf("BPF filter enabled")
			s.usingBPF = true
		}
	}

	clog.Tracef("beginning server loop on %+v", iface)

	// Start goroutine to sync device at regular intervals, and halt when
	// the Server's Close method is called.
	s.wg.Add(1)
	go func() {
		for {
			if err := s.dev.Sync(); err != nil {
				clog.Warningf("failed to sync %s: %v", s.dev, err)
			}

			select {
			case <-time.After(5 * time.Second):
			case <-s.ctx.Done():
				clog.Debugf("stopping device sync goroutine")
				s.wg.Done()
				return
			}
		}
	}()

	// broadcast ourselves
	if err := s.advertise(iface); err != nil {
		clog.Errorf("advertisement failed: %v", err)
		return err
	}

	for {
		payload := make([]byte, iface.MTU)
		n, addr, err := iface.ReadFrom(payload)
		if err != nil {
			clog.Errorf("ReadFrom failed: %v", err)
			// will be syscall.EBADF if the conn from raw closed
			if err == syscall.EBADF {
				break
			}
			return err
		}

		// resize payload
		payload = payload[:n]

		var f Frame
		if err := f.UnmarshalBinary(payload); err != nil {
			clog.Errorf("Failed to unmarshal frame: %v", err)
			continue
		}

		clog.Debugf("recv %d %s %+v", n, addr, f.Header)
		//clog.Debugf("recv arg %+v", f.Header.Arg)

		_, err = s.handleFrame(addr, iface, &f)
		if err != nil {
			clog.Errorf("failed to handle frame: %v", err)
		}
	}

	return nil
}

func (s *Server) handleFrame(from net.Addr, iface *Interface, f *Frame) (int, error) {
	hdr := &f.Header

	// If a BPF filter cannot be used, filter requests not bound for
	// this server in userspace.
	if !s.usingBPF {
		// Ignore client requests that are not being broadcast or sent to
		// our major/minor address combination.
		if hdr.Major != aoe.BroadcastMajor && hdr.Major != s.major {
			clog.Debugf("ignoring header with major address %d", hdr.Major)
			return 0, nil
		}
		if hdr.Minor != aoe.BroadcastMinor && hdr.Minor != s.minor {
			clog.Debugf("ignoring header with minor address %d", hdr.Minor)
			return 0, nil
		}
	}

	sender := &FrameSender{
		orig:  f,
		dst:   from.(*raw.Addr).HardwareAddr,
		src:   iface.HardwareAddr,
		conn:  iface.PacketConn,
		major: s.major,
		minor: s.minor,
	}

	switch hdr.Command {
	case aoe.CommandIssueATACommand:
		n, err := aoe.ServeATA(sender, hdr, s.dev)
		if err != nil {
			clog.Errorf("ServeATA failed: %v", err)
			switch err {
			case aoe.ErrInvalidATARequest:
				return sender.SendError(aoe.ErrorBadArgumentParameter)
			default:
				return sender.SendError(aoe.ErrorDeviceUnavailable)
			}
		}

		return n, nil
	case aoe.CommandQueryConfigInformation:
		cfgarg := f.Header.Arg.(*aoe.ConfigArg)
		clog.Debugf("cfgarg: %+v", cfgarg)

		switch cfgarg.Command {
		case aoe.ConfigCommandRead:
			hdr.Arg = defaultConfig(iface.MTU)

			return sender.Send(hdr)
		}

		return sender.SendError(aoe.ErrorUnrecognizedCommandCode)
	case aoe.CommandMACMaskList:
		fallthrough
	case aoe.CommandReserveRelease:
		fallthrough
	default:
		return sender.SendError(aoe.ErrorUnrecognizedCommandCode)
	}
}

func (s *Server) Close() error {
	clog.Debugf("canceling background goroutines")
	s.cancel()
	clog.Debugf("waiting for background goroutines to stop")
	s.wg.Wait()
	return s.dev.Close()
}

// defaultConfig is the default AoE server configuration.
func defaultConfig(mtu int) *aoe.ConfigArg {
	return &aoe.ConfigArg{
		// if < 2, linux aoe handles it poorly
		BufferCount:     2,
		FirmwareVersion: 0,
		// naive, but works.
		SectorCount:  uint8(mtu / 512),
		Version:      aoe.Version,
		Command:      aoe.ConfigCommandRead,
		StringLength: 0,
		String:       []byte{},
	}
}

// mustAssembleBPF assembles a BPF program to filter out packets not bound
// for this server.
func (s *Server) mustAssembleBPF(mtu int) []bpf.RawInstruction {
	// This BPF program filters out packets that are not bound for this server
	// by checking against both the AoE broadcast addresses and this server's
	// major/minor address combination.  The structure of the incoming ethernet
	// frame and AoE header is as follows:
	//
	// Offset | Length | Comment
	// -------------------------
	//   00   |   06   | Ethernet destination MAC address
	//   06   |   06   | Ethernet source MAC address
	//   12   |   02   | Ethernet EtherType
	// -------------------------
	//   14   |   01   | AoE version + flags
	//   15   |   01   | AoE error
	//   16   |   02   | AoE major address
	//   18   |   01   | AoE minor address
	//
	// Thus, our BPF program needs to check for:
	//  - major address: offset 16, length 2
	//  - minor address: offset 18, length 1
	const (
		majorOffset = 16
		majorLen    = 2
		minorOffset = 18
		minorLen    = 1
	)

	// TODO(mdlayher): this routine likely belongs in package AoE, once the server
	// component is more complete.

	prog, err := bpf.Assemble([]bpf.Instruction{
		// Load major address value from AoE header
		bpf.LoadAbsolute{
			Off:  majorOffset,
			Size: majorLen,
		},
		// If major address is equal to broadcast address, jump to minor address
		// filtering
		bpf.JumpIf{
			Cond:     bpf.JumpEqual,
			Val:      uint32(aoe.BroadcastMajor),
			SkipTrue: 2,
		},
		// If major address is equal to our server's, jump to minor address
		// filtering
		bpf.JumpIf{
			Cond:     bpf.JumpEqual,
			Val:      uint32(s.major),
			SkipTrue: 1,
		},
		// Major address is not our server's or broadcast address
		bpf.RetConstant{
			Val: 0,
		},
		// Load minor address value from AoE header
		bpf.LoadAbsolute{
			Off:  minorOffset,
			Size: minorLen,
		},
		// If minor address is equal to broadcast address, jump to accept packet
		bpf.JumpIf{
			Cond:     bpf.JumpEqual,
			Val:      uint32(aoe.BroadcastMinor),
			SkipTrue: 2,
		},
		// If minor address is equal to our server's, jump to accept packet
		bpf.JumpIf{
			Cond:     bpf.JumpEqual,
			Val:      uint32(s.minor),
			SkipTrue: 1,
		},
		// Minor address is not our server's or broadcast address
		bpf.RetConstant{
			Val: 0,
		},
		// Accept the packet bytes up to the interface's MTU
		bpf.RetConstant{
			Val: uint32(mtu),
		},
	})
	if err != nil {
		panic(fmt.Sprintf("failed to assemble BPF program: %v", err))
	}

	return prog
}
<|endoftext|>


<commit_before>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package binlogplayertest

import (
	"fmt"
	"reflect"
	"strings"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/vt/binlog/binlogplayer"
	"github.com/youtube/vitess/go/vt/key"

	binlogdatapb "github.com/youtube/vitess/go/vt/proto/binlogdata"
	querypb "github.com/youtube/vitess/go/vt/proto/query"
	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
)

// keyRangeRequest is used to make a request for StreamKeyRange.
type keyRangeRequest struct {
	Position string
	KeyRange *topodatapb.KeyRange
	Charset  *binlogdatapb.Charset
}

// tablesRequest is used to make a request for StreamTables.
type tablesRequest struct {
	Position string
	Tables   []string
	Charset  *binlogdatapb.Charset
}

// FakeBinlogStreamer is our implementation of UpdateStream
type FakeBinlogStreamer struct {
	t      *testing.T
	panics bool
}

// NewFakeBinlogStreamer returns the test instance for UpdateStream
func NewFakeBinlogStreamer(t *testing.T) *FakeBinlogStreamer {
	return &FakeBinlogStreamer{
		t:      t,
		panics: false,
	}
}

//
// ServeUpdateStream tests
//

var testUpdateStreamRequest = "UpdateStream starting position"

var testStreamEvent = &binlogdatapb.StreamEvent{
	Category:  binlogdatapb.StreamEvent_SE_DML,
	TableName: "table1",
	PrimaryKeyFields: []*querypb.Field{
		{
			Name: "id",
			Type: sqltypes.Binary,
		},
	},
	PrimaryKeyValues: []*querypb.Row{
		{
			Lengths: []int64{3},
			Values:  []byte{'1', '2', '3'},
		},
	},
	Sql:           "test sql",
	Timestamp:     372,
	TransactionId: "StreamEvent returned transaction id",
}

// ServeUpdateStream is part of the the UpdateStream interface
func (fake *FakeBinlogStreamer) ServeUpdateStream(position string, sendReply func(reply *binlogdatapb.StreamEvent) error) error {
	if fake.panics {
		panic(fmt.Errorf("test-triggered panic"))
	}
	if position != testUpdateStreamRequest {
		fake.t.Errorf("wrong ServeUpdateStream parameter, got %v want %v", position, testUpdateStreamRequest)
	}
	sendReply(testStreamEvent)
	return nil
}

func testServeUpdateStream(t *testing.T, bpc binlogplayer.Client) {
	ctx := context.Background()
	c, errFunc, err := bpc.ServeUpdateStream(ctx, testUpdateStreamRequest)
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
	if se, ok := <-c; !ok {
		t.Fatalf("got no response")
	} else {
		if !reflect.DeepEqual(*se, *testStreamEvent) {
			t.Errorf("got wrong result, got \n%#v expected \n%#v", *se, *testStreamEvent)
		}
	}
	if se, ok := <-c; ok {
		t.Fatalf("got a response when error expected: %v", se)
	}
	if err := errFunc(); err != nil {
		t.Errorf("got unexpected error: %v", err)
	}
}

func testServeUpdateStreamPanics(t *testing.T, bpc binlogplayer.Client) {
	ctx := context.Background()
	c, errFunc, err := bpc.ServeUpdateStream(ctx, testUpdateStreamRequest)
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
	if se, ok := <-c; ok {
		t.Fatalf("got a response when error expected: %v", se)
	}
	err = errFunc()
	if err == nil || !strings.Contains(err.Error(), "test-triggered panic") {
		t.Errorf("wrong error from panic: %v", err)
	}
}

//
// StreamKeyRange tests
//

var testKeyRangeRequest = &keyRangeRequest{
	Position: "KeyRange starting position",
	KeyRange: &topodatapb.KeyRange{
		Start: key.Uint64Key(0x7000000000000000).Bytes(),
		End:   key.Uint64Key(0x9000000000000000).Bytes(),
	},
	Charset: &binlogdatapb.Charset{
		Client: 12,
		Conn:   13,
		Server: 14,
	},
}

var testBinlogTransaction = &binlogdatapb.BinlogTransaction{
	Statements: []*binlogdatapb.BinlogTransaction_Statement{
		{
			Category: binlogdatapb.BinlogTransaction_Statement_BL_ROLLBACK,
			Charset: &binlogdatapb.Charset{
				Client: 120,
				Conn:   130,
				Server: 140,
			},
			Sql: "my statement",
		},
	},
	Timestamp:     78,
	TransactionId: "BinlogTransaction returned transaction id",
}

// StreamKeyRange is part of the the UpdateStream interface
func (fake *FakeBinlogStreamer) StreamKeyRange(position string, keyRange *topodatapb.KeyRange, charset *binlogdatapb.Charset, sendReply func(reply *binlogdatapb.BinlogTransaction) error) error {
	if fake.panics {
		panic(fmt.Errorf("test-triggered panic"))
	}
	req := &keyRangeRequest{
		Position: position,
		KeyRange: keyRange,
		Charset:  charset,
	}
	if !reflect.DeepEqual(req, testKeyRangeRequest) {
		fake.t.Errorf("wrong StreamKeyRange parameter, got %+v want %+v", req, testKeyRangeRequest)
	}
	sendReply(testBinlogTransaction)
	return nil
}

func testStreamKeyRange(t *testing.T, bpc binlogplayer.Client) {
	ctx := context.Background()
	c, errFunc, err := bpc.StreamKeyRange(ctx, testKeyRangeRequest.Position, testKeyRangeRequest.KeyRange, testKeyRangeRequest.Charset)
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
	if se, ok := <-c; !ok {
		t.Fatalf("got no response")
	} else {
		if !reflect.DeepEqual(*se, *testBinlogTransaction) {
			t.Errorf("got wrong result, got %v expected %v", *se, *testBinlogTransaction)
		}
	}
	if se, ok := <-c; ok {
		t.Fatalf("got a response when error expected: %v", se)
	}
	if err := errFunc(); err != nil {
		t.Errorf("got unexpected error: %v", err)
	}
}

func testStreamKeyRangePanics(t *testing.T, bpc binlogplayer.Client) {
	ctx := context.Background()
	c, errFunc, err := bpc.StreamKeyRange(ctx, testKeyRangeRequest.Position, testKeyRangeRequest.KeyRange, testKeyRangeRequest.Charset)
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
	if se, ok := <-c; ok {
		t.Fatalf("got a response when error expected: %v", se)
	}
	err = errFunc()
	if err == nil || !strings.Contains(err.Error(), "test-triggered panic") {
		t.Errorf("wrong error from panic: %v", err)
	}
}

//
// StreamTables test
//

var testTablesRequest = &tablesRequest{
	Position: "Tables starting position",
	Tables:   []string{"table1", "table2"},
	Charset: &binlogdatapb.Charset{
		Client: 12,
		Conn:   13,
		Server: 14,
	},
}

// StreamTables is part of the the UpdateStream interface
func (fake *FakeBinlogStreamer) StreamTables(position string, tables []string, charset *binlogdatapb.Charset, sendReply func(reply *binlogdatapb.BinlogTransaction) error) error {
	if fake.panics {
		panic(fmt.Errorf("test-triggered panic"))
	}
	req := &tablesRequest{
		Position: position,
		Tables:   tables,
		Charset:  charset,
	}
	if !reflect.DeepEqual(req, testTablesRequest) {
		fake.t.Errorf("wrong StreamTables parameter, got %+v want %+v", req, testTablesRequest)
	}
	sendReply(testBinlogTransaction)
	return nil
}

func testStreamTables(t *testing.T, bpc binlogplayer.Client) {
	ctx := context.Background()
	c, errFunc, err := bpc.StreamTables(ctx, testTablesRequest.Position, testTablesRequest.Tables, testTablesRequest.Charset)
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
	if se, ok := <-c; !ok {
		t.Fatalf("got no response")
	} else {
		if !reflect.DeepEqual(*se, *testBinlogTransaction) {
			t.Errorf("got wrong result, got %v expected %v", *se, *testBinlogTransaction)
		}
	}
	if se, ok := <-c; ok {
		t.Fatalf("got a response when error expected: %v", se)
	}
	if err := errFunc(); err != nil {
		t.Errorf("got unexpected error: %v", err)
	}
}

func testStreamTablesPanics(t *testing.T, bpc binlogplayer.Client) {
	ctx := context.Background()
	c, errFunc, err := bpc.StreamTables(ctx, testTablesRequest.Position, testTablesRequest.Tables, testTablesRequest.Charset)
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
	if se, ok := <-c; ok {
		t.Fatalf("got a response when error expected: %v", se)
	}
	err = errFunc()
	if err == nil || !strings.Contains(err.Error(), "test-triggered panic") {
		t.Errorf("wrong error from panic: %v", err)
	}
}

// HandlePanic is part of the the UpdateStream interface
func (fake *FakeBinlogStreamer) HandlePanic(err *error) {
	if x := recover(); x != nil {
		*err = fmt.Errorf("Caught panic: %v", x)
	}
}

// Run runs the test suite
func Run(t *testing.T, bpc binlogplayer.Client, endPoint *topodatapb.EndPoint, fake *FakeBinlogStreamer) {
	if err := bpc.Dial(endPoint, 30*time.Second); err != nil {
		t.Fatalf("Dial failed: %v", err)
	}

	// no panic
	testServeUpdateStream(t, bpc)
	testStreamKeyRange(t, bpc)
	testStreamTables(t, bpc)

	// panic now, and test
	fake.panics = true
	testServeUpdateStreamPanics(t, bpc)
	testStreamKeyRangePanics(t, bpc)
	testStreamTablesPanics(t, bpc)
}
<commit_msg>Resetting panic variable in this test so it can run multiple times.<commit_after>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package binlogplayertest

import (
	"fmt"
	"reflect"
	"strings"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/vt/binlog/binlogplayer"
	"github.com/youtube/vitess/go/vt/key"

	binlogdatapb "github.com/youtube/vitess/go/vt/proto/binlogdata"
	querypb "github.com/youtube/vitess/go/vt/proto/query"
	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
)

// keyRangeRequest is used to make a request for StreamKeyRange.
type keyRangeRequest struct {
	Position string
	KeyRange *topodatapb.KeyRange
	Charset  *binlogdatapb.Charset
}

// tablesRequest is used to make a request for StreamTables.
type tablesRequest struct {
	Position string
	Tables   []string
	Charset  *binlogdatapb.Charset
}

// FakeBinlogStreamer is our implementation of UpdateStream
type FakeBinlogStreamer struct {
	t      *testing.T
	panics bool
}

// NewFakeBinlogStreamer returns the test instance for UpdateStream
func NewFakeBinlogStreamer(t *testing.T) *FakeBinlogStreamer {
	return &FakeBinlogStreamer{
		t:      t,
		panics: false,
	}
}

//
// ServeUpdateStream tests
//

var testUpdateStreamRequest = "UpdateStream starting position"

var testStreamEvent = &binlogdatapb.StreamEvent{
	Category:  binlogdatapb.StreamEvent_SE_DML,
	TableName: "table1",
	PrimaryKeyFields: []*querypb.Field{
		{
			Name: "id",
			Type: sqltypes.Binary,
		},
	},
	PrimaryKeyValues: []*querypb.Row{
		{
			Lengths: []int64{3},
			Values:  []byte{'1', '2', '3'},
		},
	},
	Sql:           "test sql",
	Timestamp:     372,
	TransactionId: "StreamEvent returned transaction id",
}

// ServeUpdateStream is part of the the UpdateStream interface
func (fake *FakeBinlogStreamer) ServeUpdateStream(position string, sendReply func(reply *binlogdatapb.StreamEvent) error) error {
	if fake.panics {
		panic(fmt.Errorf("test-triggered panic"))
	}
	if position != testUpdateStreamRequest {
		fake.t.Errorf("wrong ServeUpdateStream parameter, got %v want %v", position, testUpdateStreamRequest)
	}
	sendReply(testStreamEvent)
	return nil
}

func testServeUpdateStream(t *testing.T, bpc binlogplayer.Client) {
	ctx := context.Background()
	c, errFunc, err := bpc.ServeUpdateStream(ctx, testUpdateStreamRequest)
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
	if se, ok := <-c; !ok {
		t.Fatalf("got no response")
	} else {
		if !reflect.DeepEqual(*se, *testStreamEvent) {
			t.Errorf("got wrong result, got \n%#v expected \n%#v", *se, *testStreamEvent)
		}
	}
	if se, ok := <-c; ok {
		t.Fatalf("got a response when error expected: %v", se)
	}
	if err := errFunc(); err != nil {
		t.Errorf("got unexpected error: %v", err)
	}
}

func testServeUpdateStreamPanics(t *testing.T, bpc binlogplayer.Client) {
	ctx := context.Background()
	c, errFunc, err := bpc.ServeUpdateStream(ctx, testUpdateStreamRequest)
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
	if se, ok := <-c; ok {
		t.Fatalf("got a response when error expected: %v", se)
	}
	err = errFunc()
	if err == nil || !strings.Contains(err.Error(), "test-triggered panic") {
		t.Errorf("wrong error from panic: %v", err)
	}
}

//
// StreamKeyRange tests
//

var testKeyRangeRequest = &keyRangeRequest{
	Position: "KeyRange starting position",
	KeyRange: &topodatapb.KeyRange{
		Start: key.Uint64Key(0x7000000000000000).Bytes(),
		End:   key.Uint64Key(0x9000000000000000).Bytes(),
	},
	Charset: &binlogdatapb.Charset{
		Client: 12,
		Conn:   13,
		Server: 14,
	},
}

var testBinlogTransaction = &binlogdatapb.BinlogTransaction{
	Statements: []*binlogdatapb.BinlogTransaction_Statement{
		{
			Category: binlogdatapb.BinlogTransaction_Statement_BL_ROLLBACK,
			Charset: &binlogdatapb.Charset{
				Client: 120,
				Conn:   130,
				Server: 140,
			},
			Sql: "my statement",
		},
	},
	Timestamp:     78,
	TransactionId: "BinlogTransaction returned transaction id",
}

// StreamKeyRange is part of the the UpdateStream interface
func (fake *FakeBinlogStreamer) StreamKeyRange(position string, keyRange *topodatapb.KeyRange, charset *binlogdatapb.Charset, sendReply func(reply *binlogdatapb.BinlogTransaction) error) error {
	if fake.panics {
		panic(fmt.Errorf("test-triggered panic"))
	}
	req := &keyRangeRequest{
		Position: position,
		KeyRange: keyRange,
		Charset:  charset,
	}
	if !reflect.DeepEqual(req, testKeyRangeRequest) {
		fake.t.Errorf("wrong StreamKeyRange parameter, got %+v want %+v", req, testKeyRangeRequest)
	}
	sendReply(testBinlogTransaction)
	return nil
}

func testStreamKeyRange(t *testing.T, bpc binlogplayer.Client) {
	ctx := context.Background()
	c, errFunc, err := bpc.StreamKeyRange(ctx, testKeyRangeRequest.Position, testKeyRangeRequest.KeyRange, testKeyRangeRequest.Charset)
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
	if se, ok := <-c; !ok {
		t.Fatalf("got no response")
	} else {
		if !reflect.DeepEqual(*se, *testBinlogTransaction) {
			t.Errorf("got wrong result, got %v expected %v", *se, *testBinlogTransaction)
		}
	}
	if se, ok := <-c; ok {
		t.Fatalf("got a response when error expected: %v", se)
	}
	if err := errFunc(); err != nil {
		t.Errorf("got unexpected error: %v", err)
	}
}

func testStreamKeyRangePanics(t *testing.T, bpc binlogplayer.Client) {
	ctx := context.Background()
	c, errFunc, err := bpc.StreamKeyRange(ctx, testKeyRangeRequest.Position, testKeyRangeRequest.KeyRange, testKeyRangeRequest.Charset)
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
	if se, ok := <-c; ok {
		t.Fatalf("got a response when error expected: %v", se)
	}
	err = errFunc()
	if err == nil || !strings.Contains(err.Error(), "test-triggered panic") {
		t.Errorf("wrong error from panic: %v", err)
	}
}

//
// StreamTables test
//

var testTablesRequest = &tablesRequest{
	Position: "Tables starting position",
	Tables:   []string{"table1", "table2"},
	Charset: &binlogdatapb.Charset{
		Client: 12,
		Conn:   13,
		Server: 14,
	},
}

// StreamTables is part of the the UpdateStream interface
func (fake *FakeBinlogStreamer) StreamTables(position string, tables []string, charset *binlogdatapb.Charset, sendReply func(reply *binlogdatapb.BinlogTransaction) error) error {
	if fake.panics {
		panic(fmt.Errorf("test-triggered panic"))
	}
	req := &tablesRequest{
		Position: position,
		Tables:   tables,
		Charset:  charset,
	}
	if !reflect.DeepEqual(req, testTablesRequest) {
		fake.t.Errorf("wrong StreamTables parameter, got %+v want %+v", req, testTablesRequest)
	}
	sendReply(testBinlogTransaction)
	return nil
}

func testStreamTables(t *testing.T, bpc binlogplayer.Client) {
	ctx := context.Background()
	c, errFunc, err := bpc.StreamTables(ctx, testTablesRequest.Position, testTablesRequest.Tables, testTablesRequest.Charset)
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
	if se, ok := <-c; !ok {
		t.Fatalf("got no response")
	} else {
		if !reflect.DeepEqual(*se, *testBinlogTransaction) {
			t.Errorf("got wrong result, got %v expected %v", *se, *testBinlogTransaction)
		}
	}
	if se, ok := <-c; ok {
		t.Fatalf("got a response when error expected: %v", se)
	}
	if err := errFunc(); err != nil {
		t.Errorf("got unexpected error: %v", err)
	}
}

func testStreamTablesPanics(t *testing.T, bpc binlogplayer.Client) {
	ctx := context.Background()
	c, errFunc, err := bpc.StreamTables(ctx, testTablesRequest.Position, testTablesRequest.Tables, testTablesRequest.Charset)
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
	if se, ok := <-c; ok {
		t.Fatalf("got a response when error expected: %v", se)
	}
	err = errFunc()
	if err == nil || !strings.Contains(err.Error(), "test-triggered panic") {
		t.Errorf("wrong error from panic: %v", err)
	}
}

// HandlePanic is part of the the UpdateStream interface
func (fake *FakeBinlogStreamer) HandlePanic(err *error) {
	if x := recover(); x != nil {
		*err = fmt.Errorf("Caught panic: %v", x)
	}
}

// Run runs the test suite
func Run(t *testing.T, bpc binlogplayer.Client, endPoint *topodatapb.EndPoint, fake *FakeBinlogStreamer) {
	if err := bpc.Dial(endPoint, 30*time.Second); err != nil {
		t.Fatalf("Dial failed: %v", err)
	}

	// no panic
	testServeUpdateStream(t, bpc)
	testStreamKeyRange(t, bpc)
	testStreamTables(t, bpc)

	// panic now, and test
	fake.panics = true
	testServeUpdateStreamPanics(t, bpc)
	testStreamKeyRangePanics(t, bpc)
	testStreamTablesPanics(t, bpc)
	fake.panics = false
}
<|endoftext|>


<commit_before>package threatexchange

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
)

const (
	// DefaultURL is the URL for the API endpoint
	DefaultURL = "https://graph.facebook.com"
	apiVersion = "v2.5"
)

// Client interacts with the services provided by TE
type Client struct {
	appId     string
	appSecret string
	logger    *log.Logger
	c         *http.Client
}

type MalwareFamiliesResults struct {
	Data   []MalwareFamilyResult `json:"data"`
	Paging ResultPaging          `json:"paging,omitempty"`
	Next   string                `json:"next,omitempty"`
}

type MalwareFamilyResult struct {
	ID             string   `json:"id"`
	Aliases        []string `json:"aliases,omitempty"`
	AddedOn        string   `json:"added_on,omitempty"`
	Description    string   `json:"description,omitempty"`
	FamilyType     string   `json:"family_type,omitempty"`
	Malicious      string   `json:"malicious,omitempty"`
	Name           string   `json:"name,omitempty"`
	SampleCount    int      `json:"sample_count,omitempty"`
	SubmitterCount int      `json:"submitter_count,omitempty"`
}
type MalwareResults struct {
	Data   []MalwareResult `json:"data"`
	Paging ResultPaging    `json:"paging,omitempty"`
	Next   string          `json:"next,omitempty"`
}

type MalwareResult struct {
	ID             string `json:"id"`
	Crx            string `json:"crx,omitempty"`
	PEHash         string `json:"imphash,omitempty"`
	MD5            string `json:"md5,omitempty"`
	Password       string `json:"password,omitempty"`
	PEHeader       string `json:"pe_rich_header,omitempty"`
	Sample         string `json:"sample,omitempty"`
	SampleType     string `json:"sample_type,omitempty"`
	Sha1           string `json:"sha1,omitempty"`
	Sha256         string `json:"sha256,omitempty"`
	ShareLevel     string `json:"share_level,omitempty"`
	Ssdeep         string `json:"ssdeep,omitempty"`
	SubmitterCount string `json:"submitter_count,omitempty"`
	VictimCount    int    `json:"victim_count,omitempty"`
	xpi            string `json:"xpi,omitempty"`
	Status         string `json:"status,omitempty"`
	AddedOn        string `json:"added_on,omitempty"`
	PrivacyType    string `json:"privacy_type,omitempty"`
}

type ThreatDescriptorResults struct {
	Data   []ThreatDescriptor `json:"data"`
	Paging ResultPaging       `json:"paging,omitempty"`
	Next   string             `json:"next,omitempty"`
}

type ResultPaging struct {
	Cursors Cursors `json:"cursors"`
}

type Cursors struct {
	Before string `json:"before,omitempty"`
	After  string `json:"after,omitempty"`
}

type ThreatDescriptor struct {
	ID           string          `json:"id"`
	Indicator    IndicatorResult `json:"indicator,omitempty"`
	Owner        TEOwner         `json:"owner,omitempty"`
	Type         string          `json:"type,omitempty"`
	RawIndicator string          `json:"raw_indicator,omitempty"`
	Description  string          `json:"description,omitempty"`
	Status       string          `json:"status,omitempty"`
	PrivacyType  string          `json:"privacy_type,omitempty"`
	ShareLevel   string          `json:"share_level,omitempty"`
	AddedOn      string          `json:"added_on,omitempty"`
	Confidence   int             `json:"confidence,omitempty"`
	ExpiredOn    int             `json:"expired_on,omitempty"`
	LastUpdated  string          `json:"last_updated,omitempty"`
	SourceUri    string          `json:"source_uri,omitempty"`
}

type IndicatorResult struct {
	ID        string `json:"id,omitempty"`
	Indicator string `json:"indicator,omitempty"`
	Type      string `json:"type,omitempty"`
}

type TEOwner struct {
	ID    string `json:"id,omitempty"`
	Email string `json:"email,omitempty"`
	Name  string `json:"name,omitempty"`
}

type TEOwnersResults struct {
	Data   []TEOwner    `json:"data"`
	Paging ResultPaging `json:"paging,omitempty"`
	Next   string       `json:"next,omitempty"`
}

// New - will create a new ThreatExchange Go client, log param may be nil.
func New(appID, appSecret string, log *log.Logger) (*Client, error) {
	c := &Client{
		appId:     appID,
		appSecret: appSecret,
		logger:    log,
		c:         http.DefaultClient,
	}
	return c, nil
}

// GetThreatIndicators - is a query to retrieve ThreatIndicators, will return ThreatDescriptorResults to hold results, and raw json of that
// limit of size <=0 will be ignored
func (c *Client) GetThreatIndicators(resourceType, text, startTime, endTime string, limit int,
	extraParams map[string]string) (*ThreatDescriptorResults, string, error) {
	res, err := c.query(apiVersion, "threat_descriptors", startTime, endTime, resourceType, text, limit, extraParams)
	if err != nil {
		return nil, "", err
	}
	var result ThreatDescriptorResults
	err = json.Unmarshal([]byte(res), &result)
	if err != nil {
		return nil, "", err
	}
	return &result, res, nil
}

// GetMalwareAnalyses - is a query to retrieve Malware Analyses, will return MalwareResults to hold results, and raw json of that
func (c *Client) GetMalwareAnalyses(text, startTime, endTime string, limit int,
	extraParams map[string]string) (*MalwareResults, string, error) {
	res, err := c.query(apiVersion, "malware_analyses", startTime, endTime, "", text, limit, extraParams)
	if err != nil {
		return nil, "", err
	}
	var result MalwareResults
	err = json.Unmarshal([]byte(res), &result)
	if err != nil {
		return nil, "", err
	}
	return &result, res, nil
}

// GetMalwareFamilies - is a query to retrieve ThreatIndicators, will return MalwareFamiliesResults to hold results, and raw json of that
func (c *Client) GetMalwareFamilies(text, startTime, endTime string, limit int,
	extraParams map[string]string) (*MalwareFamiliesResults, string, error) {
	res, err := c.query(apiVersion, "malware_families", startTime, endTime, "", text, limit, extraParams)
	if err != nil {
		return nil, "", err
	}
	var result MalwareFamiliesResults
	err = json.Unmarshal([]byte(res), &result)
	if err != nil {
		return nil, "", err
	}
	return &result, res, nil
}

// GetThreatExchangeMembers - is a query to retrieve ThreatExchange members, will return TEOwnersResults to hold results, and raw json of that
func (c *Client) GetThreatExchangeMembers() (*TEOwnersResults, string, error) {
	res, err := c.query(apiVersion, "threat_exchange_members", "", "", "", "", 0, map[string]string{})
	if err != nil {
		return nil, "", err
	}
	var result TEOwnersResults
	err = json.Unmarshal([]byte(res), &result)
	if err != nil {
		return nil, "", err
	}
	return &result, res, nil
}

func (c *Client) query(apiVersion string, resource string, startTime string, endTime string, resourceType string,
	text string, limit int, extraParams map[string]string) (string, error) {
	u, err := url.Parse(DefaultURL)
	if err != nil {
		return "", err
	}

	parameters := url.Values{}
	parameters.Add("access_token", fmt.Sprintf("%s|%s", c.appId, c.appSecret))
	if len(startTime) > 0 {
		parameters.Add("since", startTime)
	}
	if len(endTime) > 0 {
		parameters.Add("until", endTime)
	}
	if len(resourceType) > 0 {
		parameters.Add("type", resourceType)
	}
	if len(resourceType) > 0 {
		parameters.Add("text", text)
	}
	if limit > 0 {
		parameters.Add("limit", fmt.Sprintf("%d", limit))
	}
	for k, v := range extraParams {
		parameters.Add(k, v)
	}
	u.RawQuery = parameters.Encode()
	u.Path += fmt.Sprintf("/%s/%s/", apiVersion, resource)

	res, err := http.Get(u.String())
	if err != nil {
		return "", err
	}
	if res.Body == nil {
		return "", fmt.Errorf("Empty body for query :", u.String())
	}
	if res.StatusCode != http.StatusOK {
		return "", fmt.Errorf("Wrong http return code, got : %d", res.StatusCode)
	}
	defer res.Body.Close()
	byteRes, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return "", err
	}
	return string(byteRes), nil
}
<commit_msg>improve api fix missing xpi param<commit_after>package threatexchange

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
)

const (
	// DefaultURL is the URL for the API endpoint
	DefaultURL = "https://graph.facebook.com"
	apiVersion = "v2.5"
)

// Client interacts with the services provided by TE
type Client struct {
	appId     string
	appSecret string
	logger    *log.Logger
	c         *http.Client
}

type MalwareFamiliesResults struct {
	Data   []MalwareFamilyResult `json:"data"`
	Paging ResultPaging          `json:"paging,omitempty"`
	Next   string                `json:"next,omitempty"`
}

type MalwareFamilyResult struct {
	ID             string   `json:"id"`
	Aliases        []string `json:"aliases,omitempty"`
	AddedOn        string   `json:"added_on,omitempty"`
	Description    string   `json:"description,omitempty"`
	FamilyType     string   `json:"family_type,omitempty"`
	Malicious      string   `json:"malicious,omitempty"`
	Name           string   `json:"name,omitempty"`
	SampleCount    int      `json:"sample_count,omitempty"`
	SubmitterCount int      `json:"submitter_count,omitempty"`
}
type MalwareResults struct {
	Data   []MalwareResult `json:"data"`
	Paging ResultPaging    `json:"paging,omitempty"`
	Next   string          `json:"next,omitempty"`
}

type MalwareResult struct {
	ID             string `json:"id"`
	Crx            string `json:"crx,omitempty"`
	PEHash         string `json:"imphash,omitempty"`
	MD5            string `json:"md5,omitempty"`
	Password       string `json:"password,omitempty"`
	PEHeader       string `json:"pe_rich_header,omitempty"`
	Sample         string `json:"sample,omitempty"`
	SampleType     string `json:"sample_type,omitempty"`
	Sha1           string `json:"sha1,omitempty"`
	Sha256         string `json:"sha256,omitempty"`
	ShareLevel     string `json:"share_level,omitempty"`
	Ssdeep         string `json:"ssdeep,omitempty"`
	SubmitterCount string `json:"submitter_count,omitempty"`
	VictimCount    int    `json:"victim_count,omitempty"`
	Xpi            string `json:"xpi,omitempty"`
	Status         string `json:"status,omitempty"`
	AddedOn        string `json:"added_on,omitempty"`
	PrivacyType    string `json:"privacy_type,omitempty"`
}

type ThreatDescriptorResults struct {
	Data   []ThreatDescriptor `json:"data"`
	Paging ResultPaging       `json:"paging,omitempty"`
	Next   string             `json:"next,omitempty"`
}

type ResultPaging struct {
	Cursors Cursors `json:"cursors"`
}

type Cursors struct {
	Before string `json:"before,omitempty"`
	After  string `json:"after,omitempty"`
}

type ThreatDescriptor struct {
	ID           string          `json:"id"`
	Indicator    IndicatorResult `json:"indicator,omitempty"`
	Owner        TEOwner         `json:"owner,omitempty"`
	Type         string          `json:"type,omitempty"`
	RawIndicator string          `json:"raw_indicator,omitempty"`
	Description  string          `json:"description,omitempty"`
	Status       string          `json:"status,omitempty"`
	PrivacyType  string          `json:"privacy_type,omitempty"`
	ShareLevel   string          `json:"share_level,omitempty"`
	AddedOn      string          `json:"added_on,omitempty"`
	Confidence   int             `json:"confidence,omitempty"`
	ExpiredOn    int             `json:"expired_on,omitempty"`
	LastUpdated  string          `json:"last_updated,omitempty"`
	SourceUri    string          `json:"source_uri,omitempty"`
}

type IndicatorResult struct {
	ID        string `json:"id,omitempty"`
	Indicator string `json:"indicator,omitempty"`
	Type      string `json:"type,omitempty"`
}

type TEOwner struct {
	ID    string `json:"id,omitempty"`
	Email string `json:"email,omitempty"`
	Name  string `json:"name,omitempty"`
}

type TEOwnersResults struct {
	Data   []TEOwner    `json:"data"`
	Paging ResultPaging `json:"paging,omitempty"`
	Next   string       `json:"next,omitempty"`
}

// New - will create a new ThreatExchange Go client, log param may be nil.
func New(appID, appSecret string, log *log.Logger) (*Client, error) {
	c := &Client{
		appId:     appID,
		appSecret: appSecret,
		logger:    log,
		c:         http.DefaultClient,
	}
	return c, nil
}

// GetThreatIndicators - is a query to retrieve ThreatIndicators, will return ThreatDescriptorResults to hold results, and raw json of that
// limit of size <=0 will be ignored
func (c *Client) GetThreatIndicators(resourceType, text, startTime, endTime string, limit int,
	extraParams map[string]string) (*ThreatDescriptorResults, string, error) {
	var result = &ThreatDescriptorResults{}
	res, err := c.query(apiVersion, "threat_descriptors", startTime, endTime, resourceType, text, limit, extraParams, result)
	if err != nil {
		return nil, "", err
	}

	return result, res, nil
}

// GetMalwareAnalyses - is a query to retrieve Malware Analyses, will return MalwareResults to hold results, and raw json of that
func (c *Client) GetMalwareAnalyses(text, startTime, endTime string, limit int,
	extraParams map[string]string) (*MalwareResults, string, error) {
	var result = &MalwareResults{}
	res, err := c.query(apiVersion, "malware_analyses", startTime, endTime, "", text, limit, extraParams, result)
	if err != nil {
		return nil, "", err
	}
	return result, res, nil
}

// GetMalwareFamilies - is a query to retrieve ThreatIndicators, will return MalwareFamiliesResults to hold results, and raw json of that
func (c *Client) GetMalwareFamilies(text, startTime, endTime string, limit int,
	extraParams map[string]string) (*MalwareFamiliesResults, string, error) {
	var result = &MalwareFamiliesResults{}
	res, err := c.query(apiVersion, "malware_families", startTime, endTime, "", text, limit, extraParams, result)
	if err != nil {
		return nil, "", err
	}
	return result, res, nil
}

// GetThreatExchangeMembers - is a query to retrieve ThreatExchange members, will return TEOwnersResults to hold results, and raw json of that
func (c *Client) GetThreatExchangeMembers() (*TEOwnersResults, string, error) {
	var result = &TEOwnersResults{}
	res, err := c.query(apiVersion, "threat_exchange_members", "", "", "", "", 0, map[string]string{}, result)
	if err != nil {
		return nil, "", err
	}

	return result, res, nil
}

func (c *Client) query(apiVersion string, resource string, startTime string, endTime string, resourceType string,
	text string, limit int, extraParams map[string]string, result interface{}) (string, error) {
	u, err := url.Parse(DefaultURL)
	if err != nil {
		return "", err
	}

	parameters := url.Values{}
	parameters.Add("access_token", fmt.Sprintf("%s|%s", c.appId, c.appSecret))
	if len(startTime) > 0 {
		parameters.Add("since", startTime)
	}
	if len(endTime) > 0 {
		parameters.Add("until", endTime)
	}
	if len(resourceType) > 0 {
		parameters.Add("type", resourceType)
	}
	if len(resourceType) > 0 {
		parameters.Add("text", text)
	}
	if limit > 0 {
		parameters.Add("limit", fmt.Sprintf("%d", limit))
	}
	for k, v := range extraParams {
		parameters.Add(k, v)
	}
	u.RawQuery = parameters.Encode()
	u.Path += fmt.Sprintf("/%s/%s/", apiVersion, resource)

	res, err := http.Get(u.String())
	if err != nil {
		return "", err
	}
	if res.Body == nil {
		return "", fmt.Errorf("Empty body for query :", u.String())
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		return "", fmt.Errorf("Wrong http return code, got : %d", res.StatusCode)
	}
	byteRes, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return "", err
	}
	err = json.Unmarshal(byteRes, result)
	if err != nil {
		return "", err
	}
	return string(byteRes), nil
}
<|endoftext|>


<commit_before>// Copyright 2014 tsuru 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 auth

import (
	"crypto"
	"crypto/rand"
	stderrors "errors"
	"fmt"
	"time"

	"github.com/tsuru/config"
	"github.com/tsuru/go-gandalfclient"
	"github.com/tsuru/tsuru/action"
	"github.com/tsuru/tsuru/db"
	"github.com/tsuru/tsuru/errors"
	"github.com/tsuru/tsuru/quota"
	"github.com/tsuru/tsuru/repository"
	"github.com/tsuru/tsuru/validation"
	"gopkg.in/mgo.v2/bson"
)

var (
	ErrUserNotFound      = stderrors.New("user not found")
	ErrUserAlreadyHasKey = stderrors.New("user already has this key")
	ErrKeyNotFound       = stderrors.New("key not found")
)

type Key struct {
	Name    string
	Content string
}

type User struct {
	Email    string
	Password string
	Keys     []Key
	quota.Quota
	APIKey string
}

// keyToMap converts a Key array into a map maybe we should store a map
// directly instead of having a convertion
func keyToMap(keys []Key) map[string]string {
	keysMap := make(map[string]string, len(keys))
	for _, k := range keys {
		keysMap[k.Name] = k.Content
	}
	return keysMap
}

func ListUsers() ([]User, error) {
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	var users []User
	err = conn.Users().Find(nil).All(&users)
	if err != nil {
		return nil, err
	}
	return users, nil
}

func GetUserByEmail(email string) (*User, error) {
	if !validation.ValidateEmail(email) {
		return nil, &errors.ValidationError{Message: "Invalid email."}
	}
	var u User
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	err = conn.Users().Find(bson.M{"email": email}).One(&u)
	if err != nil {
		return nil, ErrUserNotFound
	}
	return &u, nil
}

func (u *User) Create() error {
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	if u.Quota.Limit == 0 {
		u.Quota = quota.Unlimited
		if limit, err := config.GetInt("quota:apps-per-user"); err == nil && limit > -1 {
			u.Quota.Limit = limit
		}
	}
	defer conn.Close()
	return conn.Users().Insert(u)
}

func (u *User) Delete() error {
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	return conn.Users().Remove(bson.M{"email": u.Email})
}

func (u *User) Update() error {
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	return conn.Users().Update(bson.M{"email": u.Email}, u)
}

// Teams returns a slice containing all teams that the user is member of.
func (u *User) Teams() ([]Team, error) {
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	var teams []Team
	err = conn.Teams().Find(bson.M{"users": u.Email}).All(&teams)
	if err != nil {
		return nil, err
	}
	return teams, nil
}

func (u *User) FindKey(key Key) (Key, int) {
	for i, k := range u.Keys {
		if k.Name == key.Name || k.Content == key.Content {
			return k, i
		}
	}
	return Key{}, -1
}

func (u *User) HasKey(key Key) bool {
	_, index := u.FindKey(key)
	return index > -1
}

func (u *User) AddKey(key Key) error {
	if key.Name == "" {
		key.Name = fmt.Sprintf("%s-%d", u.Email, len(u.Keys)+1)
	}
	if u.HasKey(key) {
		return ErrUserAlreadyHasKey
	}
	actions := []*action.Action{
		&addKeyInGandalfAction,
		&addKeyInDatabaseAction,
	}
	pipeline := action.NewPipeline(actions...)
	return pipeline.Execute(&key, u)
}

func (u *User) addKeyGandalf(key *Key) error {
	gandalfClient := gandalf.Client{Endpoint: repository.ServerURL()}
	if err := gandalfClient.AddKey(u.Email, keyToMap([]Key{*key})); err != nil {
		return fmt.Errorf("Failed to add key to git server: %s", err)
	}
	return nil
}

func (u *User) addKeyDB(key *Key) error {
	u.Keys = append(u.Keys, *key)
	return u.Update()
}

func (u *User) RemoveKey(key Key) error {
	actualKey, index := u.FindKey(key)
	if index < 0 {
		return ErrKeyNotFound
	}
	err := u.removeKeyGandalf(&actualKey)
	if err != nil {
		return err
	}
	copy(u.Keys[index:], u.Keys[index+1:])
	u.Keys = u.Keys[:len(u.Keys)-1]
	return u.Update()
}

func (u *User) removeKeyGandalf(key *Key) error {
	gandalfClient := gandalf.Client{Endpoint: repository.ServerURL()}
	if err := gandalfClient.RemoveKey(u.Email, key.Name); err != nil {
		return fmt.Errorf("Failed to remove the key from git server: %s", err)
	}
	return nil
}

func (u *User) IsAdmin() bool {
	adminTeamName, err := config.GetString("admin-team")
	if err != nil {
		return false
	}
	teams, err := u.Teams()
	if err != nil {
		return false
	}
	for _, t := range teams {
		if t.Name == adminTeamName {
			return true
		}
	}
	return false
}

func (u *User) AllowedApps() ([]string, error) {
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	var alwdApps []map[string]string
	teams, err := u.Teams()
	if err != nil {
		return nil, err
	}
	teamNames := GetTeamsNames(teams)
	q := bson.M{"teams": bson.M{"$in": teamNames}}
	if err := conn.Apps().Find(q).Select(bson.M{"name": 1}).All(&alwdApps); err != nil {
		return nil, err
	}
	appNames := make([]string, len(alwdApps))
	for i, v := range alwdApps {
		appNames[i] = v["name"]
	}
	return appNames, nil
}

func (u *User) ListKeys() (map[string]string, error) {
	gURL := repository.ServerURL()
	c := gandalf.Client{Endpoint: gURL}
	return c.ListKeys(u.Email)
}

func (u *User) CreateOnGandalf() error {
	gURL := repository.ServerURL()
	c := gandalf.Client{Endpoint: gURL}
	if _, err := c.NewUser(u.Email, keyToMap(u.Keys)); err != nil {
		return fmt.Errorf("Failed to create user in the git server: %s", err)
	}
	return nil
}

func (u *User) ShowAPIKey() (string, error) {
	if u.APIKey == "" {
		u.RegenerateAPIKey()
	}
	return u.APIKey, u.Update()
}

func (u *User) RegenerateAPIKey() (string, error) {
	random_byte := make([]byte, 32)
	_, err := rand.Read(random_byte)
	if err != nil {
		return "", err
	}
	h := crypto.SHA256.New()
	h.Write([]byte(u.Email))
	h.Write(random_byte)
	h.Write([]byte(time.Now().Format(time.RFC3339Nano)))
	u.APIKey = fmt.Sprintf("%x", h.Sum(nil))
	return u.APIKey, u.Update()
}
<commit_msg>add doc to ListUsers func<commit_after>// Copyright 2014 tsuru 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 auth

import (
	"crypto"
	"crypto/rand"
	stderrors "errors"
	"fmt"
	"time"

	"github.com/tsuru/config"
	"github.com/tsuru/go-gandalfclient"
	"github.com/tsuru/tsuru/action"
	"github.com/tsuru/tsuru/db"
	"github.com/tsuru/tsuru/errors"
	"github.com/tsuru/tsuru/quota"
	"github.com/tsuru/tsuru/repository"
	"github.com/tsuru/tsuru/validation"
	"gopkg.in/mgo.v2/bson"
)

var (
	ErrUserNotFound      = stderrors.New("user not found")
	ErrUserAlreadyHasKey = stderrors.New("user already has this key")
	ErrKeyNotFound       = stderrors.New("key not found")
)

type Key struct {
	Name    string
	Content string
}

type User struct {
	Email    string
	Password string
	Keys     []Key
	quota.Quota
	APIKey string
}

// keyToMap converts a Key array into a map maybe we should store a map
// directly instead of having a convertion
func keyToMap(keys []Key) map[string]string {
	keysMap := make(map[string]string, len(keys))
	for _, k := range keys {
		keysMap[k.Name] = k.Content
	}
	return keysMap
}

// ListUsers list all users registred in tsuru
func ListUsers() ([]User, error) {
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	var users []User
	err = conn.Users().Find(nil).All(&users)
	if err != nil {
		return nil, err
	}
	return users, nil
}

func GetUserByEmail(email string) (*User, error) {
	if !validation.ValidateEmail(email) {
		return nil, &errors.ValidationError{Message: "Invalid email."}
	}
	var u User
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	err = conn.Users().Find(bson.M{"email": email}).One(&u)
	if err != nil {
		return nil, ErrUserNotFound
	}
	return &u, nil
}

func (u *User) Create() error {
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	if u.Quota.Limit == 0 {
		u.Quota = quota.Unlimited
		if limit, err := config.GetInt("quota:apps-per-user"); err == nil && limit > -1 {
			u.Quota.Limit = limit
		}
	}
	defer conn.Close()
	return conn.Users().Insert(u)
}

func (u *User) Delete() error {
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	return conn.Users().Remove(bson.M{"email": u.Email})
}

func (u *User) Update() error {
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	return conn.Users().Update(bson.M{"email": u.Email}, u)
}

// Teams returns a slice containing all teams that the user is member of.
func (u *User) Teams() ([]Team, error) {
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	var teams []Team
	err = conn.Teams().Find(bson.M{"users": u.Email}).All(&teams)
	if err != nil {
		return nil, err
	}
	return teams, nil
}

func (u *User) FindKey(key Key) (Key, int) {
	for i, k := range u.Keys {
		if k.Name == key.Name || k.Content == key.Content {
			return k, i
		}
	}
	return Key{}, -1
}

func (u *User) HasKey(key Key) bool {
	_, index := u.FindKey(key)
	return index > -1
}

func (u *User) AddKey(key Key) error {
	if key.Name == "" {
		key.Name = fmt.Sprintf("%s-%d", u.Email, len(u.Keys)+1)
	}
	if u.HasKey(key) {
		return ErrUserAlreadyHasKey
	}
	actions := []*action.Action{
		&addKeyInGandalfAction,
		&addKeyInDatabaseAction,
	}
	pipeline := action.NewPipeline(actions...)
	return pipeline.Execute(&key, u)
}

func (u *User) addKeyGandalf(key *Key) error {
	gandalfClient := gandalf.Client{Endpoint: repository.ServerURL()}
	if err := gandalfClient.AddKey(u.Email, keyToMap([]Key{*key})); err != nil {
		return fmt.Errorf("Failed to add key to git server: %s", err)
	}
	return nil
}

func (u *User) addKeyDB(key *Key) error {
	u.Keys = append(u.Keys, *key)
	return u.Update()
}

func (u *User) RemoveKey(key Key) error {
	actualKey, index := u.FindKey(key)
	if index < 0 {
		return ErrKeyNotFound
	}
	err := u.removeKeyGandalf(&actualKey)
	if err != nil {
		return err
	}
	copy(u.Keys[index:], u.Keys[index+1:])
	u.Keys = u.Keys[:len(u.Keys)-1]
	return u.Update()
}

func (u *User) removeKeyGandalf(key *Key) error {
	gandalfClient := gandalf.Client{Endpoint: repository.ServerURL()}
	if err := gandalfClient.RemoveKey(u.Email, key.Name); err != nil {
		return fmt.Errorf("Failed to remove the key from git server: %s", err)
	}
	return nil
}

func (u *User) IsAdmin() bool {
	adminTeamName, err := config.GetString("admin-team")
	if err != nil {
		return false
	}
	teams, err := u.Teams()
	if err != nil {
		return false
	}
	for _, t := range teams {
		if t.Name == adminTeamName {
			return true
		}
	}
	return false
}

func (u *User) AllowedApps() ([]string, error) {
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	var alwdApps []map[string]string
	teams, err := u.Teams()
	if err != nil {
		return nil, err
	}
	teamNames := GetTeamsNames(teams)
	q := bson.M{"teams": bson.M{"$in": teamNames}}
	if err := conn.Apps().Find(q).Select(bson.M{"name": 1}).All(&alwdApps); err != nil {
		return nil, err
	}
	appNames := make([]string, len(alwdApps))
	for i, v := range alwdApps {
		appNames[i] = v["name"]
	}
	return appNames, nil
}

func (u *User) ListKeys() (map[string]string, error) {
	gURL := repository.ServerURL()
	c := gandalf.Client{Endpoint: gURL}
	return c.ListKeys(u.Email)
}

func (u *User) CreateOnGandalf() error {
	gURL := repository.ServerURL()
	c := gandalf.Client{Endpoint: gURL}
	if _, err := c.NewUser(u.Email, keyToMap(u.Keys)); err != nil {
		return fmt.Errorf("Failed to create user in the git server: %s", err)
	}
	return nil
}

func (u *User) ShowAPIKey() (string, error) {
	if u.APIKey == "" {
		u.RegenerateAPIKey()
	}
	return u.APIKey, u.Update()
}

func (u *User) RegenerateAPIKey() (string, error) {
	random_byte := make([]byte, 32)
	_, err := rand.Read(random_byte)
	if err != nil {
		return "", err
	}
	h := crypto.SHA256.New()
	h.Write([]byte(u.Email))
	h.Write(random_byte)
	h.Write([]byte(time.Now().Format(time.RFC3339Nano)))
	u.APIKey = fmt.Sprintf("%x", h.Sum(nil))
	return u.APIKey, u.Update()
}
<|endoftext|>


<commit_before>// Package base58 implements a human-friendly base58 encoding.
//
// As opposed to base64 and friends, base58 is typically used to
// convert integers. You can use big.Int.SetBytes to convert arbitrary
// bytes to an integer first, and big.Int.Bytes the other way around.
package base58

import (
	"math/big"
	"strconv"
)

const alphabet = "123456789abcdefghijkmnopqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ"

var decodeMap [256]byte

func init() {
	for i := 0; i < len(decodeMap); i++ {
		decodeMap[i] = 0xFF
	}
	for i := 0; i < len(alphabet); i++ {
		decodeMap[alphabet[i]] = byte(i)
	}
}

type CorruptInputError int64

func (e CorruptInputError) Error() string {
	return "illegal base58 data at input byte " + strconv.FormatInt(int64(e), 10)
}

// Decode a big integer from the bytes. Returns an error on corrupt
// input.
func DecodeToBig(src []byte) (*big.Int, error) {
	n := new(big.Int)
	radix := big.NewInt(58)
	for i := 0; i < len(src); i++ {
		b := decodeMap[src[i]]
		if b == 0xFF {
			return nil, CorruptInputError(i)
		}
		n.Mul(n, radix)
		n.Add(n, big.NewInt(int64(b)))
	}
	return n, nil
}

// Encode encodes src, appending to dst. Be sure to use the returned
// new value of dst.
func EncodeBig(dst []byte, src *big.Int) []byte {
	start := len(dst)
	n := new(big.Int)
	n.Set(src)
	radix := big.NewInt(58)
	zero := big.NewInt(0)

	for n.Cmp(zero) > 0 {
		mod := new(big.Int)
		n.DivMod(n, radix, mod)
		dst = append(dst, alphabet[mod.Int64()])
	}

	for i, j := start, len(dst)-1; i < j; i, j = i+1, j-1 {
		dst[i], dst[j] = dst[j], dst[i]
	}
	return dst
}
<commit_msg>Initialize constant values only once<commit_after>// Package base58 implements a human-friendly base58 encoding.
//
// As opposed to base64 and friends, base58 is typically used to
// convert integers. You can use big.Int.SetBytes to convert arbitrary
// bytes to an integer first, and big.Int.Bytes the other way around.
package base58

import (
	"math/big"
	"strconv"
)

const alphabet = "123456789abcdefghijkmnopqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ"

var decodeMap [256]byte

var (
	// constant-like variables

	radix = big.NewInt(58)
	zero  = big.NewInt(0)
)

func init() {
	for i := 0; i < len(decodeMap); i++ {
		decodeMap[i] = 0xFF
	}
	for i := 0; i < len(alphabet); i++ {
		decodeMap[alphabet[i]] = byte(i)
	}
}

type CorruptInputError int64

func (e CorruptInputError) Error() string {
	return "illegal base58 data at input byte " + strconv.FormatInt(int64(e), 10)
}

// Decode a big integer from the bytes. Returns an error on corrupt
// input.
func DecodeToBig(src []byte) (*big.Int, error) {
	n := new(big.Int)
	for i := 0; i < len(src); i++ {
		b := decodeMap[src[i]]
		if b == 0xFF {
			return nil, CorruptInputError(i)
		}
		n.Mul(n, radix)
		n.Add(n, big.NewInt(int64(b)))
	}
	return n, nil
}

// Encode encodes src, appending to dst. Be sure to use the returned
// new value of dst.
func EncodeBig(dst []byte, src *big.Int) []byte {
	start := len(dst)
	n := new(big.Int)
	n.Set(src)

	for n.Cmp(zero) > 0 {
		mod := new(big.Int)
		n.DivMod(n, radix, mod)
		dst = append(dst, alphabet[mod.Int64()])
	}

	for i, j := start, len(dst)-1; i < j; i, j = i+1, j-1 {
		dst[i], dst[j] = dst[j], dst[i]
	}
	return dst
}
<|endoftext|>


<commit_before>package net

import "fmt"
import "syscall"
import "unsafe"

const (
	SIOCETHTOOL     = 0x8946     // linux/sockios.h
	ETHTOOL_STXCSUM = 0x00000017 // linux/ethtool.h
	IFNAMSIZ        = 16         // linux/if.h
)

// linux/if.h 'struct ifreq'
type IFReqData struct {
	Name [IFNAMSIZ]byte
	Data uintptr
}

// linux/ethtool.h 'struct ethtool_value'
type EthtoolValue struct {
	Cmd  uint32
	Data uint32
}

// Disable TX checksum offload on specified interface
func EthtoolTXOff(name string) error {
	if len(name)+1 > IFNAMSIZ {
		return fmt.Errorf("name too long")
	}

	value := EthtoolValue{ETHTOOL_STXCSUM, 0}
	request := IFReqData{Data: uintptr(unsafe.Pointer(&value))}

	copy(request.Name[:], name)

	socket, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_DGRAM, 0)
	if err != nil {
		return err
	}
	defer syscall.Close(socket)

	_, _, errno := syscall.RawSyscall(syscall.SYS_IOCTL,
		uintptr(socket),
		uintptr(SIOCETHTOOL),
		uintptr(unsafe.Pointer(&request)))

	if errno != 0 {
		return errno
	}

	return nil
}
<commit_msg>Check current tx value before trying to set it<commit_after>package net

import "fmt"
import "syscall"
import "unsafe"

const (
	SIOCETHTOOL     = 0x8946     // linux/sockios.h
	ETHTOOL_GTXCSUM = 0x00000016 // linux/ethtool.h
	ETHTOOL_STXCSUM = 0x00000017 // linux/ethtool.h
	IFNAMSIZ        = 16         // linux/if.h
)

// linux/if.h 'struct ifreq'
type IFReqData struct {
	Name [IFNAMSIZ]byte
	Data uintptr
}

// linux/ethtool.h 'struct ethtool_value'
type EthtoolValue struct {
	Cmd  uint32
	Data uint32
}

func ioctlEthtool(fd int, argp uintptr) error {
	_, _, errno := syscall.RawSyscall(syscall.SYS_IOCTL, uintptr(fd), uintptr(SIOCETHTOOL), argp)
	if errno != 0 {
		return errno
	}
	return nil
}

// Disable TX checksum offload on specified interface
func EthtoolTXOff(name string) error {
	if len(name)+1 > IFNAMSIZ {
		return fmt.Errorf("name too long")
	}

	socket, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_DGRAM, 0)
	if err != nil {
		return err
	}
	defer syscall.Close(socket)

	// Request current value
	value := EthtoolValue{Cmd: ETHTOOL_GTXCSUM}
	request := IFReqData{Data: uintptr(unsafe.Pointer(&value))}
	copy(request.Name[:], name)

	if err := ioctlEthtool(socket, uintptr(unsafe.Pointer(&request))); err != nil {
		return err
	}
	if value.Data == 0 { // if already off, don't try to change
		return nil
	}

	value = EthtoolValue{ETHTOOL_STXCSUM, 0}
	return ioctlEthtool(socket, uintptr(unsafe.Pointer(&request)))
}
<|endoftext|>


<commit_before>package emoji

import (
	"bytes"
	"testing"
	"fmt"
)

const (
	beerKey  = ":beer:"
	beerText = " ビール!!!"
)

var testFText = "test " + emojiCodeMap[beerKey] + ReplacePadding + beerText
var testText = emojiCodeMap[beerKey] + ReplacePadding + beerText

func TestPrint(t *testing.T) {
	_, err := Print(beerKey, beerText)
	if err != nil {
		t.Error("Print ", err)
	}
}

func TestPrintln(t *testing.T) {
	_, err := Println(beerKey, beerText)
	if err != nil {
		t.Error("Println ", err)
	}
}

func TestPrintf(t *testing.T) {
	_, err := Printf("%s "+beerKey+beerText, "test")
	if err != nil {
		t.Error("Printf ", err)
	}
}

func TestFprint(t *testing.T) {
	var buf bytes.Buffer
	_, err := Fprint(&buf, beerKey+beerText)
	if err != nil {
		t.Error("Fprint ", err)
	}

	if buf.String() != testText {
		t.Error("Fprint ", buf.String(), testText)
	}
}

func TestFprintln(t *testing.T) {
	var buf bytes.Buffer
	_, err := Fprintln(&buf, beerKey+beerText)
	if err != nil {
		t.Error("Fprintln ", err)
	}

	if buf.String() != (testText + "\n") {
		t.Error("Fprintln ", buf.String(), (testText + "\n"))
	}
}

func TestFprintf(t *testing.T) {
	var buf bytes.Buffer
	_, err := Fprintf(&buf, "%s "+beerKey+beerText, "test")
	if err != nil {
		t.Error("Fprintf ", err)
	}

	if buf.String() != testFText {
		t.Error("Fprintf ", buf.String(), testFText)
	}
}

func TestSprint(t *testing.T) {
	convertBeer := Sprint(beerKey, beerText)
	if convertBeer != testText {
		t.Error("Sprint ", convertBeer, testText)
	}
}

func TestSprintf(t *testing.T) {
	convertBeer := Sprintf("%s "+beerKey+beerText, "test")
	if convertBeer != testFText {
		t.Error("Sprintf ", convertBeer, testFText)
	}
}

func TestErrorf(t *testing.T) {
	error := Errorf("%s "+beerKey+beerText, "test")
	if error.Error() != testFText {
		t.Error("Errorf ", error, testFText)
	}
}
<commit_msg>test: added emoji multi pattern<commit_after>package emoji

import (
	"bytes"
	"testing"
	"fmt"
)

const (
	beerKey  = ":beer:"
	beerText = " ビール!!!"
)

var testFText = "test " + emojiCodeMap[beerKey] + ReplacePadding + beerText
var testText = emojiCodeMap[beerKey] + ReplacePadding + beerText

func TestPrint(t *testing.T) {
	_, err := Print(beerKey, beerText)
	if err != nil {
		t.Error("Print ", err)
	}
}

func TestPrintln(t *testing.T) {
	_, err := Println(beerKey, beerText)
	if err != nil {
		t.Error("Println ", err)
	}
}

func TestPrintf(t *testing.T) {
	_, err := Printf("%s "+beerKey+beerText, "test")
	if err != nil {
		t.Error("Printf ", err)
	}
}

func TestFprint(t *testing.T) {
	var buf bytes.Buffer
	_, err := Fprint(&buf, beerKey+beerText)
	if err != nil {
		t.Error("Fprint ", err)
	}

	if buf.String() != testText {
		t.Error("Fprint ", buf.String(), testText)
	}
}

func TestFprintln(t *testing.T) {
	var buf bytes.Buffer
	_, err := Fprintln(&buf, beerKey+beerText)
	if err != nil {
		t.Error("Fprintln ", err)
	}

	if buf.String() != (testText + "\n") {
		t.Error("Fprintln ", buf.String(), (testText + "\n"))
	}
}

func TestFprintf(t *testing.T) {
	var buf bytes.Buffer
	_, err := Fprintf(&buf, "%s "+beerKey+beerText, "test")
	if err != nil {
		t.Error("Fprintf ", err)
	}

	if buf.String() != testFText {
		t.Error("Fprintf ", buf.String(), testFText)
	}
}

func TestSprint(t *testing.T) {
	convertBeer := Sprint(beerKey, beerText)
	if convertBeer != testText {
		t.Error("Sprint ", convertBeer, testText)
	}
}

func TestSprintf(t *testing.T) {
	convertBeer := Sprintf("%s "+beerKey+beerText, "test")
	if convertBeer != testFText {
		t.Error("Sprintf ", convertBeer, testFText)
	}
}

func TestErrorf(t *testing.T) {
	error := Errorf("%s "+beerKey+beerText, "test")
	if error.Error() != testFText {
		t.Error("Errorf ", error, testFText)
	}
}

func TestSprintMulti(t *testing.T) {
	convertBeer := Sprint(beerKey, beerText, beerKey, beerText)
	if convertBeer != (testText + testText) {
		t.Error("Sprint ", convertBeer, testText)
	}
	fmt.Println(convertBeer)
}
<|endoftext|>