<commit_before>package main

import (
	"errors"
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/Necroforger/dgwidgets"
	"github.com/Strum355/ytdl"
	"github.com/bwmarrin/discordgo"
	"github.com/jonas747/dca"
)

const (
	stdURL   = "https://www.youtube.com/watch"
	shortURL = "https://youtu.be/"
	embedURL = "https://www.youtube.com/embed/"
)

func init() {
	newCommand("yt", 0, false, false, msgYoutube).setHelp("Args: [play,stop] [url]\n\nWork In Progress!!! Play music from Youtube straight to your Discord Server!\n\n" +
		"Example 1: `!owo yt play https://www.youtube.com/watch?v=MvLdxtICOIY`\n" +
		"Example 2: `!owo yt stop`\n\nSubCommands:\nplay\nstop\nlist, queue, songs\npause\nresume, unpause\nskip, next").add()
}

func msgYoutube(s *discordgo.Session, m *discordgo.MessageCreate, msglist []string) {
	if len(msglist) == 1 {
		return
	}

	switch msglist[1] {
	case "play":
		addToQueue(s, m, msglist[2:])
	case "stop":
		stopQueue(s, m)
	case "list", "queue", "songs":
		listQueue(s, m)
	case "pause":
		pauseQueue(s, m)
	case "resume", "unpause":
		unpauseQueue(s, m)
	case "skip", "next":
		skipSong(s, m)
	}
}

func addToQueue(s *discordgo.Session, m *discordgo.MessageCreate, msglist []string) {
	if len(msglist) == 0 {
		return
	}

	guild, err := guildDetails(m.ChannelID, "", s)
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "There was a problem adding to queue :( please try again")
		return
	}

	srvr, ok := sMap.Server[guild.ID]
	if !ok {
		s.ChannelMessageSend(m.ChannelID, "An error occured that really shouldn't have happened...")
		log.Error("not in server map?", guild.ID)
		return
	}

	if srvr.VoiceInst == nil {
		srvr.newVoiceInstance()
	}

	srvr.VoiceInst.Lock()
	defer srvr.VoiceInst.Unlock()

	url := msglist[0]

	if !strings.HasPrefix(url, stdURL) && !strings.HasPrefix(url, shortURL) && !strings.HasPrefix(url, embedURL) {
		s.ChannelMessageSend(m.ChannelID, "Please make sure the URL is a valid YouTube URL. If I got this wrong, please let my creator know~")
		return
	}

	for _, vs := range guild.VoiceStates {
		if vs.UserID == m.Author.ID {
			if vs.ChannelID == srvr.VoiceInst.ChannelID || !srvr.VoiceInst.Playing {
				vid, err := getVideoInfo(url, s, m)
				if err != nil {
					return
				}

				vc, err := s.ChannelVoiceJoin(guild.ID, vs.ChannelID, false, true)
				if err != nil {
					s.ChannelMessageSend(m.ChannelID, "Error joining voice channel")
					log.Error("error joining voice channel", err)
					return
				}

				srvr.addSong(song{
					URL:      url,
					Name:     vid.Title,
					Duration: vid.Duration,
					Image:    vid.GetThumbnailURL(ytdl.ThumbnailQualityMedium).String(),
				})

				s.ChannelMessageSend(m.ChannelID, "Added "+vid.Title+" to the queue!")

				if !srvr.VoiceInst.Playing {
					srvr.VoiceInst.VoiceCon = vc
					srvr.VoiceInst.Playing = true
					srvr.VoiceInst.ChannelID = vc.ChannelID
					go play(s, m, srvr, vc)
				}
				return
			}
		}
	}

	s.ChannelMessageSend(m.ChannelID, "Need to be in a voice channel!")
}

func getVideoInfo(url string, s *discordgo.Session, m *discordgo.MessageCreate) (*ytdl.VideoInfo, error) {
	vid, err := ytdl.GetVideoInfo(url)
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "Error getting video info")
		log.Error("error getting video info", err)
		return nil, err
	}
	return vid, nil
}

func play(s *discordgo.Session, m *discordgo.MessageCreate, srvr *server, vc *discordgo.VoiceConnection) {
	if srvr.queueLength() == 0 {
		srvr.youtubeCleanup()
		s.ChannelMessageSend(m.ChannelID, "🔇 Done queue!")
		return
	}

	srvr.VoiceInst.Lock()

	vid, err := getVideoInfo(srvr.nextSong().URL, s, m)
	if err != nil {
		srvr.VoiceInst.Unlock()
		return
	}

	reader, writer := io.Pipe()
	defer reader.Close()

	formats := vid.Formats.Best(ytdl.FormatAudioBitrateKey)
	if len(formats) > 0 {
		go func() {
			defer writer.Close()
			if err := vid.Download(formats[0], writer); err != nil && err != io.ErrClosedPipe {
				s.ChannelMessageSend(m.ChannelID, xmark+" Error downloading the music")
				log.Error("error downloading YouTube video", err)
				srvr.VoiceInst.Done <- err
				return
			}
		}()
	}

	encSesh, err := dca.EncodeMem(reader, dca.StdEncodeOptions)
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, xmark+" Error starting the stream")
		srvr.youtubeCleanup()
		srvr.VoiceInst.Unlock()
		// will only return non nill error if options arent valid
		log.Error("error validating options", err)
		return
	}
	defer encSesh.Cleanup()

	srvr.VoiceInst.StreamingSession = dca.NewStream(encSesh, vc, srvr.VoiceInst.Done)

	s.ChannelMessageSend(m.ChannelID, "🔊 Playing: "+vid.Title)

	srvr.VoiceInst.Unlock()

Outer:
	for {
		err = <-srvr.VoiceInst.Done

		done, _ := srvr.VoiceInst.StreamingSession.Finished()

		switch {
		case err.Error() == "stop":
			srvr.youtubeCleanup()
			s.ChannelMessageSend(m.ChannelID, "🔇 Stopped")
			return
		case err.Error() == "skip":
			s.ChannelMessageSend(m.ChannelID, "⏩ Skipping")
			break Outer
		case !done && err != io.EOF:
			srvr.youtubeCleanup()
			s.ChannelMessageSend(m.ChannelID, "There was an error streaming music :(")
			log.Error("error streaming music", err)
			return
		}
	}

	go play(s, m, srvr, vc)
}

func listQueue(s *discordgo.Session, m *discordgo.MessageCreate) {
	guild, err := guildDetails(m.ChannelID, "", s)
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "There was an issue loading the list :( please try again")
		return
	}

	srvr, ok := sMap.Server[guild.ID]
	if !ok {
		s.ChannelMessageSend(m.ChannelID, "An error occured that really shouldn't have happened...")
		log.Error("not in server map?", guild.ID)
		return
	}

	if srvr.queueLength() == 0 {
		s.ChannelMessageSend(m.ChannelID, "No songs in queue!")
		return
	}

	p := dgwidgets.NewPaginator(s, m.ChannelID)
	p.Add(&discordgo.MessageEmbed{
		Title: guild.Name + "'s queue",

		Fields: func() (out []*discordgo.MessageEmbedField) {
			for i, song := range srvr.iterateQueue() {
				out = append(out, &discordgo.MessageEmbedField{
					Name:  fmt.Sprintf("%d - %s", i, song.Name),
					Value: song.Duration.String(),
				})
			}
			return
		}(),
	})

	for _, song := range srvr.iterateQueue() {
		p.Add(&discordgo.MessageEmbed{
			Title: fmt.Sprintf("Title: %s\nDuration: %s\nURL: %s", song.Name, song.Duration, song.URL),

			Image: &discordgo.MessageEmbedImage{
				URL: song.Image,
			},
		})
	}

	p.SetPageFooters()
	p.Loop = true
	p.ColourWhenDone = 0xff0000
	p.DeleteReactionsWhenDone = true
	p.Widget.Timeout = time.Minute * 2
	p.Spawn()
}

func stopQueue(s *discordgo.Session, m *discordgo.MessageCreate) {
	guild, err := guildDetails(m.ChannelID, "", s)
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "There was an error stopping the queue :( Please try again.")
		return
	}

	srvr := sMap.Server[guild.ID]
	srvr.VoiceInst.Lock()
	defer srvr.VoiceInst.Unlock()
	srvr.VoiceInst.Done <- errors.New("stop")
}

func pauseQueue(s *discordgo.Session, m *discordgo.MessageCreate) {
	guild, err := guildDetails(m.ChannelID, "", s)
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "There was an error pausing the video :( Please try again.")
		return
	}

	srvr := sMap.Server[guild.ID]

	srvr.VoiceInst.Lock()
	defer srvr.VoiceInst.Unlock()

	s.ChannelMessageSend(m.ChannelID, fmt.Sprintf("⏸ Paused. To unpause, use the command `%sunpause`", func() string {
		if srvr.Prefix == "" {
			return c.Prefix
		}
		return srvr.Prefix
	}()))

	srvr.VoiceInst.StreamingSession.SetPaused(true)
}

func unpauseQueue(s *discordgo.Session, m *discordgo.MessageCreate) {
	guild, err := guildDetails(m.ChannelID, "", s)
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "There was an error unpausing the song :( please try again")
		return
	}

	srvr := sMap.Server[guild.ID]
	srvr.VoiceInst.Lock()
	defer srvr.VoiceInst.Unlock()
	srvr.VoiceInst.StreamingSession.SetPaused(false)
}

func skipSong(s *discordgo.Session, m *discordgo.MessageCreate) {
	guild, err := guildDetails(m.ChannelID, "", s)
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "There was an error skipping the song :( please try again")
		return
	}

	srvr := sMap.Server[guild.ID]
	srvr.VoiceInst.Lock()
	defer srvr.VoiceInst.Unlock()
	srvr.VoiceInst.Done <- errors.New("skip")
}

func (s *server) youtubeCleanup() {
	s.VoiceInst.Lock()
	defer s.VoiceInst.Unlock()
	s.VoiceInst.VoiceCon.Disconnect()
	s.newVoiceInstance()
	//sMap.VoiceInsts--
}
<commit_msg>Fixed an issue where the queue wouldn't progress to the next song (#25)<commit_after>package main

import (
	"errors"
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/Necroforger/dgwidgets"
	"github.com/Strum355/ytdl"
	"github.com/bwmarrin/discordgo"
	"github.com/jonas747/dca"
)

const (
	stdURL   = "https://www.youtube.com/watch"
	shortURL = "https://youtu.be/"
	embedURL = "https://www.youtube.com/embed/"
)

func init() {
	newCommand("yt", 0, false, false, msgYoutube).setHelp("Args: [play,stop] [url]\n\nWork In Progress!!! Play music from Youtube straight to your Discord Server!\n\n" +
		"Example 1: `!owo yt play https://www.youtube.com/watch?v=MvLdxtICOIY`\n" +
		"Example 2: `!owo yt stop`\n\nSubCommands:\nplay\nstop\nlist, queue, songs\npause\nresume, unpause\nskip, next").add()
}

func msgYoutube(s *discordgo.Session, m *discordgo.MessageCreate, msglist []string) {
	if len(msglist) == 1 {
		return
	}

	switch msglist[1] {
	case "play":
		addToQueue(s, m, msglist[2:])
	case "stop":
		stopQueue(s, m)
	case "list", "queue", "songs":
		listQueue(s, m)
	case "pause":
		pauseQueue(s, m)
	case "resume", "unpause":
		unpauseQueue(s, m)
	case "skip", "next":
		skipSong(s, m)
	}
}

func addToQueue(s *discordgo.Session, m *discordgo.MessageCreate, msglist []string) {
	if len(msglist) == 0 {
		return
	}

	guild, err := guildDetails(m.ChannelID, "", s)
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "There was a problem adding to queue :( please try again")
		return
	}

	srvr, ok := sMap.Server[guild.ID]
	if !ok {
		s.ChannelMessageSend(m.ChannelID, "An error occured that really shouldn't have happened...")
		log.Error("not in server map?", guild.ID)
		return
	}

	if srvr.VoiceInst == nil {
		srvr.newVoiceInstance()
	}

	srvr.VoiceInst.Lock()
	defer srvr.VoiceInst.Unlock()

	url := msglist[0]

	if !strings.HasPrefix(url, stdURL) && !strings.HasPrefix(url, shortURL) && !strings.HasPrefix(url, embedURL) {
		s.ChannelMessageSend(m.ChannelID, "Please make sure the URL is a valid YouTube URL. If I got this wrong, please let my creator know~")
		return
	}

	for _, vs := range guild.VoiceStates {
		if vs.UserID == m.Author.ID {
			if vs.ChannelID == srvr.VoiceInst.ChannelID || !srvr.VoiceInst.Playing {
				vid, err := getVideoInfo(url, s, m)
				if err != nil {
					return
				}

				vc, err := s.ChannelVoiceJoin(guild.ID, vs.ChannelID, false, true)
				if err != nil {
					s.ChannelMessageSend(m.ChannelID, "Error joining voice channel")
					log.Error("error joining voice channel", err)
					return
				}

				srvr.addSong(song{
					URL:      url,
					Name:     vid.Title,
					Duration: vid.Duration,
					Image:    vid.GetThumbnailURL(ytdl.ThumbnailQualityMedium).String(),
				})

				s.ChannelMessageSend(m.ChannelID, "Added "+vid.Title+" to the queue!")

				if !srvr.VoiceInst.Playing {
					srvr.VoiceInst.VoiceCon = vc
					srvr.VoiceInst.Playing = true
					srvr.VoiceInst.ChannelID = vc.ChannelID
					go play(s, m, srvr, vc)
				}
				return
			}
		}
	}

	s.ChannelMessageSend(m.ChannelID, "Need to be in a voice channel!")
}

