<commit_before>package sshchat

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

	"github.com/shazow/rateio"
	"github.com/shazow/ssh-chat/chat"
	"github.com/shazow/ssh-chat/chat/message"
	"github.com/shazow/ssh-chat/sshd"
)

const maxInputLength int = 1024

// GetPrompt will render the terminal prompt string based on the user.
func GetPrompt(user *message.User) string {
	name := user.Name()
	if user.Config.Theme != nil {
		name = user.Config.Theme.ColorName(user)
	}
	return fmt.Sprintf("[%s] ", name)
}

// Host is the bridge between sshd and chat modules
// TODO: Should be easy to add support for multiple rooms, if we want.
type Host struct {
	*chat.Room
	listener *sshd.SSHListener
	commands chat.Commands
	auth     *Auth

	// Version string to print on /version
	Version string

	// Default theme
	theme message.Theme

	mu    sync.Mutex
	motd  string
	count int
}

// NewHost creates a Host on top of an existing listener.
func NewHost(listener *sshd.SSHListener, auth *Auth) *Host {
	room := chat.NewRoom()
	h := Host{
		Room:     room,
		listener: listener,
		commands: chat.Commands{},
		auth:     auth,
	}

	// Make our own commands registry instance.
	chat.InitCommands(&h.commands)
	h.InitCommands(&h.commands)
	room.SetCommands(h.commands)

	go room.Serve()
	return &h
}

// SetTheme sets the default theme for the host.
func (h *Host) SetTheme(theme message.Theme) {
	h.mu.Lock()
	h.theme = theme
	h.mu.Unlock()
}

// SetMotd sets the host's message of the day.
func (h *Host) SetMotd(motd string) {
	h.mu.Lock()
	h.motd = motd
	h.mu.Unlock()
}

func (h *Host) isOp(conn sshd.Connection) bool {
	key := conn.PublicKey()
	if key == nil {
		return false
	}
	return h.auth.IsOp(key)
}

// Connect a specific Terminal to this host and its room.
func (h *Host) Connect(term *sshd.Terminal) {
	id := NewIdentity(term.Conn)
	user := message.NewUserScreen(id, term)
	user.Config.Theme = &h.theme
	go user.Consume()

	// Close term once user is closed.
	defer user.Close()
	defer term.Close()

	h.mu.Lock()
	motd := h.motd
	count := h.count
	h.count++
	h.mu.Unlock()

	// Send MOTD
	if motd != "" {
		go user.Send(message.NewAnnounceMsg(motd))
	}

	member, err := h.Join(user)
	if err != nil {
		// Try again...
		id.SetName(fmt.Sprintf("Guest%d", count))
		member, err = h.Join(user)
	}
	if err != nil {
		logger.Errorf("Failed to join: %s", err)
		return
	}

	// Successfully joined.
	term.SetPrompt(GetPrompt(user))
	// FIXME: Re-enable once https://github.com/shazow/ssh-chat/issues/166 is fixed.
	//term.AutoCompleteCallback = h.AutoCompleteFunction(user)
	user.SetHighlight(user.Name())

	// Should the user be op'd on join?
	if h.isOp(term.Conn) {
		h.Room.Ops.Add(member)
	}
	ratelimit := rateio.NewSimpleLimiter(3, time.Second*3)

	logger.Debugf("Joined: %s", user.Name())

	for {
		line, err := term.ReadLine()
		if err == io.EOF {
			// Closed
			break
		} else if err != nil {
			logger.Errorf("Terminal reading error: %s", err)
			break
		}

		err = ratelimit.Count(1)
		if err != nil {
			user.Send(message.NewSystemMsg("Message rejected: Rate limiting is in effect.", user))
			continue
		}
		if len(line) > maxInputLength {
			user.Send(message.NewSystemMsg("Message rejected: Input too long.", user))
			continue
		}
		if line == "" {
			// Silently ignore empty lines.
			continue
		}

		m := message.ParseInput(line, user)

		// FIXME: Any reason to use h.room.Send(m) instead?
		h.HandleMsg(m)

		cmd := m.Command()
		if cmd == "/nick" || cmd == "/theme" {
			// Hijack /nick command to update terminal synchronously. Wouldn't
			// work if we use h.room.Send(m) above.
			//
			// FIXME: This is hacky, how do we improve the API to allow for
			// this? Chat module shouldn't know about terminals.
			term.SetPrompt(GetPrompt(user))
			user.SetHighlight(user.Name())
		}
	}

	err = h.Leave(user)
	if err != nil {
		logger.Errorf("Failed to leave: %s", err)
		return
	}
	logger.Debugf("Leaving: %s", user.Name())
}

// Serve our chat room onto the listener
func (h *Host) Serve() {
	h.listener.HandlerFunc = h.Connect
	h.listener.Serve()
}

func (h *Host) completeName(partial string) string {
	names := h.NamesPrefix(partial)
	if len(names) == 0 {
		// Didn't find anything
		return ""
	}

	return names[len(names)-1]
}

func (h *Host) completeCommand(partial string) string {
	for cmd := range h.commands {
		if strings.HasPrefix(cmd, partial) {
			return cmd
		}
	}
	return ""
}

// AutoCompleteFunction returns a callback for terminal autocompletion
func (h *Host) AutoCompleteFunction(u *message.User) func(line string, pos int, key rune) (newLine string, newPos int, ok bool) {
	return func(line string, pos int, key rune) (newLine string, newPos int, ok bool) {
		if key != 9 {
			return
		}

		if line == "" || strings.HasSuffix(line[:pos], " ") {
			// Don't autocomplete spaces.
			return
		}

		fields := strings.Fields(line[:pos])
		isFirst := len(fields) < 2
		partial := fields[len(fields)-1]
		posPartial := pos - len(partial)

		var completed string
		if isFirst && strings.HasPrefix(partial, "/") {
			// Command
			completed = h.completeCommand(partial)
			if completed == "/reply" {
				replyTo := u.ReplyTo()
				if replyTo != nil {
					completed = "/msg " + replyTo.Name()
				}
			}
		} else {
			// Name
			completed = h.completeName(partial)
			if completed == "" {
				return
			}
			if isFirst {
				completed += ":"
			}
		}
		completed += " "

		// Reposition the cursor
		newLine = strings.Replace(line[posPartial:], partial, completed, 1)
		newLine = line[:posPartial] + newLine
		newPos = pos + (len(completed) - len(partial))
		ok = true
		return
	}
}

// GetUser returns a message.User based on a name.
func (h *Host) GetUser(name string) (*message.User, bool) {
	m, ok := h.MemberById(name)
	if !ok {
		return nil, false
	}
	return m.User, true
}

// InitCommands adds host-specific commands to a Commands container. These will
// override any existing commands.
func (h *Host) InitCommands(c *chat.Commands) {
	c.Add(chat.Command{
		Prefix:     "/msg",
		PrefixHelp: "USER MESSAGE",
		Help:       "Send MESSAGE to USER.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			args := msg.Args()
			switch len(args) {
			case 0:
				return errors.New("must specify user")
			case 1:
				return errors.New("must specify message")
			}

			target, ok := h.GetUser(args[0])
			if !ok {
				return errors.New("user not found")
			}

			m := message.NewPrivateMsg(strings.Join(args[1:], " "), msg.From(), target)
			room.Send(&m)
			return nil
		},
	})

	c.Add(chat.Command{
		Prefix:     "/reply",
		PrefixHelp: "MESSAGE",
		Help:       "Reply with MESSAGE to the previous private message.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			args := msg.Args()
			switch len(args) {
			case 0:
				return errors.New("must specify message")
			}

			target := msg.From().ReplyTo()
			if target == nil {
				return errors.New("no message to reply to")
			}

			m := message.NewPrivateMsg(strings.Join(args, " "), msg.From(), target)
			room.Send(&m)
			return nil
		},
	})

	c.Add(chat.Command{
		Prefix:     "/whois",
		PrefixHelp: "USER",
		Help:       "Information about USER.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			args := msg.Args()
			if len(args) == 0 {
				return errors.New("must specify user")
			}

			target, ok := h.GetUser(args[0])
			if !ok {
				return errors.New("user not found")
			}

			id := target.Identifier.(*Identity)
			room.Send(message.NewSystemMsg(id.Whois(), msg.From()))

			return nil
		},
	})

	// Hidden commands
	c.Add(chat.Command{
		Prefix: "/version",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			room.Send(message.NewSystemMsg(h.Version, msg.From()))
			return nil
		},
	})

	timeStarted := time.Now()
	c.Add(chat.Command{
		Prefix: "/uptime",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			room.Send(message.NewSystemMsg(time.Now().Sub(timeStarted).String(), msg.From()))
			return nil
		},
	})

	// Op commands
	c.Add(chat.Command{
		Op:         true,
		Prefix:     "/kick",
		PrefixHelp: "USER",
		Help:       "Kick USER from the server.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			if !room.IsOp(msg.From()) {
				return errors.New("must be op")
			}

			args := msg.Args()
			if len(args) == 0 {
				return errors.New("must specify user")
			}

			target, ok := h.GetUser(args[0])
			if !ok {
				return errors.New("user not found")
			}

			body := fmt.Sprintf("%s was kicked by %s.", target.Name(), msg.From().Name())
			room.Send(message.NewAnnounceMsg(body))
			target.Close()
			return nil
		},
	})

	c.Add(chat.Command{
		Op:         true,
		Prefix:     "/ban",
		PrefixHelp: "USER [DURATION]",
		Help:       "Ban USER from the server.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			// TODO: Would be nice to specify what to ban. Key? Ip? etc.
			if !room.IsOp(msg.From()) {
				return errors.New("must be op")
			}

			args := msg.Args()
			if len(args) == 0 {
				return errors.New("must specify user")
			}

			target, ok := h.GetUser(args[0])
			if !ok {
				return errors.New("user not found")
			}

			var until time.Duration = 0
			if len(args) > 1 {
				until, _ = time.ParseDuration(args[1])
			}

			id := target.Identifier.(*Identity)
			h.auth.Ban(id.PublicKey(), until)
			h.auth.BanAddr(id.RemoteAddr(), until)

			body := fmt.Sprintf("%s was banned by %s.", target.Name(), msg.From().Name())
			room.Send(message.NewAnnounceMsg(body))
			target.Close()

			logger.Debugf("Banned: \n-> %s", id.Whois())

			return nil
		},
	})

	c.Add(chat.Command{
		Op:         true,
		Prefix:     "/motd",
		PrefixHelp: "MESSAGE",
		Help:       "Set the MESSAGE of the day.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			if !room.IsOp(msg.From()) {
				return errors.New("must be op")
			}

			motd := ""
			args := msg.Args()
			if len(args) > 0 {
				motd = strings.Join(args, " ")
			}

			h.motd = motd
			body := fmt.Sprintf("New message of the day set by %s:", msg.From().Name())
			room.Send(message.NewAnnounceMsg(body))
			if motd != "" {
				room.Send(message.NewAnnounceMsg(motd))
			}

			return nil
		},
	})

	c.Add(chat.Command{
		Op:         true,
		Prefix:     "/op",
		PrefixHelp: "USER [DURATION]",
		Help:       "Set USER as admin.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			if !room.IsOp(msg.From()) {
				return errors.New("must be op")
			}

			args := msg.Args()
			if len(args) == 0 {
				return errors.New("must specify user")
			}

			var until time.Duration = 0
			if len(args) > 1 {
				until, _ = time.ParseDuration(args[1])
			}

			member, ok := room.MemberById(args[0])
			if !ok {
				return errors.New("user not found")
			}
			room.Ops.Add(member)
			id := member.Identifier.(*Identity)
			h.auth.Op(id.PublicKey(), until)

			body := fmt.Sprintf("Made op by %s.", msg.From().Name())
			room.Send(message.NewSystemMsg(body, member.User))

			return nil
		},
	})
}
<commit_msg>Fixed autocomplete crashing bug (#166)<commit_after>package sshchat

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

	"github.com/shazow/rateio"
	"github.com/shazow/ssh-chat/chat"
	"github.com/shazow/ssh-chat/chat/message"
	"github.com/shazow/ssh-chat/sshd"
)

const maxInputLength int = 1024

// GetPrompt will render the terminal prompt string based on the user.
func GetPrompt(user *message.User) string {
	name := user.Name()
	if user.Config.Theme != nil {
		name = user.Config.Theme.ColorName(user)
	}
	return fmt.Sprintf("[%s] ", name)
}

// Host is the bridge between sshd and chat modules
// TODO: Should be easy to add support for multiple rooms, if we want.
type Host struct {
	*chat.Room
	listener *sshd.SSHListener
	commands chat.Commands
	auth     *Auth

	// Version string to print on /version
	Version string

	// Default theme
	theme message.Theme

	mu    sync.Mutex
	motd  string
	count int
}

// NewHost creates a Host on top of an existing listener.
func NewHost(listener *sshd.SSHListener, auth *Auth) *Host {
	room := chat.NewRoom()
	h := Host{
		Room:     room,
		listener: listener,
		commands: chat.Commands{},
		auth:     auth,
	}

	// Make our own commands registry instance.
	chat.InitCommands(&h.commands)
	h.InitCommands(&h.commands)
	room.SetCommands(h.commands)

	go room.Serve()
	return &h
}

