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

package main

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"time"

	"github.com/spf13/cobra"
)

var (
	killCmd = &cobra.Command{
		Use:     "kill",
		Aliases: []string{"stop", "halt"},
		Short:   "Halts one or more running CoreOS instances",
		PreRunE: func(cmd *cobra.Command, args []string) (err error) {
			vipre.BindPFlags(cmd.Flags())
			if len(args) < 1 && !vipre.GetBool("all") {
				return fmt.Errorf("This command requires either at least " +
					"one argument to work or --all.")
			}
			return nil
		},
		RunE: killCommand,
	}
)

func killCommand(cmd *cobra.Command, args []string) (err error) {
	var up []VMInfo
	if up, err = allRunningInstances(); err != nil {
		return
	}
	if vipre.GetBool("all") {
		for _, vm := range up {
			if err = vm.halt(); err != nil {
				return err
			}
		}
		return
	}
	for _, arg := range args {
		for _, vm := range up {
			if vm.Name == arg || vm.UUID == arg {
				if err = vm.halt(); err != nil {
					return err
				}
			}
		}
	}
	return
}

func (vm VMInfo) halt() (err error) {
	var sshSession *sshClient

	if sshSession, err = vm.startSSHsession(); err != nil {
		return
	}
	defer sshSession.close()
	if e := sshSession.executeRemoteCommand("sudo sync;sudo halt"); e != nil {
		// ssh messed up for some reason or target has no IP
		log.Printf("couldn't ssh to %v (%v)...\n", vm.Name, e)
		if canKill := SessionContext.allowedToRun(); canKill != nil {
			return canKill
		}
		if p, ee := os.FindProcess(vm.Pid); ee == nil {
			log.Println("hard kill...")
			if err = p.Kill(); err != nil {
				return
			}
		}
	}
	// wait until it's _really_ dead
	defer func() {
		leftover := filepath.Join(SessionContext.runDir, vm.UUID)

		for range time.Tick(500 * time.Millisecond) {
			if _, ee := os.FindProcess(vm.Pid); ee == nil {
				break
			}
		}
		if e := os.RemoveAll(leftover); e != nil {
			log.Println(e.Error())
		}
		log.Printf("successfully halted '%s'\n", vm.Name)
	}()

	return
}

func init() {
	killCmd.Flags().BoolP("all", "a", false, "halts all running instances")
	RootCmd.AddCommand(killCmd)
}
<commit_msg>fix halt logic<commit_after>// Copyright 2015 - António Meireles  <antonio.meireles@reformi.st>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

package main

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"time"

	"github.com/spf13/cobra"
)

var (
	killCmd = &cobra.Command{
		Use:     "kill",
		Aliases: []string{"stop", "halt"},
		Short:   "Halts one or more running CoreOS instances",
		PreRunE: func(cmd *cobra.Command, args []string) (err error) {
			vipre.BindPFlags(cmd.Flags())
			if len(args) < 1 && !vipre.GetBool("all") {
				return fmt.Errorf("This command requires either at least " +
					"one argument to work or --all.")
			}
			return nil
		},
		RunE: killCommand,
	}
)

func killCommand(cmd *cobra.Command, args []string) (err error) {
	var up []VMInfo
	if up, err = allRunningInstances(); err != nil {
		return
	}
	if vipre.GetBool("all") {
		for _, vm := range up {
			if err = vm.halt(); err != nil {
				return err
			}
		}
		return
	}
	for _, arg := range args {
		for _, vm := range up {
			if vm.Name == arg || vm.UUID == arg {
				if err = vm.halt(); err != nil {
					return err
				}
			}
		}
	}
	return
}

func (vm VMInfo) halt() (err error) {
	var sshSession *sshClient

	if sshSession, err = vm.startSSHsession(); err != nil {
		return
	}
	defer sshSession.close()
	if e := sshSession.executeRemoteCommand("sudo sync;sudo halt"); e != nil {
		// ssh messed up for some reason or target has no IP
		log.Printf("couldn't ssh to %v (%v)...\n", vm.Name, e)
		if canKill := SessionContext.allowedToRun(); canKill != nil {
			return canKill
		}
		if p, ee := os.FindProcess(vm.Pid); ee == nil {
			log.Println("hard kill...")
			if err = p.Kill(); err != nil {
				return
			}
		}
	}
	// wait until it's _really_ dead, but not forever
	for {
		select {
		case <-time.After(3 * time.Second):
			return fmt.Errorf("VM didn't shutdown normally after 3s (!)... ")
		case <-time.Tick(100 * time.Millisecond):
			if _, ee := os.FindProcess(vm.Pid); ee == nil {
				if e :=
					os.RemoveAll(filepath.Join(SessionContext.runDir,
						vm.UUID)); e != nil {
					log.Println(e.Error())
				}
				log.Printf("successfully halted '%s'\n", vm.Name)
				return
			}
		}
	}
}

func init() {
	killCmd.Flags().BoolP("all", "a", false, "halts all running instances")
	RootCmd.AddCommand(killCmd)
}
<|endoftext|>


<commit_before>package httphead

import (
	"bufio"
	"bytes"
)

// Version contains protocol major and minor version.
type Version struct {
	Major int
	Minor int
}

// RequestLine contains parameters parsed from the first request line.
type RequestLine struct {
	Method  []byte
	URI     []byte
	Version Version
}

// ResponseLine contains parameters parsed from the first response line.
type ResponseLine struct {
	Version Version
	Status  int
	Reason  []byte
}

// SplitRequestLine splits given slice of bytes into three chunks without
// parsing.
func SplitRequestLine(line []byte) (method, uri, version []byte) {
	return split3(line, ' ')
}

// ParseRequestLine parses http request line like "GET / HTTP/1.0".
func ParseRequestLine(line []byte) (r RequestLine, ok bool) {
	var i int
	for i = 0; i < len(line); i++ {
		c := line[i]
		if !OctetTypes[c].IsToken() {
			if i > 0 && c == ' ' {
				break
			}
			return
		}
	}
	if i == len(line) {
		return
	}

	var proto []byte
	r.Method = line[:i]
	r.URI, proto = split2(line[i+1:], ' ')
	if len(r.URI) == 0 {
		return
	}
	if major, minor, ok := ParseVersion(proto); ok {
		r.Version.Major = major
		r.Version.Minor = minor
		return r, true
	}

	return r, false
}

// SplitResponseLine splits given slice of bytes into three chunks without
// parsing.
func SplitResponseLine(line []byte) (version, status, reason []byte) {
	return split3(line, ' ')
}

// ParseResponseLine parses first response line into ResponseLine struct.
func ParseResponseLine(line []byte) (r ResponseLine, ok bool) {
	var (
		proto  []byte
		status []byte
	)
	proto, status, r.Reason = split3(line, ' ')
	if major, minor, ok := ParseVersion(proto); ok {
		r.Version.Major = major
		r.Version.Minor = minor
	} else {
		return r, false
	}
	if n, ok := IntFromASCII(status); ok {
		r.Status = n
	} else {
		return r, false
	}
	// TODO(gobwas): parse here r.Reason fot TEXT rule:
	//   TEXT = <any OCTET except CTLs,
	//           but including LWS>
	return r, true
}

var (
	httpVersion10     = []byte("HTTP/1.0")
	httpVersion11     = []byte("HTTP/1.1")
	httpVersionPrefix = []byte("HTTP/")
)

// ParseVersion parses major and minor version of HTTP protocol.
// It returns parsed values and true if parse is ok.
func ParseVersion(bts []byte) (major, minor int, ok bool) {
	switch {
	case bytes.Equal(bts, httpVersion11):
		return 1, 1, true
	case bytes.Equal(bts, httpVersion10):
		return 1, 0, true
	case len(bts) < 8:
		return
	case !bytes.Equal(bts[:5], httpVersionPrefix):
		return
	}

	bts = bts[5:]

	dot := bytes.IndexByte(bts, '.')
	if dot == -1 {
		return
	}
	major, ok = IntFromASCII(bts[:dot])
	if !ok {
		return
	}
	minor, ok = IntFromASCII(bts[dot+1:])
	if !ok {
		return
	}

	return major, minor, true
}

// ReadLine reads line from br. It reads until '\n' and returns bytes without
// '\n' or '\r\n' at the end.
// It returns err if and only if line does not end in '\n'. Note that read
// bytes returned in any case of error.
//
// It is much like the textproto/Reader.ReadLine() except the thing that it
// returns raw bytes, instead of string. That is, it avoids copying bytes read
// from br.
//
// textproto/Reader.ReadLineBytes() is also makes copy of resulting bytes to be
// safe with future I/O operations on br.
//
// We could control I/O operations on br and do not need to make additional
// copy for safety.
func ReadLine(br *bufio.Reader) ([]byte, error) {
	var line []byte
	for {
		bts, err := br.ReadSlice('\n')
		if err == bufio.ErrBufferFull {
			// Copy bytes because next read will discard them.
			line = append(line, bts...)
			continue
		}
		// Avoid copy of single read.
		if line == nil {
			line = bts
		} else {
			line = append(line, bts...)
		}
		if err != nil {
			return line, err
		}
		// Size of line is at least 1.
		// In other case bufio.ReadSlice() returns error.
		n := len(line)
		// Cut '\n' or '\r\n'.
		if n > 1 && line[n-2] == '\r' {
			line = line[:n-2]
		} else {
			line = line[:n-1]
		}
		return line, nil
	}
}

