<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 != "" {
		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("[%s] Failed to join: %s", term.Conn.RemoteAddr(), 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("[%s] Joined: %s", term.Conn.RemoteAddr(), user.Name())

	for {
		line, err := term.ReadLine()
		if err == io.EOF {
			// Closed
			break
		} else if err != nil {
			logger.Errorf("[%s] Terminal reading error: %s", term.Conn.RemoteAddr(), 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("[%s] Failed to leave: %s", term.Conn.RemoteAddr(), err)
		return
	}
	logger.Debugf("[%s] Leaving: %s", term.Conn.RemoteAddr(), 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)

			txt := fmt.Sprintf("[Sent PM to %s]", target.Name())
			ms := message.NewSystemMsg(txt, msg.From())
			room.Send(ms)
			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)

			txt := fmt.Sprintf("[Sent PM to %s]", target.Name())
			ms := message.NewSystemMsg(txt, msg.From())
			room.Send(ms)
			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)
			var whois string
			switch room.IsOp(msg.From()) {
			case true:
				whois = id.WhoisAdmin()
			case false:
				whois = id.Whois()
			}
			room.Send(message.NewSystemMsg(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 a new MESSAGE of the day, print the current motd without parameters.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			args := msg.Args()
			user := msg.From()

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

			if len(args) == 0 {
				room.Send(message.NewSystemMsg(motd, user))
				return nil
			}
			if !room.IsOp(user) {
				return errors.New("must be OP to modify the MOTD")
			}

			motd = strings.Join(args, " ")
			h.SetMotd(motd)
			fromMsg := fmt.Sprintf("New message of the day set by %s:", msg.From().Name())
			room.Send(message.NewAnnounceMsg(fromMsg + message.Newline + "-> " + 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>/reply: Enabled again.<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 != "" {
		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("[%s] Failed to join: %s", term.Conn.RemoteAddr(), 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("[%s] Joined: %s", term.Conn.RemoteAddr(), user.Name())

	for {
		line, err := term.ReadLine()
		if err == io.EOF {
			// Closed
			break
		} else if err != nil {
			logger.Errorf("[%s] Terminal reading error: %s", term.Conn.RemoteAddr(), 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("[%s] Failed to leave: %s", term.Conn.RemoteAddr(), err)
		return
	}
	logger.Debugf("[%s] Leaving: %s", term.Conn.RemoteAddr(), 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)

			txt := fmt.Sprintf("[Sent PM to %s]", target.Name())
			ms := message.NewSystemMsg(txt, msg.From())
			room.Send(ms)
			target.SetReplyTo(msg.From())
			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)

			txt := fmt.Sprintf("[Sent PM to %s]", target.Name())
			ms := message.NewSystemMsg(txt, msg.From())
			room.Send(ms)
			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)
			var whois string
			switch room.IsOp(msg.From()) {
			case true:
				whois = id.WhoisAdmin()
			case false:
				whois = id.Whois()
			}
			room.Send(message.NewSystemMsg(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 a new MESSAGE of the day, print the current motd without parameters.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			args := msg.Args()
			user := msg.From()

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

			if len(args) == 0 {
				room.Send(message.NewSystemMsg(motd, user))
				return nil
			}
			if !room.IsOp(user) {
				return errors.New("must be OP to modify the MOTD")
			}

			motd = strings.Join(args, " ")
			h.SetMotd(motd)
			fromMsg := fmt.Sprintf("New message of the day set by %s:", msg.From().Name())
			room.Send(message.NewAnnounceMsg(fromMsg + message.Newline + "-> " + 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 jsonschema

import (
	"encoding/json"
	"fmt"
	"net/url"
	"path"
)

type HRef struct {
	id      int
	value   string
	example string
}

func (h *HRef) UnmarshalJSON(data []byte) error {
	var s string
	if err := json.Unmarshal(data, &s); err != nil {
		return err
	}
	h.value = rBraceBracket.ReplaceAllStringFunc(s, func(s string) string {
		d, err := url.QueryUnescape(s)
		if err != nil {
			return s
		}
		return d
	})
	return nil
}

// func (h *HRef) Resolve(schemas *map[string]*Schema) error {
// 	h.example = h.replaceBraceBracket(func(s string) string {
// 		schema := (*schemas)[s]
// 		return schema.Example.String()
// 	})
// 	h.id = rand.Int()
// 	return nil
// }

func (h HRef) String() string {
	return h.value
}

func (h HRef) ColonString() string {
	return h.replaceBraceBracket(func(s string) string {
		return fmt.Sprintf(":%s", path.Base(s))
	})
}

func (h *HRef) ExampleString() string {
	return h.example
}

func (h HRef) replaceBraceBracket(replacer func(string) string) string {
	return rBraceBracket.ReplaceAllStringFunc(h.value, func(s string) string {
		m := rBraceBracket.FindStringSubmatch(s)
		return replacer(m[1])
	})
}
<commit_msg>Render endpoint with example id<commit_after>package jsonschema

import (
	"encoding/json"
	"fmt"
	"log"
	"net/url"
	"path"
)

type HRef struct {
	value   string
	example string
}

func (h *HRef) UnmarshalJSON(data []byte) error {
	var s string
	if err := json.Unmarshal(data, &s); err != nil {
		return err
	}
	h.value = rBraceBracket.ReplaceAllStringFunc(s, func(s string) string {
		d, err := url.QueryUnescape(s)
		if err != nil {
			return s
		}
		return d
	})
	return nil
}

func (h *HRef) Resolve(schemas *map[string]*Schema) error {
	h.example = h.replaceBraceBracket(func(s string) string {
		schema := (*schemas)[s]
		return schema.Example.RawString()
	})
	log.Println(h.example)
	return nil
}

func (h HRef) String() string {
	return h.value
}

func (h HRef) ColonString() string {
	return h.replaceBraceBracket(func(s string) string {
		return fmt.Sprintf(":%s", path.Base(s))
	})
}

func (h *HRef) ExampleString() string {
	return h.example
}

func (h HRef) replaceBraceBracket(replacer func(string) string) string {
	return rBraceBracket.ReplaceAllStringFunc(h.value, func(s string) string {
		m := rBraceBracket.FindStringSubmatch(s)
		return replacer(m[1])
	})
}
<|endoftext|>


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

package com

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

type NotFoundError struct {
	Message string
}

func (e NotFoundError) Error() string {
	return e.Message
}

type RemoteError struct {
	Host string
	Err  error
}

func (e *RemoteError) Error() string {
	return e.Err.Error()
}

var UserAgent = "Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/29.0.1541.0 Safari/537.36"

// HttpGet gets the specified resource. ErrNotFound is returned if the
// server responds with status 404.
func HttpGet(client *http.Client, url string, header http.Header) (io.ReadCloser, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", UserAgent)
	for k, vs := range header {
		req.Header[k] = vs
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, &RemoteError{req.URL.Host, err}
	}
	if resp.StatusCode == 200 {
		return resp.Body, nil
	}
	resp.Body.Close()
	if resp.StatusCode == 404 { // 403 can be rate limit error.  || resp.StatusCode == 403 {
		err = NotFoundError{"Resource not found: " + url}
	} else {
		err = &RemoteError{req.URL.Host, fmt.Errorf("get %s -> %d", url, resp.StatusCode)}
	}
	return nil, err
}

// HttpGetBytes gets the specified resource. ErrNotFound is returned if the server
// responds with status 404.
func HttpGetBytes(client *http.Client, url string, header http.Header) ([]byte, error) {
	rc, err := HttpGet(client, url, header)
	if err != nil {
		return nil, err
	}
	p, err := ioutil.ReadAll(rc)
	rc.Close()
	return p, err
}

// HttpGetJSON gets the specified resource and mapping to struct.
// ErrNotFound is returned if the server responds with status 404.
func HttpGetJSON(client *http.Client, url string, v interface{}) error {
	rc, err := HttpGet(client, url, nil)
	if err != nil {
		return err
	}
	defer rc.Close()
	err = json.NewDecoder(rc).Decode(v)
	if _, ok := err.(*json.SyntaxError); ok {
		err = NotFoundError{"JSON syntax error at " + url}
	}
	return err
}

// A RawFile describes a file that can be downloaded.
type RawFile interface {
	Name() string
	RawUrl() string
	Data() []byte
	SetData([]byte)
}

// FetchFiles fetches files specified by the rawURL field in parallel.
func FetchFiles(client *http.Client, files []RawFile, header http.Header) error {
	ch := make(chan error, len(files))
	for i := range files {
		go func(i int) {
			p, err := HttpGetBytes(client, files[i].RawUrl(), nil)
			if err != nil {
				ch <- err
				return
			}
			files[i].SetData(p)
			ch <- nil
		}(i)
	}
	for _ = range files {
		if err := <-ch; err != nil {
			return err
		}
	}
	return nil
}

// FetchFiles uses command `curl` to fetch files specified by the rawURL field in parallel.
func FetchFilesCurl(files []RawFile, curlOptions ...string) error {
	ch := make(chan error, len(files))
	for i := range files {
		go func(i int) {
			stdout, _, err := ExecCmd("curl", append(curlOptions, files[i].RawUrl())...)
			if err != nil {
				ch <- err
				return
			}

			files[i].SetData([]byte(stdout))
			ch <- nil
		}(i)
	}
	for _ = range files {
		if err := <-ch; err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>fix RawURL method bug in http.go<commit_after>// Copyright 2013 com authors
//
// Licensed under the Apache License, Version 2.0 (the "License"): you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

package com

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

type NotFoundError struct {
	Message string
}

func (e NotFoundError) Error() string {
	return e.Message
}

type RemoteError struct {
	Host string
	Err  error
}

func (e *RemoteError) Error() string {
	return e.Err.Error()
}

var UserAgent = "Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/29.0.1541.0 Safari/537.36"

// HttpGet gets the specified resource. ErrNotFound is returned if the
// server responds with status 404.
func HttpGet(client *http.Client, url string, header http.Header) (io.ReadCloser, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", UserAgent)
	for k, vs := range header {
		req.Header[k] = vs
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, &RemoteError{req.URL.Host, err}
	}
	if resp.StatusCode == 200 {
		return resp.Body, nil
	}
	resp.Body.Close()
	if resp.StatusCode == 404 { // 403 can be rate limit error.  || resp.StatusCode == 403 {
		err = NotFoundError{"Resource not found: " + url}
	} else {
		err = &RemoteError{req.URL.Host, fmt.Errorf("get %s -> %d", url, resp.StatusCode)}
	}
	return nil, err
}

// HttpGetBytes gets the specified resource. ErrNotFound is returned if the server
// responds with status 404.
func HttpGetBytes(client *http.Client, url string, header http.Header) ([]byte, error) {
	rc, err := HttpGet(client, url, header)
	if err != nil {
		return nil, err
	}
	p, err := ioutil.ReadAll(rc)
	rc.Close()
	return p, err
}

// HttpGetJSON gets the specified resource and mapping to struct.
// ErrNotFound is returned if the server responds with status 404.
func HttpGetJSON(client *http.Client, url string, v interface{}) error {
	rc, err := HttpGet(client, url, nil)
	if err != nil {
		return err
	}
	defer rc.Close()
	err = json.NewDecoder(rc).Decode(v)
	if _, ok := err.(*json.SyntaxError); ok {
		err = NotFoundError{"JSON syntax error at " + url}
	}
	return err
}

// A RawFile describes a file that can be downloaded.
type RawFile interface {
	Name() string
	RawURL() string
	Data() []byte
	SetData([]byte)
}

// FetchFiles fetches files specified by the rawURL field in parallel.
func FetchFiles(client *http.Client, files []RawFile, header http.Header) error {
	ch := make(chan error, len(files))
	for i := range files {
		go func(i int) {
			p, err := HttpGetBytes(client, files[i].RawURL(), nil)
			if err != nil {
				ch <- err
				return
			}
			files[i].SetData(p)
			ch <- nil
		}(i)
	}
	for _ = range files {
		if err := <-ch; err != nil {
			return err
		}
	}
	return nil
}

// FetchFiles uses command `curl` to fetch files specified by the rawURL field in parallel.
func FetchFilesCurl(files []RawFile, curlOptions ...string) error {
	ch := make(chan error, len(files))
	for i := range files {
		go func(i int) {
			stdout, _, err := ExecCmd("curl", append(curlOptions, files[i].RawURL())...)
			if err != nil {
				ch <- err
				return
			}

			files[i].SetData([]byte(stdout))
			ch <- nil
		}(i)
	}
	for _ = range files {
		if err := <-ch; err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package word2vec

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

type Expr struct {
	Add []string `json:"add,omitempty"`
	Sub []string `json:"sub,omitempty"`
}

func (e Expr) Eval(m *Model) (Vector, error) {
	if len(e.Add) == 0 && len(e.Sub) == 0 {
		return nil, fmt.Errorf("must specify either 'add' or 'sub'")
	}
	return m.Eval(e.Add, e.Sub)
}

type SimQuery struct {
	A Expr `json:"a,omitempty"`
	B Expr `json:"b,omitempty"`
}

type SimResponse struct {
	Value float32 `json:"value"`
}

func (q SimQuery) Eval(m *Model) (*SimResponse, error) {
	v, err := q.A.Eval(m)
	if err != nil {
		return nil, err
	}
	u, err := q.B.Eval(m)
	if err != nil {
		return nil, err
	}

	return &SimResponse{
		Value: m.Sim(u, v),
	}, nil
}

type MultiSimQuery struct {
	Queries []SimQuery `json:"queries"`
}

type MultiSimResponse struct {
	Values []SimResponse `json:"values"`
}

func (qs MultiSimQuery) Eval(m *Model) (*MultiSimResponse, error) {
	values := make([]SimResponse, len(qs.Queries))
	for i, q := range qs.Queries {
		r, err := q.Eval(m)
		if err != nil {
			return nil, err
		}
		values[i] = *r
	}
	return &MultiSimResponse{
		Values: values,
	}, nil
}

type MostSimQuery struct {
	Expr Expr `json:"expr"`
	N    int  `json:"n"`
}

type MostSimResponse struct {
	Matches []Match `json:"matches"`
}

func (q MostSimQuery) Eval(m *Model) (*MostSimResponse, error) {
	v, err := q.Expr.Eval(m)
	if err != nil {
		return nil, err
	}
	return &MostSimResponse{
		Matches: m.MostSimilar(v, q.N),
	}, nil
}

type ModelServer struct {
	*Model
}

func handleError(w http.ResponseWriter, r *http.Request, status int, msg string) {
	log.Printf(msg)
	w.WriteHeader(status)
	w.Write([]byte(msg))
	return
}

func (m *ModelServer) HandleSimQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q SimQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}

func (m *ModelServer) HandleMultiSimQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q MultiSimQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}

func (m *ModelServer) HandleMostSimQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q MostSimQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}

type Client struct {
	Addr string
}

func (c Client) Sim(x, y Expr) (float32, error) {
	req := SimQuery{A: x, B: y}

	b, err := json.Marshal(req)
	if err != nil {
		return 0.0, err
	}

	r, err := http.NewRequest("GET", "http://"+c.Addr+"/sim", bytes.NewReader(b))
	if err != nil {
		return 0.0, err
	}

	resp, err := http.DefaultClient.Do(r)
	if err != nil {
		return 0.0, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return 0.0, fmt.Errorf("non-%v status code: %v", http.StatusOK, resp.Status)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return 0.0, fmt.Errorf("error reading response: %v", err)
	}

	var data SimResponse
	err = json.Unmarshal(body, &data)
	if err != nil {
		return 0.0, fmt.Errorf("error unmarshalling result: %v", err)
	}

	return data.Value, nil
}
<commit_msg>Implemented Client.MultiSim for doing multiple similarity queries in one request.<commit_after>package word2vec

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

type Expr struct {
	Add []string `json:"add,omitempty"`
	Sub []string `json:"sub,omitempty"`
}

func (e Expr) Eval(m *Model) (Vector, error) {
	if len(e.Add) == 0 && len(e.Sub) == 0 {
		return nil, fmt.Errorf("must specify either 'add' or 'sub'")
	}
	return m.Eval(e.Add, e.Sub)
}

type SimQuery struct {
	A Expr `json:"a,omitempty"`
	B Expr `json:"b,omitempty"`
}

type SimResponse struct {
	Value float32 `json:"value"`
}

func (q SimQuery) Eval(m *Model) (*SimResponse, error) {
	v, err := q.A.Eval(m)
	if err != nil {
		return nil, err
	}
	u, err := q.B.Eval(m)
	if err != nil {
		return nil, err
	}

	return &SimResponse{
		Value: m.Sim(u, v),
	}, nil
}

type MultiSimQuery struct {
	Queries []SimQuery `json:"queries"`
}

type MultiSimResponse struct {
	Values []SimResponse `json:"values"`
}

func (qs MultiSimQuery) Eval(m *Model) (*MultiSimResponse, error) {
	values := make([]SimResponse, len(qs.Queries))
	for i, q := range qs.Queries {
		r, err := q.Eval(m)
		if err != nil {
			return nil, err
		}
		values[i] = *r
	}
	return &MultiSimResponse{
		Values: values,
	}, nil
}

type MostSimQuery struct {
	Expr Expr `json:"expr"`
	N    int  `json:"n"`
}

type MostSimResponse struct {
	Matches []Match `json:"matches"`
}

func (q MostSimQuery) Eval(m *Model) (*MostSimResponse, error) {
	v, err := q.Expr.Eval(m)
	if err != nil {
		return nil, err
	}
	return &MostSimResponse{
		Matches: m.MostSimilar(v, q.N),
	}, nil
}

type ModelServer struct {
	*Model
}

func handleError(w http.ResponseWriter, r *http.Request, status int, msg string) {
	log.Printf(msg)
	w.WriteHeader(status)
	w.Write([]byte(msg))
	return
}

func (m *ModelServer) HandleSimQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q SimQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}

func (m *ModelServer) HandleMultiSimQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q MultiSimQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}

func (m *ModelServer) HandleMostSimQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q MostSimQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}

type Client struct {
	Addr string
}

func (c Client) Sim(x, y Expr) (float32, error) {
	req := SimQuery{A: x, B: y}

	b, err := json.Marshal(req)
	if err != nil {
		return 0.0, err
	}

	r, err := http.NewRequest("GET", "http://"+c.Addr+"/sim", bytes.NewReader(b))
	if err != nil {
		return 0.0, err
	}

	resp, err := http.DefaultClient.Do(r)
	if err != nil {
		return 0.0, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return 0.0, fmt.Errorf("non-%v status code: %v", http.StatusOK, resp.Status)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return 0.0, fmt.Errorf("error reading response: %v", err)
	}

	var data SimResponse
	err = json.Unmarshal(body, &data)
	if err != nil {
		return 0.0, fmt.Errorf("error unmarshalling result: %v", err)
	}

	return data.Value, nil
}

func (c Client) MultiSim(pairs [][2]Expr) ([]float32, error) {
	req := MultiSimQuery{
		Queries: make([]SimQuery, len(pairs)),
	}
	for _, pair := range pairs {
		req.Queries = append(req.Queries, SimQuery{
			A: pair[0],
			B: pair[1],
		})
	}

	b, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	r, err := http.NewRequest("GET", "http://"+c.Addr+"/sim-multi", bytes.NewReader(b))
	if err != nil {
		return nil, err
	}

	resp, err := http.DefaultClient.Do(r)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("non-%v status code: %v", http.StatusOK, resp.Status)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("error reading response: %v", err)
	}

	var data MultiSimResponse
	err = json.Unmarshal(body, &data)
	if err != nil {
		return nil, fmt.Errorf("error unmarshalling result: %v", err)
	}

	result := make([]float32, len(data.Values))
	for i, v := range data.Values {
		result[i] = v.Value
	}
	return result, nil
}
<|endoftext|>


<commit_before>package word2vec

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

type CosineQuery struct {
	A Expr `json:"a,omitempty"`
	B Expr `json:"b,omitempty"`
}

type CosineResponse struct {
	Value float32 `json:"value"`
}

func (q CosineQuery) Eval(m *Model) (*CosineResponse, error) {
	v, err := m.Cosine(q.A, q.B)
	if err != nil {
		return nil, err
	}

	return &CosineResponse{
		Value: v,
	}, nil
}

type CosinesQuery struct {
	Queries []CosineQuery `json:"queries"`
}

type CosinesResponse struct {
	Values []CosineResponse `json:"values"`
}

func (qs CosinesQuery) Eval(m *Model) (*CosinesResponse, error) {
	values := make([]CosineResponse, len(qs.Queries))
	for i, q := range qs.Queries {
		r, err := q.Eval(m)
		if err != nil {
			return nil, err
		}
		values[i] = *r
	}
	return &CosinesResponse{
		Values: values,
	}, nil
}

type CosineNQuery struct {
	Expr Expr `json:"expr"`
	N    int  `json:"n"`
}

type CosineNResponse struct {
	Matches []Match `json:"matches"`
}

func (q CosineNQuery) Eval(m *Model) (*CosineNResponse, error) {
	r, err := m.CosineN(q.Expr, q.N)
	if err != nil {
		return nil, err
	}

	return &CosineNResponse{
		Matches: r,
	}, nil
}

type Server struct {
	*Model
	*http.ServeMux
}

func NewServer(m *Model) http.Handler {
	ms := &Server{
		Model: m,
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/cosine-n", ms.handleCosineNQuery)
	mux.HandleFunc("/cosine", ms.handleCosineQuery)
	mux.HandleFunc("/cosines", ms.handleCosinesQuery)

	ms.ServeMux = mux
	return ms
}

func handleError(w http.ResponseWriter, r *http.Request, status int, msg string) {
	log.Printf(msg)
	w.WriteHeader(status)
	w.Write([]byte(msg))
	return
}

func (m *Server) handleCosineQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q CosineQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}

func (m *Server) handleCosinesQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q CosinesQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}

func (m *Server) handleCosineNQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q CosineNQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}

type Client struct {
	Addr string
}

func (c Client) Cosine(x, y Expr) (float32, error) {
	req := CosineQuery{A: x, B: y}

	b, err := json.Marshal(req)
	if err != nil {
		return 0.0, err
	}

	r, err := http.NewRequest("GET", "http://"+c.Addr+"/cosine", bytes.NewReader(b))
	if err != nil {
		return 0.0, err
	}

	resp, err := http.DefaultClient.Do(r)
	if err != nil {
		return 0.0, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return 0.0, fmt.Errorf("non-%v status code: %v", http.StatusOK, resp.Status)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return 0.0, fmt.Errorf("error reading response: %v", err)
	}

	var data CosineResponse
	err = json.Unmarshal(body, &data)
	if err != nil {
		return 0.0, fmt.Errorf("error unmarshalling result: %v", err)
	}

	return data.Value, nil
}

func (c Client) Cosines(pairs [][2]Expr) ([]float32, error) {
	req := CosinesQuery{
		Queries: make([]CosineQuery, len(pairs)),
	}
	for _, pair := range pairs {
		req.Queries = append(req.Queries, CosineQuery{
			A: pair[0],
			B: pair[1],
		})
	}

	b, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	r, err := http.NewRequest("GET", "http://"+c.Addr+"/cosines", bytes.NewReader(b))
	if err != nil {
		return nil, err
	}

	resp, err := http.DefaultClient.Do(r)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("non-%v status code: %v", http.StatusOK, resp.Status)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("error reading response: %v", err)
	}

	var data CosinesResponse
	err = json.Unmarshal(body, &data)
	if err != nil {
		return nil, fmt.Errorf("error unmarshalling result: %v", err)
	}

	result := make([]float32, len(data.Values))
	for i, v := range data.Values {
		result[i] = v.Value
	}
	return result, nil
}

func (c Client) CosineN(e Expr, n int) ([]Match, error) {
	req := CosineNQuery{Expr: e, N: n}

	b, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	r, err := http.NewRequest("GET", "http://"+c.Addr+"/cosine-n", bytes.NewReader(b))
	if err != nil {
		return nil, err
	}

	resp, err := http.DefaultClient.Do(r)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("non-%v status code: %v", http.StatusOK, resp.Status)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("error reading response: %v", err)
	}

	var data CosineNResponse
	err = json.Unmarshal(body, &data)
	if err != nil {
		return nil, fmt.Errorf("error unmarshalling result: %v", err)
	}

	return data.Matches, nil
}
<commit_msg>Client: check for bad requests after reading body.<commit_after>package word2vec

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

type CosineQuery struct {
	A Expr `json:"a,omitempty"`
	B Expr `json:"b,omitempty"`
}

type CosineResponse struct {
	Value float32 `json:"value"`
}

func (q CosineQuery) Eval(m *Model) (*CosineResponse, error) {
	v, err := m.Cosine(q.A, q.B)
	if err != nil {
		return nil, err
	}

	return &CosineResponse{
		Value: v,
	}, nil
}

type CosinesQuery struct {
	Queries []CosineQuery `json:"queries"`
}

type CosinesResponse struct {
	Values []CosineResponse `json:"values"`
}

func (qs CosinesQuery) Eval(m *Model) (*CosinesResponse, error) {
	values := make([]CosineResponse, len(qs.Queries))
	for i, q := range qs.Queries {
		r, err := q.Eval(m)
		if err != nil {
			return nil, err
		}
		values[i] = *r
	}
	return &CosinesResponse{
		Values: values,
	}, nil
}

type CosineNQuery struct {
	Expr Expr `json:"expr"`
	N    int  `json:"n"`
}

type CosineNResponse struct {
	Matches []Match `json:"matches"`
}

func (q CosineNQuery) Eval(m *Model) (*CosineNResponse, error) {
	r, err := m.CosineN(q.Expr, q.N)
	if err != nil {
		return nil, err
	}

	return &CosineNResponse{
		Matches: r,
	}, nil
}

type Server struct {
	*Model
	*http.ServeMux
}

func NewServer(m *Model) http.Handler {
	ms := &Server{
		Model: m,
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/cosine-n", ms.handleCosineNQuery)
	mux.HandleFunc("/cosine", ms.handleCosineQuery)
	mux.HandleFunc("/cosines", ms.handleCosinesQuery)

	ms.ServeMux = mux
	return ms
}

func handleError(w http.ResponseWriter, r *http.Request, status int, msg string) {
	log.Printf(msg)
	w.WriteHeader(status)
	w.Write([]byte(msg))
	return
}

func (m *Server) handleCosineQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q CosineQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}

func (m *Server) handleCosinesQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q CosinesQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}

func (m *Server) handleCosineNQuery(w http.ResponseWriter, r *http.Request) {
	dec := json.NewDecoder(r.Body)
	defer r.Body.Close()

	var q CosineNQuery
	err := dec.Decode(&q)
	if err != nil {
		msg := fmt.Sprintf("error decoding query: %v", err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	resp, err := q.Eval(m.Model)
	if err != nil {
		msg := fmt.Sprintf("error evaluating query: %v", err)
		handleError(w, r, http.StatusBadRequest, msg)
		return
	}

	b, err := json.Marshal(resp)
	if err != nil {
		msg := fmt.Sprintf("error encoding response %#v to JSON: %v", resp, err)
		handleError(w, r, http.StatusInternalServerError, msg)
		return
	}

	_, err = w.Write(b)
	if err != nil {
		log.Printf("error writing response: %v", err)
	}
}

type Client struct {
	Addr string
}

func (c Client) Cosine(x, y Expr) (float32, error) {
	req := CosineQuery{A: x, B: y}

	b, err := json.Marshal(req)
	if err != nil {
		return 0.0, err
	}

	r, err := http.NewRequest("GET", "http://"+c.Addr+"/cosine", bytes.NewReader(b))
	if err != nil {
		return 0.0, err
	}

	resp, err := http.DefaultClient.Do(r)
	if err != nil {
		return 0.0, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return 0.0, fmt.Errorf("error reading response: %v", err)
	}

	if resp.StatusCode == http.StatusBadRequest {
		return 0.0, fmt.Errorf("error: %v", string(b))
	}

	if resp.StatusCode != http.StatusOK {
		return 0.0, fmt.Errorf("non-%v status code: %v msg: %v", http.StatusOK, resp.Status, string(b))
	}

	var data CosineResponse
	err = json.Unmarshal(body, &data)
	if err != nil {
		return 0.0, fmt.Errorf("error unmarshalling result: %v", err)
	}

	return data.Value, nil
}

func (c Client) Cosines(pairs [][2]Expr) ([]float32, error) {
	req := CosinesQuery{
		Queries: make([]CosineQuery, len(pairs)),
	}
	for _, pair := range pairs {
		req.Queries = append(req.Queries, CosineQuery{
			A: pair[0],
			B: pair[1],
		})
	}

	b, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	r, err := http.NewRequest("GET", "http://"+c.Addr+"/cosines", bytes.NewReader(b))
	if err != nil {
		return nil, err
	}

	resp, err := http.DefaultClient.Do(r)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("non-%v status code: %v", http.StatusOK, resp.Status)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("error reading response: %v", err)
	}

	var data CosinesResponse
	err = json.Unmarshal(body, &data)
	if err != nil {
		return nil, fmt.Errorf("error unmarshalling result: %v", err)
	}

	result := make([]float32, len(data.Values))
	for i, v := range data.Values {
		result[i] = v.Value
	}
	return result, nil
}

func (c Client) CosineN(e Expr, n int) ([]Match, error) {
	req := CosineNQuery{Expr: e, N: n}

	b, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	r, err := http.NewRequest("GET", "http://"+c.Addr+"/cosine-n", bytes.NewReader(b))
	if err != nil {
		return nil, err
	}

	resp, err := http.DefaultClient.Do(r)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("error reading response: %v", err)
	}

	if resp.StatusCode == http.StatusBadRequest {
		return nil, fmt.Errorf("error: %v", string(body))
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("non-%v status code: %v msg: %v", http.StatusOK, resp.Status, string(body))
	}

	var data CosineNResponse
	err = json.Unmarshal(body, &data)
	if err != nil {
		return nil, fmt.Errorf("error unmarshalling result: %v", err)
	}

	return data.Matches, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/litl/galaxy/utils"

	"github.com/litl/galaxy/log"
	gotoolslog "github.com/mailgun/gotools-log"
	"github.com/mailgun/vulcan"
	"github.com/mailgun/vulcan/endpoint"
	"github.com/mailgun/vulcan/loadbalance/roundrobin"
	"github.com/mailgun/vulcan/location/httploc"
	"github.com/mailgun/vulcan/request"
	"github.com/mailgun/vulcan/route"
	"github.com/mailgun/vulcan/route/hostroute"
)

var (
	httpRouter *HTTPRouter
)

type RequestLogger struct{}

type HTTPRouter struct {
	router    *hostroute.HostRouter
	balancers map[string]*roundrobin.RoundRobin
}

func (r *RequestLogger) ObserveRequest(req request.Request) {}

func (r *RequestLogger) ObserveResponse(req request.Request, a request.Attempt) {
	err := ""
	statusCode := ""
	if a.GetError() != nil {
		err = " err=" + a.GetError().Error()
	}

	if a.GetResponse() != nil {
		statusCode = " status=" + strconv.FormatInt(int64(a.GetResponse().StatusCode), 10)
	}

	log.Printf("id=%d method=%s clientIp=%s url=%s backend=%s%s duration=%s%s",
		req.GetId(),
		req.GetHttpRequest().Method,
		req.GetHttpRequest().RemoteAddr,
		req.GetHttpRequest().Host+req.GetHttpRequest().RequestURI,
		a.GetEndpoint(),
		statusCode, a.GetDuration(), err)
}

type SSLRedirect struct{}

func (s *SSLRedirect) ProcessRequest(r request.Request) (*http.Response, error) {
	if sslOnly && r.GetHttpRequest().Header.Get("X-Forwarded-Proto") != "https" {

		resp := &http.Response{
			Status:        "301 Moved Permanently",
			StatusCode:    301,
			Proto:         r.GetHttpRequest().Proto,
			ProtoMajor:    r.GetHttpRequest().ProtoMajor,
			ProtoMinor:    r.GetHttpRequest().ProtoMinor,
			Body:          ioutil.NopCloser(bytes.NewBufferString("")),
			ContentLength: 0,
			Request:       r.GetHttpRequest(),
			Header:        http.Header{},
		}
		resp.Header.Set("Location", "https://"+r.GetHttpRequest().Host+r.GetHttpRequest().RequestURI)
		return resp, nil
	}

	return nil, nil
}

func (s *SSLRedirect) ProcessResponse(r request.Request, a request.Attempt) {
}

func NewHTTPRouter() *HTTPRouter {
	return &HTTPRouter{
		balancers: make(map[string]*roundrobin.RoundRobin),
	}
}

func (s *HTTPRouter) AddBackend(name, vhost, url string) error {

	if vhost == "" || url == "" {
		return nil
	}

	var err error
	balancer := s.balancers[vhost]

	if balancer == nil {
		// Create a round robin load balancer with some endpoints
		balancer, err = roundrobin.NewRoundRobin()
		if err != nil {
			return err
		}

		// Create a http location with the load balancer we've just added
		loc, err := httploc.NewLocationWithOptions(name, balancer,
			httploc.Options{
				TrustForwardHeader: true,
			})
		if err != nil {
			return err
		}
		loc.GetObserverChain().Add("logger", &RequestLogger{})
		loc.GetMiddlewareChain().Add("ssl", 0, &SSLRedirect{})

		s.router.SetRouter(vhost, &route.ConstRouter{Location: loc})
		log.Printf("Starting HTTP listener for %s", vhost)
		s.balancers[vhost] = balancer
	}

	// Already registered?
	if balancer.FindEndpointByUrl(url) != nil {
		return nil
	}
	endpoint := endpoint.MustParseUrl(url)
	log.Printf("Adding HTTP endpoint %s to %s", endpoint.GetUrl(), vhost)
	err = balancer.AddEndpoint(endpoint)
	if err != nil {
		return err
	}
	return nil
}

func (s *HTTPRouter) RemoveBackend(vhost, url string) error {

	if vhost == "" || url == "" {
		return nil
	}

	balancer := s.balancers[vhost]
	if balancer == nil {
		return nil
	}

	endpoint := balancer.FindEndpointByUrl(url)
	if endpoint == nil {
		return nil
	}
	log.Printf("Removing HTTP endpoint %s from %s ", endpoint.GetUrl(), vhost)
	balancer.RemoveEndpoint(endpoint)

	endpoints := balancer.GetEndpoints()
	if len(endpoints) == 0 {
		s.RemoveRouter(vhost)
	}
	return nil
}

// Remove all backends for vhost that are not listed in addrs
func (s *HTTPRouter) RemoveBackends(vhost string, addrs []string) {

	if vhost == "" {
		return
	}

	// Remove backends that are no longer registered

	balancer := s.balancers[vhost]
	if balancer == nil {
		return
	}

	endpoints := balancer.GetEndpoints()
	for _, endpoint := range endpoints {
		if !utils.StringInSlice(endpoint.GetUrl().String(), addrs) {
			s.RemoveBackend(vhost, endpoint.GetUrl().String())
		}
	}
}

// Removes a virtual host router
func (s *HTTPRouter) RemoveRouter(vhost string) {
	if vhost == "" {
		return
	}

	log.Printf("Removing balancer for %s", vhost)
	delete(s.balancers, vhost)
	s.router.RemoveRouter(vhost)
}

func (s *HTTPRouter) adminHandler(w http.ResponseWriter, r *http.Request) {
	if len(s.balancers) == 0 {
		w.WriteHeader(503)
		return
	}

	keys := make([]string, 0, len(s.balancers))
	for key := range s.balancers {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	for _, k := range keys {
		balancer := s.balancers[k]
		endpoints := balancer.GetEndpoints()
		fmt.Fprintf(w, "%s\n", k)
		for _, endpoint := range endpoints {
			fmt.Fprintf(w, "  %s\t%d\t%d\t%0.2f\n", endpoint.GetUrl(), endpoint.GetOriginalWeight(), endpoint.GetEffectiveWeight(), endpoint.GetMeter().GetRate())
		}
	}
}

func (s *HTTPRouter) statusHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var err error
		host := r.Host
		if strings.Contains(host, ":") {
			host, _, err = net.SplitHostPort(r.Host)
			if err != nil {
				log.Warnf("%s", err)
				h.ServeHTTP(w, r)
				return
			}
		}

		if _, exists := s.balancers[host]; !exists {
			s.adminHandler(w, r)
			return
		}
		h.ServeHTTP(w, r)
	})
}

func (s *HTTPRouter) Start() {

	if debug {
		// init the vulcan logging
		gotoolslog.Init([]*gotoolslog.LogConfig{
			&gotoolslog.LogConfig{Name: "console"},
		})
	}

	log.Printf("HTTP server listening at %s", listenAddr)

	s.router = hostroute.NewHostRouter()

	proxy, err := vulcan.NewProxy(s.router)
	if err != nil {
		log.Fatalf("ERROR: %s", err)
	}

	// Proxy acts as http handler:
	server := &http.Server{
		Addr:           listenAddr,
		Handler:        s.statusHandler(proxy),
		ReadTimeout:    60 * time.Second,
		WriteTimeout:   60 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	log.Errorf("%s", server.ListenAndServe())
}

func startHTTPServer() {
	defer wg.Done()
	httpRouter = NewHTTPRouter()
	httpRouter.Start()
}
<commit_msg>Add X-Request-Id header<commit_after>package main

import (
	"bytes"
	"crypto/rand"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/litl/galaxy/utils"

	"github.com/litl/galaxy/log"
	gotoolslog "github.com/mailgun/gotools-log"
	"github.com/mailgun/vulcan"
	"github.com/mailgun/vulcan/endpoint"
	"github.com/mailgun/vulcan/loadbalance/roundrobin"
	"github.com/mailgun/vulcan/location/httploc"
	"github.com/mailgun/vulcan/request"
	"github.com/mailgun/vulcan/route"
	"github.com/mailgun/vulcan/route/hostroute"
)

var (
	httpRouter *HTTPRouter
)

type RequestLogger struct{}

type HTTPRouter struct {
	router    *hostroute.HostRouter
	balancers map[string]*roundrobin.RoundRobin
}

func (r *RequestLogger) ObserveRequest(req request.Request) {}

func (r *RequestLogger) ObserveResponse(req request.Request, a request.Attempt) {
	err := ""
	statusCode := ""
	if a.GetError() != nil {
		err = " err=" + a.GetError().Error()
	}

	if a.GetResponse() != nil {
		statusCode = " status=" + strconv.FormatInt(int64(a.GetResponse().StatusCode), 10)
	}

	log.Printf("cnt=%d id=%s method=%s clientIp=%s url=%s backend=%s%s duration=%s%s",
		req.GetId(),
		req.GetHttpRequest().Header.Get("X-Request-Id"),
		req.GetHttpRequest().Method,
		req.GetHttpRequest().RemoteAddr,
		req.GetHttpRequest().Host+req.GetHttpRequest().RequestURI,
		a.GetEndpoint(),
		statusCode, a.GetDuration(), err)
}

type SSLRedirect struct{}

func genId() string {
	b := make([]byte, 8)
	rand.Read(b)
	return fmt.Sprintf("%x", b)
}

func (s *SSLRedirect) ProcessRequest(r request.Request) (*http.Response, error) {
	r.GetHttpRequest().Header.Set("X-Request-Id", genId())

	if sslOnly && r.GetHttpRequest().Header.Get("X-Forwarded-Proto") != "https" {

		resp := &http.Response{
			Status:        "301 Moved Permanently",
			StatusCode:    301,
			Proto:         r.GetHttpRequest().Proto,
			ProtoMajor:    r.GetHttpRequest().ProtoMajor,
			ProtoMinor:    r.GetHttpRequest().ProtoMinor,
			Body:          ioutil.NopCloser(bytes.NewBufferString("")),
			ContentLength: 0,
			Request:       r.GetHttpRequest(),
			Header:        http.Header{},
		}
		resp.Header.Set("Location", "https://"+r.GetHttpRequest().Host+r.GetHttpRequest().RequestURI)
		return resp, nil
	}

	return nil, nil
}

func (s *SSLRedirect) ProcessResponse(r request.Request, a request.Attempt) {
}

func NewHTTPRouter() *HTTPRouter {
	return &HTTPRouter{
		balancers: make(map[string]*roundrobin.RoundRobin),
	}
}

func (s *HTTPRouter) AddBackend(name, vhost, url string) error {

	if vhost == "" || url == "" {
		return nil
	}

	var err error
	balancer := s.balancers[vhost]

	if balancer == nil {
		// Create a round robin load balancer with some endpoints
		balancer, err = roundrobin.NewRoundRobin()
		if err != nil {
			return err
		}

		// Create a http location with the load balancer we've just added
		loc, err := httploc.NewLocationWithOptions(name, balancer,
			httploc.Options{
				TrustForwardHeader: true,
			})
		if err != nil {
			return err
		}
		loc.GetObserverChain().Add("logger", &RequestLogger{})
		loc.GetMiddlewareChain().Add("ssl", 0, &SSLRedirect{})

		s.router.SetRouter(vhost, &route.ConstRouter{Location: loc})
		log.Printf("Starting HTTP listener for %s", vhost)
		s.balancers[vhost] = balancer
	}

	// Already registered?
	if balancer.FindEndpointByUrl(url) != nil {
		return nil
	}
	endpoint := endpoint.MustParseUrl(url)
	log.Printf("Adding HTTP endpoint %s to %s", endpoint.GetUrl(), vhost)
	err = balancer.AddEndpoint(endpoint)
	if err != nil {
		return err
	}
	return nil
}

func (s *HTTPRouter) RemoveBackend(vhost, url string) error {

	if vhost == "" || url == "" {
		return nil
	}

	balancer := s.balancers[vhost]
	if balancer == nil {
		return nil
	}

	endpoint := balancer.FindEndpointByUrl(url)
	if endpoint == nil {
		return nil
	}
	log.Printf("Removing HTTP endpoint %s from %s ", endpoint.GetUrl(), vhost)
	balancer.RemoveEndpoint(endpoint)

	endpoints := balancer.GetEndpoints()
	if len(endpoints) == 0 {
		s.RemoveRouter(vhost)
	}
	return nil
}

// Remove all backends for vhost that are not listed in addrs
func (s *HTTPRouter) RemoveBackends(vhost string, addrs []string) {

	if vhost == "" {
		return
	}

	// Remove backends that are no longer registered

	balancer := s.balancers[vhost]
	if balancer == nil {
		return
	}

	endpoints := balancer.GetEndpoints()
	for _, endpoint := range endpoints {
		if !utils.StringInSlice(endpoint.GetUrl().String(), addrs) {
			s.RemoveBackend(vhost, endpoint.GetUrl().String())
		}
	}
}

// Removes a virtual host router
func (s *HTTPRouter) RemoveRouter(vhost string) {
	if vhost == "" {
		return
	}

	log.Printf("Removing balancer for %s", vhost)
	delete(s.balancers, vhost)
	s.router.RemoveRouter(vhost)
}

func (s *HTTPRouter) adminHandler(w http.ResponseWriter, r *http.Request) {
	if len(s.balancers) == 0 {
		w.WriteHeader(503)
		return
	}

	keys := make([]string, 0, len(s.balancers))
	for key := range s.balancers {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	for _, k := range keys {
		balancer := s.balancers[k]
		endpoints := balancer.GetEndpoints()
		fmt.Fprintf(w, "%s\n", k)
		for _, endpoint := range endpoints {
			fmt.Fprintf(w, "  %s\t%d\t%d\t%0.2f\n", endpoint.GetUrl(), endpoint.GetOriginalWeight(), endpoint.GetEffectiveWeight(), endpoint.GetMeter().GetRate())
		}
	}
}

func (s *HTTPRouter) statusHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var err error
		host := r.Host
		if strings.Contains(host, ":") {
			host, _, err = net.SplitHostPort(r.Host)
			if err != nil {
				log.Warnf("%s", err)
				h.ServeHTTP(w, r)
				return
			}
		}

		if _, exists := s.balancers[host]; !exists {
			s.adminHandler(w, r)
			return
		}
		h.ServeHTTP(w, r)
	})
}

func (s *HTTPRouter) Start() {

	if debug {
		// init the vulcan logging
		gotoolslog.Init([]*gotoolslog.LogConfig{
			&gotoolslog.LogConfig{Name: "console"},
		})
	}

	log.Printf("HTTP server listening at %s", listenAddr)

	s.router = hostroute.NewHostRouter()

	proxy, err := vulcan.NewProxy(s.router)
	if err != nil {
		log.Fatalf("ERROR: %s", err)
	}

	// Proxy acts as http handler:
	server := &http.Server{
		Addr:           listenAddr,
		Handler:        s.statusHandler(proxy),
		ReadTimeout:    60 * time.Second,
		WriteTimeout:   60 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	log.Errorf("%s", server.ListenAndServe())
}

func startHTTPServer() {
	defer wg.Done()
	httpRouter = NewHTTPRouter()
	httpRouter.Start()
}
<|endoftext|>


<commit_before>package factual

import (
  "io/ioutil"
  "net/http"
  "net/url"

  "github.com/asaskevich/govalidator"
)

func (c Client) Get(path string, params url.Values) ([]byte, error) {

  fullUrl := c.BaseUri + path
  if !govalidator.IsURL(fullUrl) {
    return nil, ErrInvalidUrl(fullUrl)
  }

  resp, err := c.Oauth.Get(http.DefaultClient, nil, fullUrl, params)
  if err != nil {
    return nil, err
  }
  defer resp.Body.Close()

  body, err := ioutil.ReadAll(resp.Body)
  if err != nil {
    return nil, ErrHttpBody(fullUrl)
  }

  if resp.StatusCode != 200 {
    return nil, ErrHttpResponse(fullUrl, resp.StatusCode, body)
  }

  return body, nil
}

func (c Client) Post(url string) (string, error) {
  return "POST", nil
}
<commit_msg>Commented post out<commit_after>package factual

import (
  "io/ioutil"
  "net/http"
  "net/url"

  "github.com/asaskevich/govalidator"
)

func (c Client) Get(path string, params url.Values) ([]byte, error) {

  fullUrl := c.BaseUri + path
  if !govalidator.IsURL(fullUrl) {
    return nil, ErrInvalidUrl(fullUrl)
  }

  resp, err := c.Oauth.Get(http.DefaultClient, nil, fullUrl, params)
  if err != nil {
    return nil, err
  }
  defer resp.Body.Close()

  body, err := ioutil.ReadAll(resp.Body)
  if err != nil {
    return nil, ErrHttpBody(fullUrl)
  }

  if resp.StatusCode != 200 {
    return nil, ErrHttpResponse(fullUrl, resp.StatusCode, body)
  }

  return body, nil
}

// func (c Client) Post(url string) (string, error) {
//   return "POST", nil
// }
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"math"

	"github.com/jzacsh/netwtcpip-cmp405/parseip4"
)

var partTwoHosts = []parseip4.Addr{
	{IP: parseip4.NewAddr(9, 201, 195, 84), Mask: parseip4.NewAddr(255, 255, 240, 0)},
	{IP: parseip4.NewAddr(128, 10, 189, 215), Mask: parseip4.NewAddr(255, 255, 248, 0)},
	{IP: parseip4.NewAddr(135, 21, 243, 82), Mask: parseip4.NewAddr(255, 255, 224, 0)},
	{IP: parseip4.NewAddr(75, 149, 205, 61), Mask: parseip4.NewAddr(255, 255, 192, 0)},
	{IP: parseip4.NewAddr(7, 105, 198, 111), Mask: parseip4.NewAddr(255, 255, 252, 0)},

	// TODO(zacsh) remove this sample from the last slide
	{IP: parseip4.NewAddr(128, 10, 211, 78), Mask: parseip4.NewAddr(255, 255, 240, 0)},
}

type subnetRequisites struct {
	MaxSubnets  uint
	SubnetIndex parseip4.Octets
	HostIndex   parseip4.Octets
}

type OptimalSubnet struct {
	MinSubnetBits     uint
	MaxHostsPerSubnet parseip4.Octets
	SubnetMask        parseip4.OctsList
}

var partOneGivens = []subnetRequisites{
	{55, 51, 121},
	{55, 42, 867},
	{121, 115, 246},
	{121, 97, 443},
	{26, 19, 237},
	{26, 25, 1397},
	{261, 227, 86},
	{261, 259, 49},
	{529, 519, 33},
	{529, 510, 59},
}

func (s *subnetRequisites) String() string {
	return fmt.Sprintf(
		"max subnets: %d, subnet index: %d, host index: %d",
		s.MaxSubnets, s.SubnetIndex, s.HostIndex)
}

func maxIntWithBits(nbits uint) uint32 {
	// 1 because 2^N bits only gets 2^N-1 if all 1s. +1 more because all 1s is
	// reserved for broadcast.
	const gap float64 = 2

	maxInt := math.Pow(2, float64(nbits))
	if maxInt < gap {
		// we want to avoid underflows, so stick to the point of the API and return
		// effectively zero
		return 0
	}

	return uint32(maxInt - gap)
}

func (s *subnetRequisites) FindSolution() OptimalSubnet {
	opt := OptimalSubnet{}

	// Brute force solve for Ceil(log2(s.MaxSubnets))
	for {
		if maxIntWithBits(opt.MinSubnetBits) >= uint32(s.MaxSubnets) {
			break
		}
		opt.MinSubnetBits++
	}

	opt.MaxHostsPerSubnet = parseip4.Octets(maxIntWithBits(32 - opt.MinSubnetBits))

	// TODO: opt.SubnetMask

	return opt
}

func main() {
	fmt.Printf("part 1: analyzing %d hosts ...\n", len(partOneGivens))
	for _, req := range partOneGivens {
		sol := req.FindSolution()
		fmt.Printf(
			"  given: %s\n\tmin # of subnet bits: %d\n\tmax # hosts per subnet: %d\n\tsubnetmask: %d\n",
			req.String(),
			sol.MinSubnetBits,
			sol.MaxHostsPerSubnet,
			sol.SubnetMask)
	}

	fmt.Printf("\n\npart 2: analyzing %d hosts ...\n", len(partTwoHosts))
	for _, addr := range partTwoHosts {
		classMask, _, klass := parseip4.Classful(addr.IP)

		fmt.Printf(
			"  network: %v (class %s masked)\n\t%v\n\tnetwork id:\t%d\n\t subnet id:\t%d\n\t   host id:\t%d\n",
			(addr.IP.Pack() & classMask.Pack()).List(), klass,
			addr.String(),
			addr.NetworkIndex(),
			addr.SubnetIndex(),
			addr.HostIndex())
	}
}
<commit_msg>hw4/part1: nit(refactor) need *both* addr & mask<commit_after>package main

import (
	"fmt"
	"math"

	"github.com/jzacsh/netwtcpip-cmp405/parseip4"
)

var partTwoHosts = []parseip4.Addr{
	{IP: parseip4.NewAddr(9, 201, 195, 84), Mask: parseip4.NewAddr(255, 255, 240, 0)},
	{IP: parseip4.NewAddr(128, 10, 189, 215), Mask: parseip4.NewAddr(255, 255, 248, 0)},
	{IP: parseip4.NewAddr(135, 21, 243, 82), Mask: parseip4.NewAddr(255, 255, 224, 0)},
	{IP: parseip4.NewAddr(75, 149, 205, 61), Mask: parseip4.NewAddr(255, 255, 192, 0)},
	{IP: parseip4.NewAddr(7, 105, 198, 111), Mask: parseip4.NewAddr(255, 255, 252, 0)},

	// TODO(zacsh) remove this sample from the last slide
	{IP: parseip4.NewAddr(128, 10, 211, 78), Mask: parseip4.NewAddr(255, 255, 240, 0)},
}

type subnetRequisites struct {
	MaxSubnets  uint
	SubnetIndex parseip4.Octets
	HostIndex   parseip4.Octets
}

type OptimalSubnet struct {
	MinSubnetBits     uint
	MaxHostsPerSubnet parseip4.Octets
	Address           parseip4.Addr
}

var partOneGivens = []subnetRequisites{
	{55, 51, 121},
	{55, 42, 867},
	{121, 115, 246},
	{121, 97, 443},
	{26, 19, 237},
	{26, 25, 1397},
	{261, 227, 86},
	{261, 259, 49},
	{529, 519, 33},
	{529, 510, 59},
}

func (s *subnetRequisites) String() string {
	return fmt.Sprintf(
		"max subnets: %d, subnet index: %d, host index: %d",
		s.MaxSubnets, s.SubnetIndex, s.HostIndex)
}

func maxIntWithBits(nbits uint) uint32 {
	// 1 because 2^N bits only gets 2^N-1 if all 1s. +1 more because all 1s is
	// reserved for broadcast.
	const gap float64 = 2

	maxInt := math.Pow(2, float64(nbits))
	if maxInt < gap {
		// we want to avoid underflows, so stick to the point of the API and return
		// effectively zero
		return 0
	}

	return uint32(maxInt - gap)
}

func (s *subnetRequisites) FindSolution() OptimalSubnet {
	opt := OptimalSubnet{}

	// Brute force solve for Ceil(log2(s.MaxSubnets))
	for {
		if maxIntWithBits(opt.MinSubnetBits) >= uint32(s.MaxSubnets) {
			break
		}
		opt.MinSubnetBits++
	}

	opt.MaxHostsPerSubnet = parseip4.Octets(maxIntWithBits(32 - opt.MinSubnetBits))

	// TODO opt.Address.Mask
	opt.Address.Mask = parseip4.NewAddr(0, 0, 0, 0)

	// TODO opt.Address.IP
	opt.Address.IP = parseip4.NewAddr(0, 0, 0, 0)

	return opt
}

func main() {
	fmt.Printf("part 1: analyzing %d hosts ...\n", len(partOneGivens))
	for _, req := range partOneGivens {
		sol := req.FindSolution()
		fmt.Printf(
			"  given: %s\n\tmin # of subnet bits: %d\n\tmax # hosts per subnet: %d\n\taddress: %s\n",
			req.String(),
			sol.MinSubnetBits,
			sol.MaxHostsPerSubnet,
			sol.Address.String())
	}

	fmt.Printf("\n\npart 2: analyzing %d hosts ...\n", len(partTwoHosts))
	for _, addr := range partTwoHosts {
		classMask, _, klass := parseip4.Classful(addr.IP)

		fmt.Printf(
			"  network: %v (class %s masked)\n\t%v\n\tnetwork id:\t%d\n\t subnet id:\t%d\n\t   host id:\t%d\n",
			(addr.IP.Pack() & classMask.Pack()).List(), klass,
			addr.String(),
			addr.NetworkIndex(),
			addr.SubnetIndex(),
			addr.HostIndex())
	}
}
<|endoftext|>


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

package scanner

import (
	"context"
	"sync"
	"time"

	"github.com/golang/glog"
	ct "github.com/google/certificate-transparency-go"
	"github.com/google/certificate-transparency-go/client"
	"github.com/google/trillian/client/backoff"
)

// FetcherOptions holds configuration options for the Fetcher.
type FetcherOptions struct {
	// Number of entries to request in one batch from the Log.
	BatchSize int

	// Number of concurrent fetcher workers to run.
	ParallelFetch int

	// [StartIndex, EndIndex) is a log entry range to fetch. If EndIndex == 0,
	// then it gets reassigned to sth.TreeSize.
	StartIndex int64
	EndIndex   int64

	// Continuous determines whether Fetcher should run indefinitely after
	// reaching EndIndex.
	Continuous bool
}

// DefaultFetcherOptions returns new FetcherOptions with sensible defaults.
func DefaultFetcherOptions() *FetcherOptions {
	return &FetcherOptions{
		BatchSize:     1000,
		ParallelFetch: 1,
		StartIndex:    0,
		EndIndex:      0,
		Continuous:    false,
	}
}

// Fetcher is a tool that fetches entries from a CT Log.
type Fetcher struct {
	// Base URI of the CT log, for diagnostics.
	uri string
	// Client used to talk to the CT log instance.
	client *client.LogClient
	// Configuration options for this Fetcher instance.
	opts *FetcherOptions

	// Current STH of the Log this Fetcher sends queries to.
	sth *ct.SignedTreeHead
	// The STH retrieval backoff state. Used only in Continuous fetch mode.
	sthBackoff *backoff.Backoff

	// Stops range generator, which causes the Fetcher to terminate gracefully.
	mu     sync.Mutex
	cancel context.CancelFunc
}

// EntryBatch represents a contiguous range of entries of the Log.
type EntryBatch struct {
	Start   int64          // LeafIndex of the first entry in the range.
	Entries []ct.LeafEntry // Entries of the range.
}

// fetchRange represents a range of certs to fetch from a CT log.
type fetchRange struct {
	start int64 // inclusive
	end   int64 // inclusive
}

// NewFetcher creates a Fetcher instance using client to talk to the log,
// taking configuration options from opts.
func NewFetcher(client *client.LogClient, opts *FetcherOptions) *Fetcher {
	cancel := func() {} // Protect against calling Stop before Run.
	return &Fetcher{
		uri:    client.BaseURI(),
		client: client,
		opts:   opts,
		cancel: cancel,
	}
}

// Prepare caches the latest Log's STH if not present and returns it. It also
// adjusts the entry range to fit the size of the tree.
func (f *Fetcher) Prepare(ctx context.Context) (*ct.SignedTreeHead, error) {
	if f.sth != nil {
		return f.sth, nil
	}

	sth, err := f.client.GetSTH(ctx)
	if err != nil {
		glog.Errorf("%s: GetSTH() failed: %v", f.uri, err)
		return nil, err
	}
	glog.Infof("%s: Got STH with %d certs", f.uri, sth.TreeSize)

	if size := int64(sth.TreeSize); f.opts.EndIndex == 0 || f.opts.EndIndex > size {
		glog.Infof("%s: Reset EndIndex from %d to %d", f.uri, f.opts.EndIndex, size)
		f.opts.EndIndex = size
	}
	f.sth = sth
	return sth, nil
}

// Run performs fetching of the Log. Blocks until scanning is complete, the
// passed in context is canceled, or Stop is called (and pending work is
// finished). For each successfully fetched batch, runs the fn callback.
func (f *Fetcher) Run(ctx context.Context, fn func(EntryBatch)) error {
	glog.V(1).Infof("%s: Starting up Fetcher...", f.uri)
	if _, err := f.Prepare(ctx); err != nil {
		return err
	}

	cctx, cancel := context.WithCancel(ctx)
	defer cancel()

	f.mu.Lock()
	f.cancel = cancel
	f.mu.Unlock()

	// Use a separately-cancelable context for the range generator, so we can
	// close it down (in Stop) but still let the fetchers below run to
	// completion.
	ranges := f.genRanges(cctx)

	// Run fetcher workers.
	var wg sync.WaitGroup
	for w, cnt := 0, f.opts.ParallelFetch; w < cnt; w++ {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			glog.V(1).Infof("%s: Fetcher worker %d starting...", f.uri, idx)
			f.runWorker(ctx, ranges, fn)
			glog.V(1).Infof("%s: Fetcher worker %d finished", f.uri, idx)
		}(w)
	}
	wg.Wait()

	glog.V(1).Infof("%s: Fetcher terminated", f.uri)
	return nil
}