func getVideoInfo(url string, s *discordgo.Session, m *discordgo.MessageCreate) (*ytdl.VideoInfo, error) {
	vid, err := ytdl.GetVideoInfo(url)
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "Error getting video info")
		log.Error("error getting video info", err)
		return nil, err
	}
	return vid, nil
}

func play(s *discordgo.Session, m *discordgo.MessageCreate, srvr *server, vc *discordgo.VoiceConnection) {
	if srvr.queueLength() == 0 {
		srvr.youtubeCleanup()
		s.ChannelMessageSend(m.ChannelID, "🔇 Done queue!")
		return
	}

	srvr.VoiceInst.Lock()
	vid, err := getVideoInfo(srvr.nextSong().URL, s, m)
	if err != nil {
		srvr.VoiceInst.Unlock()
		return
	}

	reader, writer := io.Pipe()
	defer reader.Close()

	formats := vid.Formats.Best(ytdl.FormatAudioBitrateKey)
	if len(formats) > 0 {
		go func() {
			defer writer.Close()
			if err := vid.Download(formats[0], writer); err != nil && err != io.ErrClosedPipe {
				s.ChannelMessageSend(m.ChannelID, xmark+" Error downloading the music")
				log.Error("error downloading YouTube video", err)
				srvr.VoiceInst.Done <- err
				return
			}
		}()
	}

	encSesh, err := dca.EncodeMem(reader, dca.StdEncodeOptions)
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, xmark+" Error starting the stream")
		srvr.youtubeCleanup()
		srvr.VoiceInst.Unlock()
		// will only return non nill error if options arent valid
		log.Error("error validating options", err)
		return
	}
	defer encSesh.Cleanup()

	srvr.VoiceInst.StreamingSession = dca.NewStream(encSesh, vc, srvr.VoiceInst.Done)

	s.ChannelMessageSend(m.ChannelID, "🔊 Playing: "+vid.Title)

	srvr.VoiceInst.Unlock()

Outer:
	for {
		err = <-srvr.VoiceInst.Done

		done, _ := srvr.VoiceInst.StreamingSession.Finished()

		switch {
		case err.Error() == "stop":
			srvr.youtubeCleanup()
			s.ChannelMessageSend(m.ChannelID, "🔇 Stopped")
			return
		case err.Error() == "skip":
			s.ChannelMessageSend(m.ChannelID, "⏩ Skipping")
			break Outer
		case !done && err != io.EOF:
			srvr.youtubeCleanup()
			s.ChannelMessageSend(m.ChannelID, "There was an error streaming music :(")
			log.Error("error streaming music", err)
			return
		case done && err == io.EOF:
			break Outer
		}
	}

	go play(s, m, srvr, vc)
}

func listQueue(s *discordgo.Session, m *discordgo.MessageCreate) {
	guild, err := guildDetails(m.ChannelID, "", s)
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "There was an issue loading the list :( please try again")
		return
	}

	srvr, ok := sMap.Server[guild.ID]
	if !ok {
		s.ChannelMessageSend(m.ChannelID, "An error occured that really shouldn't have happened...")
		log.Error("not in server map?", guild.ID)
		return
	}

	if srvr.queueLength() == 0 {
		s.ChannelMessageSend(m.ChannelID, "No songs in queue!")
		return
	}

	p := dgwidgets.NewPaginator(s, m.ChannelID)
	p.Add(&discordgo.MessageEmbed{
		Title: guild.Name + "'s queue",

		Fields: func() (out []*discordgo.MessageEmbedField) {
			for i, song := range srvr.iterateQueue() {
				out = append(out, &discordgo.MessageEmbedField{
					Name:  fmt.Sprintf("%d - %s", i, song.Name),
					Value: song.Duration.String(),
				})
			}
			return
		}(),
	})

	for _, song := range srvr.iterateQueue() {
		p.Add(&discordgo.MessageEmbed{
			Title: fmt.Sprintf("Title: %s\nDuration: %s\nURL: %s", song.Name, song.Duration, song.URL),

			Image: &discordgo.MessageEmbedImage{
				URL: song.Image,
			},
		})
	}

	p.SetPageFooters()
	p.Loop = true
	p.ColourWhenDone = 0xff0000
	p.DeleteReactionsWhenDone = true
	p.Widget.Timeout = time.Minute * 2
	p.Spawn()
}

func stopQueue(s *discordgo.Session, m *discordgo.MessageCreate) {
	guild, err := guildDetails(m.ChannelID, "", s)
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "There was an error stopping the queue :( Please try again.")
		return
	}

	srvr := sMap.Server[guild.ID]
	srvr.VoiceInst.Lock()
	defer srvr.VoiceInst.Unlock()
	srvr.VoiceInst.Done <- errors.New("stop")
}

func pauseQueue(s *discordgo.Session, m *discordgo.MessageCreate) {
	guild, err := guildDetails(m.ChannelID, "", s)
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "There was an error pausing the video :( Please try again.")
		return
	}

	srvr := sMap.Server[guild.ID]

	srvr.VoiceInst.Lock()
	defer srvr.VoiceInst.Unlock()

	s.ChannelMessageSend(m.ChannelID, fmt.Sprintf("⏸ Paused. To unpause, use the command `%sunpause`", func() string {
		if srvr.Prefix == "" {
			return c.Prefix
		}
		return srvr.Prefix
	}()))

	srvr.VoiceInst.StreamingSession.SetPaused(true)
}

func unpauseQueue(s *discordgo.Session, m *discordgo.MessageCreate) {
	guild, err := guildDetails(m.ChannelID, "", s)
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "There was an error unpausing the song :( please try again")
		return
	}

	srvr := sMap.Server[guild.ID]
	srvr.VoiceInst.Lock()
	defer srvr.VoiceInst.Unlock()
	srvr.VoiceInst.StreamingSession.SetPaused(false)
}

func skipSong(s *discordgo.Session, m *discordgo.MessageCreate) {
	guild, err := guildDetails(m.ChannelID, "", s)
	if err != nil {
		s.ChannelMessageSend(m.ChannelID, "There was an error skipping the song :( please try again")
		return
	}

	srvr := sMap.Server[guild.ID]
	srvr.VoiceInst.Lock()
	defer srvr.VoiceInst.Unlock()
	srvr.VoiceInst.Done <- errors.New("skip")
}

func (s *server) youtubeCleanup() {
	s.VoiceInst.Lock()
	defer s.VoiceInst.Unlock()
	s.VoiceInst.VoiceCon.Disconnect()
	s.newVoiceInstance()
	//sMap.VoiceInsts--
}
<|endoftext|>


<commit_before>/*
	Utility functions for decoding mysql packets
*/

package mysql

import (
	"encoding/binary"
	"bufio"
	"os"
	"fmt"
	"crypto/sha1"
	"bytes"
)

//Read mysql packet header
func readHeader(br *bufio.Reader) (*PacketHeader, os.Error) {
	ph := new(PacketHeader)
	var i24seq [4]byte
	nn, err := br.Read(&i24seq)
	if err != nil {
		return nil, os.ErrorString(fmt.Sprintf("readHeader: %s", err))
	}
	if nn < 4 {
		return nil, os.ErrorString(fmt.Sprintf("readHeader: Packet to small (%d)", nn))
	}
	ph.Len = unpackNumber(&i24seq, 3)
	ph.Seq = i24seq[3]
	return ph, nil
}

//Decode length encoded number
//TODO: *Decode() Check Buffered bytes.
func unpackLength(br *bufio.Reader) (uint64, bool) {
	var bl uint8
	binary.Read(br, binary.LittleEndian, &bl)
	if bl < 251 {
		return uint64(bl), false
	} else if bl == 251 {
		return 0, true
	} else if bl == 252 {
		b := make([]byte, 2)
		br.Read(b)
		return unpackNumber(b, 2), false
	} else if bl == 253 {
		b := make([]byte, 3)
		br.Read(b)
		return unpackNumber(b, 3), false
	} else if bl == 254 && br.Buffered() > 8 {
		b := make([]byte, 8)
		br.Read(b)
		return unpackNumber(b, 8), false
	}
	return uint64(bl), false
}

//Special case of unpackLength where 0xfe == EOF
func unpackFieldCount(br *bufio.Reader) (uint64, bool) {
	if peekEOF(br) {
		ignoreBytes(br, 1)
		return uint64(0xfe), false
	}
	return unpackLength(br)
}

func packUint8(bw *bufio.Writer, u uint8) os.Error {
	return binary.Write(bw, binary.LittleEndian, u)
}

func packUint16(bw *bufio.Writer, u uint16) os.Error {
	return binary.Write(bw, binary.LittleEndian, u)
}

func packUint24(bw *bufio.Writer, u uint32) os.Error {
	b := make([]byte, 3)
	b[0] = byte(u)
	b[1] = byte(u >> 8)
	b[2] = byte(u >> 16)
	n, err := bw.Write(b)
	if n != 3 {
		return err
	}
	return nil
}

func packUint32(bw *bufio.Writer, u uint32) os.Error {
	return binary.Write(bw, binary.LittleEndian, u)
}

func packUint64(bw *bufio.Writer, u uint64) os.Error {
	return binary.Write(bw, binary.LittleEndian, u)
}

//Decode length encoded string
func unpackString(br *bufio.Reader) (string, bool) {
	length, isnull := unpackLength(br)
	b := make([]byte, length)
	br.Read(b)
	return string(b), isnull
}

func packString(s string) []byte {
	sb := []byte(s)
	size := make([]byte, 1)
	v := len(sb)
	if v < 250 {
		size[0] = uint8(v)
		return bytes.Add(size, sb)
	}
	size = make([]byte, 9)
	size[0] = 254
	size[1] = byte(v)
	size[2] = byte(v >> 8)
	size[3] = byte(v >> 16)
	size[4] = byte(v >> 24)
	size[5] = byte(v >> 32)
	size[6] = byte(v >> 40)
	size[7] = byte(v >> 48)
	size[8] = byte(v >> 56)
	return bytes.Add(size, sb)
}

//Peek and check if packet is EOF
func peekEOF(br *bufio.Reader) bool {
	b := make([]byte, 1)
	br.Read(b)
	br.UnreadByte()
	if b[0] == 0xfe {
		return true
	}
	return false
}

//Convert n bytes to uint64 (Little Endian)
func unpackNumber(b []byte, n uint8) uint64 {
	if n < 1 {
		return 0
	}
	var r uint64 = 0
	for i := uint8(0); i < n; i++ {
		r |= (uint64(b[i]) << (i * 8))
	}
	return r
}

//Read the field data from mysql ResultSet packet
func readFieldPacket(br *bufio.Reader) *MySQLField {
	f := new(MySQLField)
	f.Catalog, _ = unpackString(br)
	f.Db, _ = unpackString(br)
	f.Table, _ = unpackString(br)
	f.OrgTable, _ = unpackString(br)
	f.Name, _ = unpackString(br)
	f.OrgName, _ = unpackString(br)
	ignoreBytes(br, 1)
	binary.Read(br, binary.LittleEndian, &f.Charset)
	binary.Read(br, binary.LittleEndian, &f.Length)
	binary.Read(br, binary.LittleEndian, &f.Type)
	binary.Read(br, binary.LittleEndian, &f.Flags)
	binary.Read(br, binary.LittleEndian, &f.Decimals)
	ignoreBytes(br, 1)
	eb, _ := unpackLength(br)
	f.Default = eb
	return f
}

//Read error packet
func readErrorPacket(br *bufio.Reader) os.Error {
	var errcode uint16
	binary.Read(br, binary.LittleEndian, &errcode)
	status := make([]byte, 6)
	br.Read(status)
	msg := make([]byte, br.Buffered())
	br.Read(msg)
	return os.ErrorString(fmt.Sprintf("MySQL Error: (Code: %d) (Status: %s) %s", errcode, string(status), string(msg)))
}

//Read EOF packet.
//TODO: Return something useful?
func readEOFPacket(br *bufio.Reader) os.Error {
	readHeader(br)

	response := new(MySQLResponse)
	response.FieldCount, _ = unpackFieldCount(br)
	if response.FieldCount != 0xfe {
		fmt.Printf("Expected EOF! Got %#v\n", response.FieldCount)
	}
	binary.Read(br, binary.LittleEndian, &response.WarningCount)
	binary.Read(br, binary.LittleEndian, &response.ServerStatus)
	return nil
}

//Ignores n bytes in the buffer
func ignoreBytes(br *bufio.Reader, n uint64) {
	buf := make([]byte, n)
	br.Read(buf)
}

