<commit_before>// This is free and unencumbered software released into the public
// domain.  For more information, see <http://unlicense.org> or the
// accompanying UNLICENSE file.

package mocks

import (
	"go/ast"
	"go/token"
	"io"
	"golang.org/x/tools/go/ast/astutil"
	"go/parser"
	"os"
	"go/format"
	"bytes"
	"strconv"
)

const commentHeader = `// This file was generated by github.com/nelsam/hel.  Do not
// edit this code by hand unless you *really* know what you're
// doing.  Expect any changes made manually to be overwritten
// the next time hel regenerates this file.

`

//go:generate hel --type TypeFinder --output mock_type_finder_test.go

type TypeFinder interface {
	ExportedTypes() (types []*ast.TypeSpec)
	Dependencies(inter *ast.InterfaceType) (dependencies []*ast.TypeSpec)
}

type Mocks []Mock

func (m Mocks) Output(pkg, dir string, chanSize int, dest io.Writer) error {
	if _, err := dest.Write([]byte(commentHeader)); err != nil {
		return err
	}

	fset := token.NewFileSet()

	f := &ast.File{
		Name:  &ast.Ident{Name: pkg},
		Decls: m.decls(chanSize),
	}

	var b bytes.Buffer
	format.Node(&b, fset, f)

	// TODO: Determine why adding imports without creating a new ast file
	// will only allow one import to be printed to the file.
	fset = token.NewFileSet()
	file, err := parser.ParseFile(fset, pkg, &b, 0)
	if err != nil {
		return err
	}

	file, fset, err = addImports(file, fset, dir)
	if err != nil {
		return err
	}

	return format.Node(dest, fset, file)
}

func (m Mocks) PrependLocalPackage(name string) {
	for _, m := range m {
		m.PrependLocalPackage(name)
	}
}

func (m Mocks) SetBlockingReturn(blockingReturn bool) {
	for _, m := range m {
		m.SetBlockingReturn(blockingReturn)
	}
}

func (m Mocks) decls(chanSize int) (decls []ast.Decl) {
	for _, mock := range m {
		decls = append(decls, mock.Ast(chanSize)...)
	}
	return decls
}

func addImports(file *ast.File, fset *token.FileSet, dirPath string) (*ast.File, *token.FileSet, error) {
	imports, err := getImports(dirPath, fset)
	if err != nil {
		return nil, nil, err
	}

	for _, s := range imports {
		unquotedPath, err := strconv.Unquote(s.Path.Value)
		if err != nil {
			return nil, nil, err
		}

		if s.Name != nil {
			astutil.AddNamedImport(fset, file, s.Name.Name, unquotedPath)
			continue
		}

		astutil.AddImport(fset, file, unquotedPath)
	}

	return file, fset, nil
}

func getImports(dirPath string, fset *token.FileSet) ([]*ast.ImportSpec, error) {
	// Grab imports from all files except helheim_test
	pkgs, err := parser.ParseDir(fset, dirPath, func(info os.FileInfo) bool {
		return info.Name() != "helheim_test.go"
	}, parser.ImportsOnly)
	if err != nil {
		return nil, err
	}

	var imports []*ast.ImportSpec
	for _, p := range pkgs {
		files := p.Files
		for _, f := range files {
			imports = append(imports, f.Imports...)
		}
	}
	return imports, nil
}

func Generate(finder TypeFinder) (Mocks, error) {
	base := finder.ExportedTypes()
	var types []*ast.TypeSpec
	for _, typ := range base {
		types = append(types, typ)
		if inter, ok := typ.Type.(*ast.InterfaceType); ok {
			types = append(types, finder.Dependencies(inter)...)
		}
	}
	m := make(Mocks, 0, len(types))
	for _, typ := range types {
		newMock, err := For(typ)
		if err != nil {
			return nil, err
		}
		m = append(m, newMock)
	}
	return m, nil
}
<commit_msg>Changed addImports to only add imports from non test files<commit_after>// This is free and unencumbered software released into the public
// domain.  For more information, see <http://unlicense.org> or the
// accompanying UNLICENSE file.

package mocks

import (
	"go/ast"
	"go/token"
	"io"
	"golang.org/x/tools/go/ast/astutil"
	"go/parser"
	"os"
	"go/format"
	"bytes"
	"strconv"
	"strings"
)

const commentHeader = `// This file was generated by github.com/nelsam/hel.  Do not
// edit this code by hand unless you *really* know what you're
// doing.  Expect any changes made manually to be overwritten
// the next time hel regenerates this file.

`

//go:generate hel --type TypeFinder --output mock_type_finder_test.go

type TypeFinder interface {
	ExportedTypes() (types []*ast.TypeSpec)
	Dependencies(inter *ast.InterfaceType) (dependencies []*ast.TypeSpec)
}

type Mocks []Mock

func (m Mocks) Output(pkg, dir string, chanSize int, dest io.Writer) error {
	if _, err := dest.Write([]byte(commentHeader)); err != nil {
		return err
	}

	fset := token.NewFileSet()

	f := &ast.File{
		Name:  &ast.Ident{Name: pkg},
		Decls: m.decls(chanSize),
	}

	var b bytes.Buffer
	format.Node(&b, fset, f)

	// TODO: Determine why adding imports without creating a new ast file
	// will only allow one import to be printed to the file.
	fset = token.NewFileSet()
	file, err := parser.ParseFile(fset, pkg, &b, 0)
	if err != nil {
		return err
	}

	file, fset, err = addImports(file, fset, dir)
	if err != nil {
		return err
	}

	return format.Node(dest, fset, file)
}

func (m Mocks) PrependLocalPackage(name string) {
	for _, m := range m {
		m.PrependLocalPackage(name)
	}
}

func (m Mocks) SetBlockingReturn(blockingReturn bool) {
	for _, m := range m {
		m.SetBlockingReturn(blockingReturn)
	}
}

func (m Mocks) decls(chanSize int) (decls []ast.Decl) {
	for _, mock := range m {
		decls = append(decls, mock.Ast(chanSize)...)
	}
	return decls
}

func addImports(file *ast.File, fset *token.FileSet, dirPath string) (*ast.File, *token.FileSet, error) {
	imports, err := getImports(dirPath, fset)
	if err != nil {
		return nil, nil, err
	}

	for _, s := range imports {
		unquotedPath, err := strconv.Unquote(s.Path.Value)
		if err != nil {
			return nil, nil, err
		}

		if s.Name != nil {
			astutil.AddNamedImport(fset, file, s.Name.Name, unquotedPath)
			continue
		}

		astutil.AddImport(fset, file, unquotedPath)
	}

	return file, fset, nil
}

func getImports(dirPath string, fset *token.FileSet) ([]*ast.ImportSpec, error) {
	// Grab imports from all files except helheim_test
	pkgs, err := parser.ParseDir(fset, dirPath, func(info os.FileInfo) bool {
		return !strings.Contains(info.Name(), "_test.go")
	}, parser.ImportsOnly)
	if err != nil {
		return nil, err
	}

	var imports []*ast.ImportSpec
	for _, p := range pkgs {
		files := p.Files
		for _, f := range files {
			imports = append(imports, f.Imports...)
		}
	}
	return imports, nil
}

func Generate(finder TypeFinder) (Mocks, error) {
	base := finder.ExportedTypes()
	var types []*ast.TypeSpec
	for _, typ := range base {
		types = append(types, typ)
		if inter, ok := typ.Type.(*ast.InterfaceType); ok {
			types = append(types, finder.Dependencies(inter)...)
		}
	}
	m := make(Mocks, 0, len(types))
	for _, typ := range types {
		newMock, err := For(typ)
		if err != nil {
			return nil, err
		}
		m = append(m, newMock)
	}
	return m, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
)

const MSG_BUFFER int = 50
const MAX_MSG_LENGTH int = 512

const HELP_TEXT string = SYSTEM_MESSAGE_FORMAT + `-> Available commands:
   /about               - About this chat.
   /exit                - Exit the chat.
   /help                - Show this help text.
   /list                - List the users that are currently connected.
   /beep                - Enable BEL notifications on mention.
   /me $ACTION          - Show yourself doing an action.
   /nick $NAME          - Rename yourself to a new name.
   /whois $NAME         - Display information about another connected user.
   /msg $NAME $MESSAGE  - Sends a private message to a user.
   /motd                - Prints the Message of the Day
` + RESET

const OP_HELP_TEXT string = SYSTEM_MESSAGE_FORMAT + `-> Available operator commands:
   /ban $NAME       - Banish a user from the chat
   /kick $NAME      - Kick em' out.
   /op $NAME        - Promote a user to server operator
   /silence $NAME   - Revoke a user's ability to speak
   /ban $NAME           - Banish a user from the chat
   /kick $NAME          - Kick em' out.
   /op $NAME            - Promote a user to server operator.
   /silence $NAME       - Revoke a user's ability to speak.
   /motd $MESSAGE    - Sets the Message of the Day
` + RESET

const ABOUT_TEXT string = SYSTEM_MESSAGE_FORMAT + `-> ssh-chat is made by @shazow.

   It is a custom ssh server built in Go to serve a chat experience
   instead of a shell.

   Source: https://github.com/shazow/ssh-chat

   For more, visit shazow.net or follow at twitter.com/shazow
` + RESET

const REQUIRED_WAIT time.Duration = time.Second / 2

type Client struct {
	Server        *Server
	Conn          *ssh.ServerConn
	Msg           chan string
	Name          string
	Color         string
	Op            bool
	ready         chan struct{}
	term          *terminal.Terminal
	termWidth     int
	termHeight    int
	silencedUntil time.Time
	lastTX        time.Time
	beepMe        bool
}

func NewClient(server *Server, conn *ssh.ServerConn) *Client {
	return &Client{
		Server: server,
		Conn:   conn,
		Name:   conn.User(),
		Color:  RandomColor256(),
		Msg:    make(chan string, MSG_BUFFER),
		ready:  make(chan struct{}, 1),
		lastTX: time.Now(),
	}
}

func (c *Client) ColoredName() string {
	return ColorString(c.Color, c.Name)
}

func (c *Client) SysMsg(msg string, args ...interface{}) {
	c.Msg <- ContinuousFormat(SYSTEM_MESSAGE_FORMAT, "-> "+fmt.Sprintf(msg, args...))
}

func (c *Client) Write(msg string) {
	c.term.Write([]byte(msg + "\r\n"))
}

func (c *Client) WriteLines(msg []string) {
	for _, line := range msg {
		c.Write(line)
	}
}

func (c *Client) Send(msg string) {
	if len(msg) > MAX_MSG_LENGTH {
		return
	}
	select {
	case c.Msg <- msg:
	default:
		logger.Errorf("Msg buffer full, dropping: %s (%s)", c.Name, c.Conn.RemoteAddr())
		c.Conn.Close()
	}
}

func (c *Client) SendLines(msg []string) {
	for _, line := range msg {
		c.Send(line)
	}
}

func (c *Client) IsSilenced() bool {
	return c.silencedUntil.After(time.Now())
}

func (c *Client) Silence(d time.Duration) {
	c.silencedUntil = time.Now().Add(d)
}

func (c *Client) Resize(width int, height int) error {
	err := c.term.SetSize(width, height)
	if err != nil {
		logger.Errorf("Resize failed: %dx%d", width, height)
		return err
	}
	c.termWidth, c.termHeight = width, height
	return nil
}

func (c *Client) Rename(name string) {
	c.Name = name
	c.term.SetPrompt(fmt.Sprintf("[%s] ", c.ColoredName()))
}

func (c *Client) Fingerprint() string {
	return c.Conn.Permissions.Extensions["fingerprint"]
}