// ParseHeaderLine parses HTTP header as key-value pair. It returns parsed
// values and true if parse is ok.
func ParseHeaderLine(line []byte) (k, v []byte, ok bool) {
	colon := bytes.IndexByte(line, ':')
	if colon == -1 {
		return
	}
	k = trim(line[:colon])
	for _, c := range k {
		if !OctetTypes[c].IsToken() {
			return nil, nil, false
		}
	}
	v = trim(line[colon+1:])
	return k, v, true
}

// IntFromASCII converts bytes met insidie HTTP heads to int.
func IntFromASCII(bts []byte) (ret int, ok bool) {
	// ASCII numbers all start with the high-order bits 0011.
	// If you see that, and the next bits are 0-9 (0000 - 1001) you can grab those
	// bits and interpret them directly as an integer.
	var n int
	if n = len(bts); n < 1 {
		return 0, false
	}
	for i := 0; i < n; i++ {
		if bts[i]&0xf0 != 0x30 {
			return 0, false
		}
		ret += int(bts[i]&0xf) * pow(10, n-i-1)
	}
	return ret, true
}

// pow for integers implementation.
// See Donald Knuth, The Art of Computer Programming, Volume 2, Section 4.6.3
func pow(a, b int) int {
	p := 1
	for b > 0 {
		if b&1 != 0 {
			p *= a
		}
		b >>= 1
		a *= a
	}
	return p
}

func split3(p []byte, sep byte) (p1, p2, p3 []byte) {
	a := bytes.IndexByte(p, sep)
	b := bytes.IndexByte(p[a+1:], sep)
	if a == -1 || b == -1 {
		return p, nil, nil
	}
	b += a + 1
	return p[:a], p[a+1 : b], p[b+1:]
}

func split2(p []byte, sep byte) (p1, p2 []byte) {
	i := bytes.IndexByte(p, sep)
	if i == -1 {
		return p, nil
	}
	return p[:i], p[i+1:]
}

func trim(p []byte) []byte {
	var i, j int
	for i = 0; i < len(p) && (p[i] == ' ' || p[i] == '\t'); {
		i++
	}
	for j = len(p); j > i && (p[j-1] == ' ' || p[j-1] == '\t'); {
		j--
	}
	return p[i:j]
}
<commit_msg>comments<commit_after>package httphead

import (
	"bufio"
	"bytes"
)

// Version contains protocol major and minor version.
type Version struct {
	Major int
	Minor int
}

// RequestLine contains parameters parsed from the first request line.
type RequestLine struct {
	Method  []byte
	URI     []byte
	Version Version
}

// ResponseLine contains parameters parsed from the first response line.
type ResponseLine struct {
	Version Version
	Status  int
	Reason  []byte
}

// SplitRequestLine splits given slice of bytes into three chunks without
// parsing.
func SplitRequestLine(line []byte) (method, uri, version []byte) {
	return split3(line, ' ')
}

// ParseRequestLine parses http request line like "GET / HTTP/1.0".
func ParseRequestLine(line []byte) (r RequestLine, ok bool) {
	var i int
	for i = 0; i < len(line); i++ {
		c := line[i]
		if !OctetTypes[c].IsToken() {
			if i > 0 && c == ' ' {
				break
			}
			return
		}
	}
	if i == len(line) {
		return
	}

	var proto []byte
	r.Method = line[:i]
	r.URI, proto = split2(line[i+1:], ' ')
	if len(r.URI) == 0 {
		return
	}
	if major, minor, ok := ParseVersion(proto); ok {
		r.Version.Major = major
		r.Version.Minor = minor
		return r, true
	}

	return r, false
}

// SplitResponseLine splits given slice of bytes into three chunks without
// parsing.
func SplitResponseLine(line []byte) (version, status, reason []byte) {
	return split3(line, ' ')
}

// ParseResponseLine parses first response line into ResponseLine struct.
func ParseResponseLine(line []byte) (r ResponseLine, ok bool) {
	var (
		proto  []byte
		status []byte
	)
	proto, status, r.Reason = split3(line, ' ')
	if major, minor, ok := ParseVersion(proto); ok {
		r.Version.Major = major
		r.Version.Minor = minor
	} else {
		return r, false
	}
	if n, ok := IntFromASCII(status); ok {
		r.Status = n
	} else {
		return r, false
	}
	// TODO(gobwas): parse here r.Reason fot TEXT rule:
	//   TEXT = <any OCTET except CTLs,
	//           but including LWS>
	return r, true
}

var (
	httpVersion10     = []byte("HTTP/1.0")
	httpVersion11     = []byte("HTTP/1.1")
	httpVersionPrefix = []byte("HTTP/")
)

// ParseVersion parses major and minor version of HTTP protocol.
// It returns parsed values and true if parse is ok.
func ParseVersion(bts []byte) (major, minor int, ok bool) {
	switch {
	case bytes.Equal(bts, httpVersion11):
		return 1, 1, true
	case bytes.Equal(bts, httpVersion10):
		return 1, 0, true
	case len(bts) < 8:
		return
	case !bytes.Equal(bts[:5], httpVersionPrefix):
		return
	}

	bts = bts[5:]

	dot := bytes.IndexByte(bts, '.')
	if dot == -1 {
		return
	}
	major, ok = IntFromASCII(bts[:dot])
	if !ok {
		return
	}
	minor, ok = IntFromASCII(bts[dot+1:])
	if !ok {
		return
	}

	return major, minor, true
}

// ReadLine reads line from br. It reads until '\n' and returns bytes without
// '\n' or '\r\n' at the end.
// It returns err if and only if line does not end in '\n'. Note that read
// bytes returned in any case of error.
//
// It is much like the textproto/Reader.ReadLine() except the thing that it
// returns raw bytes, instead of string. That is, it avoids copying bytes read
// from br.
//
// textproto/Reader.ReadLineBytes() is also makes copy of resulting bytes to be
// safe with future I/O operations on br.
//
// We could control I/O operations on br and do not need to make additional
// copy for safety.
func ReadLine(br *bufio.Reader) ([]byte, error) {
	var line []byte
	for {
		bts, err := br.ReadSlice('\n')
		if err == bufio.ErrBufferFull {
			// Copy bytes because next read will discard them.
			line = append(line, bts...)
			continue
		}
		// Avoid copy of single read.
		if line == nil {
			line = bts
		} else {
			line = append(line, bts...)
		}
		if err != nil {
			return line, err
		}
		// Size of line is at least 1.
		// In other case bufio.ReadSlice() returns error.
		n := len(line)
		// Cut '\n' or '\r\n'.
		if n > 1 && line[n-2] == '\r' {
			line = line[:n-2]
		} else {
			line = line[:n-1]
		}
		return line, nil
	}
}

// ParseHeaderLine parses HTTP header as key-value pair. It returns parsed
// values and true if parse is ok.
func ParseHeaderLine(line []byte) (k, v []byte, ok bool) {
	colon := bytes.IndexByte(line, ':')
	if colon == -1 {
		return
	}
	k = trim(line[:colon])
	for _, c := range k {
		if !OctetTypes[c].IsToken() {
			return nil, nil, false
		}
	}
	v = trim(line[colon+1:])
	return k, v, true
}

// IntFromASCII converts ascii encoded decimal numeric value from HTTP entities
// to an integer.
func IntFromASCII(bts []byte) (ret int, ok bool) {
	// ASCII numbers all start with the high-order bits 0011.
	// If you see that, and the next bits are 0-9 (0000 - 1001) you can grab those
	// bits and interpret them directly as an integer.
	var n int
	if n = len(bts); n < 1 {
		return 0, false
	}
	for i := 0; i < n; i++ {
		if bts[i]&0xf0 != 0x30 {
			return 0, false
		}
		ret += int(bts[i]&0xf) * pow(10, n-i-1)
	}
	return ret, true
}

// pow for integers implementation.
// See Donald Knuth, The Art of Computer Programming, Volume 2, Section 4.6.3
func pow(a, b int) int {
	p := 1
	for b > 0 {
		if b&1 != 0 {
			p *= a
		}
		b >>= 1
		a *= a
	}
	return p
}

func split3(p []byte, sep byte) (p1, p2, p3 []byte) {
	a := bytes.IndexByte(p, sep)
	b := bytes.IndexByte(p[a+1:], sep)
	if a == -1 || b == -1 {
		return p, nil, nil
	}
	b += a + 1
	return p[:a], p[a+1 : b], p[b+1:]
}

func split2(p []byte, sep byte) (p1, p2 []byte) {
	i := bytes.IndexByte(p, sep)
	if i == -1 {
		return p, nil
	}
	return p[:i], p[i+1:]
}

func trim(p []byte) []byte {
	var i, j int
	for i = 0; i < len(p) && (p[i] == ' ' || p[i] == '\t'); {
		i++
	}
	for j = len(p); j > i && (p[j-1] == ' ' || p[j-1] == '\t'); {
		j--
	}
	return p[i:j]
}
<|endoftext|>


<commit_before>package flag

import (
	"bytes"
	"fmt"
	"strings"
)

const (
	minInfoLen = 12
	maxInfoLen = 24
)

type writer struct {
	buf           *bytes.Buffer
	isTop         bool
	inheritIndent string
	forceVerbose  bool
	maxInfoLen    int
}