// Stop causes the Fetcher to terminate gracefully. After this call Run will
// try to finish all the started fetches, and then return. Does nothing if
// there was no preceding Run invocation.
func (f *Fetcher) Stop() {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.cancel()
}

// genRanges returns a channel of ranges to fetch, and starts a goroutine that
// sends things down this channel. The goroutine terminates when all ranges
// have been generated, or if context is cancelled.
func (f *Fetcher) genRanges(ctx context.Context) <-chan fetchRange {
	batch := int64(f.opts.BatchSize)
	ranges := make(chan fetchRange)

	go func() {
		glog.V(1).Infof("%s: Range generator starting", f.uri)
		defer glog.V(1).Infof("%s: Range generator finished", f.uri)
		defer close(ranges)
		start, end := f.opts.StartIndex, f.opts.EndIndex

		for start < end || f.opts.Continuous {
			// In continuous mode wait for bigger STH every time we reach the end,
			// including, possibly, the very first iteration.
			if start == end { // Implies f.opts.Continuous == true.
				if err := f.updateSTH(ctx); err != nil {
					glog.Warningf("%s: Failed to obtain bigger STH: %v", f.uri, err)
					return
				}
				end = f.opts.EndIndex
			}

			batchEnd := start + min(end-start, batch)
			next := fetchRange{start, batchEnd - 1}
			select {
			case <-ctx.Done():
				glog.Warningf("%s: Cancelling genRanges: %v", f.uri, ctx.Err())
				return
			case ranges <- next:
			}
			start = batchEnd
		}
	}()

	return ranges
}