func (c *Client) handleShell(channel ssh.Channel) {
	defer channel.Close()

	// FIXME: This shouldn't live here, need to restructure the call chaining.
	c.Server.Add(c)
	go func() {
		// Block until done, then remove.
		c.Conn.Wait()
		c.Server.Remove(c)
	}()

	go func() {
		for msg := range c.Msg {
			c.Write(msg)
		}
	}()

	for {
		line, err := c.term.ReadLine()
		if err != nil {
			break
		}

		parts := strings.SplitN(line, " ", 3)
		isCmd := strings.HasPrefix(parts[0], "/")

		if isCmd {
			// TODO: Factor this out.
			switch parts[0] {
			case "/test-colors": // Shh, this command is a secret!
				c.Write(ColorString("32", "Lorem ipsum dolor sit amet,"))
				c.Write("consectetur " + ColorString("31;1", "adipiscing") + " elit.")
			case "/exit":
				channel.Close()
			case "/help":
				c.WriteLines(strings.Split(HELP_TEXT, "\n"))
				if c.Server.IsOp(c) {
					c.WriteLines(strings.Split(OP_HELP_TEXT, "\n"))
				}
			case "/about":
				c.WriteLines(strings.Split(ABOUT_TEXT, "\n"))
			case "/uptime":
				c.Write(c.Server.Uptime())
			case "/beep":
				c.beepMe = !c.beepMe
				if c.beepMe {
					c.SysMsg("I'll beep you good.")
				} else {
					c.SysMsg("No more beeps. :(")
				}
			case "/me":
				me := strings.TrimLeft(line, "/me")
				if me == "" {
					me = " is at a loss for words."
				}
				msg := fmt.Sprintf("** %s%s", c.ColoredName(), me)
				if c.IsSilenced() || len(msg) > 1000 {
					c.SysMsg("Message rejected.")
				} else {
					c.Server.Broadcast(msg, nil)
				}
			case "/nick":
				if len(parts) == 2 {
					c.Server.Rename(c, parts[1])
				} else {
					c.SysMsg("Missing $NAME from: /nick $NAME")
				}
			case "/whois":
				if len(parts) == 2 {
					client := c.Server.Who(parts[1])
					if client != nil {
						version := RE_STRIP_TEXT.ReplaceAllString(string(client.Conn.ClientVersion()), "")
						if len(version) > 100 {
							version = "Evil Jerk with a superlong string"
						}
						c.SysMsg("%s is %s via %s", client.ColoredName(), client.Fingerprint(), version)
					} else {
						c.SysMsg("No such name: %s", parts[1])
					}
				} else {
					c.SysMsg("Missing $NAME from: /whois $NAME")
				}
			case "/list":
				names := ""
				nameList := c.Server.List(nil)
				for _, name := range nameList {
					names += c.Server.Who(name).ColoredName() + SYSTEM_MESSAGE_FORMAT + ", "
				}
				if len(names) > 2 {
					names = names[:len(names)-2]
				}
				c.SysMsg("%d connected: %s", len(nameList), names)
			case "/ban":
				if !c.Server.IsOp(c) {
					c.SysMsg("You're not an admin.")
				} else if len(parts) != 2 {
					c.SysMsg("Missing $NAME from: /ban $NAME")
				} else {
					client := c.Server.Who(parts[1])
					if client == nil {
						c.SysMsg("No such name: %s", parts[1])
					} else {
						fingerprint := client.Fingerprint()
						client.SysMsg("Banned by %s.", c.ColoredName())
						c.Server.Ban(fingerprint, nil)
						client.Conn.Close()
						c.Server.Broadcast(fmt.Sprintf("* %s was banned by %s", parts[1], c.ColoredName()), nil)
					}
				}
			case "/op":
				if !c.Server.IsOp(c) {
					c.SysMsg("You're not an admin.")
				} else if len(parts) != 2 {
					c.SysMsg("Missing $NAME from: /op $NAME")
				} else {
					client := c.Server.Who(parts[1])
					if client == nil {
						c.SysMsg("No such name: %s", parts[1])
					} else {
						fingerprint := client.Fingerprint()
						client.SysMsg("Made op by %s.", c.ColoredName())
						c.Server.Op(fingerprint)
					}
				}
			case "/kick":
				if !c.Server.IsOp(c) {
					c.SysMsg("You're not an admin.")
				} else if len(parts) != 2 {
					c.SysMsg("Missing $NAME from: /kick $NAME")
				} else {
					client := c.Server.Who(parts[1])
					if client == nil {
						c.SysMsg("No such name: %s", parts[1])
					} else {
						client.SysMsg("Kicked by %s.", c.ColoredName())
						client.Conn.Close()
						c.Server.Broadcast(fmt.Sprintf("* %s was kicked by %s", parts[1], c.ColoredName()), nil)
					}
				}
			case "/silence":
				if !c.Server.IsOp(c) {
					c.SysMsg("You're not an admin.")
				} else if len(parts) < 2 {
					c.SysMsg("Missing $NAME from: /silence $NAME")
				} else {
					duration := time.Duration(5) * time.Minute
					if len(parts) >= 3 {
						parsedDuration, err := time.ParseDuration(parts[2])
						if err == nil {
							duration = parsedDuration
						}
					}
					client := c.Server.Who(parts[1])
					if client == nil {
						c.SysMsg("No such name: %s", parts[1])
					} else {
						client.Silence(duration)
						client.SysMsg("Silenced for %s by %s.", duration, c.ColoredName())
					}
				}
			case "/msg": /* Send a PM */
				/* Make sure we have a recipient and a message */
				if len(parts) < 2 {
					c.SysMsg("Missing $NAME from: /msg $NAME $MESSAGE")
					break
				} else if len(parts) < 3 {
					c.SysMsg("Missing $MESSAGE from: /msg $NAME $MESSAGE")
					break
				}
				/* Ask the server to send the message */
				if err := c.Server.Privmsg(parts[1], parts[2], c); nil != err {
					c.SysMsg("Unable to send message to %v: %v", parts[1], err)
				}
			case "/motd": /* print motd */
				if !c.Server.IsOp(c) {
					c.Server.MotdUnicast(c)
				} else if len(parts) < 2 {
					c.Server.MotdUnicast(c)
				} else {
					var newmotd string
					if (len(parts) == 2) {
						newmotd = parts[1]
					} else {
						newmotd = parts[1] + " " + parts[2]
					}
					c.Server.SetMotd(c, newmotd)
					c.Server.MotdBroadcast(c)
				}

			default:
				c.SysMsg("Invalid command: %s", line)
			}
			continue
		}

		msg := fmt.Sprintf("%s: %s", c.ColoredName(), line)
		/* Rate limit */
		if time.Now().Sub(c.lastTX) < REQUIRED_WAIT {
			c.SysMsg("Rate limiting in effect.")
			continue
		}
		if c.IsSilenced() || len(msg) > 1000 || len(line) < 1 {
			c.SysMsg("Message rejected.")
			continue
		}
		c.Server.Broadcast(msg, c)
		c.lastTX = time.Now()
	}

}

func (c *Client) handleChannels(channels <-chan ssh.NewChannel) {
	prompt := fmt.Sprintf("[%s] ", c.ColoredName())

	hasShell := false

	for ch := range channels {
		if t := ch.ChannelType(); t != "session" {
			ch.Reject(ssh.UnknownChannelType, fmt.Sprintf("unknown channel type: %s", t))
			continue
		}

		channel, requests, err := ch.Accept()
		if err != nil {
			logger.Errorf("Could not accept channel: %v", err)
			continue
		}
		defer channel.Close()

		c.term = terminal.NewTerminal(channel, prompt)
		c.term.AutoCompleteCallback = c.Server.AutoCompleteFunction

		for req := range requests {
			var width, height int
			var ok bool

			switch req.Type {
			case "shell":
				if c.term != nil && !hasShell {
					go c.handleShell(channel)
					ok = true
					hasShell = true
				}
			case "pty-req":
				width, height, ok = parsePtyRequest(req.Payload)
				if ok {
					err := c.Resize(width, height)
					ok = err == nil
				}
			case "window-change":
				width, height, ok = parseWinchRequest(req.Payload)
				if ok {
					err := c.Resize(width, height)
					ok = err == nil
				}
			}

			if req.WantReply {
				req.Reply(ok, nil)
			}
		}
	}
}
<commit_msg>clean up duplicate commands in OP_HELP_TEXT<commit_after>package main

import (
	"fmt"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
)

const MSG_BUFFER int = 50
const MAX_MSG_LENGTH int = 512

const HELP_TEXT string = SYSTEM_MESSAGE_FORMAT + `-> Available commands:
   /about               - About this chat.
   /exit                - Exit the chat.
   /help                - Show this help text.
   /list                - List the users that are currently connected.
   /beep                - Enable BEL notifications on mention.
   /me $ACTION          - Show yourself doing an action.
   /nick $NAME          - Rename yourself to a new name.
   /whois $NAME         - Display information about another connected user.
   /msg $NAME $MESSAGE  - Sends a private message to a user.
   /motd                - Prints the Message of the Day
` + RESET

const OP_HELP_TEXT string = SYSTEM_MESSAGE_FORMAT + `-> Available operator commands:
   /ban $NAME           - Banish a user from the chat
   /kick $NAME          - Kick em' out.
   /op $NAME            - Promote a user to server operator.
   /silence $NAME       - Revoke a user's ability to speak.
   /motd $MESSAGE       - Sets the Message of the Day
` + RESET

const ABOUT_TEXT string = SYSTEM_MESSAGE_FORMAT + `-> ssh-chat is made by @shazow.

   It is a custom ssh server built in Go to serve a chat experience
   instead of a shell.

   Source: https://github.com/shazow/ssh-chat

   For more, visit shazow.net or follow at twitter.com/shazow
` + RESET

const REQUIRED_WAIT time.Duration = time.Second / 2

type Client struct {
	Server        *Server
	Conn          *ssh.ServerConn
	Msg           chan string
	Name          string
	Color         string
	Op            bool
	ready         chan struct{}
	term          *terminal.Terminal
	termWidth     int
	termHeight    int
	silencedUntil time.Time
	lastTX        time.Time
	beepMe        bool
}

func NewClient(server *Server, conn *ssh.ServerConn) *Client {
	return &Client{
		Server: server,
		Conn:   conn,
		Name:   conn.User(),
		Color:  RandomColor256(),
		Msg:    make(chan string, MSG_BUFFER),
		ready:  make(chan struct{}, 1),
		lastTX: time.Now(),
	}
}

func (c *Client) ColoredName() string {
	return ColorString(c.Color, c.Name)
}

func (c *Client) SysMsg(msg string, args ...interface{}) {
	c.Msg <- ContinuousFormat(SYSTEM_MESSAGE_FORMAT, "-> "+fmt.Sprintf(msg, args...))
}

func (c *Client) Write(msg string) {
	c.term.Write([]byte(msg + "\r\n"))
}

func (c *Client) WriteLines(msg []string) {
	for _, line := range msg {
		c.Write(line)
	}
}

func (c *Client) Send(msg string) {
	if len(msg) > MAX_MSG_LENGTH {
		return
	}
	select {
	case c.Msg <- msg:
	default:
		logger.Errorf("Msg buffer full, dropping: %s (%s)", c.Name, c.Conn.RemoteAddr())
		c.Conn.Close()
	}
}

func (c *Client) SendLines(msg []string) {
	for _, line := range msg {
		c.Send(line)
	}
}

func (c *Client) IsSilenced() bool {
	return c.silencedUntil.After(time.Now())
}

func (c *Client) Silence(d time.Duration) {
	c.silencedUntil = time.Now().Add(d)
}

func (c *Client) Resize(width int, height int) error {
	err := c.term.SetSize(width, height)
	if err != nil {
		logger.Errorf("Resize failed: %dx%d", width, height)
		return err
	}
	c.termWidth, c.termHeight = width, height
	return nil
}

func (c *Client) Rename(name string) {
	c.Name = name
	c.term.SetPrompt(fmt.Sprintf("[%s] ", c.ColoredName()))
}

func (c *Client) Fingerprint() string {
	return c.Conn.Permissions.Extensions["fingerprint"]
}