func (w *writer) maxFlagInfoLen(f *FlagSet) int {
	var maxLen int
	for i := range f.flags {
		l := len(f.flags[i].Names)
		if f.flags[i].Arglist != "" {
			l += 1 + len(f.flags[i].Arglist)
		}
		if maxLen < l {
			maxLen = l
		}
	}
	if maxLen < minInfoLen {
		maxLen = minInfoLen
	}
	if maxLen > maxInfoLen {
		maxLen = maxInfoLen
	}
	return maxLen
}

func (w *writer) maxSubsetInfoLen(f *FlagSet, needArglist bool) int {
	var maxLen int
	for i := range f.subsets {
		l := len(f.subsets[i].self.Names)
		if needArglist {
			args := w.arglist(&f.subsets[i])
			if args != "" {
				l += 1 + len(args)
			}
		}
		if maxLen < l {
			maxLen = l
		}
	}
	if maxLen < minInfoLen {
		maxLen = minInfoLen
	}
	if maxLen > maxInfoLen {
		maxLen = maxInfoLen
	}
	return maxLen
}

func (w *writer) arglist(f *FlagSet) string {
	if f.self.Arglist == "-" {
		return ""
	}
	if f.self.Arglist != "" {
		return f.self.Arglist
	}
	var (
		arglist             string
		flagCount, setCount = len(f.flags), len(f.subsets)
	)
	if flagCount != 0 {
		if setCount != 0 {
			arglist = "[FLAG|SET]..."
		} else {
			arglist = "[FLAG]..."
		}
	} else {
		if setCount != 0 {
			arglist = "[SET]..."
		}
	}
	return arglist
}

func (w *writer) write(elem ...string) {
	for _, s := range elem {
		w.buf.WriteString(s)
	}
}

func (w *writer) nextIndent(curr string) string {
	const indent = "      "
	return curr + indent
}

func (w *writer) writeln(elem ...string) {
	w.write(elem...)
	w.buf.WriteByte('\n')
}

func (w *writer) writeWithPads(names string, maxLen int) {
	w.write(names)
	if padlen := maxLen - len(names); padlen > 0 {
		w.write(strings.Repeat(" ", padlen))
	}
}

func (w *writer) writeLines(indent string, lines []string) {
	for _, line := range lines {
		w.writeln(indent, line)
	}
}

func (w *writer) parseFlagInfo(flag *Flag, args string) string {
	info := flag.Names
	if args != "" {
		info += " " + args
	}
	return info
}

func (w *writer) writeFlagInfo(currIndent string, flag *Flag, isTop bool, args string, maxInfoLen int) {
	w.write(currIndent)
	if isTop {
		if flag.Usage != "" {
			w.writeln(currIndent, flag.Usage)
			w.writeln()
		}
		w.writeln(currIndent, "Usage:")
		w.write(w.nextIndent(currIndent))
	}
	flagInfo := w.parseFlagInfo(flag, args)
	w.writeWithPads(flagInfo, maxInfoLen)
	if !isTop && flag.Usage != "" {
		w.write(" ", flag.Usage)
	}
}

func (w *writer) writeFlagValueInfo(flag *Flag) {
	if flag.Usage != "" && (flag.Env != "" || flag.Default != nil || flag.Selects != nil) {
		w.write(" (")
		var hasPrev bool
		if flag.ShowType {
			w.write(typeName(flag.Ptr))
			hasPrev = true
		}
		if flag.Env != "" {
			if hasPrev {
				w.write("; ")
			}
			fmt.Fprintf(w.buf, "env: %s", flag.Env)
			if isSlicePtr(flag.Ptr) {
				fmt.Fprintf(w.buf, ", splitted by '%s'", flag.ValSep)
			}
		}
		if flag.Default != nil {
			if hasPrev {
				w.write("; ")
			}
			fmt.Fprintf(w.buf, "default: %v", flag.Default)
		}
		if flag.Selects != nil {
			if hasPrev {
				w.write("; ")
			}
			fmt.Fprintf(w.buf, "selects: %v", flag.Selects)
		}
		w.write(")")
	}
}

func (w *writer) writeSet(f *FlagSet) {
	var (
		currIndent       = w.inheritIndent
		flagIndent       = w.nextIndent(currIndent)
		outline          = !f.self.Expand && !w.forceVerbose
		flagCount        = len(f.flags)
		subsetCount      = len(f.subsets)
		versionLineCount = len(f.self.versionLines)
		descLineCount    = len(f.self.descLines)
	)

	var arglist string
	if w.isTop {
		arglist = w.arglist(f)
	}
	w.writeFlagInfo(currIndent, &f.self, w.isTop, arglist, w.maxInfoLen)
	w.writeln()

	if outline && !w.isTop {
		return
	}
	if versionLineCount > 0 {
		if w.isTop {
			w.writeln()
			w.writeln(currIndent, "Version:")
			w.writeLines(flagIndent, f.self.versionLines)
		}
	}

	if descLineCount > 0 {
		if w.isTop {
			w.writeln()
			w.writeln(currIndent, "Description:")
		}
		w.writeLines(flagIndent, f.self.descLines)
	}

	if flagCount > 0 {
		if versionLineCount > 0 || descLineCount > 0 || w.isTop {
			w.writeln()
		}
		if w.isTop {
			w.writeln(currIndent, "Flags:")
		}
		var (
			maxFlagInfoLen = w.maxFlagInfoLen(f)
			nextFlagIndent = w.nextIndent(flagIndent)
			important      = true
			hasImportant   bool
		)
		for {
			for i := range f.flags {
				flag := &f.flags[i]
				if important != flag.Important {
					continue
				}
				if flag.Important {
					hasImportant = true
				} else if hasImportant {
					w.writeln()
					hasImportant = false
				}

				w.writeFlagInfo(flagIndent, flag, false, flag.Arglist, maxFlagInfoLen)
				w.writeFlagValueInfo(flag)
				w.writeln()
				w.writeLines(nextFlagIndent, flag.descLines)
			}
			if important {
				important = false
			} else {
				break
			}
		}
	}

	if subsetCount > 0 {
		if w.isTop || descLineCount > 0 || flagCount > 0 {
			w.writeln()
		}
		if w.isTop {
			w.writeln(currIndent, "Sets:")
		}
		var (
			maxSubsetLen = w.maxSubsetInfoLen(f, !outline)
			subsetIndent = flagIndent
			important    = true
			hasImportant bool
		)
		for {
			for i := range f.subsets {
				set := &f.subsets[i]
				if important != set.self.Important {
					continue
				}

				if set.self.Important {
					hasImportant = true
				} else if hasImportant {
					w.writeln()
					hasImportant = false
				}
				nw := writer{
					buf:           w.buf,
					inheritIndent: subsetIndent,
					forceVerbose:  w.forceVerbose,
					maxInfoLen:    maxSubsetLen,
				}
				nw.writeSet(set)
			}
			if important {
				important = false
			} else {
				break
			}
		}
	}
}
<commit_msg>fix help message<commit_after>package flag

import (
	"bytes"
	"fmt"
	"strings"
)

const (
	minInfoLen = 12
	maxInfoLen = 24
)

type writer struct {
	buf           *bytes.Buffer
	isTop         bool
	inheritIndent string
	forceVerbose  bool
	maxInfoLen    int
}

func (w *writer) maxFlagInfoLen(f *FlagSet) int {
	var maxLen int
	for i := range f.flags {
		l := len(f.flags[i].Names)
		if f.flags[i].Arglist != "" {
			l += 1 + len(f.flags[i].Arglist)
		}
		if maxLen < l {
			maxLen = l
		}
	}
	if maxLen < minInfoLen {
		maxLen = minInfoLen
	}
	if maxLen > maxInfoLen {
		maxLen = maxInfoLen
	}
	return maxLen
}

func (w *writer) maxSubsetInfoLen(f *FlagSet, needArglist bool) int {
	var maxLen int
	for i := range f.subsets {
		l := len(f.subsets[i].self.Names)
		if needArglist {
			args := w.arglist(&f.subsets[i])
			if args != "" {
				l += 1 + len(args)
			}
		}
		if maxLen < l {
			maxLen = l
		}
	}
	if maxLen < minInfoLen {
		maxLen = minInfoLen
	}
	if maxLen > maxInfoLen {
		maxLen = maxInfoLen
	}
	return maxLen
}

func (w *writer) arglist(f *FlagSet) string {
	if f.self.Arglist == "-" {
		return ""
	}
	if f.self.Arglist != "" {
		return f.self.Arglist
	}
	var (
		arglist             string
		flagCount, setCount = len(f.flags), len(f.subsets)
	)
	if flagCount != 0 {
		if setCount != 0 {
			arglist = "[FLAG|SET]..."
		} else {
			arglist = "[FLAG]..."
		}
	} else {
		if setCount != 0 {
			arglist = "[SET]..."
		}
	}
	return arglist
}

func (w *writer) write(elem ...string) {
	for _, s := range elem {
		w.buf.WriteString(s)
	}
}

func (w *writer) nextIndent(curr string) string {
	const indent = "      "
	return curr + indent
}

func (w *writer) writeln(elem ...string) {
	w.write(elem...)
	w.buf.WriteByte('\n')
}

func (w *writer) writeWithPads(names string, maxLen int) {
	w.write(names)
	if padlen := maxLen - len(names); padlen > 0 {
		w.write(strings.Repeat(" ", padlen))
	}
}

func (w *writer) writeLines(indent string, lines []string) {
	for _, line := range lines {
		w.writeln(indent, line)
	}
}