//Generate scrabled password using password and scramble buffer.
func mysqlPassword(password []byte, scrambleBuffer []byte) []byte {
	ctx := sha1.New()
	ctx.Write(password)
	stage1 := ctx.Sum()

	ctx.Reset()
	ctx.Write(stage1)
	stage2 := ctx.Sum()

	ctx.Reset()
	ctx.Write(scrambleBuffer)
	ctx.Write(stage2)
	result := ctx.Sum()

	token := new([21]byte)
	token_t := new([20]byte)
	for i := 0; i < 20; i++ {
		token[i+1] = result[i] ^ stage1[i]
	}
	for i := 0; i < 20; i++ {
		token_t[i] = token[i+1] ^ result[i]
	}
	token[0] = 20
	return token
}
<commit_msg>It's only a warning<commit_after>/*
	Utility functions for decoding mysql packets
*/

package mysql

import (
	"encoding/binary"
	"bufio"
	"os"
	"fmt"
	"crypto/sha1"
	"bytes"
)

//Read mysql packet header
func readHeader(br *bufio.Reader) (*PacketHeader, os.Error) {
	ph := new(PacketHeader)
	var i24seq [4]byte
	nn, err := br.Read(&i24seq)
	if err != nil {
		return nil, os.ErrorString(fmt.Sprintf("readHeader: %s", err))
	}
	if nn < 4 {
		return nil, os.ErrorString(fmt.Sprintf("readHeader: Packet to small (%d)", nn))
	}
	ph.Len = unpackNumber(&i24seq, 3)
	ph.Seq = i24seq[3]
	return ph, nil
}

//Decode length encoded number
//TODO: *Decode() Check Buffered bytes.
func unpackLength(br *bufio.Reader) (uint64, bool) {
	var bl uint8
	binary.Read(br, binary.LittleEndian, &bl)
	if bl < 251 {
		return uint64(bl), false
	} else if bl == 251 {
		return 0, true
	} else if bl == 252 {
		b := make([]byte, 2)
		br.Read(b)
		return unpackNumber(b, 2), false
	} else if bl == 253 {
		b := make([]byte, 3)
		br.Read(b)
		return unpackNumber(b, 3), false
	} else if bl == 254 && br.Buffered() > 8 {
		b := make([]byte, 8)
		br.Read(b)
		return unpackNumber(b, 8), false
	}
	return uint64(bl), false
}

//Special case of unpackLength where 0xfe == EOF
func unpackFieldCount(br *bufio.Reader) (uint64, bool) {
	if peekEOF(br) {
		ignoreBytes(br, 1)
		return uint64(0xfe), false
	}
	return unpackLength(br)
}

func packUint8(bw *bufio.Writer, u uint8) os.Error {
	return binary.Write(bw, binary.LittleEndian, u)
}

func packUint16(bw *bufio.Writer, u uint16) os.Error {
	return binary.Write(bw, binary.LittleEndian, u)
}

func packUint24(bw *bufio.Writer, u uint32) os.Error {
	b := make([]byte, 3)
	b[0] = byte(u)
	b[1] = byte(u >> 8)
	b[2] = byte(u >> 16)
	n, err := bw.Write(b)
	if n != 3 {
		return err
	}
	return nil
}

func packUint32(bw *bufio.Writer, u uint32) os.Error {
	return binary.Write(bw, binary.LittleEndian, u)
}

func packUint64(bw *bufio.Writer, u uint64) os.Error {
	return binary.Write(bw, binary.LittleEndian, u)
}

//Decode length encoded string
func unpackString(br *bufio.Reader) (string, bool) {
	length, isnull := unpackLength(br)
	b := make([]byte, length)
	br.Read(b)
	return string(b), isnull
}

func packString(s string) []byte {
	sb := []byte(s)
	size := make([]byte, 1)
	v := len(sb)
	if v < 250 {
		size[0] = uint8(v)
		return bytes.Add(size, sb)
	}
	size = make([]byte, 9)
	size[0] = 254
	size[1] = byte(v)
	size[2] = byte(v >> 8)
	size[3] = byte(v >> 16)
	size[4] = byte(v >> 24)
	size[5] = byte(v >> 32)
	size[6] = byte(v >> 40)
	size[7] = byte(v >> 48)
	size[8] = byte(v >> 56)
	return bytes.Add(size, sb)
}

//Peek and check if packet is EOF
func peekEOF(br *bufio.Reader) bool {
	b := make([]byte, 1)
	br.Read(b)
	br.UnreadByte()
	if b[0] == 0xfe {
		return true
	}
	return false
}

//Convert n bytes to uint64 (Little Endian)
func unpackNumber(b []byte, n uint8) uint64 {
	if n < 1 {
		return 0
	}
	var r uint64 = 0
	for i := uint8(0); i < n; i++ {
		r |= (uint64(b[i]) << (i * 8))
	}
	return r
}

//Read the field data from mysql ResultSet packet
func readFieldPacket(br *bufio.Reader) *MySQLField {
	f := new(MySQLField)
	f.Catalog, _ = unpackString(br)
	f.Db, _ = unpackString(br)
	f.Table, _ = unpackString(br)
	f.OrgTable, _ = unpackString(br)
	f.Name, _ = unpackString(br)
	f.OrgName, _ = unpackString(br)
	ignoreBytes(br, 1)
	binary.Read(br, binary.LittleEndian, &f.Charset)
	binary.Read(br, binary.LittleEndian, &f.Length)
	binary.Read(br, binary.LittleEndian, &f.Type)
	binary.Read(br, binary.LittleEndian, &f.Flags)
	binary.Read(br, binary.LittleEndian, &f.Decimals)
	ignoreBytes(br, 1)
	eb, _ := unpackLength(br)
	f.Default = eb
	return f
}

//Read error packet
func readErrorPacket(br *bufio.Reader) os.Error {
	var errcode uint16
	binary.Read(br, binary.LittleEndian, &errcode)
	status := make([]byte, 6)
	br.Read(status)
	msg := make([]byte, br.Buffered())
	br.Read(msg)
	return os.ErrorString(fmt.Sprintf("MySQL Error: (Code: %d) (Status: %s) %s", errcode, string(status), string(msg)))
}

//Read EOF packet.
//TODO: Return something useful?
func readEOFPacket(br *bufio.Reader) os.Error {
	readHeader(br)

	response := new(MySQLResponse)
	response.FieldCount, _ = unpackFieldCount(br)
	if response.FieldCount != 0xfe {
		fmt.Printf("Warning: Expected EOF! Got %#v\n", response.FieldCount)
	}
	binary.Read(br, binary.LittleEndian, &response.WarningCount)
	binary.Read(br, binary.LittleEndian, &response.ServerStatus)
	return nil
}

//Ignores n bytes in the buffer
func ignoreBytes(br *bufio.Reader, n uint64) {
	buf := make([]byte, n)
	br.Read(buf)
}

//Generate scrabled password using password and scramble buffer.
func mysqlPassword(password []byte, scrambleBuffer []byte) []byte {
	ctx := sha1.New()
	ctx.Write(password)
	stage1 := ctx.Sum()

	ctx.Reset()
	ctx.Write(stage1)
	stage2 := ctx.Sum()

	ctx.Reset()
	ctx.Write(scrambleBuffer)
	ctx.Write(stage2)
	result := ctx.Sum()

	token := new([21]byte)
	token_t := new([20]byte)
	for i := 0; i < 20; i++ {
		token[i+1] = result[i] ^ stage1[i]
	}
	for i := 0; i < 20; i++ {
		token_t[i] = token[i+1] ^ result[i]
	}
	token[0] = 20
	return token
}
<|endoftext|>


<commit_before>package sqlx

import (
	"database/sql"
	"fmt"
	"testing"
)

func TestCompileQuery(t *testing.T) {
	table := []struct {
		Q, R, D, T, N string
		V             []string
	}{
		// basic test for named parameters, invalid char ',' terminating
		{
			Q: `INSERT INTO foo (a,b,c,d) VALUES (:name, :age, :first, :last)`,
			R: `INSERT INTO foo (a,b,c,d) VALUES (?, ?, ?, ?)`,
			D: `INSERT INTO foo (a,b,c,d) VALUES ($1, $2, $3, $4)`,
			T: `INSERT INTO foo (a,b,c,d) VALUES (@p1, @p2, @p3, @p4)`,
			N: `INSERT INTO foo (a,b,c,d) VALUES (:name, :age, :first, :last)`,
			V: []string{"name", "age", "first", "last"},
		},
		// This query tests a named parameter ending the string as well as numbers
		{
			Q: `SELECT * FROM a WHERE first_name=:name1 AND last_name=:name2`,
			R: `SELECT * FROM a WHERE first_name=? AND last_name=?`,
			D: `SELECT * FROM a WHERE first_name=$1 AND last_name=$2`,
			T: `SELECT * FROM a WHERE first_name=@p1 AND last_name=@p2`,
			N: `SELECT * FROM a WHERE first_name=:name1 AND last_name=:name2`,
			V: []string{"name1", "name2"},
		},
		{
			Q: `SELECT "::foo" FROM a WHERE first_name=:name1 AND last_name=:name2`,
			R: `SELECT ":foo" FROM a WHERE first_name=? AND last_name=?`,
			D: `SELECT ":foo" FROM a WHERE first_name=$1 AND last_name=$2`,
			T: `SELECT ":foo" FROM a WHERE first_name=@p1 AND last_name=@p2`,
			N: `SELECT ":foo" FROM a WHERE first_name=:name1 AND last_name=:name2`,
			V: []string{"name1", "name2"},
		},
		{
			Q: `SELECT 'a::b::c' || first_name, '::::ABC::_::' FROM person WHERE first_name=:first_name AND last_name=:last_name`,
			R: `SELECT 'a:b:c' || first_name, '::ABC:_:' FROM person WHERE first_name=? AND last_name=?`,
			D: `SELECT 'a:b:c' || first_name, '::ABC:_:' FROM person WHERE first_name=$1 AND last_name=$2`,
			T: `SELECT 'a:b:c' || first_name, '::ABC:_:' FROM person WHERE first_name=@p1 AND last_name=@p2`,
			N: `SELECT 'a:b:c' || first_name, '::ABC:_:' FROM person WHERE first_name=:first_name AND last_name=:last_name`,
			V: []string{"first_name", "last_name"},
		},
		{
			Q: `SELECT @name := "name", :age, :first, :last`,
			R: `SELECT @name := "name", ?, ?, ?`,
			D: `SELECT @name := "name", $1, $2, $3`,
			N: `SELECT @name := "name", :age, :first, :last`,
			T: `SELECT @name := "name", @p1, @p2, @p3`,
			V: []string{"age", "first", "last"},
		},
		/* This unicode awareness test sadly fails, because of our byte-wise worldview.
		 * We could certainly iterate by Rune instead, though it's a great deal slower,
		 * it's probably the RightWay(tm)
		{
			Q: `INSERT INTO foo (a,b,c,d) VALUES (:あ, :b, :キコ, :名前)`,
			R: `INSERT INTO foo (a,b,c,d) VALUES (?, ?, ?, ?)`,
			D: `INSERT INTO foo (a,b,c,d) VALUES ($1, $2, $3, $4)`,
			N: []string{"name", "age", "first", "last"},
		},
		*/
	}

	for _, test := range table {
		qr, names, err := compileNamedQuery([]byte(test.Q), QUESTION)
		if err != nil {
			t.Error(err)
		}
		if qr != test.R {
			t.Errorf("expected %s, got %s", test.R, qr)
		}
		if len(names) != len(test.V) {
			t.Errorf("expected %#v, got %#v", test.V, names)
		} else {
			for i, name := range names {
				if name != test.V[i] {
					t.Errorf("expected %dth name to be %s, got %s", i+1, test.V[i], name)
				}
			}
		}
		qd, _, _ := compileNamedQuery([]byte(test.Q), DOLLAR)
		if qd != test.D {
			t.Errorf("\nexpected: `%s`\ngot:      `%s`", test.D, qd)
		}

		qt, _, _ := compileNamedQuery([]byte(test.Q), AT)
		if qt != test.T {
			t.Errorf("\nexpected: `%s`\ngot:      `%s`", test.T, qt)
		}

		qq, _, _ := compileNamedQuery([]byte(test.Q), NAMED)
		if qq != test.N {
			t.Errorf("\nexpected: `%s`\ngot:      `%s`\n(len: %d vs %d)", test.N, qq, len(test.N), len(qq))
		}
	}
}

type Test struct {
	t *testing.T
}

func (t Test) Error(err error, msg ...interface{}) {
	if err != nil {
		if len(msg) == 0 {
			t.t.Error(err)
		} else {
			t.t.Error(msg...)
		}
	}
}

func (t Test) Errorf(err error, format string, args ...interface{}) {
	if err != nil {
		t.t.Errorf(format, args...)
	}
}

func TestEscapedColons(t *testing.T) {
	t.Skip("not sure it is possible to support this in general case without an SQL parser")
	var qs = `SELECT * FROM testtable WHERE timeposted BETWEEN (now() AT TIME ZONE 'utc') AND
	(now() AT TIME ZONE 'utc') - interval '01:30:00') AND name = '\'this is a test\'' and id = :id`
	_, _, err := compileNamedQuery([]byte(qs), DOLLAR)
	if err != nil {
		t.Error("Didn't handle colons correctly when inside a string")
	}
}