// SetTheme sets the default theme for the host.
func (h *Host) SetTheme(theme message.Theme) {
	h.mu.Lock()
	h.theme = theme
	h.mu.Unlock()
}

// SetMotd sets the host's message of the day.
func (h *Host) SetMotd(motd string) {
	h.mu.Lock()
	h.motd = motd
	h.mu.Unlock()
}

func (h *Host) isOp(conn sshd.Connection) bool {
	key := conn.PublicKey()
	if key == nil {
		return false
	}
	return h.auth.IsOp(key)
}

// Connect a specific Terminal to this host and its room.
func (h *Host) Connect(term *sshd.Terminal) {
	id := NewIdentity(term.Conn)
	user := message.NewUserScreen(id, term)
	user.Config.Theme = &h.theme
	go user.Consume()

	// Close term once user is closed.
	defer user.Close()
	defer term.Close()

	h.mu.Lock()
	motd := h.motd
	count := h.count
	h.count++
	h.mu.Unlock()

	// Send MOTD
	if motd != "" {
		go user.Send(message.NewAnnounceMsg(motd))
	}

	member, err := h.Join(user)
	if err != nil {
		// Try again...
		id.SetName(fmt.Sprintf("Guest%d", count))
		member, err = h.Join(user)
	}
	if err != nil {
		logger.Errorf("Failed to join: %s", err)
		return
	}

	// Successfully joined.
	term.SetPrompt(GetPrompt(user))
	term.AutoCompleteCallback = h.AutoCompleteFunction(user)
	user.SetHighlight(user.Name())

	// Should the user be op'd on join?
	if h.isOp(term.Conn) {
		h.Room.Ops.Add(member)
	}
	ratelimit := rateio.NewSimpleLimiter(3, time.Second*3)

	logger.Debugf("Joined: %s", user.Name())

	for {
		line, err := term.ReadLine()
		if err == io.EOF {
			// Closed
			break
		} else if err != nil {
			logger.Errorf("Terminal reading error: %s", err)
			break
		}

		err = ratelimit.Count(1)
		if err != nil {
			user.Send(message.NewSystemMsg("Message rejected: Rate limiting is in effect.", user))
			continue
		}
		if len(line) > maxInputLength {
			user.Send(message.NewSystemMsg("Message rejected: Input too long.", user))
			continue
		}
		if line == "" {
			// Silently ignore empty lines.
			continue
		}

		m := message.ParseInput(line, user)

		// FIXME: Any reason to use h.room.Send(m) instead?
		h.HandleMsg(m)

		cmd := m.Command()
		if cmd == "/nick" || cmd == "/theme" {
			// Hijack /nick command to update terminal synchronously. Wouldn't
			// work if we use h.room.Send(m) above.
			//
			// FIXME: This is hacky, how do we improve the API to allow for
			// this? Chat module shouldn't know about terminals.
			term.SetPrompt(GetPrompt(user))
			user.SetHighlight(user.Name())
		}
	}

	err = h.Leave(user)
	if err != nil {
		logger.Errorf("Failed to leave: %s", err)
		return
	}
	logger.Debugf("Leaving: %s", user.Name())
}

// Serve our chat room onto the listener
func (h *Host) Serve() {
	h.listener.HandlerFunc = h.Connect
	h.listener.Serve()
}

func (h *Host) completeName(partial string) string {
	names := h.NamesPrefix(partial)
	if len(names) == 0 {
		// Didn't find anything
		return ""
	}

	return names[len(names)-1]
}

func (h *Host) completeCommand(partial string) string {
	for cmd := range h.commands {
		if strings.HasPrefix(cmd, partial) {
			return cmd
		}
	}
	return ""
}

// AutoCompleteFunction returns a callback for terminal autocompletion
func (h *Host) AutoCompleteFunction(u *message.User) func(line string, pos int, key rune) (newLine string, newPos int, ok bool) {
	return func(line string, pos int, key rune) (newLine string, newPos int, ok bool) {
		if key != 9 {
			return
		}

		if line == "" || strings.HasSuffix(line[:pos], " ") {
			// Don't autocomplete spaces.
			return
		}

		fields := strings.Fields(line[:pos])
		isFirst := len(fields) < 2
		partial := ""
		if len(fields) > 0 {
			partial = fields[len(fields)-1]
		}
		posPartial := pos - len(partial)

		var completed string
		if isFirst && strings.HasPrefix(partial, "/") {
			// Command
			completed = h.completeCommand(partial)
			if completed == "/reply" {
				replyTo := u.ReplyTo()
				if replyTo != nil {
					completed = "/msg " + replyTo.Name()
				}
			}
		} else {
			// Name
			completed = h.completeName(partial)
			if completed == "" {
				return
			}
			if isFirst {
				completed += ":"
			}
		}
		completed += " "

		// Reposition the cursor
		newLine = strings.Replace(line[posPartial:], partial, completed, 1)
		newLine = line[:posPartial] + newLine
		newPos = pos + (len(completed) - len(partial))
		ok = true
		return
	}
}

// GetUser returns a message.User based on a name.
func (h *Host) GetUser(name string) (*message.User, bool) {
	m, ok := h.MemberById(name)
	if !ok {
		return nil, false
	}
	return m.User, true
}

// InitCommands adds host-specific commands to a Commands container. These will
// override any existing commands.
func (h *Host) InitCommands(c *chat.Commands) {
	c.Add(chat.Command{
		Prefix:     "/msg",
		PrefixHelp: "USER MESSAGE",
		Help:       "Send MESSAGE to USER.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			args := msg.Args()
			switch len(args) {
			case 0:
				return errors.New("must specify user")
			case 1:
				return errors.New("must specify message")
			}

			target, ok := h.GetUser(args[0])
			if !ok {
				return errors.New("user not found")
			}

			m := message.NewPrivateMsg(strings.Join(args[1:], " "), msg.From(), target)
			room.Send(&m)
			return nil
		},
	})

	c.Add(chat.Command{
		Prefix:     "/reply",
		PrefixHelp: "MESSAGE",
		Help:       "Reply with MESSAGE to the previous private message.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			args := msg.Args()
			switch len(args) {
			case 0:
				return errors.New("must specify message")
			}

			target := msg.From().ReplyTo()
			if target == nil {
				return errors.New("no message to reply to")
			}

			m := message.NewPrivateMsg(strings.Join(args, " "), msg.From(), target)
			room.Send(&m)
			return nil
		},
	})

	c.Add(chat.Command{
		Prefix:     "/whois",
		PrefixHelp: "USER",
		Help:       "Information about USER.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			args := msg.Args()
			if len(args) == 0 {
				return errors.New("must specify user")
			}

			target, ok := h.GetUser(args[0])
			if !ok {
				return errors.New("user not found")
			}

			id := target.Identifier.(*Identity)
			room.Send(message.NewSystemMsg(id.Whois(), msg.From()))

			return nil
		},
	})

	// Hidden commands
	c.Add(chat.Command{
		Prefix: "/version",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			room.Send(message.NewSystemMsg(h.Version, msg.From()))
			return nil
		},
	})

	timeStarted := time.Now()
	c.Add(chat.Command{
		Prefix: "/uptime",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			room.Send(message.NewSystemMsg(time.Now().Sub(timeStarted).String(), msg.From()))
			return nil
		},
	})

	// Op commands
	c.Add(chat.Command{
		Op:         true,
		Prefix:     "/kick",
		PrefixHelp: "USER",
		Help:       "Kick USER from the server.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			if !room.IsOp(msg.From()) {
				return errors.New("must be op")
			}

			args := msg.Args()
			if len(args) == 0 {
				return errors.New("must specify user")
			}

			target, ok := h.GetUser(args[0])
			if !ok {
				return errors.New("user not found")
			}

			body := fmt.Sprintf("%s was kicked by %s.", target.Name(), msg.From().Name())
			room.Send(message.NewAnnounceMsg(body))
			target.Close()
			return nil
		},
	})

	c.Add(chat.Command{
		Op:         true,
		Prefix:     "/ban",
		PrefixHelp: "USER [DURATION]",
		Help:       "Ban USER from the server.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			// TODO: Would be nice to specify what to ban. Key? Ip? etc.
			if !room.IsOp(msg.From()) {
				return errors.New("must be op")
			}

			args := msg.Args()
			if len(args) == 0 {
				return errors.New("must specify user")
			}

			target, ok := h.GetUser(args[0])
			if !ok {
				return errors.New("user not found")
			}

			var until time.Duration = 0
			if len(args) > 1 {
				until, _ = time.ParseDuration(args[1])
			}

			id := target.Identifier.(*Identity)
			h.auth.Ban(id.PublicKey(), until)
			h.auth.BanAddr(id.RemoteAddr(), until)

			body := fmt.Sprintf("%s was banned by %s.", target.Name(), msg.From().Name())
			room.Send(message.NewAnnounceMsg(body))
			target.Close()

			logger.Debugf("Banned: \n-> %s", id.Whois())

			return nil
		},
	})

	c.Add(chat.Command{
		Op:         true,
		Prefix:     "/motd",
		PrefixHelp: "MESSAGE",
		Help:       "Set the MESSAGE of the day.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			if !room.IsOp(msg.From()) {
				return errors.New("must be op")
			}

			motd := ""
			args := msg.Args()
			if len(args) > 0 {
				motd = strings.Join(args, " ")
			}

			h.motd = motd
			body := fmt.Sprintf("New message of the day set by %s:", msg.From().Name())
			room.Send(message.NewAnnounceMsg(body))
			if motd != "" {
				room.Send(message.NewAnnounceMsg(motd))
			}

			return nil
		},
	})

	c.Add(chat.Command{
		Op:         true,
		Prefix:     "/op",
		PrefixHelp: "USER [DURATION]",
		Help:       "Set USER as admin.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			if !room.IsOp(msg.From()) {
				return errors.New("must be op")
			}

			args := msg.Args()
			if len(args) == 0 {
				return errors.New("must specify user")
			}

			var until time.Duration = 0
			if len(args) > 1 {
				until, _ = time.ParseDuration(args[1])
			}

			member, ok := room.MemberById(args[0])
			if !ok {
				return errors.New("user not found")
			}
			room.Ops.Add(member)
			id := member.Identifier.(*Identity)
			h.auth.Op(id.PublicKey(), until)

			body := fmt.Sprintf("Made op by %s.", msg.From().Name())
			room.Send(message.NewSystemMsg(body, member.User))

			return nil
		},
	})
}
<|endoftext|>


<commit_before>package model

import (
	"testing"

	"fmt"

	"strings"

	"github.com/evandroflores/claimr/messages"
	"github.com/stretchr/testify/assert"
)

func TestGetContainerNoTeam(t *testing.T) {
	container, err := GetContainer("", "", "")

	assert.ObjectsAreEqual(Container{}, container)
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "teamID"))
}

func TestGetContainerNoChannel(t *testing.T) {
	container, err := GetContainer("TestTeam", "", "")

	assert.ObjectsAreEqual(Container{}, container)
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "channelID"))
}

func TestGetContainerNoName(t *testing.T) {
	container, err := GetContainer("TestTeam", "TestChannel", "")

	assert.ObjectsAreEqual(Container{}, container)
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "container name"))
}

func TestGetContainerBigName(t *testing.T) {
	container, err := GetContainer("TestTeam", "TestChannel",
		"LoremIpsumDolorSitAmetConsecteturAdipiscingElit")

	assert.ObjectsAreEqual(Container{}, container)
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-name-too-big"), MaxNameSize))
}

func TestGetContainerNotFound(t *testing.T) {
	containerName := "TestDoesNotExist"
	container, err := GetContainer("TestTeam", "TestChannel", containerName)

	assert.ObjectsAreEqual(Container{}, container)
	assert.NoError(t, err, fmt.Sprintf("Container %s not found", containerName))
}

func TestGetContainerIgnoreCase(t *testing.T) {
	containerName := "UPPERCASE_NAME"
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: containerName}
	err := container.Add()
	assert.NoError(t, err)

	containerFromDB, err2 := GetContainer("TestTeam", "TestChannel", containerName)

	assert.NoError(t, err2)
	assert.Equal(t, containerFromDB.Name, strings.ToLower(containerName))

	container.Delete()
}

func TestAddContainer(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: "Name"}
	err := container.Add()
	assert.NoError(t, err)

	containerFromDB, err2 := GetContainer("TestTeam", "TestChannel", "Name")

	assert.NoError(t, err2)
	assert.ObjectsAreEqual(container, containerFromDB)
}

func TestAddContainerValidateTeamID(t *testing.T) {
	container := Container{TeamID: "", ChannelID: "", Name: ""}
	err := container.Add()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "teamID"))
}

func TestAddContainerValidateChannelID(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "", Name: ""}
	err := container.Add()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "channelID"))
}

func TestAddContainerValidateName(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: ""}
	err := container.Add()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "container name"))
}