func (w *writer) parseFlagInfo(flag *Flag, args string) string {
	info := flag.Names
	if args != "" {
		info += " " + args
	}
	return info
}

func (w *writer) writeFlagInfo(currIndent string, flag *Flag, isTop bool, args string, maxInfoLen int) {
	w.write(currIndent)
	if isTop {
		if flag.Usage != "" {
			w.writeln(currIndent, flag.Usage)
			w.writeln()
		}
		w.writeln(currIndent, "Usage:")
		w.write(w.nextIndent(currIndent))
	}
	flagInfo := w.parseFlagInfo(flag, args)
	w.writeWithPads(flagInfo, maxInfoLen)
	if !isTop && flag.Usage != "" {
		w.write(" ", flag.Usage)
	}
}

func (w *writer) writeFlagValueInfo(flag *Flag) {
	if flag.Usage != "" && (flag.Env != "" || flag.Default != nil || flag.Selects != nil) {
		w.write(" (")
		var hasPrev bool
		if flag.ShowType {
			w.write(typeName(flag.Ptr))
			hasPrev = true
		}
		if flag.Env != "" {
			if hasPrev {
				w.write("; ")
			}
			fmt.Fprintf(w.buf, "env: %s", flag.Env)
			if isSlicePtr(flag.Ptr) {
				fmt.Fprintf(w.buf, ", splitted by '%s'", flag.ValSep)
			}
			hasPrev = true
		}
		if flag.Default != nil {
			if hasPrev {
				w.write("; ")
			}
			fmt.Fprintf(w.buf, "default: %v", flag.Default)
			hasPrev = true
		}
		if flag.Selects != nil {
			if hasPrev {
				w.write("; ")
			}
			fmt.Fprintf(w.buf, "selects: %v", flag.Selects)
		}
		w.write(")")
	}
}

func (w *writer) writeSet(f *FlagSet) {
	var (
		currIndent       = w.inheritIndent
		flagIndent       = w.nextIndent(currIndent)
		outline          = !f.self.Expand && !w.forceVerbose
		flagCount        = len(f.flags)
		subsetCount      = len(f.subsets)
		versionLineCount = len(f.self.versionLines)
		descLineCount    = len(f.self.descLines)
	)

	var arglist string
	if w.isTop {
		arglist = w.arglist(f)
	}
	w.writeFlagInfo(currIndent, &f.self, w.isTop, arglist, w.maxInfoLen)
	w.writeln()

	if outline && !w.isTop {
		return
	}
	if versionLineCount > 0 {
		if w.isTop {
			w.writeln()
			w.writeln(currIndent, "Version:")
			w.writeLines(flagIndent, f.self.versionLines)
		}
	}

	if descLineCount > 0 {
		if w.isTop {
			w.writeln()
			w.writeln(currIndent, "Description:")
		}
		w.writeLines(flagIndent, f.self.descLines)
	}

	if flagCount > 0 {
		if versionLineCount > 0 || descLineCount > 0 || w.isTop {
			w.writeln()
		}
		if w.isTop {
			w.writeln(currIndent, "Flags:")
		}
		var (
			maxFlagInfoLen = w.maxFlagInfoLen(f)
			nextFlagIndent = w.nextIndent(flagIndent)
			important      = true
			hasImportant   bool
		)
		for {
			for i := range f.flags {
				flag := &f.flags[i]
				if important != flag.Important {
					continue
				}
				if flag.Important {
					hasImportant = true
				} else if hasImportant {
					w.writeln()
					hasImportant = false
				}

				w.writeFlagInfo(flagIndent, flag, false, flag.Arglist, maxFlagInfoLen)
				w.writeFlagValueInfo(flag)
				w.writeln()
				w.writeLines(nextFlagIndent, flag.descLines)
			}
			if important {
				important = false
			} else {
				break
			}
		}
	}

	if subsetCount > 0 {
		if w.isTop || descLineCount > 0 || flagCount > 0 {
			w.writeln()
		}
		if w.isTop {
			w.writeln(currIndent, "Sets:")
		}
		var (
			maxSubsetLen = w.maxSubsetInfoLen(f, !outline)
			subsetIndent = flagIndent
			important    = true
			hasImportant bool
		)
		for {
			for i := range f.subsets {
				set := &f.subsets[i]
				if important != set.self.Important {
					continue
				}

				if set.self.Important {
					hasImportant = true
				} else if hasImportant {
					w.writeln()
					hasImportant = false
				}
				nw := writer{
					buf:           w.buf,
					inheritIndent: subsetIndent,
					forceVerbose:  w.forceVerbose,
					maxInfoLen:    maxSubsetLen,
				}
				nw.writeSet(set)
			}
			if important {
				important = false
			} else {
				break
			}
		}
	}
}
<|endoftext|>


<commit_before>package sshchat

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

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

const maxInputLength int = 1024

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

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

	// Version string to print on /version
	Version string

	// Default theme
	theme message.Theme

	mu    sync.Mutex
	motd  string
	count int
}

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

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

	go room.Serve()
	return &h
}

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

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

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

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

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

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

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

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

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

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

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

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

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

		m := message.ParseInput(line, user)

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

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

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

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

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

	return names[len(names)-1]
}

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

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

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

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

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

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

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

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

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

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

			txt := fmt.Sprintf("[Sent PM to %s]", target.Name())
			ms := message.NewSystemMsg(txt, msg.From())
			room.Send(ms)
			return nil
		},
	})

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

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

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

			txt := fmt.Sprintf("[Sent PM to %s]", target.Name())
			ms := message.NewSystemMsg(txt, msg.From())
			room.Send(ms)
			return nil
		},
	})

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

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

			id := target.Identifier.(*Identity)
			var whois string
			switch room.IsOp(msg.From()) {
			case true:
				whois = id.WhoisAdmin()
			case false:
				whois = id.Whois()
			}
			room.Send(message.NewSystemMsg(whois, msg.From()))

			return nil
		},
	})

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

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

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

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

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

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

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

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

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

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

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

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

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

			return nil
		},
	})

	c.Add(chat.Command{
		Op:         true,
		Prefix:     "/motd",
		PrefixHelp: "[MESSAGE]",
		Help:       "Set a new MESSAGE of the day, print the current motd without parameters.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			args := msg.Args()
			user := msg.From()
			motd := h.motd

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

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

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

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

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

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

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

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

			return nil
		},
	})
}
<commit_msg>/motd: Fixed set message and actual motd coming out of order.<commit_after>package sshchat

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

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

const maxInputLength int = 1024

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

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

	// Version string to print on /version
	Version string

	// Default theme
	theme message.Theme

	mu    sync.Mutex
	motd  string
	count int
}

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

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

	go room.Serve()
	return &h
}

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

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

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

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

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

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

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

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

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

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

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

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

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

		m := message.ParseInput(line, user)

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

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

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

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

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

	return names[len(names)-1]
}

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

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

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

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

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

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

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

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

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

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

			txt := fmt.Sprintf("[Sent PM to %s]", target.Name())
			ms := message.NewSystemMsg(txt, msg.From())
			room.Send(ms)
			return nil
		},
	})

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

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

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

			txt := fmt.Sprintf("[Sent PM to %s]", target.Name())
			ms := message.NewSystemMsg(txt, msg.From())
			room.Send(ms)
			return nil
		},
	})

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

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

			id := target.Identifier.(*Identity)
			var whois string
			switch room.IsOp(msg.From()) {
			case true:
				whois = id.WhoisAdmin()
			case false:
				whois = id.Whois()
			}
			room.Send(message.NewSystemMsg(whois, msg.From()))

			return nil
		},
	})

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

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

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

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

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

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

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

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

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

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

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

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

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

			return nil
		},
	})

	c.Add(chat.Command{
		Op:         true,
		Prefix:     "/motd",
		PrefixHelp: "[MESSAGE]",
		Help:       "Set a new MESSAGE of the day, print the current motd without parameters.",
		Handler: func(room *chat.Room, msg message.CommandMsg) error {
			args := msg.Args()
			user := msg.From()

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

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

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

			return nil
		},
	})

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

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

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

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

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

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


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
package vix

/*
#include "vix.h"
#include "helper.h"
*/
import "C"

import "unsafe"

type Host struct {
	Provider Provider
	handle   C.VixHandle
}

// Destroys the state for a particular host instance
//
// Call this function to disconnect the host. After you call this function the
// Host object is no longer valid and you should not longer use it.
// Similarly, you should not use any other object instances obtained from the
// Host object while it was connected.
//
// Since VMware Server 1.0
func (h *Host) Disconnect() {
	if h.handle != C.VIX_INVALID_HANDLE {
		C.VixHost_Disconnect(h.handle)
		h.handle = C.VIX_INVALID_HANDLE
	}
}

//export go_callback_char
func go_callback_char(callbackPtr unsafe.Pointer, item *C.char) {
	callback := *(*func(*C.char))(callbackPtr)
	callback(item)
}

// This function finds Vix objects. For example, when used to find all
// running virtual machines, Host.FindItems() returns a series of virtual
// machine file path names.
func (h *Host) FindItems(options SearchType) ([]string, error) {
	var jobHandle C.VixHandle = C.VIX_INVALID_HANDLE
	var err C.VixError = C.VIX_OK
	var items []string

	callback := func(item *C.char) {
		items = append(items, C.GoString(item))
	}

	jobHandle = C.VixHost_FindItems(h.handle,
		C.VixFindItemType(options), //searchType
		C.VIX_INVALID_HANDLE,       //searchCriteria
		-1,                         //timeout
		(*C.VixEventProc)(C.find_items_callback), //callbackProc
		unsafe.Pointer(&callback))                //clientData

	defer C.Vix_ReleaseHandle(jobHandle)

	err = C.vix_job_wait(jobHandle)
	if C.VIX_OK != err {
		return nil, &VixError{
			Operation: "host.FindItems",
			Code:      int(err & 0xFFFF),
			Text:      C.GoString(C.Vix_GetErrorText(err, nil)),
		}
	}

	return items, nil
}