func (c *Client) handleShell(channel ssh.Channel) {
	defer channel.Close()

	// FIXME: This shouldn't live here, need to restructure the call chaining.
	c.Server.Add(c)
	go func() {
		// Block until done, then remove.
		c.Conn.Wait()
		c.Server.Remove(c)
	}()

	go func() {
		for msg := range c.Msg {
			c.Write(msg)
		}
	}()

	for {
		line, err := c.term.ReadLine()
		if err != nil {
			break
		}

		parts := strings.SplitN(line, " ", 3)
		isCmd := strings.HasPrefix(parts[0], "/")

		if isCmd {
			// TODO: Factor this out.
			switch parts[0] {
			case "/test-colors": // Shh, this command is a secret!
				c.Write(ColorString("32", "Lorem ipsum dolor sit amet,"))
				c.Write("consectetur " + ColorString("31;1", "adipiscing") + " elit.")
			case "/exit":
				channel.Close()
			case "/help":
				c.WriteLines(strings.Split(HELP_TEXT, "\n"))
				if c.Server.IsOp(c) {
					c.WriteLines(strings.Split(OP_HELP_TEXT, "\n"))
				}
			case "/about":
				c.WriteLines(strings.Split(ABOUT_TEXT, "\n"))
			case "/uptime":
				c.Write(c.Server.Uptime())
			case "/beep":
				c.beepMe = !c.beepMe
				if c.beepMe {
					c.SysMsg("I'll beep you good.")
				} else {
					c.SysMsg("No more beeps. :(")
				}
			case "/me":
				me := strings.TrimLeft(line, "/me")
				if me == "" {
					me = " is at a loss for words."
				}
				msg := fmt.Sprintf("** %s%s", c.ColoredName(), me)
				if c.IsSilenced() || len(msg) > 1000 {
					c.SysMsg("Message rejected.")
				} else {
					c.Server.Broadcast(msg, nil)
				}
			case "/nick":
				if len(parts) == 2 {
					c.Server.Rename(c, parts[1])
				} else {
					c.SysMsg("Missing $NAME from: /nick $NAME")
				}
			case "/whois":
				if len(parts) == 2 {
					client := c.Server.Who(parts[1])
					if client != nil {
						version := RE_STRIP_TEXT.ReplaceAllString(string(client.Conn.ClientVersion()), "")
						if len(version) > 100 {
							version = "Evil Jerk with a superlong string"
						}
						c.SysMsg("%s is %s via %s", client.ColoredName(), client.Fingerprint(), version)
					} else {
						c.SysMsg("No such name: %s", parts[1])
					}
				} else {
					c.SysMsg("Missing $NAME from: /whois $NAME")
				}
			case "/list":
				names := ""
				nameList := c.Server.List(nil)
				for _, name := range nameList {
					names += c.Server.Who(name).ColoredName() + SYSTEM_MESSAGE_FORMAT + ", "
				}
				if len(names) > 2 {
					names = names[:len(names)-2]
				}
				c.SysMsg("%d connected: %s", len(nameList), names)
			case "/ban":
				if !c.Server.IsOp(c) {
					c.SysMsg("You're not an admin.")
				} else if len(parts) != 2 {
					c.SysMsg("Missing $NAME from: /ban $NAME")
				} else {
					client := c.Server.Who(parts[1])
					if client == nil {
						c.SysMsg("No such name: %s", parts[1])
					} else {
						fingerprint := client.Fingerprint()
						client.SysMsg("Banned by %s.", c.ColoredName())
						c.Server.Ban(fingerprint, nil)
						client.Conn.Close()
						c.Server.Broadcast(fmt.Sprintf("* %s was banned by %s", parts[1], c.ColoredName()), nil)
					}
				}
			case "/op":
				if !c.Server.IsOp(c) {
					c.SysMsg("You're not an admin.")
				} else if len(parts) != 2 {
					c.SysMsg("Missing $NAME from: /op $NAME")
				} else {
					client := c.Server.Who(parts[1])
					if client == nil {
						c.SysMsg("No such name: %s", parts[1])
					} else {
						fingerprint := client.Fingerprint()
						client.SysMsg("Made op by %s.", c.ColoredName())
						c.Server.Op(fingerprint)
					}
				}
			case "/kick":
				if !c.Server.IsOp(c) {
					c.SysMsg("You're not an admin.")
				} else if len(parts) != 2 {
					c.SysMsg("Missing $NAME from: /kick $NAME")
				} else {
					client := c.Server.Who(parts[1])
					if client == nil {
						c.SysMsg("No such name: %s", parts[1])
					} else {
						client.SysMsg("Kicked by %s.", c.ColoredName())
						client.Conn.Close()
						c.Server.Broadcast(fmt.Sprintf("* %s was kicked by %s", parts[1], c.ColoredName()), nil)
					}
				}
			case "/silence":
				if !c.Server.IsOp(c) {
					c.SysMsg("You're not an admin.")
				} else if len(parts) < 2 {
					c.SysMsg("Missing $NAME from: /silence $NAME")
				} else {
					duration := time.Duration(5) * time.Minute
					if len(parts) >= 3 {
						parsedDuration, err := time.ParseDuration(parts[2])
						if err == nil {
							duration = parsedDuration
						}
					}
					client := c.Server.Who(parts[1])
					if client == nil {
						c.SysMsg("No such name: %s", parts[1])
					} else {
						client.Silence(duration)
						client.SysMsg("Silenced for %s by %s.", duration, c.ColoredName())
					}
				}
			case "/msg": /* Send a PM */
				/* Make sure we have a recipient and a message */
				if len(parts) < 2 {
					c.SysMsg("Missing $NAME from: /msg $NAME $MESSAGE")
					break
				} else if len(parts) < 3 {
					c.SysMsg("Missing $MESSAGE from: /msg $NAME $MESSAGE")
					break
				}
				/* Ask the server to send the message */
				if err := c.Server.Privmsg(parts[1], parts[2], c); nil != err {
					c.SysMsg("Unable to send message to %v: %v", parts[1], err)
				}
			case "/motd": /* print motd */
				if !c.Server.IsOp(c) {
					c.Server.MotdUnicast(c)
				} else if len(parts) < 2 {
					c.Server.MotdUnicast(c)
				} else {
					var newmotd string
					if (len(parts) == 2) {
						newmotd = parts[1]
					} else {
						newmotd = parts[1] + " " + parts[2]
					}
					c.Server.SetMotd(c, newmotd)
					c.Server.MotdBroadcast(c)
				}

			default:
				c.SysMsg("Invalid command: %s", line)
			}
			continue
		}

		msg := fmt.Sprintf("%s: %s", c.ColoredName(), line)
		/* Rate limit */
		if time.Now().Sub(c.lastTX) < REQUIRED_WAIT {
			c.SysMsg("Rate limiting in effect.")
			continue
		}
		if c.IsSilenced() || len(msg) > 1000 || len(line) < 1 {
			c.SysMsg("Message rejected.")
			continue
		}
		c.Server.Broadcast(msg, c)
		c.lastTX = time.Now()
	}

}

func (c *Client) handleChannels(channels <-chan ssh.NewChannel) {
	prompt := fmt.Sprintf("[%s] ", c.ColoredName())

	hasShell := false

	for ch := range channels {
		if t := ch.ChannelType(); t != "session" {
			ch.Reject(ssh.UnknownChannelType, fmt.Sprintf("unknown channel type: %s", t))
			continue
		}

		channel, requests, err := ch.Accept()
		if err != nil {
			logger.Errorf("Could not accept channel: %v", err)
			continue
		}
		defer channel.Close()

		c.term = terminal.NewTerminal(channel, prompt)
		c.term.AutoCompleteCallback = c.Server.AutoCompleteFunction

		for req := range requests {
			var width, height int
			var ok bool

			switch req.Type {
			case "shell":
				if c.term != nil && !hasShell {
					go c.handleShell(channel)
					ok = true
					hasShell = true
				}
			case "pty-req":
				width, height, ok = parsePtyRequest(req.Payload)
				if ok {
					err := c.Resize(width, height)
					ok = err == nil
				}
			case "window-change":
				width, height, ok = parseWinchRequest(req.Payload)
				if ok {
					err := c.Resize(width, height)
					ok = err == nil
				}
			}

			if req.WantReply {
				req.Reply(ok, nil)
			}
		}
	}
}
<|endoftext|>


<commit_before>package paypalsdk

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

// NewClient returns new Client struct
// APIBase is a base API URL, for testing you can use paypalsdk.APIBaseSandBox
func NewClient(clientID string, secret string, APIBase string) (*Client, error) {
	if clientID == "" || secret == "" || APIBase == "" {
		return &Client{}, errors.New("ClientID, Secret and APIBase are required to create a Client")
	}

	return &Client{
		&http.Client{},
		clientID,
		secret,
		APIBase,
		nil,
		nil,
	}, nil
}

// SetLog will set/change the output destination.
// If log file is set paypalsdk will log all requests and responses to this Writer
func (c *Client) SetLog(log io.Writer) error {
	c.Log = log
	return nil
}

// SetAccessToken sets saved token to current client
func (c *Client) SetAccessToken(token string) error {
	c.Token = &TokenResponse{
		Token: token,
	}

	return nil
}