// updateSTH waits until a bigger STH is discovered, and updates the Fetcher
// accordingly. It is optimized for both bulk-load (new STH is way bigger then
// the last one) and keep-up (STH grows slowly) modes of operation. Waits for
// some time until the STH grows enough to request a full batch, but falls back
// to *any* STH bigger than the old one if it takes too long.
// Returns error only if the context is cancelled.
func (f *Fetcher) updateSTH(ctx context.Context) error {
	// TODO(pavelkalinnikov): Make these parameters tunable.
	const quickDur = 45 * time.Second
	if f.sthBackoff == nil {
		f.sthBackoff = &backoff.Backoff{
			Min:    1 * time.Second,
			Max:    30 * time.Second,
			Factor: 2,
			Jitter: true,
		}
	}

	lastSize := uint64(f.opts.EndIndex)
	targetSize := lastSize + uint64(f.opts.BatchSize)
	quickDeadline := time.Now().Add(quickDur)

	return f.sthBackoff.Retry(ctx, func() error {
		sth, err := f.client.GetSTH(ctx)
		if err != nil {
			return err
		}
		glog.V(2).Infof("%s: Got STH with %d certs", f.uri, sth.TreeSize)

		quick := time.Now().Before(quickDeadline)
		if sth.TreeSize <= lastSize || quick && sth.TreeSize < targetSize {
			return backoff.RetriableErrorf("wait for bigger STH than %d (last=%d, target=%d)", sth.TreeSize, lastSize, targetSize)
		}

		if quick {
			f.sthBackoff.Reset() // Growth is presumably fast, set next pause to Min.
		}
		f.sth = sth
		f.opts.EndIndex = int64(sth.TreeSize)
		return nil
	})
}