func TestAddContainerDuplicate(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: "Name"}
	err := container.Add()
	assert.EqualError(t, err, messages.Get("same-name"))
}

func TestDeleteContainer(t *testing.T) {
	container, err := GetContainer("TestTeam", "TestChannel", "Name")
	assert.NoError(t, err)

	err2 := container.Delete()
	assert.NoError(t, err2)
}

func TestDeleteContainerValidateTeamID(t *testing.T) {
	container := Container{TeamID: "", ChannelID: "", Name: ""}
	err := container.Delete()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "teamID"))
}

func TestDeleteContainerValidateChannelID(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "", Name: ""}
	err := container.Delete()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "channelID"))
}

func TestDeleteContainerValidateName(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: ""}
	err := container.Delete()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "container name"))
}

func TestDeleteContainerNotFound(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: "Name"}
	err := container.Delete()
	assert.EqualError(t, err, messages.Get("container-not-found"))

}

func TestUpdateContainer(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: "Name"}
	err := container.Add()
	assert.NoError(t, err)

	container.InUseBy = "me"
	err2 := container.Update()
	assert.NoError(t, err2)

	containerUpdated := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: "Name"}

	assert.ObjectsAreEqual(containerUpdated.InUseBy, "me")
	assert.ObjectsAreEqual(container, containerUpdated)
	container.Delete()
}

func TestUpdateContainerValidateTeamID(t *testing.T) {
	container := Container{TeamID: "", ChannelID: "", Name: ""}
	err := container.Update()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "teamID"))
}

func TestUpdateContainerValidateChannelID(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "", Name: ""}
	err := container.Update()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "channelID"))
}

func TestUpdateContainerValidateName(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: ""}
	err := container.Update()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "container name"))
}

func TestUpdateContainerNotFound(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: "Name"}
	err := container.Update()
	assert.EqualError(t, err, messages.Get("container-not-found"))
}

func TestListContainers(t *testing.T) {
	names := [4]string{"A", "B", "C", "D"}
	for _, name := range names {
		container := Container{TeamID: "TestList", ChannelID: "TestChannel", Name: name}
		err := container.Add()
		assert.NoError(t, err)
	}

	containers, err2 := GetContainers("TestList", "TestChannel")
	assert.NoError(t, err2)

	assert.Len(t, containers, len(names))
	for idx, container := range containers {
		assert.ObjectsAreEqual(container.Name, names[idx])
		container.Delete()
	}
}

func TestListContainersValidateTeamID(t *testing.T) {
	containers, err := GetContainers("", "")
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "teamID"))
	assert.ObjectsAreEqual(containers, []Container{})
}

func TestListContainersValidateChannelID(t *testing.T) {
	containers, err := GetContainers("TestTeam", "")
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "channelID"))
	assert.ObjectsAreEqual(containers, []Container{})
}

func TestRemoveInUseData(t *testing.T) {
	team := "TestTeam"
	channel := "TestChannel"
	containerName := "name"
	user := "User"
	reason := "testing"

	container := Container{TeamID: team, ChannelID: channel, Name: containerName, InUseBy: user, InUseForReason: reason}
	err := container.Add()
	assert.NoError(t, err)

	err2 := container.ClearInUse()
	assert.NoError(t, err2)

	containerExpected := Container{TeamID: team, ChannelID: channel, Name: containerName, InUseBy: "", InUseForReason: ""}
	containerFromDB, err3 := GetContainer(team, channel, containerName)
	assert.NoError(t, err3)

	// Ignoring the difference for this fields
	containerExpected.ID = containerFromDB.ID
	containerExpected.CreatedAt = containerFromDB.CreatedAt
	containerExpected.UpdatedAt = containerFromDB.UpdatedAt
	containerExpected.DeletedAt = containerFromDB.DeletedAt

	assert.Empty(t, containerFromDB.InUseBy)
	assert.Empty(t, containerFromDB.InUseForReason)
	assert.Equal(t, containerExpected, containerFromDB)
	container.Delete()
}

func TestSetInUseData(t *testing.T) {
	team := "TestTeam"
	channel := "TestChannel"
	containerName := "name"
	user := "User"
	reason := "testing"

	container := Container{TeamID: team, ChannelID: channel, Name: containerName, InUseBy: "", InUseForReason: ""}
	err := container.Add()
	assert.NoError(t, err)

	err2 := container.SetInUse(user, reason)
	assert.NoError(t, err2)

	containerExpected := Container{TeamID: team, ChannelID: channel, Name: containerName, InUseBy: user, InUseForReason: reason}
	containerFromDB, err3 := GetContainer(team, channel, containerName)
	assert.NoError(t, err3)

	// Ignoring the difference for this fields
	containerExpected.ID = containerFromDB.ID
	containerExpected.CreatedAt = containerFromDB.CreatedAt
	containerExpected.UpdatedAt = containerFromDB.UpdatedAt
	containerExpected.DeletedAt = containerFromDB.DeletedAt

	assert.Equal(t, containerExpected, containerFromDB)
	container.Delete()
}
<commit_msg>in-use-text-invalid test<commit_after>package model

import (
	"testing"

	"fmt"

	"strings"

	"github.com/evandroflores/claimr/messages"
	"github.com/stretchr/testify/assert"
)

func TestGetContainerNoTeam(t *testing.T) {
	container, err := GetContainer("", "", "")

	assert.ObjectsAreEqual(Container{}, container)
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "teamID"))
}

func TestGetContainerNoChannel(t *testing.T) {
	container, err := GetContainer("TestTeam", "", "")

	assert.ObjectsAreEqual(Container{}, container)
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "channelID"))
}

func TestGetContainerNoName(t *testing.T) {
	container, err := GetContainer("TestTeam", "TestChannel", "")

	assert.ObjectsAreEqual(Container{}, container)
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "container name"))
}

func TestGetContainerBigName(t *testing.T) {
	container, err := GetContainer("TestTeam", "TestChannel",
		"LoremIpsumDolorSitAmetConsecteturAdipiscingElit")

	assert.ObjectsAreEqual(Container{}, container)
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-name-too-big"), MaxNameSize))
}

func TestGetContainerNotFound(t *testing.T) {
	containerName := "TestDoesNotExist"
	container, err := GetContainer("TestTeam", "TestChannel", containerName)

	assert.ObjectsAreEqual(Container{}, container)
	assert.NoError(t, err, fmt.Sprintf("Container %s not found", containerName))
}

func TestGetContainerIgnoreCase(t *testing.T) {
	containerName := "UPPERCASE_NAME"
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: containerName}
	err := container.Add()
	assert.NoError(t, err)

	containerFromDB, err2 := GetContainer("TestTeam", "TestChannel", containerName)

	assert.NoError(t, err2)
	assert.Equal(t, containerFromDB.Name, strings.ToLower(containerName))

	container.Delete()
}

func TestAddContainer(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: "Name"}
	err := container.Add()
	assert.NoError(t, err)

	containerFromDB, err2 := GetContainer("TestTeam", "TestChannel", "Name")

	assert.NoError(t, err2)
	assert.ObjectsAreEqual(container, containerFromDB)
}

func TestAddContainerValidateTeamID(t *testing.T) {
	container := Container{TeamID: "", ChannelID: "", Name: ""}
	err := container.Add()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "teamID"))
}

func TestAddContainerValidateChannelID(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "", Name: ""}
	err := container.Add()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "channelID"))
}

func TestAddContainerValidateName(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: ""}
	err := container.Add()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "container name"))
}

func TestAddContainerDuplicate(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: "Name"}
	err := container.Add()
	assert.EqualError(t, err, messages.Get("same-name"))
}

func TestDeleteContainer(t *testing.T) {
	container, err := GetContainer("TestTeam", "TestChannel", "Name")
	assert.NoError(t, err)

	err2 := container.Delete()
	assert.NoError(t, err2)
}

func TestDeleteContainerValidateTeamID(t *testing.T) {
	container := Container{TeamID: "", ChannelID: "", Name: ""}
	err := container.Delete()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "teamID"))
}

func TestDeleteContainerValidateChannelID(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "", Name: ""}
	err := container.Delete()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "channelID"))
}

func TestDeleteContainerValidateName(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: ""}
	err := container.Delete()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "container name"))
}

func TestDeleteContainerNotFound(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: "Name"}
	err := container.Delete()
	assert.EqualError(t, err, messages.Get("container-not-found"))

}

func TestUpdateContainer(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: "Name"}
	err := container.Add()
	assert.NoError(t, err)

	container.InUseBy = "me"
	err2 := container.Update()
	assert.NoError(t, err2)

	containerUpdated := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: "Name"}

	assert.ObjectsAreEqual(containerUpdated.InUseBy, "me")
	assert.ObjectsAreEqual(container, containerUpdated)
	container.Delete()
}

func TestUpdateContainerValidateTeamID(t *testing.T) {
	container := Container{TeamID: "", ChannelID: "", Name: ""}
	err := container.Update()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "teamID"))
}

func TestUpdateContainerValidateChannelID(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "", Name: ""}
	err := container.Update()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "channelID"))
}

func TestUpdateContainerValidateName(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: ""}
	err := container.Update()
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "container name"))
}

func TestUpdateContainerNotFound(t *testing.T) {
	container := Container{TeamID: "TestTeam", ChannelID: "TestChannel", Name: "Name"}
	err := container.Update()
	assert.EqualError(t, err, messages.Get("container-not-found"))
}

func TestListContainers(t *testing.T) {
	names := [4]string{"A", "B", "C", "D"}
	for _, name := range names {
		container := Container{TeamID: "TestList", ChannelID: "TestChannel", Name: name}
		err := container.Add()
		assert.NoError(t, err)
	}

	containers, err2 := GetContainers("TestList", "TestChannel")
	assert.NoError(t, err2)

	assert.Len(t, containers, len(names))
	for idx, container := range containers {
		assert.ObjectsAreEqual(container.Name, names[idx])
		container.Delete()
	}
}

func TestListContainersValidateTeamID(t *testing.T) {
	containers, err := GetContainers("", "")
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "teamID"))
	assert.ObjectsAreEqual(containers, []Container{})
}

func TestListContainersValidateChannelID(t *testing.T) {
	containers, err := GetContainers("TestTeam", "")
	assert.EqualError(t, err, fmt.Sprintf(messages.Get("field-required"), "channelID"))
	assert.ObjectsAreEqual(containers, []Container{})
}

func TestRemoveInUseData(t *testing.T) {
	team := "TestTeam"
	channel := "TestChannel"
	containerName := "name"
	user := "User"
	reason := "testing"

	container := Container{TeamID: team, ChannelID: channel, Name: containerName, InUseBy: user, InUseForReason: reason}
	err := container.Add()
	assert.NoError(t, err)

	err2 := container.ClearInUse()
	assert.NoError(t, err2)

	containerExpected := Container{TeamID: team, ChannelID: channel, Name: containerName, InUseBy: "", InUseForReason: ""}
	containerFromDB, err3 := GetContainer(team, channel, containerName)
	assert.NoError(t, err3)

	// Ignoring the difference for this fields
	containerExpected.ID = containerFromDB.ID
	containerExpected.CreatedAt = containerFromDB.CreatedAt
	containerExpected.UpdatedAt = containerFromDB.UpdatedAt
	containerExpected.DeletedAt = containerFromDB.DeletedAt

	assert.Empty(t, containerFromDB.InUseBy)
	assert.Empty(t, containerFromDB.InUseForReason)
	assert.Equal(t, containerExpected, containerFromDB)
	container.Delete()
}

func TestSetInUseData(t *testing.T) {
	team := "TestTeam"
	channel := "TestChannel"
	containerName := "name"
	user := "User"
	reason := "testing"

	container := Container{TeamID: team, ChannelID: channel, Name: containerName, InUseBy: "", InUseForReason: ""}
	err := container.Add()
	assert.NoError(t, err)

	err2 := container.SetInUse(user, reason)
	assert.NoError(t, err2)

	containerExpected := Container{TeamID: team, ChannelID: channel, Name: containerName, InUseBy: user, InUseForReason: reason}
	containerFromDB, err3 := GetContainer(team, channel, containerName)
	assert.NoError(t, err3)

	// Ignoring the difference for this fields
	containerExpected.ID = containerFromDB.ID
	containerExpected.CreatedAt = containerFromDB.CreatedAt
	containerExpected.UpdatedAt = containerFromDB.UpdatedAt
	containerExpected.DeletedAt = containerFromDB.DeletedAt

	assert.Equal(t, containerExpected, containerFromDB)
	container.Delete()
}

func TestInUseTextInvalid(t *testing.T) {
	container := Container{}
	assert.Equal(t, messages.Get("in-use-text-invalid"), container.InUseText("whatever"))
}
<|endoftext|>


<commit_before>package models_test

import (
	"github.com/herald-it/goncord/models"
	. "github.com/smartystreets/goconvey/convey"
	"testing"
)

