<commit_before>package telebot

import (
	"strconv"
	"time"
)

// Message object represents a message.
type Message struct {
	ID int `json:"message_id"`

	InlineID string `json:"-"`

	// For message sent to channels, Sender will be nil
	Sender *User `json:"from"`

	// Unixtime, use Message.Time() to get time.Time
	Unixtime int64 `json:"date"`

	// Conversation the message belongs to.
	Chat *Chat `json:"chat"`

	// For forwarded messages, sender of the original message.
	OriginalSender *User `json:"forward_from"`

	// For forwarded messages, chat of the original message when
	// forwarded from a channel.
	OriginalChat *Chat `json:"forward_from_chat"`

	// For forwarded messages, unixtime of the original message.
	OriginalUnixtime int `json:"forward_date"`

	// For replies, ReplyTo represents the original message.
	//
	// Note that the Message object in this field will not
	// contain further ReplyTo fields even if it
	// itself is a reply.
	ReplyTo *Message `json:"reply_to_message"`

	// (Optional) Time of last edit in Unix
	LastEdit int64 `json:"edit_date"`

	// AlbumID is the unique identifier of a media message group
	// this message belongs to.
	AlbumID string `json:"media_group_id"`

	// Author signature (in channels).
	Signature string `json:"author_signature"`

	// For a text message, the actual UTF-8 text of the message.
	Text string `json:"text"`

	// For registered commands, will contain the string payload:
	//
	// Ex: `/command <payload>` or `/command@botname <payload>`
	Payload string `json:"-"`

	// For text messages, special entities like usernames, URLs, bot commands,
	// etc. that appear in the text.
	Entities []MessageEntity `json:"entities,omitempty"`

	// Some messages containing media, may as well have a caption.
	Caption string `json:"caption,omitempty"`

	// For messages with a caption, special entities like usernames, URLs,
	// bot commands, etc. that appear in the caption.
	CaptionEntities []MessageEntity `json:"caption_entities,omitempty"`

	// For an audio recording, information about it.
	Audio *Audio `json:"audio"`

	// For a gneral file, information about it.
	Document *Document `json:"document"`

	// For a photo, all available sizes (thumbnails).
	Photo *Photo `json:"photo"`

	// For a sticker, information about it.
	Sticker *Sticker `json:"sticker"`

	// For a voice message, information about it.
	Voice *Voice `json:"voice"`

	// For a video note, information about it.
	VideoNote *VideoNote `json:"video_note"`

	// For a video, information about it.
	Video *Video `json:"video"`

	// For a contact, contact information itself.
	Contact *Contact `json:"contact"`

	// For a location, its longitude and latitude.
	Location *Location `json:"location"`

	// For a venue, information about it.
	Venue *Venue `json:"venue"`

	// For a service message, represents a user,
	// that just got added to chat, this message came from.
	//
	// Sender leads to User, capable of invite.
	//
	// UserJoined might be the Bot itself.
	UserJoined *User `json:"new_chat_member"`

	// For a service message, represents a user,
	// that just left chat, this message came from.
	//
	// If user was kicked, Sender leads to a User,
	// capable of this kick.
	//
	// UserLeft might be the Bot itself.
	UserLeft *User `json:"left_chat_member"`

	// For a service message, represents a new title
	// for chat this message came from.
	//
	// Sender would lead to a User, capable of change.
	NewGroupTitle string `json:"new_chat_title"`

	// For a service message, represents all available
	// thumbnails of the new chat photo.
	//
	// Sender would lead to a User, capable of change.
	NewGroupPhoto *Photo `json:"new_chat_photo"`

	// For a service message, new members that were added to
	// the group or supergroup and information about them
	// (the bot itself may be one of these members).
	UsersJoined []User `json:"new_chat_members"`

	// For a service message, true if chat photo just
	// got removed.
	//
	// Sender would lead to a User, capable of change.
	GroupPhotoDeleted bool `json:"delete_chat_photo"`

	// For a service message, true if group has been created.
	//
	// You would recieve such a message if you are one of
	// initial group chat members.
	//
	// Sender would lead to creator of the chat.
	GroupCreated bool `json:"group_chat_created"`

	// For a service message, true if super group has been created.
	//
	// You would recieve such a message if you are one of
	// initial group chat members.
	//
	// Sender would lead to creator of the chat.
	SuperGroupCreated bool `json:"supergroup_chat_created"`

	// For a service message, true if channel has been created.
	//
	// You would recieve such a message if you are one of
	// initial channel administrators.
	//
	// Sender would lead to creator of the chat.
	ChannelCreated bool `json:"channel_chat_created"`

	// For a service message, the destination (super group) you
	// migrated to.
	//
	// You would recieve such a message when your chat has migrated
	// to a super group.
	//
	// Sender would lead to creator of the migration.
	MigrateTo int64 `json:"migrate_to_chat_id"`

	// For a service message, the Origin (normal group) you migrated
	// from.
	//
	// You would recieve such a message when your chat has migrated
	// to a super group.
	//
	// Sender would lead to creator of the migration.
	MigrateFrom int64 `json:"migrate_from_chat_id"`

	// Specified message was pinned. Note that the Message object
	// in this field will not contain further ReplyTo fields even
	// if it is itself a reply.
	PinnedMessage *Message `json:"pinned_message"`
}

// MessageEntity object represents "special" parts of text messages,
// including hashtags, usernames, URLs, etc.
type MessageEntity struct {
	// Specifies entity type.
	Type EntityType `json:"type"`

	// Offset in UTF-16 code units to the start of the entity.
	Offset int `json:"offset"`

	// Length of the entity in UTF-16 code units.
	Length int `json:"length"`

	// (Optional) For EntityTextLink entity type only.
	//
	// URL will be opened after user taps on the text.
	URL string `json:"url,omitempty"`

	// (Optional) For EntityTMention entity type only.
	User *User `json:"user,omitempty"`
}

// MessageSig satisfies Editable interface (see Editable.)
func (m *Message) MessageSig() (string, int64) {
	if m.InlineID != "" {
		return m.InlineID, 0
	}
	return strconv.Itoa(m.ID), m.Chat.ID
}

// Time returns the moment of message creation in local time.
func (m *Message) Time() time.Time {
	return time.Unix(m.Unixtime, 0)
}

// LastEdited returns time.Time of last edit.
func (m *Message) LastEdited() time.Time {
	return time.Unix(m.LastEdit, 0)
}

// IsForwarded says whether message is forwarded copy of another
// message or not.
func (m *Message) IsForwarded() bool {
	return m.OriginalSender != nil || m.OriginalChat != nil
}

// IsReply says whether message is a reply to another message.
func (m *Message) IsReply() bool {
	return m.ReplyTo != nil
}

// Private returns true, if it's a personal message.
func (m *Message) Private() bool {
	return m.Chat.Type == ChatPrivate
}

// FromGroup returns true, if message came from a group OR
// a super group.
func (m *Message) FromGroup() bool {
	return m.Chat.Type == ChatGroup || m.Chat.Type == ChatSuperGroup
}

// FromChannel returns true, if message came from a channel.
func (m *Message) FromChannel() bool {
	return m.Chat.Type == ChatChannel
}

// IsService returns true, if message is a service message,
// returns false otherwise.
//
// Service messages are automatically sent messages, which
// typically occur on some global action. For instance, when
// anyone leaves the chat or chat title changes.
func (m *Message) IsService() bool {
	fact := false

	fact = fact || m.UserJoined != nil
	fact = fact || len(m.UsersJoined) > 0
	fact = fact || m.UserLeft != nil
	fact = fact || m.NewGroupTitle != ""
	fact = fact || m.NewGroupPhoto != nil
	fact = fact || m.GroupPhotoDeleted
	fact = fact || m.GroupCreated || m.SuperGroupCreated
	fact = fact || (m.MigrateTo != m.MigrateFrom)

	return fact
}
<commit_msg>Add ReplyMarkup field to Message<commit_after>package telebot

import (
	"strconv"
	"time"
)