// This function opens a virtual machine on the host
// and returns a VM instance.
//
// Parameters:
//
//   VmxFile: The path name of the virtual machine configuration file on the
//   local host.
//
//   Password: If VM is encrypted, this is the password for VIX to be able to
//   open it.
//
// Remarks:
//
//   * This function opens a virtual machine on the host instance
//     The virtual machine is identified by vmxFile, which is a path name to the
//     configuration file (.VMX file) for that virtual machine.
//
//   * The format of the path name depends on the host operating system.
//     For example, a path name for a Windows host requires backslash as a
//     directory separator, whereas a Linux host requires a forward slash. If the
//     path name includes backslash characters, you need to precede each one with
//     an escape character. For VMware Server 2.x, the path contains a preceeding
//     data store, for example [storage1] vm/vm.vmx.
//
//   * For VMware Server hosts, a virtual machine must be registered before you
//     can open it. You can register a virtual machine by opening it with the
//     VMware Server Console, through the vmware-cmd command with the register
//     parameter, or with Host.RegisterVM().
//
//   * For vSphere, the virtual machine opened may not be the one desired if more
//     than one Datacenter contains VmxFile.
//
//   * To open an encrypted virtual machine, pass its correspondent password.
//
// Since VMware Workstation 7.0
func (h *Host) OpenVm(vmxFile, password string) (*VM, error) {
	var jobHandle C.VixHandle = C.VIX_INVALID_HANDLE
	var propertyHandle C.VixHandle = C.VIX_INVALID_HANDLE
	var vmHandle C.VixHandle = C.VIX_INVALID_HANDLE
	var err C.VixError = C.VIX_OK

	defer C.Vix_ReleaseHandle(propertyHandle)
	defer C.Vix_ReleaseHandle(jobHandle)

	if password != "" {
		cpassword := C.CString(password)
		defer C.free(unsafe.Pointer(cpassword))

		err = C.alloc_vm_pwd_proplist(h.handle,
			&propertyHandle,
			cpassword)

		if C.VIX_OK != err {
			return nil, &VixError{
				Operation: "host.OpenVM",
				Code:      int(err & 0xFFFF),
				Text:      C.GoString(C.Vix_GetErrorText(err, nil)),
			}
		}
	}

	cVmxFile := C.CString(vmxFile)
	defer C.free(unsafe.Pointer(cVmxFile))

	jobHandle = C.VixHost_OpenVM(h.handle,
		cVmxFile,
		C.VIX_VMOPEN_NORMAL,
		propertyHandle,
		nil, // callbackProc
		nil) // clientData

	err = C.get_vix_handle(jobHandle,
		C.VIX_PROPERTY_JOB_RESULT_HANDLE,
		&vmHandle,
		C.VIX_PROPERTY_NONE)

	if C.VIX_OK != err {
		return nil, &VixError{
			Operation: "host.OpenVM.get_vix_handle",
			Code:      int(err & 0xFFFF),
			Text:      C.GoString(C.Vix_GetErrorText(err, nil)),
		}
	}

	return NewVirtualMachine(vmHandle, vmxFile)
}

// This function adds a virtual machine to the host's inventory.
//
// Parameters:
//
//   vmxFile: The path name of the .vmx file on the host.
//
// Remarks:
//
//   * This function registers the virtual machine identified by vmxFile, which
//     is a storage path to the configuration file (.vmx) for that virtual machine.
//     You can register a virtual machine regardless of its power state.
//
//   * The format of the path name depends on the host operating system.
//     If the path name includes backslash characters, you need to precede each
//     one with an escape character. Path to storage [standard] or [storage1] may
//     vary.
//
//   * For VMware Server 1.x, supply the full path name instead of storage path,
//     and specify provider VMWARE_SERVER to connect.
//
//   * This function has no effect on Workstation or Player, which lack a virtual
//     machine inventory.
//
//   * It is not a Vix error to register an already-registered virtual machine,
//     although the VMware Server UI shows an error icon in the Task pane.
//     Trying to register a non-existent virtual machine results in error 2000,
//     VIX_E_NOT_FOUND.
//
// Since VMware Server 1.0
func (h *Host) RegisterVm(vmxFile string) error {
	var jobHandle C.VixHandle = C.VIX_INVALID_HANDLE
	var err C.VixError = C.VIX_OK

	cVmxFile := C.CString(vmxFile)
	defer C.free(unsafe.Pointer(cVmxFile))

	jobHandle = C.VixHost_RegisterVM(h.handle,
		cVmxFile,
		nil, // callbackProc
		nil) // clientData

	defer C.Vix_ReleaseHandle(jobHandle)

	err = C.vix_job_wait(jobHandle)
	if C.VIX_OK != err {
		return &VixError{
			Operation: "host.RegisterVM",
			Code:      int(err & 0xFFFF),
			Text:      C.GoString(C.Vix_GetErrorText(err, nil)),
		}
	}

	return nil
}

// This function removes a virtual machine from the host's inventory.
//
// Parameters:
//
//   VmxFile: The path name of the .vmx file on the host.
//
// Remarks:
//
//   * This function unregisters the virtual machine identified by vmxFile,
//     which is a storage path to the configuration file (.vmx) for that virtual
//     machine. A virtual machine must be powered off to unregister it.
//   * The format of the storage path depends on the host operating system.
//     If the storage path includes backslash characters, you need to precede each
//     one with an escape character. Path to storage [standard] or [storage1] may
//     vary.
//   * For VMware Server 1.x, supply the full path name instead of storage path,
//     and specify VMWARE_SERVER provider to connect.
//   * This function has no effect on Workstation or Player, which lack a virtual
//     machine inventory.
//   * It is not a Vix error to unregister an already-unregistered virtual machine,
//     nor is it a Vix error to unregister a non-existent virtual machine.
//
// Since VMware Server 1.0
func (h *Host) UnregisterVm(vmxFile string) error {
	var jobHandle C.VixHandle = C.VIX_INVALID_HANDLE
	var err C.VixError = C.VIX_OK

	cVmxFile := C.CString(vmxFile)
	defer C.free(unsafe.Pointer(cVmxFile))

	jobHandle = C.VixHost_UnregisterVM(h.handle,
		cVmxFile,
		nil, // callbackProc
		nil) // clientData

	defer C.Vix_ReleaseHandle(jobHandle)

	err = C.vix_job_wait(jobHandle)
	if C.VIX_OK != err {
		return &VixError{
			Operation: "host.UnregisterVM",
			Code:      int(err & 0xFFFF),
			Text:      C.GoString(C.Vix_GetErrorText(err, nil)),
		}
	}

	return nil
}

// Copies a file or directory from the local system (where the Vix client is
// running) to the guest operating system.
//
// Parameters:
//
//   src: The path name of a file on a file system available to the Vix client.
//   guest: Guest instance where the file is going to be copied to
//   dest: The path name of a file on a file system available to the guest.
//
// Remarks:
//
//   * The virtual machine must be running while the file is copied from the Vix
//     client machine to the guest operating system.
//
//   * Existing files of the same name are overwritten, and folder contents are
//     merged.
//
//   * The copy operation requires VMware Tools to be installed and running in
//     the guest operating system.
//
//   * You must call VM.LoginInGuest() before calling this function in order
//     to get a Guest instance.
//
//   * The format of the file name depends on the guest or local operating system.
//     For example, a path name for a Microsoft Windows guest or host requires
//     backslash as a directory separator, whereas a Linux guest or host requires
//     a forward slash. If the path name includes backslash characters,
//     you need to precede each one with an escape character.
//
//   * Only absolute paths should be used for files in the guest; the resolution
//     of relative paths is not specified.
//
//   * If any file fails to be copied, Vix aborts the operation, does not attempt
//     to copy the remaining files, and returns an error.
//
//   * In order to copy a file to a mapped network drive in a Windows guest
//     operating system, it is necessary to call VixVM_LoginInGuest() with the
//     LOGIN_IN_GUEST_REQUIRE_INTERACTIVE_ENVIRONMENT flag set.
//     Using the interactive session option incurs an overhead in file transfer
//     speed.
//
// Since VMware Server 1.0
// Minimum Supported Guest OS: Microsoft Windows NT Series, Linux
//
func (h *Host) CopyFileToGuest(src string, guest *Guest, dest string) error {
	var jobHandle C.VixHandle = C.VIX_INVALID_HANDLE
	var err C.VixError = C.VIX_OK

	csrc := C.CString(src)
	cdest := C.CString(dest)
	defer C.free(unsafe.Pointer(csrc))
	defer C.free(unsafe.Pointer(cdest))

	jobHandle = C.VixVM_CopyFileFromHostToGuest(
		guest.handle,         // VM handle
		csrc,                 // src name
		cdest,                // dest name
		C.int(0),             // options
		C.VIX_INVALID_HANDLE, // propertyListHandle
		nil,                  // callbackProc
		nil)                  // clientData

	defer C.Vix_ReleaseHandle(jobHandle)

	err = C.vix_job_wait(jobHandle)
	if C.VIX_OK != err {
		return &VixError{
			Operation: "host.CopyFileToGuest",
			Code:      int(err & 0xFFFF),
			Text:      C.GoString(C.Vix_GetErrorText(err, nil)),
		}
	}

	return nil
}
<commit_msg>Makes sure host handle is nil or invalid before attempting to disconnect<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
package vix