func TestNamedQueries(t *testing.T) {
	RunWithSchema(defaultSchema, t, func(db *DB, t *testing.T, now string) {
		loadDefaultFixture(db, t)
		test := Test{t}
		var ns *NamedStmt
		var err error

		// Check that invalid preparations fail
		ns, err = db.PrepareNamed("SELECT * FROM person WHERE first_name=:first:name")
		if err == nil {
			t.Error("Expected an error with invalid prepared statement.")
		}

		ns, err = db.PrepareNamed("invalid sql")
		if err == nil {
			t.Error("Expected an error with invalid prepared statement.")
		}

		// Check closing works as anticipated
		ns, err = db.PrepareNamed("SELECT * FROM person WHERE first_name=:first_name")
		test.Error(err)
		err = ns.Close()
		test.Error(err)

		ns, err = db.PrepareNamed(`
			SELECT first_name, last_name, email
			FROM person WHERE first_name=:first_name AND email=:email`)
		test.Error(err)

		// test Queryx w/ uses Query
		p := Person{FirstName: "Jason", LastName: "Moiron", Email: "jmoiron@jmoiron.net"}

		rows, err := ns.Queryx(p)
		test.Error(err)
		for rows.Next() {
			var p2 Person
			rows.StructScan(&p2)
			if p.FirstName != p2.FirstName {
				t.Errorf("got %s, expected %s", p.FirstName, p2.FirstName)
			}
			if p.LastName != p2.LastName {
				t.Errorf("got %s, expected %s", p.LastName, p2.LastName)
			}
			if p.Email != p2.Email {
				t.Errorf("got %s, expected %s", p.Email, p2.Email)
			}
		}

		// test Select
		people := make([]Person, 0, 5)
		err = ns.Select(&people, p)
		test.Error(err)

		if len(people) != 1 {
			t.Errorf("got %d results, expected %d", len(people), 1)
		}
		if p.FirstName != people[0].FirstName {
			t.Errorf("got %s, expected %s", p.FirstName, people[0].FirstName)
		}
		if p.LastName != people[0].LastName {
			t.Errorf("got %s, expected %s", p.LastName, people[0].LastName)
		}
		if p.Email != people[0].Email {
			t.Errorf("got %s, expected %s", p.Email, people[0].Email)
		}

		// test struct batch inserts
		sls := []Person{
			{FirstName: "Ardie", LastName: "Savea", Email: "asavea@ab.co.nz"},
			{FirstName: "Sonny Bill", LastName: "Williams", Email: "sbw@ab.co.nz"},
			{FirstName: "Ngani", LastName: "Laumape", Email: "nlaumape@ab.co.nz"},
		}

		insert := fmt.Sprintf(
			"INSERT INTO person (first_name, last_name, email, added_at) VALUES (:first_name, :last_name, :email, %v)\n",
			now,
		)
		_, err = db.NamedExec(insert, sls)
		test.Error(err)

		// test map batch inserts
		slsMap := []map[string]interface{}{
			{"first_name": "Ardie", "last_name": "Savea", "email": "asavea@ab.co.nz"},
			{"first_name": "Sonny Bill", "last_name": "Williams", "email": "sbw@ab.co.nz"},
			{"first_name": "Ngani", "last_name": "Laumape", "email": "nlaumape@ab.co.nz"},
		}

		_, err = db.NamedExec(`INSERT INTO person (first_name, last_name, email)
			VALUES (:first_name, :last_name, :email) ;--`, slsMap)
		test.Error(err)

		type A map[string]interface{}

		typedMap := []A{
			{"first_name": "Ardie", "last_name": "Savea", "email": "asavea@ab.co.nz"},
			{"first_name": "Sonny Bill", "last_name": "Williams", "email": "sbw@ab.co.nz"},
			{"first_name": "Ngani", "last_name": "Laumape", "email": "nlaumape@ab.co.nz"},
		}

		_, err = db.NamedExec(`INSERT INTO person (first_name, last_name, email)
			VALUES (:first_name, :last_name, :email) ;--`, typedMap)
		test.Error(err)

		for _, p := range sls {
			dest := Person{}
			err = db.Get(&dest, db.Rebind("SELECT * FROM person WHERE email=?"), p.Email)
			test.Error(err)
			if dest.Email != p.Email {
				t.Errorf("expected %s, got %s", p.Email, dest.Email)
			}
		}

		// test Exec
		ns, err = db.PrepareNamed(`
			INSERT INTO person (first_name, last_name, email)
			VALUES (:first_name, :last_name, :email)`)
		test.Error(err)

		js := Person{
			FirstName: "Julien",
			LastName:  "Savea",
			Email:     "jsavea@ab.co.nz",
		}
		_, err = ns.Exec(js)
		test.Error(err)

		// Make sure we can pull him out again
		p2 := Person{}
		db.Get(&p2, db.Rebind("SELECT * FROM person WHERE email=?"), js.Email)
		if p2.Email != js.Email {
			t.Errorf("expected %s, got %s", js.Email, p2.Email)
		}

		// test Txn NamedStmts
		tx := db.MustBegin()
		txns := tx.NamedStmt(ns)

		// We're going to add Steven in this txn
		sl := Person{
			FirstName: "Steven",
			LastName:  "Luatua",
			Email:     "sluatua@ab.co.nz",
		}

		_, err = txns.Exec(sl)
		test.Error(err)
		// then rollback...
		tx.Rollback()
		// looking for Steven after a rollback should fail
		err = db.Get(&p2, db.Rebind("SELECT * FROM person WHERE email=?"), sl.Email)
		if err != sql.ErrNoRows {
			t.Errorf("expected no rows error, got %v", err)
		}

		// now do the same, but commit
		tx = db.MustBegin()
		txns = tx.NamedStmt(ns)
		_, err = txns.Exec(sl)
		test.Error(err)
		tx.Commit()

		// looking for Steven after a Commit should succeed
		err = db.Get(&p2, db.Rebind("SELECT * FROM person WHERE email=?"), sl.Email)
		test.Error(err)
		if p2.Email != sl.Email {
			t.Errorf("expected %s, got %s", sl.Email, p2.Email)
		}

	})
}
<commit_msg>add basic happy tests<commit_after>package sqlx

import (
	"database/sql"
	"fmt"
	"testing"
)

func TestCompileQuery(t *testing.T) {
	table := []struct {
		Q, R, D, T, N string
		V             []string
	}{
		// basic test for named parameters, invalid char ',' terminating
		{
			Q: `INSERT INTO foo (a,b,c,d) VALUES (:name, :age, :first, :last)`,
			R: `INSERT INTO foo (a,b,c,d) VALUES (?, ?, ?, ?)`,
			D: `INSERT INTO foo (a,b,c,d) VALUES ($1, $2, $3, $4)`,
			T: `INSERT INTO foo (a,b,c,d) VALUES (@p1, @p2, @p3, @p4)`,
			N: `INSERT INTO foo (a,b,c,d) VALUES (:name, :age, :first, :last)`,
			V: []string{"name", "age", "first", "last"},
		},
		// This query tests a named parameter ending the string as well as numbers
		{
			Q: `SELECT * FROM a WHERE first_name=:name1 AND last_name=:name2`,
			R: `SELECT * FROM a WHERE first_name=? AND last_name=?`,
			D: `SELECT * FROM a WHERE first_name=$1 AND last_name=$2`,
			T: `SELECT * FROM a WHERE first_name=@p1 AND last_name=@p2`,
			N: `SELECT * FROM a WHERE first_name=:name1 AND last_name=:name2`,
			V: []string{"name1", "name2"},
		},
		{
			Q: `SELECT "::foo" FROM a WHERE first_name=:name1 AND last_name=:name2`,
			R: `SELECT ":foo" FROM a WHERE first_name=? AND last_name=?`,
			D: `SELECT ":foo" FROM a WHERE first_name=$1 AND last_name=$2`,
			T: `SELECT ":foo" FROM a WHERE first_name=@p1 AND last_name=@p2`,
			N: `SELECT ":foo" FROM a WHERE first_name=:name1 AND last_name=:name2`,
			V: []string{"name1", "name2"},
		},
		{
			Q: `SELECT 'a::b::c' || first_name, '::::ABC::_::' FROM person WHERE first_name=:first_name AND last_name=:last_name`,
			R: `SELECT 'a:b:c' || first_name, '::ABC:_:' FROM person WHERE first_name=? AND last_name=?`,
			D: `SELECT 'a:b:c' || first_name, '::ABC:_:' FROM person WHERE first_name=$1 AND last_name=$2`,
			T: `SELECT 'a:b:c' || first_name, '::ABC:_:' FROM person WHERE first_name=@p1 AND last_name=@p2`,
			N: `SELECT 'a:b:c' || first_name, '::ABC:_:' FROM person WHERE first_name=:first_name AND last_name=:last_name`,
			V: []string{"first_name", "last_name"},
		},
		{
			Q: `SELECT @name := "name", :age, :first, :last`,
			R: `SELECT @name := "name", ?, ?, ?`,
			D: `SELECT @name := "name", $1, $2, $3`,
			N: `SELECT @name := "name", :age, :first, :last`,
			T: `SELECT @name := "name", @p1, @p2, @p3`,
			V: []string{"age", "first", "last"},
		},
		/* This unicode awareness test sadly fails, because of our byte-wise worldview.
		 * We could certainly iterate by Rune instead, though it's a great deal slower,
		 * it's probably the RightWay(tm)
		{
			Q: `INSERT INTO foo (a,b,c,d) VALUES (:あ, :b, :キコ, :名前)`,
			R: `INSERT INTO foo (a,b,c,d) VALUES (?, ?, ?, ?)`,
			D: `INSERT INTO foo (a,b,c,d) VALUES ($1, $2, $3, $4)`,
			N: []string{"name", "age", "first", "last"},
		},
		*/
	}

	for _, test := range table {
		qr, names, err := compileNamedQuery([]byte(test.Q), QUESTION)
		if err != nil {
			t.Error(err)
		}
		if qr != test.R {
			t.Errorf("expected %s, got %s", test.R, qr)
		}
		if len(names) != len(test.V) {
			t.Errorf("expected %#v, got %#v", test.V, names)
		} else {
			for i, name := range names {
				if name != test.V[i] {
					t.Errorf("expected %dth name to be %s, got %s", i+1, test.V[i], name)
				}
			}
		}
		qd, _, _ := compileNamedQuery([]byte(test.Q), DOLLAR)
		if qd != test.D {
			t.Errorf("\nexpected: `%s`\ngot:      `%s`", test.D, qd)
		}

		qt, _, _ := compileNamedQuery([]byte(test.Q), AT)
		if qt != test.T {
			t.Errorf("\nexpected: `%s`\ngot:      `%s`", test.T, qt)
		}

		qq, _, _ := compileNamedQuery([]byte(test.Q), NAMED)
		if qq != test.N {
			t.Errorf("\nexpected: `%s`\ngot:      `%s`\n(len: %d vs %d)", test.N, qq, len(test.N), len(qq))
		}
	}
}

type Test struct {
	t *testing.T
}

func (t Test) Error(err error, msg ...interface{}) {
	if err != nil {
		if len(msg) == 0 {
			t.t.Error(err)
		} else {
			t.t.Error(msg...)
		}
	}
}

func (t Test) Errorf(err error, format string, args ...interface{}) {
	if err != nil {
		t.t.Errorf(format, args...)
	}
}

func TestEscapedColons(t *testing.T) {
	t.Skip("not sure it is possible to support this in general case without an SQL parser")
	var qs = `SELECT * FROM testtable WHERE timeposted BETWEEN (now() AT TIME ZONE 'utc') AND
	(now() AT TIME ZONE 'utc') - interval '01:30:00') AND name = '\'this is a test\'' and id = :id`
	_, _, err := compileNamedQuery([]byte(qs), DOLLAR)
	if err != nil {
		t.Error("Didn't handle colons correctly when inside a string")
	}
}