// Message object represents a message.
type Message struct {
	ID int `json:"message_id"`

	InlineID string `json:"-"`

	// For message sent to channels, Sender will be nil
	Sender *User `json:"from"`

	// Unixtime, use Message.Time() to get time.Time
	Unixtime int64 `json:"date"`

	// Conversation the message belongs to.
	Chat *Chat `json:"chat"`

	// For forwarded messages, sender of the original message.
	OriginalSender *User `json:"forward_from"`

	// For forwarded messages, chat of the original message when
	// forwarded from a channel.
	OriginalChat *Chat `json:"forward_from_chat"`

	// For forwarded messages, unixtime of the original message.
	OriginalUnixtime int `json:"forward_date"`

	// For replies, ReplyTo represents the original message.
	//
	// Note that the Message object in this field will not
	// contain further ReplyTo fields even if it
	// itself is a reply.
	ReplyTo *Message `json:"reply_to_message"`

	// (Optional) Time of last edit in Unix
	LastEdit int64 `json:"edit_date"`

	// AlbumID is the unique identifier of a media message group
	// this message belongs to.
	AlbumID string `json:"media_group_id"`

	// Author signature (in channels).
	Signature string `json:"author_signature"`

	// For a text message, the actual UTF-8 text of the message.
	Text string `json:"text"`

	// For registered commands, will contain the string payload:
	//
	// Ex: `/command <payload>` or `/command@botname <payload>`
	Payload string `json:"-"`

	// For text messages, special entities like usernames, URLs, bot commands,
	// etc. that appear in the text.
	Entities []MessageEntity `json:"entities,omitempty"`

	// Some messages containing media, may as well have a caption.
	Caption string `json:"caption,omitempty"`

	// For messages with a caption, special entities like usernames, URLs,
	// bot commands, etc. that appear in the caption.
	CaptionEntities []MessageEntity `json:"caption_entities,omitempty"`

	// For an audio recording, information about it.
	Audio *Audio `json:"audio"`

	// For a gneral file, information about it.
	Document *Document `json:"document"`

	// For a photo, all available sizes (thumbnails).
	Photo *Photo `json:"photo"`

	// For a sticker, information about it.
	Sticker *Sticker `json:"sticker"`

	// For a voice message, information about it.
	Voice *Voice `json:"voice"`

	// For a video note, information about it.
	VideoNote *VideoNote `json:"video_note"`

	// For a video, information about it.
	Video *Video `json:"video"`

	// For a contact, contact information itself.
	Contact *Contact `json:"contact"`

	// For a location, its longitude and latitude.
	Location *Location `json:"location"`

	// For a venue, information about it.
	Venue *Venue `json:"venue"`

	// For a service message, represents a user,
	// that just got added to chat, this message came from.
	//
	// Sender leads to User, capable of invite.
	//
	// UserJoined might be the Bot itself.
	UserJoined *User `json:"new_chat_member"`

	// For a service message, represents a user,
	// that just left chat, this message came from.
	//
	// If user was kicked, Sender leads to a User,
	// capable of this kick.
	//
	// UserLeft might be the Bot itself.
	UserLeft *User `json:"left_chat_member"`

	// For a service message, represents a new title
	// for chat this message came from.
	//
	// Sender would lead to a User, capable of change.
	NewGroupTitle string `json:"new_chat_title"`

	// For a service message, represents all available
	// thumbnails of the new chat photo.
	//
	// Sender would lead to a User, capable of change.
	NewGroupPhoto *Photo `json:"new_chat_photo"`

	// For a service message, new members that were added to
	// the group or supergroup and information about them
	// (the bot itself may be one of these members).
	UsersJoined []User `json:"new_chat_members"`

	// For a service message, true if chat photo just
	// got removed.
	//
	// Sender would lead to a User, capable of change.
	GroupPhotoDeleted bool `json:"delete_chat_photo"`

	// For a service message, true if group has been created.
	//
	// You would recieve such a message if you are one of
	// initial group chat members.
	//
	// Sender would lead to creator of the chat.
	GroupCreated bool `json:"group_chat_created"`

	// For a service message, true if super group has been created.
	//
	// You would recieve such a message if you are one of
	// initial group chat members.
	//
	// Sender would lead to creator of the chat.
	SuperGroupCreated bool `json:"supergroup_chat_created"`

	// For a service message, true if channel has been created.
	//
	// You would recieve such a message if you are one of
	// initial channel administrators.
	//
	// Sender would lead to creator of the chat.
	ChannelCreated bool `json:"channel_chat_created"`

	// For a service message, the destination (super group) you
	// migrated to.
	//
	// You would recieve such a message when your chat has migrated
	// to a super group.
	//
	// Sender would lead to creator of the migration.
	MigrateTo int64 `json:"migrate_to_chat_id"`

	// For a service message, the Origin (normal group) you migrated
	// from.
	//
	// You would recieve such a message when your chat has migrated
	// to a super group.
	//
	// Sender would lead to creator of the migration.
	MigrateFrom int64 `json:"migrate_from_chat_id"`

	// Specified message was pinned. Note that the Message object
	// in this field will not contain further ReplyTo fields even
	// if it is itself a reply.
	PinnedMessage *Message `json:"pinned_message"`

	// Inline keyboard attached to the message.
	ReplyMarkup InlineKeyboardMarkup `json:"reply_markup"`
}

// MessageEntity object represents "special" parts of text messages,
// including hashtags, usernames, URLs, etc.
type MessageEntity struct {
	// Specifies entity type.
	Type EntityType `json:"type"`

	// Offset in UTF-16 code units to the start of the entity.
	Offset int `json:"offset"`

	// Length of the entity in UTF-16 code units.
	Length int `json:"length"`

	// (Optional) For EntityTextLink entity type only.
	//
	// URL will be opened after user taps on the text.
	URL string `json:"url,omitempty"`

	// (Optional) For EntityTMention entity type only.
	User *User `json:"user,omitempty"`
}

// MessageSig satisfies Editable interface (see Editable.)
func (m *Message) MessageSig() (string, int64) {
	if m.InlineID != "" {
		return m.InlineID, 0
	}
	return strconv.Itoa(m.ID), m.Chat.ID
}

// Time returns the moment of message creation in local time.
func (m *Message) Time() time.Time {
	return time.Unix(m.Unixtime, 0)
}

// LastEdited returns time.Time of last edit.
func (m *Message) LastEdited() time.Time {
	return time.Unix(m.LastEdit, 0)
}

// IsForwarded says whether message is forwarded copy of another
// message or not.
func (m *Message) IsForwarded() bool {
	return m.OriginalSender != nil || m.OriginalChat != nil
}

// IsReply says whether message is a reply to another message.
func (m *Message) IsReply() bool {
	return m.ReplyTo != nil
}

// Private returns true, if it's a personal message.
func (m *Message) Private() bool {
	return m.Chat.Type == ChatPrivate
}

// FromGroup returns true, if message came from a group OR
// a super group.
func (m *Message) FromGroup() bool {
	return m.Chat.Type == ChatGroup || m.Chat.Type == ChatSuperGroup
}

// FromChannel returns true, if message came from a channel.
func (m *Message) FromChannel() bool {
	return m.Chat.Type == ChatChannel
}

// IsService returns true, if message is a service message,
// returns false otherwise.
//
// Service messages are automatically sent messages, which
// typically occur on some global action. For instance, when
// anyone leaves the chat or chat title changes.
func (m *Message) IsService() bool {
	fact := false

	fact = fact || m.UserJoined != nil
	fact = fact || len(m.UsersJoined) > 0
	fact = fact || m.UserLeft != nil
	fact = fact || m.NewGroupTitle != ""
	fact = fact || m.NewGroupPhoto != nil
	fact = fact || m.GroupPhotoDeleted
	fact = fact || m.GroupCreated || m.SuperGroupCreated
	fact = fact || (m.MigrateTo != m.MigrateFrom)

	return fact
}
<|endoftext|>


<commit_before>package netlink

import (
	"errors"
	"fmt"

	"github.com/mdlayher/netlink/nlenc"
)

// Flags which may apply to netlink attribute types when communicating with
// certain netlink families.
const (
	Nested       = 0x8000
	NetByteOrder = 0x4000

	// attrTypeMask masks off Type bits used for the above flags.
	attrTypeMask = 0x3fff
)

// Various errors which may occur when attempting to marshal or unmarshal
// a Message to and from its binary form.
var (
	errIncorrectMessageLength = errors.New("netlink message header length incorrect")
	errShortMessage           = errors.New("not enough data to create a netlink message")
	errUnalignedMessage       = errors.New("input data is not properly aligned for netlink message")
)

// HeaderFlags specify flags which may be present in a Header.
type HeaderFlags uint16

const (
	// General netlink communication flags.

	// Request indicates a request to netlink.
	Request HeaderFlags = 1

	// Multi indicates a multi-part message, terminated by Done on the
	// last message.
	Multi HeaderFlags = 2

	// Acknowledge requests that netlink reply with an acknowledgement
	// using Error and, if needed, an error code.
	Acknowledge HeaderFlags = 4

	// Echo requests that netlink echo this request back to the sender.
	Echo HeaderFlags = 8

	// DumpInterrupted indicates that a dump was inconsistent due to a
	// sequence change.
	DumpInterrupted HeaderFlags = 16

	// DumpFiltered indicates that a dump was filtered as requested.
	DumpFiltered HeaderFlags = 32

	// Flags used to retrieve data from netlink.

	// Root requests that netlink return a complete table instead of a
	// single entry.
	Root HeaderFlags = 0x100

	// Match requests that netlink return a list of all matching entries.
	Match HeaderFlags = 0x200

	// Atomic requests that netlink send an atomic snapshot of its entries.
	// Requires CAP_NET_ADMIN or an effective UID of 0.
	Atomic HeaderFlags = 0x400

	// Dump requests that netlink return a complete list of all entries.
	Dump HeaderFlags = Root | Match

	// Flags used to create objects.

	// Replace indicates request replaces an existing matching object.
	Replace HeaderFlags = 0x100

	// Excl indicates request does not replace the object if it already exists.
	Excl HeaderFlags = 0x200

	// Create indicates request creates an object if it doesn't already exist.
	Create HeaderFlags = 0x400

	// Append indicates request adds to the end of the object list.
	Append HeaderFlags = 0x800
)

// String returns the string representation of a HeaderFlags.
func (f HeaderFlags) String() string {
	names := []string{
		"request",
		"multi",
		"acknowledge",
		"echo",
		"dumpinterrupted",
		"dumpfiltered",
	}

	var s string

	left := uint(f)

	for i, name := range names {
		if f&(1<<uint(i)) != 0 {
			if s != "" {
				s += "|"
			}

			s += name

			left ^= (1 << uint(i))
		}
	}

	if s == "" && left == 0 {
		s = "0"
	}

	if left > 0 {
		if s != "" {
			s += "|"
		}
		s += fmt.Sprintf("%#x", left)
	}

	return s
}

// HeaderType specifies the type of a Header.
type HeaderType uint16

const (
	// Noop indicates that no action was taken.
	Noop HeaderType = 0x1

	// Error indicates an error code is present, which is also used to indicate
	// success when the code is 0.
	Error HeaderType = 0x2

	// Done indicates the end of a multi-part message.
	Done HeaderType = 0x3

	// Overrun indicates that data was lost from this message.
	Overrun HeaderType = 0x4
)

// String returns the string representation of a HeaderType.
func (t HeaderType) String() string {
	switch t {
	case Noop:
		return "noop"
	case Error:
		return "error"
	case Done:
		return "done"
	case Overrun:
		return "overrun"
	default:
		return fmt.Sprintf("unknown(%d)", t)
	}
}

// NB: the memory layout of Header and Linux's syscall.NlMsgHdr must be
// exactly the same.  Cannot reorder, change data type, add, or remove fields.
// Named types of the same size (e.g. HeaderFlags is a uint16) are okay.