// Send makes a request to the API, the response body will be
// unmarshaled into v, or if v is an io.Writer, the response will
// be written to it without decoding
func (c *Client) Send(req *http.Request, v interface{}) error {
	var (
		err  error
		resp *http.Response
		data []byte
	)

	// Set default headers
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Accept-Language", "en_US")

	// Default values for headers
	if req.Header.Get("Content-type") == "" {
		req.Header.Set("Content-type", "application/json")
	}

	resp, err = c.client.Do(req)
	c.log(req, resp)

	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		errResp := &ErrorResponse{Response: resp}
		data, err = ioutil.ReadAll(resp.Body)

		if err == nil && len(data) > 0 {
			json.Unmarshal(data, errResp)
		}

		return errResp
	}

	if v != nil {
		if w, ok := v.(io.Writer); ok {
			io.Copy(w, resp.Body)
		} else {
			err = json.NewDecoder(resp.Body).Decode(v)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// SendWithAuth makes a request to the API and apply OAuth2 header automatically.
// If the access token soon to be expired or already expired, it will try to get a new one before
// making the main request
// client.Token will be updated when changed
func (c *Client) SendWithAuth(req *http.Request, v interface{}) error {
	if c.Token != nil {
		if c.Token.ExpiresIn < RequestNewTokenBeforeExpiresIn {
			// c.Token willbe updated in GetAccessToken call
			_, err := c.GetAccessToken()
			if err != nil {
				return err
			}
		}

		req.Header.Set("Authorization", "Bearer "+c.Token.Token)
	}

	return c.Send(req, v)
}

// NewRequest constructs a request
// Convert payload to a JSON
func (c *Client) NewRequest(method, url string, payload interface{}) (*http.Request, error) {
	var buf io.Reader
	if payload != nil {
		var b []byte
		b, err := json.Marshal(&payload)
		if err != nil {
			return nil, err
		}
		buf = bytes.NewBuffer(b)
	}
	return http.NewRequest(method, url, buf)
}

// log will dump request and response to the log file
func (c *Client) log(r *http.Request, resp *http.Response) {
	if c.Log != nil {
		reqDump := fmt.Sprintf("%s %s. Data: %s", r.Method, r.URL.String(), r.Form.Encode())
		dump, _ := ioutil.ReadAll(r.Body)
		fmt.Println(string(dump))
		respDump, _ := httputil.DumpResponse(resp, true)

		c.Log.Write([]byte("Request: " + reqDump + "\nResponse: " + string(respDump) + "\n\n"))
	}
}
<commit_msg>Fixed panic on logging requests without body (GET)<commit_after>package paypalsdk

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

// NewClient returns new Client struct
// APIBase is a base API URL, for testing you can use paypalsdk.APIBaseSandBox
func NewClient(clientID string, secret string, APIBase string) (*Client, error) {
	if clientID == "" || secret == "" || APIBase == "" {
		return &Client{}, errors.New("ClientID, Secret and APIBase are required to create a Client")
	}

	return &Client{
		&http.Client{},
		clientID,
		secret,
		APIBase,
		nil,
		nil,
	}, nil
}

// SetLog will set/change the output destination.
// If log file is set paypalsdk will log all requests and responses to this Writer
func (c *Client) SetLog(log io.Writer) error {
	c.Log = log
	return nil
}

// SetAccessToken sets saved token to current client
func (c *Client) SetAccessToken(token string) error {
	c.Token = &TokenResponse{
		Token: token,
	}

	return nil
}

// Send makes a request to the API, the response body will be
// unmarshaled into v, or if v is an io.Writer, the response will
// be written to it without decoding
func (c *Client) Send(req *http.Request, v interface{}) error {
	var (
		err  error
		resp *http.Response
		data []byte
	)

	// Set default headers
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Accept-Language", "en_US")

	// Default values for headers
	if req.Header.Get("Content-type") == "" {
		req.Header.Set("Content-type", "application/json")
	}

	resp, err = c.client.Do(req)
	c.log(req, resp)

	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		errResp := &ErrorResponse{Response: resp}
		data, err = ioutil.ReadAll(resp.Body)

		if err == nil && len(data) > 0 {
			json.Unmarshal(data, errResp)
		}

		return errResp
	}

	if v != nil {
		if w, ok := v.(io.Writer); ok {
			io.Copy(w, resp.Body)
		} else {
			err = json.NewDecoder(resp.Body).Decode(v)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// SendWithAuth makes a request to the API and apply OAuth2 header automatically.
// If the access token soon to be expired or already expired, it will try to get a new one before
// making the main request
// client.Token will be updated when changed
func (c *Client) SendWithAuth(req *http.Request, v interface{}) error {
	if c.Token != nil {
		if c.Token.ExpiresIn < RequestNewTokenBeforeExpiresIn {
			// c.Token willbe updated in GetAccessToken call
			_, err := c.GetAccessToken()
			if err != nil {
				return err
			}
		}

		req.Header.Set("Authorization", "Bearer "+c.Token.Token)
	}

	return c.Send(req, v)
}

// NewRequest constructs a request
// Convert payload to a JSON
func (c *Client) NewRequest(method, url string, payload interface{}) (*http.Request, error) {
	var buf io.Reader
	if payload != nil {
		var b []byte
		b, err := json.Marshal(&payload)
		if err != nil {
			return nil, err
		}
		buf = bytes.NewBuffer(b)
	}
	return http.NewRequest(method, url, buf)
}

// log will dump request and response to the log file
func (c *Client) log(r *http.Request, resp *http.Response) {
	if c.Log != nil {
		reqDump := fmt.Sprintf("%s %s. Data: %s", r.Method, r.URL.String(), r.Form.Encode())
		respDump, _ := httputil.DumpResponse(resp, true)

		c.Log.Write([]byte("Request: " + reqDump + "\nResponse: " + string(respDump) + "\n\n"))
	}
}
<|endoftext|>


<commit_before>/*
A smart client for go.

Usage:

 client, err := couchbase.Connect("http://myserver:8091/")
 handleError(err)
 pool, err := client.GetPool("default")
 handleError(err)
 bucket, err := pool.getBucket("MyAwesomeBucket")
 handleError(err)
 ...

or a shortcut for the bucket directly

 bucket, err := couchbase.GetBucket("http://myserver:8091/", "default", "default")
*/
package couchbase

import (
	"encoding/json"
	"errors"
	"fmt"
	"math/rand"
	"net/url"
	"sync"
	"sync/atomic"

	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

type connectionPool struct {
	host        string
	name        string
	connections []*memcached.Client
	mutex       sync.Mutex
}

func (cp *connectionPool) Close() error {
	cp.mutex.Lock()
	defer cp.mutex.Unlock()
	for _, c := range cp.connections {
		c.Close()
	}
	cp.connections = []*memcached.Client{}
	return nil
}

func (cp *connectionPool) Get() (*memcached.Client, error) {
	cp.mutex.Lock()
	defer cp.mutex.Unlock()

	if len(cp.connections) == 0 {
		conn, err := memcached.Connect("tcp", cp.host)
		if err != nil {
			return nil, err
		}
		if cp.name != "default" {
			conn.Auth(cp.name, "")
		}

		cp.connections = append(cp.connections, conn)
	}

	rv := cp.connections[0]
	cp.connections = cp.connections[1:]

	return rv, nil
}

func (cp *connectionPool) Return(c *memcached.Client) {
	cp.mutex.Lock()
	defer cp.mutex.Unlock()

	if c != nil {
		if c.IsHealthy() {
			cp.connections = append(cp.connections, c)
		} else {
			c.Close()
		}
	}
}

// Execute a function on a memcached connection to the node owning key "k"
//
// Note that this automatically handles transient errors by replaying
// your function on a "not-my-vbucket" error, so don't assume
// your command will only be executed only once.
func (b *Bucket) Do(k string, f func(mc *memcached.Client, vb uint16) error) error {
	vb := b.VBHash(k)
	for {
		masterId := b.VBucketServerMap.VBucketMap[vb][0]
		conn, err := b.connections[masterId].Get()
		defer b.connections[masterId].Return(conn)
		if err != nil {
			return err
		}

		err = f(conn, uint16(vb))
		switch err.(type) {
		default:
			return err
		case gomemcached.MCResponse:
			st := err.(gomemcached.MCResponse).Status
			atomic.AddUint64(&b.pool.client.Statuses[st], 1)
			if st == gomemcached.NOT_MY_VBUCKET {
				b.refresh()
			} else {
				return err
			}
		}
	}
	panic("Unreachable.")
}

type gathered_stats struct {
	sn   string
	vals map[string]string
}

func getStatsParallel(b *Bucket, offset int, which string, ch chan<- gathered_stats) {
	sn := b.VBucketServerMap.ServerList[offset]

	results := map[string]string{}
	conn, err := b.connections[offset].Get()
	defer b.connections[offset].Return(conn)
	if err != nil {
		ch <- gathered_stats{sn, results}
	} else {
		st, err := conn.StatsMap(which)
		if err == nil {
			ch <- gathered_stats{sn, st}
		} else {
			ch <- gathered_stats{sn, results}
		}
	}
}

// Get a set of stats from all servers.
//
// Returns a map of server ID -> map of stat key to map value.
func (b *Bucket) GetStats(which string) map[string]map[string]string {
	rv := map[string]map[string]string{}

	if b.VBucketServerMap.ServerList == nil {
		return rv
	}
	// Go grab all the things at once.
	todo := len(b.VBucketServerMap.ServerList)
	ch := make(chan gathered_stats, todo)

	for offset, _ := range b.VBucketServerMap.ServerList {
		go getStatsParallel(b, offset, which, ch)
	}

	// Gather the results
	for i := 0; i < len(b.VBucketServerMap.ServerList); i++ {
		g := <-ch
		if len(g.vals) > 0 {
			rv[g.sn] = g.vals
		}
	}

	return rv
}

func (b *Bucket) doBulkGet(vb uint16, keys []string,
	ch chan<- map[string]*gomemcached.MCResponse) {

	masterId := b.VBucketServerMap.VBucketMap[vb][0]
	conn, err := b.connections[masterId].Get()
	if err != nil {
		ch <- map[string]*gomemcached.MCResponse{}
	}
	defer b.connections[masterId].Return(conn)

	m, err := conn.GetBulk(vb, keys)
	switch err.(type) {
	default:
		ch <- m
	case *gomemcached.MCResponse:
		fmt.Printf("Got a memcached error")
		st := err.(gomemcached.MCResponse).Status
		atomic.AddUint64(&b.pool.client.Statuses[st], 1)
		if st == gomemcached.NOT_MY_VBUCKET {
			b.refresh()
		}
		ch <- map[string]*gomemcached.MCResponse{}
	}
}

func (b *Bucket) processBulkGet(kdm map[uint16][]string,
	ch chan map[string]*gomemcached.MCResponse) {

	wch := make(chan uint16)

	worker := func() {
		for k := range wch {
			b.doBulkGet(k, kdm[k], ch)
		}
	}

	for i := 0; i < 4; i++ {
		go worker()
	}

	for k := range kdm {
		wch <- k
	}
	close(wch)

}
func (b *Bucket) GetBulk(keys []string) map[string]*gomemcached.MCResponse {
	// Organize by vbucket
	kdm := map[uint16][]string{}
	for _, k := range keys {
		vb := uint16(b.VBHash(k))
		a, ok := kdm[vb]
		if !ok {
			a = []string{}
		}
		kdm[vb] = append(a, k)
	}

	ch := make(chan map[string]*gomemcached.MCResponse)
	defer close(ch)

	go b.processBulkGet(kdm, ch)

	rv := map[string]*gomemcached.MCResponse{}
	for _ = range kdm {
		m := <-ch
		for k, v := range m {
			rv[k] = v
		}
	}

	return rv
}

// Set a value in this bucket.
// The value will be serialized into a JSON document.
func (b *Bucket) Set(k string, exp int, v interface{}) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}
	return b.SetRaw(k, exp, data)
}

// Set a value in this bucket.
// The value will be stored as raw bytes.
func (b *Bucket) SetRaw(k string, exp int, v []byte) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Set(vb, k, 0, exp, v)
		if err != nil {
			return err
		}
		if res.Status != gomemcached.SUCCESS {
			return res
		}
		return nil
	})
}

// Adds a value to this bucket; like Set except that nothing happens if the key exists.
// The value will be serialized into a JSON document.
func (b *Bucket) Add(k string, exp int, v interface{}) (added bool, err error) {
	data, err := json.Marshal(v)
	if err != nil {
		return false, err
	}
	return b.AddRaw(k, exp, data)
}

// Adds a value to this bucket; like SetRaw except that nothing happens if the key exists.
// The value will be stored as raw bytes.
func (b *Bucket) AddRaw(k string, exp int, v []byte) (added bool, err error) {
	err = b.Do(k, func(mc *memcached.Client, vb uint16) error {
		switch res, err := mc.Add(vb, k, 0, exp, v); {
		case err != nil:
			return err
		case res.Status == gomemcached.SUCCESS:
			added = true
		case res.Status != gomemcached.KEY_EEXISTS:
			return res
		}
		return nil
	})
	return
}

// Get a raw value from this bucket, including its CAS counter.
func (b *Bucket) GetsRaw(k string, cas *uint64) ([]byte, error) {
	var data []byte
	err := b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Get(vb, k)
		if err != nil {
			return err
		}
		if res.Status != gomemcached.SUCCESS {
			return res
		}
		if cas != nil {
			*cas = res.Cas
		}
		data = res.Body
		return nil
	})
	return data, err
}

// Get a value from this bucket, including its CAS counter.
// The value is expected to be a JSON stream and will be deserialized
// into rv.
func (b *Bucket) Gets(k string, rv interface{}, cas *uint64) error {
	data, err := b.GetsRaw(k, cas)
	if err != nil {
		return err
	}
	return json.Unmarshal(data, rv)
}

// Get a value from this bucket.
// The value is expected to be a JSON stream and will be deserialized
// into rv.
func (b *Bucket) Get(k string, rv interface{}) error {
	return b.Gets(k, rv, nil)
}

// Get a raw value from this bucket.
func (b *Bucket) GetRaw(k string) ([]byte, error) {
	return b.GetsRaw(k, nil)
}

// Delete a key from this bucket.
func (b *Bucket) Delete(k string) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Del(vb, k)
		if err != nil {
			return err
		}
		if res.Status != gomemcached.SUCCESS {
			return res
		}
		return nil
	})
}

// Increment a key
func (b *Bucket) Incr(k string, amt, def uint64, exp int) (uint64, error) {
	var rv uint64
	err := b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Incr(vb, k, amt, def, exp)
		if err != nil {
			return err
		}
		rv = res
		return nil
	})
	return rv, err
}

type ViewRow struct {
	ID    string
	Key   interface{}
	Value interface{}
	Doc   *interface{}
}

type ViewError struct {
	From   string
	Reason string
}

func (ve ViewError) Error() string {
	return fmt.Sprintf("Node: %v, reason: %v", ve.From, ve.Reason)
}

type ViewResult struct {
	TotalRows int `json:"total_rows"`
	Rows      []ViewRow
	Errors    []ViewError
}

// Document ID type for the startkey_docid parameter in views.
type DocId string

// Perform a view request that can map row values to a custom type.
//
// See the source to View for an example usage.
func (b *Bucket) ViewCustom(ddoc, name string, params map[string]interface{},
	vres interface{}) error {

	// Pick a random node to service our request.
	node := b.Nodes[rand.Intn(len(b.Nodes))]
	u, err := url.Parse(node.CouchAPIBase)
	if err != nil {
		return err
	}

	values := url.Values{}
	for k, v := range params {
		switch t := v.(type) {
		case DocId:
			values[k] = []string{string(t)}
		case string:
			values[k] = []string{fmt.Sprintf(`"%s"`, t)}
		case int:
			values[k] = []string{fmt.Sprintf(`%d`, t)}
		case bool:
			values[k] = []string{fmt.Sprintf(`%v`, t)}
		default:
			b, err := json.Marshal(v)
			if err != nil {
				panic(fmt.Sprintf("unsupported value-type %T in Query, json encoder said %v", t, err))
			}
			values[k] = []string{fmt.Sprintf(`%v`, string(b))}
		}
	}

	u.Path = fmt.Sprintf("/%s/_design/%s/_view/%s", b.Name, ddoc, name)
	u.RawQuery = values.Encode()

	res, err := HttpClient.Get(u.String())
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.StatusCode != 200 {
		return errors.New(res.Status)
	}

	d := json.NewDecoder(res.Body)
	if err := d.Decode(vres); err != nil {
		return err
	}
	return nil
}

// Execute a view
func (b *Bucket) View(ddoc, name string, params map[string]interface{}) (ViewResult, error) {
	vres := ViewResult{}
	return vres, b.ViewCustom(ddoc, name, params, &vres)
}
<commit_msg>Handle Get() and Return() on nil pools.<commit_after>/*
A smart client for go.

Usage:

 client, err := couchbase.Connect("http://myserver:8091/")
 handleError(err)
 pool, err := client.GetPool("default")
 handleError(err)
 bucket, err := pool.getBucket("MyAwesomeBucket")
 handleError(err)
 ...

or a shortcut for the bucket directly

 bucket, err := couchbase.GetBucket("http://myserver:8091/", "default", "default")
*/
package couchbase