/*
#include "vix.h"
#include "helper.h"
*/
import "C"

import "unsafe"

type Host struct {
	Provider Provider
	handle   C.VixHandle
}

// Destroys the state for a particular host instance
//
// Call this function to disconnect the host. After you call this function the
// Host object is no longer valid and you should not longer use it.
// Similarly, you should not use any other object instances obtained from the
// Host object while it was connected.
//
// Since VMware Server 1.0
func (h *Host) Disconnect() {
	if &h.handle != nil && h.handle != C.VIX_INVALID_HANDLE {
		C.VixHost_Disconnect(h.handle)
		h.handle = C.VIX_INVALID_HANDLE
	}
}

//export go_callback_char
func go_callback_char(callbackPtr unsafe.Pointer, item *C.char) {
	callback := *(*func(*C.char))(callbackPtr)
	callback(item)
}

// This function finds Vix objects. For example, when used to find all
// running virtual machines, Host.FindItems() returns a series of virtual
// machine file path names.
func (h *Host) FindItems(options SearchType) ([]string, error) {
	var jobHandle C.VixHandle = C.VIX_INVALID_HANDLE
	var err C.VixError = C.VIX_OK
	var items []string

	callback := func(item *C.char) {
		items = append(items, C.GoString(item))
	}

	jobHandle = C.VixHost_FindItems(h.handle,
		C.VixFindItemType(options), //searchType
		C.VIX_INVALID_HANDLE,       //searchCriteria
		-1,                         //timeout
		(*C.VixEventProc)(C.find_items_callback), //callbackProc
		unsafe.Pointer(&callback))                //clientData

	defer C.Vix_ReleaseHandle(jobHandle)

	err = C.vix_job_wait(jobHandle)
	if C.VIX_OK != err {
		return nil, &VixError{
			Operation: "host.FindItems",
			Code:      int(err & 0xFFFF),
			Text:      C.GoString(C.Vix_GetErrorText(err, nil)),
		}
	}

	return items, nil
}

// This function opens a virtual machine on the host
// and returns a VM instance.
//
// Parameters:
//
//   VmxFile: The path name of the virtual machine configuration file on the
//   local host.
//
//   Password: If VM is encrypted, this is the password for VIX to be able to
//   open it.
//
// Remarks:
//
//   * This function opens a virtual machine on the host instance
//     The virtual machine is identified by vmxFile, which is a path name to the
//     configuration file (.VMX file) for that virtual machine.
//
//   * The format of the path name depends on the host operating system.
//     For example, a path name for a Windows host requires backslash as a
//     directory separator, whereas a Linux host requires a forward slash. If the
//     path name includes backslash characters, you need to precede each one with
//     an escape character. For VMware Server 2.x, the path contains a preceeding
//     data store, for example [storage1] vm/vm.vmx.
//
//   * For VMware Server hosts, a virtual machine must be registered before you
//     can open it. You can register a virtual machine by opening it with the
//     VMware Server Console, through the vmware-cmd command with the register
//     parameter, or with Host.RegisterVM().
//
//   * For vSphere, the virtual machine opened may not be the one desired if more
//     than one Datacenter contains VmxFile.
//
//   * To open an encrypted virtual machine, pass its correspondent password.
//
// Since VMware Workstation 7.0
func (h *Host) OpenVm(vmxFile, password string) (*VM, error) {
	var jobHandle C.VixHandle = C.VIX_INVALID_HANDLE
	var propertyHandle C.VixHandle = C.VIX_INVALID_HANDLE
	var vmHandle C.VixHandle = C.VIX_INVALID_HANDLE
	var err C.VixError = C.VIX_OK

	defer C.Vix_ReleaseHandle(propertyHandle)
	defer C.Vix_ReleaseHandle(jobHandle)

	if password != "" {
		cpassword := C.CString(password)
		defer C.free(unsafe.Pointer(cpassword))

		err = C.alloc_vm_pwd_proplist(h.handle,
			&propertyHandle,
			cpassword)

		if C.VIX_OK != err {
			return nil, &VixError{
				Operation: "host.OpenVM",
				Code:      int(err & 0xFFFF),
				Text:      C.GoString(C.Vix_GetErrorText(err, nil)),
			}
		}
	}

	cVmxFile := C.CString(vmxFile)
	defer C.free(unsafe.Pointer(cVmxFile))

	jobHandle = C.VixHost_OpenVM(h.handle,
		cVmxFile,
		C.VIX_VMOPEN_NORMAL,
		propertyHandle,
		nil, // callbackProc
		nil) // clientData

	err = C.get_vix_handle(jobHandle,
		C.VIX_PROPERTY_JOB_RESULT_HANDLE,
		&vmHandle,
		C.VIX_PROPERTY_NONE)

	if C.VIX_OK != err {
		return nil, &VixError{
			Operation: "host.OpenVM.get_vix_handle",
			Code:      int(err & 0xFFFF),
			Text:      C.GoString(C.Vix_GetErrorText(err, nil)),
		}
	}

	return NewVirtualMachine(vmHandle, vmxFile)
}

// This function adds a virtual machine to the host's inventory.
//
// Parameters:
//
//   vmxFile: The path name of the .vmx file on the host.
//
// Remarks:
//
//   * This function registers the virtual machine identified by vmxFile, which
//     is a storage path to the configuration file (.vmx) for that virtual machine.
//     You can register a virtual machine regardless of its power state.
//
//   * The format of the path name depends on the host operating system.
//     If the path name includes backslash characters, you need to precede each
//     one with an escape character. Path to storage [standard] or [storage1] may
//     vary.
//
//   * For VMware Server 1.x, supply the full path name instead of storage path,
//     and specify provider VMWARE_SERVER to connect.
//
//   * This function has no effect on Workstation or Player, which lack a virtual
//     machine inventory.
//
//   * It is not a Vix error to register an already-registered virtual machine,
//     although the VMware Server UI shows an error icon in the Task pane.
//     Trying to register a non-existent virtual machine results in error 2000,
//     VIX_E_NOT_FOUND.
//
// Since VMware Server 1.0
func (h *Host) RegisterVm(vmxFile string) error {
	var jobHandle C.VixHandle = C.VIX_INVALID_HANDLE
	var err C.VixError = C.VIX_OK

	cVmxFile := C.CString(vmxFile)
	defer C.free(unsafe.Pointer(cVmxFile))

	jobHandle = C.VixHost_RegisterVM(h.handle,
		cVmxFile,
		nil, // callbackProc
		nil) // clientData

	defer C.Vix_ReleaseHandle(jobHandle)

	err = C.vix_job_wait(jobHandle)
	if C.VIX_OK != err {
		return &VixError{
			Operation: "host.RegisterVM",
			Code:      int(err & 0xFFFF),
			Text:      C.GoString(C.Vix_GetErrorText(err, nil)),
		}
	}

	return nil
}

// This function removes a virtual machine from the host's inventory.
//
// Parameters:
//
//   VmxFile: The path name of the .vmx file on the host.
//
// Remarks:
//
//   * This function unregisters the virtual machine identified by vmxFile,
//     which is a storage path to the configuration file (.vmx) for that virtual
//     machine. A virtual machine must be powered off to unregister it.
//   * The format of the storage path depends on the host operating system.
//     If the storage path includes backslash characters, you need to precede each
//     one with an escape character. Path to storage [standard] or [storage1] may
//     vary.
//   * For VMware Server 1.x, supply the full path name instead of storage path,
//     and specify VMWARE_SERVER provider to connect.
//   * This function has no effect on Workstation or Player, which lack a virtual
//     machine inventory.
//   * It is not a Vix error to unregister an already-unregistered virtual machine,
//     nor is it a Vix error to unregister a non-existent virtual machine.
//
// Since VMware Server 1.0
func (h *Host) UnregisterVm(vmxFile string) error {
	var jobHandle C.VixHandle = C.VIX_INVALID_HANDLE
	var err C.VixError = C.VIX_OK

	cVmxFile := C.CString(vmxFile)
	defer C.free(unsafe.Pointer(cVmxFile))

	jobHandle = C.VixHost_UnregisterVM(h.handle,
		cVmxFile,
		nil, // callbackProc
		nil) // clientData

	defer C.Vix_ReleaseHandle(jobHandle)

	err = C.vix_job_wait(jobHandle)
	if C.VIX_OK != err {
		return &VixError{
			Operation: "host.UnregisterVM",
			Code:      int(err & 0xFFFF),
			Text:      C.GoString(C.Vix_GetErrorText(err, nil)),
		}
	}

	return nil
}