// A Header is a netlink header.  A Header is sent and received with each
// Message to indicate metadata regarding a Message.
type Header struct {
	// Length of a Message, including this Header.
	Length uint32

	// Contents of a Message.
	Type HeaderType

	// Flags which may be used to modify a request or response.
	Flags HeaderFlags

	// The sequence number of a Message.
	Sequence uint32

	// The process ID of the sending process.
	PID uint32
}

// A Message is a netlink message.  It contains a Header and an arbitrary
// byte payload, which may be decoded using information from the Header.
//
// Data is encoded in the native endianness of the host system.  For easier
// of encoding and decoding of integers, use package nlenc.
type Message struct {
	Header Header
	Data   []byte
}

// MarshalBinary marshals a Message into a byte slice.
func (m Message) MarshalBinary() ([]byte, error) {
	ml := nlmsgAlign(int(m.Header.Length))
	if ml < nlmsgHeaderLen || ml != int(m.Header.Length) {
		return nil, errIncorrectMessageLength
	}

	b := make([]byte, ml)

	nlenc.PutUint32(b[0:4], m.Header.Length)
	nlenc.PutUint16(b[4:6], uint16(m.Header.Type))
	nlenc.PutUint16(b[6:8], uint16(m.Header.Flags))
	nlenc.PutUint32(b[8:12], m.Header.Sequence)
	nlenc.PutUint32(b[12:16], m.Header.PID)
	copy(b[16:], m.Data)

	return b, nil
}

// UnmarshalBinary unmarshals the contents of a byte slice into a Message.
func (m *Message) UnmarshalBinary(b []byte) error {
	if len(b) < nlmsgHeaderLen {
		return errShortMessage
	}
	if len(b) != nlmsgAlign(len(b)) {
		return errUnalignedMessage
	}

	// Don't allow misleading length
	m.Header.Length = nlenc.Uint32(b[0:4])
	if int(m.Header.Length) != len(b) {
		return errShortMessage
	}

	m.Header.Type = HeaderType(nlenc.Uint16(b[4:6]))
	m.Header.Flags = HeaderFlags(nlenc.Uint16(b[6:8]))
	m.Header.Sequence = nlenc.Uint32(b[8:12])
	m.Header.PID = nlenc.Uint32(b[12:16])
	m.Data = b[16:]

	return nil
}

// checkMessage checks a single Message for netlink errors.
func checkMessage(m Message) error {
	// NB: All non-nil errors returned from this function *must* be of type
	// OpError in order to maintain the appropriate contract with callers of
	// this package.

	const success = 0

	// Per libnl documentation, only messages that indicate type error can
	// contain error codes:
	// https://www.infradead.org/~tgr/libnl/doc/core.html#core_errmsg.
	//
	// However, at one point, this package checked both done and error for
	// error codes.  Because there was no issue associated with the change,
	// it is unknown whether this change was correct or not.  If you run into
	// a problem with your application because of this change, please file
	// an issue.
	if m.Header.Type != Error {
		return nil
	}

	if len(m.Data) < 4 {
		return newOpError("receive", errShortErrorMessage)
	}

	if c := nlenc.Int32(m.Data[0:4]); c != success {
		// Error code is a negative integer, convert it into an OS-specific raw
		// system call error, but do not wrap with os.NewSyscallError to signify
		// that this error was produced by a netlink message; not a system call.
		return newOpError("receive", newError(-1*int(c)))
	}

	return nil
}
<commit_msg>netlink: set uint16 type on type flags<commit_after>package netlink

import (
	"errors"
	"fmt"

	"github.com/mdlayher/netlink/nlenc"
)

// Flags which may apply to netlink attribute types when communicating with
// certain netlink families.
const (
	Nested       uint16 = 0x8000
	NetByteOrder uint16 = 0x4000

	// attrTypeMask masks off Type bits used for the above flags.
	attrTypeMask uint16 = 0x3fff
)

// Various errors which may occur when attempting to marshal or unmarshal
// a Message to and from its binary form.
var (
	errIncorrectMessageLength = errors.New("netlink message header length incorrect")
	errShortMessage           = errors.New("not enough data to create a netlink message")
	errUnalignedMessage       = errors.New("input data is not properly aligned for netlink message")
)

// HeaderFlags specify flags which may be present in a Header.
type HeaderFlags uint16

const (
	// General netlink communication flags.

	// Request indicates a request to netlink.
	Request HeaderFlags = 1

	// Multi indicates a multi-part message, terminated by Done on the
	// last message.
	Multi HeaderFlags = 2

	// Acknowledge requests that netlink reply with an acknowledgement
	// using Error and, if needed, an error code.
	Acknowledge HeaderFlags = 4

	// Echo requests that netlink echo this request back to the sender.
	Echo HeaderFlags = 8

	// DumpInterrupted indicates that a dump was inconsistent due to a
	// sequence change.
	DumpInterrupted HeaderFlags = 16

	// DumpFiltered indicates that a dump was filtered as requested.
	DumpFiltered HeaderFlags = 32

	// Flags used to retrieve data from netlink.

	// Root requests that netlink return a complete table instead of a
	// single entry.
	Root HeaderFlags = 0x100

	// Match requests that netlink return a list of all matching entries.
	Match HeaderFlags = 0x200

	// Atomic requests that netlink send an atomic snapshot of its entries.
	// Requires CAP_NET_ADMIN or an effective UID of 0.
	Atomic HeaderFlags = 0x400

	// Dump requests that netlink return a complete list of all entries.
	Dump HeaderFlags = Root | Match

	// Flags used to create objects.

	// Replace indicates request replaces an existing matching object.
	Replace HeaderFlags = 0x100

	// Excl indicates request does not replace the object if it already exists.
	Excl HeaderFlags = 0x200

	// Create indicates request creates an object if it doesn't already exist.
	Create HeaderFlags = 0x400

	// Append indicates request adds to the end of the object list.
	Append HeaderFlags = 0x800
)

// String returns the string representation of a HeaderFlags.
func (f HeaderFlags) String() string {
	names := []string{
		"request",
		"multi",
		"acknowledge",
		"echo",
		"dumpinterrupted",
		"dumpfiltered",
	}

	var s string

	left := uint(f)

	for i, name := range names {
		if f&(1<<uint(i)) != 0 {
			if s != "" {
				s += "|"
			}

			s += name

			left ^= (1 << uint(i))
		}
	}

	if s == "" && left == 0 {
		s = "0"
	}

	if left > 0 {
		if s != "" {
			s += "|"
		}
		s += fmt.Sprintf("%#x", left)
	}

	return s
}

// HeaderType specifies the type of a Header.
type HeaderType uint16

const (
	// Noop indicates that no action was taken.
	Noop HeaderType = 0x1

	// Error indicates an error code is present, which is also used to indicate
	// success when the code is 0.
	Error HeaderType = 0x2

	// Done indicates the end of a multi-part message.
	Done HeaderType = 0x3

	// Overrun indicates that data was lost from this message.
	Overrun HeaderType = 0x4
)

// String returns the string representation of a HeaderType.
func (t HeaderType) String() string {
	switch t {
	case Noop:
		return "noop"
	case Error:
		return "error"
	case Done:
		return "done"
	case Overrun:
		return "overrun"
	default:
		return fmt.Sprintf("unknown(%d)", t)
	}
}

// NB: the memory layout of Header and Linux's syscall.NlMsgHdr must be
// exactly the same.  Cannot reorder, change data type, add, or remove fields.
// Named types of the same size (e.g. HeaderFlags is a uint16) are okay.

// A Header is a netlink header.  A Header is sent and received with each
// Message to indicate metadata regarding a Message.
type Header struct {
	// Length of a Message, including this Header.
	Length uint32

	// Contents of a Message.
	Type HeaderType

	// Flags which may be used to modify a request or response.
	Flags HeaderFlags

	// The sequence number of a Message.
	Sequence uint32

	// The process ID of the sending process.
	PID uint32
}

// A Message is a netlink message.  It contains a Header and an arbitrary
// byte payload, which may be decoded using information from the Header.
//
// Data is encoded in the native endianness of the host system.  For easier
// of encoding and decoding of integers, use package nlenc.
type Message struct {
	Header Header
	Data   []byte
}

// MarshalBinary marshals a Message into a byte slice.
func (m Message) MarshalBinary() ([]byte, error) {
	ml := nlmsgAlign(int(m.Header.Length))
	if ml < nlmsgHeaderLen || ml != int(m.Header.Length) {
		return nil, errIncorrectMessageLength
	}

	b := make([]byte, ml)

	nlenc.PutUint32(b[0:4], m.Header.Length)
	nlenc.PutUint16(b[4:6], uint16(m.Header.Type))
	nlenc.PutUint16(b[6:8], uint16(m.Header.Flags))
	nlenc.PutUint32(b[8:12], m.Header.Sequence)
	nlenc.PutUint32(b[12:16], m.Header.PID)
	copy(b[16:], m.Data)

	return b, nil
}

// UnmarshalBinary unmarshals the contents of a byte slice into a Message.
func (m *Message) UnmarshalBinary(b []byte) error {
	if len(b) < nlmsgHeaderLen {
		return errShortMessage
	}
	if len(b) != nlmsgAlign(len(b)) {
		return errUnalignedMessage
	}

	// Don't allow misleading length
	m.Header.Length = nlenc.Uint32(b[0:4])
	if int(m.Header.Length) != len(b) {
		return errShortMessage
	}

	m.Header.Type = HeaderType(nlenc.Uint16(b[4:6]))
	m.Header.Flags = HeaderFlags(nlenc.Uint16(b[6:8]))
	m.Header.Sequence = nlenc.Uint32(b[8:12])
	m.Header.PID = nlenc.Uint32(b[12:16])
	m.Data = b[16:]

	return nil
}