import (
	"encoding/json"
	"errors"
	"fmt"
	"math/rand"
	"net/url"
	"sync"
	"sync/atomic"

	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

type connectionPool struct {
	host        string
	name        string
	connections []*memcached.Client
	mutex       sync.Mutex
}

func (cp *connectionPool) Close() error {
	cp.mutex.Lock()
	defer cp.mutex.Unlock()
	for _, c := range cp.connections {
		c.Close()
	}
	cp.connections = []*memcached.Client{}
	return nil
}

func (cp *connectionPool) Get() (*memcached.Client, error) {
	if cp == nil {
		return nil, errors.New("no pool")
	}
	cp.mutex.Lock()
	defer cp.mutex.Unlock()

	if len(cp.connections) == 0 {
		conn, err := memcached.Connect("tcp", cp.host)
		if err != nil {
			return nil, err
		}
		if cp.name != "default" {
			conn.Auth(cp.name, "")
		}

		cp.connections = append(cp.connections, conn)
	}

	rv := cp.connections[0]
	cp.connections = cp.connections[1:]

	return rv, nil
}

func (cp *connectionPool) Return(c *memcached.Client) {
	if cp == nil {
		return
	}
	cp.mutex.Lock()
	defer cp.mutex.Unlock()

	if c != nil {
		if c.IsHealthy() {
			cp.connections = append(cp.connections, c)
		} else {
			c.Close()
		}
	}
}

// Execute a function on a memcached connection to the node owning key "k"
//
// Note that this automatically handles transient errors by replaying
// your function on a "not-my-vbucket" error, so don't assume
// your command will only be executed only once.
func (b *Bucket) Do(k string, f func(mc *memcached.Client, vb uint16) error) error {
	vb := b.VBHash(k)
	for {
		masterId := b.VBucketServerMap.VBucketMap[vb][0]
		conn, err := b.connections[masterId].Get()
		defer b.connections[masterId].Return(conn)
		if err != nil {
			return err
		}

		err = f(conn, uint16(vb))
		switch err.(type) {
		default:
			return err
		case gomemcached.MCResponse:
			st := err.(gomemcached.MCResponse).Status
			atomic.AddUint64(&b.pool.client.Statuses[st], 1)
			if st == gomemcached.NOT_MY_VBUCKET {
				b.refresh()
			} else {
				return err
			}
		}
	}
	panic("Unreachable.")
}

type gathered_stats struct {
	sn   string
	vals map[string]string
}

func getStatsParallel(b *Bucket, offset int, which string, ch chan<- gathered_stats) {
	sn := b.VBucketServerMap.ServerList[offset]

	results := map[string]string{}
	conn, err := b.connections[offset].Get()
	defer b.connections[offset].Return(conn)
	if err != nil {
		ch <- gathered_stats{sn, results}
	} else {
		st, err := conn.StatsMap(which)
		if err == nil {
			ch <- gathered_stats{sn, st}
		} else {
			ch <- gathered_stats{sn, results}
		}
	}
}

// Get a set of stats from all servers.
//
// Returns a map of server ID -> map of stat key to map value.
func (b *Bucket) GetStats(which string) map[string]map[string]string {
	rv := map[string]map[string]string{}

	if b.VBucketServerMap.ServerList == nil {
		return rv
	}
	// Go grab all the things at once.
	todo := len(b.VBucketServerMap.ServerList)
	ch := make(chan gathered_stats, todo)

	for offset, _ := range b.VBucketServerMap.ServerList {
		go getStatsParallel(b, offset, which, ch)
	}

	// Gather the results
	for i := 0; i < len(b.VBucketServerMap.ServerList); i++ {
		g := <-ch
		if len(g.vals) > 0 {
			rv[g.sn] = g.vals
		}
	}

	return rv
}

func (b *Bucket) doBulkGet(vb uint16, keys []string,
	ch chan<- map[string]*gomemcached.MCResponse) {

	masterId := b.VBucketServerMap.VBucketMap[vb][0]
	conn, err := b.connections[masterId].Get()
	if err != nil {
		ch <- map[string]*gomemcached.MCResponse{}
	}
	defer b.connections[masterId].Return(conn)

	m, err := conn.GetBulk(vb, keys)
	switch err.(type) {
	default:
		ch <- m
	case *gomemcached.MCResponse:
		fmt.Printf("Got a memcached error")
		st := err.(gomemcached.MCResponse).Status
		atomic.AddUint64(&b.pool.client.Statuses[st], 1)
		if st == gomemcached.NOT_MY_VBUCKET {
			b.refresh()
		}
		ch <- map[string]*gomemcached.MCResponse{}
	}
}

func (b *Bucket) processBulkGet(kdm map[uint16][]string,
	ch chan map[string]*gomemcached.MCResponse) {

	wch := make(chan uint16)

	worker := func() {
		for k := range wch {
			b.doBulkGet(k, kdm[k], ch)
		}
	}

	for i := 0; i < 4; i++ {
		go worker()
	}

	for k := range kdm {
		wch <- k
	}
	close(wch)

}
func (b *Bucket) GetBulk(keys []string) map[string]*gomemcached.MCResponse {
	// Organize by vbucket
	kdm := map[uint16][]string{}
	for _, k := range keys {
		vb := uint16(b.VBHash(k))
		a, ok := kdm[vb]
		if !ok {
			a = []string{}
		}
		kdm[vb] = append(a, k)
	}

	ch := make(chan map[string]*gomemcached.MCResponse)
	defer close(ch)

	go b.processBulkGet(kdm, ch)

	rv := map[string]*gomemcached.MCResponse{}
	for _ = range kdm {
		m := <-ch
		for k, v := range m {
			rv[k] = v
		}
	}

	return rv
}

// Set a value in this bucket.
// The value will be serialized into a JSON document.
func (b *Bucket) Set(k string, exp int, v interface{}) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}
	return b.SetRaw(k, exp, data)
}

// Set a value in this bucket.
// The value will be stored as raw bytes.
func (b *Bucket) SetRaw(k string, exp int, v []byte) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Set(vb, k, 0, exp, v)
		if err != nil {
			return err
		}
		if res.Status != gomemcached.SUCCESS {
			return res
		}
		return nil
	})
}

// Adds a value to this bucket; like Set except that nothing happens if the key exists.
// The value will be serialized into a JSON document.
func (b *Bucket) Add(k string, exp int, v interface{}) (added bool, err error) {
	data, err := json.Marshal(v)
	if err != nil {
		return false, err
	}
	return b.AddRaw(k, exp, data)
}

// Adds a value to this bucket; like SetRaw except that nothing happens if the key exists.
// The value will be stored as raw bytes.
func (b *Bucket) AddRaw(k string, exp int, v []byte) (added bool, err error) {
	err = b.Do(k, func(mc *memcached.Client, vb uint16) error {
		switch res, err := mc.Add(vb, k, 0, exp, v); {
		case err != nil:
			return err
		case res.Status == gomemcached.SUCCESS:
			added = true
		case res.Status != gomemcached.KEY_EEXISTS:
			return res
		}
		return nil
	})
	return
}

// Get a raw value from this bucket, including its CAS counter.
func (b *Bucket) GetsRaw(k string, cas *uint64) ([]byte, error) {
	var data []byte
	err := b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Get(vb, k)
		if err != nil {
			return err
		}
		if res.Status != gomemcached.SUCCESS {
			return res
		}
		if cas != nil {
			*cas = res.Cas
		}
		data = res.Body
		return nil
	})
	return data, err
}

// Get a value from this bucket, including its CAS counter.
// The value is expected to be a JSON stream and will be deserialized
// into rv.
func (b *Bucket) Gets(k string, rv interface{}, cas *uint64) error {
	data, err := b.GetsRaw(k, cas)
	if err != nil {
		return err
	}
	return json.Unmarshal(data, rv)
}

// Get a value from this bucket.
// The value is expected to be a JSON stream and will be deserialized
// into rv.
func (b *Bucket) Get(k string, rv interface{}) error {
	return b.Gets(k, rv, nil)
}

// Get a raw value from this bucket.
func (b *Bucket) GetRaw(k string) ([]byte, error) {
	return b.GetsRaw(k, nil)
}

// Delete a key from this bucket.
func (b *Bucket) Delete(k string) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Del(vb, k)
		if err != nil {
			return err
		}
		if res.Status != gomemcached.SUCCESS {
			return res
		}
		return nil
	})
}

// Increment a key
func (b *Bucket) Incr(k string, amt, def uint64, exp int) (uint64, error) {
	var rv uint64
	err := b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Incr(vb, k, amt, def, exp)
		if err != nil {
			return err
		}
		rv = res
		return nil
	})
	return rv, err
}

type ViewRow struct {
	ID    string
	Key   interface{}
	Value interface{}
	Doc   *interface{}
}

type ViewError struct {
	From   string
	Reason string
}

func (ve ViewError) Error() string {
	return fmt.Sprintf("Node: %v, reason: %v", ve.From, ve.Reason)
}

type ViewResult struct {
	TotalRows int `json:"total_rows"`
	Rows      []ViewRow
	Errors    []ViewError
}

// Document ID type for the startkey_docid parameter in views.
type DocId string

// Perform a view request that can map row values to a custom type.
//
// See the source to View for an example usage.
func (b *Bucket) ViewCustom(ddoc, name string, params map[string]interface{},
	vres interface{}) error {

	// Pick a random node to service our request.
	node := b.Nodes[rand.Intn(len(b.Nodes))]
	u, err := url.Parse(node.CouchAPIBase)
	if err != nil {
		return err
	}

	values := url.Values{}
	for k, v := range params {
		switch t := v.(type) {
		case DocId:
			values[k] = []string{string(t)}
		case string:
			values[k] = []string{fmt.Sprintf(`"%s"`, t)}
		case int:
			values[k] = []string{fmt.Sprintf(`%d`, t)}
		case bool:
			values[k] = []string{fmt.Sprintf(`%v`, t)}
		default:
			b, err := json.Marshal(v)
			if err != nil {
				panic(fmt.Sprintf("unsupported value-type %T in Query, json encoder said %v", t, err))
			}
			values[k] = []string{fmt.Sprintf(`%v`, string(b))}
		}
	}

	u.Path = fmt.Sprintf("/%s/_design/%s/_view/%s", b.Name, ddoc, name)
	u.RawQuery = values.Encode()

	res, err := HttpClient.Get(u.String())
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.StatusCode != 200 {
		return errors.New(res.Status)
	}

	d := json.NewDecoder(res.Body)
	if err := d.Decode(vres); err != nil {
		return err
	}
	return nil
}

// Execute a view
func (b *Bucket) View(ddoc, name string, params map[string]interface{}) (ViewResult, error) {
	vres := ViewResult{}
	return vres, b.ViewCustom(ddoc, name, params, &vres)
}
<|endoftext|>


<commit_before>package dockertest

import (
	"context"
	"errors"
	"io"
	"io/ioutil"
	"strings"

	"github.com/crewjam/errset"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/client"
)

var (
	// ErrContainerNotFound is returned by GetContainer if we were
	// unable to find the requested container.
	ErrContainerNotFound = errors.New("failed to locate the container")
)

// DockerClient provides a wrapper for the standard docker client. The intent
// is to wrap common operations so the internal of docker's own client are
// abstracted. Use NewClient() to construct and produce this struct.
type DockerClient struct {
	docker *client.Client
	ctx    context.Context
}

// NewClient produces a *DockerClient struct.
func NewClient(ctx context.Context) (*DockerClient, error) {
	docker, err := client.NewEnvClient()
	return &DockerClient{docker: docker, ctx: ctx}, err
}

// ContainerInfo retrieves a single container by id and returns a
// *ContainerInfo struct.
func (d *DockerClient) ContainerInfo(id string) (*ContainerInfo, error) {
	args := filters.NewArgs()
	args.Add("id", id)

	options := types.ContainerListOptions{Filters: args, All: true}
	containers, err := d.docker.ContainerList(d.ctx, options)
	if err != nil {
		return nil, err
	}

	if len(containers) == 0 {
		return nil, ErrContainerNotFound
	}

	inspection, err := d.docker.ContainerInspect(d.ctx, id)
	if err != nil {
		return nil, err
	}

	return &ContainerInfo{
		Data:     containers[0],
		State:    inspection.State,
		JSON:     inspection,
		Warnings: []string{},
		client:   d,
	}, nil
}

// ListContainers will return a list of *ContainerInfo structs based on the
// provided input.
func (d *DockerClient) ListContainers(input *ClientInput) ([]*ContainerInfo, error) {
	options := types.ContainerListOptions{
		All:     input.All,
		Since:   input.Since,
		Before:  input.Before,
		Filters: input.FilterArgs(),
	}

	listed, err := d.docker.ContainerList(d.ctx, options)
	if err != nil {
		return nil, err
	}

	containers := make(chan *ContainerInfo)
	errs := make(chan error)

	for _, entry := range listed {
		go func(c types.Container) {
			info, err := d.ContainerInfo(c.ID)
			if err != nil {
				errs <- err
				return
			}
			containers <- info
		}(entry)
	}

	results := []*ContainerInfo{}
	errout := errset.ErrSet{}
	for i := 0; i < len(listed); i++ {
		select {
		case err := <-errs:
			errout = append(errout, err)
		case info := <-containers:
			results = append(results, info)
		}
	}

	return results, errout.ReturnValue()
}

// RemoveContainer will delete the requested Container, force terminating
// it if necessary.
func (d *DockerClient) RemoveContainer(id string) error {
	err := d.docker.ContainerRemove(d.ctx, id, types.ContainerRemoveOptions{Force: true})

	// Docker's API does not expose their error structs and their
	// IsErrNotFound does not seem to work.
	if err != nil && strings.Contains(err.Error(), "No such container") {
		return nil
	}
	return err
}