func TestNamedQueries(t *testing.T) {
	RunWithSchema(defaultSchema, t, func(db *DB, t *testing.T, now string) {
		loadDefaultFixture(db, t)
		test := Test{t}
		var ns *NamedStmt
		var err error

		// Check that invalid preparations fail
		ns, err = db.PrepareNamed("SELECT * FROM person WHERE first_name=:first:name")
		if err == nil {
			t.Error("Expected an error with invalid prepared statement.")
		}

		ns, err = db.PrepareNamed("invalid sql")
		if err == nil {
			t.Error("Expected an error with invalid prepared statement.")
		}

		// Check closing works as anticipated
		ns, err = db.PrepareNamed("SELECT * FROM person WHERE first_name=:first_name")
		test.Error(err)
		err = ns.Close()
		test.Error(err)

		ns, err = db.PrepareNamed(`
			SELECT first_name, last_name, email
			FROM person WHERE first_name=:first_name AND email=:email`)
		test.Error(err)

		// test Queryx w/ uses Query
		p := Person{FirstName: "Jason", LastName: "Moiron", Email: "jmoiron@jmoiron.net"}

		rows, err := ns.Queryx(p)
		test.Error(err)
		for rows.Next() {
			var p2 Person
			rows.StructScan(&p2)
			if p.FirstName != p2.FirstName {
				t.Errorf("got %s, expected %s", p.FirstName, p2.FirstName)
			}
			if p.LastName != p2.LastName {
				t.Errorf("got %s, expected %s", p.LastName, p2.LastName)
			}
			if p.Email != p2.Email {
				t.Errorf("got %s, expected %s", p.Email, p2.Email)
			}
		}

		// test Select
		people := make([]Person, 0, 5)
		err = ns.Select(&people, p)
		test.Error(err)

		if len(people) != 1 {
			t.Errorf("got %d results, expected %d", len(people), 1)
		}
		if p.FirstName != people[0].FirstName {
			t.Errorf("got %s, expected %s", p.FirstName, people[0].FirstName)
		}
		if p.LastName != people[0].LastName {
			t.Errorf("got %s, expected %s", p.LastName, people[0].LastName)
		}
		if p.Email != people[0].Email {
			t.Errorf("got %s, expected %s", p.Email, people[0].Email)
		}

		// test struct batch inserts
		sls := []Person{
			{FirstName: "Ardie", LastName: "Savea", Email: "asavea@ab.co.nz"},
			{FirstName: "Sonny Bill", LastName: "Williams", Email: "sbw@ab.co.nz"},
			{FirstName: "Ngani", LastName: "Laumape", Email: "nlaumape@ab.co.nz"},
		}

		insert := fmt.Sprintf(
			"INSERT INTO person (first_name, last_name, email, added_at) VALUES (:first_name, :last_name, :email, %v)\n",
			now,
		)
		_, err = db.NamedExec(insert, sls)
		test.Error(err)

		// test map batch inserts
		slsMap := []map[string]interface{}{
			{"first_name": "Ardie", "last_name": "Savea", "email": "asavea@ab.co.nz"},
			{"first_name": "Sonny Bill", "last_name": "Williams", "email": "sbw@ab.co.nz"},
			{"first_name": "Ngani", "last_name": "Laumape", "email": "nlaumape@ab.co.nz"},
		}

		_, err = db.NamedExec(`INSERT INTO person (first_name, last_name, email)
			VALUES (:first_name, :last_name, :email) ;--`, slsMap)
		test.Error(err)

		type A map[string]interface{}

		typedMap := []A{
			{"first_name": "Ardie", "last_name": "Savea", "email": "asavea@ab.co.nz"},
			{"first_name": "Sonny Bill", "last_name": "Williams", "email": "sbw@ab.co.nz"},
			{"first_name": "Ngani", "last_name": "Laumape", "email": "nlaumape@ab.co.nz"},
		}

		_, err = db.NamedExec(`INSERT INTO person (first_name, last_name, email)
			VALUES (:first_name, :last_name, :email) ;--`, typedMap)
		test.Error(err)

		for _, p := range sls {
			dest := Person{}
			err = db.Get(&dest, db.Rebind("SELECT * FROM person WHERE email=?"), p.Email)
			test.Error(err)
			if dest.Email != p.Email {
				t.Errorf("expected %s, got %s", p.Email, dest.Email)
			}
		}

		// test Exec
		ns, err = db.PrepareNamed(`
			INSERT INTO person (first_name, last_name, email)
			VALUES (:first_name, :last_name, :email)`)
		test.Error(err)

		js := Person{
			FirstName: "Julien",
			LastName:  "Savea",
			Email:     "jsavea@ab.co.nz",
		}
		_, err = ns.Exec(js)
		test.Error(err)

		// Make sure we can pull him out again
		p2 := Person{}
		db.Get(&p2, db.Rebind("SELECT * FROM person WHERE email=?"), js.Email)
		if p2.Email != js.Email {
			t.Errorf("expected %s, got %s", js.Email, p2.Email)
		}

		// test Txn NamedStmts
		tx := db.MustBegin()
		txns := tx.NamedStmt(ns)

		// We're going to add Steven in this txn
		sl := Person{
			FirstName: "Steven",
			LastName:  "Luatua",
			Email:     "sluatua@ab.co.nz",
		}

		_, err = txns.Exec(sl)
		test.Error(err)
		// then rollback...
		tx.Rollback()
		// looking for Steven after a rollback should fail
		err = db.Get(&p2, db.Rebind("SELECT * FROM person WHERE email=?"), sl.Email)
		if err != sql.ErrNoRows {
			t.Errorf("expected no rows error, got %v", err)
		}

		// now do the same, but commit
		tx = db.MustBegin()
		txns = tx.NamedStmt(ns)
		_, err = txns.Exec(sl)
		test.Error(err)
		tx.Commit()

		// looking for Steven after a Commit should succeed
		err = db.Get(&p2, db.Rebind("SELECT * FROM person WHERE email=?"), sl.Email)
		test.Error(err)
		if p2.Email != sl.Email {
			t.Errorf("expected %s, got %s", sl.Email, p2.Email)
		}

	})
}

func TestFixBounds(t *testing.T) {
	table := []struct {
		name, query, expect string
		loop                int
	}{
		{
			name:   `named syntax`,
			query:  `INSERT INTO foo (a,b,c,d) VALUES (:name, :age, :first, :last)`,
			expect: `INSERT INTO foo (a,b,c,d) VALUES (:name, :age, :first, :last),(:name, :age, :first, :last)`,
			loop:   2,
		},
		{
			name:   `mysql syntax`,
			query:  `INSERT INTO foo (a,b,c,d) VALUES (?, ?, ?, ?)`,
			expect: `INSERT INTO foo (a,b,c,d) VALUES (?, ?, ?, ?),(?, ?, ?, ?)`,
			loop:   2,
		},
		{
			name:   `named syntax w/ trailer`,
			query:  `INSERT INTO foo (a,b,c,d) VALUES (:name, :age, :first, :last) ;--`,
			expect: `INSERT INTO foo (a,b,c,d) VALUES (:name, :age, :first, :last),(:name, :age, :first, :last) ;--`,
			loop:   2,
		},
		{
			name:   `mysql syntax w/ trailer`,
			query:  `INSERT INTO foo (a,b,c,d) VALUES (?, ?, ?, ?) ;--`,
			expect: `INSERT INTO foo (a,b,c,d) VALUES (?, ?, ?, ?),(?, ?, ?, ?) ;--`,
			loop:   2,
		},
		{
			name:   `not found test`,
			query:  `INSERT INTO foo (a,b,c,d) (:name, :age, :first, :last)`,
			expect: `INSERT INTO foo (a,b,c,d) (:name, :age, :first, :last)`,
			loop:   2,
		},
		{
			name:   `found twice test`,
			query:  `INSERT INTO foo (a,b,c,d) VALUES (:name, :age, :first, :last) VALUES (:name, :age, :first, :last)`,
			expect: `INSERT INTO foo (a,b,c,d) VALUES (:name, :age, :first, :last) VALUES (:name, :age, :first, :last)`,
			loop:   2,
		},
	}

	for _, tc := range table {
		t.Run(tc.name, func(t *testing.T) {
			res := fixBound(tc.query, tc.loop)
			if res != tc.expect {
				t.Errorf("mismatched results")
			}
		})
	}

}
<|endoftext|>


<commit_before>package jsonrpc

import "github.com/nbusy/neptulon"

// Sender is a JSON-RPC request/notification sending middleware.
type Sender struct {
	pendinRequests map[string]chan *Response
}

// NewSender creates a JSON-RPC sender instance and registers it with the Neptulon JSON-RPC app.
func NewSender(app *App) (*Sender, error) {
	s := Sender{
		pendinRequests: make(map[string]chan *Response),
	}

	app.Middleware(s.middleware)
	return &s, nil
}

// Request sends a JSON-RPC request throught the connection denoted by the session ID.
func (s *Sender) Request(sessionID string, req *Request) chan<- *Response {
	ch := make(chan *Response)
	s.pendinRequests[req.ID] = ch
	return ch
}

// Notification sends a JSON-RPC notification throught the connection denoted by the session ID.
func (s *Sender) Notification(sessionID string, not *Notification) {

}

func (s *Sender) middleware(conn *neptulon.Conn, msg *Message) (result interface{}, resErr *ResError) {
	if ch, ok := s.pendinRequests[msg.ID]; ok {
		ch <- &Response{ID: msg.ID, Result: msg.Result, Error: msg.Error}
		delete(s.pendinRequests, msg.ID)
	}

	return nil, nil
}
<commit_msg>add delegation note<commit_after>package jsonrpc

import (
	"encoding/json"
	"log"

	"github.com/nbusy/neptulon"
)

// Sender is a JSON-RPC request/notification sending middleware.
type Sender struct {
	pendinRequests map[string]chan *Response
}

// NewSender creates a JSON-RPC sender instance and registers it with the Neptulon JSON-RPC app.
func NewSender(app *App) (*Sender, error) {
	s := Sender{
		pendinRequests: make(map[string]chan *Response),
	}

	app.Middleware(s.middleware)
	return &s, nil
}

// Request sends a JSON-RPC request throught the connection denoted by the session ID.
func (s *Sender) Request(sessionID string, req *Request) chan<- *Response {
	data, err := json.Marshal(req)
	if err != nil {
		log.Fatalln("Cannot serialize outgoing request:", err)
	}

	_, err = conn.Write(data) // todo: delegate serialization to jsonrpc app which should delegate writing to neptulon app
	if err != nil {
		log.Fatalln("Errored while writing response to connection:", err)
	}

	ch := make(chan *Response)
	s.pendinRequests[req.ID] = ch
	return ch
}

// Notification sends a JSON-RPC notification throught the connection denoted by the session ID.
func (s *Sender) Notification(sessionID string, not *Notification) {

}

func (s *Sender) middleware(conn *neptulon.Conn, msg *Message) (result interface{}, resErr *ResError) {
	if ch, ok := s.pendinRequests[msg.ID]; ok {
		ch <- &Response{ID: msg.ID, Result: msg.Result, Error: msg.Error}
		delete(s.pendinRequests, msg.ID)
	}

	return nil, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strings"

	"github.com/spf13/cobra"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/shared/api"
	cli "github.com/lxc/lxd/shared/cmd"
	"github.com/lxc/lxd/shared/i18n"

	"github.com/lxc/lxd/shared"
)

type cmdPublish struct {
	global *cmdGlobal

	flagAliases              []string
	flagCompressionAlgorithm string
	flagMakePublic           bool
	flagForce                bool
}

func (c *cmdPublish) showByDefault() bool {
	return true
}

func (c *cmdPublish) Command() *cobra.Command {
	cmd := &cobra.Command{}
	cmd.Use = i18n.G("publish [<remote>:]<container>[/<snapshot>] [<remote>:] [flags] [key=value...]")
	cmd.Short = i18n.G("Publish containers as images")
	cmd.Long = cli.FormatSection(i18n.G("Description"), i18n.G(
		`Publish containers as images`))

	cmd.RunE = c.Run
	cmd.Flags().BoolVar(&c.flagMakePublic, "public", false, i18n.G("Make the image public"))
	cmd.Flags().StringArrayVar(&c.flagAliases, "alias", nil, i18n.G("New alias to define at target")+"``")
	cmd.Flags().BoolVarP(&c.flagForce, "force", "f", false, i18n.G("Stop the container if currently running"))
	cmd.Flags().StringVar(&c.flagCompressionAlgorithm, "compression", "", i18n.G("Define a compression algorithm: for image or none")+"``")

	return cmd
}