// checkMessage checks a single Message for netlink errors.
func checkMessage(m Message) error {
	// NB: All non-nil errors returned from this function *must* be of type
	// OpError in order to maintain the appropriate contract with callers of
	// this package.

	const success = 0

	// Per libnl documentation, only messages that indicate type error can
	// contain error codes:
	// https://www.infradead.org/~tgr/libnl/doc/core.html#core_errmsg.
	//
	// However, at one point, this package checked both done and error for
	// error codes.  Because there was no issue associated with the change,
	// it is unknown whether this change was correct or not.  If you run into
	// a problem with your application because of this change, please file
	// an issue.
	if m.Header.Type != Error {
		return nil
	}

	if len(m.Data) < 4 {
		return newOpError("receive", errShortErrorMessage)
	}

	if c := nlenc.Int32(m.Data[0:4]); c != success {
		// Error code is a negative integer, convert it into an OS-specific raw
		// system call error, but do not wrap with os.NewSyscallError to signify
		// that this error was produced by a netlink message; not a system call.
		return newOpError("receive", newError(-1*int(c)))
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/sinbad/git-lfs-ssh-serve/Godeps/_workspace/src/github.com/github/git-lfs/lfs"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

func upload(req *lfs.JsonRequest, in io.Reader, out io.Writer, config *Config, path string) *lfs.JsonResponse {
	upreq := lfs.UploadRequest{}
	err := lfs.ExtractStructFromJsonRawMessage(req.Params, &upreq)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	logf("Request %d: Upload %v %d\n", req.Id, upreq.Oid, upreq.Size)
	// Build destination path
	filename, err := mediaPath(upreq.Oid, config, path)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Error determining media path. %v", err))
	}
	startresult := lfs.UploadResponse{}
	_, staterr := os.Stat(filename)
	if staterr != nil && os.IsNotExist(staterr) {
		startresult.OkToSend = true
	}
	// Send start response immediately
	resp, err := lfs.NewJsonResponse(req.Id, startresult)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	err = sendResponse(resp, out)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	if !startresult.OkToSend {
		logf("Upload %d: content already exists for %v\n", req.Id, upreq.Oid)
		return nil
	}

	logf("Upload %d: waiting for content %v\n", req.Id, upreq.Oid)
	// Next from client should be byte stream of exactly the stated number of bytes
	// Now open temp file to write to
	tempf, err := ioutil.TempFile("", "tempupload")
	defer os.Remove(tempf.Name())
	defer tempf.Close()
	n, err := io.CopyN(tempf, in, upreq.Size)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Unable to read data: %v", err.Error()))
	} else if n != upreq.Size {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Received wrong number of bytes %d (expected %d)", n, upreq.Size))
	}

	receivedresult := lfs.UploadCompleteResponse{}
	receivedresult.ReceivedOk = true
	var receiveerr string
	// force close now before defer so we can copy
	err = tempf.Close()
	if err != nil {
		receivedresult.ReceivedOk = false
		receiveerr = fmt.Sprintf("Error when closing temp file: %v", err.Error())
	} else {
		// ensure final directory exists
		ensureDirExists(filepath.Dir(filename), config)
		// Move temp file to final location
		err = os.Rename(tempf.Name(), filename)
		if err != nil {
			receivedresult.ReceivedOk = false
			receiveerr = fmt.Sprintf("Error when closing temp file: %v", err.Error())
		}

	}

	resp, _ = lfs.NewJsonResponse(req.Id, receivedresult)
	if receiveerr != "" {
		logf("Upload %d: error in content for %v: %v\n", req.Id, upreq.Oid, receiveerr)
		resp.Error = receiveerr
	} else {
		logf("Upload %d: content for %v received\n", req.Id, upreq.Oid)
	}

	return resp

}

func ensureDirExists(dir string, cfg *Config) error {
	s, err := os.Stat(dir)
	if err == nil {
		if !s.IsDir() {
			return fmt.Errorf("%v exists but isn't a dir", dir)
		}
	} else {
		// Get permissions from base path & match (or default to user/group write)
		mode := os.FileMode(0775)
		s, err := os.Stat(cfg.BasePath)
		if err == nil {
			mode = s.Mode()
		}
		return os.MkdirAll(dir, mode)
	}
	return nil
}

func uploadCheck(req *lfs.JsonRequest, in io.Reader, out io.Writer, config *Config, path string) *lfs.JsonResponse {
	upreq := lfs.UploadRequest{}
	err := lfs.ExtractStructFromJsonRawMessage(req.Params, &upreq)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	// Build destination path
	filename, err := mediaPath(upreq.Oid, config, path)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Error determining media path. %v", err))
	}
	startresult := lfs.UploadResponse{}
	_, staterr := os.Stat(filename)
	if staterr != nil && os.IsNotExist(staterr) {
		startresult.OkToSend = true
	}
	// Send start response immediately
	resp, err := lfs.NewJsonResponse(req.Id, startresult)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	return resp

}

func downloadCheck(req *lfs.JsonRequest, in io.Reader, out io.Writer, config *Config, path string) *lfs.JsonResponse {
	downreq := lfs.DownloadCheckRequest{}
	err := lfs.ExtractStructFromJsonRawMessage(req.Params, &downreq)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	filename, err := mediaPath(downreq.Oid, config, path)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Problem determining media path: %v", err))
	}
	result := lfs.DownloadCheckResponse{}
	s, err := os.Stat(filename)
	if err == nil {
		// file exists
		result.Size = s.Size()
	} else {
		result.Size = -1
	}
	resp, err := lfs.NewJsonResponse(req.Id, result)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	return resp
}
func download(req *lfs.JsonRequest, in io.Reader, out io.Writer, config *Config, path string) *lfs.JsonResponse {
	downreq := lfs.DownloadRequest{}
	err := lfs.ExtractStructFromJsonRawMessage(req.Params, &downreq)
	if err != nil {
		// Serve() copes with converting this to stderr rather than JSON response
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	filename, err := mediaPath(downreq.Oid, config, path)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Problem determining the media path: %v", err))
	}
	// check size
	s, err := os.Stat(filename)
	if err != nil {
		// file doesn't exist, this should not have been called
		return lfs.NewJsonErrorResponse(req.Id, "File doesn't exist")
	}
	if s.Size() != downreq.Size {
		// This won't work!
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("File sizes disagree (client: %d server: %d)", downreq.Size, s.Size()))
	}

	f, err := os.OpenFile(filename, os.O_RDONLY, 0644)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	defer f.Close()

	n, err := io.Copy(out, f)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Error copying data to output: %v", err.Error()))
	}
	if n != s.Size() {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Amount of data copied disagrees (expected: %d actual: %d)", s.Size(), n))
	}
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Error copying data to output: %v", err.Error()))
	}
	if n != s.Size() {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Amount of data copied disagrees (expected: %d actual: %d)", s.Size(), n))
	}

	// Don't return a response, only response is byte stream above except in error cases
	return nil
}

func batch(req *lfs.JsonRequest, in io.Reader, out io.Writer, config *Config, path string) *lfs.JsonResponse {
	batchreq := lfs.BatchRequest{}
	err := lfs.ExtractStructFromJsonRawMessage(req.Params, &batchreq)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	result := lfs.BatchResponse{}
	for _, o := range batchreq.Objects {
		filename, err := mediaPath(o.Oid, config, path)
		if err != nil {
			return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Problem determining the media path: %v", err))
		}
		resultObj := lfs.BatchResponseObject{Oid: o.Oid}
		s, err := os.Stat(filename)
		if err == nil {
			// file exists
			resultObj.Action = "download"
			resultObj.Size = s.Size()
		} else {
			resultObj.Action = "upload"
			resultObj.Size = o.Size
		}
		result.Results = append(result.Results, resultObj)
	}

	resp, err := lfs.NewJsonResponse(req.Id, result)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	return resp

}

// Store in the same structure as client, just under BasePath
func mediaPath(sha string, config *Config, path string) (string, error) {
	abspath := filepath.Join(config.BasePath, path, sha[0:2], sha[2:4])
	if err := os.MkdirAll(abspath, 0744); err != nil {
		return "", fmt.Errorf("Error trying to create local media directory in '%s': %s", abspath, err)
	}
	return filepath.Join(abspath, sha), nil
}
<commit_msg>A little more diagnostic logging<commit_after>package main

import (
	"fmt"
	"github.com/sinbad/git-lfs-ssh-serve/Godeps/_workspace/src/github.com/github/git-lfs/lfs"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

func upload(req *lfs.JsonRequest, in io.Reader, out io.Writer, config *Config, path string) *lfs.JsonResponse {
	upreq := lfs.UploadRequest{}
	err := lfs.ExtractStructFromJsonRawMessage(req.Params, &upreq)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	logf("Upload %d: requested %v %d\n", req.Id, upreq.Oid, upreq.Size)
	// Build destination path
	filename, err := mediaPath(upreq.Oid, config, path)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Error determining media path. %v", err))
	}
	startresult := lfs.UploadResponse{}
	_, staterr := os.Stat(filename)
	if staterr != nil && os.IsNotExist(staterr) {
		startresult.OkToSend = true
	}
	// Send start response immediately
	resp, err := lfs.NewJsonResponse(req.Id, startresult)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	err = sendResponse(resp, out)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	if !startresult.OkToSend {
		logf("Upload %d: content already exists for %v\n", req.Id, upreq.Oid)
		return nil
	}

	logf("Upload %d: waiting for content %v\n", req.Id, upreq.Oid)
	// Next from client should be byte stream of exactly the stated number of bytes
	// Now open temp file to write to
	tempf, err := ioutil.TempFile("", "tempupload")
	defer os.Remove(tempf.Name())
	defer tempf.Close()
	n, err := io.CopyN(tempf, in, upreq.Size)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Unable to read data: %v", err.Error()))
	} else if n != upreq.Size {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Received wrong number of bytes %d (expected %d)", n, upreq.Size))
	}

	receivedresult := lfs.UploadCompleteResponse{}
	receivedresult.ReceivedOk = true
	var receiveerr string
	// force close now before defer so we can copy
	err = tempf.Close()
	if err != nil {
		receivedresult.ReceivedOk = false
		receiveerr = fmt.Sprintf("Error when closing temp file: %v", err.Error())
	} else {
		// ensure final directory exists
		ensureDirExists(filepath.Dir(filename), config)
		// Move temp file to final location
		err = os.Rename(tempf.Name(), filename)
		if err != nil {
			receivedresult.ReceivedOk = false
			receiveerr = fmt.Sprintf("Error when closing temp file: %v", err.Error())
		}

	}

	resp, _ = lfs.NewJsonResponse(req.Id, receivedresult)
	if receiveerr != "" {
		logf("Upload %d: error in content for %v: %v\n", req.Id, upreq.Oid, receiveerr)
		resp.Error = receiveerr
	} else {
		logf("Upload %d: content for %v received\n", req.Id, upreq.Oid)
	}

	return resp

}