// RunContainer will run a new c and return the results. By default
// all ports that are exposed by the c will be published to the host
// randomly. The published ports will be accessible using functions on the
// struct:
//    client, err := NewClient()
//    c := client.RunContainer("testimage", "testing", nil)
//    port, err := c.Port(80)
//    port.External
func (d *DockerClient) RunContainer(input *ClientInput) (*ContainerInfo, error) {
	bindings, err := input.Ports.Bindings()
	if err != nil {
		return nil, err
	}
	ctx, cancel := context.WithTimeout(d.ctx, DefaultServiceTimeout)
	defer cancel()

	if input.Timeout.Nanoseconds() > 0 {
		cancel()
		ctx, cancel = context.WithTimeout(d.ctx, input.Timeout)
	}

	for {
		created, err := d.docker.ContainerCreate(
			ctx,
			input.ContainerConfig(),
			&container.HostConfig{PortBindings: bindings}, &network.NetworkingConfig{}, "")
		if client.IsErrNotFound(err) {
			reader, err := d.docker.ImagePull(ctx, input.Image, types.ImagePullOptions{})
			if err != nil {
				return nil, err
			}
			if _, err := io.Copy(ioutil.Discard, reader); err != nil {
				return nil, err
			}
			continue
		}
		if err != nil {
			return nil, err
		}
		if err := d.docker.ContainerStart(ctx, created.ID, types.ContainerStartOptions{}); err != nil {
			return nil, err
		}
		info, err := d.ContainerInfo(created.ID)
		info.Warnings = created.Warnings
		return info, err
	}
}

// Service will return a *Service struct that may be used to spin up
// a specific service. See the documentation present on the Service struct
// for more information.
func (d *DockerClient) Service(input *ClientInput) *Service {
	timeout := input.Timeout
	if timeout.Nanoseconds() == 0 {
		timeout = DefaultServiceTimeout
	}
	ctx, _ := context.WithTimeout(d.ctx, timeout)
	return &Service{Context: ctx, Input: input, Client: d}
}
<commit_msg>minor move<commit_after>package dockertest

import (
	"context"
	"errors"
	"io"
	"io/ioutil"
	"strings"

	"github.com/crewjam/errset"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/client"
)

var (
	// ErrContainerNotFound is returned by GetContainer if we were
	// unable to find the requested container.
	ErrContainerNotFound = errors.New("failed to locate the container")
)

// DockerClient provides a wrapper for the standard docker client. The intent
// is to wrap common operations so the internal of docker's own client are
// abstracted. Use NewClient() to construct and produce this struct.
type DockerClient struct {
	docker *client.Client
	ctx    context.Context
}

// ContainerInfo retrieves a single container by id and returns a
// *ContainerInfo struct.
func (d *DockerClient) ContainerInfo(id string) (*ContainerInfo, error) {
	args := filters.NewArgs()
	args.Add("id", id)

	options := types.ContainerListOptions{Filters: args, All: true}
	containers, err := d.docker.ContainerList(d.ctx, options)
	if err != nil {
		return nil, err
	}

	if len(containers) == 0 {
		return nil, ErrContainerNotFound
	}

	inspection, err := d.docker.ContainerInspect(d.ctx, id)
	if err != nil {
		return nil, err
	}

	return &ContainerInfo{
		Data:     containers[0],
		State:    inspection.State,
		JSON:     inspection,
		Warnings: []string{},
		client:   d,
	}, nil
}

// ListContainers will return a list of *ContainerInfo structs based on the
// provided input.
func (d *DockerClient) ListContainers(input *ClientInput) ([]*ContainerInfo, error) {
	options := types.ContainerListOptions{
		All:     input.All,
		Since:   input.Since,
		Before:  input.Before,
		Filters: input.FilterArgs(),
	}

	listed, err := d.docker.ContainerList(d.ctx, options)
	if err != nil {
		return nil, err
	}

	containers := make(chan *ContainerInfo)
	errs := make(chan error)

	for _, entry := range listed {
		go func(c types.Container) {
			info, err := d.ContainerInfo(c.ID)
			if err != nil {
				errs <- err
				return
			}
			containers <- info
		}(entry)
	}

	results := []*ContainerInfo{}
	errout := errset.ErrSet{}
	for i := 0; i < len(listed); i++ {
		select {
		case err := <-errs:
			errout = append(errout, err)
		case info := <-containers:
			results = append(results, info)
		}
	}

	return results, errout.ReturnValue()
}

// RemoveContainer will delete the requested Container, force terminating
// it if necessary.
func (d *DockerClient) RemoveContainer(id string) error {
	err := d.docker.ContainerRemove(d.ctx, id, types.ContainerRemoveOptions{Force: true})

	// Docker's API does not expose their error structs and their
	// IsErrNotFound does not seem to work.
	if err != nil && strings.Contains(err.Error(), "No such container") {
		return nil
	}
	return err
}

// RunContainer will run a new c and return the results. By default
// all ports that are exposed by the c will be published to the host
// randomly. The published ports will be accessible using functions on the
// struct:
//    client, err := NewClient()
//    c := client.RunContainer("testimage", "testing", nil)
//    port, err := c.Port(80)
//    port.External
func (d *DockerClient) RunContainer(input *ClientInput) (*ContainerInfo, error) {
	bindings, err := input.Ports.Bindings()
	if err != nil {
		return nil, err
	}
	ctx, cancel := context.WithTimeout(d.ctx, DefaultServiceTimeout)
	defer cancel()

	if input.Timeout.Nanoseconds() > 0 {
		cancel()
		ctx, cancel = context.WithTimeout(d.ctx, input.Timeout)
	}

	for {
		created, err := d.docker.ContainerCreate(
			ctx,
			input.ContainerConfig(),
			&container.HostConfig{PortBindings: bindings}, &network.NetworkingConfig{}, "")
		if client.IsErrNotFound(err) {
			reader, err := d.docker.ImagePull(ctx, input.Image, types.ImagePullOptions{})
			if err != nil {
				return nil, err
			}
			if _, err := io.Copy(ioutil.Discard, reader); err != nil {
				return nil, err
			}
			continue
		}
		if err != nil {
			return nil, err
		}
		if err := d.docker.ContainerStart(ctx, created.ID, types.ContainerStartOptions{}); err != nil {
			return nil, err
		}
		info, err := d.ContainerInfo(created.ID)
		info.Warnings = created.Warnings
		return info, err
	}
}

// Service will return a *Service struct that may be used to spin up
// a specific service. See the documentation present on the Service struct
// for more information.
func (d *DockerClient) Service(input *ClientInput) *Service {
	timeout := input.Timeout
	if timeout.Nanoseconds() == 0 {
		timeout = DefaultServiceTimeout
	}
	ctx, _ := context.WithTimeout(d.ctx, timeout)
	return &Service{Context: ctx, Input: input, Client: d}
}

// NewClient produces a *DockerClient struct.
func NewClient(ctx context.Context) (*DockerClient, error) {
	docker, err := client.NewEnvClient()
	return &DockerClient{docker: docker, ctx: ctx}, err
}
<|endoftext|>


<commit_before>/*******************************************************************************
The MIT License (MIT)

Copyright (c) 2015 Hajime Nakagami

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*******************************************************************************/

package toybroker

import (
	"net"
	"sync"
)

type Client struct {
	clientID         string
	loginName        string
	conn             net.Conn
	currentMessageID uint16
	sync.RWMutex
}

func NewClient(id string, name string, c net.Conn, oldClient *Client) *Client {
	client := &Client{
		clientID:  id,
		loginName: name,
		conn:      c,
	}
	return client
}

func (c *Client) GetClientID() string {
	return c.clientID
}

func (c *Client) GetConn() net.Conn {
	return c.conn
}

func (c *Client) getNextMessageID() uint16 {
	c.Lock()
	defer c.Unlock()
	c.currentMessageID++
	if c.currentMessageID == 0 {
		c.currentMessageID++
	}
	return c.currentMessageID
}

func (c *Client) Publish(dup bool, qos int, topic string, payload []byte) {
	c.conn.Write(packPUBLISH(dup, qos, topic, c.getNextMessageID(), payload))
}

func (c *Client) Send(data []byte) {
	c.conn.Write(data)
}
<commit_msg>add GetLoginName() method<commit_after>/*******************************************************************************
The MIT License (MIT)

Copyright (c) 2015 Hajime Nakagami

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*******************************************************************************/

package toybroker

import (
	"net"
	"sync"
)

type Client struct {
	clientID         string
	loginName        string
	conn             net.Conn
	currentMessageID uint16
	sync.RWMutex
}

func NewClient(id string, name string, c net.Conn, oldClient *Client) *Client {
	client := &Client{
		clientID:  id,
		loginName: name,
		conn:      c,
	}
	return client
}

func (c *Client) GetClientID() string {
	return c.clientID
}

func (c *Client) GetLoginName() string {
	return c.loginName
}

func (c *Client) GetConn() net.Conn {
	return c.conn
}

func (c *Client) getNextMessageID() uint16 {
	c.Lock()
	defer c.Unlock()
	c.currentMessageID++
	if c.currentMessageID == 0 {
		c.currentMessageID++
	}
	return c.currentMessageID
}

func (c *Client) Publish(dup bool, qos int, topic string, payload []byte) {
	c.conn.Write(packPUBLISH(dup, qos, topic, c.getNextMessageID(), payload))
}

func (c *Client) Send(data []byte) {
	c.conn.Write(data)
}
<|endoftext|>


<commit_before>//
// Copyright (c) 2014 MessageBird B.V.
// All rights reserved.
//
// Author: Maurice Nonnekes <maurice@messagebird.com>

package messagebird

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"runtime"
	"strings"
)

const (
	ClientVersion = "2.2.0"
	Endpoint      = "https://rest.messagebird.com"
)

var (
	ErrResponse           = errors.New("The MessageBird API returned an error")
	ErrUnexpectedResponse = errors.New("The MessageBird API is currently unavailable")
)

type Client struct {
	AccessKey  string       // The API access key
	HTTPClient *http.Client // The HTTP client to send requests on
	DebugLog   *log.Logger  // Optional logger for debugging purposes
}

// New creates a new MessageBird client object.
func New(AccessKey string) *Client {
	return &Client{AccessKey: AccessKey, HTTPClient: &http.Client{}}
}

func (c *Client) request(v interface{}, path string, params *url.Values) error {
	uri, err := url.Parse(Endpoint + "/" + path)
	if err != nil {
		return err
	}

	var request *http.Request
	if params != nil {
		body := params.Encode()
		if request, err = http.NewRequest("POST", uri.String(), strings.NewReader(body)); err != nil {
			return err
		}

		if c.DebugLog != nil {
			if unescapedBody, err := url.QueryUnescape(body); err == nil {
				log.Printf("HTTP REQUEST: POST %s %s", uri.String(), unescapedBody)
			} else {
				log.Printf("HTTP REQUEST: POST %s %s", uri.String(), body)
			}
		}

		request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	} else {
		if request, err = http.NewRequest("GET", uri.String(), nil); err != nil {
			return err
		}

		if c.DebugLog != nil {
			log.Printf("HTTP REQUEST: GET %s", uri.String())
		}
	}

	request.Header.Add("Accept", "application/json")
	request.Header.Add("Authorization", "AccessKey "+c.AccessKey)
	request.Header.Add("User-Agent", "MessageBird/ApiClient/"+ClientVersion+" Go/"+runtime.Version())

	response, err := c.HTTPClient.Do(request)
	if err != nil {
		return err
	}

	defer response.Body.Close()

	responseBody, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}

	if c.DebugLog != nil {
		log.Printf("HTTP RESPONSE: %s", string(responseBody))
	}

	// Status code 500 is a server error and means nothing can be done at this
	// point.
	if response.StatusCode == 500 {
		return ErrUnexpectedResponse
	}

	if err = json.Unmarshal(responseBody, &v); err != nil {
		return err
	}

	// Status codes 200 and 201 are indicative of being able to convert the
	// response body to the struct that was specified.
	if response.StatusCode == 200 || response.StatusCode == 201 {
		return nil
	}

	// Anything else than a 200/201/500 should be a JSON error.
	return ErrResponse
}

// Balance returns the balance information for the account that is associated
// with the access key.
func (c *Client) Balance() (*Balance, error) {
	balance := &Balance{}
	if err := c.request(balance, "balance", nil); err != nil {
		if err == ErrResponse {
			return balance, err
		}

		return nil, err
	}

	return balance, nil
}

// HLR looks up an existing HLR object for the specified id that was previously
// created by the NewHLR function.
func (c *Client) HLR(id string) (*HLR, error) {
	hlr := &HLR{}
	if err := c.request(hlr, "hlr/"+id, nil); err != nil {
		if err == ErrResponse {
			return hlr, err
		}

		return nil, err
	}

	return hlr, nil
}

// NewHLR retrieves the information of an existing HLR.
func (c *Client) NewHLR(msisdn, reference string) (*HLR, error) {
	params := &url.Values{
		"msisdn":    {msisdn},
		"reference": {reference}}

	hlr := &HLR{}
	if err := c.request(hlr, "hlr", params); err != nil {
		if err == ErrResponse {
			return hlr, err
		}

		return nil, err
	}

	return hlr, nil
}