func TestLoadSettings(t *testing.T) {
	Convey("Load settings from test", t, func() {
		err := models.LoadSettings()

		So(err, ShouldNotBeNil)
		So(models.Set.Database, ShouldNotBeNil)
		So(models.Set.Database.DbName, ShouldNotBeNil)
		So(models.Set.Database.Host, ShouldNotBeNil)
		So(models.Set.Database.TokenTable, ShouldNotBeNil)
		So(models.Set.Database.UserTable, ShouldNotBeNil)
	})
}
<commit_msg>Fix bad test.<commit_after>package models_test

import (
	"testing"

	"github.com/herald-it/goncord/models"
	. "github.com/smartystreets/goconvey/convey"
)

func TestLoadSettings(t *testing.T) {
	Convey("Load settings from test", t, func() {
		err := models.LoadSettings()

		So(err, ShouldBeNil)
		So(models.Set.Database, ShouldNotBeNil)
		So(models.Set.Database.DbName, ShouldNotBeNil)
		So(models.Set.Database.Host, ShouldNotBeNil)
		So(models.Set.Database.TokenTable, ShouldNotBeNil)
		So(models.Set.Database.UserTable, ShouldNotBeNil)
	})
}
<|endoftext|>


<commit_before>// Copyright 2015 ThoughtWorks, Inc.

// This file is part of getgauge/html-report.

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

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

// You should have received a copy of the GNU General Public License
// along with getgauge/html-report.  If not, see <http://www.gnu.org/licenses/>.
package generator

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/template"
	"time"

	"github.com/getgauge/common"
	gm "github.com/getgauge/html-report/gauge_messages"
	"github.com/russross/blackfriday"
)

type summary struct {
	Total   int
	Failed  int
	Passed  int
	Skipped int
}

type overview struct {
	ProjectName string
	Env         string
	Tags        string
	SuccRate    float32
	ExecTime    string
	Timestamp   string
	Summary     *summary
	BasePath    string
}

type specsMeta struct {
	SpecName   string
	ExecTime   string
	Failed     bool
	Skipped    bool
	Tags       []string
	ReportFile string
}

type sidebar struct {
	IsBeforeHookFailure bool
	Specs               []*specsMeta
}

type hookFailure struct {
	HookName   string
	ErrMsg     string
	Screenshot string
	StackTrace string
}

type specHeader struct {
	SpecName string
	ExecTime string
	FileName string
	Tags     []string
	Summary  *summary
}

type row struct {
	Cells []string
	Res   status
}

type table struct {
	Headers []string
	Rows    []*row
}

type spec struct {
	CommentsBeforeTable []string
	Table               *table
	CommentsAfterTable  []string
	Scenarios           []*scenario
	BeforeHookFailure   *hookFailure
	AfterHookFailure    *hookFailure
}

type scenario struct {
	Heading           string
	ExecTime          string
	Tags              []string
	ExecStatus        status
	Contexts          []item
	Items             []item
	Teardown          []item
	BeforeHookFailure *hookFailure
	AfterHookFailure  *hookFailure
	TableRowIndex     int
}

const (
	stepKind kind = iota
	commentKind
	conceptKind
)

type kind int

type item interface {
	kind() kind
}

type step struct {
	Fragments       []*fragment
	Res             *result
	PreHookFailure  *hookFailure
	PostHookFailure *hookFailure
}

func (s *step) kind() kind {
	return stepKind
}

type concept struct {
	CptStep *step
	Items   []item
}

func (c *concept) kind() kind {
	return conceptKind
}

type comment struct {
	Text string
}

func (c *comment) kind() kind {
	return commentKind
}

type result struct {
	Status        status
	StackTrace    string
	Screenshot    string
	ErrorMessage  string
	ExecTime      string
	SkippedReason string
	Messages      []string
}

type searchIndex struct {
	Tags  map[string][]string `json:"tags"`
	Specs map[string][]string `json:"specs"`
}

type status int

const (
	pass status = iota
	fail
	skip
	notExecuted
)

var parsedTemplates = make(map[string]*template.Template, 0)

// Any new templates that are added in file `templates.go` should be registered here
var templates = []string{bodyFooterTag, reportOverviewTag, sidebarDiv, congratsDiv, hookFailureDiv, tagsDiv, messageDiv, skippedReasonDiv,
	specsStartDiv, specsItemsContainerDiv, specsItemsContentsDiv, specHeaderStartTag, scenarioContainerStartDiv, scenarioHeaderStartDiv, specCommentsAndTableTag,
	htmlPageStartTag, headerEndTag, mainEndTag, endDiv, conceptStartDiv, stepStartDiv,
	stepMetaDiv, stepBodyDiv, stepFailureDiv, stepEndDiv, conceptSpan, contextOrTeardownStartDiv, commentSpan, conceptStepsStartDiv, nestedConceptDiv, htmlPageEndWithJS,
}

func init() {
	var encodeNewLine = func(s string) string {
		return strings.Replace(s, "\n", "
", -1)
	}
	var parseMarkdown = func(args ...interface{}) string {
		s := blackfriday.MarkdownCommon([]byte(fmt.Sprintf("%s", args...)))
		return string(s)
	}

	var funcs = template.FuncMap{"parseMarkdown": parseMarkdown, "escapeHTML": template.HTMLEscapeString, "encodeNewLine": encodeNewLine}
	for _, tmpl := range templates {
		t, err := template.New("Reports").Funcs(funcs).Parse(tmpl)
		if err != nil {
			log.Fatalf(err.Error())
		}
		parsedTemplates[tmpl] = t
	}
}

func execTemplate(tmplName string, w io.Writer, data interface{}) {
	tmpl := parsedTemplates[tmplName]
	if tmpl == nil {
		log.Fatal(tmplName)
	}
	err := tmpl.Execute(w, data)
	if err != nil {
		log.Fatalf(err.Error())
	}
}

// ProjectRoot is root dir of current project
var ProjectRoot string

// GenerateReports generates HTML report in the given report dir location
func GenerateReports(suiteRes *gm.ProtoSuiteResult, reportDir string) error {
	t := time.Now()
	f, err := os.Create(filepath.Join(reportDir, "index.html"))
	if err != nil {
		return err
	}
	if suiteRes.GetPreHookFailure() != nil {
		overview := toOverview(suiteRes, nil)
		generateOverview(overview, f)
		execTemplate(hookFailureDiv, f, toHookFailure(suiteRes.GetPreHookFailure(), "Before Suite"))
		if suiteRes.GetPostHookFailure() != nil {
			execTemplate(hookFailureDiv, f, toHookFailure(suiteRes.GetPostHookFailure(), "After Suite"))
		}
		generatePageFooter(overview, f)
	} else {
		go generateIndexPage(suiteRes, f)
		specRes := suiteRes.GetSpecResults()
		done := make(chan bool, len(specRes))
		for _, res := range specRes {
			relPath, _ := filepath.Rel(ProjectRoot, res.GetProtoSpec().GetFileName())
			CreateDirectory(filepath.Join(reportDir, filepath.Dir(relPath)))
			sf, err := os.Create(filepath.Join(reportDir, toHTMLFileName(res.GetProtoSpec().GetFileName(), ProjectRoot)))
			if err != nil {
				return err
			}
			go generateSpecPage(suiteRes, res, sf, done)
		}
		for _ = range specRes {
			<-done
		}
		close(done)
	}
	err = generateSearchIndex(suiteRes, reportDir)
	if err != nil {
		return err
	}
	fmt.Println("To generate reports: ", time.Since(t))
	return nil
}

func newSearchIndex() *searchIndex {
	var i searchIndex
	i.Tags = make(map[string][]string)
	i.Specs = make(map[string][]string)
	return &i
}

func (i *searchIndex) hasValueForTag(tag string, spec string) bool {
	for _, s := range i.Tags[tag] {
		if s == spec {
			return true
		}
	}
	return false
}

func (i *searchIndex) hasSpec(specHeading string, specFileName string) bool {
	for _, s := range i.Specs[specHeading] {
		if s == specFileName {
			return true
		}
	}
	return false
}

func generateSearchIndex(suiteRes *gm.ProtoSuiteResult, reportDir string) error {
	CreateDirectory(filepath.Join(reportDir, "js"))
	f, err := os.Create(filepath.Join(reportDir, "js", "search_index.js"))
	if err != nil {
		return err
	}
	index := newSearchIndex()
	for _, r := range suiteRes.GetSpecResults() {
		spec := r.GetProtoSpec()
		specFileName := toHTMLFileName(spec.GetFileName(), ProjectRoot)
		for _, t := range spec.GetTags() {
			if !index.hasValueForTag(t, specFileName) {
				index.Tags[t] = append(index.Tags[t], specFileName)
			}
		}
		var addTagsFromScenario = func(s *gm.ProtoScenario) {
			for _, t := range s.GetTags() {
				if !index.hasValueForTag(t, specFileName) {
					index.Tags[t] = append(index.Tags[t], specFileName)
				}
			}
		}
		for _, i := range spec.GetItems() {
			if s := i.GetScenario(); s != nil {
				addTagsFromScenario(s)
			}
			if tds := i.GetTableDrivenScenario(); tds != nil {
				addTagsFromScenario(tds.GetScenario())
			}
		}
		specHeading := spec.GetSpecHeading()
		if !index.hasSpec(specHeading, specFileName) {
			index.Specs[specHeading] = append(index.Specs[specHeading], specFileName)
		}
	}
	s, err := json.Marshal(index)
	if err != nil {
		return err
	}
	f.WriteString(fmt.Sprintf("var index = %s;", s))
	return nil
}

func generateIndexPage(suiteRes *gm.ProtoSuiteResult, w io.Writer) {
	overview := toOverview(suiteRes, nil)
	generateOverview(overview, w)
	if suiteRes.GetPostHookFailure() != nil {
		execTemplate(hookFailureDiv, w, toHookFailure(suiteRes.GetPostHookFailure(), "After Suite"))
	}
	execTemplate(specsStartDiv, w, nil)
	execTemplate(sidebarDiv, w, toSidebar(suiteRes, nil))
	if !suiteRes.GetFailed() {
		execTemplate(congratsDiv, w, nil)
	}
	execTemplate(endDiv, w, nil)
	generatePageFooter(overview, w)
}

func generateSpecPage(suiteRes *gm.ProtoSuiteResult, specRes *gm.ProtoSpecResult, w io.Writer, done chan bool) {
	overview := toOverview(suiteRes, specRes)

	generateOverview(overview, w)

	if suiteRes.GetPreHookFailure() != nil {
		execTemplate(hookFailureDiv, w, toHookFailure(suiteRes.GetPreHookFailure(), "Before Suite"))
	}

	if suiteRes.GetPostHookFailure() != nil {
		execTemplate(hookFailureDiv, w, toHookFailure(suiteRes.GetPostHookFailure(), "After Suite"))
	}

	if suiteRes.GetPreHookFailure() == nil {
		execTemplate(specsStartDiv, w, nil)
		execTemplate(sidebarDiv, w, toSidebar(suiteRes, specRes))
		generateSpecDiv(w, specRes)
		execTemplate(endDiv, w, nil)
	}

	generatePageFooter(overview, w)
	done <- true
}

func generateOverview(overview *overview, w io.Writer) {
	execTemplate(htmlPageStartTag, w, overview)
	execTemplate(reportOverviewTag, w, overview)
}

func generatePageFooter(overview *overview, w io.Writer) {
	execTemplate(endDiv, w, nil)
	execTemplate(mainEndTag, w, nil)
	execTemplate(bodyFooterTag, w, nil)
	execTemplate(htmlPageEndWithJS, w, overview)
}

func generateSpecDiv(w io.Writer, res *gm.ProtoSpecResult) {
	specHeader := toSpecHeader(res)
	spec := toSpec(res)

	execTemplate(specHeaderStartTag, w, specHeader)
	execTemplate(tagsDiv, w, specHeader)
	execTemplate(headerEndTag, w, nil)
	execTemplate(specsItemsContainerDiv, w, nil)

	if spec.BeforeHookFailure != nil {
		execTemplate(hookFailureDiv, w, spec.BeforeHookFailure)
	}

	execTemplate(specsItemsContentsDiv, w, nil)
	execTemplate(specCommentsAndTableTag, w, spec)

	if spec.BeforeHookFailure == nil {
		for _, scn := range spec.Scenarios {
			generateScenario(w, scn)
		}
	}

	execTemplate(endDiv, w, nil)
	execTemplate(endDiv, w, nil)

	if spec.AfterHookFailure != nil {
		execTemplate(hookFailureDiv, w, spec.AfterHookFailure)
	}

	execTemplate(endDiv, w, nil)
}

func generateScenario(w io.Writer, scn *scenario) {
	execTemplate(scenarioContainerStartDiv, w, scn)
	execTemplate(scenarioHeaderStartDiv, w, scn)
	execTemplate(tagsDiv, w, scn)
	execTemplate(endDiv, w, nil)
	if scn.BeforeHookFailure != nil {
		execTemplate(hookFailureDiv, w, scn.BeforeHookFailure)
	}

	generateItems(w, scn.Contexts, generateContextOrTeardown)
	generateItems(w, scn.Items, generateItem)
	generateItems(w, scn.Teardown, generateContextOrTeardown)

	if scn.AfterHookFailure != nil {
		execTemplate(hookFailureDiv, w, scn.AfterHookFailure)
	}
	execTemplate(endDiv, w, nil)
}

func generateItems(w io.Writer, items []item, predicate func(w io.Writer, item item)) {
	for _, item := range items {
		predicate(w, item)
	}
}

func generateContextOrTeardown(w io.Writer, item item) {
	execTemplate(contextOrTeardownStartDiv, w, nil)
	generateItem(w, item)
	execTemplate(endDiv, w, nil)
}

func generateItem(w io.Writer, item item) {
	switch item.kind() {
	case stepKind:
		execTemplate(stepStartDiv, w, item.(*step))
		execTemplate(stepBodyDiv, w, item.(*step))

		if item.(*step).PreHookFailure != nil {
			execTemplate(hookFailureDiv, w, item.(*step).PreHookFailure)
		}

		stepRes := item.(*step).Res
		if stepRes.Status == fail && stepRes.ErrorMessage != "" && stepRes.StackTrace != "" {
			execTemplate(stepFailureDiv, w, stepRes)
		}

		if item.(*step).PostHookFailure != nil {
			execTemplate(hookFailureDiv, w, item.(*step).PostHookFailure)
		}
		execTemplate(messageDiv, w, stepRes)
		execTemplate(stepEndDiv, w, item.(*step))
		if stepRes.Status == skip && stepRes.SkippedReason != "" {
			execTemplate(skippedReasonDiv, w, stepRes)
		}
	case commentKind:
		execTemplate(commentSpan, w, item.(*comment))
	case conceptKind:
		execTemplate(conceptStartDiv, w, item.(*concept).CptStep)
		execTemplate(conceptSpan, w, nil)
		execTemplate(stepBodyDiv, w, item.(*concept).CptStep)
		execTemplate(stepEndDiv, w, item.(*concept).CptStep)
		execTemplate(conceptStepsStartDiv, w, nil)
		generateItems(w, item.(*concept).Items, generateItem)
		execTemplate(endDiv, w, nil)
	}
}

// CreateDirectory creates given directory if it doesn't exist
func CreateDirectory(dir string) {
	if common.DirExists(dir) {
		return
	}
	if err := os.MkdirAll(dir, common.NewDirectoryPermissions); err != nil {
		fmt.Printf("Failed to create directory %s: %s\n", dir, err)
		os.Exit(1)
	}
}
<commit_msg>Formatting list of templates | Removed print statement<commit_after>// Copyright 2015 ThoughtWorks, Inc.

// This file is part of getgauge/html-report.

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

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

// You should have received a copy of the GNU General Public License
// along with getgauge/html-report.  If not, see <http://www.gnu.org/licenses/>.
package generator

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/getgauge/common"
	gm "github.com/getgauge/html-report/gauge_messages"
	"github.com/russross/blackfriday"
)