func (c *cmdPublish) Run(cmd *cobra.Command, args []string) error {
	conf := c.global.conf

	iName := ""
	iRemote := ""
	properties := map[string]string{}
	firstprop := 1 // first property is arg[2] if arg[1] is image remote, else arg[1]

	// Sanity checks
	exit, err := c.global.CheckArgs(cmd, args, 1, -1)
	if exit {
		return err
	}

	cRemote, cName, err := conf.ParseRemote(args[0])
	if err != nil {
		return err
	}

	if len(args) >= 2 && !strings.Contains(args[1], "=") {
		firstprop = 2
		iRemote, iName, err = conf.ParseRemote(args[1])
		if err != nil {
			return err
		}
	} else {
		iRemote, iName, err = conf.ParseRemote("")
		if err != nil {
			return err
		}
	}

	if cName == "" {
		return fmt.Errorf(i18n.G("Container name is mandatory"))
	}
	if iName != "" {
		return fmt.Errorf(i18n.G("There is no \"image name\".  Did you want an alias?"))
	}

	d, err := conf.GetContainerServer(iRemote)
	if err != nil {
		return err
	}

	s := d
	if cRemote != iRemote {
		s, err = conf.GetContainerServer(cRemote)
		if err != nil {
			return err
		}
	}

	if !shared.IsSnapshot(cName) {
		ct, etag, err := s.GetContainer(cName)
		if err != nil {
			return err
		}

		wasRunning := ct.StatusCode != 0 && ct.StatusCode != api.Stopped
		wasEphemeral := ct.Ephemeral

		if wasRunning {
			if !c.flagForce {
				return fmt.Errorf(i18n.G("The container is currently running. Use --force to have it stopped and restarted"))
			}

			if ct.Ephemeral {
				ct.Ephemeral = false
				op, err := s.UpdateContainer(cName, ct.Writable(), etag)
				if err != nil {
					return err
				}

				err = op.Wait()
				if err != nil {
					return err
				}

				// Refresh the ETag
				_, etag, err = s.GetContainer(cName)
				if err != nil {
					return err
				}
			}

			req := api.ContainerStatePut{
				Action:  string(shared.Stop),
				Timeout: -1,
				Force:   true,
			}

			op, err := s.UpdateContainerState(cName, req, "")
			if err != nil {
				return err
			}

			err = op.Wait()
			if err != nil {
				return fmt.Errorf(i18n.G("Stopping container failed!"))
			}

			defer func() {
				req.Action = string(shared.Start)
				op, err = s.UpdateContainerState(cName, req, "")
				if err != nil {
					return
				}

				op.Wait()
			}()

			if wasEphemeral {
				ct.Ephemeral = true
				op, err := s.UpdateContainer(cName, ct.Writable(), etag)
				if err != nil {
					return err
				}

				err = op.Wait()
				if err != nil {
					return err
				}
			}
		}
	}

	for i := firstprop; i < len(args); i++ {
		entry := strings.SplitN(args[i], "=", 2)
		if len(entry) < 2 {
			return fmt.Errorf(i18n.G("Bad key=value pair: %s"), entry)
		}
		properties[entry[0]] = entry[1]
	}

	// We should only set the properties field if there actually are any.
	// Otherwise we will only delete any existing properties on publish.
	// This is something which only direct callers of the API are allowed to
	// do.
	if len(properties) == 0 {
		properties = nil
	}

	// Reformat aliases
	aliases := []api.ImageAlias{}
	for _, entry := range c.flagAliases {
		alias := api.ImageAlias{}
		alias.Name = entry
		aliases = append(aliases, alias)
	}

	// Create the image
	req := api.ImagesPost{
		Source: &api.ImagesPostSource{
			Type: "container",
			Name: cName,
		},
		CompressionAlgorithm: c.flagCompressionAlgorithm,
	}
	req.Properties = properties

	if shared.IsSnapshot(cName) {
		req.Source.Type = "snapshot"
	}

	if cRemote == iRemote {
		req.Public = c.flagMakePublic
	}

	op, err := s.CreateImage(req, nil)
	if err != nil {
		return err
	}

	err = op.Wait()
	if err != nil {
		return err
	}
	opAPI := op.Get()

	// Grab the fingerprint
	fingerprint := opAPI.Metadata["fingerprint"].(string)

	// For remote publish, copy to target now
	if cRemote != iRemote {
		defer s.DeleteImage(fingerprint)

		// Get the source image
		image, _, err := s.GetImage(fingerprint)
		if err != nil {
			return err
		}

		// Image copy arguments
		args := lxd.ImageCopyArgs{
			Public: c.flagMakePublic,
		}

		// Copy the image to the destination host
		op, err := d.CopyImage(s, *image, &args)
		if err != nil {
			return err
		}

		err = op.Wait()
		if err != nil {
			return err
		}
	}

	err = ensureImageAliases(d, aliases, fingerprint)
	if err != nil {
		return err
	}
	fmt.Printf(i18n.G("Container published with fingerprint: %s")+"\n", fingerprint)

	return nil
}
<commit_msg>lxc/publish: Add progress reporting<commit_after>package main

import (
	"fmt"
	"strings"

	"github.com/spf13/cobra"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxc/utils"
	"github.com/lxc/lxd/shared/api"
	cli "github.com/lxc/lxd/shared/cmd"
	"github.com/lxc/lxd/shared/i18n"

	"github.com/lxc/lxd/shared"
)

type cmdPublish struct {
	global *cmdGlobal

	flagAliases              []string
	flagCompressionAlgorithm string
	flagMakePublic           bool
	flagForce                bool
}

func (c *cmdPublish) showByDefault() bool {
	return true
}

func (c *cmdPublish) Command() *cobra.Command {
	cmd := &cobra.Command{}
	cmd.Use = i18n.G("publish [<remote>:]<container>[/<snapshot>] [<remote>:] [flags] [key=value...]")
	cmd.Short = i18n.G("Publish containers as images")
	cmd.Long = cli.FormatSection(i18n.G("Description"), i18n.G(
		`Publish containers as images`))

	cmd.RunE = c.Run
	cmd.Flags().BoolVar(&c.flagMakePublic, "public", false, i18n.G("Make the image public"))
	cmd.Flags().StringArrayVar(&c.flagAliases, "alias", nil, i18n.G("New alias to define at target")+"``")
	cmd.Flags().BoolVarP(&c.flagForce, "force", "f", false, i18n.G("Stop the container if currently running"))
	cmd.Flags().StringVar(&c.flagCompressionAlgorithm, "compression", "", i18n.G("Define a compression algorithm: for image or none")+"``")

	return cmd
}

func (c *cmdPublish) Run(cmd *cobra.Command, args []string) error {
	conf := c.global.conf

	iName := ""
	iRemote := ""
	properties := map[string]string{}
	firstprop := 1 // first property is arg[2] if arg[1] is image remote, else arg[1]

	// Sanity checks
	exit, err := c.global.CheckArgs(cmd, args, 1, -1)
	if exit {
		return err
	}

	cRemote, cName, err := conf.ParseRemote(args[0])
	if err != nil {
		return err
	}

	if len(args) >= 2 && !strings.Contains(args[1], "=") {
		firstprop = 2
		iRemote, iName, err = conf.ParseRemote(args[1])
		if err != nil {
			return err
		}
	} else {
		iRemote, iName, err = conf.ParseRemote("")
		if err != nil {
			return err
		}
	}

	if cName == "" {
		return fmt.Errorf(i18n.G("Container name is mandatory"))
	}
	if iName != "" {
		return fmt.Errorf(i18n.G("There is no \"image name\".  Did you want an alias?"))
	}

	d, err := conf.GetContainerServer(iRemote)
	if err != nil {
		return err
	}

	s := d
	if cRemote != iRemote {
		s, err = conf.GetContainerServer(cRemote)
		if err != nil {
			return err
		}
	}

	if !shared.IsSnapshot(cName) {
		ct, etag, err := s.GetContainer(cName)
		if err != nil {
			return err
		}

		wasRunning := ct.StatusCode != 0 && ct.StatusCode != api.Stopped
		wasEphemeral := ct.Ephemeral

		if wasRunning {
			if !c.flagForce {
				return fmt.Errorf(i18n.G("The container is currently running. Use --force to have it stopped and restarted"))
			}

			if ct.Ephemeral {
				ct.Ephemeral = false
				op, err := s.UpdateContainer(cName, ct.Writable(), etag)
				if err != nil {
					return err
				}

				err = op.Wait()
				if err != nil {
					return err
				}

				// Refresh the ETag
				_, etag, err = s.GetContainer(cName)
				if err != nil {
					return err
				}
			}

			req := api.ContainerStatePut{
				Action:  string(shared.Stop),
				Timeout: -1,
				Force:   true,
			}

			op, err := s.UpdateContainerState(cName, req, "")
			if err != nil {
				return err
			}

			err = op.Wait()
			if err != nil {
				return fmt.Errorf(i18n.G("Stopping container failed!"))
			}

			defer func() {
				req.Action = string(shared.Start)
				op, err = s.UpdateContainerState(cName, req, "")
				if err != nil {
					return
				}

				op.Wait()
			}()

			if wasEphemeral {
				ct.Ephemeral = true
				op, err := s.UpdateContainer(cName, ct.Writable(), etag)
				if err != nil {
					return err
				}

				err = op.Wait()
				if err != nil {
					return err
				}
			}
		}
	}

	for i := firstprop; i < len(args); i++ {
		entry := strings.SplitN(args[i], "=", 2)
		if len(entry) < 2 {
			return fmt.Errorf(i18n.G("Bad key=value pair: %s"), entry)
		}
		properties[entry[0]] = entry[1]
	}

	// We should only set the properties field if there actually are any.
	// Otherwise we will only delete any existing properties on publish.
	// This is something which only direct callers of the API are allowed to
	// do.
	if len(properties) == 0 {
		properties = nil
	}

	// Reformat aliases
	aliases := []api.ImageAlias{}
	for _, entry := range c.flagAliases {
		alias := api.ImageAlias{}
		alias.Name = entry
		aliases = append(aliases, alias)
	}

	// Create the image
	req := api.ImagesPost{
		Source: &api.ImagesPostSource{
			Type: "container",
			Name: cName,
		},
		CompressionAlgorithm: c.flagCompressionAlgorithm,
	}
	req.Properties = properties

	if shared.IsSnapshot(cName) {
		req.Source.Type = "snapshot"
	}

	if cRemote == iRemote {
		req.Public = c.flagMakePublic
	}

	op, err := s.CreateImage(req, nil)
	if err != nil {
		return err
	}

	// Watch the background operation
	progress := utils.ProgressRenderer{
		Format: i18n.G("Publishing container: %s"),
		Quiet:  c.global.flagQuiet,
	}

	_, err = op.AddHandler(progress.UpdateOp)
	if err != nil {
		progress.Done("")
		return err
	}

	// Wait for the copy to complete
	err = utils.CancelableWait(op, &progress)
	if err != nil {
		progress.Done("")
		return err
	}
	progress.Done("")

	opAPI := op.Get()

	// Grab the fingerprint
	fingerprint := opAPI.Metadata["fingerprint"].(string)

	// For remote publish, copy to target now
	if cRemote != iRemote {
		defer s.DeleteImage(fingerprint)

		// Get the source image
		image, _, err := s.GetImage(fingerprint)
		if err != nil {
			return err
		}

		// Image copy arguments
		args := lxd.ImageCopyArgs{
			Public: c.flagMakePublic,
		}

		// Copy the image to the destination host
		op, err := d.CopyImage(s, *image, &args)
		if err != nil {
			return err
		}

		err = op.Wait()
		if err != nil {
			return err
		}
	}

	err = ensureImageAliases(d, aliases, fingerprint)
	if err != nil {
		return err
	}
	fmt.Printf(i18n.G("Container published with fingerprint: %s")+"\n", fingerprint)

	return nil
}
<|endoftext|>


<commit_before>package plugins

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"

	"github.com/praqma/git-phlow/githandler"
	"github.com/praqma/git-phlow/options"
	"os"
)

//GitHub ...
var GitHub GitHubRequest

//AuthURL ...
var authURL = "https://api.github.com/authorizations"

//RepoURL ...
var repoURL = "https://api.github.com/repos/%s/%s"

//IssueUrl
var issueURL = "https://api.github.com/repos/%s/%s/issues"

//labelURL
var labelURL = "https://api.github.com/repos/%s/%s/issues/%d/labels"

var assigneeURL = "https://api.github.com/repos/%s/%s/issues/%d/assignees"

//GitHubRequest ...
//request object for github
type GitHubRequest struct {
	Assignee AssigneeRequest
	Issue    IssueRequest
	Label    LabelRequest
	Branch   BranchRequest
	Auth     AuthRequest
}

type requestData struct {
	URL    string
	repo   string
	org    string
	client *http.Client
}

//IssueRequest ...
//request for GitHub issues
type IssueRequest struct {
	requestData
}

//AssigneeRequest ...
//request for assignees on GitHub
type AssigneeRequest struct {
	requestData
	token string
}

//BranchRequest ...
//request for default branch on GitHub repository
type BranchRequest struct {
	requestData
}

//LabelRequest ...
//request for getting and setting labels on GitHub issues
type LabelRequest struct {
	requestData
	token string
}

//AuthRequest ...
//request for authentication on GitHub
type AuthRequest struct {
	requestData
}

func init() {
	info, _ := githandler.Remote()
	org := info.Organisation
	repo := info.Repository
	token := githandler.ConfigGet("token", "phlow")

	GitHub.Issue = IssueRequest{requestData{issueURL, repo, org, http.DefaultClient}}
	GitHub.Branch = BranchRequest{requestData{repoURL, repo, org, http.DefaultClient}}
	GitHub.Label = LabelRequest{requestData{labelURL, repo, org, http.DefaultClient}, token}
	GitHub.Assignee = AssigneeRequest{requestData{assigneeURL, repo, org, http.DefaultClient}, token}
	GitHub.Auth = AuthRequest{requestData{authURL, repo, org, http.DefaultClient}}
}

//Get ...
//Get All Issues from GitHub
func (i *IssueRequest) Get() ([]Issues, error) {
	var resp *http.Response
	var err error
	var body []byte

	i.URL = fmt.Sprintf(i.URL, i.org, i.repo)

	token := githandler.ConfigGet("token", "phlow")
	request, err := http.NewRequest("GET", i.URL, nil)
	if err != nil {
		return nil, err
	}
	q := request.URL.Query()
	q.Add("access_token", token)
	request.URL.RawQuery = q.Encode()

	if options.GlobalFlagVerbose {
		fmt.Println("github uri: " + request.URL.String())
	}

	if resp, err = i.client.Do(request); err != nil {
		return nil, err
	}
	if err = requestStatus(resp); err != nil {
		return nil, err
	}

	defer resp.Body.Close()

	if body, err = ioutil.ReadAll(resp.Body); err == nil {
		re := []Issues{}
		if err = json.Unmarshal(body, &re); err != nil {
			return nil, err
		}
		return re, nil
	}
	return nil, err
}