// Copies a file or directory from the local system (where the Vix client is
// running) to the guest operating system.
//
// Parameters:
//
//   src: The path name of a file on a file system available to the Vix client.
//   guest: Guest instance where the file is going to be copied to
//   dest: The path name of a file on a file system available to the guest.
//
// Remarks:
//
//   * The virtual machine must be running while the file is copied from the Vix
//     client machine to the guest operating system.
//
//   * Existing files of the same name are overwritten, and folder contents are
//     merged.
//
//   * The copy operation requires VMware Tools to be installed and running in
//     the guest operating system.
//
//   * You must call VM.LoginInGuest() before calling this function in order
//     to get a Guest instance.
//
//   * The format of the file name depends on the guest or local operating system.
//     For example, a path name for a Microsoft Windows guest or host requires
//     backslash as a directory separator, whereas a Linux guest or host requires
//     a forward slash. If the path name includes backslash characters,
//     you need to precede each one with an escape character.
//
//   * Only absolute paths should be used for files in the guest; the resolution
//     of relative paths is not specified.
//
//   * If any file fails to be copied, Vix aborts the operation, does not attempt
//     to copy the remaining files, and returns an error.
//
//   * In order to copy a file to a mapped network drive in a Windows guest
//     operating system, it is necessary to call VixVM_LoginInGuest() with the
//     LOGIN_IN_GUEST_REQUIRE_INTERACTIVE_ENVIRONMENT flag set.
//     Using the interactive session option incurs an overhead in file transfer
//     speed.
//
// Since VMware Server 1.0
// Minimum Supported Guest OS: Microsoft Windows NT Series, Linux
//
func (h *Host) CopyFileToGuest(src string, guest *Guest, dest string) error {
	var jobHandle C.VixHandle = C.VIX_INVALID_HANDLE
	var err C.VixError = C.VIX_OK

	csrc := C.CString(src)
	cdest := C.CString(dest)
	defer C.free(unsafe.Pointer(csrc))
	defer C.free(unsafe.Pointer(cdest))

	jobHandle = C.VixVM_CopyFileFromHostToGuest(
		guest.handle,         // VM handle
		csrc,                 // src name
		cdest,                // dest name
		C.int(0),             // options
		C.VIX_INVALID_HANDLE, // propertyListHandle
		nil,                  // callbackProc
		nil)                  // clientData

	defer C.Vix_ReleaseHandle(jobHandle)

	err = C.vix_job_wait(jobHandle)
	if C.VIX_OK != err {
		return &VixError{
			Operation: "host.CopyFileToGuest",
			Code:      int(err & 0xFFFF),
			Text:      C.GoString(C.Vix_GetErrorText(err, nil)),
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
	"sort"
	"strings"
	"time"

	"github.com/gorilla/mux"
	"github.com/nitro/memberlist"
	"github.com/newrelic/sidecar/catalog"
	"github.com/newrelic/sidecar/output"
	"github.com/newrelic/sidecar/service"
)

func makeHandler(fn func(http.ResponseWriter, *http.Request,
	*memberlist.Memberlist, *catalog.ServicesState),
	list *memberlist.Memberlist, state *catalog.ServicesState) http.HandlerFunc {

	return func(response http.ResponseWriter, req *http.Request) {
		fn(response, req, list, state)
	}
}

func watchHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	defer req.Body.Close()

	response.Header().Set("Content-Type", "application/json")

	lastChange := time.Unix(0, 0)

	for {
		if state.LastChanged.After(lastChange) {
			lastChange = state.LastChanged
			jsonStr, _ := json.Marshal(state.ByService())
			response.Write(jsonStr)
			// In order to flush immediately, we have to cast to a Flusher.
			// The normal HTTP library supports this but not all do, so we
			// check just in case.
			if f, ok := response.(http.Flusher); ok {
				f.Flush()
			}
		}
		time.Sleep(250 * time.Millisecond)
	}
}

func servicesHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	params := mux.Vars(req)

	defer req.Body.Close()

	if params["extension"] == ".json" {
		response.Header().Set("Content-Type", "application/json")
		jsonStr, _ := json.MarshalIndent(state.ByService(), "", "  ")
		response.Write(jsonStr)
		return
	}
}

func serversHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	defer req.Body.Close()

	response.Header().Set("Content-Type", "text/html")
	response.Write(
		[]byte(`

 			<meta http-equiv="refresh" content="4">

` + state.Format(list) + "
"))
}

func stateHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	defer req.Body.Close()

	response.Header().Set("Content-Type", "application/json")
	response.Write(state.Encode())
	return
}

func statusStr(status int) string {
	switch status {
	case 0:
		return "Alive"
	case 1:
		return "Tombstone"
	case 2:
		return "Unhealthy"
	default:
		return "Unknown"
	}
}

func portsStr(svcPorts []service.Port) string {
	var ports []string

	for _, port := range svcPorts {
	    if port.ServicePort != 0 {
			ports = append(ports, fmt.Sprintf("%v->%v", port.ServicePort, port.Port))
		} else {
			ports = append(ports, fmt.Sprintf("%v", port.Port))
		}
	}

	return strings.Join(ports, ", ")
}

type listByName []*memberlist.Node

func (a listByName) Len() int           { return len(a) }
func (a listByName) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a listByName) Less(i, j int) bool { return a[i].Name < a[j].Name }

type Member struct {
	Node    *memberlist.Node
	Updated time.Time
}

func lineWrapMembers(cols int, fields []*Member) [][]*Member {
	if len(fields) < cols {
		return [][]*Member{fields}
	}

	retval := make([][]*Member, len(fields)/cols+1)
	for i := 0; i < len(fields); i++ {
		row := i / cols
		retval[row] = append(retval[row], fields[i])
	}

	return retval
}

func viewHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	timeAgo := func(when time.Time) string { return output.TimeAgo(when, time.Now().UTC()) }

	funcMap := template.FuncMap{
		"statusStr":   statusStr,
		"timeAgo":     timeAgo,
		"portsStr":    portsStr,
		"clusterName": func() string { return list.ClusterName() },
	}

	t, err := template.New("services").Funcs(funcMap).ParseFiles("views/services.html")
	if err != nil {
		println("Error parsing template: " + err.Error())
	}

	members := list.Members()
	sort.Sort(listByName(members))

	compiledMembers := make([]*Member, len(members))
	for i, member := range members {
		if _, ok := state.Servers[member.Name]; ok {
			compiledMembers[i] = &Member{member, state.Servers[member.Name].LastUpdated}
		} else {
			compiledMembers[i] = &Member{Node: member}
			println("No updated time for " + member.Name)
		}
	}

	wrappedMembers := lineWrapMembers(5, compiledMembers)

	viewData := struct {
		Services map[string][]*service.Service
		Members  [][]*Member
	}{
		Services: state.ByService(),
		Members:  wrappedMembers,
	}

	t.ExecuteTemplate(response, "services.html", viewData)
}

func serveHttp(list *memberlist.Memberlist, state *catalog.ServicesState) {
	router := mux.NewRouter()

	router.HandleFunc(
		"/services{extension}", makeHandler(servicesHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/servers", makeHandler(serversHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/services", makeHandler(viewHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/state", makeHandler(stateHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/watch", makeHandler(watchHandler, list, state),
	).Methods("GET")

	fs := http.FileServer(http.Dir("views/static/"))

	router.Handle("/static/{file}", http.StripPrefix("/static/", fs))

	http.Handle("/", router)

	err := http.ListenAndServe("0.0.0.0:7777", nil)
	exitWithError(err, "Can't start HTTP server")
}
<commit_msg>Fix logging levels in http.go.<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
	"sort"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/gorilla/mux"
	"github.com/nitro/memberlist"
	"github.com/newrelic/sidecar/catalog"
	"github.com/newrelic/sidecar/output"
	"github.com/newrelic/sidecar/service"
)

func makeHandler(fn func(http.ResponseWriter, *http.Request,
	*memberlist.Memberlist, *catalog.ServicesState),
	list *memberlist.Memberlist, state *catalog.ServicesState) http.HandlerFunc {

	return func(response http.ResponseWriter, req *http.Request) {
		fn(response, req, list, state)
	}
}

func watchHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	defer req.Body.Close()

	response.Header().Set("Content-Type", "application/json")

	lastChange := time.Unix(0, 0)

	for {
		if state.LastChanged.After(lastChange) {
			lastChange = state.LastChanged
			jsonStr, _ := json.Marshal(state.ByService())
			response.Write(jsonStr)
			// In order to flush immediately, we have to cast to a Flusher.
			// The normal HTTP library supports this but not all do, so we
			// check just in case.
			if f, ok := response.(http.Flusher); ok {
				f.Flush()
			}
		}
		time.Sleep(250 * time.Millisecond)
	}
}

func servicesHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	params := mux.Vars(req)

	defer req.Body.Close()

	if params["extension"] == ".json" {
		response.Header().Set("Content-Type", "application/json")
		jsonStr, _ := json.MarshalIndent(state.ByService(), "", "  ")
		response.Write(jsonStr)
		return
	}
}

func serversHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	defer req.Body.Close()

	response.Header().Set("Content-Type", "text/html")
	response.Write(
		[]byte(`

 			<meta http-equiv="refresh" content="4">

` + state.Format(list) + "
"))
}

func stateHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	defer req.Body.Close()

	response.Header().Set("Content-Type", "application/json")
	response.Write(state.Encode())
	return
}

func statusStr(status int) string {
	switch status {
	case 0:
		return "Alive"
	case 1:
		return "Tombstone"
	case 2:
		return "Unhealthy"
	default:
		return "Unknown"
	}
}

func portsStr(svcPorts []service.Port) string {
	var ports []string

	for _, port := range svcPorts {
	    if port.ServicePort != 0 {
			ports = append(ports, fmt.Sprintf("%v->%v", port.ServicePort, port.Port))
		} else {
			ports = append(ports, fmt.Sprintf("%v", port.Port))
		}
	}

	return strings.Join(ports, ", ")
}