func ensureDirExists(dir string, cfg *Config) error {
	s, err := os.Stat(dir)
	if err == nil {
		if !s.IsDir() {
			return fmt.Errorf("%v exists but isn't a dir", dir)
		}
	} else {
		// Get permissions from base path & match (or default to user/group write)
		mode := os.FileMode(0775)
		s, err := os.Stat(cfg.BasePath)
		if err == nil {
			mode = s.Mode()
		}
		return os.MkdirAll(dir, mode)
	}
	return nil
}

func uploadCheck(req *lfs.JsonRequest, in io.Reader, out io.Writer, config *Config, path string) *lfs.JsonResponse {
	upreq := lfs.UploadRequest{}
	err := lfs.ExtractStructFromJsonRawMessage(req.Params, &upreq)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	logf("UploadCheck %d: %v %d requested\n", req.Id, upreq.Oid, upreq.Size)
	// Build destination path
	filename, err := mediaPath(upreq.Oid, config, path)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Error determining media path. %v", err))
	}
	startresult := lfs.UploadResponse{}
	_, staterr := os.Stat(filename)
	if staterr != nil && os.IsNotExist(staterr) {
		startresult.OkToSend = true
	}
	logf("UploadCheck %d: OK to send %v? %v\n", req.Id, upreq.Oid, startresult.OkToSend)
	// Send start response immediately
	resp, err := lfs.NewJsonResponse(req.Id, startresult)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	return resp

}

func downloadCheck(req *lfs.JsonRequest, in io.Reader, out io.Writer, config *Config, path string) *lfs.JsonResponse {
	downreq := lfs.DownloadCheckRequest{}
	err := lfs.ExtractStructFromJsonRawMessage(req.Params, &downreq)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	logf("DownloadCheck %d: %v requested\n", req.Id, downreq.Oid)
	filename, err := mediaPath(downreq.Oid, config, path)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Problem determining media path: %v", err))
	}
	result := lfs.DownloadCheckResponse{}
	s, err := os.Stat(filename)
	if err == nil {
		// file exists
		result.Size = s.Size()
		logf("DownloadCheck %d: %v response size %d\n", req.Id, downreq.Oid, result.Size)
	} else {
		result.Size = -1
		logf("DownloadCheck %d: %v does not exist\n", req.Id, downreq.Oid)
	}
	resp, err := lfs.NewJsonResponse(req.Id, result)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	return resp
}
func download(req *lfs.JsonRequest, in io.Reader, out io.Writer, config *Config, path string) *lfs.JsonResponse {
	downreq := lfs.DownloadRequest{}
	err := lfs.ExtractStructFromJsonRawMessage(req.Params, &downreq)
	if err != nil {
		// Serve() copes with converting this to stderr rather than JSON response
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	logf("Download %d: %v requested\n", req.Id, downreq.Oid)
	filename, err := mediaPath(downreq.Oid, config, path)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Problem determining the media path: %v", err))
	}
	// check size
	s, err := os.Stat(filename)
	if err != nil {
		// file doesn't exist, this should not have been called
		return lfs.NewJsonErrorResponse(req.Id, "File doesn't exist")
	}
	if s.Size() != downreq.Size {
		// This won't work!
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("File sizes disagree (client: %d server: %d)", downreq.Size, s.Size()))
	}

	f, err := os.OpenFile(filename, os.O_RDONLY, 0644)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	defer f.Close()

	logf("Download %d: sending content for %v\n", req.Id, downreq.Oid)
	n, err := io.Copy(out, f)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Error copying data to output: %v", err.Error()))
	}
	if n != s.Size() {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Amount of data copied disagrees (expected: %d actual: %d)", s.Size(), n))
	}
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Error copying data to output: %v", err.Error()))
	}
	if n != s.Size() {
		return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Amount of data copied disagrees (expected: %d actual: %d)", s.Size(), n))
	}
	logf("Download %d: successfully sent content for %v\n", req.Id, downreq.Oid)

	// Don't return a response, only response is byte stream above except in error cases
	return nil
}

func batch(req *lfs.JsonRequest, in io.Reader, out io.Writer, config *Config, path string) *lfs.JsonResponse {
	batchreq := lfs.BatchRequest{}
	err := lfs.ExtractStructFromJsonRawMessage(req.Params, &batchreq)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	logf("Batch %d: %d objects requested\n", req.Id, len(batchreq.Objects))
	result := lfs.BatchResponse{}
	for _, o := range batchreq.Objects {
		filename, err := mediaPath(o.Oid, config, path)
		if err != nil {
			return lfs.NewJsonErrorResponse(req.Id, fmt.Sprintf("Problem determining the media path: %v", err))
		}
		resultObj := lfs.BatchResponseObject{Oid: o.Oid}
		s, err := os.Stat(filename)
		if err == nil {
			// file exists
			resultObj.Action = "download"
			resultObj.Size = s.Size()
		} else {
			resultObj.Action = "upload"
			resultObj.Size = o.Size
		}
		logf("Batch %d: %v response is %v (%d)\n", req.Id, o.Oid, resultObj.Action, resultObj.Size)
		result.Results = append(result.Results, resultObj)
	}

	resp, err := lfs.NewJsonResponse(req.Id, result)
	if err != nil {
		return lfs.NewJsonErrorResponse(req.Id, err.Error())
	}
	return resp

}

// Store in the same structure as client, just under BasePath
func mediaPath(sha string, config *Config, path string) (string, error) {
	abspath := filepath.Join(config.BasePath, path, sha[0:2], sha[2:4])
	if err := os.MkdirAll(abspath, 0744); err != nil {
		return "", fmt.Errorf("Error trying to create local media directory in '%s': %s", abspath, err)
	}
	return filepath.Join(abspath, sha), nil
}
<|endoftext|>


<commit_before>package updown

import (
	"fmt"
	"net/http"
)

// ResponseTime represents the response times in milliseconds
type ResponseTime struct {
	Under125  int `json:"under125,omitempty"`
	Under250  int `json:"under250,omitempty"`
	Under500  int `json:"under500,omitempty"`
	Under1000 int `json:"under1000,omitempty"`
	Under2000 int `json:"under2000,omitempty"`
	Under4000 int `json:"under4000,omitempty"`
}

// Requests gives statistics about requests made to check the status
type Requests struct {
	Samples      int          `json:"samples,omitempty"`
	Failures     int          `json:"failures,omitempty"`
	Satisfied    int          `json:"satisfied,omitempty"`
	Tolerated    int          `json:"tolerated,omitempty"`
	ResponseTime ResponseTime `json:"by_response_time,omitempty"`
}

// Host represents the host where the check was made
type Host struct {
	IP          string `json:"ip,omitempty"`
	City        string `json:"city,omitempty"`
	Country     string `json:"country,omitempty"`
	CountryCode string `json:"country_code,omitempty"`
}

// Timings represents the amount of time taken by each part of the connection
type Timings struct {
	Redirect   int `json:"redirect,omitempty"`
	NameLookup int `json:"namelookup,omitempty"`
	Connection int `json:"connection,omitempty"`
	Handshake  int `json:"handshake,omitempty"`
	Response   int `json:"response,omitempty"`
	Total      int `json:"total,omitempty"`
}

// MetricItem is basically the core metric
type MetricItem struct {
	Apdex    float64  `json:"apdex,omitempty"`
	Requests Requests `json:"requests,omitempty"`
	Timings  Timings  `json:"timings,omitempty"`
	Host     Host     `json:"host,omitempty"`
}

// Metrics represents multiple metrics
type Metrics map[string]MetricItem

// MetricService interacts with the metrics section of the API
type MetricService struct {
	client *Client
}

// List lists metrics available for a check identified by a taken, grouped by the given group
// (host|time) over a period
func (s *MetricService) List(token, group, from, to string) (Metrics, *http.Response, error) {
	path := fmt.Sprintf(pathForToken(token)+"/metrics?group=%s", group)
	// Optional from and to parameters
	if from != "" {
		path += "&from=" + from
	}
	if to != "" {
		path += "&to=" + to
	}
	req, err := s.client.NewRequest("GET", path, nil)
	if err != nil {
		return nil, nil, err
	}

	var res Metrics
	resp, err := s.client.Do(req, &res)
	if err != nil {
		return nil, resp, err
	}

	return res, resp, err
}
<commit_msg>Setting query strings manually was clearly a mistake<commit_after>package updown

import (
	"net/http"
	"net/url"
)

// ResponseTime represents the response times in milliseconds
type ResponseTime struct {
	Under125  int `json:"under125,omitempty"`
	Under250  int `json:"under250,omitempty"`
	Under500  int `json:"under500,omitempty"`
	Under1000 int `json:"under1000,omitempty"`
	Under2000 int `json:"under2000,omitempty"`
	Under4000 int `json:"under4000,omitempty"`
}

// Requests gives statistics about requests made to check the status
type Requests struct {
	Samples      int          `json:"samples,omitempty"`
	Failures     int          `json:"failures,omitempty"`
	Satisfied    int          `json:"satisfied,omitempty"`
	Tolerated    int          `json:"tolerated,omitempty"`
	ResponseTime ResponseTime `json:"by_response_time,omitempty"`
}