// Message retrieves the information of an existing Message.
func (c *Client) Message(id string) (*Message, error) {
	message := &Message{}
	if err := c.request(message, "messages/"+id, nil); err != nil {
		if err == ErrResponse {
			return message, err
		}

		return nil, err
	}

	return message, nil
}

// NewMessage creates a new message for one or more recipients.
func (c *Client) NewMessage(originator string, recipients []string, body string, msgParams *MessageParams) (*Message, error) {
	params, err := paramsForMessage(msgParams)
	if err != nil {
		return nil, err
	}

	params.Set("originator", originator)
	params.Set("body", body)
	params.Set("recipients", strings.Join(recipients, ","))

	message := &Message{}
	if err := c.request(message, "messages", params); err != nil {
		if err == ErrResponse {
			return message, err
		}

		return nil, err
	}

	return message, nil
}

// VoiceMessage retrieves the information of an existing VoiceMessage.
func (c *Client) VoiceMessage(id string) (*VoiceMessage, error) {
	message := &VoiceMessage{}
	if err := c.request(message, "voicemessages/"+id, nil); err != nil {
		if err == ErrResponse {
			return message, err
		}

		return nil, err
	}

	return message, nil
}

// NewVoiceMessage creates a new voice message for one or more recipients.
func (c *Client) NewVoiceMessage(recipients []string, body string, params *VoiceMessageParams) (*VoiceMessage, error) {
	urlParams := paramsForVoiceMessage(params)
	urlParams.Set("body", body)
	urlParams.Set("recipients", strings.Join(recipients, ","))

	message := &VoiceMessage{}
	if err := c.request(message, "voicemessages", urlParams); err != nil {
		if err == ErrResponse {
			return message, err
		}

		return nil, err
	}

	return message, nil
}

// OtpGenerate generates a new One-Time-Password for one recipient.
func (c *Client) OtpGenerate(recipient string, params *OtpParams) (*OtpMessage, error) {
	urlParams := paramsForOtp(params)
	urlParams.Set("recipient", recipient)

	message := &OtpMessage{}
	if err := c.request(message, "otp/generate", urlParams); err != nil {
		if err == ErrResponse {
			return message, err
		}

		return nil, err
	}

	return message, nil
}

// OtpVerify verifies the token that was generated with OtpGenerate.
func (c *Client) OtpVerify(recipient string, token string, params *OtpParams) (*OtpMessage, error) {
	urlParams := paramsForOtp(params)
	urlParams.Set("recipient", recipient)
	urlParams.Set("token", token)

	path := "otp/verify?" + urlParams.Encode()

	message := &OtpMessage{}
	if err := c.request(message, path, nil); err != nil {
		if err == ErrResponse {
			return message, err
		}

		return nil, err
	}

	return message, nil
}

// Lookup performs a new lookup for the specified number.
func (c *Client) Lookup(phoneNumber string, params *LookupParams) (*Lookup, error) {
	urlParams := paramsForLookup(params)
	path := "lookup/" + phoneNumber + "?" + urlParams.Encode()

	lookup := &Lookup{}
	if err := c.request(lookup, path, nil); err != nil {
		if err == ErrResponse {
			return lookup, err
		}

		return nil, err
	}

	return lookup, nil
}

// NewLookupHLR creates a new HLR lookup for the specified number.
func (c *Client) NewLookupHLR(phoneNumber string, params *LookupParams) (*HLR, error) {
	urlParams := paramsForLookup(params)
	path := "lookup/" + phoneNumber + "/hlr"

	hlr := &HLR{}
	if err := c.request(hlr, path, urlParams); err != nil {
		if err == ErrResponse {
			return hlr, err
		}

		return nil, err
	}

	return hlr, nil
}

// LookupHLR performs a HLR lookup for the specified number.
func (c *Client) LookupHLR(phoneNumber string, params *LookupParams) (*HLR, error) {
	urlParams := paramsForLookup(params)
	path := "lookup/" + phoneNumber + "/hlr?" + urlParams.Encode()

	hlr := &HLR{}
	if err := c.request(hlr, path, nil); err != nil {
		if err == ErrResponse {
			return hlr, err
		}

		return nil, err
	}

	return hlr, nil
}
<commit_msg>bump version to 2.2.1<commit_after>//
// Copyright (c) 2014 MessageBird B.V.
// All rights reserved.
//
// Author: Maurice Nonnekes <maurice@messagebird.com>

package messagebird

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"runtime"
	"strings"
)

const (
	ClientVersion = "2.2.1"
	Endpoint      = "https://rest.messagebird.com"
)

var (
	ErrResponse           = errors.New("The MessageBird API returned an error")
	ErrUnexpectedResponse = errors.New("The MessageBird API is currently unavailable")
)

type Client struct {
	AccessKey  string       // The API access key
	HTTPClient *http.Client // The HTTP client to send requests on
	DebugLog   *log.Logger  // Optional logger for debugging purposes
}

// New creates a new MessageBird client object.
func New(AccessKey string) *Client {
	return &Client{AccessKey: AccessKey, HTTPClient: &http.Client{}}
}

func (c *Client) request(v interface{}, path string, params *url.Values) error {
	uri, err := url.Parse(Endpoint + "/" + path)
	if err != nil {
		return err
	}

	var request *http.Request
	if params != nil {
		body := params.Encode()
		if request, err = http.NewRequest("POST", uri.String(), strings.NewReader(body)); err != nil {
			return err
		}

		if c.DebugLog != nil {
			if unescapedBody, err := url.QueryUnescape(body); err == nil {
				log.Printf("HTTP REQUEST: POST %s %s", uri.String(), unescapedBody)
			} else {
				log.Printf("HTTP REQUEST: POST %s %s", uri.String(), body)
			}
		}

		request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	} else {
		if request, err = http.NewRequest("GET", uri.String(), nil); err != nil {
			return err
		}

		if c.DebugLog != nil {
			log.Printf("HTTP REQUEST: GET %s", uri.String())
		}
	}

	request.Header.Add("Accept", "application/json")
	request.Header.Add("Authorization", "AccessKey "+c.AccessKey)
	request.Header.Add("User-Agent", "MessageBird/ApiClient/"+ClientVersion+" Go/"+runtime.Version())

	response, err := c.HTTPClient.Do(request)
	if err != nil {
		return err
	}

	defer response.Body.Close()

	responseBody, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}

	if c.DebugLog != nil {
		log.Printf("HTTP RESPONSE: %s", string(responseBody))
	}

	// Status code 500 is a server error and means nothing can be done at this
	// point.
	if response.StatusCode == 500 {
		return ErrUnexpectedResponse
	}

	if err = json.Unmarshal(responseBody, &v); err != nil {
		return err
	}

	// Status codes 200 and 201 are indicative of being able to convert the
	// response body to the struct that was specified.
	if response.StatusCode == 200 || response.StatusCode == 201 {
		return nil
	}

	// Anything else than a 200/201/500 should be a JSON error.
	return ErrResponse
}

// Balance returns the balance information for the account that is associated
// with the access key.
func (c *Client) Balance() (*Balance, error) {
	balance := &Balance{}
	if err := c.request(balance, "balance", nil); err != nil {
		if err == ErrResponse {
			return balance, err
		}

		return nil, err
	}

	return balance, nil
}

// HLR looks up an existing HLR object for the specified id that was previously
// created by the NewHLR function.
func (c *Client) HLR(id string) (*HLR, error) {
	hlr := &HLR{}
	if err := c.request(hlr, "hlr/"+id, nil); err != nil {
		if err == ErrResponse {
			return hlr, err
		}

		return nil, err
	}

	return hlr, nil
}

// NewHLR retrieves the information of an existing HLR.
func (c *Client) NewHLR(msisdn, reference string) (*HLR, error) {
	params := &url.Values{
		"msisdn":    {msisdn},
		"reference": {reference}}

	hlr := &HLR{}
	if err := c.request(hlr, "hlr", params); err != nil {
		if err == ErrResponse {
			return hlr, err
		}

		return nil, err
	}

	return hlr, nil
}

// Message retrieves the information of an existing Message.
func (c *Client) Message(id string) (*Message, error) {
	message := &Message{}
	if err := c.request(message, "messages/"+id, nil); err != nil {
		if err == ErrResponse {
			return message, err
		}

		return nil, err
	}

	return message, nil
}

// NewMessage creates a new message for one or more recipients.
func (c *Client) NewMessage(originator string, recipients []string, body string, msgParams *MessageParams) (*Message, error) {
	params, err := paramsForMessage(msgParams)
	if err != nil {
		return nil, err
	}

	params.Set("originator", originator)
	params.Set("body", body)
	params.Set("recipients", strings.Join(recipients, ","))

	message := &Message{}
	if err := c.request(message, "messages", params); err != nil {
		if err == ErrResponse {
			return message, err
		}

		return nil, err
	}

	return message, nil
}

// VoiceMessage retrieves the information of an existing VoiceMessage.
func (c *Client) VoiceMessage(id string) (*VoiceMessage, error) {
	message := &VoiceMessage{}
	if err := c.request(message, "voicemessages/"+id, nil); err != nil {
		if err == ErrResponse {
			return message, err
		}

		return nil, err
	}

	return message, nil
}

// NewVoiceMessage creates a new voice message for one or more recipients.
func (c *Client) NewVoiceMessage(recipients []string, body string, params *VoiceMessageParams) (*VoiceMessage, error) {
	urlParams := paramsForVoiceMessage(params)
	urlParams.Set("body", body)
	urlParams.Set("recipients", strings.Join(recipients, ","))

	message := &VoiceMessage{}
	if err := c.request(message, "voicemessages", urlParams); err != nil {
		if err == ErrResponse {
			return message, err
		}

		return nil, err
	}

	return message, nil
}

// OtpGenerate generates a new One-Time-Password for one recipient.
func (c *Client) OtpGenerate(recipient string, params *OtpParams) (*OtpMessage, error) {
	urlParams := paramsForOtp(params)
	urlParams.Set("recipient", recipient)

	message := &OtpMessage{}
	if err := c.request(message, "otp/generate", urlParams); err != nil {
		if err == ErrResponse {
			return message, err
		}

		return nil, err
	}

	return message, nil
}

// OtpVerify verifies the token that was generated with OtpGenerate.
func (c *Client) OtpVerify(recipient string, token string, params *OtpParams) (*OtpMessage, error) {
	urlParams := paramsForOtp(params)
	urlParams.Set("recipient", recipient)
	urlParams.Set("token", token)

	path := "otp/verify?" + urlParams.Encode()

	message := &OtpMessage{}
	if err := c.request(message, path, nil); err != nil {
		if err == ErrResponse {
			return message, err
		}

		return nil, err
	}

	return message, nil
}

// Lookup performs a new lookup for the specified number.
func (c *Client) Lookup(phoneNumber string, params *LookupParams) (*Lookup, error) {
	urlParams := paramsForLookup(params)
	path := "lookup/" + phoneNumber + "?" + urlParams.Encode()

	lookup := &Lookup{}
	if err := c.request(lookup, path, nil); err != nil {
		if err == ErrResponse {
			return lookup, err
		}

		return nil, err
	}

	return lookup, nil
}

// NewLookupHLR creates a new HLR lookup for the specified number.
func (c *Client) NewLookupHLR(phoneNumber string, params *LookupParams) (*HLR, error) {
	urlParams := paramsForLookup(params)
	path := "lookup/" + phoneNumber + "/hlr"

	hlr := &HLR{}
	if err := c.request(hlr, path, urlParams); err != nil {
		if err == ErrResponse {
			return hlr, err
		}

		return nil, err
	}

	return hlr, nil
}

// LookupHLR performs a HLR lookup for the specified number.
func (c *Client) LookupHLR(phoneNumber string, params *LookupParams) (*HLR, error) {
	urlParams := paramsForLookup(params)
	path := "lookup/" + phoneNumber + "/hlr?" + urlParams.Encode()

	hlr := &HLR{}
	if err := c.request(hlr, path, nil); err != nil {
		if err == ErrResponse {
			return hlr, err
		}

		return nil, err
	}

	return hlr, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"go/build"
	"io/ioutil"
	"log"
	"net/rpc"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"github.com/mdempsky/gocode/gbimporter"
	"github.com/mdempsky/gocode/suggest"
)