//Permissions ...
//data struct for permissions
type Permissions struct {
	Scopes []string `json:"scopes"`
	Note   string   `json:"note"`
}

//createPermissions ...
func createPermissions() (string, error) {
	hostname, err := os.Hostname()
	if err != nil {
		return "", err
	}

	note := "git phlow " + hostname
	if options.GlobalFlagVerbose {
		fmt.Println("github plugin: " + note)
	}

	perm := Permissions{
		Scopes: []string{"public_repo", "repo", "repo_deployment", "repo:status"},
		Note:   note,
	}
	b2b, err := json.Marshal(&perm)
	if err != nil {
		return "", err
	}
	return string(b2b), nil
}

//Auth ...
//Auth request to github
func (a *AuthRequest) Auth(user, pass string) (string, error) {
	var auth Auth

	perm, err := createPermissions()
	if err != nil {
		return "", err
	}

	if options.GlobalFlagVerbose {
		fmt.Println(a.URL)
	}
	req, err := http.NewRequest("POST", a.URL, bytes.NewBuffer([]byte(perm)))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.SetBasicAuth(user, pass)

	resp, err := a.client.Do(req)
	if err != nil {
		return "", err
	}
	if err = requestStatus(resp); err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if output, err := ioutil.ReadAll(resp.Body); err == nil {
		if err = json.Unmarshal(output, &auth); err != nil {
			return "", err
		}
		return auth.Token, nil
	}

	return "", errors.New("Could not parse json request")
}

//Set ...
//Set assignee on a GitHub Issue
func (a *AssigneeRequest) Set(assignee string, issue int) error {
	var resp *http.Response
	var err error
	a.URL = fmt.Sprintf(a.URL, a.org, a.repo, issue)

	jsonBytes, _ := json.Marshal(Assignee{Assignees: []string{assignee}})
	req, err := http.NewRequest("POST", a.URL, bytes.NewBuffer(jsonBytes))
	if err != nil {
		return err
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "token "+a.token)

	if resp, err = a.client.Do(req); err != nil {
		return err
	}
	if err = requestStatus(resp); err != nil {
		return err
	}
	return nil
}

//Set ...
//Set label on a GitHub issue
func (l *LabelRequest) Set(label string, issue int) ([]Label, error) {
	client := &http.Client{}
	var body = `[ "` + label + `" ]`
	var err error
	var resp *http.Response
	var req *http.Request

	l.URL = fmt.Sprintf(l.URL, l.org, l.repo, issue)
	if options.GlobalFlagVerbose {
		fmt.Println("github uri: " + l.URL)
	}

	if req, err = http.NewRequest("POST", l.URL, bytes.NewBuffer([]byte(body))); err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "token "+l.token)

	resp, err = client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if err = requestStatus(resp); err != nil {
		return nil, err
	}

	var output []byte
	if output, err = ioutil.ReadAll(resp.Body); err == nil {
		re := []Label{}
		if err = json.Unmarshal(output, &re); err == nil {
			return re, nil
		}
	}
	return nil, err
}

//Default ...
//Get default branch of a GitHub issue
func (b *BranchRequest) Default() (string, error) {
	var resp *http.Response
	var err error
	var body []byte

	b.URL = fmt.Sprintf(b.URL, b.org, b.repo)
	if options.GlobalFlagVerbose {
		fmt.Println("github uri: " + b.URL)
	}
	if resp, err = http.Get(b.URL); err != nil {
		return "", err
	}
	if err = requestStatus(resp); err != nil {
		return "", err
	}

	defer resp.Body.Close()

	if body, err = ioutil.ReadAll(resp.Body); err == nil {
		re := Repo{}
		if err = json.Unmarshal(body, &re); err != nil {
			return "", err
		}
		return re.DefaultBranch, nil
	}
	return "", err

}

//requestStatus ...
//Determines the status of the request
func requestStatus(res *http.Response) error {
	switch res.StatusCode {
	case http.StatusCreated:
		return nil
	case http.StatusOK:
		return nil
	case http.StatusUnprocessableEntity:
		//For POST requests
		return errors.New("git-phlow token already exists")
	case http.StatusNotFound:
		//For GET and POST
		return fmt.Errorf("responded with %s - malformed url", res.Status)
	default:
		//Default behaviour if status is not OK
		return fmt.Errorf("request did not respond with 200 OK, but %s", res.Status)
	}
}
<commit_msg>close #73 err when establishing build environment<commit_after>package plugins

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"

	"github.com/praqma/git-phlow/githandler"
	"github.com/praqma/git-phlow/options"
	"os"
)

//GitHub ...
var GitHub GitHubRequest

//AuthURL ...
var authURL = "https://api.github.com/authorizations"

//RepoURL ...
var repoURL = "https://api.github.com/repos/%s/%s"

//IssueUrl
var issueURL = "https://api.github.com/repos/%s/%s/issues"

//labelURL
var labelURL = "https://api.github.com/repos/%s/%s/issues/%d/labels"

var assigneeURL = "https://api.github.com/repos/%s/%s/issues/%d/assignees"

//GitHubRequest ...
//request object for github
type GitHubRequest struct {
	Assignee AssigneeRequest
	Issue    IssueRequest
	Label    LabelRequest
	Branch   BranchRequest
	Auth     AuthRequest
}

type requestData struct {
	URL    string
	repo   string
	org    string
	client *http.Client
}

//IssueRequest ...
//request for GitHub issues
type IssueRequest struct {
	requestData
}

//AssigneeRequest ...
//request for assignees on GitHub
type AssigneeRequest struct {
	requestData
	token string
}

//BranchRequest ...
//request for default branch on GitHub repository
type BranchRequest struct {
	requestData
}

//LabelRequest ...
//request for getting and setting labels on GitHub issues
type LabelRequest struct {
	requestData
	token string
}

//AuthRequest ...
//request for authentication on GitHub
type AuthRequest struct {
	requestData
}

func init() {
	info, _ := githandler.Remote()
	org := info.Organisation
	repo := info.Repository
	token := githandler.ConfigGet("token", "phlow")

	GitHub.Issue = IssueRequest{requestData{issueURL, repo, org, http.DefaultClient}}
	GitHub.Branch = BranchRequest{requestData{repoURL, repo, org, http.DefaultClient}}
	GitHub.Label = LabelRequest{requestData{labelURL, repo, org, http.DefaultClient}, token}
	GitHub.Assignee = AssigneeRequest{requestData{assigneeURL, repo, org, http.DefaultClient}, token}
	GitHub.Auth = AuthRequest{requestData{authURL, repo, org, http.DefaultClient}}
}

//Get ...
//Get All Issues from GitHub
func (i *IssueRequest) Get() ([]Issues, error) {
	var resp *http.Response
	var err error
	var body []byte

	i.URL = fmt.Sprintf(i.URL, i.org, i.repo)

	token := githandler.ConfigGet("token", "phlow")
	request, err := http.NewRequest("GET", i.URL, nil)
	if err != nil {
		return nil, err
	}
	q := request.URL.Query()
	q.Add("access_token", token)
	request.URL.RawQuery = q.Encode()

	if options.GlobalFlagVerbose {
		fmt.Println("github uri: " + request.URL.String())
	}

	if resp, err = i.client.Do(request); err != nil {
		return nil, err
	}
	if err = requestStatus(resp); err != nil {
		return nil, err
	}

	defer resp.Body.Close()

	if body, err = ioutil.ReadAll(resp.Body); err == nil {
		re := []Issues{}
		if err = json.Unmarshal(body, &re); err != nil {
			return nil, err
		}
		return re, nil
	}
	return nil, err
}

//Permissions ...
//data struct for permissions
type Permissions struct {
	Scopes []string `json:"scopes"`
	Note   string   `json:"note"`
}

//createPermissions ...
func createPermissions() (string, error) {
	hostname, err := os.Hostname()
	if err != nil {
		return "", err
	}

	note := "git phlow " + hostname
	if options.GlobalFlagVerbose {
		fmt.Println("github plugin: " + note)
	}

	perm := Permissions{
		Scopes: []string{"public_repo", "repo", "repo_deployment", "repo:status"},
		Note:   note,
	}
	b2b, err := json.Marshal(&perm)
	if err != nil {
		return "", err
	}
	return string(b2b), nil
}

//Auth ...
//Auth request to github
func (a *AuthRequest) Auth(user, pass string) (string, error) {
	var auth Auth

	perm, err := createPermissions()
	if err != nil {
		return "", err
	}

	if options.GlobalFlagVerbose {
		fmt.Println(a.URL)
	}
	req, err := http.NewRequest("POST", a.URL, bytes.NewBuffer([]byte(perm)))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.SetBasicAuth(user, pass)

	resp, err := a.client.Do(req)
	if err != nil {
		return "", err
	}
	if err = requestStatus(resp); err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if output, err := ioutil.ReadAll(resp.Body); err == nil {
		if err = json.Unmarshal(output, &auth); err != nil {
			return "", err
		}
		return auth.Token, nil
	}

	return "", errors.New("Could not parse json request")
}

//Set ...
//Set assignee on a GitHub Issue
func (a *AssigneeRequest) Set(assignee string, issue int) error {
	var resp *http.Response
	var err error
	a.URL = fmt.Sprintf(a.URL, a.org, a.repo, issue)

	jsonBytes, _ := json.Marshal(Assignee{Assignees: []string{assignee}})
	req, err := http.NewRequest("POST", a.URL, bytes.NewBuffer(jsonBytes))
	if err != nil {
		return err
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "token "+a.token)

	if resp, err = a.client.Do(req); err != nil {
		return err
	}
	if err = requestStatus(resp); err != nil {
		return err
	}
	return nil
}

//Set ...
//Set label on a GitHub issue
func (l *LabelRequest) Set(label string, issue int) ([]Label, error) {
	client := &http.Client{}
	var body = `[ "` + label + `" ]`
	var err error
	var resp *http.Response
	var req *http.Request

	l.URL = fmt.Sprintf(l.URL, l.org, l.repo, issue)
	if options.GlobalFlagVerbose {
		fmt.Println("github uri: " + l.URL)
	}

	if req, err = http.NewRequest("POST", l.URL, bytes.NewBuffer([]byte(body))); err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "token "+l.token)

	resp, err = client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if err = requestStatus(resp); err != nil {
		return nil, err
	}

	var output []byte
	if output, err = ioutil.ReadAll(resp.Body); err == nil {
		re := []Label{}
		if err = json.Unmarshal(output, &re); err == nil {
			return re, nil
		}
	}
	return nil, err
}

//Default ...
//Get default branch of a GitHub issue
func (b *BranchRequest) Default() (string, error) {
	var resp *http.Response
	var err error
	var body []byte

	b.URL = fmt.Sprintf(b.URL, b.org, b.repo)
	if options.GlobalFlagVerbose {
		fmt.Println("github uri: " + b.URL)
	}
	if resp, err = http.Get(b.URL); err != nil {
		return "", err
	}
	if err = requestStatus(resp); err != nil {
		return "", err
	}

	defer resp.Body.Close()

	if body, err = ioutil.ReadAll(resp.Body); err == nil {
		re := Repo{}
		if err = json.Unmarshal(body, &re); err != nil {
			return "", err
		}
		return re.DefaultBranch, nil
	}
	return "", err

}

//requestStatus ...
//Determines the status of the request
func requestStatus(res *http.Response) error {
	switch res.StatusCode {
	case http.StatusCreated:
		return nil
	case http.StatusOK:
		return nil
	case 422: //StatusUnprocessableEntity
		//For POST requests
		return errors.New("git-phlow token already exists")
	case http.StatusNotFound:
		//For GET and POST
		return fmt.Errorf("responded with %s - malformed url", res.Status)
	default:
		//Default behaviour if status is not OK
		return fmt.Errorf("request did not respond with 200 OK, but %s", res.Status)
	}
}
<|endoftext|>


<commit_before>package poll