// Host represents the host where the check was made
type Host struct {
	IP          string `json:"ip,omitempty"`
	City        string `json:"city,omitempty"`
	Country     string `json:"country,omitempty"`
	CountryCode string `json:"country_code,omitempty"`
}

// Timings represents the amount of time taken by each part of the connection
type Timings struct {
	Redirect   int `json:"redirect,omitempty"`
	NameLookup int `json:"namelookup,omitempty"`
	Connection int `json:"connection,omitempty"`
	Handshake  int `json:"handshake,omitempty"`
	Response   int `json:"response,omitempty"`
	Total      int `json:"total,omitempty"`
}

// MetricItem is basically the core metric
type MetricItem struct {
	Apdex    float64  `json:"apdex,omitempty"`
	Requests Requests `json:"requests,omitempty"`
	Timings  Timings  `json:"timings,omitempty"`
	Host     Host     `json:"host,omitempty"`
}

// Metrics represents multiple metrics
type Metrics map[string]MetricItem

// MetricService interacts with the metrics section of the API
type MetricService struct {
	client *Client
}

// List lists metrics available for a check identified by a taken, grouped by the given group
// (host|time) over a period
func (s *MetricService) List(token, group, from, to string) (Metrics, *http.Response, error) {
	u, _ := url.Parse(pathForToken(token) + "/metrics")
	q := u.Query()
	q.Add("group", group)

	// Optional from and to parameters
	if from != "" {
		q.Add("from", from)
	}
	if to != "" {
		q.Add("to", to)
	}
	u.RawQuery = q.Encode()

	req, err := s.client.NewRequest("GET", u.String(), nil)
	if err != nil {
		return nil, nil, err
	}

	var res Metrics
	resp, err := s.client.Do(req, &res)
	if err != nil {
		return nil, resp, err
	}

	return res, resp, err
}
<|endoftext|>


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

package main

import (
	"database/sql"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"time"
)

var cmdMigrate = &Command{
	UsageLine: "migrate [Command]",
	Short:     "run database migrations",
	Long: `
bee migrate
    run all outstanding migrations
    -driver: [mysql | postgresql | sqlite], the default is mysql
    -conn:   the connection string used by the driver, the default is root:@tcp(127.0.0.1:3306)/test

bee migrate rollback
    rollback the last migration operation
    -driver: [mysql | postgresql | sqlite], the default is mysql
    -conn:   the connection string used by the driver, the default is root:@tcp(127.0.0.1:3306)/test

bee migrate reset
    rollback all migrations
    -driver: [mysql | postgresql | sqlite], the default is mysql
    -conn:   the connection string used by the driver, the default is root:@tcp(127.0.0.1:3306)/test

bee migrate refresh
    rollback all migrations and run them all again
    -driver: [mysql | postgresql | sqlite], the default is mysql
    -conn:   the connection string used by the driver, the default is root:@tcp(127.0.0.1:3306)/test
`,
}

var mDriver docValue
var mConn docValue

func init() {
	cmdMigrate.Run = runMigration
	cmdMigrate.Flag.Var(&mDriver, "driver", "database driver: mysql, postgresql, etc.")
	cmdMigrate.Flag.Var(&mConn, "conn", "connection string used by the driver to connect to a database instance")
}

// runMigration is the entry point for starting a migration
func runMigration(cmd *Command, args []string) {
	crupath, _ := os.Getwd()

	gopath := os.Getenv("GOPATH")
	Debugf("gopath:%s", gopath)
	if gopath == "" {
		ColorLog("[ERRO] $GOPATH not found\n")
		ColorLog("[HINT] Set $GOPATH in your environment vairables\n")
		os.Exit(2)
	}
	// load config
	err := loadConfig()
	if err != nil {
		ColorLog("[ERRO] Fail to parse bee.json[ %s ]\n", err)
	}
	// getting command line arguments
	if len(args) != 0 {
		cmd.Flag.Parse(args[1:])
	}
	if mDriver == "" {
		mDriver = docValue(conf.Database.Driver)
		if mDriver == "" {
			mDriver = "mysql"
		}
	}
	if mConn == "" {
		mConn = docValue(conf.Database.Conn)
		if mConn == "" {
			mConn = "root:@tcp(127.0.0.1:3306)/test"
		}
	}
	ColorLog("[INFO] Using '%s' as 'driver'\n", mDriver)
	ColorLog("[INFO] Using '%s' as 'conn'\n", mConn)
	driverStr, connStr := string(mDriver), string(mConn)
	if len(args) == 0 {
		// run all outstanding migrations
		ColorLog("[INFO] Running all outstanding migrations\n")
		migrateUpdate(crupath, driverStr, connStr)
	} else {
		mcmd := args[0]
		switch mcmd {
		case "rollback":
			ColorLog("[INFO] Rolling back the last migration operation\n")
			migrateRollback(crupath, driverStr, connStr)
		case "reset":
			ColorLog("[INFO] Reseting all migrations\n")
			migrateReset(crupath, driverStr, connStr)
		case "refresh":
			ColorLog("[INFO] Refreshing all migrations\n")
			migrateReset(crupath, driverStr, connStr)
		default:
			ColorLog("[ERRO] Command is missing\n")
			os.Exit(2)
		}
	}
	ColorLog("[SUCC] Migration successful!\n")
}

// migrateUpdate does the schema update
func migrateUpdate(crupath, driver, connStr string) {
	migrate("upgrade", crupath, driver, connStr)
}

// migrateRollback rolls back the latest migration
func migrateRollback(crupath, driver, connStr string) {
	migrate("rollback", crupath, driver, connStr)
}

// migrateReset rolls back all migrations
func migrateReset(crupath, driver, connStr string) {
	migrate("reset", crupath, driver, connStr)
}

// migrationRefresh rolls back all migrations and start over again
func migrateRefresh(crupath, driver, connStr string) {
	migrate("refresh", crupath, driver, connStr)
}

// migrate generates source code, build it, and invoke the binary who does the actual migration
func migrate(goal, crupath, driver, connStr string) {
	dir := path.Join(crupath, "database", "migrations")
	binary := "m"
	source := binary + ".go"
	// connect to database
	db, err := sql.Open(driver, connStr)
	if err != nil {
		ColorLog("[ERRO] Could not connect to %s: %s\n", driver, connStr)
		os.Exit(2)
	}
	defer db.Close()
	checkForSchemaUpdateTable(db)
	latestName, latestTime := getLatestMigration(db)
	writeMigrationSourceFile(dir, source, driver, connStr, latestTime, latestName, goal)
	buildMigrationBinary(dir, binary)
	runMigrationBinary(dir, binary)
	removeTempFile(dir, source)
	removeTempFile(dir, binary)
}

// checkForSchemaUpdateTable checks the existence of migrations table.
// It checks for the proper table structures and creates the table using MYSQL_MIGRATION_DDL if it does not exist.
func checkForSchemaUpdateTable(db *sql.DB) {
	if rows, err := db.Query("SHOW TABLES LIKE 'migrations'"); err != nil {
		ColorLog("[ERRO] Could not show migrations table: %s\n", err)
		os.Exit(2)
	} else if !rows.Next() {
		// no migrations table, create anew
		ColorLog("[INFO] Creating 'migrations' table...\n")
		if _, err := db.Query(MYSQL_MIGRATION_DDL); err != nil {
			ColorLog("[ERRO] Could not create migrations table: %s\n", err)
			os.Exit(2)
		}
	}
	// checking that migrations table schema are expected
	if rows, err := db.Query("DESC migrations"); err != nil {
		ColorLog("[ERRO] Could not show columns of migrations table: %s\n", err)
		os.Exit(2)
	} else {
		for rows.Next() {
			var fieldBytes, typeBytes, nullBytes, keyBytes, defaultBytes, extraBytes []byte
			if err := rows.Scan(&fieldBytes, &typeBytes, &nullBytes, &keyBytes, &defaultBytes, &extraBytes); err != nil {
				ColorLog("[ERRO] Could not read column information: %s\n", err)
				os.Exit(2)
			}
			fieldStr, typeStr, nullStr, keyStr, defaultStr, extraStr :=
				string(fieldBytes), string(typeBytes), string(nullBytes), string(keyBytes), string(defaultBytes), string(extraBytes)
			if fieldStr == "id_migration" {
				if keyStr != "PRI" || extraStr != "auto_increment" {
					ColorLog("[ERRO] Column migration.id_migration type mismatch: KEY: %s, EXTRA: %s\n", keyStr, extraStr)
					ColorLog("[HINT] Expecting KEY: PRI, EXTRA: auto_increment\n")
					os.Exit(2)
				}
			} else if fieldStr == "name" {
				if !strings.HasPrefix(typeStr, "varchar") || nullStr != "YES" {
					ColorLog("[ERRO] Column migration.name type mismatch: TYPE: %s, NULL: %s\n", typeStr, nullStr)
					ColorLog("[HINT] Expecting TYPE: varchar, NULL: YES\n")
					os.Exit(2)
				}

			} else if fieldStr == "created_at" {
				if typeStr != "timestamp" || defaultStr != "CURRENT_TIMESTAMP" {
					ColorLog("[ERRO] Column migration.timestamp type mismatch: TYPE: %s, DEFAULT: %s\n", typeStr, defaultStr)
					ColorLog("[HINT] Expecting TYPE: timestamp, DEFAULT: CURRENT_TIMESTAMP\n")
					os.Exit(2)
				}
			}
		}
	}
}