type summary struct {
	Total   int
	Failed  int
	Passed  int
	Skipped int
}

type overview struct {
	ProjectName string
	Env         string
	Tags        string
	SuccRate    float32
	ExecTime    string
	Timestamp   string
	Summary     *summary
	BasePath    string
}

type specsMeta struct {
	SpecName   string
	ExecTime   string
	Failed     bool
	Skipped    bool
	Tags       []string
	ReportFile string
}

type sidebar struct {
	IsBeforeHookFailure bool
	Specs               []*specsMeta
}

type hookFailure struct {
	HookName   string
	ErrMsg     string
	Screenshot string
	StackTrace string
}

type specHeader struct {
	SpecName string
	ExecTime string
	FileName string
	Tags     []string
	Summary  *summary
}

type row struct {
	Cells []string
	Res   status
}

type table struct {
	Headers []string
	Rows    []*row
}

type spec struct {
	CommentsBeforeTable []string
	Table               *table
	CommentsAfterTable  []string
	Scenarios           []*scenario
	BeforeHookFailure   *hookFailure
	AfterHookFailure    *hookFailure
}

type scenario struct {
	Heading           string
	ExecTime          string
	Tags              []string
	ExecStatus        status
	Contexts          []item
	Items             []item
	Teardown          []item
	BeforeHookFailure *hookFailure
	AfterHookFailure  *hookFailure
	TableRowIndex     int
}

const (
	stepKind kind = iota
	commentKind
	conceptKind
)

type kind int

type item interface {
	kind() kind
}

type step struct {
	Fragments       []*fragment
	Res             *result
	PreHookFailure  *hookFailure
	PostHookFailure *hookFailure
}

func (s *step) kind() kind {
	return stepKind
}

type concept struct {
	CptStep *step
	Items   []item
}

func (c *concept) kind() kind {
	return conceptKind
}

type comment struct {
	Text string
}

func (c *comment) kind() kind {
	return commentKind
}

type result struct {
	Status        status
	StackTrace    string
	Screenshot    string
	ErrorMessage  string
	ExecTime      string
	SkippedReason string
	Messages      []string
}

type searchIndex struct {
	Tags  map[string][]string `json:"tags"`
	Specs map[string][]string `json:"specs"`
}

type status int

const (
	pass status = iota
	fail
	skip
	notExecuted
)

var parsedTemplates = make(map[string]*template.Template, 0)

// Any new templates that are added in file `templates.go` should be registered here
var templates = []string{bodyFooterTag, reportOverviewTag, sidebarDiv, congratsDiv, hookFailureDiv, tagsDiv, messageDiv, skippedReasonDiv,
	specsStartDiv, specsItemsContainerDiv, specsItemsContentsDiv, specHeaderStartTag, scenarioContainerStartDiv, scenarioHeaderStartDiv, specCommentsAndTableTag,
	htmlPageStartTag, headerEndTag, mainEndTag, endDiv, conceptStartDiv, stepStartDiv, stepMetaDiv, stepBodyDiv, stepFailureDiv, stepEndDiv, conceptSpan,
	contextOrTeardownStartDiv, commentSpan, conceptStepsStartDiv, nestedConceptDiv, htmlPageEndWithJS,
}