// runWorker is a worker function for handling fetcher ranges.
// Accepts cert ranges to fetch over the ranges channel, and if the fetch is
// successful sends the corresponding EntryBatch through the fn callback. Will
// retry failed attempts to retrieve ranges until the context is cancelled.
func (f *Fetcher) runWorker(ctx context.Context, ranges <-chan fetchRange, fn func(EntryBatch)) {
	for r := range ranges {
		// Logs MAY return fewer than the number of leaves requested. Only complete
		// if we actually got all the leaves we were expecting.
		for r.start <= r.end {
			// Fetcher.Run() can be cancelled while we are looping over this job.
			if err := ctx.Err(); err != nil {
				glog.Warningf("%s: Worker context closed: %v", f.uri, err)
				return
			}
			resp, err := f.client.GetRawEntries(ctx, r.start, r.end)
			if err != nil {
				glog.Errorf("%s: GetRawEntries() failed: %v", f.uri, err)
				// TODO(pavelkalinnikov): Introduce backoff policy and pause here.
				continue
			}
			fn(EntryBatch{Start: r.start, Entries: resp.Entries})
			r.start += int64(len(resp.Entries))
		}
	}
}

func min(a, b int64) int64 {
	if a < b {
		return a
	}
	return b
}
<commit_msg>Added backoff to fetcher worker to partially address #746 (#747)<commit_after>// Copyright 2018 Google LLC. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package scanner