// getLatestMigration retrives latest migration with status 'update'
func getLatestMigration(db *sql.DB) (file string, createdAt int64) {
	sql := "SELECT name, created_at FROM migrations where status = 'update' ORDER BY id_migration DESC LIMIT 1"
	if rows, err := db.Query(sql); err != nil {
		ColorLog("[ERRO] Could not retrieve migrations: %s\n", err)
		os.Exit(2)
	} else {
		var createdAtStr string
		if rows.Next() {
			if err := rows.Scan(&file, &createdAtStr); err != nil {
				ColorLog("[ERRO] Could not read migrations in database: %s\n", err)
				os.Exit(2)
			}
			if t, err := time.Parse("2006-01-02 15:04:05", createdAtStr); err != nil {
				ColorLog("[ERRO] Could not parse time: %s\n", err)
				os.Exit(2)
			} else {
				createdAt = t.Unix()
			}
		} else {
			file, createdAt = "", 0
		}
	}
	return
}

// writeMigrationSourceFile create the source file based on MIGRATION_MAIN_TPL
func writeMigrationSourceFile(dir, source, driver, connStr string, latestTime int64, latestName string, task string) {
	changeDir(dir)
	if f, err := os.OpenFile(source, os.O_CREATE|os.O_EXCL|os.O_RDWR, 0666); err != nil {
		ColorLog("[ERRO] Could not create file: %s\n", err)
		os.Exit(2)
	} else {
		content := strings.Replace(MIGRATION_MAIN_TPL, "{{DBDriver}}", driver, -1)
		content = strings.Replace(content, "{{ConnStr}}", connStr, -1)
		content = strings.Replace(content, "{{LatestTime}}", strconv.FormatInt(latestTime, 10), -1)
		content = strings.Replace(content, "{{LatestName}}", latestName, -1)
		content = strings.Replace(content, "{{Task}}", task, -1)
		if _, err := f.WriteString(content); err != nil {
			ColorLog("[ERRO] Could not write to file: %s\n", err)
			os.Exit(2)
		}
		f.Close()
	}
}

// buildMigrationBinary changes directory to database/migrations folder and go-build the source
func buildMigrationBinary(dir, binary string) {
	changeDir(dir)
	cmd := exec.Command("go", "build", "-o", binary)
	if out, err := cmd.CombinedOutput(); err != nil {
		ColorLog("[ERRO] Could not build migration binary: %s\n", err)
		formatShellErrOutput(string(out))
		removeTempFile(dir, binary)
		removeTempFile(dir, binary+".go")
		os.Exit(2)
	}
}

// runMigrationBinary runs the migration program who does the actual work
func runMigrationBinary(dir, binary string) {
	changeDir(dir)
	cmd := exec.Command("./" + binary)
	if out, err := cmd.CombinedOutput(); err != nil {
		formatShellOutput(string(out))
		ColorLog("[ERRO] Could not run migration binary: %s\n", err)
		removeTempFile(dir, binary)
		removeTempFile(dir, binary+".go")
		os.Exit(2)
	} else {
		formatShellOutput(string(out))
	}
}

// changeDir changes working directory to dir.
// It exits the system when encouter an error
func changeDir(dir string) {
	if err := os.Chdir(dir); err != nil {
		ColorLog("[ERRO] Could not find migration directory: %s\n", err)
		os.Exit(2)
	}
}

// removeTempFile removes a file in dir
func removeTempFile(dir, file string) {
	os.Chdir(dir)
	if err := os.Remove(file); err != nil {
		ColorLog("[ERRO] Could not remove temporary migration files: %s\n", err)
		os.Exit(2)
	}
}

// formatShellErrOutput formats the error shell output
func formatShellErrOutput(o string) {
	for _, line := range strings.Split(o, "\n") {
		if line != "" {
			ColorLog("[ERRO] -| %s\n", line)
		}
	}
}

// formatShellOutput formats the normal shell output
func formatShellOutput(o string) {
	for _, line := range strings.Split(o, "\n") {
		if line != "" {
			ColorLog("[INFO] -| %s\n", line)
		}
	}
}

const (
	MIGRATION_MAIN_TPL = `package main

import(
	"os"

	"github.com/astaxie/beego/orm"
	"github.com/astaxie/beego/migration"

	_ "github.com/go-sql-driver/mysql"
)

func init(){
	orm.RegisterDataBase("default", "{{DBDriver}}","{{ConnStr}}")
}

func main(){
	task := "{{Task}}"
	switch task {
	case "upgrade":
		if err := migration.Upgrade({{LatestTime}}); err != nil {
			os.Exit(2)
		}
	case "rollback":
		if err := migration.Rollback("{{LatestName}}"); err != nil {
			os.Exit(2)
		}
	case "reset":
		if err := migration.Reset(); err != nil {
			os.Exit(2)
		}
	case "refresh":
		if err := migration.Refresh(); err != nil {
			os.Exit(2)
		}
	}
}

`
	MYSQL_MIGRATION_DDL = `
CREATE TABLE migrations (
	id_migration int(10) unsigned NOT NULL AUTO_INCREMENT COMMENT 'surrogate key',
	name varchar(255) DEFAULT NULL COMMENT 'migration name, unique',
	created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT 'date migrated or rolled back',
	statements longtext COMMENT 'SQL statements for this migration',
	rollback_statements longtext COMMENT 'SQL statment for rolling back migration',
	status ENUM('update', 'rollback') COMMENT 'update indicates it is a normal migration while rollback means this migration is rolled back',
	PRIMARY KEY (id_migration)
) ENGINE=InnoDB DEFAULT CHARSET=utf8
`
)
<commit_msg>bugfix: call refresh function when refreshing<commit_after>// Copyright 2013 bee authors
//
// Licensed under the Apache License, Version 2.0 (the "License"): you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

package main

import (
	"database/sql"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"time"
)

var cmdMigrate = &Command{
	UsageLine: "migrate [Command]",
	Short:     "run database migrations",
	Long: `
bee migrate
    run all outstanding migrations
    -driver: [mysql | postgresql | sqlite], the default is mysql
    -conn:   the connection string used by the driver, the default is root:@tcp(127.0.0.1:3306)/test

bee migrate rollback
    rollback the last migration operation
    -driver: [mysql | postgresql | sqlite], the default is mysql
    -conn:   the connection string used by the driver, the default is root:@tcp(127.0.0.1:3306)/test

bee migrate reset
    rollback all migrations
    -driver: [mysql | postgresql | sqlite], the default is mysql
    -conn:   the connection string used by the driver, the default is root:@tcp(127.0.0.1:3306)/test

bee migrate refresh
    rollback all migrations and run them all again
    -driver: [mysql | postgresql | sqlite], the default is mysql
    -conn:   the connection string used by the driver, the default is root:@tcp(127.0.0.1:3306)/test
`,
}

var mDriver docValue
var mConn docValue

func init() {
	cmdMigrate.Run = runMigration
	cmdMigrate.Flag.Var(&mDriver, "driver", "database driver: mysql, postgresql, etc.")
	cmdMigrate.Flag.Var(&mConn, "conn", "connection string used by the driver to connect to a database instance")
}

// runMigration is the entry point for starting a migration
func runMigration(cmd *Command, args []string) {
	crupath, _ := os.Getwd()

	gopath := os.Getenv("GOPATH")
	Debugf("gopath:%s", gopath)
	if gopath == "" {
		ColorLog("[ERRO] $GOPATH not found\n")
		ColorLog("[HINT] Set $GOPATH in your environment vairables\n")
		os.Exit(2)
	}
	// load config
	err := loadConfig()
	if err != nil {
		ColorLog("[ERRO] Fail to parse bee.json[ %s ]\n", err)
	}
	// getting command line arguments
	if len(args) != 0 {
		cmd.Flag.Parse(args[1:])
	}
	if mDriver == "" {
		mDriver = docValue(conf.Database.Driver)
		if mDriver == "" {
			mDriver = "mysql"
		}
	}
	if mConn == "" {
		mConn = docValue(conf.Database.Conn)
		if mConn == "" {
			mConn = "root:@tcp(127.0.0.1:3306)/test"
		}
	}
	ColorLog("[INFO] Using '%s' as 'driver'\n", mDriver)
	ColorLog("[INFO] Using '%s' as 'conn'\n", mConn)
	driverStr, connStr := string(mDriver), string(mConn)
	if len(args) == 0 {
		// run all outstanding migrations
		ColorLog("[INFO] Running all outstanding migrations\n")
		migrateUpdate(crupath, driverStr, connStr)
	} else {
		mcmd := args[0]
		switch mcmd {
		case "rollback":
			ColorLog("[INFO] Rolling back the last migration operation\n")
			migrateRollback(crupath, driverStr, connStr)
		case "reset":
			ColorLog("[INFO] Reseting all migrations\n")
			migrateReset(crupath, driverStr, connStr)
		case "refresh":
			ColorLog("[INFO] Refreshing all migrations\n")
			migrateRefresh(crupath, driverStr, connStr)
		default:
			ColorLog("[ERRO] Command is missing\n")
			os.Exit(2)
		}
	}
	ColorLog("[SUCC] Migration successful!\n")
}

// migrateUpdate does the schema update
func migrateUpdate(crupath, driver, connStr string) {
	migrate("upgrade", crupath, driver, connStr)
}

// migrateRollback rolls back the latest migration
func migrateRollback(crupath, driver, connStr string) {
	migrate("rollback", crupath, driver, connStr)
}

// migrateReset rolls back all migrations
func migrateReset(crupath, driver, connStr string) {
	migrate("reset", crupath, driver, connStr)
}

// migrationRefresh rolls back all migrations and start over again
func migrateRefresh(crupath, driver, connStr string) {
	migrate("refresh", crupath, driver, connStr)
}

// migrate generates source code, build it, and invoke the binary who does the actual migration
func migrate(goal, crupath, driver, connStr string) {
	dir := path.Join(crupath, "database", "migrations")
	binary := "m"
	source := binary + ".go"
	// connect to database
	db, err := sql.Open(driver, connStr)
	if err != nil {
		ColorLog("[ERRO] Could not connect to %s: %s\n", driver, connStr)
		os.Exit(2)
	}
	defer db.Close()
	checkForSchemaUpdateTable(db)
	latestName, latestTime := getLatestMigration(db)
	writeMigrationSourceFile(dir, source, driver, connStr, latestTime, latestName, goal)
	buildMigrationBinary(dir, binary)
	runMigrationBinary(dir, binary)
	removeTempFile(dir, source)
	removeTempFile(dir, binary)
}