func init() {
	var encodeNewLine = func(s string) string {
		return strings.Replace(s, "\n", "
", -1)
	}
	var parseMarkdown = func(args ...interface{}) string {
		s := blackfriday.MarkdownCommon([]byte(fmt.Sprintf("%s", args...)))
		return string(s)
	}
	var funcs = template.FuncMap{"parseMarkdown": parseMarkdown, "escapeHTML": template.HTMLEscapeString, "encodeNewLine": encodeNewLine}
	for _, tmpl := range templates {
		t, err := template.New("Reports").Funcs(funcs).Parse(tmpl)
		if err != nil {
			log.Fatalf(err.Error())
		}
		parsedTemplates[tmpl] = t
	}
}

func execTemplate(tmplName string, w io.Writer, data interface{}) {
	tmpl := parsedTemplates[tmplName]
	if tmpl == nil {
		log.Fatal(tmplName)
	}
	err := tmpl.Execute(w, data)
	if err != nil {
		log.Fatalf(err.Error())
	}
}

// ProjectRoot is root dir of current project
var ProjectRoot string

// GenerateReports generates HTML report in the given report dir location
func GenerateReports(suiteRes *gm.ProtoSuiteResult, reportDir string) error {
	f, err := os.Create(filepath.Join(reportDir, "index.html"))
	if err != nil {
		return err
	}
	if suiteRes.GetPreHookFailure() != nil {
		overview := toOverview(suiteRes, nil)
		generateOverview(overview, f)
		execTemplate(hookFailureDiv, f, toHookFailure(suiteRes.GetPreHookFailure(), "Before Suite"))
		if suiteRes.GetPostHookFailure() != nil {
			execTemplate(hookFailureDiv, f, toHookFailure(suiteRes.GetPostHookFailure(), "After Suite"))
		}
		generatePageFooter(overview, f)
	} else {
		go generateIndexPage(suiteRes, f)
		specRes := suiteRes.GetSpecResults()
		done := make(chan bool, len(specRes))
		for _, res := range specRes {
			relPath, _ := filepath.Rel(ProjectRoot, res.GetProtoSpec().GetFileName())
			CreateDirectory(filepath.Join(reportDir, filepath.Dir(relPath)))
			sf, err := os.Create(filepath.Join(reportDir, toHTMLFileName(res.GetProtoSpec().GetFileName(), ProjectRoot)))
			if err != nil {
				return err
			}
			go generateSpecPage(suiteRes, res, sf, done)
		}
		for _ = range specRes {
			<-done
		}
		close(done)
	}
	err = generateSearchIndex(suiteRes, reportDir)
	if err != nil {
		return err
	}
	return nil
}

func newSearchIndex() *searchIndex {
	var i searchIndex
	i.Tags = make(map[string][]string)
	i.Specs = make(map[string][]string)
	return &i
}

func (i *searchIndex) hasValueForTag(tag string, spec string) bool {
	for _, s := range i.Tags[tag] {
		if s == spec {
			return true
		}
	}
	return false
}

func (i *searchIndex) hasSpec(specHeading string, specFileName string) bool {
	for _, s := range i.Specs[specHeading] {
		if s == specFileName {
			return true
		}
	}
	return false
}

func generateSearchIndex(suiteRes *gm.ProtoSuiteResult, reportDir string) error {
	CreateDirectory(filepath.Join(reportDir, "js"))
	f, err := os.Create(filepath.Join(reportDir, "js", "search_index.js"))
	if err != nil {
		return err
	}
	index := newSearchIndex()
	for _, r := range suiteRes.GetSpecResults() {
		spec := r.GetProtoSpec()
		specFileName := toHTMLFileName(spec.GetFileName(), ProjectRoot)
		for _, t := range spec.GetTags() {
			if !index.hasValueForTag(t, specFileName) {
				index.Tags[t] = append(index.Tags[t], specFileName)
			}
		}
		var addTagsFromScenario = func(s *gm.ProtoScenario) {
			for _, t := range s.GetTags() {
				if !index.hasValueForTag(t, specFileName) {
					index.Tags[t] = append(index.Tags[t], specFileName)
				}
			}
		}
		for _, i := range spec.GetItems() {
			if s := i.GetScenario(); s != nil {
				addTagsFromScenario(s)
			}
			if tds := i.GetTableDrivenScenario(); tds != nil {
				addTagsFromScenario(tds.GetScenario())
			}
		}
		specHeading := spec.GetSpecHeading()
		if !index.hasSpec(specHeading, specFileName) {
			index.Specs[specHeading] = append(index.Specs[specHeading], specFileName)
		}
	}
	s, err := json.Marshal(index)
	if err != nil {
		return err
	}
	f.WriteString(fmt.Sprintf("var index = %s;", s))
	return nil
}

func generateIndexPage(suiteRes *gm.ProtoSuiteResult, w io.Writer) {
	overview := toOverview(suiteRes, nil)
	generateOverview(overview, w)
	if suiteRes.GetPostHookFailure() != nil {
		execTemplate(hookFailureDiv, w, toHookFailure(suiteRes.GetPostHookFailure(), "After Suite"))
	}
	execTemplate(specsStartDiv, w, nil)
	execTemplate(sidebarDiv, w, toSidebar(suiteRes, nil))
	if !suiteRes.GetFailed() {
		execTemplate(congratsDiv, w, nil)
	}
	execTemplate(endDiv, w, nil)
	generatePageFooter(overview, w)
}

func generateSpecPage(suiteRes *gm.ProtoSuiteResult, specRes *gm.ProtoSpecResult, w io.Writer, done chan bool) {
	overview := toOverview(suiteRes, specRes)

	generateOverview(overview, w)

	if suiteRes.GetPreHookFailure() != nil {
		execTemplate(hookFailureDiv, w, toHookFailure(suiteRes.GetPreHookFailure(), "Before Suite"))
	}

	if suiteRes.GetPostHookFailure() != nil {
		execTemplate(hookFailureDiv, w, toHookFailure(suiteRes.GetPostHookFailure(), "After Suite"))
	}

	if suiteRes.GetPreHookFailure() == nil {
		execTemplate(specsStartDiv, w, nil)
		execTemplate(sidebarDiv, w, toSidebar(suiteRes, specRes))
		generateSpecDiv(w, specRes)
		execTemplate(endDiv, w, nil)
	}

	generatePageFooter(overview, w)
	done <- true
}

func generateOverview(overview *overview, w io.Writer) {
	execTemplate(htmlPageStartTag, w, overview)
	execTemplate(reportOverviewTag, w, overview)
}

func generatePageFooter(overview *overview, w io.Writer) {
	execTemplate(endDiv, w, nil)
	execTemplate(mainEndTag, w, nil)
	execTemplate(bodyFooterTag, w, nil)
	execTemplate(htmlPageEndWithJS, w, overview)
}

func generateSpecDiv(w io.Writer, res *gm.ProtoSpecResult) {
	specHeader := toSpecHeader(res)
	spec := toSpec(res)

	execTemplate(specHeaderStartTag, w, specHeader)
	execTemplate(tagsDiv, w, specHeader)
	execTemplate(headerEndTag, w, nil)
	execTemplate(specsItemsContainerDiv, w, nil)

	if spec.BeforeHookFailure != nil {
		execTemplate(hookFailureDiv, w, spec.BeforeHookFailure)
	}

	execTemplate(specsItemsContentsDiv, w, nil)
	execTemplate(specCommentsAndTableTag, w, spec)

	if spec.BeforeHookFailure == nil {
		for _, scn := range spec.Scenarios {
			generateScenario(w, scn)
		}
	}

	execTemplate(endDiv, w, nil)
	execTemplate(endDiv, w, nil)

	if spec.AfterHookFailure != nil {
		execTemplate(hookFailureDiv, w, spec.AfterHookFailure)
	}

	execTemplate(endDiv, w, nil)
}

func generateScenario(w io.Writer, scn *scenario) {
	execTemplate(scenarioContainerStartDiv, w, scn)
	execTemplate(scenarioHeaderStartDiv, w, scn)
	execTemplate(tagsDiv, w, scn)
	execTemplate(endDiv, w, nil)
	if scn.BeforeHookFailure != nil {
		execTemplate(hookFailureDiv, w, scn.BeforeHookFailure)
	}

	generateItems(w, scn.Contexts, generateContextOrTeardown)
	generateItems(w, scn.Items, generateItem)
	generateItems(w, scn.Teardown, generateContextOrTeardown)

	if scn.AfterHookFailure != nil {
		execTemplate(hookFailureDiv, w, scn.AfterHookFailure)
	}
	execTemplate(endDiv, w, nil)
}

func generateItems(w io.Writer, items []item, predicate func(w io.Writer, item item)) {
	for _, item := range items {
		predicate(w, item)
	}
}

func generateContextOrTeardown(w io.Writer, item item) {
	execTemplate(contextOrTeardownStartDiv, w, nil)
	generateItem(w, item)
	execTemplate(endDiv, w, nil)
}

func generateItem(w io.Writer, item item) {
	switch item.kind() {
	case stepKind:
		execTemplate(stepStartDiv, w, item.(*step))
		execTemplate(stepBodyDiv, w, item.(*step))

		if item.(*step).PreHookFailure != nil {
			execTemplate(hookFailureDiv, w, item.(*step).PreHookFailure)
		}

		stepRes := item.(*step).Res
		if stepRes.Status == fail && stepRes.ErrorMessage != "" && stepRes.StackTrace != "" {
			execTemplate(stepFailureDiv, w, stepRes)
		}

		if item.(*step).PostHookFailure != nil {
			execTemplate(hookFailureDiv, w, item.(*step).PostHookFailure)
		}
		execTemplate(messageDiv, w, stepRes)
		execTemplate(stepEndDiv, w, item.(*step))
		if stepRes.Status == skip && stepRes.SkippedReason != "" {
			execTemplate(skippedReasonDiv, w, stepRes)
		}
	case commentKind:
		execTemplate(commentSpan, w, item.(*comment))
	case conceptKind:
		execTemplate(conceptStartDiv, w, item.(*concept).CptStep)
		execTemplate(conceptSpan, w, nil)
		execTemplate(stepBodyDiv, w, item.(*concept).CptStep)
		execTemplate(stepEndDiv, w, item.(*concept).CptStep)
		execTemplate(conceptStepsStartDiv, w, nil)
		generateItems(w, item.(*concept).Items, generateItem)
		execTemplate(endDiv, w, nil)
	}
}

// CreateDirectory creates given directory if it doesn't exist
func CreateDirectory(dir string) {
	if common.DirExists(dir) {
		return
	}
	if err := os.MkdirAll(dir, common.NewDirectoryPermissions); err != nil {
		fmt.Printf("Failed to create directory %s: %s\n", dir, err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"context"
	"crypto/tls"
	"io"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cybozu-go/log"
)

const (
	defaultHTTPReadTimeout = 30 * time.Second

	// our request tracking header.
	uuidHeaderName = "X-Cybozu-Request-ID"
)

// HTTPServer is a wrapper for http.Server.
//
// This struct overrides Serve and ListenAndServe* methods.
//
// http.Server members are replaced as following:
//    - Handler is replaced with a wrapper handler.
//    - ReadTimeout is set to 30 seconds if it is zero.
//    - ConnState is replaced with the one provided by the framework.
type HTTPServer struct {
	*http.Server

	// AccessLog is a logger for access logs.
	// If this is nil, the default logger is used.
	AccessLog *log.Logger

	// ShutdownTimeout is the maximum duration the server waits for
	// all connections to be closed before shutdown.
	//
	// Zero duration disables timeout.
	ShutdownTimeout time.Duration

	// Env is the environment where this server runs.
	//
	// The global environment is used if Env is nil.
	Env *Environment

	handler  http.Handler
	wg       sync.WaitGroup
	timedout int32

	mu        sync.Mutex
	idleConns map[net.Conn]struct{}

	initOnce sync.Once
}

// StdResponseWriter is the interface implemented by
// the ResponseWriter from http.Server.
//
// HTTPServer's ResponseWriter implements this as well.
type StdResponseWriter interface {
	http.ResponseWriter
	io.ReaderFrom
	http.Flusher
	http.CloseNotifier
	http.Hijacker
	WriteString(data string) (int, error)
}

type logResponseWriter struct {
	StdResponseWriter
	status int
	size   int64
}

func (w *logResponseWriter) WriteHeader(status int) {
	w.status = status
	w.StdResponseWriter.WriteHeader(status)
}

func (w *logResponseWriter) Write(data []byte) (int, error) {
	n, err := w.StdResponseWriter.Write(data)
	w.size += int64(n)
	return n, err
}

// ServeHTTP implements http.Handler interface.
func (s *HTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	startTime := time.Now()

	lw := &logResponseWriter{w.(StdResponseWriter), http.StatusOK, 0}
	ctx, cancel := context.WithCancel(s.Env.ctx)
	defer cancel()
	s.handler.ServeHTTP(lw, r.WithContext(ctx))

	fields := map[string]interface{}{
		log.FnType:           "access",
		log.FnResponseTime:   time.Since(startTime).Seconds(),
		log.FnProtocol:       r.Proto,
		log.FnHTTPStatusCode: lw.status,
		log.FnHTTPMethod:     r.Method,
		log.FnURL:            r.RequestURI,
		log.FnHTTPHost:       r.Host,
		log.FnRequestSize:    r.ContentLength,
		log.FnResponseSize:   lw.size,
	}
	ip, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		fields[log.FnRemoteAddress] = ip
	}
	ua := r.Header.Get("User-Agent")
	if len(ua) > 0 {
		fields[log.FnHTTPUserAgent] = ua
	}
	reqid := r.Header.Get(uuidHeaderName)
	if len(reqid) > 0 {
		fields[log.FnRequestID] = reqid
	}

	lv := log.LvInfo
	switch {
	case 500 <= lw.status:
		lv = log.LvError
	case 400 <= lw.status:
		lv = log.LvWarn
	}
	s.AccessLog.Log(lv, "cmd: "+http.StatusText(lw.status), fields)
}

func (s *HTTPServer) init() {
	if s.handler != nil {
		return
	}

	s.idleConns = make(map[net.Conn]struct{}, 100000)

	if s.Server.Handler == nil {
		panic("Handler must not be nil")
	}
	s.handler = s.Server.Handler
	s.Server.Handler = s
	if s.Server.ReadTimeout == 0 {
		s.Server.ReadTimeout = defaultHTTPReadTimeout
	}
	s.Server.ConnState = func(c net.Conn, state http.ConnState) {
		s.mu.Lock()
		if state == http.StateIdle {
			s.idleConns[c] = struct{}{}
		} else {
			delete(s.idleConns, c)
		}
		s.mu.Unlock()

		if state == http.StateNew {
			s.wg.Add(1)
			return
		}
		if state == http.StateHijacked || state == http.StateClosed {
			s.wg.Done()
		}
	}

	if s.AccessLog == nil {
		s.AccessLog = log.DefaultLogger()
	}

	if s.Env == nil {
		s.Env = defaultEnv
	}
	s.Env.Go(s.wait)
}

func (s *HTTPServer) wait(ctx context.Context) error {
	<-ctx.Done()

	s.Server.SetKeepAlivesEnabled(false)

	ch := make(chan struct{})

	// Interrupt conn.Read for idle connections.
	//
	// This must be run inside for-loop to catch connections
	// going idle at critical timing to acquire s.mu
	go func() {
	AGAIN:
		s.mu.Lock()
		for conn := range s.idleConns {
			conn.SetReadDeadline(time.Now())
		}
		s.mu.Unlock()
		select {
		case <-ch:
			return
		default:
		}
		time.Sleep(10 * time.Millisecond)
		goto AGAIN
	}()

	go func() {
		s.wg.Wait()
		close(ch)
	}()

	if s.ShutdownTimeout == 0 {
		<-ch
		return nil
	}

	select {
	case <-ch:
	case <-time.After(s.ShutdownTimeout):
		log.Warn("cmd: timeout waiting for shutdown", nil)
		atomic.StoreInt32(&s.timedout, 1)
	}
	return nil
}

// TimedOut returns true if the server shut down before all connections
// got closed.
func (s *HTTPServer) TimedOut() bool {
	return atomic.LoadInt32(&s.timedout) != 0
}

// Serve overrides http.Server's Serve method.
//
// Unlike the original, this method returns immediately just after
// starting a goroutine to accept connections.
//
// The framework automatically closes l when the environment's Cancel
// is called.
//
// Serve always returns nil.
func (s *HTTPServer) Serve(l net.Listener) error {
	s.initOnce.Do(s.init)

	go func() {
		<-s.Env.ctx.Done()
		l.Close()
	}()

	go func() {
		s.Server.Serve(l)
	}()

	return nil
}

type tcpKeepAliveListener struct {
	*net.TCPListener
}

func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) {
	tc, err := ln.AcceptTCP()
	if err != nil {
		return
	}
	tc.SetKeepAlive(true)
	tc.SetKeepAlivePeriod(3 * time.Minute)
	return tc, nil
}

// ListenAndServe overrides http.Server's method.
//
// Unlike the original, this method returns immediately just after
// starting a goroutine to accept connections.  To stop listening,
// call the environment's Cancel.
//
// ListenAndServe returns non-nil error if and only if net.Listen failed.
func (s *HTTPServer) ListenAndServe() error {
	addr := s.Server.Addr
	if addr == "" {
		addr = ":http"
	}
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	return s.Serve(tcpKeepAliveListener{ln.(*net.TCPListener)})
}

// ListenAndServeTLS overrides http.Server's method.
//
// Unlike the original, this method returns immediately just after
// starting a goroutine to accept connections.  To stop listening,
// call the environment's Cancel.
//
// Another difference from the original is that certFile and keyFile
// must be specified.  If not, configure http.Server.TLSConfig
// manually and use Serve().
//
// HTTP/2 is always enabled.
//
// ListenAndServeTLS returns non-nil error if net.Listen failed
// or failed to load certificate files.
func (s *HTTPServer) ListenAndServeTLS(certFile, keyFile string) error {
	addr := s.Server.Addr
	if addr == "" {
		addr = ":https"
	}

	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}

	config := &tls.Config{
		NextProtos:               []string{"h2", "http/1.1"},
		Certificates:             []tls.Certificate{cert},
		PreferServerCipherSuites: true,
		ClientSessionCache:       tls.NewLRUClientSessionCache(0),
	}
	s.Server.TLSConfig = config

	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	tlsListener := tls.NewListener(tcpKeepAliveListener{ln.(*net.TCPListener)}, config)
	return s.Serve(tlsListener)
}
<commit_msg>[logResponseWriter] intercept io.ReaderFrom and WriteString.<commit_after>package cmd

import (
	"context"
	"crypto/tls"
	"io"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cybozu-go/log"
)

const (
	defaultHTTPReadTimeout = 30 * time.Second

	// our request tracking header.
	uuidHeaderName = "X-Cybozu-Request-ID"
)

// HTTPServer is a wrapper for http.Server.
//
// This struct overrides Serve and ListenAndServe* methods.
//
// http.Server members are replaced as following:
//    - Handler is replaced with a wrapper handler.
//    - ReadTimeout is set to 30 seconds if it is zero.
//    - ConnState is replaced with the one provided by the framework.
type HTTPServer struct {
	*http.Server

	// AccessLog is a logger for access logs.
	// If this is nil, the default logger is used.
	AccessLog *log.Logger

	// ShutdownTimeout is the maximum duration the server waits for
	// all connections to be closed before shutdown.
	//
	// Zero duration disables timeout.
	ShutdownTimeout time.Duration

	// Env is the environment where this server runs.
	//
	// The global environment is used if Env is nil.
	Env *Environment

	handler  http.Handler
	wg       sync.WaitGroup
	timedout int32

	mu        sync.Mutex
	idleConns map[net.Conn]struct{}

	initOnce sync.Once
}

// StdResponseWriter is the interface implemented by
// the ResponseWriter from http.Server.
//
// HTTPServer's ResponseWriter implements this as well.
type StdResponseWriter interface {
	http.ResponseWriter
	io.ReaderFrom
	http.Flusher
	http.CloseNotifier
	http.Hijacker
	WriteString(data string) (int, error)
}