import (
	"context"
	"sync"
	"time"

	"github.com/golang/glog"
	ct "github.com/google/certificate-transparency-go"
	"github.com/google/certificate-transparency-go/client"
	"github.com/google/trillian/client/backoff"
)

// FetcherOptions holds configuration options for the Fetcher.
type FetcherOptions struct {
	// Number of entries to request in one batch from the Log.
	BatchSize int

	// Number of concurrent fetcher workers to run.
	ParallelFetch int

	// [StartIndex, EndIndex) is a log entry range to fetch. If EndIndex == 0,
	// then it gets reassigned to sth.TreeSize.
	StartIndex int64
	EndIndex   int64

	// Continuous determines whether Fetcher should run indefinitely after
	// reaching EndIndex.
	Continuous bool
}

// DefaultFetcherOptions returns new FetcherOptions with sensible defaults.
func DefaultFetcherOptions() *FetcherOptions {
	return &FetcherOptions{
		BatchSize:     1000,
		ParallelFetch: 1,
		StartIndex:    0,
		EndIndex:      0,
		Continuous:    false,
	}
}

// Fetcher is a tool that fetches entries from a CT Log.
type Fetcher struct {
	// Base URI of the CT log, for diagnostics.
	uri string
	// Client used to talk to the CT log instance.
	client *client.LogClient
	// Configuration options for this Fetcher instance.
	opts *FetcherOptions

	// Current STH of the Log this Fetcher sends queries to.
	sth *ct.SignedTreeHead
	// The STH retrieval backoff state. Used only in Continuous fetch mode.
	sthBackoff *backoff.Backoff

	// Stops range generator, which causes the Fetcher to terminate gracefully.
	mu     sync.Mutex
	cancel context.CancelFunc
}