// checkForSchemaUpdateTable checks the existence of migrations table.
// It checks for the proper table structures and creates the table using MYSQL_MIGRATION_DDL if it does not exist.
func checkForSchemaUpdateTable(db *sql.DB) {
	if rows, err := db.Query("SHOW TABLES LIKE 'migrations'"); err != nil {
		ColorLog("[ERRO] Could not show migrations table: %s\n", err)
		os.Exit(2)
	} else if !rows.Next() {
		// no migrations table, create anew
		ColorLog("[INFO] Creating 'migrations' table...\n")
		if _, err := db.Query(MYSQL_MIGRATION_DDL); err != nil {
			ColorLog("[ERRO] Could not create migrations table: %s\n", err)
			os.Exit(2)
		}
	}
	// checking that migrations table schema are expected
	if rows, err := db.Query("DESC migrations"); err != nil {
		ColorLog("[ERRO] Could not show columns of migrations table: %s\n", err)
		os.Exit(2)
	} else {
		for rows.Next() {
			var fieldBytes, typeBytes, nullBytes, keyBytes, defaultBytes, extraBytes []byte
			if err := rows.Scan(&fieldBytes, &typeBytes, &nullBytes, &keyBytes, &defaultBytes, &extraBytes); err != nil {
				ColorLog("[ERRO] Could not read column information: %s\n", err)
				os.Exit(2)
			}
			fieldStr, typeStr, nullStr, keyStr, defaultStr, extraStr :=
				string(fieldBytes), string(typeBytes), string(nullBytes), string(keyBytes), string(defaultBytes), string(extraBytes)
			if fieldStr == "id_migration" {
				if keyStr != "PRI" || extraStr != "auto_increment" {
					ColorLog("[ERRO] Column migration.id_migration type mismatch: KEY: %s, EXTRA: %s\n", keyStr, extraStr)
					ColorLog("[HINT] Expecting KEY: PRI, EXTRA: auto_increment\n")
					os.Exit(2)
				}
			} else if fieldStr == "name" {
				if !strings.HasPrefix(typeStr, "varchar") || nullStr != "YES" {
					ColorLog("[ERRO] Column migration.name type mismatch: TYPE: %s, NULL: %s\n", typeStr, nullStr)
					ColorLog("[HINT] Expecting TYPE: varchar, NULL: YES\n")
					os.Exit(2)
				}

			} else if fieldStr == "created_at" {
				if typeStr != "timestamp" || defaultStr != "CURRENT_TIMESTAMP" {
					ColorLog("[ERRO] Column migration.timestamp type mismatch: TYPE: %s, DEFAULT: %s\n", typeStr, defaultStr)
					ColorLog("[HINT] Expecting TYPE: timestamp, DEFAULT: CURRENT_TIMESTAMP\n")
					os.Exit(2)
				}
			}
		}
	}
}

// getLatestMigration retrives latest migration with status 'update'
func getLatestMigration(db *sql.DB) (file string, createdAt int64) {
	sql := "SELECT name, created_at FROM migrations where status = 'update' ORDER BY id_migration DESC LIMIT 1"
	if rows, err := db.Query(sql); err != nil {
		ColorLog("[ERRO] Could not retrieve migrations: %s\n", err)
		os.Exit(2)
	} else {
		var createdAtStr string
		if rows.Next() {
			if err := rows.Scan(&file, &createdAtStr); err != nil {
				ColorLog("[ERRO] Could not read migrations in database: %s\n", err)
				os.Exit(2)
			}
			if t, err := time.Parse("2006-01-02 15:04:05", createdAtStr); err != nil {
				ColorLog("[ERRO] Could not parse time: %s\n", err)
				os.Exit(2)
			} else {
				createdAt = t.Unix()
			}
		} else {
			file, createdAt = "", 0
		}
	}
	return
}

// writeMigrationSourceFile create the source file based on MIGRATION_MAIN_TPL
func writeMigrationSourceFile(dir, source, driver, connStr string, latestTime int64, latestName string, task string) {
	changeDir(dir)
	if f, err := os.OpenFile(source, os.O_CREATE|os.O_EXCL|os.O_RDWR, 0666); err != nil {
		ColorLog("[ERRO] Could not create file: %s\n", err)
		os.Exit(2)
	} else {
		content := strings.Replace(MIGRATION_MAIN_TPL, "{{DBDriver}}", driver, -1)
		content = strings.Replace(content, "{{ConnStr}}", connStr, -1)
		content = strings.Replace(content, "{{LatestTime}}", strconv.FormatInt(latestTime, 10), -1)
		content = strings.Replace(content, "{{LatestName}}", latestName, -1)
		content = strings.Replace(content, "{{Task}}", task, -1)
		if _, err := f.WriteString(content); err != nil {
			ColorLog("[ERRO] Could not write to file: %s\n", err)
			os.Exit(2)
		}
		f.Close()
	}
}

// buildMigrationBinary changes directory to database/migrations folder and go-build the source
func buildMigrationBinary(dir, binary string) {
	changeDir(dir)
	cmd := exec.Command("go", "build", "-o", binary)
	if out, err := cmd.CombinedOutput(); err != nil {
		ColorLog("[ERRO] Could not build migration binary: %s\n", err)
		formatShellErrOutput(string(out))
		removeTempFile(dir, binary)
		removeTempFile(dir, binary+".go")
		os.Exit(2)
	}
}

// runMigrationBinary runs the migration program who does the actual work
func runMigrationBinary(dir, binary string) {
	changeDir(dir)
	cmd := exec.Command("./" + binary)
	if out, err := cmd.CombinedOutput(); err != nil {
		formatShellOutput(string(out))
		ColorLog("[ERRO] Could not run migration binary: %s\n", err)
		removeTempFile(dir, binary)
		removeTempFile(dir, binary+".go")
		os.Exit(2)
	} else {
		formatShellOutput(string(out))
	}
}

// changeDir changes working directory to dir.
// It exits the system when encouter an error
func changeDir(dir string) {
	if err := os.Chdir(dir); err != nil {
		ColorLog("[ERRO] Could not find migration directory: %s\n", err)
		os.Exit(2)
	}
}

// removeTempFile removes a file in dir
func removeTempFile(dir, file string) {
	os.Chdir(dir)
	if err := os.Remove(file); err != nil {
		ColorLog("[ERRO] Could not remove temporary migration files: %s\n", err)
		os.Exit(2)
	}
}

// formatShellErrOutput formats the error shell output
func formatShellErrOutput(o string) {
	for _, line := range strings.Split(o, "\n") {
		if line != "" {
			ColorLog("[ERRO] -| %s\n", line)
		}
	}
}

// formatShellOutput formats the normal shell output
func formatShellOutput(o string) {
	for _, line := range strings.Split(o, "\n") {
		if line != "" {
			ColorLog("[INFO] -| %s\n", line)
		}
	}
}

const (
	MIGRATION_MAIN_TPL = `package main

import(
	"os"

	"github.com/astaxie/beego/orm"
	"github.com/astaxie/beego/migration"

	_ "github.com/go-sql-driver/mysql"
)

func init(){
	orm.RegisterDataBase("default", "{{DBDriver}}","{{ConnStr}}")
}

func main(){
	task := "{{Task}}"
	switch task {
	case "upgrade":
		if err := migration.Upgrade({{LatestTime}}); err != nil {
			os.Exit(2)
		}
	case "rollback":
		if err := migration.Rollback("{{LatestName}}"); err != nil {
			os.Exit(2)
		}
	case "reset":
		if err := migration.Reset(); err != nil {
			os.Exit(2)
		}
	case "refresh":
		if err := migration.Refresh(); err != nil {
			os.Exit(2)
		}
	}
}

`
	MYSQL_MIGRATION_DDL = `
CREATE TABLE migrations (
	id_migration int(10) unsigned NOT NULL AUTO_INCREMENT COMMENT 'surrogate key',
	name varchar(255) DEFAULT NULL COMMENT 'migration name, unique',
	created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT 'date migrated or rolled back',
	statements longtext COMMENT 'SQL statements for this migration',
	rollback_statements longtext COMMENT 'SQL statment for rolling back migration',
	status ENUM('update', 'rollback') COMMENT 'update indicates it is a normal migration while rollback means this migration is rolled back',
	PRIMARY KEY (id_migration)
) ENGINE=InnoDB DEFAULT CHARSET=utf8
`
)
<|endoftext|>


<commit_before>package transport

import (
	"errors"
	"net"

	"github.com/elastic/beats/libbeat/logp"
)

type Dialer interface {
	Dial(network, address string) (net.Conn, error)
}

type DialerFunc func(netowrk, address string) (net.Conn, error)

var (
	ErrNotConnected = errors.New("client is not connected")

	debugf = logp.MakeDebug("transport")
)

func (d DialerFunc) Dial(network, address string) (net.Conn, error) {
	return d(network, address)
}

func Dial(c *Config, network, address string) (net.Conn, error) {
	d, err := MakeDialer(c)
	if err != nil {
		return nil, err
	}
	return d.Dial(network, address)
}
<commit_msg>fix typo (#1563)<commit_after>package transport

import (
	"errors"
	"net"

	"github.com/elastic/beats/libbeat/logp"
)

type Dialer interface {
	Dial(network, address string) (net.Conn, error)
}

type DialerFunc func(network, address string) (net.Conn, error)

var (
	ErrNotConnected = errors.New("client is not connected")

	debugf = logp.MakeDebug("transport")
)

func (d DialerFunc) Dial(network, address string) (net.Conn, error) {
	return d(network, address)
}

func Dial(c *Config, network, address string) (net.Conn, error) {
	d, err := MakeDialer(c)
	if err != nil {
		return nil, err
	}
	return d.Dial(network, address)
}
<|endoftext|>