import (
	"github.com/marcw/poller/check"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

type successPollHandler struct {
}
type errorPollHandler struct {
}
type timeoutPollHandler struct {
}

func (p successPollHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
}

func (p errorPollHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	http.Error(w, "error", 500)
}

func (p timeoutPollHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	time.Sleep(200 * time.Millisecond)
}

func TestSuccessfullPoll(t *testing.T) {
	server := httptest.NewServer(successPollHandler{})
	defer server.Close()

	poll := NewHttpPoller("foobar", 10*time.Second)

	c, _ := check.NewCheck(server.URL, "foobar", "10s", make(map[string]string))
	event := poll.Poll(c)
	if event.StatusCode != 200 {
		t.Error("statusCode should be 200")
	}
	if event.IsUp() != true {
		t.Error("IsUp() should be true")
	}
	if event.Duration.Nanoseconds() == 0 {
		t.Error("Duration can't be equals to 0 nanosecond")
	}
}

func TestFailedPoll(t *testing.T) {
	server := httptest.NewServer(errorPollHandler{})
	defer server.Close()

	poll := NewHttpPoller("foobar", 10*time.Second)

	c, _ := check.NewCheck(server.URL, "foobar", "10s", make(map[string]string))
	event := poll.Poll(c)
	if event.StatusCode != 500 {
		t.Error("statusCode should be 500")
	}
	if event.IsUp() != false {
		t.Error("IsUp() should be false")
	}
}

func TestTimeoutedPoll(t *testing.T) {
	server := httptest.NewServer(timeoutPollHandler{})
	defer server.Close()

	poll := NewHttpPoller("foobar", 100*time.Millisecond)

	c, _ := check.NewCheck(server.URL, "foobar", "10s", make(map[string]string))
	event := poll.Poll(c)
	if event.StatusCode != 0 {
		t.Error("statusCode should be 0")
	}
	if event.IsUp() != false {
		t.Error("IsUp() should be false")
	}
}
<commit_msg>Fixes test<commit_after>package poll

import (
	"github.com/marcw/poller/check"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

type successPollHandler struct {
}
type errorPollHandler struct {
}
type timeoutPollHandler struct {
}

func (p successPollHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
}

func (p errorPollHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	http.Error(w, "error", 500)
}

func (p timeoutPollHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	time.Sleep(200 * time.Millisecond)
}

func TestSuccessfullPoll(t *testing.T) {
	server := httptest.NewServer(successPollHandler{})
	defer server.Close()

	poll := NewHttpPoller("foobar", 10*time.Second)

	c, _ := check.NewCheck(server.URL, "foobar", "10s", false, "", make(map[string]string))
	event := poll.Poll(c)
	if event.StatusCode != 200 {
		t.Error("statusCode should be 200")
	}
	if event.IsUp() != true {
		t.Error("IsUp() should be true")
	}
	if event.Duration.Nanoseconds() == 0 {
		t.Error("Duration can't be equals to 0 nanosecond")
	}
}

func TestFailedPoll(t *testing.T) {
	server := httptest.NewServer(errorPollHandler{})
	defer server.Close()

	poll := NewHttpPoller("foobar", 10*time.Second)

	c, _ := check.NewCheck(server.URL, "foobar", "10s", false, "", make(map[string]string))
	event := poll.Poll(c)
	if event.StatusCode != 500 {
		t.Error("statusCode should be 500")
	}
	if event.IsUp() != false {
		t.Error("IsUp() should be false")
	}
}

func TestTimeoutedPoll(t *testing.T) {
	server := httptest.NewServer(timeoutPollHandler{})
	defer server.Close()

	poll := NewHttpPoller("foobar", 100*time.Millisecond)

	c, _ := check.NewCheck(server.URL, "foobar", "10s", false, "", make(map[string]string))
	event := poll.Poll(c)
	if event.StatusCode != 0 {
		t.Error("statusCode should be 0")
	}
	if event.IsUp() != false {
		t.Error("IsUp() should be false")
	}
}
<|endoftext|>


<commit_before>// Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

package gerror

// Reserved internal error code of framework: code < 1000.

const (
	CodeNil                  = -1 // No error code specified.
	CodeOk                   = 0  // It is OK without error.
	CodeInternalError        = 50 // An error occurred internally.
	CodeValidationFailed     = 51 // Data validation failed.
	CodeDbOperationError     = 52 // Database operation error.
	CodeInvalidParameter     = 53 // The given parameter for current operation is invalid.
	CodeMissingParameter     = 54 // Parameter for current operation is missing.
	CodeInvalidOperation     = 55 // The function cannot be used like this.
	CodeInvalidConfiguration = 56 // The configuration is invalid for current operation.
	CodeMissingConfiguration = 57 // The configuration is missing for current operation.
	CodeNotImplemented       = 58 // The operation is not implemented yet.
	CodeNotSupported         = 59 // The operation is not supported yet.
	CodeOperationFailed      = 60 // I tried, but I cannot give you what you want.
)
<commit_msg>improve error code for package gerror<commit_after>// Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

package gerror

// Reserved internal error code of framework: code < 1000.

const (
	// ===============================================================================
	// Common system codes.
	// ===============================================================================

	CodeNil                  = -1        // No error code specified.
	CodeOk                   = 0         // It is OK.
	CodeInternalError        = 50 + iota // An error occurred internally.
	CodeValidationFailed                 // Data validation failed.
	CodeDbOperationError                 // Database operation error.
	CodeInvalidParameter                 // The given parameter for current operation is invalid.
	CodeMissingParameter                 // Parameter for current operation is missing.
	CodeInvalidOperation                 // The function cannot be used like this.
	CodeInvalidConfiguration             // The configuration is invalid for current operation.
	CodeMissingConfiguration             // The configuration is missing for current operation.
	CodeNotImplemented                   // The operation is not implemented yet.
	CodeNotSupported                     // The operation is not supported yet.
	CodeOperationFailed                  // I tried, but I cannot give you what you want.
	CodeNotAuthorized                    // Not Authorized.
	CodeSecurityReason                   // Security Reason.
	CodeServerBusy                       // Server is busy, please try again later.
	CodeUnknown                          // Unknown error.
	CodeResourceNotExist                 // Resource does not exist.

	// ===============================================================================
	// Common business codes.
	// ===============================================================================

	CodeBusinessValidationFailed = 300 + iota // Business validation failed.
)
<|endoftext|>


<commit_before>/*
Copyright 2019 The Jetstack cert-manager contributors.

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 venafi

import (
	"context"

	"github.com/Venafi/vcert/pkg/endpoint"
	k8sErrors "k8s.io/apimachinery/pkg/api/errors"
	corelisters "k8s.io/client-go/listers/core/v1"

	apiutil "github.com/jetstack/cert-manager/pkg/api/util"
	cmapi "github.com/jetstack/cert-manager/pkg/apis/certmanager/v1alpha2"
	controllerpkg "github.com/jetstack/cert-manager/pkg/controller"
	"github.com/jetstack/cert-manager/pkg/controller/certificaterequests"
	crutil "github.com/jetstack/cert-manager/pkg/controller/certificaterequests/util"
	venafiinternal "github.com/jetstack/cert-manager/pkg/internal/venafi"
	issuerpkg "github.com/jetstack/cert-manager/pkg/issuer"
	logf "github.com/jetstack/cert-manager/pkg/logs"
)

const (
	CRControllerName = "certificaterequests-issuer-venafi"
)

type Venafi struct {
	issuerOptions controllerpkg.IssuerOptions
	secretsLister corelisters.SecretLister
	reporter      *crutil.Reporter

	clientBuilder venafiinternal.VenafiClientBuilder
}

func init() {
	// create certificate request controller for venafi issuer
	controllerpkg.Register(CRControllerName, func(ctx *controllerpkg.Context) (controllerpkg.Interface, error) {
		return controllerpkg.NewBuilder(ctx, CRControllerName).
			For(certificaterequests.New(apiutil.IssuerVenafi, NewVenafi(ctx))).
			Complete()
	})
}

func NewVenafi(ctx *controllerpkg.Context) *Venafi {
	return &Venafi{
		issuerOptions: ctx.IssuerOptions,
		secretsLister: ctx.KubeSharedInformerFactory.Core().V1().Secrets().Lister(),
		reporter:      crutil.NewReporter(ctx.Clock, ctx.Recorder),
		clientBuilder: venafiinternal.New,
	}
}

func (v *Venafi) Sign(ctx context.Context, cr *cmapi.CertificateRequest, issuerObj cmapi.GenericIssuer) (*issuerpkg.IssueResponse, error) {
	log := logf.FromContext(ctx, "sign")
	log = logf.WithRelatedResource(log, issuerObj)

	client, err := v.clientBuilder(cr.Namespace, v.secretsLister, issuerObj)
	if k8sErrors.IsNotFound(err) {
		message := "Required secret resource not found"

		v.reporter.Pending(cr, err, "SecretMissing", message)
		log.Error(err, message)

		return nil, nil
	}

	if err != nil {
		message := "Failed to initialise venafi client for signing"

		v.reporter.Pending(cr, err, "VenafiInitError", message)
		log.Error(err, message)

		return nil, err
	}

	duration := apiutil.DefaultCertDuration(cr.Spec.Duration)

	certPem, err := client.Sign(cr.Spec.CSRPEM, duration)

	// Check some known error types
	if err != nil {
		switch err.(type) {

		case endpoint.ErrCertificatePending:
			message := "Venafi certificate still in a pending state, the request will be retried"

			v.reporter.Pending(cr, err, "IssuancePending", message)
			log.Error(err, message)
			return nil, err

		case endpoint.ErrRetrieveCertificateTimeout:
			message := "Timed out waiting for venafi certificate, the request will be retried"

			v.reporter.Failed(cr, err, "Timeout", message)
			log.Error(err, message)
			return nil, nil

		default:
			message := "Failed to obtain venafi certificate"

			v.reporter.Failed(cr, err, "RetrieveError", message)
			log.Error(err, message)

			return nil, err
		}
	}

	log.Info("certificate issued")

	return &issuerpkg.IssueResponse{
		Certificate: certPem,
	}, nil
}
<commit_msg>fixed which namespace to find secret in based on if its a clusterissuer for venafi<commit_after>/*
Copyright 2019 The Jetstack cert-manager contributors.

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 venafi

import (
	"context"

	"github.com/Venafi/vcert/pkg/endpoint"
	k8sErrors "k8s.io/apimachinery/pkg/api/errors"
	corelisters "k8s.io/client-go/listers/core/v1"

	apiutil "github.com/jetstack/cert-manager/pkg/api/util"
	cmapi "github.com/jetstack/cert-manager/pkg/apis/certmanager/v1alpha2"
	controllerpkg "github.com/jetstack/cert-manager/pkg/controller"
	"github.com/jetstack/cert-manager/pkg/controller/certificaterequests"
	crutil "github.com/jetstack/cert-manager/pkg/controller/certificaterequests/util"
	venafiinternal "github.com/jetstack/cert-manager/pkg/internal/venafi"
	issuerpkg "github.com/jetstack/cert-manager/pkg/issuer"
	logf "github.com/jetstack/cert-manager/pkg/logs"
)

const (
	CRControllerName = "certificaterequests-issuer-venafi"
)

type Venafi struct {
	issuerOptions controllerpkg.IssuerOptions
	secretsLister corelisters.SecretLister
	reporter      *crutil.Reporter

	clientBuilder venafiinternal.VenafiClientBuilder
}

func init() {
	// create certificate request controller for venafi issuer
	controllerpkg.Register(CRControllerName, func(ctx *controllerpkg.Context) (controllerpkg.Interface, error) {
		return controllerpkg.NewBuilder(ctx, CRControllerName).
			For(certificaterequests.New(apiutil.IssuerVenafi, NewVenafi(ctx))).
			Complete()
	})
}

func NewVenafi(ctx *controllerpkg.Context) *Venafi {
	return &Venafi{
		issuerOptions: ctx.IssuerOptions,
		secretsLister: ctx.KubeSharedInformerFactory.Core().V1().Secrets().Lister(),
		reporter:      crutil.NewReporter(ctx.Clock, ctx.Recorder),
		clientBuilder: venafiinternal.New,
	}
}

func (v *Venafi) Sign(ctx context.Context, cr *cmapi.CertificateRequest, issuerObj cmapi.GenericIssuer) (*issuerpkg.IssueResponse, error) {
	log := logf.FromContext(ctx, "sign")
	log = logf.WithRelatedResource(log, issuerObj)
	ns := issuerObj.GetObjectMeta().Namespace
	if ns == "" {
		ns = v.issuerOptions.ClusterResourceNamespace
	}

	client, err := v.clientBuilder(ns, v.secretsLister, issuerObj)
	if k8sErrors.IsNotFound(err) {
		message := "Required secret resource not found"

		v.reporter.Pending(cr, err, "SecretMissing", message)
		log.Error(err, message)

		return nil, nil
	}

	if err != nil {
		message := "Failed to initialise venafi client for signing"

		v.reporter.Pending(cr, err, "VenafiInitError", message)
		log.Error(err, message)

		return nil, err
	}

	duration := apiutil.DefaultCertDuration(cr.Spec.Duration)

	certPem, err := client.Sign(cr.Spec.CSRPEM, duration)

	// Check some known error types
	if err != nil {
		switch err.(type) {

		case endpoint.ErrCertificatePending:
			message := "Venafi certificate still in a pending state, the request will be retried"

			v.reporter.Pending(cr, err, "IssuancePending", message)
			log.Error(err, message)
			return nil, err

		case endpoint.ErrRetrieveCertificateTimeout:
			message := "Timed out waiting for venafi certificate, the request will be retried"

			v.reporter.Failed(cr, err, "Timeout", message)
			log.Error(err, message)
			return nil, nil

		default:
			message := "Failed to obtain venafi certificate"

			v.reporter.Failed(cr, err, "RetrieveError", message)
			log.Error(err, message)

			return nil, err
		}
	}

	log.Info("certificate issued")

	return &issuerpkg.IssueResponse{
		Certificate: certPem,
	}, nil
}
<|endoftext|>