// EntryBatch represents a contiguous range of entries of the Log.
type EntryBatch struct {
	Start   int64          // LeafIndex of the first entry in the range.
	Entries []ct.LeafEntry // Entries of the range.
}

// fetchRange represents a range of certs to fetch from a CT log.
type fetchRange struct {
	start int64 // inclusive
	end   int64 // inclusive
}

// NewFetcher creates a Fetcher instance using client to talk to the log,
// taking configuration options from opts.
func NewFetcher(client *client.LogClient, opts *FetcherOptions) *Fetcher {
	cancel := func() {} // Protect against calling Stop before Run.
	return &Fetcher{
		uri:    client.BaseURI(),
		client: client,
		opts:   opts,
		cancel: cancel,
	}
}

// Prepare caches the latest Log's STH if not present and returns it. It also
// adjusts the entry range to fit the size of the tree.
func (f *Fetcher) Prepare(ctx context.Context) (*ct.SignedTreeHead, error) {
	if f.sth != nil {
		return f.sth, nil
	}

	sth, err := f.client.GetSTH(ctx)
	if err != nil {
		glog.Errorf("%s: GetSTH() failed: %v", f.uri, err)
		return nil, err
	}
	glog.Infof("%s: Got STH with %d certs", f.uri, sth.TreeSize)

	if size := int64(sth.TreeSize); f.opts.EndIndex == 0 || f.opts.EndIndex > size {
		glog.Infof("%s: Reset EndIndex from %d to %d", f.uri, f.opts.EndIndex, size)
		f.opts.EndIndex = size
	}
	f.sth = sth
	return sth, nil
}