type logResponseWriter struct {
	StdResponseWriter
	status int
	size   int64
}

func (w *logResponseWriter) WriteHeader(status int) {
	w.status = status
	w.StdResponseWriter.WriteHeader(status)
}

func (w *logResponseWriter) Write(data []byte) (int, error) {
	n, err := w.StdResponseWriter.Write(data)
	w.size += int64(n)
	return n, err
}

func (w *logResponseWriter) ReadFrom(r io.Reader) (int64, error) {
	n, err := w.StdResponseWriter.ReadFrom(r)
	w.size += n
	return n, err
}

func (w *logResponseWriter) WriteString(data string) (int, error) {
	n, err := w.StdResponseWriter.WriteString(data)
	w.size += int64(n)
	return n, err
}

// ServeHTTP implements http.Handler interface.
func (s *HTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	startTime := time.Now()

	lw := &logResponseWriter{w.(StdResponseWriter), http.StatusOK, 0}
	ctx, cancel := context.WithCancel(s.Env.ctx)
	defer cancel()
	s.handler.ServeHTTP(lw, r.WithContext(ctx))

	fields := map[string]interface{}{
		log.FnType:           "access",
		log.FnResponseTime:   time.Since(startTime).Seconds(),
		log.FnProtocol:       r.Proto,
		log.FnHTTPStatusCode: lw.status,
		log.FnHTTPMethod:     r.Method,
		log.FnURL:            r.RequestURI,
		log.FnHTTPHost:       r.Host,
		log.FnRequestSize:    r.ContentLength,
		log.FnResponseSize:   lw.size,
	}
	ip, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		fields[log.FnRemoteAddress] = ip
	}
	ua := r.Header.Get("User-Agent")
	if len(ua) > 0 {
		fields[log.FnHTTPUserAgent] = ua
	}
	reqid := r.Header.Get(uuidHeaderName)
	if len(reqid) > 0 {
		fields[log.FnRequestID] = reqid
	}

	lv := log.LvInfo
	switch {
	case 500 <= lw.status:
		lv = log.LvError
	case 400 <= lw.status:
		lv = log.LvWarn
	}
	s.AccessLog.Log(lv, "cmd: "+http.StatusText(lw.status), fields)
}

func (s *HTTPServer) init() {
	if s.handler != nil {
		return
	}

	s.idleConns = make(map[net.Conn]struct{}, 100000)

	if s.Server.Handler == nil {
		panic("Handler must not be nil")
	}
	s.handler = s.Server.Handler
	s.Server.Handler = s
	if s.Server.ReadTimeout == 0 {
		s.Server.ReadTimeout = defaultHTTPReadTimeout
	}
	s.Server.ConnState = func(c net.Conn, state http.ConnState) {
		s.mu.Lock()
		if state == http.StateIdle {
			s.idleConns[c] = struct{}{}
		} else {
			delete(s.idleConns, c)
		}
		s.mu.Unlock()

		if state == http.StateNew {
			s.wg.Add(1)
			return
		}
		if state == http.StateHijacked || state == http.StateClosed {
			s.wg.Done()
		}
	}

	if s.AccessLog == nil {
		s.AccessLog = log.DefaultLogger()
	}

	if s.Env == nil {
		s.Env = defaultEnv
	}
	s.Env.Go(s.wait)
}

func (s *HTTPServer) wait(ctx context.Context) error {
	<-ctx.Done()

	s.Server.SetKeepAlivesEnabled(false)

	ch := make(chan struct{})

	// Interrupt conn.Read for idle connections.
	//
	// This must be run inside for-loop to catch connections
	// going idle at critical timing to acquire s.mu
	go func() {
	AGAIN:
		s.mu.Lock()
		for conn := range s.idleConns {
			conn.SetReadDeadline(time.Now())
		}
		s.mu.Unlock()
		select {
		case <-ch:
			return
		default:
		}
		time.Sleep(10 * time.Millisecond)
		goto AGAIN
	}()

	go func() {
		s.wg.Wait()
		close(ch)
	}()

	if s.ShutdownTimeout == 0 {
		<-ch
		return nil
	}

	select {
	case <-ch:
	case <-time.After(s.ShutdownTimeout):
		log.Warn("cmd: timeout waiting for shutdown", nil)
		atomic.StoreInt32(&s.timedout, 1)
	}
	return nil
}

// TimedOut returns true if the server shut down before all connections
// got closed.
func (s *HTTPServer) TimedOut() bool {
	return atomic.LoadInt32(&s.timedout) != 0
}

// Serve overrides http.Server's Serve method.
//
// Unlike the original, this method returns immediately just after
// starting a goroutine to accept connections.
//
// The framework automatically closes l when the environment's Cancel
// is called.
//
// Serve always returns nil.
func (s *HTTPServer) Serve(l net.Listener) error {
	s.initOnce.Do(s.init)

	go func() {
		<-s.Env.ctx.Done()
		l.Close()
	}()

	go func() {
		s.Server.Serve(l)
	}()

	return nil
}

type tcpKeepAliveListener struct {
	*net.TCPListener
}

func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) {
	tc, err := ln.AcceptTCP()
	if err != nil {
		return
	}
	tc.SetKeepAlive(true)
	tc.SetKeepAlivePeriod(3 * time.Minute)
	return tc, nil
}

// ListenAndServe overrides http.Server's method.
//
// Unlike the original, this method returns immediately just after
// starting a goroutine to accept connections.  To stop listening,
// call the environment's Cancel.
//
// ListenAndServe returns non-nil error if and only if net.Listen failed.
func (s *HTTPServer) ListenAndServe() error {
	addr := s.Server.Addr
	if addr == "" {
		addr = ":http"
	}
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	return s.Serve(tcpKeepAliveListener{ln.(*net.TCPListener)})
}

// ListenAndServeTLS overrides http.Server's method.
//
// Unlike the original, this method returns immediately just after
// starting a goroutine to accept connections.  To stop listening,
// call the environment's Cancel.
//
// Another difference from the original is that certFile and keyFile
// must be specified.  If not, configure http.Server.TLSConfig
// manually and use Serve().
//
// HTTP/2 is always enabled.
//
// ListenAndServeTLS returns non-nil error if net.Listen failed
// or failed to load certificate files.
func (s *HTTPServer) ListenAndServeTLS(certFile, keyFile string) error {
	addr := s.Server.Addr
	if addr == "" {
		addr = ":https"
	}

	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}

	config := &tls.Config{
		NextProtos:               []string{"h2", "http/1.1"},
		Certificates:             []tls.Certificate{cert},
		PreferServerCipherSuites: true,
		ClientSessionCache:       tls.NewLRUClientSessionCache(0),
	}
	s.Server.TLSConfig = config

	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	tlsListener := tls.NewListener(tcpKeepAliveListener{ln.(*net.TCPListener)}, config)
	return s.Serve(tlsListener)
}
<|endoftext|>


<commit_before>package wxpay

import (
	"bytes"
	"crypto/tls"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
)

// AppTrans is abstact of Transaction handler. With AppTrans, we can get prepay id
type AppTrans struct {
	Config *WxConfig
}

// Initialized the AppTrans with specific config
func NewAppTrans(cfg *WxConfig) (*AppTrans, error) {
	if cfg.AppId == "" ||
		cfg.MchId == "" ||
		cfg.AppKey == "" ||
		cfg.NotifyUrl == "" ||
		cfg.QueryOrderUrl == "" ||
		cfg.PlaceOrderUrl == "" ||
		cfg.TradeType == "" {
		return &AppTrans{Config: cfg}, errors.New("config field canot empty string")
	}

	return &AppTrans{Config: cfg}, nil
}

// Submit the order to weixin pay and return the prepay id if success,
// Prepay id is used for app to start a payment
// If fail, error is not nil, check error for more information
func (this *AppTrans) Submit(orderId string, amount float64, desc string, clientIp string) (string, error) {

	odrInXml := this.signedOrderRequestXmlString(orderId, fmt.Sprintf("%.0f", amount), desc, clientIp)
	resp, err := doHttpPost(this.Config.PlaceOrderUrl, []byte(odrInXml))
	if err != nil {
		return "", err
	}

	placeOrderResult, err := ParsePlaceOrderResult(resp)
	if err != nil {
		return "", err
	}

	//Verify the sign of response
	resultInMap := placeOrderResult.ToMap()
	wantSign := Sign(resultInMap, this.Config.AppKey)
	gotSign := resultInMap["sign"]
	if wantSign != gotSign {
		return "", fmt.Errorf("sign not match, want:%s, got:%s", wantSign, gotSign)
	}

	if placeOrderResult.ReturnCode != "SUCCESS" {
		return "", fmt.Errorf("return code:%s, return desc:", placeOrderResult.ReturnCode, placeOrderResult.ReturnMsg)
	}

	if placeOrderResult.ResultCode != "SUCCESS" {
		return "", fmt.Errorf("resutl code:%s, result desc:%s", placeOrderResult.ErrCode, placeOrderResult.ErrCodeDesc)
	}

	return placeOrderResult.PrepayId, nil
}

func (this *AppTrans) newQueryXml(transId string) string {
	param := make(map[string]string)
	param["appid"] = this.Config.AppId
	param["mch_id"] = this.Config.MchId
	param["transaction_id"] = transId
	param["nonce_str"] = NewNonceString()

	sign := Sign(param, this.Config.AppKey)
	param["sign"] = sign

	return ToXmlString(param)
}

// Query the order from weixin pay server by transaction id of weixin pay
func (this *AppTrans) Query(transId string) (QueryOrderResult, error) {
	queryOrderResult := QueryOrderResult{}

	queryXml := this.newQueryXml(transId)
	// fmt.Println(queryXml)
	resp, err := doHttpPost(this.Config.QueryOrderUrl, []byte(queryXml))
	if err != nil {
		return queryOrderResult, nil
	}

	queryOrderResult, err = ParseQueryOrderResult(resp)
	if err != nil {
		return queryOrderResult, err
	}

	//verity sign of response
	resultInMap := queryOrderResult.ToMap()
	wantSign := Sign(resultInMap, this.Config.AppKey)
	gotSign := resultInMap["sign"]
	if wantSign != gotSign {
		return queryOrderResult, fmt.Errorf("sign not match, want:%s, got:%s", wantSign, gotSign)
	}

	return queryOrderResult, nil
}

// NewPaymentRequest build the payment request structure for app to start a payment.
// Return stuct of PaymentRequest, please refer to http://pay.weixin.qq.com/wiki/doc/api/app.php?chapter=9_12&index=2
func (this *AppTrans) NewPaymentRequest(prepayId string) PaymentRequest {
	param := make(map[string]string)
	param["appid"] = this.Config.AppId
	param["partnerid"] = this.Config.MchId
	param["prepayid"] = prepayId
	param["package"] = "Sign=WXPay"
	param["noncestr"] = NewNonceString()
	param["timestamp"] = NewTimestampString()

	sign := Sign(param, this.Config.AppKey)

	payRequest := PaymentRequest{
		AppId:     this.Config.AppId,
		PartnerId: this.Config.MchId,
		PrepayId:  prepayId,
		Package:   "Sign=WXPay",
		NonceStr:  NewNonceString(),
		Timestamp: NewTimestampString(),
		Sign:      sign,
	}

	return payRequest
}

func (this *AppTrans) newOrderRequest(orderId, amount, desc, clientIp string) map[string]string {
	param := make(map[string]string)
	param["appid"] = this.Config.AppId
	param["attach"] = "透传字段" //optional
	param["body"] = desc
	param["mch_id"] = this.Config.MchId
	param["nonce_str"] = NewNonceString()
	param["notify_url"] = this.Config.NotifyUrl
	param["out_trade_no"] = orderId
	param["spbill_create_ip"] = clientIp
	param["total_fee"] = amount
	param["trade_type"] = "APP"

	return param
}

func (this *AppTrans) signedOrderRequestXmlString(orderId, amount, desc, clientIp string) string {
	order := this.newOrderRequest(orderId, amount, desc, clientIp)
	sign := Sign(order, this.Config.AppKey)
	// fmt.Println(sign)

	order["sign"] = sign

	return ToXmlString(order)
}

// doRequest post the order in xml format with a sign
func doHttpPost(targetUrl string, body []byte) ([]byte, error) {
	req, err := http.NewRequest("POST", targetUrl, bytes.NewBuffer([]byte(body)))
	if err != nil {
		return []byte(""), err
	}
	req.Header.Add("Content-type", "application/x-www-form-urlencoded;charset=UTF-8")

	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: false},
	}
	client := &http.Client{Transport: tr}

	resp, err := client.Do(req)
	if err != nil {
		return []byte(""), err
	}

	defer resp.Body.Close()
	respData, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return []byte(""), err
	}

	return respData, nil
}
<commit_msg>fix typo<commit_after>package wxpay

import (
	"bytes"
	"crypto/tls"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
)

// AppTrans is abstact of Transaction handler. With AppTrans, we can get prepay id
type AppTrans struct {
	Config *WxConfig
}