type listByName []*memberlist.Node

func (a listByName) Len() int           { return len(a) }
func (a listByName) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a listByName) Less(i, j int) bool { return a[i].Name < a[j].Name }

type Member struct {
	Node    *memberlist.Node
	Updated time.Time
}

func lineWrapMembers(cols int, fields []*Member) [][]*Member {
	if len(fields) < cols {
		return [][]*Member{fields}
	}

	retval := make([][]*Member, len(fields)/cols+1)
	for i := 0; i < len(fields); i++ {
		row := i / cols
		retval[row] = append(retval[row], fields[i])
	}

	return retval
}

func viewHandler(response http.ResponseWriter, req *http.Request, list *memberlist.Memberlist, state *catalog.ServicesState) {
	timeAgo := func(when time.Time) string { return output.TimeAgo(when, time.Now().UTC()) }

	funcMap := template.FuncMap{
		"statusStr":   statusStr,
		"timeAgo":     timeAgo,
		"portsStr":    portsStr,
		"clusterName": func() string { return list.ClusterName() },
	}

	t, err := template.New("services").Funcs(funcMap).ParseFiles("views/services.html")
	if err != nil {
		log.Errorf("Error parsing template: %s", err.Error())
	}

	members := list.Members()
	sort.Sort(listByName(members))

	compiledMembers := make([]*Member, len(members))
	for i, member := range members {
		if _, ok := state.Servers[member.Name]; ok {
			compiledMembers[i] = &Member{member, state.Servers[member.Name].LastUpdated}
		} else {
			compiledMembers[i] = &Member{Node: member}
			log.Debug("No updated time for " + member.Name)
		}
	}

	wrappedMembers := lineWrapMembers(5, compiledMembers)

	viewData := struct {
		Services map[string][]*service.Service
		Members  [][]*Member
	}{
		Services: state.ByService(),
		Members:  wrappedMembers,
	}

	t.ExecuteTemplate(response, "services.html", viewData)
}

func serveHttp(list *memberlist.Memberlist, state *catalog.ServicesState) {
	router := mux.NewRouter()

	router.HandleFunc(
		"/services{extension}", makeHandler(servicesHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/servers", makeHandler(serversHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/services", makeHandler(viewHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/state", makeHandler(stateHandler, list, state),
	).Methods("GET")

	router.HandleFunc(
		"/watch", makeHandler(watchHandler, list, state),
	).Methods("GET")

	fs := http.FileServer(http.Dir("views/static/"))

	router.Handle("/static/{file}", http.StripPrefix("/static/", fs))

	http.Handle("/", router)

	err := http.ListenAndServe("0.0.0.0:7777", nil)
	exitWithError(err, "Can't start HTTP server")
}
<|endoftext|>


<commit_before>package yiigo

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"time"
)

type httpConf struct {
	ConnTimeout         int `toml:"connTimeout"`
	KeepAlive           int `toml:"keepAlive"`
	MaxConnsPerHost     int `toml:"maxConnsPerHost"`
	MaxIdleConnsPerHost int `toml:"maxIdleConnsPerHost"`
	MaxIdleConns        int `toml:"maxIdleConns"`
	IdleConnTimeout     int `toml:"idleConnTimeout"`
}

// httpClient HTTP request client
var httpClient *http.Client

func initHTTPClient() {
	conf := &httpConf{
		ConnTimeout:         30,
		KeepAlive:           60,
		MaxConnsPerHost:     200,
		MaxIdleConnsPerHost: 100,
		MaxIdleConns:        100,
		IdleConnTimeout:     60,
	}

	Env.Unmarshal("http", conf)

	httpClient = &http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			DialContext: (&net.Dialer{
				Timeout:   time.Duration(conf.ConnTimeout) * time.Second,
				KeepAlive: time.Duration(conf.KeepAlive) * time.Second,
				DualStack: true,
			}).DialContext,
			MaxConnsPerHost:       conf.MaxConnsPerHost,
			MaxIdleConnsPerHost:   conf.MaxIdleConnsPerHost,
			MaxIdleConns:          conf.MaxIdleConns,
			IdleConnTimeout:       time.Duration(conf.IdleConnTimeout) * time.Second,
			TLSHandshakeTimeout:   10 * time.Second,
			ExpectContinueTimeout: 1 * time.Second,
		},
		Timeout: 10 * time.Second,
	}
}

// HTTPGet http get request
func HTTPGet(url string, headers map[string]string, timeout ...time.Duration) ([]byte, error) {
	req, err := http.NewRequest("GET", url, nil)

	if err != nil {
		return nil, err
	}

	// custom headers
	if len(headers) != 0 {
		for k, v := range headers {
			req.Header.Set(k, v)
		}
	}

	// custom timeout
	if len(timeout) > 0 {
		httpClient.Timeout = timeout[0]
	}

	resp, err := httpClient.Do(req)

	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		io.Copy(ioutil.Discard, resp.Body)

		return nil, fmt.Errorf("error http code: %d", resp.StatusCode)
	}

	b, err := ioutil.ReadAll(resp.Body)

	if err != nil {
		return nil, err
	}

	return b, nil
}

// HTTPPost http post request, default content-type is 'application/json'
func HTTPPost(url string, body []byte, headers map[string]string, timeout ...time.Duration) ([]byte, error) {
	reader := bytes.NewReader(body)

	req, err := http.NewRequest("POST", url, reader)

	if err != nil {
		return nil, err
	}

	req.Header.Set("Content-Type", "application/json; charset=utf-8")

	// custom headers
	if len(headers) != 0 {
		for k, v := range headers {
			req.Header.Set(k, v)
		}
	}

	// custom timeout
	if len(timeout) > 0 {
		httpClient.Timeout = timeout[0]
	}

	resp, err := httpClient.Do(req)

	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		io.Copy(ioutil.Discard, resp.Body)

		return nil, fmt.Errorf("error http code: %d", resp.StatusCode)
	}

	b, err := ioutil.ReadAll(resp.Body)

	if err != nil {
		return nil, err
	}

	return b, nil
}
<commit_msg>update http client<commit_after>package yiigo

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"time"
)

type httpConf struct {
	ConnTimeout         int `toml:"connTimeout"`
	KeepAlive           int `toml:"keepAlive"`
	MaxConnsPerHost     int `toml:"maxConnsPerHost"`
	MaxIdleConnsPerHost int `toml:"maxIdleConnsPerHost"`
	MaxIdleConns        int `toml:"maxIdleConns"`
	IdleConnTimeout     int `toml:"idleConnTimeout"`
}

// httpClient HTTP request client
var httpClient *http.Client

func initHTTPClient() {
	conf := &httpConf{
		ConnTimeout:         30,
		KeepAlive:           60,
		MaxConnsPerHost:     200,
		MaxIdleConnsPerHost: 100,
		MaxIdleConns:        100,
		IdleConnTimeout:     60,
	}

	Env.Unmarshal("http", conf)

	httpClient = &http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			DialContext: (&net.Dialer{
				Timeout:   time.Duration(conf.ConnTimeout) * time.Second,
				KeepAlive: time.Duration(conf.KeepAlive) * time.Second,
				DualStack: true,
			}).DialContext,
			MaxConnsPerHost:       conf.MaxConnsPerHost,
			MaxIdleConnsPerHost:   conf.MaxIdleConnsPerHost,
			MaxIdleConns:          conf.MaxIdleConns,
			IdleConnTimeout:       time.Duration(conf.IdleConnTimeout) * time.Second,
			TLSHandshakeTimeout:   10 * time.Second,
			ExpectContinueTimeout: 1 * time.Second,
		},
		Timeout: 10 * time.Second,
	}
}

// HTTPGet http get request
func HTTPGet(url string, headers map[string]string, timeout ...time.Duration) ([]byte, error) {
	req, err := http.NewRequest("GET", url, nil)

	if err != nil {
		return nil, err
	}

	// custom headers
	if len(headers) != 0 {
		for k, v := range headers {
			req.Header.Set(k, v)
		}
	}

	// custom timeout
	if len(timeout) > 0 {
		httpClient.Timeout = timeout[0]
	}

	resp, err := httpClient.Do(req)

	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		io.Copy(ioutil.Discard, resp.Body)

		return nil, fmt.Errorf("error http code: %d", resp.StatusCode)
	}

	b, err := ioutil.ReadAll(resp.Body)

	if err != nil {
		return nil, err
	}

	return b, nil
}

// HTTPPost http post request, default content-type is 'application/json'.
func HTTPPost(url string, body []byte, headers map[string]string, timeout ...time.Duration) ([]byte, error) {
	reader := bytes.NewReader(body)

	req, err := http.NewRequest("POST", url, reader)

	if err != nil {
		return nil, err
	}

	req.Header.Set("Content-Type", "application/json; charset=utf-8")

	// custom headers
	if len(headers) != 0 {
		for k, v := range headers {
			req.Header.Set(k, v)
		}
	}

	// custom timeout
	if len(timeout) > 0 {
		httpClient.Timeout = timeout[0]
	}

	resp, err := httpClient.Do(req)

	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		io.Copy(ioutil.Discard, resp.Body)

		return nil, fmt.Errorf("error http code: %d", resp.StatusCode)
	}

	b, err := ioutil.ReadAll(resp.Body)

	if err != nil {
		return nil, err
	}

	return b, nil
}
<|endoftext|>