// Run performs fetching of the Log. Blocks until scanning is complete, the
// passed in context is canceled, or Stop is called (and pending work is
// finished). For each successfully fetched batch, runs the fn callback.
func (f *Fetcher) Run(ctx context.Context, fn func(EntryBatch)) error {
	glog.V(1).Infof("%s: Starting up Fetcher...", f.uri)
	if _, err := f.Prepare(ctx); err != nil {
		return err
	}

	cctx, cancel := context.WithCancel(ctx)
	defer cancel()

	f.mu.Lock()
	f.cancel = cancel
	f.mu.Unlock()

	// Use a separately-cancelable context for the range generator, so we can
	// close it down (in Stop) but still let the fetchers below run to
	// completion.
	ranges := f.genRanges(cctx)

	// Run fetcher workers.
	var wg sync.WaitGroup
	for w, cnt := 0, f.opts.ParallelFetch; w < cnt; w++ {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			glog.V(1).Infof("%s: Fetcher worker %d starting...", f.uri, idx)
			f.runWorker(ctx, ranges, fn)
			glog.V(1).Infof("%s: Fetcher worker %d finished", f.uri, idx)
		}(w)
	}
	wg.Wait()

	glog.V(1).Infof("%s: Fetcher terminated", f.uri)
	return nil
}