// Initialized the AppTrans with specific config
func NewAppTrans(cfg *WxConfig) (*AppTrans, error) {
	if cfg.AppId == "" ||
		cfg.MchId == "" ||
		cfg.AppKey == "" ||
		cfg.NotifyUrl == "" ||
		cfg.QueryOrderUrl == "" ||
		cfg.PlaceOrderUrl == "" ||
		cfg.TradeType == "" {
		return &AppTrans{Config: cfg}, errors.New("config field canot empty string")
	}

	return &AppTrans{Config: cfg}, nil
}

// Submit the order to weixin pay and return the prepay id if success,
// Prepay id is used for app to start a payment
// If fail, error is not nil, check error for more information
func (this *AppTrans) Submit(orderId string, amount float64, desc string, clientIp string) (string, error) {

	odrInXml := this.signedOrderRequestXmlString(orderId, fmt.Sprintf("%.0f", amount), desc, clientIp)
	resp, err := doHttpPost(this.Config.PlaceOrderUrl, []byte(odrInXml))
	if err != nil {
		return "", err
	}

	placeOrderResult, err := ParsePlaceOrderResult(resp)
	if err != nil {
		return "", err
	}

	//Verify the sign of response
	resultInMap := placeOrderResult.ToMap()
	wantSign := Sign(resultInMap, this.Config.AppKey)
	gotSign := resultInMap["sign"]
	if wantSign != gotSign {
		return "", fmt.Errorf("sign not match, want:%s, got:%s", wantSign, gotSign)
	}

	if placeOrderResult.ReturnCode != "SUCCESS" {
		return "", fmt.Errorf("return code:%s, return desc:%s", placeOrderResult.ReturnCode, placeOrderResult.ReturnMsg)
	}

	if placeOrderResult.ResultCode != "SUCCESS" {
		return "", fmt.Errorf("resutl code:%s, result desc:%s", placeOrderResult.ErrCode, placeOrderResult.ErrCodeDesc)
	}

	return placeOrderResult.PrepayId, nil
}

func (this *AppTrans) newQueryXml(transId string) string {
	param := make(map[string]string)
	param["appid"] = this.Config.AppId
	param["mch_id"] = this.Config.MchId
	param["transaction_id"] = transId
	param["nonce_str"] = NewNonceString()

	sign := Sign(param, this.Config.AppKey)
	param["sign"] = sign

	return ToXmlString(param)
}

// Query the order from weixin pay server by transaction id of weixin pay
func (this *AppTrans) Query(transId string) (QueryOrderResult, error) {
	queryOrderResult := QueryOrderResult{}

	queryXml := this.newQueryXml(transId)
	// fmt.Println(queryXml)
	resp, err := doHttpPost(this.Config.QueryOrderUrl, []byte(queryXml))
	if err != nil {
		return queryOrderResult, nil
	}

	queryOrderResult, err = ParseQueryOrderResult(resp)
	if err != nil {
		return queryOrderResult, err
	}

	//verity sign of response
	resultInMap := queryOrderResult.ToMap()
	wantSign := Sign(resultInMap, this.Config.AppKey)
	gotSign := resultInMap["sign"]
	if wantSign != gotSign {
		return queryOrderResult, fmt.Errorf("sign not match, want:%s, got:%s", wantSign, gotSign)
	}

	return queryOrderResult, nil
}

// NewPaymentRequest build the payment request structure for app to start a payment.
// Return stuct of PaymentRequest, please refer to http://pay.weixin.qq.com/wiki/doc/api/app.php?chapter=9_12&index=2
func (this *AppTrans) NewPaymentRequest(prepayId string) PaymentRequest {
	param := make(map[string]string)
	param["appid"] = this.Config.AppId
	param["partnerid"] = this.Config.MchId
	param["prepayid"] = prepayId
	param["package"] = "Sign=WXPay"
	param["noncestr"] = NewNonceString()
	param["timestamp"] = NewTimestampString()

	sign := Sign(param, this.Config.AppKey)

	payRequest := PaymentRequest{
		AppId:     this.Config.AppId,
		PartnerId: this.Config.MchId,
		PrepayId:  prepayId,
		Package:   "Sign=WXPay",
		NonceStr:  NewNonceString(),
		Timestamp: NewTimestampString(),
		Sign:      sign,
	}

	return payRequest
}

func (this *AppTrans) newOrderRequest(orderId, amount, desc, clientIp string) map[string]string {
	param := make(map[string]string)
	param["appid"] = this.Config.AppId
	param["attach"] = "透传字段" //optional
	param["body"] = desc
	param["mch_id"] = this.Config.MchId
	param["nonce_str"] = NewNonceString()
	param["notify_url"] = this.Config.NotifyUrl
	param["out_trade_no"] = orderId
	param["spbill_create_ip"] = clientIp
	param["total_fee"] = amount
	param["trade_type"] = "APP"

	return param
}

func (this *AppTrans) signedOrderRequestXmlString(orderId, amount, desc, clientIp string) string {
	order := this.newOrderRequest(orderId, amount, desc, clientIp)
	sign := Sign(order, this.Config.AppKey)
	// fmt.Println(sign)

	order["sign"] = sign

	return ToXmlString(order)
}

// doRequest post the order in xml format with a sign
func doHttpPost(targetUrl string, body []byte) ([]byte, error) {
	req, err := http.NewRequest("POST", targetUrl, bytes.NewBuffer([]byte(body)))
	if err != nil {
		return []byte(""), err
	}
	req.Header.Add("Content-type", "application/x-www-form-urlencoded;charset=UTF-8")

	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: false},
	}
	client := &http.Client{Transport: tr}

	resp, err := client.Do(req)
	if err != nil {
		return []byte(""), err
	}

	defer resp.Body.Close()
	respData, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return []byte(""), err
	}

	return respData, nil
}
<|endoftext|>


<commit_before>package disk

import (
	"context"
	"io"

	"github.com/buchgr/bazel-remote/cache"

	pb "github.com/buchgr/bazel-remote/genproto/build/bazel/remote/execution/v2"

	"github.com/prometheus/client_golang/prometheus"
)

type metricsDecorator struct {
	counter *prometheus.CounterVec
	*diskCache
}

const (
	hitStatus  = "hit"
	missStatus = "miss"

	containsMethod = "contains"
	getMethod      = "get"
	//putMethod      = "put"

	acKind  = "ac" // This must be lowercase to match cache.EntryKind.String()
	casKind = "cas"
	rawKind = "raw"
)

func (m *metricsDecorator) RegisterMetrics() {
	prometheus.MustRegister(m.counter)
	m.diskCache.RegisterMetrics()
}

func (m *metricsDecorator) Get(ctx context.Context, kind cache.EntryKind, hash string, size int64, offset int64) (io.ReadCloser, int64, error) {
	rc, size, err := m.diskCache.Get(ctx, kind, hash, size, offset)

	lbls := prometheus.Labels{"method": getMethod, "kind": kind.String()}
	if rc != nil {
		lbls["status"] = hitStatus
	} else if err == nil {
		lbls["status"] = missStatus
	}
	m.counter.With(lbls).Inc()

	return rc, size, err
}

func (m *metricsDecorator) GetValidatedActionResult(ctx context.Context, hash string) (*pb.ActionResult, []byte, error) {
	ar, data, err := m.diskCache.GetValidatedActionResult(ctx, hash)

	lbls := prometheus.Labels{"method": getMethod, "kind": acKind}
	if ar != nil {
		lbls["status"] = hitStatus
	} else if err == nil {
		lbls["status"] = missStatus
	}
	m.counter.With(lbls).Inc()

	return ar, data, err
}

func (m *metricsDecorator) GetZstd(ctx context.Context, hash string, size int64, offset int64) (io.ReadCloser, int64, error) {
	rc, size, err := m.diskCache.GetZstd(ctx, hash, size, offset)

	lbls := prometheus.Labels{
		"method": getMethod,
		"kind":   "cas",
	}
	if rc != nil {
		lbls["status"] = hitStatus
	} else if err == nil {
		lbls["status"] = missStatus
	}
	m.counter.With(lbls).Inc()

	return rc, size, err
}

func (m *metricsDecorator) Contains(ctx context.Context, kind cache.EntryKind, hash string, size int64) (bool, int64) {
	ok, size := m.diskCache.Contains(ctx, kind, hash, size)

	lbls := prometheus.Labels{"method": containsMethod, "kind": kind.String()}
	if ok {
		lbls["status"] = hitStatus
	} else {
		lbls["status"] = missStatus
	}
	m.counter.With(lbls).Inc()

	return ok, size
}

func (m *metricsDecorator) FindMissingCasBlobs(ctx context.Context, blobs []*pb.Digest) ([]*pb.Digest, error) {
	numLooking := len(blobs)
	digests, err := m.diskCache.FindMissingCasBlobs(ctx, blobs)
	numFound := len(digests)

	numMissing := numLooking - numFound

	hitLabels := prometheus.Labels{
		"method": containsMethod,
		"kind":   "cas",
		"status": hitStatus,
	}
	hits := m.counter.With(hitLabels)

	missLabels := prometheus.Labels{
		"method": containsMethod,
		"kind":   "cas",
		"status": missStatus,
	}
	misses := m.counter.With(missLabels)

	hits.Add(float64(numFound))
	misses.Add(float64(numMissing))

	return digests, err
}
<commit_msg>Avoid prometheus counter panic<commit_after>package disk

import (
	"context"
	"io"

	"github.com/buchgr/bazel-remote/cache"

	pb "github.com/buchgr/bazel-remote/genproto/build/bazel/remote/execution/v2"

	"github.com/prometheus/client_golang/prometheus"
)

type metricsDecorator struct {
	counter *prometheus.CounterVec
	*diskCache
}

const (
	hitStatus  = "hit"
	missStatus = "miss"

	containsMethod = "contains"
	getMethod      = "get"
	//putMethod      = "put"

	acKind  = "ac" // This must be lowercase to match cache.EntryKind.String()
	casKind = "cas"
	rawKind = "raw"
)

func (m *metricsDecorator) RegisterMetrics() {
	prometheus.MustRegister(m.counter)
	m.diskCache.RegisterMetrics()
}

func (m *metricsDecorator) Get(ctx context.Context, kind cache.EntryKind, hash string, size int64, offset int64) (io.ReadCloser, int64, error) {
	rc, size, err := m.diskCache.Get(ctx, kind, hash, size, offset)
	if err != nil {
		return rc, size, err
	}

	lbls := prometheus.Labels{"method": getMethod, "kind": kind.String()}
	if rc != nil {
		lbls["status"] = hitStatus
	} else {
		lbls["status"] = missStatus
	}
	m.counter.With(lbls).Inc()

	return rc, size, nil
}

func (m *metricsDecorator) GetValidatedActionResult(ctx context.Context, hash string) (*pb.ActionResult, []byte, error) {
	ar, data, err := m.diskCache.GetValidatedActionResult(ctx, hash)
	if err != nil {
		return ar, data, err
	}

	lbls := prometheus.Labels{"method": getMethod, "kind": acKind}
	if ar != nil {
		lbls["status"] = hitStatus
	} else {
		lbls["status"] = missStatus
	}
	m.counter.With(lbls).Inc()

	return ar, data, err
}

func (m *metricsDecorator) GetZstd(ctx context.Context, hash string, size int64, offset int64) (io.ReadCloser, int64, error) {
	rc, size, err := m.diskCache.GetZstd(ctx, hash, size, offset)
	if err != nil {
		return rc, size, err
	}

	lbls := prometheus.Labels{
		"method": getMethod,
		"kind":   "cas",
	}
	if rc != nil {
		lbls["status"] = hitStatus
	} else {
		lbls["status"] = missStatus
	}
	m.counter.With(lbls).Inc()

	return rc, size, nil
}

func (m *metricsDecorator) Contains(ctx context.Context, kind cache.EntryKind, hash string, size int64) (bool, int64) {
	ok, size := m.diskCache.Contains(ctx, kind, hash, size)

	lbls := prometheus.Labels{"method": containsMethod, "kind": kind.String()}
	if ok {
		lbls["status"] = hitStatus
	} else {
		lbls["status"] = missStatus
	}
	m.counter.With(lbls).Inc()

	return ok, size
}

func (m *metricsDecorator) FindMissingCasBlobs(ctx context.Context, blobs []*pb.Digest) ([]*pb.Digest, error) {
	numLooking := len(blobs)
	digests, err := m.diskCache.FindMissingCasBlobs(ctx, blobs)
	if err != nil {
		return digests, err
	}

	numFound := len(digests)

	numMissing := numLooking - numFound

	hitLabels := prometheus.Labels{
		"method": containsMethod,
		"kind":   "cas",
		"status": hitStatus,
	}
	hits := m.counter.With(hitLabels)

	missLabels := prometheus.Labels{
		"method": containsMethod,
		"kind":   "cas",
		"status": missStatus,
	}
	misses := m.counter.With(missLabels)

	hits.Add(float64(numFound))
	misses.Add(float64(numMissing))

	return digests, nil
}
<|endoftext|>