func doClient() {
	if *g_debug {
		start := time.Now()
		defer func() {
			elapsed := time.Since(start)
			log.Printf("Elapsed duration: %v\n", elapsed)
		}()
	}

	var command string
	if flag.NArg() > 0 {
		command = flag.Arg(0)
		switch command {
		case "autocomplete", "exit":
			// these are valid commands
		case "close":
			// "close" is an alias for "exit"
			command = "exit"
		default:
			fmt.Printf("gocode: unknown subcommand: %q\nRun 'gocode -help' for usage.\n", command)
			os.Exit(2)
		}
	}

	// client
	var client *rpc.Client
	if *g_sock != "none" {
		addr := *g_addr
		if *g_sock == "unix" {
			addr = getSocketPath()
		}

		var err error
		client, err = rpc.Dial(*g_sock, addr)
		if err != nil {
			if command == "exit" {
				log.Fatal(err)
			}

			if *g_sock == "unix" {
				_ = os.Remove(addr)
			}
			err = tryStartServer()
			if err != nil {
				log.Fatalf("Failed to start server: %s\n", err)
			}
			client, err = tryToConnect(*g_sock, addr)
			if err != nil {
				log.Fatalf("Failed to connect to %q: %s\n", addr, err)
			}
		}
		defer client.Close()
	}

	switch command {
	case "autocomplete":
		cmdAutoComplete(client)
	case "exit":
		cmdExit(client)
	}
}

func tryStartServer() error {
	path := get_executable_filename()
	args := []string{os.Args[0], "-s", "-sock", *g_sock, "-addr", *g_addr}
	cwd, _ := os.Getwd()

	var err error
	stdin, err := os.Open(os.DevNull)
	if err != nil {
		return err
	}
	stdout, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0)
	if err != nil {
		return err
	}
	stderr, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0)
	if err != nil {
		return err
	}

	procattr := os.ProcAttr{Dir: cwd, Env: os.Environ(), Files: []*os.File{stdin, stdout, stderr}}
	p, err := os.StartProcess(path, args, &procattr)
	if err != nil {
		return err
	}

	return p.Release()
}

func tryToConnect(network, address string) (*rpc.Client, error) {
	start := time.Now()
	for {
		client, err := rpc.Dial(network, address)
		if err != nil && time.Since(start) < time.Second {
			continue
		}
		return client, err
	}
}

func cmdAutoComplete(c *rpc.Client) {
	var req AutoCompleteRequest
	req.Filename, req.Data, req.Cursor = prepareFilenameDataCursor()
	req.Context = gbimporter.PackContext(&build.Default)

	var res AutoCompleteReply
	var err error
	if c == nil {
		s := Server{}
		err = s.AutoComplete(&req, &res)
	} else {
		err = c.Call("Server.AutoComplete", &req, &res)
	}
	if err != nil {
		log.Fatal(err)
	}

	fmt := suggest.Formatters[*g_format]
	if fmt == nil {
		fmt = suggest.NiceFormat
	}
	fmt(os.Stdout, res.Candidates, res.Len)
}

func cmdExit(c *rpc.Client) {
	if c == nil {
		return
	}
	var req ExitRequest
	var res ExitReply
	if err := c.Call("Server.Exit", &req, &res); err != nil {
		log.Fatal(err)
	}
}

func prepareFilenameDataCursor() (string, []byte, int) {
	var file []byte
	var err error

	if *g_input != "" {
		file, err = ioutil.ReadFile(*g_input)
	} else {
		file, err = ioutil.ReadAll(os.Stdin)
	}

	if err != nil {
		log.Fatal(err)
	}

	filename := *g_input
	offset := ""
	switch flag.NArg() {
	case 2:
		offset = flag.Arg(1)
	case 3:
		filename = flag.Arg(1) // Override default filename
		offset = flag.Arg(2)
	}

	if filename != "" {
		filename, _ = filepath.Abs(filename)
	}

	cursor := -1
	if offset != "" {
		if offset[0] == 'c' || offset[0] == 'C' {
			cursor, _ = strconv.Atoi(offset[1:])
			cursor = runeToByteOffset(file, cursor)
		} else {
			cursor, _ = strconv.Atoi(offset)
		}
	}

	return filename, file, cursor
}
<commit_msg>gocode: speed up client by disabling GC<commit_after>package main

import (
	"flag"
	"fmt"
	"go/build"
	"io/ioutil"
	"log"
	"net/rpc"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"runtime/debug"

	"github.com/mdempsky/gocode/gbimporter"
	"github.com/mdempsky/gocode/suggest"
)

func doClient() {
	// Client is a short-lived program.
	// Disable GC to make it faster
	debug.SetGCPercent(-1)

	if *g_debug {
		start := time.Now()
		defer func() {
			elapsed := time.Since(start)
			log.Printf("Elapsed duration: %v\n", elapsed)
		}()
	}

	var command string
	if flag.NArg() > 0 {
		command = flag.Arg(0)
		switch command {
		case "autocomplete", "exit":
			// these are valid commands
		case "close":
			// "close" is an alias for "exit"
			command = "exit"
		default:
			fmt.Printf("gocode: unknown subcommand: %q\nRun 'gocode -help' for usage.\n", command)
			os.Exit(2)
		}
	}

	// client
	var client *rpc.Client
	if *g_sock != "none" {
		addr := *g_addr
		if *g_sock == "unix" {
			addr = getSocketPath()
		}

		var err error
		client, err = rpc.Dial(*g_sock, addr)
		if err != nil {
			if command == "exit" {
				log.Fatal(err)
			}

			if *g_sock == "unix" {
				_ = os.Remove(addr)
			}
			err = tryStartServer()
			if err != nil {
				log.Fatalf("Failed to start server: %s\n", err)
			}
			client, err = tryToConnect(*g_sock, addr)
			if err != nil {
				log.Fatalf("Failed to connect to %q: %s\n", addr, err)
			}
		}
		defer client.Close()
	}

	switch command {
	case "autocomplete":
		cmdAutoComplete(client)
	case "exit":
		cmdExit(client)
	}
}

func tryStartServer() error {
	path := get_executable_filename()
	args := []string{os.Args[0], "-s", "-sock", *g_sock, "-addr", *g_addr}
	cwd, _ := os.Getwd()

	var err error
	stdin, err := os.Open(os.DevNull)
	if err != nil {
		return err
	}
	stdout, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0)
	if err != nil {
		return err
	}
	stderr, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0)
	if err != nil {
		return err
	}

	procattr := os.ProcAttr{Dir: cwd, Env: os.Environ(), Files: []*os.File{stdin, stdout, stderr}}
	p, err := os.StartProcess(path, args, &procattr)
	if err != nil {
		return err
	}

	return p.Release()
}

func tryToConnect(network, address string) (*rpc.Client, error) {
	start := time.Now()
	for {
		client, err := rpc.Dial(network, address)
		if err != nil && time.Since(start) < time.Second {
			continue
		}
		return client, err
	}
}

func cmdAutoComplete(c *rpc.Client) {
	var req AutoCompleteRequest
	req.Filename, req.Data, req.Cursor = prepareFilenameDataCursor()
	req.Context = gbimporter.PackContext(&build.Default)

	var res AutoCompleteReply
	var err error
	if c == nil {
		s := Server{}
		err = s.AutoComplete(&req, &res)
	} else {
		err = c.Call("Server.AutoComplete", &req, &res)
	}
	if err != nil {
		log.Fatal(err)
	}

	fmt := suggest.Formatters[*g_format]
	if fmt == nil {
		fmt = suggest.NiceFormat
	}
	fmt(os.Stdout, res.Candidates, res.Len)
}

func cmdExit(c *rpc.Client) {
	if c == nil {
		return
	}
	var req ExitRequest
	var res ExitReply
	if err := c.Call("Server.Exit", &req, &res); err != nil {
		log.Fatal(err)
	}
}

func prepareFilenameDataCursor() (string, []byte, int) {
	var file []byte
	var err error

	if *g_input != "" {
		file, err = ioutil.ReadFile(*g_input)
	} else {
		file, err = ioutil.ReadAll(os.Stdin)
	}

	if err != nil {
		log.Fatal(err)
	}

	filename := *g_input
	offset := ""
	switch flag.NArg() {
	case 2:
		offset = flag.Arg(1)
	case 3:
		filename = flag.Arg(1) // Override default filename
		offset = flag.Arg(2)
	}

	if filename != "" {
		filename, _ = filepath.Abs(filename)
	}

	cursor := -1
	if offset != "" {
		if offset[0] == 'c' || offset[0] == 'C' {
			cursor, _ = strconv.Atoi(offset[1:])
			cursor = runeToByteOffset(file, cursor)
		} else {
			cursor, _ = strconv.Atoi(offset)
		}
	}

	return filename, file, cursor
}
<|endoftext|>


<commit_before>package atlas // import "github.com/keltia/ripe-atlas"

import (
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"time"

	"github.com/keltia/proxy"
	"github.com/pkg/errors"
)

// NewClient is the first function to call.
// Yes, it does take multiple config
// and the last one wins.
func NewClient(cfgs ...Config) (*Client, error) {
	c := &Client{}
	for _, cfg := range cfgs {
		c.config = cfg
	}

	// This holds the global options
	c.opts = make(map[string]string)

	// If no log output is specified, use the default one
	if c.config.Log == nil {
		c.log = log.New(os.Stderr, "", log.LstdFlags|log.LUTC)
	} else {
		c.log = c.config.Log
	}

	// Set log levels
	if c.config.Verbose == true {
		c.level = 1
	}

	if c.config.Level > 2 {
		c.level = 2
	}

	// Ensure this is not empty
	if c.config.endpoint == "" {
		c.config.endpoint = apiEndpoint
	}
	c.verbose("c.config=%#v", c.config)

	// Create and save the http.Client
	return c.addHTTPClient()
}

// HasAPIKey returns whether an API key is stored
func (c *Client) HasAPIKey() (string, bool) {
	if c.config.APIKey == "" {
		return "", false
	}
	return c.config.APIKey, true
}

// call is s shortcut
func (c *Client) call(req *http.Request) (*http.Response, error) {
	c.verbose("Full URL:\n%v", req.URL)

	myurl, _ := url.Parse(apiEndpoint)
	req.Header.Set("Host", myurl.Host)
	req.Header.Set("User-Agent", fmt.Sprintf("ripe-atlas/%s", ourVersion))

	return c.client.Do(req)
}

func (c *Client) addHTTPClient() (*Client, error) {
	_, transport := proxy.SetupTransport(apiEndpoint)
	if transport == nil {
		return c, errors.New("addhttpclient")
	}
	c.client = &http.Client{Transport: transport, Timeout: 20 * time.Second}
	return c, nil
}

// SetOption sets a global option
func (c *Client) SetOption(name, value string) *Client {
	if value != "" {
		c.opts[name] = value
	}
	return c
}

// GetVersion returns the API wrapper version
func GetVersion() string {
	return ourVersion
}
<commit_msg>Pass NewClient3() and NewClient4().<commit_after>package atlas // import "github.com/keltia/ripe-atlas"

import (
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"time"

	"github.com/keltia/proxy"
	"github.com/pkg/errors"
)

// NewClient is the first function to call.
// Yes, it does take multiple config
// and the last one wins.
func NewClient(cfgs ...Config) (*Client, error) {
	c := &Client{}
	for _, cfg := range cfgs {
		c.config = cfg
	}

	// This holds the global options
	c.opts = make(map[string]string)

	// If no log output is specified, use the default one
	if c.config.Log == nil {
		c.log = log.New(os.Stderr, "", log.LstdFlags|log.LUTC)
	} else {
		c.log = c.config.Log
	}

	// Set log levels
	if c.config.Verbose == true {
		c.level = 1
	}

	if c.config.Level != 0 {
		c.level = c.config.Level
	}

	// Final check
	if c.config.Level > 2 {
		c.level = 2
	}

	// Ensure this is not empty
	if c.config.endpoint == "" {
		c.config.endpoint = apiEndpoint
	}
	c.verbose("c.config=%#v", c.config)

	// Create and save the http.Client
	return c.addHTTPClient()
}

// HasAPIKey returns whether an API key is stored
func (c *Client) HasAPIKey() (string, bool) {
	if c.config.APIKey == "" {
		return "", false
	}
	return c.config.APIKey, true
}

// call is s shortcut
func (c *Client) call(req *http.Request) (*http.Response, error) {
	c.verbose("Full URL:\n%v", req.URL)

	myurl, _ := url.Parse(apiEndpoint)
	req.Header.Set("Host", myurl.Host)
	req.Header.Set("User-Agent", fmt.Sprintf("ripe-atlas/%s", ourVersion))

	return c.client.Do(req)
}

func (c *Client) addHTTPClient() (*Client, error) {
	_, transport := proxy.SetupTransport(apiEndpoint)
	if transport == nil {
		return c, errors.New("addhttpclient")
	}
	c.client = &http.Client{Transport: transport, Timeout: 20 * time.Second}
	return c, nil
}

// SetOption sets a global option
func (c *Client) SetOption(name, value string) *Client {
	if value != "" {
		c.opts[name] = value
	}
	return c
}

// GetVersion returns the API wrapper version
func GetVersion() string {
	return ourVersion
}
<|endoftext|>