// Stop causes the Fetcher to terminate gracefully. After this call Run will
// try to finish all the started fetches, and then return. Does nothing if
// there was no preceding Run invocation.
func (f *Fetcher) Stop() {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.cancel()
}

// genRanges returns a channel of ranges to fetch, and starts a goroutine that
// sends things down this channel. The goroutine terminates when all ranges
// have been generated, or if context is cancelled.
func (f *Fetcher) genRanges(ctx context.Context) <-chan fetchRange {
	batch := int64(f.opts.BatchSize)
	ranges := make(chan fetchRange)

	go func() {
		glog.V(1).Infof("%s: Range generator starting", f.uri)
		defer glog.V(1).Infof("%s: Range generator finished", f.uri)
		defer close(ranges)
		start, end := f.opts.StartIndex, f.opts.EndIndex

		for start < end || f.opts.Continuous {
			// In continuous mode wait for bigger STH every time we reach the end,
			// including, possibly, the very first iteration.
			if start == end { // Implies f.opts.Continuous == true.
				if err := f.updateSTH(ctx); err != nil {
					glog.Warningf("%s: Failed to obtain bigger STH: %v", f.uri, err)
					return
				}
				end = f.opts.EndIndex
			}

			batchEnd := start + min(end-start, batch)
			next := fetchRange{start, batchEnd - 1}
			select {
			case <-ctx.Done():
				glog.Warningf("%s: Cancelling genRanges: %v", f.uri, ctx.Err())
				return
			case ranges <- next:
			}
			start = batchEnd
		}
	}()

	return ranges
}

// updateSTH waits until a bigger STH is discovered, and updates the Fetcher
// accordingly. It is optimized for both bulk-load (new STH is way bigger then
// the last one) and keep-up (STH grows slowly) modes of operation. Waits for
// some time until the STH grows enough to request a full batch, but falls back
// to *any* STH bigger than the old one if it takes too long.
// Returns error only if the context is cancelled.
func (f *Fetcher) updateSTH(ctx context.Context) error {
	// TODO(pavelkalinnikov): Make these parameters tunable.
	const quickDur = 45 * time.Second
	if f.sthBackoff == nil {
		f.sthBackoff = &backoff.Backoff{
			Min:    1 * time.Second,
			Max:    30 * time.Second,
			Factor: 2,
			Jitter: true,
		}
	}

	lastSize := uint64(f.opts.EndIndex)
	targetSize := lastSize + uint64(f.opts.BatchSize)
	quickDeadline := time.Now().Add(quickDur)

	return f.sthBackoff.Retry(ctx, func() error {
		sth, err := f.client.GetSTH(ctx)
		if err != nil {
			return err
		}
		glog.V(2).Infof("%s: Got STH with %d certs", f.uri, sth.TreeSize)

		quick := time.Now().Before(quickDeadline)
		if sth.TreeSize <= lastSize || quick && sth.TreeSize < targetSize {
			return backoff.RetriableErrorf("wait for bigger STH than %d (last=%d, target=%d)", sth.TreeSize, lastSize, targetSize)
		}

		if quick {
			f.sthBackoff.Reset() // Growth is presumably fast, set next pause to Min.
		}
		f.sth = sth
		f.opts.EndIndex = int64(sth.TreeSize)
		return nil
	})
}

// runWorker is a worker function for handling fetcher ranges.
// Accepts cert ranges to fetch over the ranges channel, and if the fetch is
// successful sends the corresponding EntryBatch through the fn callback. Will
// retry failed attempts to retrieve ranges until the context is cancelled.
func (f *Fetcher) runWorker(ctx context.Context, ranges <-chan fetchRange, fn func(EntryBatch)) {
	for r := range ranges {
		// Logs MAY return fewer than the number of leaves requested. Only complete
		// if we actually got all the leaves we were expecting.
		for r.start <= r.end {
			// TODO(pavelkalinnikov): Make these parameters tunable.
			// This backoff will only apply to a single request and be reset for the next one.
			// This precludes reaching some kind of stability in request rate, but means that
			// an intermittent problem won't harm long-term running of the worker.
			bo := &backoff.Backoff{
				Min:    1 * time.Second,
				Max:    30 * time.Second,
				Factor: 2,
				Jitter: true,
			}

			var resp *ct.GetEntriesResponse
			// TODO(pavelkalinnikov): Report errors in a LogClient decorator on failure.
			if err := bo.Retry(ctx, func() error {
				var err error
				resp, err = f.client.GetRawEntries(ctx, r.start, r.end)
				return err
			}); err != nil {
				glog.Errorf("%s: GetRawEntries() failed: %v", f.uri, err)
				// There is no error reporting yet for this worker, so just retry again.
				continue
			}
			fn(EntryBatch{Start: r.start, Entries: resp.Entries})
			r.start += int64(len(resp.Entries))
		}
	}
}

func min(a, b int64) int64 {
	if a < b {
		return a
	}
	return b
}
<|endoftext|>